[mlir][Linalg] Improve region support in Linalg ops

This revision takes advantage of the newly extended `ref` directive in assembly format
to allow better region handling for LinalgOps. Specifically, FillOp and CopyOp now build their regions explicitly which allows retiring older behavior that relied on specific op knowledge in both lowering to loops and vectorization.

This reverts commit 3f22547 and reland 973e133 with a workaround for
a gcc bug that does not accept lambda default parameters:
https://gcc.gnu.org/bugzilla/show_bug.cgi?id=59949

Differential Revision: https://reviews.llvm.org/D96598

Author: joker-eph

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

@@ -1056,20 +1056,6 @@ def LinalgStructuredInterface : OpInterface<"LinalgOp"> {
     //===------------------------------------------------------------------===//
     // Other static interface methods.
     //===------------------------------------------------------------------===//
-    StaticInterfaceMethod<
-      /*desc=*/[{
-        Create an operation of the current type with the given location,
-        operands, and attributes.
-      }],
-      /*retTy=*/"Operation *",
-      /*methodName=*/"create",
-      (ins "OpBuilder &":$builder, "Location":$loc, "TypeRange":$resultTypes,
-           "ValueRange":$operands,
-           "ArrayRef<NamedAttribute>":$attributes), [{
-        return builder.create<ConcreteOp>(
-          loc, resultTypes, operands, attributes);
-      }]
-    >,
     InterfaceMethod<
       /*desc=*/[{
         Clone the current operation with the given location and operands. This
@@ -1082,14 +1068,13 @@ def LinalgStructuredInterface : OpInterface<"LinalgOp"> {
       (ins "OpBuilder &":$b, "Location":$loc, "TypeRange":$resultTypes,
            "ValueRange":$operands),
       [{
-        BlockAndValueMapping map;
-        unsigned numRegions = $_op->getNumRegions();
-        Operation *res = create(b, loc, resultTypes, operands, $_op->getAttrs());
-        assert(res->getNumRegions() == numRegions && "inconsistent # regions");
-        for (unsigned ridx = 0; ridx < numRegions; ++ridx)
-          $_op->getRegion(ridx).cloneInto(
-            &res->getRegion(ridx), map);
-        return res;
+        BlockAndValueMapping bvm;
+        OperationState state(
+          loc, ConcreteOp::getOperationName(), operands, resultTypes,
+          $_op->getAttrs());
+        for (Region &r : $_op->getRegions())
+          r.cloneInto(state.addRegion(), bvm);
+        return b.createOperation(state);
       }]
     >,
     StaticInterfaceMethod<
@@ -1098,7 +1083,7 @@ def LinalgStructuredInterface : OpInterface<"LinalgOp"> {
         Returns a null function if this named op does not define a region
         builder.
       }],
-      /*retTy=*/"std::function<void(Block &)>",
+      /*retTy=*/"std::function<void(Block &, ValueRange)>",
       /*methodName=*/"getRegionBuilder",
       (ins),
       [{ return ConcreteOp::getRegionBuilder(); }]

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

@@ -110,14 +110,13 @@ def CopyOp : LinalgStructured_Op<"copy", [CopyOpInterface]> {
     AnyStridedMemRef:$output,
     OptionalAttr<AffineMapAttr>:$inputPermutation,
     OptionalAttr<AffineMapAttr>:$outputPermutation);
+  let regions = (region AnyRegion:$region);
 
-  // TODO: this should go away once the usage of OptionalAttr triggers emission
-  // of builders with default arguments left unspecified.
-  let builders = [OpBuilderDAG<(ins "Value":$input, "Value":$output),
-    [{
-      return build(
-        $_builder, $_state, input, output, AffineMapAttr(), AffineMapAttr());
-    }]>];
+  let builders = [
+    OpBuilderDAG<(ins "Value":$input, "Value":$output,
+      CArg<"AffineMap", "AffineMap()">:$inputPermutation,
+      CArg<"AffineMap", "AffineMap()">:$outputPermutation,
+      CArg<"ArrayRef<NamedAttribute>", "{}">:$attrs)>];
 
   let extraClassDeclaration = structuredOpsDecls # [{
     ValueRange inputs() { return getOperands().take_front(); }
@@ -146,24 +145,31 @@ def CopyOp : LinalgStructured_Op<"copy", [CopyOpInterface]> {
     Value getSource() { return input();}
     Value getTarget() { return output(); }
 
-    static std::function<void(Block &)> getRegionBuilder() {
-      return nullptr;
+    static void regionBuilder(Block &block, ValueRange captures);
+    static std::function<void(Block &block, ValueRange captures)>
+    getRegionBuilder() {
+      return &regionBuilder;
     }
+    static unsigned getNumRegionArgs() { return 2; }
   }];
   let verifier = [{ return ::verify(*this); }];
 
   let assemblyFormat = [{
-    `(` operands `)` attr-dict `:` type(operands)
+    `(` $input `,` $output `)` attr-dict `:`
+        type($input) `,` type($output)
+      custom<CopyOpRegion>($region, ref(type($input)), ref(type($input)))
   }];
 
   let hasFolder = 1;
   let hasCanonicalizer = 1;
+  let skipDefaultBuilders = 1;
 }
 
 def FillOp : LinalgStructured_Op<"fill", []> {
   let arguments = (ins AnyShaped:$output,
                    AnyTypeOf<[AnyFloat, AnySignlessInteger, AnyVector]>:$value);
   let results = (outs Optional<AnyRankedTensor>:$result);
+  let regions = (region AnyRegion:$region);
   let extraClassDeclaration = structuredOpsDecls # [{
     ValueRange inputs() { return {}; }
     ValueRange outputs() { return getOperands().take_front(); }
@@ -183,13 +189,18 @@ def FillOp : LinalgStructured_Op<"fill", []> {
           extractOrIdentityMap(llvm::None, getNumParallelLoops(), context)});
     }
 
-    static std::function<void(Block &)> getRegionBuilder() {
-      return nullptr;
+    static void regionBuilder(Block &block, ValueRange captures);
+    static std::function<void(Block &block, ValueRange captures)>
+    getRegionBuilder() {
+      return &regionBuilder;
     }
+    static unsigned getNumRegionArgs() { return 1; }
   }];
 
   let assemblyFormat = [{
-    `(` operands `)` attr-dict `:` type(operands) (`->` type($result)^)?
+    `(` $output `,` $value `)` attr-dict `:`
+        type($output) `,` type($value) (`->` type($result)^)?
+      custom<FillOpRegion>($region, ref(type($output)), ref($value))
   }];
 
   let builders = [
@@ -268,7 +279,8 @@ class PoolingBase_Op<string mnemonic, list<OpTrait> props>
       return padding().getValue().getValue<int64_t>({i, 1});
     }
 
-    static std::function<void(Block &)> getRegionBuilder() {
+    static std::function<void(Block &, ValueRange captures)> getRegionBuilder()
+    {
       return nullptr;
     }
   }];
@@ -519,7 +531,7 @@ class GenericOpBase<string mnemonic> : LinalgStructuredBase_Op<mnemonic, [
         library_call()->str() : "op_has_no_registered_library_name";
     }
 
-    static std::function<void(Block &)> getRegionBuilder() {
+    static std::function<void(Block &, ValueRange)> getRegionBuilder() {
       return nullptr;
     }
   }];

mlir/lib/Conversion/LinalgToStandard/LinalgToStandard.cpp

@@ -154,7 +154,13 @@ LogicalResult mlir::linalg::CopyTransposeRewrite::matchAndRewrite(
   if (in == op.input() && out == op.output())
     return failure();
 
-  rewriter.replaceOpWithNewOp<CopyOp>(op, in, out);
+  auto libraryCallName = getLibraryCallSymbolRef(op, rewriter);
+  if (!libraryCallName)
+    return failure();
+
+  rewriter.replaceOpWithNewOp<mlir::CallOp>(
+      op, libraryCallName.getValue(), TypeRange(),
+      createTypeCanonicalizedMemRefOperands(rewriter, op.getLoc(), {in, out}));
   return success();
 }
 

mlir/lib/Dialect/Linalg/EDSC/Builders.cpp

@@ -27,8 +27,6 @@ Operation *mlir::edsc::makeGenericLinalgOp(
     ArrayRef<StructuredIndexed> outputs, TypeRange resultTensorTypes,
     function_ref<void(ValueRange)> regionBuilder, ArrayRef<Value> otherValues,
     ArrayRef<Attribute> otherAttributes) {
-  OpBuilder &builder = edsc::ScopedContext::getBuilderRef();
-
   // Build maps
   SmallVector<SmallVector<AffineExpr, 4>, 4> exprsList;
   exprsList.reserve(inputs.size() + outputs.size());
@@ -54,13 +52,10 @@ Operation *mlir::edsc::makeGenericLinalgOp(
               resultTensorTypes,
               inputValues,
               outputValues,
-              builder.getAffineMapArrayAttr(maps),
-              builder.getStrArrayAttr(iteratorStrTypes),
-              StringAttr() /*doc*/,
-              StringAttr() /*library_call*/,
-              ArrayAttr() /*sparse*/
-              /* TODO: other attributes in op */
-              )
+              maps,
+              iteratorStrTypes,
+              ""/*doc*/,
+              ""/*library_call*/)
           .getOperation();
   // clang-format on