| // softmax(matmul(arg0, arg1)) |
| func.func @main(%arg0: tensor<?x16xf32>, %arg1: tensor<16x16xf32>) |
| -> tensor<?x16xf32> { |
| %c0 = arith.constant 0 : index |
| %m = tensor.dim %arg0, %c0 : tensor<?x16xf32> |
| %empty = tensor.empty(%m) : tensor<?x16xf32> |
| %zero = arith.constant 0.0 : f32 |
| %init = linalg.fill ins(%zero : f32) outs(%empty : tensor<?x16xf32>) |
| -> tensor<?x16xf32> |
| %mat = linalg.matmul ins(%arg0, %arg1 : tensor<?x16xf32>, tensor<16x16xf32>) |
| outs(%init : tensor<?x16xf32>) -> tensor<?x16xf32> |
| %out = linalg.softmax dimension(1) ins(%mat : tensor<?x16xf32>) |
| outs(%empty : tensor<?x16xf32>) -> tensor<?x16xf32> |
| return %out : tensor<?x16xf32> |
| } |