[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

261

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

1106

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.

1124

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

1150

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

1150

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

1089

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

 ///       representation. This order ensures that all the operands of a given
 ///       operation have been vectorized before the operation itself in a single
 ///       traversal, except for operands defined outside of the loop nest. The
 ///       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,
+///         * Uniform operands (only induction variables of loops not mapped to
+///           a vector dimension, or operands defined outside of the loop nest
 ///           for now) are broadcasted to a vector.
 ///           TODO: Support more uniform cases.
 ///         * 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
 ///       original scalar loop.
 ///    c. vectorization is applied on the next pattern in the list. Because
LLVM_DEBUG(dbgs() << "\n[early-vect]+++++ creating a mask:\n"
                     << mask << "\n");
 
   state.vecLoopToMask[vecForOp] = mask;
   return mask;
 }
 
 /// 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.
+// TODO: For now, only values that are induction variables of loops not in
+// `loopToVectorDim` or invariants to all the loops in the vectorization
+// strategy are considered vector uniforms.
 static bool isUniformDefinition(Value value,
                                 const VectorizationStrategy *strategy) {
+  AffineForOp forOp = getForInductionVarOwner(value);
+  if (forOp &&
+      strategy->loopToVectorDim.find(forOp) == strategy->loopToVectorDim.end())
+    return true;
+
   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 uniformScalarRepl =
+      state.valueScalarReplacement.lookupOrDefault(uniformVal);
+  state.builder.setInsertionPointAfterValue(uniformScalarRepl);
 
   auto vectorTy = getVectorType(uniformVal.getType(), state.strategy);
   auto bcastOp = state.builder.create<BroadcastOp>(uniformVal.getLoc(),
-                                                   vectorTy, uniformVal);
+                                                   vectorTy, uniformScalarRepl);
   state.registerValueVectorReplacement(uniformVal, 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.
 /// 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.
 ///
 /// TODO: handle more complex cases.
 static Value vectorizeOperand(Value operand, VectorizationState &state) {
   LLVM_DEBUG(dbgs() << "\n[early-vect]+++++ vectorize operand: " << operand);
   // If this value is already vectorized, we are done.
   if (Value vecRepl = state.valueVectorReplacement.lookupOrNull(operand)) {
     LLVM_DEBUG(dbgs() << " -> already vectorized: " << vecRepl);
     return vecRepl;
   }
 
   // An vector operand that is not in the replacement map should never reach
   // this point. Reaching this point could mean that the code was already
   // vectorized and we shouldn't try to vectorize already vectorized code.
   assert(!operand.getType().isa<VectorType>() &&
          "Vector op not found in replacement map");
 
   // Vectorize constant.
   if (auto constOp = operand.getDefiningOp<ConstantOp>()) {
if (isUniformDefinition(operand, state.strategy)) {
     return vecUniform->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");
 
   return nullptr;
 }
 
 /// Vectorizes an affine load with the vectorization strategy in 'state' by
 /// generating a 'vector.transfer_read' op with the proper permutation map
 /// inferred from the indices of the load. The new 'vector.transfer_read' is

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-DAG: #[[$map0:map[0-9]+]] = affine_map<(d0, d1, d2) -> (d0 * 2 + d1 - 1)>
+// CHECK-DAG: #[[$map1:map[0-9]+]] = affine_map<(d0, d1, d2) -> (d2)>
+// CHECK-LABEL: func @vec_block_arg_2
+func @vec_block_arg_2(%A : memref<?x512xindex>) {
+  %c0 = constant 0 : index
+  %N = memref.dim %A, %c0 : memref<?x512xindex>
+  // 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 {
+  // CHECK-NEXT:       %[[BROADCAST2:.*]] = vector.broadcast %[[IV2]] : index to vector<128xindex>
+  // CHECK-NEXT:       %[[INDEX1:.*]] = affine.apply #[[$map0]](%[[IV0]], %[[IV2]], %[[IV1]])
+  // CHECK-NEXT:       %[[INDEX2:.*]] = affine.apply #[[$map1]](%[[IV0]], %[[IV2]], %[[IV1]])
+  // CHECK:            %[[LOAD:.*]] = vector.transfer_read %{{.*}}[%[[INDEX1]], %[[INDEX2]]], %{{.*}} : memref<?x512xindex>, vector<128xindex>
+  // CHECK-NEXT:       muli %[[BROADCAST1]], %[[LOAD]] : vector<128xindex>
+  // CHECK-NEXT:       addi %{{.*}}, %[[BROADCAST2]] : vector<128xindex>
+  // CHECK:          }
+  affine.for %i0 = 0 to %N {
+    affine.for %i1 = 0 to 512 { // vectorized
+      affine.for %i2 = 0 to 2 {
+        %0 = affine.load %A[%i0 * 2 + %i2 - 1, %i1] : memref<?x512xindex>
+        %mul = muli %i0, %0 : index
+        %add = addi %mul, %i2 : index
+      }
+    }  
+  }  
+  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>