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

1058–1061

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

1064

Please use proper sentences for comments.

[mehdi_amini]

Improve description

[mehdi_amini]

mlir/lib/IR/AsmPrinter.cpp

1058–1061

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...

1064

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

1058–1061

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

1064

(Just meant the upper case)

[mehdi_amini]

mlir/lib/IR/AsmPrinter.cpp

1058–1061

(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

1058–1061

(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.

1051–1052

This looks like an inlined copy of SSANameState::uniqueValueName

1053

[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, ...).

1051–1052

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

1051–1052

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

2704

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

1043

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

[Mogball]

mlir/lib/IR/AsmPrinter.cpp

1043

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

978

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

1037–1041

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:
+              ...
+          })
+        ```
+      }],
+      "std::function<std::string(::mlir::Block*)>", "getAsmBlockNameFactory",
+      (ins), "", "return nullptr;"
+    >,
     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/lib/IR/AsmPrinter.cpp

public:
   /// 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);
+  StringRef getBlockID(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;
+  /// These are the block IDs for each block in the current region.
+  DenseMap<Block *, std::string> blockIDs;
 
   /// 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) {
+StringRef SSANameState::getBlockID(Block *block) {
   auto it = blockIDs.find(block);
-  return it != blockIDs.end() ? it->second : NameSentinel;
+  return it != blockIDs.end() ? it->second : StringRef("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");
 
   SmallVector<char, 16> nameStr;
for (unsigned i = 0, e = namesToUse.size(); i != e; ++i) {
 
     // Overwrite the name.
     valueNames[nameToReplace] = name.copy(usedNameAllocator);
   }
 }
 
 void SSANameState::numberValuesInRegion(Region &region) {
   // Number the values within this region in a breadth-first order.
+
+  // Operations can customize the printing of block names in OpAsmOpInterface.
+  std::function<std::string(Block *)> blockNameFactory;
+  if (!printerFlags.shouldPrintGenericOpForm()) {
+    auto parentOp = dyn_cast_or_null<OpAsmOpInterface>(region.getParentOp());
+    if (parentOp)
+      blockNameFactory = parentOp.getAsmBlockNameFactory();
+  }
+  if (!blockNameFactory) {
+    // This is the generic case, number the blocks in sequence: bb0, bb1, ...
+    unsigned nextBlockID = 0;
+    std::string blockName;
+    for (auto &block : region) {
+      std::string blockName = "bb" + std::to_string(nextBlockID++);
+      blockIDs.insert(std::make_pair(&block, std::move(blockName)));
+
+      // Number all of the operations within the block as well.
+      numberValuesInBlock(block);
+    }
+    return;
+  }
+  // Use the factory provided by the operation to customize block names.
+  StringSet<> blockNames;
   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++;
+    std::string baseBlockName = blockNameFactory(&block);
+    std::string blockName = baseBlockName;
+    if (baseBlockName.empty()) {
+      blockName = "bb" + std::to_string(nextBlockID++);
+      baseBlockName = "bb";
+    }
+    // Each block gets a unique ID, try to use the provided name but detects
+    // name collision and append a increasing counter.
+    int count = 0;
+    do {
+      auto insertIt = blockNames.insert(blockName);
+      if (insertIt.second) {
+        // No collision, the name is unique in this region!
+        blockIDs.insert(std::make_pair(&block, insertIt.first->first().str()));
+        break;
+      }
+      // We hit a collision, increment `count` and append to the name.
+      blockName = baseBlockName + std::to_string(++count);
+    } while (1);
+
+    // Number all of the operations within the block as well.
     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);
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().getBlockID(block);
 }
 
 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<StringRef, 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().getBlockID(pred));
+      llvm::sort(predIDs.begin(), predIDs.end());
 
       os << "  // " << predIDs.size() << " preds: ";
 
-      interleaveComma(predIDs, [&](std::pair<unsigned, Block *> pred) {
-        printBlockName(pred.second);
-      });
+      interleaveComma(predIDs, [&](StringRef pred) { os << "^" << pred; });
     }
     os << newLine;
   }
 
   currentIndent += indentWidth;
   bool hasTerminator =
       !block->empty() && block->back().hasTrait<OpTrait::IsTerminator>();
   auto range = llvm::make_range(

mlir/test/Conversion/StandardToLLVM/standard-to-llvm.mlir

 // CHECK-NEXT:  llvm.br ^bb5({{.*}} : i64)
 // CHECK-NEXT:^bb7: // pred: ^bb5
 // CHECK-NEXT:  llvm.call @mid({{.*}}) : (i64) -> ()
 // CHECK-NEXT:  llvm.br ^bb8
 // CHECK-NEXT:^bb8: // pred: ^bb7
 // CHECK-NEXT:  {{.*}} = llvm.mlir.constant(18 : index) : i64
 // CHECK-NEXT:  {{.*}} = llvm.mlir.constant(37 : index) : i64
 // CHECK-NEXT:  llvm.br ^bb9({{.*}} : i64)
-// CHECK-NEXT:^bb9({{.*}}: i64):  // 2 preds: ^bb8, ^bb10
+// CHECK-NEXT:^bb9({{.*}}: i64):  // 2 preds: ^bb10, ^bb8
 // CHECK-NEXT:  {{.*}} = llvm.icmp "slt" {{.*}}, {{.*}} : i64
 // CHECK-NEXT:  llvm.cond_br {{.*}}, ^bb10, ^bb11
 // CHECK-NEXT:^bb10:  // pred: ^bb9
 // CHECK-NEXT:  llvm.call @body3({{.*}}, {{.*}}) : (i64, i64) -> ()
 // CHECK-NEXT:  {{.*}} = llvm.mlir.constant(3 : index) : i64
 // CHECK-NEXT:  {{.*}} = llvm.add {{.*}}, {{.*}} : i64
 // CHECK-NEXT:  llvm.br ^bb9({{.*}} : i64)
 // CHECK-NEXT:^bb11:  // pred: ^bb9

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 = [{
+    std::function<std::string(::mlir::Block*)> getAsmBlockNameFactory() {
+      return [] (::mlir::Block *block) { return "foo"; };
+    }
+  }];
+  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'");