[mlir][Affine][Vector] Add initial support for 'iter_args' to Affine vectorizer.

[dcaballe]

Summary

This patch adds support for vectorizing loops with 'iter_args' when those loops
are not a vector dimension. This allows vectorizing outer loops with an inner
'iter_args' loop (e.g., reductions). Vectorizing scenarios where 'iter_args'
loops are vector dimensions would require more work (e.g., analysis,
generating horizontal reduction, etc.) not included in this patch.

[aartbik]

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

624

The second sentence reads almost like a dup of the first.

728

same

1041

add TODO: for this

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

518

(1) space after // before CHECK
(2 can we move this to L505

(or move the CHECK-LABEL here)

I know that was have many different styles throughout, but I have a slight preference on keeping all CHECKs together when this form is used

[dcaballe]

Addressed feedback. Thanks!

[dcaballe]

Thanks! I'll commit this tomorrow if no more comments.

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.
+///         * 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.
-///         * TODO: Add vectorization support for loops with 'iter_args' and
-///           more complex loops with divergent lbs and/or ubs.
 ///    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
 ///       pattern interference avoidance is not yet implemented and that we do
 ///       not support further vectorizing an already vector load we need to
 ///       re-verify that the pattern is still vectorizable. This is expected to
struct VectorizationState {
   /// scalar version still exists after vectorization), like uniforms.
   ///
   /// Example:
   ///   * 'replaced': block argument or operation outside of the vectorized
   ///     loop.
   ///   * 'replacement': %0 = vector.broadcast %1 : f32 to vector<128xf32>
   void registerValueVectorReplacement(Value replaced, Operation *replacement);
 
+  /// Registers the vector replacement of a block argument (e.g., iter_args).
+  ///
+  /// Example:
+  ///   * 'replaced': 'iter_arg' block argument.
+  ///   * 'replacement': vectorized 'iter_arg' block argument.
+  void registerBlockArgVectorReplacement(BlockArgument replaced,
+                                         BlockArgument replacement);
+
   /// Registers the scalar replacement of a scalar value. 'replacement' must be
   /// scalar. Both values must be block arguments. Operation results should be
   /// replaced using the 'registerOp*' utilitites.
   ///
   /// This utility is used to register the replacement of block arguments
   /// that are within the loop to be vectorized and will continue being scalar
   /// within the vector loop.
   ///

 ///   * 'replacement': %0 = addf %1, %2 : vector<128xf32>
 void VectorizationState::registerOpVectorReplacement(Operation *replaced,
                                                      Operation *replacement) {
   LLVM_DEBUG(dbgs() << "\n[early-vect]+++++ commit vectorized op:\n");
   LLVM_DEBUG(dbgs() << *replaced << "\n");
   LLVM_DEBUG(dbgs() << "into\n");
   LLVM_DEBUG(dbgs() << *replacement << "\n");
 
-  assert(replaced->getNumResults() <= 1 && "Unsupported multi-result op");
   assert(replaced->getNumResults() == replacement->getNumResults() &&
          "Unexpected replaced and replacement results");
   assert(opVectorReplacement.count(replaced) == 0 && "already registered");
   opVectorReplacement[replaced] = replacement;
 
-  if (replaced->getNumResults() > 0)
-    registerValueVectorReplacementImpl(replaced->getResult(0),
-                                       replacement->getResult(0));
+  for (auto resultTuple :
+       llvm::zip(replaced->getResults(), replacement->getResults()))
+    registerValueVectorReplacementImpl(std::get<0>(resultTuple),
+                                       std::get<1>(resultTuple));
 }
 
 /// Registers the vector replacement of a scalar value. The replacement
 /// operation should have a single result, which replaces the scalar value.
 ///
 /// This utility is used to register the vector replacement of block arguments
 /// and operation results which are not directly vectorized (i.e., their
 /// scalar version still exists after vectorization), like uniforms.
 ///
 /// Example:
 ///   * 'replaced': block argument or operation outside of the vectorized loop.
 ///   * 'replacement': %0 = vector.broadcast %1 : f32 to vector<128xf32>
 void VectorizationState::registerValueVectorReplacement(
     Value replaced, Operation *replacement) {
   assert(replacement->getNumResults() == 1 &&
          "Expected single-result replacement");
   if (Operation *defOp = replaced.getDefiningOp())
     registerOpVectorReplacement(defOp, replacement);
   else
     registerValueVectorReplacementImpl(replaced, replacement->getResult(0));
 }
 
+/// Registers the vector replacement of a block argument (e.g., iter_args).
+///
+/// Example:
+///   * 'replaced': 'iter_arg' block argument.
+///   * 'replacement': vectorized 'iter_arg' block argument.
+void VectorizationState::registerBlockArgVectorReplacement(
+    BlockArgument replaced, BlockArgument replacement) {
+  registerValueVectorReplacementImpl(replaced, replacement);
+}
+
 void VectorizationState::registerValueVectorReplacementImpl(Value replaced,
                                                             Value replacement) {
   assert(!valueVectorReplacement.contains(replaced) &&
          "Vector replacement already registered");
   assert(replacement.getType().isa<VectorType>() &&
          "Expected vector type in vector replacement");
   valueVectorReplacement.map(replaced, replacement);
 }

 /// insertion point is set to the new loop's body so that subsequent vectorized
 /// operations are inserted into the new loop. If the loop is a vector
 /// dimension, the step of the newly created loop will reflect the vectorization
 /// factor used to vectorized that dimension.
 // TODO: Add support for 'iter_args'. Related operands and results will be
 // vectorized at this point.
 static Operation *vectorizeAffineForOp(AffineForOp forOp,
                                        VectorizationState &state) {
-  // 'iter_args' not supported yet.
-  if (forOp.getNumIterOperands() > 0)
+  const VectorizationStrategy &strategy = *state.strategy;
+  auto loopToVecDimIt = strategy.loopToVectorDim.find(forOp);
+  bool isLoopVecDim = loopToVecDimIt != strategy.loopToVectorDim.end();
+
+  // We only support 'iter_args' when the loop is not one of the vector
+  // dimensions.
+  // TODO: Support vector dimension loops. They require special handling:
+  // generate horizontal reduction, last-value extraction, etc.
+  if (forOp.getNumIterOperands() > 0 && isLoopVecDim)
     return nullptr;
 
   // If we are vectorizing a vector dimension, compute a new step for the new
   // vectorized loop using the vectorization factor for the vector dimension.
   // Otherwise, propagate the step of the scalar loop.
-  const VectorizationStrategy &strategy = *state.strategy;
-  auto loopToVecDimIt = strategy.loopToVectorDim.find(forOp);
-  bool isLoopVecDim = loopToVecDimIt != strategy.loopToVectorDim.end();
   unsigned newStep;
   if (isLoopVecDim) {
     unsigned vectorDim = loopToVecDimIt->second;
     assert(vectorDim < strategy.vectorSizes.size() && "vector dim overflow");
     int64_t forOpVecFactor = strategy.vectorSizes[vectorDim];
     newStep = forOp.getStep() * forOpVecFactor;
   } else {
     newStep = forOp.getStep();
   }
 
+  // Vectorize 'iter_args'.
+  SmallVector<Value, 8> vecIterOperands;
+  for (auto operand : forOp.getIterOperands())
+    vecIterOperands.push_back(vectorizeOperand(operand, state));
+
   auto vecForOp = state.builder.create<AffineForOp>(
       forOp.getLoc(), forOp.getLowerBoundOperands(), forOp.getLowerBoundMap(),
       forOp.getUpperBoundOperands(), forOp.getUpperBoundMap(), newStep,
-      forOp.getIterOperands(),
+      vecIterOperands,
       /*bodyBuilder=*/[](OpBuilder &, Location, Value, ValueRange) {
         // Make sure we don't create a default terminator in the loop body as
         // the proper terminator will be added during vectorization.
         return;
       });
 
   // Register loop-related replacements:
   //   1) The new vectorized loop is registered as vector replacement of the
   //      scalar loop.
   //      TODO: Support reductions along the vector dimension.
   //   2) The new iv of the vectorized loop is registered as scalar replacement
   //      since a scalar copy of the iv will prevail in the vectorized loop.
   //      TODO: A vector replacement will also be added in the future when
   //      vectorization of linear ops is supported.
-  //   3) TODO: Support 'iter_args' along non-vector dimensions.
+  //   3) The new 'iter_args' region arguments are registered as vector
+  //      replacements since they have been vectorized.
   state.registerOpVectorReplacement(forOp, vecForOp);
   state.registerValueScalarReplacement(forOp.getInductionVar(),
                                        vecForOp.getInductionVar());
-  // Map the new vectorized loop to its vector dimension.
+  for (auto iterTuple :
+       llvm ::zip(forOp.getRegionIterArgs(), vecForOp.getRegionIterArgs()))
+    state.registerBlockArgVectorReplacement(std::get<0>(iterTuple),
+                                            std::get<1>(iterTuple));
+
   if (isLoopVecDim)
     state.vecLoopToVecDim[vecForOp] = loopToVecDimIt->second;
 
   // Change insertion point so that upcoming vectorized instructions are
   // inserted into the vectorized loop's body.
   state.builder.setInsertionPointToStart(vecForOp.getBody());
   return vecForOp;
 }
static Operation *widenOp(Operation *op, VectorizationState &state) {
   return vecOp;
 }
 
 /// Vectorizes a yield operation by widening its types. The builder's insertion
 /// point is set after the vectorized parent op to continue vectorizing the
 /// operations after the parent op.
 static Operation *vectorizeAffineYieldOp(AffineYieldOp yieldOp,
                                          VectorizationState &state) {
-  // 'iter_args' not supported yet.
-  if (yieldOp.getNumOperands() > 0)
-    return nullptr;
-
-  // Vectorize the yield op and change the insertion point right after the new
-  // parent op.
   Operation *newYieldOp = widenOp(yieldOp, state);
   Operation *newParentOp = state.builder.getInsertionBlock()->getParentOp();
   state.builder.setInsertionPointAfter(newParentOp);
   return newYieldOp;
 }
 
 /// Encodes Operation-specific behavior for vectorization. In general we
 /// assume that all operands of an op must be vectorized but this is not

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

affine.for %i = 0 to 512 {
     %idx = std.index_cast %i : index to i32
     affine.store %idx, %A[%i] : memref<512xi32>
   }
   return
 }
 
 // -----
 
-// CHECK-LABEL: @vec_rejected_unsupported_reduction
-func @vec_rejected_unsupported_reduction(%in: memref<128x256xf32>, %out: memref<256xf32>) {
+// '%i' loop is vectorized, including the inner reduction over '%j'.
+
+func @vec_non_vecdim_reduction(%in: memref<128x256xf32>, %out: memref<256xf32>) {
  %cst = constant 0.000000e+00 : f32
  affine.for %i = 0 to 256 {
-   // CHECK-NOT: vector
    %final_red = affine.for %j = 0 to 128 iter_args(%red_iter = %cst) -> (f32) {
      %ld = affine.load %in[%j, %i] : memref<128x256xf32>
      %add = addf %red_iter, %ld : f32
      affine.yield %add : f32
    }
    affine.store %final_red, %out[%i] : memref<256xf32>
  }
  return
 }
 
+// CHECK-LABEL: @vec_non_vecdim_reduction
+// CHECK:       affine.for %{{.*}} = 0 to 256 step 128 {
+// CHECK:         %[[vzero:.*]] = constant dense<0.000000e+00> : vector<128xf32>
+// CHECK:         %[[final_red:.*]] = affine.for %{{.*}} = 0 to 128 iter_args(%[[red_iter:.*]] = %[[vzero]]) -> (vector<128xf32>) {
+// CHECK:           %[[ld:.*]] = vector.transfer_read %{{.*}} : memref<128x256xf32>, vector<128xf32>
+// CHECK:           %[[add:.*]] = addf %[[red_iter]], %[[ld]] : vector<128xf32>
+// CHECK:           affine.yield %[[add]] : vector<128xf32>
+// CHECK:         }
+// CHECK:         vector.transfer_write %[[final_red]], %{{.*}} : vector<128xf32>, memref<256xf32>
+// CHECK:       }
+
 // -----
 
-// CHECK-LABEL: @vec_rejected_unsupported_last_value
-func @vec_rejected_unsupported_last_value(%in: memref<128x256xf32>, %out: memref<256xf32>) {
+// '%i' loop is vectorized, including the inner reductions over '%j'.
+
+func @vec_non_vecdim_reductions(%in0: memref<128x256xf32>, %in1: memref<128x256xi32>,
+                                %out0: memref<256xf32>, %out1: memref<256xi32>) {
+ %zero = constant 0.000000e+00 : f32
+ %one = constant 1 : i32
+ affine.for %i = 0 to 256 {
+   %red0, %red1 = affine.for %j = 0 to 128
+     iter_args(%red_iter0 = %zero, %red_iter1 = %one) -> (f32, i32) {
+     %ld0 = affine.load %in0[%j, %i] : memref<128x256xf32>
+     %add = addf %red_iter0, %ld0 : f32
+     %ld1 = affine.load %in1[%j, %i] : memref<128x256xi32>
+     %mul = muli %red_iter1, %ld1 : i32
+     affine.yield %add, %mul : f32, i32
+   }
+   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:         %[[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:         }
+// CHECK:         vector.transfer_write %[[reds]]#0, %{{.*}} : vector<128xf32>, memref<256xf32>
+// CHECK:         vector.transfer_write %[[reds]]#1, %{{.*}} : vector<128xi32>, memref<256xi32>
+// CHECK:       }
+
+// -----
+
+// '%i' loop is vectorized, including the inner last value computation over '%j'.
+
+func @vec_no_vecdim_last_value(%in: memref<128x256xf32>, %out: memref<256xf32>) {
  %cst = constant 0.000000e+00 : f32
  affine.for %i = 0 to 256 {
-   // CHECK-NOT: vector
    %last_val = affine.for %j = 0 to 128 iter_args(%last_iter = %cst) -> (f32) {
      %ld = affine.load %in[%j, %i] : memref<128x256xf32>
      affine.yield %ld : f32
    }
    affine.store %last_val, %out[%i] : memref<256xf32>
  }
  return
 }
+
+// CHECK-LABEL: @vec_no_vecdim_last_value
+// CHECK:       affine.for %{{.*}} = 0 to 256 step 128 {
+// CHECK:         %[[vzero:.*]] = constant dense<0.000000e+00> : vector<128xf32>
+// CHECK:         %[[last_val:.*]] = affine.for %{{.*}} = 0 to 128 iter_args(%[[last_iter:.*]] = %[[vzero]]) -> (vector<128xf32>) {
+// CHECK:           %[[ld:.*]] = vector.transfer_read %{{.*}} : memref<128x256xf32>, vector<128xf32>
+// CHECK:           affine.yield %[[ld]] : vector<128xf32>
+// CHECK:         }
+// CHECK:         vector.transfer_write %[[last_val]], %{{.*}} : vector<128xf32>, memref<256xf32>
+// CHECK:       }