// 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>
}