[mlir] Set insertion point of vector constant to the top of the vectorized loop body

[ayzhuang]

Summary

When we vectorize a scalar constant, the vector constant is inserted before its first user if the scalar constant is defined outside the loops to be vectorized. It is possible that the vector constant does not dominate all its users. To fix the problem, we find the innermost vectorized loop that encloses that first user and insert the vector constant at the top of the loop body.

[vinayaka-polymage]

Thank you, changes look good. Can you please add a brief summary for the fix (unless you think it is generally obvious) ?

Without this fix, the output for the example would be (violating SSA dominance for %cst_0) :

mlir-opt  -affine-super-vectorize="virtual-vector-size=128 test-fastest-varying=0" -split-input-file ../mlir/test/Dialect/Affine/input.mlir -verify-each=0
module  {
  func @vec_constant_with_two_users(%arg0: index, %arg1: index) -> (f32, f32) {
    %0 = memref.alloc(%arg0, %arg1) : memref<?x?xf32>
    %1 = memref.alloc(%arg0) : memref<?xf32>
    %cst = constant 1.000000e+00 : f32
    affine.for %arg2 = 0 to %arg0 step 128 {
      affine.for %arg3 = 0 to %arg1 {
        %cst_0 = constant dense<1.000000e+00> : vector<128xf32>
        vector.transfer_write %cst_0, %0[%arg3, %arg2] : vector<128xf32>, memref<?x?xf32>
      }
      vector.transfer_write %cst_0, %1[%arg2] : vector<128xf32>, memref<?xf32>
    }
    %c12 = constant 12 : index
    %2 = affine.load %0[%c12, %c12] : memref<?x?xf32>
    %3 = affine.load %1[%c12] : memref<?xf32>
    return %2, %3 : f32, f32
  }
}

[ayzhuang]

@vinayaka-polymage Thanks for the review. I have added a brief summary.

[nicolasvasilache]

thanks for fixing this!

[ayzhuang]

@nicolasvasilache Thank you for the review. We'll merge it tomorrow if there are no more comments.

mlir/lib/Dialect/Affine/Transforms/SuperVectorize.cpp

 static ConstantOp vectorizeConstant(ConstantOp constOp,
                                     VectorizationState &state) {
   Type scalarTy = constOp.getType();
   if (!VectorType::isValidElementType(scalarTy))
     return nullptr;
 
   auto vecTy = getVectorType(scalarTy, state.strategy);
   auto vecAttr = DenseElementsAttr::get(vecTy, constOp.getValue());
+
+  OpBuilder::InsertionGuard guard(state.builder);
+  Operation *parentOp = state.builder.getInsertionBlock()->getParentOp();
+  // Find the innermost vectorized ancestor loop to insert the vector constant.
+  while (parentOp && !state.vecLoopToVecDim.count(parentOp))
+    parentOp = parentOp->getParentOp();
+  assert(parentOp && state.vecLoopToVecDim.count(parentOp) &&
+         isa<AffineForOp>(parentOp) && "Expected a vectorized for op");
+  auto vecForOp = cast<AffineForOp>(parentOp);
+  state.builder.setInsertionPointToStart(vecForOp.getBody());
   auto newConstOp = state.builder.create<ConstantOp>(constOp.getLoc(), vecAttr);
 
   // Register vector replacement for future uses in the scope.
   state.registerOpVectorReplacement(constOp, newConstOp);
   return newConstOp;
 }
 
 /// Creates a constant vector filled with the neutral elements of the given

mlir/test/Dialect/Affine/SuperVectorize/vectorize_1d.mlir

affine.for %i3 = 0 to %N {
       // CHECK: %[[C3:.*]] = constant dense<2.000000e+00> : vector<128xf32>
       // CHECK: vector.transfer_write %[[C3]], {{.*}} : vector<128xf32>, memref<?x?xf32>
       // non-scoped %f2
       affine.store %f2, %B[%i2, %i3] : memref<?x?xf32, 0>
     }
   }
   affine.for %i4 = 0 to %M {
     affine.for %i5 = 0 to %N {
+      // CHECK: %[[SPLAT2:.*]] = constant dense<2.000000e+00> : vector<128xf32>
+      // CHECK: %[[SPLAT1:.*]] = constant dense<1.000000e+00> : vector<128xf32>
       // CHECK: %[[A5:.*]] = vector.transfer_read %{{.*}}[{{.*}}], %{{[a-zA-Z0-9_]*}} : memref<?x?xf32>, vector<128xf32>
       // CHECK: %[[B5:.*]] = vector.transfer_read %{{.*}}[{{.*}}], %{{[a-zA-Z0-9_]*}} : memref<?x?xf32>, vector<128xf32>
       // CHECK: %[[S5:.*]] = addf %[[A5]], %[[B5]] : vector<128xf32>
-      // CHECK: %[[SPLAT1:.*]] = constant dense<1.000000e+00> : vector<128xf32>
       // CHECK: %[[S6:.*]] = addf %[[S5]], %[[SPLAT1]] : vector<128xf32>
-      // CHECK: %[[SPLAT2:.*]] = constant dense<2.000000e+00> : vector<128xf32>
       // CHECK: %[[S7:.*]] = addf %[[S5]], %[[SPLAT2]] : vector<128xf32>
       // CHECK: %[[S8:.*]] = addf %[[S7]], %[[S6]] : vector<128xf32>
       // CHECK: vector.transfer_write %[[S8]], {{.*}} : vector<128xf32>, memref<?x?xf32>
       %a5 = affine.load %A[%i4, %i5] : memref<?x?xf32, 0>
       %b5 = affine.load %B[%i4, %i5] : memref<?x?xf32, 0>
       %s5 = addf %a5, %b5 : f32
       // non-scoped %f1
       %s6 = addf %s5, %f1 : f32
       // non-scoped %f2
       %s7 = addf %s5, %f2 : f32
       // diamond dependency.
       %s8 = addf %s7, %s6 : f32
       affine.store %s8, %C[%i4, %i5] : memref<?x?xf32, 0>
     }
   }
   %c7 = constant 7 : index
   %c42 = constant 42 : index
   %res = affine.load %C[%c7, %c42] : memref<?x?xf32, 0>
   return %res : f32
 }
 
 // -----
 
+// CHECK-LABEL: func @vec_constant_with_two_users
+func @vec_constant_with_two_users(%M : index, %N : index) -> (f32, f32) {
+  %A = memref.alloc (%M, %N) : memref<?x?xf32, 0>
+  %B = memref.alloc (%M) : memref<?xf32, 0>
+  %f1 = constant 1.0 : f32
+  affine.for %i0 = 0 to %M { // vectorized
+    // CHECK:      %[[C1:.*]] = constant dense<1.000000e+00> : vector<128xf32>
+    // CHECK-NEXT: affine.for
+    // CHECK-NEXT:   vector.transfer_write %[[C1]], {{.*}} : vector<128xf32>, memref<?x?xf32>
+    affine.for %i1 = 0 to %N {
+      affine.store %f1, %A[%i1, %i0] : memref<?x?xf32, 0>
+    }
+    // CHECK: vector.transfer_write %[[C1]], {{.*}} : vector<128xf32>, memref<?xf32>
+    affine.store %f1, %B[%i0] : memref<?xf32, 0>
+  }
+  %c12 = constant 12 : index
+  %res1 = affine.load %A[%c12, %c12] : memref<?x?xf32, 0>
+  %res2 = affine.load %B[%c12] : memref<?xf32, 0>
+  return %res1, %res2 : f32, f32
+}
+
+// -----
+
 // CHECK-LABEL: func @vec_rejected_1
 func @vec_rejected_1(%A : memref<?x?xf32>, %B : memref<?x?x?xf32>) {
 // CHECK-DAG: %[[C0:.*]] = constant 0 : index
 // CHECK-DAG: %[[C1:.*]] = constant 1 : index
 // CHECK-DAG: %[[C2:.*]] = constant 2 : index
 // CHECK-DAG: [[ARG_M:%[0-9]+]] = memref.dim %{{.*}}, %[[C0]] : memref<?x?xf32>
 // CHECK-DAG: [[ARG_N:%[0-9]+]] = memref.dim %{{.*}}, %[[C1]] : memref<?x?xf32>
 // CHECK-DAG: [[ARG_P:%[0-9]+]] = memref.dim %{{.*}}, %[[C2]] : memref<?x?x?xf32>
affine.for %i = 0 to 256 {
    affine.store %red0, %out0[%i] : memref<256xf32>
    affine.store %red1, %out1[%i] : memref<256xi32>
  }
  return
 }
 
 // CHECK-LABEL: @vec_non_vecdim_reductions
 // CHECK:       affine.for %{{.*}} = 0 to 256 step 128 {
-// CHECK:         %[[vzero:.*]] = constant dense<0.000000e+00> : vector<128xf32>
 // CHECK:         %[[vone:.*]] = constant dense<1> : vector<128xi32>
+// CHECK:         %[[vzero:.*]] = constant dense<0.000000e+00> : vector<128xf32>
 // CHECK:         %[[reds:.*]]:2 = affine.for %{{.*}} = 0 to 128
 // CHECK-SAME:      iter_args(%[[red_iter0:.*]] = %[[vzero]], %[[red_iter1:.*]] = %[[vone]]) -> (vector<128xf32>, vector<128xi32>) {
 // CHECK:           %[[ld0:.*]] = vector.transfer_read %{{.*}} : memref<128x256xf32>, vector<128xf32>
 // CHECK:           %[[add:.*]] = addf %[[red_iter0]], %[[ld0]] : vector<128xf32>
 // CHECK:           %[[ld1:.*]] = vector.transfer_read %{{.*}} : memref<128x256xi32>, vector<128xi32>
 // CHECK:           %[[mul:.*]] = muli %[[red_iter1]], %[[ld1]] : vector<128xi32>
 // CHECK:           affine.yield %[[add]], %[[mul]] : vector<128xf32>, vector<128xi32>
 // CHECK:         }

mlir/test/Dialect/Affine/SuperVectorize/vectorize_2d.mlir

affine.for %i3 = 0 to %N {
       // CHECK: [[C3:%.*]] = constant dense<2.000000e+00> : vector<32x256xf32>
       // CHECK: vector.transfer_write [[C3]], {{.*}}  : vector<32x256xf32>, memref<?x?xf32>
       // non-scoped %f2
       affine.store %f2, %B[%i2, %i3] : memref<?x?xf32, 0>
     }
   }
   affine.for %i4 = 0 to %M {
     affine.for %i5 = 0 to %N {
+      // CHECK: [[SPLAT2:%.*]] = constant dense<2.000000e+00> : vector<32x256xf32>
+      // CHECK: [[SPLAT1:%.*]] = constant dense<1.000000e+00> : vector<32x256xf32>
       // CHECK: [[A5:%.*]] = vector.transfer_read %{{.*}}[{{.*}}], %{{.*}} : memref<?x?xf32>, vector<32x256xf32>
       // CHECK: [[B5:%.*]] = vector.transfer_read %{{.*}}[{{.*}}], %{{.*}} : memref<?x?xf32>, vector<32x256xf32>
       // CHECK: [[S5:%.*]] = addf [[A5]], [[B5]] : vector<32x256xf32>
-      // CHECK: [[SPLAT1:%.*]] = constant dense<1.000000e+00> : vector<32x256xf32>
       // CHECK: [[S6:%.*]] = addf [[S5]], [[SPLAT1]] : vector<32x256xf32>
-      // CHECK: [[SPLAT2:%.*]] = constant dense<2.000000e+00> : vector<32x256xf32>
       // CHECK: [[S7:%.*]] = addf [[S5]], [[SPLAT2]] : vector<32x256xf32>
       // CHECK: [[S8:%.*]] = addf [[S7]], [[S6]] : vector<32x256xf32>
       // CHECK: vector.transfer_write [[S8]], {{.*}} : vector<32x256xf32>, memref<?x?xf32>
       //
       %a5 = affine.load %A[%i4, %i5] : memref<?x?xf32, 0>
       %b5 = affine.load %B[%i4, %i5] : memref<?x?xf32, 0>
       %s5 = addf %a5, %b5 : f32
       // non-scoped %f1