Add a new interface method `getAsmBlockName()` on OpAsmOpInterface to control block names

[mehdi_amini]

Summary

This allows operations to control the block ids used by the printer in nested regions.

[rriddle]

mlir/include/mlir/IR/OpAsmInterface.td

55–56

The phrasing of this is not the same as others, can you restructure this message to have a bit more detail? and e.g. show an example?

mlir/lib/IR/AsmPrinter.cpp

1046–1049

What's the runtime cost of using std::string for all blocks now?

1052

Please use proper sentences for comments.

[mehdi_amini]

Improve description

[mehdi_amini]

mlir/lib/IR/AsmPrinter.cpp

1046–1049

I'd say "not measurable compared to the rest of the infra involved in printing stuff"?

The real cost isn't in the block you highlighted actually, it is above: the loop with the while(1) will be invoked when printing the blocks in a function for example. That's an O(N^2) loop in terms of number of blocks...

1052

Except the leading upper case, what's missing?
(this is left from the LHS of the diff FYI)

[mehdi_amini]

Add the O(N2) behavior

[rriddle]

mlir/lib/IR/AsmPrinter.cpp

1046–1049

Have you measured it? Can we avoid all of this complexity in the common case (no override)?

1052

(Just meant the upper case)

[mehdi_amini]

mlir/lib/IR/AsmPrinter.cpp

1046–1049

(fixed the O(N2) right now)

Have you measured it?

No, but .... really?

Can we avoid all of this complexity in the common case (no override)?

Sure, I'll change the interface to return a callback...

[mehdi_amini]

mlir/lib/IR/AsmPrinter.cpp

1046–1049

(Actually that won't change the fact that we'll be using strings instead of unsigned in the map)

[mehdi_amini]

Use a factory to detect cases where the interface isn't overloaded

[Mogball]

mlir/include/mlir/IR/OpAsmInterface.td

81

Why is using std::string necessary?

mlir/lib/IR/AsmPrinter.cpp

851

Values have a separate map for numeric IDs and string names. Why not do the same for blocks? Then it might be possible to avoid some of the penalty for ops that don't implement the interface method.

1039–1040

This looks like an inlined copy of SSANameState::uniqueValueName

1041

[mehdi_amini]

mlir/include/mlir/IR/OpAsmInterface.td

81

It is an ownership problem: if you build a string dynamically on the other side of the interface, who owns it?

mlir/lib/IR/AsmPrinter.cpp

851

It just seems overly complex for something that seems extremely unlikely to me to every be a bottleneck anywhere.

Note also the tradeoff: here we precompute the string for the ID, while before you recompute a string for the ID on the fly. It's not clear to me that recomputing on the fly the integer serialization is necessarily better. But in any case these are not the bulk of the processing where we're talking about streaming a lot of string to an abstract stream implementation (regardless of the stream destination: string, file, console, socket, ...).

1039–1040

Yes: uniqueValueName has the "SSA value scope" I believe while this is for block names (which don't conflict with SSA value name I believe?).
Are you thinking that this logic should be refactored and shared somehow?

[mehdi_amini]

mlir/include/mlir/IR/OpAsmInterface.td

81

(the alternative would be a "StrAttr" instead I guess?)

[Mogball]

mlir/include/mlir/IR/OpAsmInterface.td

81

Why not use the usedNameAllocator?

mlir/lib/IR/AsmPrinter.cpp

1039–1040

Are you thinking that this logic should be refactored and shared somehow?

If possible, yeah.

[mehdi_amini]

mlir/include/mlir/IR/OpAsmInterface.td

81

You mean pass the allocator as an argument? And then return a StringRef?

Sure we could do this as well, that seems just overly convolute to me for "block names".

[Mogball]

mlir/include/mlir/IR/OpAsmInterface.td

81

I'm okay with std::string, but given that value names have been done a particular way, I find it strange that block names are done in a different way.

[mehdi_amini]

mlir/include/mlir/IR/OpAsmInterface.td

81

We can add more indirection and pass a callback to the callback so that the code would call setNameFn like the getAsmResultNames above?

81

(some non-negligible difference I make here is that the number block IDs in a typical IR is a small fraction of the number of SSA values)

[mehdi_amini]

Rework the flow entirely, align the callback with getAsmResultNames()

[Mogball]

Thanks! Otherwise LGTM

mlir/lib/IR/AsmPrinter.cpp

2700

Should this also sort by block name as well as ordering? Named blocks with ordering := -1 might be printed in an unstable order.

[mehdi_amini]

mlir/lib/IR/AsmPrinter.cpp

1031

Here we always define an ordering Jeff: even for block that have a name.

[Mogball]

mlir/lib/IR/AsmPrinter.cpp

1031

Ah, I missed this. Thanks.

[mehdi_amini]

Rebase.

I don't need this right now, but it is ready to land if we think it is a desirable addition!

[rriddle]

I think I would be okay landing this given that it is not super intrusive, and opens up interesting possibilities for users. That being said, I don't have a concrete use in mind right now.

mlir/lib/IR/AsmPrinter.cpp

984

Do the other "INVALID" things have an underscore in the name? I can't remember.

1025–1029

We could just say if the explicit name is empty, we use the ordering for a ^bbN name instead.

mlir/include/mlir/IR/OpAsmInterface.td

InterfaceMethod<[{
         ```mlir
           %first_result, %middle_results:2, %0 = "my.op" ...
         ```
       }],
       "void", "getAsmResultNames",
       (ins "::mlir::OpAsmSetValueNameFn":$setNameFn),
       "", ";"
     >,
+    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:
+              ...
+            ^bb1:
+              ...
+          })
+        ```
+
+        the method will be invoked on each of the blocks allowing the op to
+        print:
+
+        ```mlir
+          some.op() ({
+            ^custom_foo_name:
+              ...
+            ^custom_bar_name:
+              ...
+          })
+        ```
+      }],
+      "void", "getAsmBlockNames",
+      (ins "::mlir::OpAsmSetBlockNameFn":$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

 //===--------------------------------------------------------------------===//
 // 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.
 using OpAsmSetValueNameFn = function_ref<void(Value, StringRef)>;
 
+/// 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:
   /// Holds the result of `getAlias` hook call.
   enum class AliasResult {
     /// The object (type or attribute) is not supported by the hook
     /// and an alias was not provided.
     NoAlias,

mlir/lib/IR/AsmPrinter.cpp

void AliasState::printAliases(raw_ostream &os, NewLineCounter &newLine,
   }
 }
 
 //===----------------------------------------------------------------------===//
 // SSANameState
 //===----------------------------------------------------------------------===//
 
 namespace {
+/// Info about block printing: a number which is its position in the visitation
+/// order, and a name that is used to print reference to it, e.g. ^bb42.
+struct BlockInfo {
+  int ordering;
+  StringRef name;
+};
+
 /// This class manages the state of SSA value names.
 class SSANameState {
 public:
   /// A sentinel value used for values with names set.
   enum : unsigned { NameSentinel = ~0U };
 
   SSANameState(Operation *op, const OpPrintingFlags &printerFlags,
                DialectInterfaceCollection<OpAsmDialectInterface> &interfaces);
 
   /// Print the SSA identifier for the given value to 'stream'. If
   /// 'printResultNo' is true, it also presents the result number ('#' number)
   /// of this value.
   void printValueID(Value value, bool printResultNo, raw_ostream &stream) const;
 
   /// 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 ID for the given block.
-  unsigned getBlockID(Block *block);
+  /// 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.
private:
   DenseMap<Value, unsigned> valueIDs;
   DenseMap<Value, StringRef> valueNames;
 
   /// This is a map of operations that contain multiple named result groups,
   /// i.e. there may be multiple names for the results of the operation. The
   /// value of this map are the result numbers that start a result group.
   DenseMap<Operation *, SmallVector<int, 1>> opResultGroups;
 
-  /// This is the block ID for each block in the current.
-  DenseMap<Block *, unsigned> blockIDs;
+  /// This maps blocks to there visitation number in the current region as well
+  /// as the string representing their name.
+  DenseMap<Block *, BlockInfo> blockNames;
 
   /// This keeps track of all of the non-numeric names that are in flight,
   /// allowing us to check for duplicates.
   /// Note: the value of the map is unused.
   llvm::ScopedHashTable<StringRef, char> usedNames;
   llvm::BumpPtrAllocator usedNameAllocator;
 
   /// This is the next value ID to assign in numbering.
if (resultNo.hasValue() && printResultNo)
     stream << '#' << resultNo;
 }
 
 ArrayRef<int> SSANameState::getOpResultGroups(Operation *op) {
   auto it = opResultGroups.find(op);
   return it == opResultGroups.end() ? ArrayRef<int>() : it->second;
 }
 
-unsigned SSANameState::getBlockID(Block *block) {
-  auto it = blockIDs.find(block);
-  return it != blockIDs.end() ? it->second : NameSentinel;
+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");
   assert(region.getParentOp()->hasTrait<OpTrait::IsIsolatedFromAbove>() &&
          "only KnownIsolatedFromAbove ops can shadow names");

 }
 
 void SSANameState::numberValuesInRegion(Region &region) {
   // 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++;
+    auto blockInfoIt = blockNames.insert({&block, {-1, ""}});
+    if (blockInfoIt.second) {
+      // This block hasn't been named through `getAsmBlockArgumentNames`, use
+      // default `^bbNNN` format.
+      std::string name;
+      llvm::raw_string_ostream(name) << "^bb" << nextBlockID;
+      blockInfoIt.first->second.name = StringRef(name).copy(usedNameAllocator);
+    }
+    blockInfoIt.first->second.ordering = 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'.
   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++;
     }
     setValueName(arg, specialName.str());
   }
 
   // Number the operations in this block.
   for (auto &op : block)
     numberValuesInOp(op);
 }
 
 void SSANameState::numberValuesInOp(Operation &op) {
-  unsigned numResults = op.getNumResults();
-  if (numResults == 0)
-    return;
-  Value resultBegin = op.getResult(0);
-
   // Function used to set the special result names for the operation.
   SmallVector<int, 2> resultGroups(/*Size=*/1, /*Value=*/0);
   auto setResultNameFn = [&](Value result, StringRef name) {
     assert(!valueIDs.count(result) && "result numbered multiple times");
     assert(result.getDefiningOp() == &op && "result not defined by 'op'");
     setValueName(result, name);
 
     // Record the result number for groups not anchored at 0.
     if (int resultNo = result.cast<OpResult>().getResultNumber())
       resultGroups.push_back(resultNo);
   };
+  // Operations can customize the printing of block names in OpAsmOpInterface.
+  auto setBlockNameFn = [&](Block *block, StringRef name) {
+    assert(block->getParentOp() == &op &&
+           "getAsmBlockArgumentNames callback invoked on a block not directly "
+           "nested under the current operation");
+    assert(!blockNames.count(block) && "block numbered multiple times");
+    SmallString<16> tmpBuffer{"^"};
+    name = sanitizeIdentifier(name, tmpBuffer);
+    if (name.data() != tmpBuffer.data()) {
+      tmpBuffer.append(name);
+      name = tmpBuffer.str();
+    }
+    name = name.copy(usedNameAllocator);
+    blockNames[block] = {-1, name};
+  };
+
   if (!printerFlags.shouldPrintGenericOpForm()) {
-    if (OpAsmOpInterface asmInterface = dyn_cast<OpAsmOpInterface>(&op))
+    if (OpAsmOpInterface asmInterface = dyn_cast<OpAsmOpInterface>(&op)) {
+      asmInterface.getAsmBlockNames(setBlockNameFn);
       asmInterface.getAsmResultNames(setResultNameFn);
-    else if (auto *asmInterface = interfaces.getInterfaceFor(op.getDialect()))
+    } else if (auto *asmInterface =
+                   interfaces.getInterfaceFor(op.getDialect())) {
       asmInterface->getAsmResultNames(&op, setResultNameFn);
     }
+  }
+
+  unsigned numResults = op.getNumResults();
+  if (numResults == 0)
+    return;
+  Value resultBegin = op.getResult(0);
 
   // If the first result wasn't numbered, give it a default number.
   if (valueIDs.try_emplace(resultBegin, nextValueID).second)
     ++nextValueID;
 
   // If this operation has multiple result groups, mark it.
   if (resultGroups.size() != 1) {
     llvm::array_pod_sort(resultGroups.begin(), resultGroups.end());
void OperationPrinter::printGenericOp(Operation *op, bool printOpName) {
   printOptionalAttrDict(attrs);
 
   // Print the type signature of the operation.
   os << " : ";
   printFunctionalType(op);
 }
 
 void OperationPrinter::printBlockName(Block *block) {
-  auto id = state->getSSANameState().getBlockID(block);
-  if (id != SSANameState::NameSentinel)
-    os << "^bb" << id;
-  else
-    os << "^INVALIDBLOCK";
+  os << state->getSSANameState().getBlockInfo(block).name;
 }
 
 void OperationPrinter::print(Block *block, bool printBlockArgs,
                              bool printBlockTerminator) {
   // Print the block label and argument list if requested.
   if (printBlockArgs) {
     os.indent(currentIndent);
     printBlockName(block);
if (printBlockArgs) {
     if (!block->getParent()) {
       os << "  // block is not in a region!";
     } else if (block->hasNoPredecessors()) {
       os << "  // no predecessors";
     } else if (auto *pred = block->getSinglePredecessor()) {
       os << "  // pred: ";
       printBlockName(pred);
     } else {
-      // We want to print the predecessors in increasing numeric order, not in
+      // We want to print the predecessors in a stable order, not in
       // whatever order the use-list is in, so gather and sort them.
-      SmallVector<std::pair<unsigned, Block *>, 4> predIDs;
+      SmallVector<BlockInfo, 4> predIDs;
       for (auto *pred : block->getPredecessors())
-        predIDs.push_back({state->getSSANameState().getBlockID(pred), pred});
-      llvm::array_pod_sort(predIDs.begin(), predIDs.end());
+        predIDs.push_back(state->getSSANameState().getBlockInfo(pred));
+      llvm::sort(predIDs, [](BlockInfo lhs, BlockInfo rhs) {
+        return lhs.ordering < rhs.ordering;
+      });
 
       os << "  // " << predIDs.size() << " preds: ";
 
-      interleaveComma(predIDs, [&](std::pair<unsigned, Block *> pred) {
-        printBlockName(pred.second);
-      });
+      interleaveComma(predIDs, [&](BlockInfo pred) { os << pred.name; });
     }
     os << newLine;
   }
 
   currentIndent += indentWidth;
   bool hasTerminator =
       !block->empty() && block->back().hasTrait<OpTrait::IsTerminator>();
   auto range = llvm::make_range(

mlir/test/IR/pretty_printed_region_op.mlir

// CHECK:          : (f32, f32) -> f32
   %0 = test.pretty_printed_region %arg1, %arg0 ( {
     ^bb0(%arg2: f32, %arg3: f32):
       %1 = "non.special.op"(%arg2, %arg3) : (f32, f32) -> f32
       "test.return"(%1) : (f32) -> ()
     }) : (f32, f32) -> f32
     return %0 : f32
 }
 
+// -----
+
+// This tests the behavior of custom block names:
+// operations like `test.block_names` can define custom names for blocks in
+// nested regions.
+// CHECK-CUSTOM-LABEL: func @block_names
+func @block_names(%bool : i1) {
+  // CHECK: test.block_names
+  test.block_names {
+    // CHECK-CUSTOM: br ^foo1
+    // CHECK-GENERIC: std.br{{.*}}^bb1
+    br ^foo1
+  // CHECK-CUSTOM: ^foo1:
+  // CHECK-GENERIC: ^bb1:
+  ^foo1:
+    // CHECK-CUSTOM: br ^foo2
+    // CHECK-GENERIC: std.br{{.*}}^bb2
+    br ^foo2
+  // CHECK-CUSTOM: ^foo2:
+  // CHECK-GENERIC: ^bb2:
+  ^foo2:
+     "test.return"() : () -> ()
+  }
+  return
+}

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

let extraClassDeclaration = [{
     static ::llvm::StringRef getDefaultDialect() {
       return "test";
     }
     void getAsmResultNames(::llvm::function_ref<void(::mlir::Value, ::llvm::StringRef)> setNameFn) {}
   }];
   let assemblyFormat = "regions attr-dict-with-keyword";
 }
 
+// This is used to test the OpAsmOpInterface::getAsmBlockName() feature:
+// blocks nested in a region under this op will have a name defined by the
+// interface.
+def AsmBlockNameOp : TEST_Op<"block_names", [OpAsmOpInterface]> {
+ let regions = (region AnyRegion:$body);
+  let extraClassDeclaration = [{
+    void getAsmBlockNames(mlir::OpAsmSetBlockNameFn setNameFn) {
+      std::string name;
+      int count = 0;
+      for (::mlir::Block &block : getRegion().getBlocks()) {
+        name = "foo" + std::to_string(count++);
+        setNameFn(&block, name);
+      }
+    }
+  }];
+  let assemblyFormat = "regions attr-dict-with-keyword";
+}
+
 // This operation requires its return type to have the trait 'TestTypeTrait'.
 def ResultTypeWithTraitOp : TEST_Op<"result_type_with_trait", []> {
   let results = (outs AnyType);
 
   let verifier = [{
     if((*this)->getResultTypes()[0].hasTrait<TypeTrait::TestTypeTrait>())
       return success();
     return this->emitError("result type should have trait 'TestTypeTrait'");