[mlir][linalg] Remove the StructuredOp capture mechanism.

[gysit]

Summary

After https://reviews.llvm.org/D104109, structured ops support scalar inputs. As a result, the capture mechanism meant to pass non-shaped parameters got redundant. The patch removes the capture semantics after the FillOp migrated to use scalar operands https://reviews.llvm.org/D104121.

[gysit]

Rebase.

mlir/include/mlir-c/Dialect/Linalg.h

 #include "mlir-c/Registration.h"
 #include "mlir-c/Support.h"
 
 #ifdef __cplusplus
 extern "C" {
 #endif
 
 /// Apply the special region builder for the builtin named Linalg op.
-/// The list of `capture` MlirValue is passed as-is to the region builder.
 /// Assert that `op` is a builtin named Linalg op.
 MLIR_CAPI_EXPORTED void
-mlirLinalgFillBuiltinNamedOpRegion(MlirDialect linalgDialect, MlirOperation op,
-                                   intptr_t n, MlirValue const *mlirCaptures);
+mlirLinalgFillBuiltinNamedOpRegion(MlirDialect linalgDialect, MlirOperation op);
 
 MLIR_DECLARE_CAPI_DIALECT_REGISTRATION(Linalg, linalg);
 
 #ifdef __cplusplus
 }
 #endif
 
 #include "mlir/Dialect/Linalg/Passes.capi.h.inc"
 
 #endif // MLIR_C_DIALECT_LINALG_H

mlir/include/mlir/Dialect/Linalg/IR/LinalgBase.td

constexpr const static ::llvm::StringLiteral
         kMemoizedIndexingMapsAttrName = "linalg.memoized_indexing_maps";
 
     /// Attribute name used to mark region arguments that can be bufferized
     /// in-place during linalg comprehensive bufferization.
     constexpr const static ::llvm::StringLiteral
       kInplaceableAttrName = "linalg.inplaceable";
 
     using RegionBuilderFunType =
-      llvm::function_ref<void(ImplicitLocOpBuilder &b, Block &, ValueRange)>;
+      llvm::function_ref<void(ImplicitLocOpBuilder &b, Block &)>;
     RegionBuilderFunType getRegionBuilder(StringRef name) {
       return namedStructuredOpRegionBuilders.lookup(name);
     }
     private:
       llvm::StringMap<RegionBuilderFunType> namedStructuredOpRegionBuilders;
   }];
 }
 
 // Whether a type is a RangeType.
 def LinalgIsRangeTypePred : CPred<"$_self.isa<RangeType>()">;
 def Range : DialectType<Linalg_Dialect, LinalgIsRangeTypePred, "range">;
 
 #endif // LINALG_BASE

mlir/include/mlir/Dialect/Linalg/IR/LinalgInterfaces.td

InterfaceMethod<
       }]
     >,
     StaticInterfaceMethod<
       /*desc=*/[{
         Returns the region builder for constructing the body for linalg.generic.
         Returns a null function if this named op does not define a region
         builder.
       }],
-      /*retTy=*/"std::function<void(ImplicitLocOpBuilder &, Block &, ValueRange)>",
+      /*retTy=*/"std::function<void(ImplicitLocOpBuilder &, Block &)>",
       /*methodName=*/"getRegionBuilder",
       (ins),
       [{ return ConcreteOp::getRegionBuilder(); }]
     >
   ];
 
   let extraClassDeclaration = [{
     /// Return the flat list of all operand dimension sizes in the order they

mlir/include/mlir/Dialect/Linalg/IR/LinalgStructuredOps.td

ArrayAttr indexing_maps() {
       return Builder(getContext()).getAffineMapArrayAttr({
           extractOrIdentityMap(maybeInputMap, inputRank, context),
           extractOrIdentityMap(maybeOutputMap, outputRank, context)});
     }
 
     Value getSource() { return input();}
     Value getTarget() { return output(); }
 
-    static void regionBuilder(
-      ImplicitLocOpBuilder &b, Block &block, ValueRange captures);
-    static std::function<
-      void(ImplicitLocOpBuilder &b, Block &block, ValueRange captures)>
+    static void regionBuilder(ImplicitLocOpBuilder &b, Block &block);
+    static std::function<void(ImplicitLocOpBuilder &b, Block &block)>
     getRegionBuilder() {
       return &regionBuilder;
     }
     static unsigned getNumRegionArgs() { return 2; }
   }];
   let verifier = [{ return ::verify(*this); }];
 
   let assemblyFormat = [{
let extraClassDeclaration = structuredOpsDecls # [{
     ArrayAttr indexing_maps() {
       MLIRContext *context = getContext();
       // filling_value -> O(ivs)
       return Builder(getContext()).getAffineMapArrayAttr({
           AffineMap::get(getNumParallelLoops(), 0, {}, getContext()),
           extractOrIdentityMap(llvm::None, getNumParallelLoops(), context)});
     }
 
-    static void regionBuilder(
-      ImplicitLocOpBuilder &b, Block &block, ValueRange captures);
-    static std::function<
-      void(ImplicitLocOpBuilder &b, Block &block, ValueRange captures)>
+    static void regionBuilder(ImplicitLocOpBuilder &b, Block &block);
+    static std::function<void(ImplicitLocOpBuilder &b, Block &block)>
     getRegionBuilder() {
       return &regionBuilder;
     }
     static unsigned getNumRegionArgs() { return 2; }
   }];
 
   let assemblyFormat = [{
     `(` $value `,` $output `)` attr-dict `:`
code commonUtils = structuredOpsDecls # [{
 
     int64_t getHighPad(unsigned i) {
       assert(i < getNumWindowLoops());
       if (!padding().hasValue()) return 0;
       return padding().getValue().getValue<int64_t>({i, 1});
     }
 
     static std::function<
-      void(ImplicitLocOpBuilder &b, Block &block, ValueRange captures)>
+      void(ImplicitLocOpBuilder &b, Block &block)>
     getRegionBuilder() {
       return nullptr;
     }
   }];
 }
 
 // Only support buffer semantics.
 def ConvOp : PoolingBase_Op<"conv", []> {
SmallVector<StringRef, 8> linalgTraitAttrNames() {
       };
     }
     std::string getLibraryCallName() {
       return library_call().hasValue() ?
         library_call()->str() : "op_has_no_registered_library_name";
     }
 
     static std::function<
-      void(ImplicitLocOpBuilder &b, Block &block, ValueRange captures)>
+      void(ImplicitLocOpBuilder &b, Block &block)>
     getRegionBuilder() {
       return nullptr;
     }
   }];
   let printer = [{ return ::print(p, *this); }];
   let parser = [{ return ::parseGenericOp(parser, result); }];
 }
 

mlir/lib/Bindings/Python/DialectLinalg.cpp

 
 namespace py = pybind11;
 using namespace mlir;
 using namespace mlir::python;
 
 void mlir::python::populateDialectLinalgSubmodule(py::module m) {
   m.def(
       "fill_builtin_region",
-      [](PyDialectDescriptor &dialect, PyOperation &op, py::list captures) {
-        llvm::SmallVector<MlirValue, 4> mlirOperands;
-        mlirOperands.reserve(captures.size());
-        for (auto v : captures)
-          mlirOperands.push_back(py::cast<PyValue *>(v)->get());
-        mlirLinalgFillBuiltinNamedOpRegion(
-            dialect.get(), op.get(), mlirOperands.size(), mlirOperands.data());
+      [](PyDialectDescriptor &dialect, PyOperation &op) {
+        mlirLinalgFillBuiltinNamedOpRegion(dialect.get(), op.get());
       },
-      py::arg("dialect"), py::arg("op"), py::arg("captures") = py::list(),
+      py::arg("dialect"), py::arg("op"),
       "Fill the region for `op`, which is assumed to be a builtin named Linalg "
       "op.");
 }

mlir/lib/CAPI/Dialect/Linalg.cpp

 #include "mlir/Dialect/Linalg/IR/LinalgOps.h"
 
 using namespace mlir;
 using namespace mlir::linalg;
 
 /// Apply the special region builder for the builtin named Linalg op.
 /// Assert that `op` is a builtin named Linalg op.
 void mlirLinalgFillBuiltinNamedOpRegion(MlirDialect linalgDialect,
-                                        MlirOperation mlirOp, intptr_t n,
-                                        MlirValue const *mlirCaptures) {
+                                        MlirOperation mlirOp) {
   Operation *op = unwrap(mlirOp);
-  SmallVector<Value> captures;
-  captures.reserve(n);
-  for (unsigned idx = 0; idx < n; ++idx)
-    captures.push_back(unwrap(mlirCaptures[idx]));
 
   LinalgDialect::RegionBuilderFunType fun =
       static_cast<LinalgDialect *>(unwrap(linalgDialect))
           ->getRegionBuilder(op->getName().getStringRef());
   assert(fun && "Expected a builtin named Linalg op.");
   assert(op->getNumRegions() == 1 && "Expected Linalg op with 1 region");
   assert(op->getRegion(0).getBlocks().empty() &&
          "Expected Linalg op with 0 blocks");
 
   SmallVector<Type, 8> argTypes;
   auto linalgOp = cast<LinalgOp>(op);
   for (OpOperand *opOperand : linalgOp.getInputAndOutputOperands())
     argTypes.push_back(getElementTypeOrSelf(opOperand->get().getType()));
 
   ImplicitLocOpBuilder b(op->getLoc(), op->getContext());
   Region &region = op->getRegion(0);
   Block *body = b.createBlock(&region, /*insertPt=*/{}, argTypes);
   b.setInsertionPointToStart(body);
-  fun(b, *body, captures);
+  fun(b, *body);
 }
 
 MLIR_DEFINE_CAPI_DIALECT_REGISTRATION(Linalg, linalg, LinalgDialect)

mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp

 
 /// Forward declarations.
 
 /// Generic entry point to create the block for the region of a LinalgOp.
 /// This is used by both named structured ops created by ods-gen and by manually
 /// defined C++ ops.
 /// This is used by both builders and parsers.
 /// This function creates the block in the region with arguments corresponding
-/// to the elemental types of `inputTypes` and `outputTypes`, which are asserted
-/// to be ShapedType.
+/// to the elemental types of `inputTypes` and `outputTypes`. The latter are
+/// asserted to be of ShapedType.
 template <typename NamedStructuredOpType>
 static void fillStructuredOpRegion(
     OpBuilder &opBuilder, Region &region, TypeRange inputTypes,
-    TypeRange outputTypes, ValueRange captures = {},
+    TypeRange outputTypes,
     std::function<void(unsigned, unsigned)> errorHandler = nullptr);
 
 /// Generic entry point to create both the region and the block of a LinalgOp.
 template <typename NamedStructuredOpType>
 static void
 createAndFillStructuredOpRegion(OpBuilder &opBuilder, OperationState &result,
-                                TypeRange inputTypes, TypeRange outputTypes,
-                                ValueRange captures = {});
+                                TypeRange inputTypes, TypeRange outputTypes);
 
 /// Common parsing and printing used for both named structured ops created by
 /// ods-gen and by manually defined C++ ops. Does not handle regions.
 static ParseResult
 parseCommonStructuredOpParts(OpAsmParser &parser, OperationState &result,
                              SmallVectorImpl<Type> &inputTypes,
                              SmallVectorImpl<Type> &outputTypes);
 template <typename NamedStructuredOpType>
 static void printCommonStructuredOpParts(OpAsmPrinter &p,
                                          NamedStructuredOpType op);
 
 /// Specific parsing and printing for named structured ops created by ods-gen.
 template <typename NamedStructuredOpType>
 static ParseResult
 parseNamedStructuredOpRegion(OpAsmParser &parser, Region &region,
-                             TypeRange inputTypes, TypeRange outputTypes,
-                             ArrayRef<OpAsmParser::OperandType> captures = {});
+                             TypeRange inputTypes, TypeRange outputTypes);
 
 static ParseResult
 parseNamedStructuredOpResults(OpAsmParser &parser,
                               SmallVectorImpl<Type> &resultTypes);
 
 template <typename NamedStructuredOpType>
-static ParseResult
-parseNamedStructuredOp(OpAsmParser &parser, OperationState &result,
-                       ArrayRef<OpAsmParser::OperandType> captures = {});
+static ParseResult parseNamedStructuredOp(OpAsmParser &parser,
+                                          OperationState &result);
 
 static void printNamedStructuredOpResults(OpAsmPrinter &p,
                                           TypeRange resultTypes);
 
 template <typename NamedStructuredOpType>
 static void printNamedStructuredOp(OpAsmPrinter &p, NamedStructuredOpType op);
 
 /// Helper function to convert a Value into an OpFoldResult, if the Value is
private:
   }
 };
 
 } // namespace
 
 //===----------------------------------------------------------------------===//
 // CopyOp
 //===----------------------------------------------------------------------===//
-void CopyOp::regionBuilder(ImplicitLocOpBuilder &b, Block &block,
-                           ValueRange captures) {
+void CopyOp::regionBuilder(ImplicitLocOpBuilder &b, Block &block) {
   assert(block.getNumArguments() == 2 && "CopyOp regionBuilder expects 2 args");
   b.create<linalg::YieldOp>(block.getArgument(0));
 }
 
 void CopyOp::build(OpBuilder &builder, OperationState &result, Value input,
                    Value output, AffineMap inputPermutation,
                    AffineMap outputPermutation,
                    ArrayRef<NamedAttribute> namedAttrs) {
effects.emplace_back(MemoryEffects::Read::get(), input(),
                        SideEffects::DefaultResource::get());
   effects.emplace_back(MemoryEffects::Write::get(), output(),
                        SideEffects::DefaultResource::get());
 }
 
 //===----------------------------------------------------------------------===//
 // FillOp
 //===----------------------------------------------------------------------===//
-void FillOp::regionBuilder(ImplicitLocOpBuilder &b, Block &block,
-                           ValueRange captures) {
+void FillOp::regionBuilder(ImplicitLocOpBuilder &b, Block &block) {
   assert(block.getNumArguments() == 2 && "FillOp regionBuilder expects 2 args");
   b.create<linalg::YieldOp>(block.getArgument(0));
 }
 
 void FillOp::build(OpBuilder &builder, OperationState &result, Value value,
                    Value output) {
   build(builder, result, output.getType().dyn_cast<RankedTensorType>(), value,
         output);

 /// This is used by both builders and parsers.
 /// This function creates the block in the region with arguments corresponding
 /// to the elemental types of `inputTypes` and `outputTypes`, which are asserted
 /// to be ShapedType.
 template <typename NamedStructuredOpType>
 static void
 fillStructuredOpRegion(OpBuilder &opBuilder, Region &region,
                        TypeRange inputTypes, TypeRange outputTypes,
-                       ValueRange captures,
                        std::function<void(unsigned, unsigned)> errorHandler) {
   assert(llvm::all_of(outputTypes, [](Type t) { return t.isa<ShapedType>(); }));
 
   // TODO: atm all operands go through getElementTypeOrSelf,
   // reconsider when we have evidence we need to.
   SmallVector<Type, 8> argTypes;
   for (auto containers : {inputTypes, outputTypes})
     for (auto t : containers)
       argTypes.push_back(getElementTypeOrSelf(t));
 
   // RAII.
   OpBuilder::InsertionGuard guard(opBuilder);
   Block *body = opBuilder.createBlock(&region, /*insertPt=*/{}, argTypes);
   unsigned actual = body->getNumArguments();
   unsigned expected = NamedStructuredOpType::getNumRegionArgs();
   if (expected != actual) {
     if (errorHandler)
       errorHandler(expected, actual);
     return;
   }
 
   opBuilder.setInsertionPointToStart(body);
   ImplicitLocOpBuilder b(opBuilder.getUnknownLoc(), opBuilder);
-  NamedStructuredOpType::regionBuilder(b, *body, captures);
+  NamedStructuredOpType::regionBuilder(b, *body);
 
   // indexing_maps is an auto-generated method.
 
   // iterator_types is an auto-generated method.
 }
 
 /// Generic entry point to create both the region and the block of a LinalgOp.
 template <typename NamedStructuredOpType>
 void createAndFillStructuredOpRegion(OpBuilder &opBuilder,
                                      OperationState &result,
                                      TypeRange inputTypes,
-                                     TypeRange outputTypes,
-                                     ValueRange captures) {
+                                     TypeRange outputTypes) {
   Region &region = *result.addRegion();
   fillStructuredOpRegion<NamedStructuredOpType>(
-      opBuilder, region, inputTypes, outputTypes, captures,
+      opBuilder, region, inputTypes, outputTypes,
       [&](unsigned expected, unsigned actual) {
         assert(expected != actual && "incorrect number of arguments");
       });
 }
 
 /// Common parsing used for both named structured ops created by ods-gen and by
 /// manually defined C++ ops. Does not handle regions.
 static ParseResult

 
 //===----------------------------------------------------------------------===//
 // Specific parsing and printing for named structured ops created by ods-gen.
 //===----------------------------------------------------------------------===//
 
 template <typename NamedStructuredOpType>
 static ParseResult
 parseNamedStructuredOpRegion(OpAsmParser &parser, Region &region,
-                             TypeRange inputTypes, TypeRange outputTypes,
-                             ArrayRef<OpAsmParser::OperandType> captures) {
+                             TypeRange inputTypes, TypeRange outputTypes) {
   ParseResult res = success();
   OpBuilder opBuilder(parser.getBuilder().getContext());
   // Resolve `captures` into `capturedValues` at parse time so we can build the
   // region with captures.
   SmallVector<Value> capturedValues;
   fillStructuredOpRegion<NamedStructuredOpType>(
-      opBuilder, region, inputTypes, outputTypes, capturedValues,
+      opBuilder, region, inputTypes, outputTypes,
       [&](unsigned expected, unsigned actual) {
         res = parser.emitError(
             parser.getCurrentLocation(),
             llvm::formatv("[parseNamedStructuredOpRegion] ods-gen generated "
                           "region expects {0} args, got {1}",
                           expected, actual));
         region.front().dump();
       });
   return res;
 }
 
 static ParseResult
 parseNamedStructuredOpResults(OpAsmParser &parser,
                               SmallVectorImpl<Type> &resultTypes) {
   if (parser.parseOptionalArrowTypeList(resultTypes))
     return failure();
   return success();
 }
 
 template <typename NamedStructuredOpType>
-static ParseResult
-parseNamedStructuredOp(OpAsmParser &parser, OperationState &result,
-                       ArrayRef<OpAsmParser::OperandType> captures) {
+static ParseResult parseNamedStructuredOp(OpAsmParser &parser,
+                                          OperationState &result) {
   // TODO: Enable when ods-gen supports captures.
-  assert(captures.empty() && "unexpected captures for named structured ops");
   SmallVector<Type, 1> inputTypes, outputTypes;
   if (parseCommonStructuredOpParts(parser, result, inputTypes, outputTypes))
     return failure();
 
   // TODO: consider merging results parsing into region parsing.
   // Need to wait for declarative assembly resolution to decide.
   SmallVector<Type, 1> outputTensorsTypes;
   if (parseNamedStructuredOpResults(parser, outputTensorsTypes))
     return failure();
   result.addTypes(outputTensorsTypes);
 
   std::unique_ptr<Region> region = std::make_unique<Region>();
   if (parseNamedStructuredOpRegion<NamedStructuredOpType>(
-          parser, *region, inputTypes, outputTypes, captures))
+          parser, *region, inputTypes, outputTypes))
     return failure();
   result.addRegion(std::move(region));
 
   return success();
 }
 
 static void printNamedStructuredOpResults(OpAsmPrinter &p,
                                           TypeRange resultTypes) {

mlir/lib/Dialect/Linalg/Transforms/Generalization.cpp

if (!regionBuilder) {
     LLVM_DEBUG(llvm::dbgs() << "no region builder for op: " << namedOp << "\n");
     return nullptr;
   }
   return rewriter.create<GenericOp>(
       namedOp.getLoc(), types, inputOperands, outputOperands, indexingMaps,
       iterators,
       [&regionBuilder](OpBuilder &bodyBuilder, Location loc, ValueRange) {
         ImplicitLocOpBuilder b(loc, bodyBuilder);
-        regionBuilder(b, *bodyBuilder.getBlock(),
-                      /*captures=*/{});
+        regionBuilder(b, *bodyBuilder.getBlock());
       });
 }
 
 namespace {
 
 /// Base class for all linalg generalization patterns. A subclass must provide
 /// the following method:
 ///   GenericOp createGenericOp(RootOp, PatternRewriter &)

mlir/python/mlir/dialects/_linalg_ops_ext.py

def __init__(self, output: Value, value: Value, *, loc=None, ip=None):
     op = self.build_generic(
         results=results,
         operands=[value, output],
         attributes=None,
         loc=loc,
         ip=ip)
     OpView.__init__(self, op)
     linalgDialect = Context.current.get_dialect_descriptor("linalg")
-    fill_builtin_region(linalgDialect, self.operation, [])
+    fill_builtin_region(linalgDialect, self.operation)
     # TODO: self.result is None. When len(results) == 1 we expect it to be
     # results[0] as per _linalg_ops_gen.py. This seems like an orthogonal bug
     # in the generator of _linalg_ops_gen.py where we have:
     # ```
     # def result(self):
     #   return self.operation.results[0] \
     #     if len(self.operation.results) > 1 else None
     # ```

mlir/test/mlir-linalg-ods-gen/test-linalg-ods-gen.tc

 //  IMPL-NEXT: map1 = map1.replaceDimsAndSymbols({}, { s0, s1 }, 2, 0);
 //  IMPL-NEXT: map1 = simplifyAffineMap(map1);
 //  IMPL-NEXT: auto map2 = AffineMap::get(2, 2, {d0}, context);
 //  IMPL-NEXT: map2 = map2.replaceDimsAndSymbols({}, { s0, s1 }, 2, 0);
 //  IMPL-NEXT: map2 = simplifyAffineMap(map2);
 //  IMPL-NEXT: return {{.+}}.getAffineMapArrayAttr({ map0, map1, map2 });
 //
 //       IMPL:  void Test1Op::regionBuilder(ImplicitLocOpBuilder &b,
-//       IMPL:    Block &block, ValueRange captures) {
+//       IMPL:    Block &block) {
 //       IMPL:  Value [[a:.*]](args[0]), [[b:.*]](args[1]), [[c:.*]](args[2]);
 //       IMPL:  Value [[d:.*]] = b.create<MulFOp>([[a]], [[b]]);
 //       IMPL:  Value [[e:.*]] = b.create<AddFOp>([[c]], [[d]]);
 //       IMPL:  b.create<linalg::YieldOp>(ValueRange{ [[e]] });
 //
 ods_def<Test1Op> :
 def test1(A: f32(M, K), B: f32(K)) -> (C: f32(M)) {
   C(m) = AddFOp<k>(C(m), MulFOp(A(m, k), B(k)));
 }
 
 // ODS-LABEL: def Test2Op : LinalgStructuredBase_Op<"test2", [
 //  ODS-NEXT:   AttrSizedOperandSegments
 //  ODS-NEXT:   DeclareOpInterfaceMethods<MemoryEffectsOpInterface>,
 //  ODS-NEXT:   SingleBlockImplicitTerminator<"YieldOp">
 //
 // IMPL-LABEL:  ArrayAttr Test2Op::iterator_types() {
 //       IMPL:  { {{.*}}Parallel{{.*}}, {{.*}}Parallel{{.*}}, {{.*}}Reduction{{.*}} }
 //
 //       IMPL:  ArrayAttr Test2Op::indexing_maps() {
 //       IMPL:  AffineMap::get(3, 3, {d0, d2}, context)
 //       IMPL:  AffineMap::get(3, 3, {d2, d1}, context)
 //       IMPL:  AffineMap::get(3, 3, {d0, d1}, context)
 //
 //       IMPL:  Test2Op::regionBuilder(ImplicitLocOpBuilder &b,
-//       IMPL:    Block &block, ValueRange captures) {
+//       IMPL:    Block &block) {
 //       IMPL:  Value [[a:.*]](args[0]), [[b:.*]](args[1]), [[c:.*]](args[2]);
 //       IMPL:  Value [[d:.*]] = b.create<MulFOp>([[a]], [[b]]);
 //       IMPL:  Value [[e:.*]] = b.create<AddFOp>([[c]], [[d]]);
 //       IMPL:  b.create<linalg::YieldOp>(ValueRange{ [[e]] });
 //
 ods_def<Test2Op> :
 def test2(A: f32(M, K), B: f32(K, N)) -> (C: f32(M, N)) {
   C(m, n) = AddFOp<k>(C(m, n), MulFOp(A(m, k), B(k, n)));
 }
 
 // ODS-LABEL: def Test3Op : LinalgStructuredBase_Op<"test3", [
 //  ODS-NEXT:   AttrSizedOperandSegments
 //  ODS-NEXT:   DeclareOpInterfaceMethods<MemoryEffectsOpInterface>,
 //  ODS-NEXT:   SingleBlockImplicitTerminator<"YieldOp">
 //
 // IMPL-LABEL:  ArrayAttr Test3Op::iterator_types() {
 //       IMPL:  { {{.*}}Parallel{{.*}}, {{.*}}Parallel{{.*}}, {{.*}}Reduction{{.*}} }
 //
 //       IMPL:  ArrayAttr Test3Op::indexing_maps() {
 //       IMPL:  AffineMap::get(4, 4, {d0, d1, d3}, context)
 //       IMPL:  AffineMap::get(4, 4, {d3, d2}, context)
 //       IMPL:  AffineMap::get(4, 4, {d0, d1, d2}, context)
 //
 //       IMPL:  Test3Op::regionBuilder(ImplicitLocOpBuilder &b,
-//       IMPL:    Block &block, ValueRange captures) {
+//       IMPL:    Block &block) {
 //       IMPL:  Value [[a:.*]](args[0]), [[b:.*]](args[1]), [[c:.*]](args[2]);
 //       IMPL:  Value [[d:.*]] = b.create<MulFOp>([[a]], [[b]]);
 //       IMPL:  Value [[e:.*]] = b.create<AddFOp>([[c]], [[d]]);
 //       IMPL:  b.create<linalg::YieldOp>(ValueRange{ [[e]] });
 //
 ods_def<Test3Op> :
 def test3(A: f32(Batch, M, K), B: f32(K, N)) -> (C: f32(Batch, M, N)) {
   C(b, m, n) = AddFOp<k>(C(b, m, n), MulFOp(A(b, m, k), B(k, n)));

 def test7(A: f32(M, K), B: f32(K)) -> (C: f32(M))
      attr(attr_a: f32, attr_b: 4xi32)
 {
   C(m) = AddFOp<k>(C(m), MulFOp(A(m, k), B(k)));
 }
 
 // Test output arg order.
 // IMPL-LABEL:  void Test8Op::regionBuilder(ImplicitLocOpBuilder &b,
-//       IMPL:    Block &block, ValueRange captures) {
+//       IMPL:    Block &block) {
 //       IMPL:  Value [[a:.*]](args[0]), [[b:.*]](args[1]), [[c:.*]](args[2]);
 //       IMPL:  Value [[d:.*]] = b.create<MulFOp>([[a]], [[b]]);
 //       IMPL:  Value [[e:.*]] = b.create<SubFOp>([[d]], [[c]]);
 //       IMPL:  b.create<linalg::YieldOp>(ValueRange{ [[e]] });
 ods_def<Test8Op>:
 def test8(A: f32(M, K), B: f32(K)) -> (C: f32(M))
 {
   C(m) = SubFOp<k>(MulFOp(A(m, k), B(k)), C(m));
 }
 
 // Test shape-only operand.
 // IMPL-LABEL:  ArrayAttr Test9Op::indexing_maps() {
 //       IMPL:    auto map0 = AffineMap::get(2, 2, {d0, d1}, context);
 //       IMPL:    auto map1 = AffineMap::get(2, 2, {d1}, context);
 //       IMPL:    auto map2 = AffineMap::get(2, 2, {d0}, context);
 // IMPL-LABEL:  void Test9Op::regionBuilder(ImplicitLocOpBuilder &b,
-//       IMPL:    Block &block, ValueRange captures) {
+//       IMPL:    Block &block) {
 //       IMPL:  Value [[a:.*]](args[0]), [[c:.*]](args[2]);
 ods_def<Test9Op>:
 def test9(A: f32(M, K), B: f32(K)) -> (C: f32(M))
 {
   C(m) = AddFOp<k>(C(m), A(m, k));
 }

mlir/test/mlir-linalg-ods-gen/test-linalg-ods-yaml-gen.yaml

 #  ODS-NEXT:        static_cast<int32_t>(outputs.size())}));
 #  ODS-NEXT:    createAndFillStructuredOpRegion<Test1Op>(
 #  ODS-NEXT:      $_builder,
 #  ODS-NEXT:      $_state,
 #  ODS-NEXT:      TypeRange(inputs),
 #  ODS-NEXT:      TypeRange(outputs)
 
 # IMPL-LABEL:  void Test1Op::regionBuilder(
-#       IMPL:    ImplicitLocOpBuilder &b, Block &block, ValueRange captures)
+#       IMPL:    ImplicitLocOpBuilder &b, Block &block)
 #       IMPL:  Value [[VAL0:[a-z0-9]+]] = helper.constant("42 : i64");
 #   IMPL-DAG:  Value [[VAL1:[a-z0-9]+]] = helper.cast(block.getArgument(0).getType(), [[VAL0]]);
 #   IMPL-DAG:  Value [[VAL2:[a-z0-9]+]] = helper.index(1);
 #   IMPL-DAG:  Value [[VAL3:[a-z0-9]+]] = helper.cast(block.getArgument(0).getType(), [[VAL2]]);
 #   IMPL-DAG:  Value [[VAL4:[a-z0-9]+]] = helper.applyfn__add([[VAL1]], [[VAL3]]);
 
 
 # @linalg_structured_op

 #  IMPL-NEXT:  { getParallelIteratorTypeName(), getParallelIteratorTypeName() }
 
 #       IMPL:  Test2Op::indexing_maps()
 #       IMPL:  "affine_map<(d0, d1)[s0, s1] -> ()>"
 #       IMPL:  "affine_map<(d0, d1)[s0, s1] -> (d1, d0)>"
 #       IMPL:  "affine_map<(d0, d1)[s0, s1] -> (d0, d1)>"
 
 #       IMPL:  void Test2Op::regionBuilder(
-#       IMPL:    ImplicitLocOpBuilder &b, Block &block, ValueRange captures)
+#       IMPL:    ImplicitLocOpBuilder &b, Block &block)
 
 #       IMPL:   = helper.applyfn__add(block.getArgument(0), block.getArgument(1));

mlir/tools/mlir-linalg-ods-gen/mlir-linalg-ods-gen.cpp

/*extraInterfaces=*/{2}]> {
             "operand_segment_sizes",
             $_builder.getI32VectorAttr({{
               static_cast<int32_t>(inputs.size()),
               static_cast<int32_t>(outputs.size())}));
           createAndFillStructuredOpRegion<{0}>(
             $_builder,
             $_state,
             TypeRange(inputs),
-            TypeRange(outputs)/*, TODO: support captures*/);
+            TypeRange(outputs));
         }]>,
         OpBuilder<
         (ins "TypeRange":$resultTensorTypes, "ValueRange":$inputs,
              "ValueRange":$outputs),
         [{{
           $_state.addOperands(inputs);
           $_state.addOperands(outputs);
           $_state.addTypes(resultTensorTypes);
           $_state.addAttribute(
             "operand_segment_sizes",
             $_builder.getI32VectorAttr({{
               static_cast<int32_t>(inputs.size()),
               static_cast<int32_t>(outputs.size())}));
           createAndFillStructuredOpRegion<{0}>(
             $_builder,
             $_state,
             TypeRange(inputs),
-            TypeRange(outputs)/*, TODO: support captures*/);
+            TypeRange(outputs));
         }]>,
         OpBuilder<
         (ins "TypeRange":$resultTensorTypes, "ValueRange":$operands,
              CArg<"ArrayRef<NamedAttribute>", "{{}">:$attributes),
         [{{
           $_state.addOperands(operands);
           $_state.addAttributes(attributes);
           $_state.addTypes(resultTensorTypes);
           (void)$_state.addRegion();
         }]>
         {6}
       ];
       let printer = [{{ return ::printNamedStructuredOp(p, *this); }];
       let parser = [{{
-        return ::parseNamedStructuredOp<{0}>(parser, result/*TODO:, captures*/);
+        return ::parseNamedStructuredOp<{0}>(parser, result);
       }];
       let hasFolder = 1;
 
       let extraClassDeclaration = structuredOpsBaseDecls # [{{
         // Auto-generated.
         ArrayAttr iterator_types();
         ArrayAttr indexing_maps();
-        static void regionBuilder(ImplicitLocOpBuilder &b,
-                                  Block &block, ValueRange captures);
-        static std::function<void(ImplicitLocOpBuilder &b,
-                                  Block &, ValueRange)> getRegionBuilder() {{
+        static void regionBuilder(ImplicitLocOpBuilder &b, Block &block);
+        static std::function<void(ImplicitLocOpBuilder &b, Block &)>
+        getRegionBuilder() {{
           return regionBuilder;
         }
 
         // Generic methods.
         static unsigned getNumRegionArgs() {{ return {5}; }
         std::string getLibraryCallName() {{
           return generateLibraryCallName(getOperation());
         }
const char *builderFmt = R"FMT(
           "operand_segment_sizes",
           $_builder.getI32VectorAttr({{
             static_cast<int32_t>(inputs.size()),
             static_cast<int32_t>(outputs.size())}));
         createAndFillStructuredOpRegion<{0}>(
           $_builder,
           $_state,
           TypeRange(inputs),
-          TypeRange(outputs)/*, TODO: support captures*/);
+          TypeRange(outputs));
         {2}
       }]>
     )FMT";
     attrBuilder =
         llvm::formatv(builderFmt, cppOpName, attrParamsList, attrStmtsList);
   }
 
   std::string attrMethods;
if (subExprsMap.count(pTensorExpr) > 0) {
       const char *tensorExprFmt = "\n    Value _{0} = b.create<{1}>({2});";
       os << llvm::formatv(tensorExprFmt, ++count, pTensorExpr->operationName,
                           subExprs);
       subExprsMap[pTensorExpr] = count;
     }
   };
 
   const char *regionBuilderFmt = R"FMT(
-  void {0}::regionBuilder(ImplicitLocOpBuilder &b,
-                          Block &block, ValueRange captures) {
+  void {0}::regionBuilder(ImplicitLocOpBuilder &b, Block &block) {
     auto args = block.getArguments();
     Value {1};
     {2}
     b.create<linalg::YieldOp>(ValueRange{ {3} });
   })FMT";
 
   std::string valueHandleStr;
   llvm::raw_string_ostream valueHandleStringStream(valueHandleStr);

mlir/tools/mlir-linalg-ods-gen/mlir-linalg-ods-yaml-gen.cpp

let parser = [{{
       return ::parseNamedStructuredOp<{0}>(parser, result);
     }];
     let hasFolder = 1;
 
     let extraClassDeclaration = structuredOpsBaseDecls # [{{
       // Auto-generated.
       ArrayAttr iterator_types();
       ArrayAttr indexing_maps();
-      static void regionBuilder(
-        ImplicitLocOpBuilder &b, Block &block, ValueRange captures);
-      static std::function<
-        void(ImplicitLocOpBuilder &b, Block &, ValueRange)>
+      static void regionBuilder(ImplicitLocOpBuilder &b, Block &block);
+      static std::function<void(ImplicitLocOpBuilder &b, Block &)>
       getRegionBuilder() {{
         return regionBuilder;
       }
 
       // Generic methods.
       static unsigned getNumRegionArgs();
       std::string getLibraryCallName();
       {7}
)FMT";
 
   // regionBuilder()
   {
     // Generates a regionBuilder method. Parameters.
     // {0}: Class name
     // {1}: Number of args
     // {2}: Statements
     static const char structuredOpRegionBuilderFormat[] = R"FMT(
-void {0}::regionBuilder(
-    ImplicitLocOpBuilder &b, Block &block, ValueRange captures) {{
+void {0}::regionBuilder(ImplicitLocOpBuilder &b, Block &block) {{
   assert({1} > 0 && block.getNumArguments() == {1} &&
          "{0} regionBuilder expects {1} (>=0) args");
   RegionBuilderHelper helper(block.getArgument(0).getContext(), block);
   SmallVector<Value> yields;
   {2}
   helper.yieldOutputs(yields);
 }
 )FMT";