[mlir] Move value name shadowing functionality into OpAsmOpInterface

[rkayaith]

Summary

The OpAsmPrinter::shadowRegionArgs API is problematic as it allows
value ids to be changed by custom op printers, making it impossible to
know the ID of a particular value without printing ops. Moving the
shadowing functionality to OpAsmOpInterface allows the final value IDs
to be determined during AsmState construction.

depends on D150080

[rkayaith]

friendly ping

mlir/include/mlir/IR/OpAsmInterface.td

InterfaceMethod<[{
       "void", "getAsmBlockArgumentNames",
       (ins
         "::mlir::Region&":$region,
         "::mlir::OpAsmSetValueNameFn":$setNameFn
       ),
       "", "return;"
     >,
     InterfaceMethod<[{
+        Get the value names that a region's arguments should shadow.
+
+        For regions under an `IsolatedFromAbove` operation, the names of
+        `Values` defined above the region may be re-used for region argument
+        names. For example,
+        ```c++
+          setShadowingFn(region.getArgument(0), getOperand(0));
+        ```
+        would print as:
+        ```mlir
+          %foo = ...
+          my.isolated_from_above_op(%foo) ({
+            ^bb0(%foo):  // the outer %foo is shadowed
+              ...
+          })
+        ```
+        This should be combined with a custom printer/parser which elides the
+        shadowing region arguments.
+      }],
+      "void", "getAsmRegionArgumentShadowing",
+      (ins
+        "::mlir::Region&":$region,
+        "::mlir::OpAsmSetRegionArgumentShadowingFn":$setShadowingFn
+      ),
+      "", "return;"
+    >,
+    InterfaceMethod<[{
         Get the name to use for a given block inside a region attached to this
         operation.
 
         For example if this operation has multiple blocks:
 
         ```mlir
           some.op() ({
             ^bb0:

mlir/include/mlir/IR/OpImplementation.h

public:
   /// If 'printEntryBlockArgs' is false, the arguments of the
   /// block are not printed. If 'printBlockTerminator' is false, the terminator
   /// operation of the block is not printed. If printEmptyBlock is true, then
   /// the block header is printed even if the block is empty.
   virtual void printRegion(Region &blocks, bool printEntryBlockArgs = true,
                            bool printBlockTerminators = true,
                            bool printEmptyBlock = false) = 0;
 
-  /// Renumber the arguments for the specified region to the same names as the
-  /// SSA values in namesToUse.  This may only be used for IsolatedFromAbove
-  /// operations.  If any entry in namesToUse is null, the corresponding
-  /// argument name is left alone.
-  virtual void shadowRegionArgs(Region &region, ValueRange namesToUse) = 0;
-
   /// Prints an affine map of SSA ids, where SSA id names are used in place
   /// of dims/symbols.
   /// Operand values must come from single-result sources, and be valid
   /// dimensions/symbol identifiers according to mlir::isValidDim/Symbol.
   virtual void printAffineMapOfSSAIds(AffineMapAttr mapAttr,
                                       ValueRange operands) = 0;
 
   /// Prints an affine expression of SSA ids with SSA id names used instead of
parseOptionalAssignmentList(SmallVectorImpl<Argument> &lhs,
                               SmallVectorImpl<UnresolvedOperand> &rhs) = 0;
 };
 
 //===--------------------------------------------------------------------===//
 // Dialect OpAsm interface.
 //===--------------------------------------------------------------------===//
 
 /// A functor used to set the name of the start of a result group of an
-/// operation. See 'getAsmResultNames' below for more details.
+/// operation. See `OpAsmOpInterface::getAsmResultNames` for more details.
 using OpAsmSetValueNameFn = function_ref<void(Value, StringRef)>;
 
+/// A functor used to set the name shadowing of a region argument. See
+/// `OpAsmOpInterface::getRegionArgumentShadowing` for more details.
+using OpAsmSetRegionArgumentShadowingFn =
+    function_ref<void(BlockArgument, Value)>;
+
 /// A functor used to set the name of blocks in regions directly nested under
 /// an operation.
 using OpAsmSetBlockNameFn = function_ref<void(Block *, StringRef)>;
 
 class OpAsmDialectInterface
     : public DialectInterface::Base<OpAsmDialectInterface> {
 public:
   OpAsmDialectInterface(Dialect *dialect) : Base(dialect) {}

mlir/lib/IR/AsmPrinter.cpp

void printOperand(Value, raw_ostream &os) override {
     // guaranteed to go unused.
     os << "%";
   }
   void printKeywordOrString(StringRef) override {}
   void printResourceHandle(const AsmDialectResourceHandle &) override {}
   void printSymbolName(StringRef) override {}
   void printSuccessor(Block *) override {}
   void printSuccessorAndUseList(Block *, ValueRange) override {}
-  void shadowRegionArgs(Region &, ValueRange) override {}
 
   /// The printer flags to use when determining potential aliases.
   const OpPrintingFlags &printerFlags;
 
   /// The initializer to use when identifying aliases.
   AliasInitializer &initializer;
 
   /// A dummy output stream.
public:
 
   /// Return the result indices for each of the result groups registered by this
   /// operation, or empty if none exist.
   ArrayRef<int> getOpResultGroups(Operation *op);
 
   /// Get the info for the given block.
   BlockInfo getBlockInfo(Block *block);
 
-  /// Renumber the arguments for the specified region to the same names as the
-  /// SSA values in namesToUse. See OperationPrinter::shadowRegionArgs for
-  /// details.
-  void shadowRegionArgs(Region &region, ValueRange namesToUse);
-
 private:
   /// Number the SSA values within the given IR unit.
   void numberValuesInRegion(Region &region);
   void numberValuesInBlock(Block &block);
   void numberValuesInOp(Operation &op);
 
   /// Given a result of an operation 'result', find the result group head
   /// 'lookupValue' and the result of 'result' within that group in

 }
 
 BlockInfo SSANameState::getBlockInfo(Block *block) {
   auto it = blockNames.find(block);
   BlockInfo invalidBlock{-1, "INVALIDBLOCK"};
   return it != blockNames.end() ? it->second : invalidBlock;
 }
 
-void SSANameState::shadowRegionArgs(Region &region, ValueRange namesToUse) {
-  assert(!region.empty() && "cannot shadow arguments of an empty region");
-  assert(region.getNumArguments() == namesToUse.size() &&
-         "incorrect number of names passed in");
+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);
+  };
+
+  auto setRegionArgShadowingFn = [&](BlockArgument arg, Value valToShadow) {
+    assert(!valueIDs.count(arg) && "arg numbered multiple times");
+    assert(arg.getOwner() == &region.front() &&
+           "arg is not an argument of the current region");
     assert(region.getParentOp()->hasTrait<OpTrait::IsIsolatedFromAbove>() &&
-         "only KnownIsolatedFromAbove ops can shadow names");
+           "only known IsolatedFromAbove ops can shadow names");
+    assert(valToShadow.getParentRegion()->isProperAncestor(&region) &&
+           "shadowed value must be defined outside the region");
 
     SmallVector<char, 16> nameStr;
-  for (unsigned i = 0, e = namesToUse.size(); i != e; ++i) {
-    auto nameToUse = namesToUse[i];
-    if (nameToUse == nullptr)
-      continue;
-    auto nameToReplace = region.getArgument(i);
-
-    nameStr.clear();
     llvm::raw_svector_ostream nameStream(nameStr);
-    printValueID(nameToUse, /*printResultNo=*/true, nameStream);
-
-    // Entry block arguments should already have a pretty "arg" name.
-    assert(valueIDs[nameToReplace] == NameSentinel);
+    printValueID(valToShadow, /*printResultNo=*/true, nameStream);
 
     // Use the name without the leading %.
     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);
+    // NOTE: We bypass `setValueName` here because the name is likely not
+    //       unique, and may contain invalid characters (e.g. if `valToShadow`
+    //       is part of a result group, we may get a name like 'result#1').
+    valueIDs[arg] = NameSentinel;
+    valueNames[arg] = name.copy(usedNameAllocator);
   };
 
   if (!printerFlags.shouldPrintGenericOpForm()) {
     if (Operation *op = region.getParentOp()) {
-      if (auto asmInterface = dyn_cast<OpAsmOpInterface>(op))
+      if (auto asmInterface = dyn_cast<OpAsmOpInterface>(op)) {
         asmInterface.getAsmBlockArgumentNames(region, setBlockArgNameFn);
+        asmInterface.getAsmRegionArgumentShadowing(region,
+                                                   setRegionArgShadowingFn);
+      }
     }
   }
 
   // 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.
public:
   /// arguments.
   void printSuccessorAndUseList(Block *successor,
                                 ValueRange succOperands) override;
 
   /// Print the given region.
   void printRegion(Region &region, bool printEntryBlockArgs,
                    bool printBlockTerminators, bool printEmptyBlock) override;
 
-  /// Renumber the arguments for the specified region to the same names as the
-  /// SSA values in namesToUse. This may only be used for IsolatedFromAbove
-  /// operations. If any entry in namesToUse is null, the corresponding
-  /// argument name is left alone.
-  void shadowRegionArgs(Region &region, ValueRange namesToUse) override {
-    state.getSSANameState().shadowRegionArgs(region, namesToUse);
-  }
-
   /// Print the given affine map with the symbol and dimension operands printed
   /// inline with the map.
   void printAffineMapOfSSAIds(AffineMapAttr mapAttr,
                               ValueRange operands) override;
 
   /// Print the given affine expression with the symbol and dimension operands
   /// printed inline with the expression.
   void printAffineExprOfSSAIds(AffineExpr expr, ValueRange dimOperands,

 }
 
 void Value::dump() {
   print(llvm::errs());
   llvm::errs() << "\n";
 }
 
 void Value::printAsOperand(raw_ostream &os, AsmState &state) {
-  // TODO: This doesn't necessarily capture all potential cases.
-  // Currently, region arguments can be shadowed when printing the main
-  // operation. If the IR hasn't been printed, this will produce the old SSA
-  // name and not the shadowed name.
   state.getImpl().getSSANameState().printValueID(*this, /*printResultNo=*/true,
                                                  os);
 }
 
 void Operation::print(raw_ostream &os, const OpPrintingFlags &printerFlags) {
   // Find the operation to number from based upon the provided flags.
   Operation *op = this;
   bool shouldUseLocalScope = printerFlags.shouldUseLocalScope();

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

ParseResult IsolatedRegionOp::parse(OpAsmParser &parser,
   // Parse the body region, and reuse the operand info as the argument info.
   Region *body = result.addRegion();
   return parser.parseRegion(*body, argInfo, /*enableNameShadowing=*/true);
 }
 
 void IsolatedRegionOp::print(OpAsmPrinter &p) {
   p << ' ';
   p.printOperand(getOperand());
-  p.shadowRegionArgs(getRegion(), getOperand());
   p << ' ';
   p.printRegion(getRegion(), /*printEntryBlockArgs=*/false);
 }
 
+void IsolatedRegionOp::getAsmRegionArgumentShadowing(
+    Region &region, OpAsmSetRegionArgumentShadowingFn setShadowingFn) {
+  setShadowingFn(region.getArgument(0), getOperand());
+}
+
 //===----------------------------------------------------------------------===//
 // Test SSACFGRegionOp
 //===----------------------------------------------------------------------===//
 
 RegionKind SSACFGRegionOp::getRegionKind(unsigned index) {
   return RegionKind::SSACFG;
 }
 

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

 def ParseB64BytesOp : TEST_Op<"parse_b64"> {
   let arguments = (ins StrAttr:$b64);
   let hasCustomAssemblyFormat = 1;
 }
 
 //===----------------------------------------------------------------------===//
 // Test region argument list parsing.
 
-def IsolatedRegionOp : TEST_Op<"isolated_region", [IsolatedFromAbove]> {
+def IsolatedRegionOp : TEST_Op<"isolated_region", [
+    IsolatedFromAbove,
+    DeclareOpInterfaceMethods<OpAsmOpInterface, ["getAsmRegionArgumentShadowing"]>]> {
   let summary =  "isolated region operation";
   let description = [{
     Test op with an isolated region, to test passthrough region arguments. Each
     argument is of index type.
   }];
 
   let arguments = (ins Index);
   let regions = (region SizedRegion<1>:$region);