Move getAsmBlockArgumentNames from OpAsmDialectInterface to OpAsmOpInterface

[mehdi_amini]

Summary

This method is more suitable as an opinterface: it seems intrinsic to
individual instances of the operation instead of the dialect.
Also remove the restriction on the interface being applicable to the entry block only.

[rriddle]

LG. The behavior change of allowing non-entry blocks also LGTM, no valid reason not to allow that IMO.

mlir/include/mlir/IR/OpAsmInterface.td

63

? Makes it a little more obvious.

mlir/include/mlir/IR/OpImplementation.h

1353

We also have getAsmResultNames nested in the dialect interface, should we just remove that now that dialects can provide fallback op interfaces (no need to resolve that in this patch)?

mlir/lib/IR/AsmPrinter.cpp

1017

1017–1019

The behavior change of allowing non-entry blocks is not captured in the patch description.

[mehdi_amini]

mlir/include/mlir/IR/OpAsmInterface.td

63

I copied from the one above, will update both!

mlir/include/mlir/IR/OpImplementation.h

1353

Yes absolutely! I noticed this as well and was planning to look into this later.

mlir/lib/IR/AsmPrinter.cpp

1017–1019

Good point!

[mehdi_amini]

Added comments

[mehdi_amini]

Update description

mlir/include/mlir/IR/OpAsmInterface.td

InterfaceMethod<[{
         ```mlir
           %first_result, %middle_results:2, %0 = "my.op" ...
         ```
       }],
       "void", "getAsmResultNames",
       (ins "::mlir::OpAsmSetValueNameFn":$setNameFn),
       "", ";"
     >,
+    InterfaceMethod<[{
+        Get a special name to use when printing the block arguments for a region
+        immediately nested under this operation.
+      }],
+      "void", "getAsmBlockArgumentNames",
+      (ins
+        "::mlir::Region&":$region,
+        "::mlir::OpAsmSetValueNameFn":$setNameFn
+      ),
+      "", ";"
+    >,
     StaticInterfaceMethod<[{
       Return the default dialect used when printing/parsing operations in
       regions nested under this operation. This allows for eliding the dialect
       prefix from the operation name, for example it would be possible to omit
       the `spv.` prefix from all operations within a SpirV module if this method
       returned `spv`. The default implementation returns an empty string which
       is ignored.
       }],
       "StringRef", "getDefaultDialect", (ins), "", "return \"\";"
     >,
   ];
 }
 
 #endif // MLIR_OPASMINTERFACE

mlir/include/mlir/IR/OpImplementation.h

public:
   virtual AliasResult getAlias(Type type, raw_ostream &os) const {
     return AliasResult::NoAlias;
   }
 
   /// Get a special name to use when printing the given operation. See
   /// OpAsmInterface.td#getAsmResultNames for usage details and documentation.
   virtual void getAsmResultNames(Operation *op,
                                  OpAsmSetValueNameFn setNameFn) const {}
-
-  /// Get a special name to use when printing the entry block arguments of the
-  /// region contained by an operation in this dialect.
-  virtual void getAsmBlockArgumentNames(Block *block,
-                                        OpAsmSetValueNameFn setNameFn) const {}
 };
 } // namespace mlir
 
 //===--------------------------------------------------------------------===//
 // Operation OpAsm interface.
 //===--------------------------------------------------------------------===//
 
 /// The OpAsmOpInterface, see OpAsmInterface.td for more details.
 #include "mlir/IR/OpAsmInterface.h.inc"
 
 #endif

mlir/lib/IR/AsmPrinter.cpp

for (unsigned i = 0, e = namesToUse.size(); i != e; ++i) {
     auto name = StringRef(nameStream.str()).drop_front();
 
     // Overwrite the name.
     valueNames[nameToReplace] = name.copy(usedNameAllocator);
   }
 }
 
 void SSANameState::numberValuesInRegion(Region &region) {
+  auto setBlockArgNameFn = [&](Value arg, StringRef name) {
+    assert(!valueIDs.count(arg) && "arg numbered multiple times");
+    assert(arg.cast<BlockArgument>().getOwner()->getParent() == &region &&
+           "arg not defined in current region");
+    setValueName(arg, name);
+  };
+
+  if (!printerFlags.shouldPrintGenericOpForm()) {
+    if (auto *op = region.getParentOp()) {
+      if (auto asmInterface = dyn_cast<OpAsmOpInterface>(op))
+        asmInterface.getAsmBlockArgumentNames(region, setBlockArgNameFn);
+    }
+  }
+
   // Number the values within this region in a breadth-first order.
   unsigned nextBlockID = 0;
   for (auto &block : region) {
     // Each block gets a unique ID, and all of the operations within it get
     // numbered as well.
     blockIDs[&block] = nextBlockID++;
     numberValuesInBlock(block);
   }
 }
 
 void SSANameState::numberValuesInBlock(Block &block) {
-  auto setArgNameFn = [&](Value arg, StringRef name) {
-    assert(!valueIDs.count(arg) && "arg numbered multiple times");
-    assert(arg.cast<BlockArgument>().getOwner() == &block &&
-           "arg not defined in 'block'");
-    setValueName(arg, name);
-  };
-
-  bool isEntryBlock = block.isEntryBlock();
-  if (isEntryBlock && !printerFlags.shouldPrintGenericOpForm()) {
-    if (auto *op = block.getParentOp()) {
-      if (auto asmInterface = interfaces.getInterfaceFor(op->getDialect()))
-        asmInterface->getAsmBlockArgumentNames(&block, setArgNameFn);
-    }
-  }
-
   // Number the block arguments. We give entry block arguments a special name
   // 'arg'.
+  bool isEntryBlock = block.isEntryBlock();
   SmallString<32> specialNameBuffer(isEntryBlock ? "arg" : "");
   llvm::raw_svector_ostream specialName(specialNameBuffer);
   for (auto arg : block.getArguments()) {
     if (valueIDs.count(arg))
       continue;
     if (isEntryBlock) {
       specialNameBuffer.resize(strlen("arg"));
       specialName << nextArgumentID++;

mlir/test/lib/Dialect/Test/TestDialect.cpp

AliasResult getAlias(Type type, raw_ostream &os) const final {
     return AliasResult::NoAlias;
   }
 
   void getAsmResultNames(Operation *op,
                          OpAsmSetValueNameFn setNameFn) const final {
     if (auto asmOp = dyn_cast<AsmDialectInterfaceOp>(op))
       setNameFn(asmOp, "result");
   }
-
-  void getAsmBlockArgumentNames(Block *block,
-                                OpAsmSetValueNameFn setNameFn) const final {
-    auto op = block->getParentOp();
-    auto arrayAttr = op->getAttrOfType<ArrayAttr>("arg_names");
-    if (!arrayAttr)
-      return;
-    auto args = block->getArguments();
-    auto e = std::min(arrayAttr.size(), args.size());
-    for (unsigned i = 0; i < e; ++i) {
-      if (auto strAttr = arrayAttr[i].dyn_cast<StringAttr>())
-        setNameFn(args[i], strAttr.getValue());
-    }
-  }
 };
 
 struct TestDialectFoldInterface : public DialectFoldInterface {
   using DialectFoldInterface::DialectFoldInterface;
 
   /// Registered hook to check if the given region, which is attached to an
   /// operation that is *not* isolated from above, should be used when
   /// materializing constants.
static ParseResult parsePolyForOp(OpAsmParser &parser, OperationState &result) {
 
   // Parse the body region.
   Region *body = result.addRegion();
   auto &builder = parser.getBuilder();
   SmallVector<Type, 4> argTypes(ivsInfo.size(), builder.getIndexType());
   return parser.parseRegion(*body, ivsInfo, argTypes);
 }
 
+void PolyForOp::getAsmBlockArgumentNames(Region &region,
+                                         OpAsmSetValueNameFn setNameFn) {
+  auto arrayAttr = getOperation()->getAttrOfType<ArrayAttr>("arg_names");
+  if (!arrayAttr)
+    return;
+  auto args = getRegion().front().getArguments();
+  auto e = std::min(arrayAttr.size(), args.size());
+  for (unsigned i = 0; i < e; ++i) {
+    if (auto strAttr = arrayAttr[i].dyn_cast<StringAttr>())
+      setNameFn(args[i], strAttr.getValue());
+  }
+}
+
 //===----------------------------------------------------------------------===//
 // Test removing op with inner ops.
 //===----------------------------------------------------------------------===//
 
 namespace {
 struct TestRemoveOpWithInnerOps
     : public OpRewritePattern<TestOpWithRegionPattern> {
   using OpRewritePattern<TestOpWithRegionPattern>::OpRewritePattern;

mlir/test/lib/Dialect/Test/TestOps.td

def PrettyPrintedRegionOp : TEST_Op<"pretty_printed_region",
   );
 
   let results = (outs AnyType);
   let regions = (region SizedRegion<1>:$region);
   let parser = [{ return ::parse$cppClass(parser, result); }];
   let printer = [{ return ::print(p, *this); }];
 }
 
-def PolyForOp : TEST_Op<"polyfor">
+def PolyForOp : TEST_Op<"polyfor", [OpAsmOpInterface]>
 {
   let summary =  "polyfor operation";
   let description = [{
     Test op with multiple region arguments, each argument of index type.
   }];
-
+  let extraClassDeclaration = [{
+    void getAsmBlockArgumentNames(mlir::Region &region,
+                                  mlir::OpAsmSetValueNameFn setNameFn);
+  }];
   let regions = (region SizedRegion<1>:$region);
   let parser = [{ return ::parse$cppClass(parser, result); }];
 }
 
 //===----------------------------------------------------------------------===//
 // Test OpAsmInterface.
 
 def AsmInterfaceOp : TEST_Op<"asm_interface_op"> {