[mlir] Update Toy tutorial to use callback-based loop constructors

[ftynse]

Summary

We recently introduced support for building loops or loop nests using callbacks
that populate the body. Use these in the tutorial instead of setInsertionPoint
manipulations.

Depends On D81955

[rriddle]

mlir/examples/toy/Ch5/mlir/LowerToAffineLoops.cpp

74

nit: One of these is Value and the other is value.

mlir/examples/toy/Ch6/mlir/LowerToAffineLoops.cpp

75

nit: One of these is Value and the other is value.

[rriddle]

mlir/examples/toy/Ch7/mlir/LowerToAffineLoops.cpp

74

Same here.

[ftynse]

Thanks for the review! Addressed.

mlir/examples/toy/Ch5/mlir/LowerToAffineLoops.cpp

static Value insertAllocAndDealloc(MemRefType type, Location loc,
   // Make sure to deallocate this alloc at the end of the block. This is fine
   // as toy functions have no control flow.
   auto dealloc = rewriter.create<DeallocOp>(loc, alloc);
   dealloc.getOperation()->moveBefore(&parentBlock->back());
   return alloc;
 }
 
 /// This defines the function type used to process an iteration of a lowered
-/// loop. It takes as input a rewriter, an array of memRefOperands corresponding
-/// to the operands of the input operation, and the set of loop induction
-/// variables for the iteration. It returns a value to store at the current
-/// index of the iteration.
-using LoopIterationFn = function_ref<Value(PatternRewriter &rewriter,
-                                           ArrayRef<Value> memRefOperands,
-                                           ArrayRef<Value> loopIvs)>;
+/// loop. It takes as input an OpBuilder, an range of memRefOperands
+/// corresponding to the operands of the input operation, and the range of loop
+/// induction variables for the iteration. It returns a value to store at the
+/// current index of the iteration.
+using LoopIterationFn = function_ref<Value(
+    OpBuilder &rewriter, ValueRange memRefOperands, ValueRange loopIvs)>;
 
-static void lowerOpToLoops(Operation *op, ArrayRef<Value> operands,
+static void lowerOpToLoops(Operation *op, ValueRange operands,
                            PatternRewriter &rewriter,
                            LoopIterationFn processIteration) {
   auto tensorType = (*op->result_type_begin()).cast<TensorType>();
   auto loc = op->getLoc();
 
   // Insert an allocation and deallocation for the result of this operation.
   auto memRefType = convertTensorToMemRef(tensorType);
   auto alloc = insertAllocAndDealloc(memRefType, loc, rewriter);
 
-  // Create an empty affine loop for each of the dimensions within the shape.
-  SmallVector<Value, 4> loopIvs;
-  for (auto dim : tensorType.getShape()) {
-    auto loop = rewriter.create<AffineForOp>(loc, /*lb=*/0, dim, /*step=*/1);
-    loopIvs.push_back(loop.getInductionVar());
-
-    // Update the rewriter insertion point to the beginning of the loop.
-    rewriter.setInsertionPointToStart(loop.getBody());
-  }
-
-  // Generate a call to the processing function with the rewriter, the memref
-  // operands, and the loop induction variables. This function will return the
-  // value to store at the current index.
-  Value valueToStore = processIteration(rewriter, operands, loopIvs);
-  rewriter.create<AffineStoreOp>(loc, valueToStore, alloc,
-                                 llvm::makeArrayRef(loopIvs));
+  // Create a nest of affine loops, with one loop per dimension of the shape.
+  // The buildAffineLoopNest function takes a callback that is used to construct
+  // the body of the innermost loop given a builder, a location and a range of
+  // loop induction variables.
+  SmallVector<int64_t, 4> lowerBounds(tensorType.getRank(), /*Value=*/0);
+  SmallVector<int64_t, 4> steps(tensorType.getRank(), /*Value=*/1);
+  buildAffineLoopNest(
+      rewriter, loc, lowerBounds, tensorType.getShape(), steps,
+      [&](OpBuilder &nestedBuilder, Location loc, ValueRange ivs) {
+        // Call the processing function with the rewriter, the memref operands,
+        // and the loop induction variables. This function will return the value
+        // to store at the current index.
+        Value valueToStore = processIteration(nestedBuilder, operands, ivs);
+        nestedBuilder.create<AffineStoreOp>(loc, valueToStore, alloc, ivs);
+      });
 
   // Replace this operation with the generated alloc.
   rewriter.replaceOp(op, alloc);
 }
 
 namespace {
 //===----------------------------------------------------------------------===//
 // ToyToAffine RewritePatterns: Binary operations
 //===----------------------------------------------------------------------===//
 
 template <typename BinaryOp, typename LoweredBinaryOp>
 struct BinaryOpLowering : public ConversionPattern {
   BinaryOpLowering(MLIRContext *ctx)
       : ConversionPattern(BinaryOp::getOperationName(), 1, ctx) {}
 
   LogicalResult
   matchAndRewrite(Operation *op, ArrayRef<Value> operands,
                   ConversionPatternRewriter &rewriter) const final {
     auto loc = op->getLoc();
     lowerOpToLoops(
         op, operands, rewriter,
-        [loc](PatternRewriter &rewriter, ArrayRef<Value> memRefOperands,
-              ArrayRef<Value> loopIvs) {
+        [loc](OpBuilder &builder, ValueRange memRefOperands,
+              ValueRange loopIvs) {
           // Generate an adaptor for the remapped operands of the BinaryOp. This
           // allows for using the nice named accessors that are generated by the
           // ODS.
           typename BinaryOp::Adaptor binaryAdaptor(memRefOperands);
 
           // Generate loads for the element of 'lhs' and 'rhs' at the inner
           // loop.
           auto loadedLhs =
-              rewriter.create<AffineLoadOp>(loc, binaryAdaptor.lhs(), loopIvs);
+              builder.create<AffineLoadOp>(loc, binaryAdaptor.lhs(), loopIvs);
           auto loadedRhs =
-              rewriter.create<AffineLoadOp>(loc, binaryAdaptor.rhs(), loopIvs);
+              builder.create<AffineLoadOp>(loc, binaryAdaptor.rhs(), loopIvs);
 
           // Create the binary operation performed on the loaded values.
-          return rewriter.create<LoweredBinaryOp>(loc, loadedLhs, loadedRhs);
+          return builder.create<LoweredBinaryOp>(loc, loadedLhs, loadedRhs);
         });
     return success();
   }
 };
 using AddOpLowering = BinaryOpLowering<toy::AddOp, AddFOp>;
 using MulOpLowering = BinaryOpLowering<toy::MulOp, MulFOp>;
 
 //===----------------------------------------------------------------------===//

 struct TransposeOpLowering : public ConversionPattern {
   TransposeOpLowering(MLIRContext *ctx)
       : ConversionPattern(toy::TransposeOp::getOperationName(), 1, ctx) {}
 
   LogicalResult
   matchAndRewrite(Operation *op, ArrayRef<Value> operands,
                   ConversionPatternRewriter &rewriter) const final {
     auto loc = op->getLoc();
-    lowerOpToLoops(
-        op, operands, rewriter,
-        [loc](PatternRewriter &rewriter, ArrayRef<Value> memRefOperands,
-              ArrayRef<Value> loopIvs) {
-          // Generate an adaptor for the remapped operands of the TransposeOp.
-          // This allows for using the nice named accessors that are generated
-          // by the ODS.
+    lowerOpToLoops(op, operands, rewriter,
+                   [loc](OpBuilder &builder, ValueRange memRefOperands,
+                         ValueRange loopIvs) {
+                     // Generate an adaptor for the remapped operands of the
+                     // TransposeOp. This allows for using the nice named
+                     // accessors that are generated by the ODS.
                      toy::TransposeOpAdaptor transposeAdaptor(memRefOperands);
                      Value input = transposeAdaptor.input();
 
-          // Transpose the elements by generating a load from the reverse
-          // indices.
+                     // Transpose the elements by generating a load from the
+                     // reverse indices.
                      SmallVector<Value, 2> reverseIvs(llvm::reverse(loopIvs));
-          return rewriter.create<AffineLoadOp>(loc, input, reverseIvs);
+                     return builder.create<AffineLoadOp>(loc, input,
+                                                         reverseIvs);
                    });
     return success();
   }
 };
 
 } // end anonymous namespace.
 
 //===----------------------------------------------------------------------===//
 // ToyToAffineLoweringPass

mlir/examples/toy/Ch6/mlir/LowerToAffineLoops.cpp

static Value insertAllocAndDealloc(MemRefType type, Location loc,
   // Make sure to deallocate this alloc at the end of the block. This is fine
   // as toy functions have no control flow.
   auto dealloc = rewriter.create<DeallocOp>(loc, alloc);
   dealloc.getOperation()->moveBefore(&parentBlock->back());
   return alloc;
 }
 
 /// This defines the function type used to process an iteration of a lowered
-/// loop. It takes as input a rewriter, an array of memRefOperands corresponding
-/// to the operands of the input operation, and the set of loop induction
-/// variables for the iteration. It returns a value to store at the current
-/// index of the iteration.
-using LoopIterationFn = function_ref<Value(PatternRewriter &rewriter,
-                                           ArrayRef<Value> memRefOperands,
-                                           ArrayRef<Value> loopIvs)>;
+/// loop. It takes as input an OpBuilder, an range of memRefOperands
+/// corresponding to the operands of the input operation, and the range of loop
+/// induction variables for the iteration. It returns a value to store at the
+/// current index of the iteration.
+using LoopIterationFn = function_ref<Value(
+    OpBuilder &builder, ValueRange memRefOperands, ValueRange loopIvs)>;
 
 static void lowerOpToLoops(Operation *op, ArrayRef<Value> operands,
                            PatternRewriter &rewriter,
                            LoopIterationFn processIteration) {
   auto tensorType = (*op->result_type_begin()).cast<TensorType>();
   auto loc = op->getLoc();
 
   // Insert an allocation and deallocation for the result of this operation.
   auto memRefType = convertTensorToMemRef(tensorType);
   auto alloc = insertAllocAndDealloc(memRefType, loc, rewriter);
 
-  // Create an empty affine loop for each of the dimensions within the shape.
-  SmallVector<Value, 4> loopIvs;
-  for (auto dim : tensorType.getShape()) {
-    auto loop = rewriter.create<AffineForOp>(loc, /*lb=*/0, dim, /*step=*/1);
-    loopIvs.push_back(loop.getInductionVar());
-
-    // Update the rewriter insertion point to the beginning of the loop.
-    rewriter.setInsertionPointToStart(loop.getBody());
-  }
-
-  // Generate a call to the processing function with the rewriter, the memref
-  // operands, and the loop induction variables. This function will return the
-  // value to store at the current index.
-  Value valueToStore = processIteration(rewriter, operands, loopIvs);
-  rewriter.create<AffineStoreOp>(loc, valueToStore, alloc,
-                                 llvm::makeArrayRef(loopIvs));
+  // Create a nest of affine loops, with one loop per dimension of the shape.
+  // The buildAffineLoopNest function takes a callback that is used to construct
+  // the body of the innermost loop given a builder, a location and a range of
+  // loop induction variables.
+  SmallVector<int64_t, 4> lowerBounds(tensorType.getRank(), /*Value=*/0);
+  SmallVector<int64_t, 4> steps(tensorType.getRank(), /*Value=*/1);
+  buildAffineLoopNest(
+      rewriter, loc, lowerBounds, tensorType.getShape(), steps,
+      [&](OpBuilder &nestedBuilder, Location loc, ValueRange ivs) {
+        // Call the processing function with the rewriter, the memref operands,
+        // and the loop induction variables. This function will return the value
+        // to store at the current index.
+        Value valueToStore = processIteration(nestedBuilder, operands, ivs);
+        nestedBuilder.create<AffineStoreOp>(loc, valueToStore, alloc, ivs);
+      });
 
   // Replace this operation with the generated alloc.
   rewriter.replaceOp(op, alloc);
 }
 
 namespace {
 //===----------------------------------------------------------------------===//
 // ToyToAffine RewritePatterns: Binary operations
 //===----------------------------------------------------------------------===//
 
 template <typename BinaryOp, typename LoweredBinaryOp>
 struct BinaryOpLowering : public ConversionPattern {
   BinaryOpLowering(MLIRContext *ctx)
       : ConversionPattern(BinaryOp::getOperationName(), 1, ctx) {}
 
   LogicalResult
   matchAndRewrite(Operation *op, ArrayRef<Value> operands,
                   ConversionPatternRewriter &rewriter) const final {
     auto loc = op->getLoc();
     lowerOpToLoops(
         op, operands, rewriter,
-        [loc](PatternRewriter &rewriter, ArrayRef<Value> memRefOperands,
-              ArrayRef<Value> loopIvs) {
+        [loc](OpBuilder &builder, ValueRange memRefOperands,
+              ValueRange loopIvs) {
           // Generate an adaptor for the remapped operands of the BinaryOp. This
           // allows for using the nice named accessors that are generated by the
           // ODS.
           typename BinaryOp::Adaptor binaryAdaptor(memRefOperands);
 
           // Generate loads for the element of 'lhs' and 'rhs' at the inner
           // loop.
           auto loadedLhs =
-              rewriter.create<AffineLoadOp>(loc, binaryAdaptor.lhs(), loopIvs);
+              builder.create<AffineLoadOp>(loc, binaryAdaptor.lhs(), loopIvs);
           auto loadedRhs =
-              rewriter.create<AffineLoadOp>(loc, binaryAdaptor.rhs(), loopIvs);
+              builder.create<AffineLoadOp>(loc, binaryAdaptor.rhs(), loopIvs);
 
           // Create the binary operation performed on the loaded values.
-          return rewriter.create<LoweredBinaryOp>(loc, loadedLhs, loadedRhs);
+          return builder.create<LoweredBinaryOp>(loc, loadedLhs, loadedRhs);
         });
     return success();
   }
 };
 using AddOpLowering = BinaryOpLowering<toy::AddOp, AddFOp>;
 using MulOpLowering = BinaryOpLowering<toy::MulOp, MulFOp>;
 
 //===----------------------------------------------------------------------===//

 struct TransposeOpLowering : public ConversionPattern {
   TransposeOpLowering(MLIRContext *ctx)
       : ConversionPattern(toy::TransposeOp::getOperationName(), 1, ctx) {}
 
   LogicalResult
   matchAndRewrite(Operation *op, ArrayRef<Value> operands,
                   ConversionPatternRewriter &rewriter) const final {
     auto loc = op->getLoc();
-    lowerOpToLoops(
-        op, operands, rewriter,
-        [loc](PatternRewriter &rewriter, ArrayRef<Value> memRefOperands,
-              ArrayRef<Value> loopIvs) {
-          // Generate an adaptor for the remapped operands of the TransposeOp.
-          // This allows for using the nice named accessors that are generated
-          // by the ODS.
+    lowerOpToLoops(op, operands, rewriter,
+                   [loc](OpBuilder &builder, ValueRange memRefOperands,
+                         ValueRange loopIvs) {
+                     // Generate an adaptor for the remapped operands of the
+                     // TransposeOp. This allows for using the nice named
+                     // accessors that are generated by the ODS.
                      toy::TransposeOpAdaptor transposeAdaptor(memRefOperands);
                      Value input = transposeAdaptor.input();
 
-          // Transpose the elements by generating a load from the reverse
-          // indices.
+                     // Transpose the elements by generating a load from the
+                     // reverse indices.
                      SmallVector<Value, 2> reverseIvs(llvm::reverse(loopIvs));
-          return rewriter.create<AffineLoadOp>(loc, input, reverseIvs);
+                     return builder.create<AffineLoadOp>(loc, input,
+                                                         reverseIvs);
                    });
     return success();
   }
 };
 
 } // end anonymous namespace.
 
 //===----------------------------------------------------------------------===//
 // ToyToAffineLoweringPass

mlir/examples/toy/Ch7/mlir/LowerToAffineLoops.cpp

static Value insertAllocAndDealloc(MemRefType type, Location loc,
   // Make sure to deallocate this alloc at the end of the block. This is fine
   // as toy functions have no control flow.
   auto dealloc = rewriter.create<DeallocOp>(loc, alloc);
   dealloc.getOperation()->moveBefore(&parentBlock->back());
   return alloc;
 }
 
 /// This defines the function type used to process an iteration of a lowered
-/// loop. It takes as input a rewriter, an array of memRefOperands corresponding
-/// to the operands of the input operation, and the set of loop induction
-/// variables for the iteration. It returns a value to store at the current
-/// index of the iteration.
-using LoopIterationFn = function_ref<Value(PatternRewriter &rewriter,
-                                           ArrayRef<Value> memRefOperands,
-                                           ArrayRef<Value> loopIvs)>;
+/// loop. It takes as input an OpBuilder, an range of memRefOperands
+/// corresponding to the operands of the input operation, and the range of loop
+/// induction variables for the iteration. It returns a value to store at the
+/// current index of the iteration.
+using LoopIterationFn = function_ref<Value(
+    OpBuilder &rewriter, ValueRange memRefOperands, ValueRange loopIvs)>;
 
-static void lowerOpToLoops(Operation *op, ArrayRef<Value> operands,
+static void lowerOpToLoops(Operation *op, ValueRange operands,
                            PatternRewriter &rewriter,
                            LoopIterationFn processIteration) {
   auto tensorType = (*op->result_type_begin()).cast<TensorType>();
   auto loc = op->getLoc();
 
   // Insert an allocation and deallocation for the result of this operation.
   auto memRefType = convertTensorToMemRef(tensorType);
   auto alloc = insertAllocAndDealloc(memRefType, loc, rewriter);
 
-  // Create an empty affine loop for each of the dimensions within the shape.
-  SmallVector<Value, 4> loopIvs;
-  for (auto dim : tensorType.getShape()) {
-    auto loop = rewriter.create<AffineForOp>(loc, /*lb=*/0, dim, /*step=*/1);
-    loopIvs.push_back(loop.getInductionVar());
-
-    // Update the rewriter insertion point to the beginning of the loop.
-    rewriter.setInsertionPointToStart(loop.getBody());
-  }
-
-  // Generate a call to the processing function with the rewriter, the memref
-  // operands, and the loop induction variables. This function will return the
-  // value to store at the current index.
-  Value valueToStore = processIteration(rewriter, operands, loopIvs);
-  rewriter.create<AffineStoreOp>(loc, valueToStore, alloc,
-                                 llvm::makeArrayRef(loopIvs));
+  // Create a nest of affine loops, with one loop per dimension of the shape.
+  // The buildAffineLoopNest function takes a callback that is used to construct
+  // the body of the innermost loop given a builder, a location and a range of
+  // loop induction variables.
+  SmallVector<int64_t, 4> lowerBounds(tensorType.getRank(), /*Value=*/0);
+  SmallVector<int64_t, 4> steps(tensorType.getRank(), /*Value=*/1);
+  buildAffineLoopNest(
+      rewriter, loc, lowerBounds, tensorType.getShape(), steps,
+      [&](OpBuilder &nestedBuilder, Location loc, ValueRange ivs) {
+        // Call the processing function with the rewriter, the memref operands,
+        // and the loop induction variables. This function will return the value
+        // to store at the current index.
+        Value valueToStore = processIteration(nestedBuilder, operands, ivs);
+        nestedBuilder.create<AffineStoreOp>(loc, valueToStore, alloc, ivs);
+      });
 
   // Replace this operation with the generated alloc.
   rewriter.replaceOp(op, alloc);
 }
 
 namespace {
 //===----------------------------------------------------------------------===//
 // ToyToAffine RewritePatterns: Binary operations
 //===----------------------------------------------------------------------===//
 
 template <typename BinaryOp, typename LoweredBinaryOp>
 struct BinaryOpLowering : public ConversionPattern {
   BinaryOpLowering(MLIRContext *ctx)
       : ConversionPattern(BinaryOp::getOperationName(), 1, ctx) {}
 
   LogicalResult
   matchAndRewrite(Operation *op, ArrayRef<Value> operands,
                   ConversionPatternRewriter &rewriter) const final {
     auto loc = op->getLoc();
     lowerOpToLoops(
         op, operands, rewriter,
-        [loc](PatternRewriter &rewriter, ArrayRef<Value> memRefOperands,
-              ArrayRef<Value> loopIvs) {
+        [loc](OpBuilder &builder, ValueRange memRefOperands,
+              ValueRange loopIvs) {
           // Generate an adaptor for the remapped operands of the BinaryOp. This
           // allows for using the nice named accessors that are generated by the
           // ODS.
           typename BinaryOp::Adaptor binaryAdaptor(memRefOperands);
 
           // Generate loads for the element of 'lhs' and 'rhs' at the inner
           // loop.
           auto loadedLhs =
-              rewriter.create<AffineLoadOp>(loc, binaryAdaptor.lhs(), loopIvs);
+              builder.create<AffineLoadOp>(loc, binaryAdaptor.lhs(), loopIvs);
           auto loadedRhs =
-              rewriter.create<AffineLoadOp>(loc, binaryAdaptor.rhs(), loopIvs);
+              builder.create<AffineLoadOp>(loc, binaryAdaptor.rhs(), loopIvs);
 
           // Create the binary operation performed on the loaded values.
-          return rewriter.create<LoweredBinaryOp>(loc, loadedLhs, loadedRhs);
+          return builder.create<LoweredBinaryOp>(loc, loadedLhs, loadedRhs);
         });
     return success();
   }
 };
 using AddOpLowering = BinaryOpLowering<toy::AddOp, AddFOp>;
 using MulOpLowering = BinaryOpLowering<toy::MulOp, MulFOp>;
 
 //===----------------------------------------------------------------------===//

 struct TransposeOpLowering : public ConversionPattern {
   TransposeOpLowering(MLIRContext *ctx)
       : ConversionPattern(toy::TransposeOp::getOperationName(), 1, ctx) {}
 
   LogicalResult
   matchAndRewrite(Operation *op, ArrayRef<Value> operands,
                   ConversionPatternRewriter &rewriter) const final {
     auto loc = op->getLoc();
-    lowerOpToLoops(
-        op, operands, rewriter,
-        [loc](PatternRewriter &rewriter, ArrayRef<Value> memRefOperands,
-              ArrayRef<Value> loopIvs) {
-          // Generate an adaptor for the remapped operands of the TransposeOp.
-          // This allows for using the nice named accessors that are generated
-          // by the ODS.
+    lowerOpToLoops(op, operands, rewriter,
+                   [loc](OpBuilder &builder, ValueRange memRefOperands,
+                         ValueRange loopIvs) {
+                     // Generate an adaptor for the remapped operands of the
+                     // TransposeOp. This allows for using the nice named
+                     // accessors that are generated by the ODS.
                      toy::TransposeOpAdaptor transposeAdaptor(memRefOperands);
                      Value input = transposeAdaptor.input();
 
-          // Transpose the elements by generating a load from the reverse
-          // indices.
+                     // Transpose the elements by generating a load from the
+                     // reverse indices.
                      SmallVector<Value, 2> reverseIvs(llvm::reverse(loopIvs));
-          return rewriter.create<AffineLoadOp>(loc, input, reverseIvs);
+                     return builder.create<AffineLoadOp>(loc, input,
+                                                         reverseIvs);
                    });
     return success();
   }
 };
 
 } // end anonymous namespace.
 
 //===----------------------------------------------------------------------===//
 // ToyToAffineLoweringPass