[VectorLayout] Fix insertion of new constOp for non dominate issue. (#18894)
Main motivation of this patch is to resolve issue where we have the same
constOp being used by multiple operations.
But with a twist where first time the constOp needs a layout is on a op
that topologically comes after other ops that use constOp. This will
generate a copy of constOp in the location right before the latter op,
which is problematic because this constOp will be used by other ops
before it.
Previously for the test added we get this error:
```
within split at contraction_layout.mlir:1 offset :24:10: note: see current operation: %9 = "arith.addf"(%8, %6) <{fastmath = #arith.fastmath<none>}> : (vector<96x64xf16>, vector<96x64xf16>) -> vector<96x64xf16>
within split at contraction_layout.mlir:1 offset :22:19: error: operand #1 does not dominate this use
%scaled_rhs = arith.mulf %read_1, %cst_1 : vector<96x64xf16>
```
While minor, this is also problematic because this error seem to stopped
layout analysis (but not fatally) S.T it fails to vector distribute in
some cases, making it hard to debug.
Signed-off-by: Stanley Winata <stanley.winata@amd.com>diff --git a/compiler/src/iree/compiler/Codegen/Common/GPU/test/gpu_vector_distribution.mlir b/compiler/src/iree/compiler/Codegen/Common/GPU/test/gpu_vector_distribution.mlir
index da40806..cf47ca9 100644
--- a/compiler/src/iree/compiler/Codegen/Common/GPU/test/gpu_vector_distribution.mlir
+++ b/compiler/src/iree/compiler/Codegen/Common/GPU/test/gpu_vector_distribution.mlir
@@ -666,10 +666,10 @@
}
// CHECK-LABEL: func.func @resolve_constant_with_multiple_layout_uses
// CHECK-SAME: (%[[ARG0:.+]]: vector<64x64xf16>, %[[ARG0:.+]]: vector<64x64xf16>)
-// CHECK: %[[V0:.+]] = arith.constant dense<0.000000e+00> : vector<2x2x16xf16>
-// CHECK: %[[V1:.+]] = arith.constant dense<0.000000e+00> : vector<2x2x8xf16>
-// CHECK: %[[ADD0:.+]] = arith.addf %{{.+}}, %[[V1]]{{.*}} : vector<2x2x8xf16>
-// CHECK: %[[ADD1:.+]] = arith.addf %{{.+}}, %[[V0]]{{.*}} : vector<2x2x16xf16>
+// CHECK: %[[V0:.+]] = arith.constant dense<0.000000e+00> : vector<2x2x8xf16>
+// CHECK: %[[V1:.+]] = arith.constant dense<0.000000e+00> : vector<2x2x16xf16>
+// CHECK: %[[ADD0:.+]] = arith.addf %{{.+}}, %[[V0]]{{.*}} : vector<2x2x8xf16>
+// CHECK: %[[ADD1:.+]] = arith.addf %{{.+}}, %[[V1]]{{.*}} : vector<2x2x16xf16>
// CHECK: arith.addf %{{.+}}, %[[ADD0]]{{.*}} : vector<2x2x8xf16>
transform.named_sequence @__transform_main(%variant_op: !transform.any_op {transform.readonly}) {
diff --git a/compiler/src/iree/compiler/Codegen/Common/VectorLayoutAnalysis.cpp b/compiler/src/iree/compiler/Codegen/Common/VectorLayoutAnalysis.cpp
index fa3786c..deabc58 100644
--- a/compiler/src/iree/compiler/Codegen/Common/VectorLayoutAnalysis.cpp
+++ b/compiler/src/iree/compiler/Codegen/Common/VectorLayoutAnalysis.cpp
@@ -205,6 +205,7 @@
if (!opOperand.get().hasOneUse() && !vectorLayout &&
llvm::dyn_cast_or_null<arith::ConstantOp>(
opOperand.get().getDefiningOp())) {
+ builder.setInsertionPoint(opOperand.get().getDefiningOp());
Operation *copiedConstOp = builder.clone(*opOperand.get().getDefiningOp());
Value copiedConst = copiedConstOp->getResult(0);
builder.replaceAllUsesExcept(opOperand.get(), copiedConst,
diff --git a/compiler/src/iree/compiler/Codegen/Common/test/vector_layout_analysis.mlir b/compiler/src/iree/compiler/Codegen/Common/test/vector_layout_analysis.mlir
index 09c4d27..28e8ab0 100644
--- a/compiler/src/iree/compiler/Codegen/Common/test/vector_layout_analysis.mlir
+++ b/compiler/src/iree/compiler/Codegen/Common/test/vector_layout_analysis.mlir
@@ -464,6 +464,47 @@
// -----
+#contract_layout = #iree_vector_ext.nested_layout<
+ subgroup_tile = [1, 1],
+ batch_tile = [3, 2],
+ outer_tile = [4, 1],
+ thread_tile = [2, 32],
+ element_tile = [4, 1],
+
+ subgroup_strides = [0, 0],
+ thread_strides = [32, 1]
+>
+
+// This test ensures that we are not running into ops not dominating constantOp operands after layout analysis.
+// We simulate that by doing elmentwise op on the value with "layout" i.e scaled_lhs after scaled_rhs.
+// If not handled properly, will generate constOp before "scaled_lhs", but would get used also by "scaled_rhs".
+builtin.module attributes { transform.with_named_sequence } {
+ func.func @handle_multiuse_constant(%lhs: vector<96x64xf16>, %rhs: vector<96x64xf16>, %arr: memref<96x64xf16>) -> () {
+ %c0 = arith.constant 0 : index
+ %cst = arith.constant 0.000000e+00 : f16
+ %cst_1 = arith.constant dense<1.562500e-02> : vector<96x64xf16>
+ // expected-remark @above {{thread_strides = [32, 1]}}
+ %lhs_layout = iree_vector_ext.to_layout %lhs to layout(#contract_layout) : vector<96x64xf16>
+
+ %scaled_rhs = arith.mulf %rhs, %cst_1 : vector<96x64xf16>
+ // expected-remark @above {{thread_strides = [32, 1]}}
+ %scaled_lhs = arith.mulf %lhs_layout, %cst_1 : vector<96x64xf16>
+ // expected-remark @above {{thread_strides = [32, 1]}}
+ %add = arith.addf %scaled_lhs, %scaled_rhs : vector<96x64xf16>
+ // expected-remark @above {{thread_strides = [32, 1]}}
+ vector.transfer_write %add, %arr[%c0, %c0] {in_bounds = [true, true]} : vector<96x64xf16>, memref<96x64xf16>
+ func.return
+ }
+
+ transform.named_sequence @__transform_main(%variant_op: !transform.any_op {transform.readonly}) {
+ %top_level_func = transform.structured.match ops{["func.func"]} in %variant_op : (!transform.any_op) -> !transform.any_op
+ transform.iree.test_vector_layout_analysis %top_level_func : !transform.any_op
+ transform.yield
+ }
+}
+
+// -----
+
#layout = #iree_vector_ext.nested_layout<
subgroup_tile = [2, 1, 1],
batch_tile = [1, 2, 4],
