[mlir] Vectorize induction variables of loops not mapped to a vector dimension

[ayzhuang]

Summary

1. Add support to vectorize induction variables of loops that are not mapped to any vector dimension in SuperVectorize pass.
2. Fix a bug in getForInductionVarOwner.

[dcaballe]

Thanks for adding support for this missing case, Amy! The implementation looks great! I only think we shouldn't use a special case for uniform IVs, they are just like any other uniform value. I would extend the uniform utilities to support this case instead of adding a new one. More details inline. Other than that, it looks great!

Thanks!
Diego

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

260

Instead of adding a new bullet, I would extend the previous ones about "Uniform operands". Uniform IVs are just uniform values.

1117

I think the vectorizeIV utility shouldn't be needed. We should be able to use vectorizeUniform for this case.

1125

I think this comment wouldn't be necessary. Uniform IVs are just uniform values so they would fall into #3. It just happen that some cases are missing in the uniform check, which is what you are adding here.

1172

I think we should move this check into isUniformDefinition. There shouldn't be any difference between vectorizing a uniform IV and any other uniform value. Actually, if the IV is outside the vectorize loop(s), it should be hitting that case already (line 1161), right? This extra check would handle uniform IVs inside the vectorized loop(s). Am I missing something here?

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

198

nit: we can drop %{{.*}} =

200

Since you are capturing the IVs, it would interesting to check that the transfer read is actually using the scalar version of the IVs

211–217

Since this is only testing the vectorization of IVs, could we simplify the test and remove all the ops not related to the IVs, including iter_args?

[ayzhuang]

Address review comments to extend the uniform utilities to support the induction variable vectorization.

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

1172

Yes, you're right. I have moved this check into isUniformDefinition.

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

200

This transfer read is using the result of affine.apply. I have added checks to see if affine.apply is using the scalar version of the IVs.

[dcaballe]

Thanks, Amy! LGTM. Just a couple of nits. Please, fix them before landing.

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

1073

nit: you could do loopToVectorDim.count(forOp) == 0

1084

nit: spell out auto, please

[ayzhuang]

Address review comments.

[ayzhuang]

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

mlir/lib/Dialect/Affine/IR/AffineOps.cpp

 
 /// Returns the loop parent of an induction variable. If the provided value is
 /// not an induction variable, then return nullptr.
 AffineForOp mlir::getForInductionVarOwner(Value val) {
   auto ivArg = val.dyn_cast<BlockArgument>();
   if (!ivArg || !ivArg.getOwner())
     return AffineForOp();
   auto *containingInst = ivArg.getOwner()->getParent()->getParentOp();
-  return dyn_cast<AffineForOp>(containingInst);
+  if (auto forOp = dyn_cast<AffineForOp>(containingInst))
+    // Check to make sure `val` is the induction variable, not an iter_arg.
+    return forOp.getInductionVar() == val ? forOp : AffineForOp();
+  return AffineForOp();
 }
 
 /// Extracts the induction variables from a list of AffineForOps and returns
 /// them.
 void mlir::extractForInductionVars(ArrayRef<AffineForOp> forInsts,
                                    SmallVectorImpl<Value> *ivs) {
   ivs->reserve(forInsts.size());
   for (auto forInst : forInsts)

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

 ///       algorithm can convert the following operations to their vector form:
 ///         * Affine load and store operations are converted to opaque vector
 ///           transfer read and write operations.
 ///         * Scalar constant operations/operands are converted to vector
 ///           constant operations (splat).
 ///         * Uniform operands (only operands defined outside of the loop nest,
 ///           for now) are broadcasted to a vector.
 ///           TODO: Support more uniform cases.
+///         * Induction variables of loops not mapped to a vector dimension are
+///           broadcasted to a vector.
 ///         * Affine for operations with 'iter_args' are vectorized by
 ///           vectorizing their 'iter_args' operands and results.
 ///           TODO: Support more complex loops with divergent lbs and/or ubs.
 ///         * The remaining operations in the loop nest are vectorized by
 ///           widening their scalar types to vector types.
 ///    b. if everything under the root AffineForOp in the current pattern
 ///       is vectorized properly, we commit that loop to the IR and remove the
 ///       scalar loop. Otherwise, we discard the vectorized loop and keep the

 /// Returns true if the provided value is vector uniform given the vectorization
 /// strategy.
 // TODO: For now, only values that are invariants to all the loops in the
 // vectorization strategy are considered vector uniforms.
 static bool isUniformDefinition(Value value,
                                 const VectorizationStrategy *strategy) {
   for (auto loopToDim : strategy->loopToVectorDim) {
     auto loop = cast<AffineForOp>(loopToDim.first);
     if (!loop.isDefinedOutsideOfLoop(value))
       return false;
   }
   return true;
 }
 
 /// Generates a broadcast op for the provided uniform value using the
 /// vectorization strategy in 'state'.
 static Operation *vectorizeUniform(Value uniformVal,
                                    VectorizationState &state) {
   OpBuilder::InsertionGuard guard(state.builder);
   state.builder.setInsertionPointAfterValue(uniformVal);
 
   auto vectorTy = getVectorType(uniformVal.getType(), state.strategy);
   auto bcastOp = state.builder.create<BroadcastOp>(uniformVal.getLoc(),
                                                    vectorTy, uniformVal);
   state.registerValueVectorReplacement(uniformVal, bcastOp);
   return bcastOp;
 }
 
+/// Returns true if the provided value is an induction variable of a loop
+/// not in `loopToVectorDim` given the vectorization strategy.
+static bool isNonVecDimLoopIV(Value value,
+                              const VectorizationStrategy *strategy) {
+  AffineForOp forOp = getForInductionVarOwner(value);
+  if (!forOp ||
+      strategy->loopToVectorDim.find(forOp) != strategy->loopToVectorDim.end())
+    return false;
+
+  return true;
+}
+
+/// Generates a broadcast op for the provided induction variable or its scalar
+/// replacement using the vectorization strategy in 'state'.
+static Operation *vectorizeIV(Value iv, VectorizationState &state) {
+  OpBuilder::InsertionGuard guard(state.builder);
+  auto ivScalarRepl = state.valueScalarReplacement.lookupOrDefault(iv);
+  state.builder.setInsertionPointAfterValue(ivScalarRepl);
+
+  auto vectorTy = getVectorType(iv.getType(), state.strategy);
+  auto bcastOp =
+      state.builder.create<BroadcastOp>(iv.getLoc(), vectorTy, ivScalarRepl);
+  state.registerValueVectorReplacement(iv, bcastOp);
+  return bcastOp;
+}
+
 /// Tries to vectorize a given `operand` by applying the following logic:
 /// 1. if the defining operation has been already vectorized, `operand` is
 ///    already in the proper vector form;
 /// 2. if the `operand` is a constant, returns the vectorized form of the
 ///    constant;
 /// 3. if the `operand` is uniform, returns a vector broadcast of the `op`;
-/// 4. otherwise, the vectorization of `operand` is not supported.
+/// 4. if the `operand` is an induction variable of a loop not in
+///    `loopToVectorDim`, returns a vector broadcast of the `operand` or its
+///    scalar replacement;
+/// 5. otherwise, the vectorization of `operand` is not supported.
 /// Newly created vector operations are registered in `state` as replacement
 /// for their scalar counterparts.
 /// In particular this logic captures some of the use cases where definitions
 /// that are not scoped under the current pattern are needed to vectorize.
 /// One such example is top level function constants that need to be splatted.
 ///
 /// Returns an operand that has been vectorized to match `state`'s strategy if
 /// vectorization is possible with the above logic. Returns nullptr otherwise.
static Value vectorizeOperand(Value operand, VectorizationState &state) {
 
   // Vectorize uniform values.
   if (isUniformDefinition(operand, state.strategy)) {
     Operation *vecUniform = vectorizeUniform(operand, state);
     LLVM_DEBUG(dbgs() << "-> uniform: " << *vecUniform);
     return vecUniform->getResult(0);
   }
 
+  // Vectorize induction variables of loops not in `loopToVectorDim`.
+  if (isNonVecDimLoopIV(operand, state.strategy)) {
+    Operation *vecIV = vectorizeIV(operand, state);
+    LLVM_DEBUG(dbgs() << "-> iv: " << *vecIV);
+    return vecIV->getResult(0);
+  }
+
   // Check for unsupported block argument scenarios. A supported block argument
   // should have been vectorized already.
   if (!operand.getDefiningOp())
     LLVM_DEBUG(dbgs() << "-> unsupported block argument\n");
   else
     // Generic unsupported case.
     LLVM_DEBUG(dbgs() << "-> non-vectorizable\n");
 

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

func @vec_constant_with_two_users(%M : index, %N : index) -> (f32, f32) {
   %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_block_arg
+func @vec_block_arg(%A : memref<32x512xi32>) {
+  // CHECK:      affine.for %[[IV0:[arg0-9]+]] = 0 to 512 step 128 {
+  // CHECK-NEXT:   affine.for %[[IV1:[arg0-9]+]] = 0 to 32 {
+  // CHECK-NEXT:     %[[BROADCAST:.*]] = vector.broadcast %[[IV1]] : index to vector<128xindex>
+  // CHECK-NEXT:     %[[CAST:.*]] = index_cast %[[BROADCAST]] : vector<128xindex> to vector<128xi32>
+  // CHECK-NEXT:     vector.transfer_write %[[CAST]], {{.*}}[%[[IV1]], %[[IV0]]] : vector<128xi32>, memref<32x512xi32>
+  affine.for %i = 0 to 512 {  // vectorized
+    affine.for %j = 0 to 32 {
+      %idx = std.index_cast %j : index to i32
+      affine.store %idx, %A[%j, %i] : memref<32x512xi32>
+    }
+  }
+  return
+}
+
+// -----
+
+// CHECK-LABEL: func @vec_block_arg_2
+func @vec_block_arg_2(%A : memref<?x512xf32>, %B : memref<?x512xf32>) {
+  %c0 = constant 0 : index
+  %M = memref.dim %A, %c0 : memref<?x512xf32>
+  %N = memref.dim %B, %c0 : memref<?x512xf32>
+  %cst_0 = constant 0.0 : f32
+  %c2 = constant 2 : index
+  %c1 = constant 1 : index
+  // CHECK:      affine.for %[[IV0:[arg0-9]+]] = 0 to %{{.*}} {
+  // CHECK-NEXT:   %[[BROADCAST1:.*]] = vector.broadcast %[[IV0]] : index to vector<128xindex>
+  // CHECK-NEXT:   affine.for %[[IV1:[arg0-9]+]] = 0 to 512 step 128 {
+  // CHECK-NOT:      vector.broadcast %[[IV1]]
+  // CHECK:          %{{.*}} = affine.for %[[IV2:[arg0-9]+]] = 0 to 2 {{.*}} -> (vector<128xf32>) {
+  // CHECK-NEXT:       %[[BROADCAST2:.*]] = vector.broadcast %[[IV2]] : index to vector<128xindex>
+  // CHECK:             vector.transfer_read %{{.*}}[%{{.*}}, %{{.*}}], %{{.*}} : memref<?x512xf32>, vector<128xf32>
+  // CHECK-NEXT:        muli %[[BROADCAST1]], %{{.*}} : vector<128xindex>
+  // CHECK-NEXT:        addi %{{.*}}, %[[BROADCAST2]] : vector<128xindex>
+  // CHECK:          }
+  // CHECK-NEXT:     vector.transfer_write %{{.*}}, %{{.*}} : vector<128xf32>, memref<?x512xf32>
+  affine.for %i0 = 0 to %N {
+    affine.for %i1 = 0 to 512 { // vectorized
+      %0 = affine.for %i2 = 0 to 2 iter_args(%accum = %cst_0) -> f32 {
+        %2 = affine.load %A[%i0 * 2 + %i2 - 1, %i1] : memref<?x512xf32>
+        %mul = muli %i0, %c2 : index
+        %add = addi %mul, %i2 : index
+        %sub = subi %add, %c1 : index
+        %3 = cmpi uge, %sub, %c0 : index
+        %4 = cmpi ule, %sub, %M : index
+        %5 = and %3, %4 : i1
+        %6 = select %5, %2, %cst_0 : f32
+        %7 = addf %accum, %6 : f32
+        affine.yield %7 : f32
+      }
+      affine.store %0, %B[%i0, %i1] : memref<?x512xf32>
+    }
+  }
+  return
+}
+
+// -----
+
 // 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>