MLIR Linalg short form format is used when it should not be #117528
Description
Both linalg.map and linalg.reduce sometimes print using "short" form even though the payload is not in short form. This means that writing the op, and then reading it back, will result in an op that does something completely different from the original operation.
The reason is that findPayloadOp() does not verify that the YieldOp actually yields the result(s) of the payload operation.
So, the following (nonsensical) linalg.map will, incorrectly, print in short form, which, when read back will do something else than the original.
$ cat ~/llvm-builds/linalg-map-bug/bug.mlir
"builtin.module"() ({
"func.func"() <{function_type = (tensor<1x32xf32>, tensor<1x32xf32>) -> tensor<1x32xf32>, sym_name = "map_maximumf"}> ({
^bb0(%arg0: tensor<1x32xf32>, %arg1: tensor<1x32xf32>):
%0 = "tensor.empty"() : () -> tensor<1x32xf32>
%1 = "linalg.map"(%arg0, %arg1, %0) ({
^bb0(%arg2: f32, %arg3: f32):
%2 = "arith.maximumf"(%arg2, %arg3) <{fastmath = #arith.fastmath<none>}> : (f32, f32) -> f32
// NOTE: This does not yield the output of arith.maximumf
"linalg.yield"(%arg2) : (f32) -> ()
}) : (tensor<1x32xf32>, tensor<1x32xf32>, tensor<1x32xf32>) -> tensor<1x32xf32>
"func.return"(%1) : (tensor<1x32xf32>) -> ()
}) : () -> ()
}) : () -> ()
The default printing gives the incorrect:
$ mlir-opt -- ~/llvm-builds/linalg-map-bug/bug.mlir
module {
func.func @map_maximumf(%arg0: tensor<1x32xf32>, %arg1: tensor<1x32xf32>) -> tensor<1x32xf32> {
%0 = tensor.empty() : tensor<1x32xf32>
%mapped = linalg.map { arith.maximumf } ins(%arg0, %arg1 : tensor<1x32xf32>, tensor<1x32xf32>) outs(%0 : tensor<1x32xf32>)
return %mapped : tensor<1x32xf32>
}
}
and round-tripping through the default printer does not bring back the original operation:
$ mlir-opt -- ~/llvm-builds/linalg-map-bug/bug.mlir | mlir-opt --mlir-print-op-generic
"builtin.module"() ({
"func.func"() <{function_type = (tensor<1x32xf32>, tensor<1x32xf32>) -> tensor<1x32xf32>, sym_name = "map_maximumf"}> ({
^bb0(%arg0: tensor<1x32xf32>, %arg1: tensor<1x32xf32>):
%0 = "tensor.empty"() : () -> tensor<1x32xf32>
%1 = "linalg.map"(%arg0, %arg1, %0) ({
^bb0(%arg2: f32, %arg3: f32):
%2 = "arith.maximumf"(%arg2, %arg3) <{fastmath = #arith.fastmath<none>}> : (f32, f32) -> f32
"linalg.yield"(%2) : (f32) -> ()
}) : (tensor<1x32xf32>, tensor<1x32xf32>, tensor<1x32xf32>) -> tensor<1x32xf32>
"func.return"(%1) : (tensor<1x32xf32>) -> ()
}) : () -> ()
}) : () -> ()
Note that the resulting payload block now yields the result of the arith.maximumf, which the original operation did not.
The same problem most likely affects linalg.reduce, which also uses findPayloadOp() in the same way.
Also, the documentation for findPayloadOp() says "If initFirst flag is enabled, we check that init takes the first position in operands of payload." but this does not correspond to what the code does (the code does special things with the first argument of the block and the last operand of the payload operation).
This is llvm-project at:
commit 73bebf96bc21dcc01a8eccfb4ece200c1c665931 (HEAD -> main, origin/main, origin/HEAD)
Date: Mon Nov 25 14:47:50 2024 +0800