[VectorLayout] Fix insertion of new constOp for non dominate issue. (i…

…ree-org#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 <[email protected]>

compiler/src/iree/compiler/Codegen/Common/GPU/test/gpu_vector_distribution.mlir

@@ -666,10 +666,10 @@ builtin.module attributes { transform.with_named_sequence } {
   }
 // 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}) {

compiler/src/iree/compiler/Codegen/Common/VectorLayoutAnalysis.cpp

@@ -205,6 +205,7 @@ ChangeResult DistributionLayout::resolveWithPossibleConflict(
   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,

compiler/src/iree/compiler/Codegen/Common/test/vector_layout_analysis.mlir

@@ -464,6 +464,47 @@ builtin.module attributes { transform.with_named_sequence } {
 
 // -----
 
+#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],