[mlir][AsmPrinter] Change value numbering for local scope to be the next isolated operation.

[rriddle]

Summary

This revision updates the value numbering when printing to number from the next parent operation that is isolated from above. This is the highest level to number from that still ensures thread-safety. This revision also changes the behavior of Operator::operator<< to use local scope to avoid thread races when numbering operations.

[bondhugula]

mlir/include/mlir/IR/Operation.h

608

Did this race condition exist all along (since pass threading) or was it introduced lately? It looks like when pass threading is on, the printing has to always be "isolated from above" op local?

mlir/lib/IR/AsmPrinter.cpp

2347

Nit: The flags.shouldUseLocalScope() check can be hoisted out.

[rriddle]

Resolve comments

[rriddle]

mlir/include/mlir/IR/Operation.h

608

This has always existed, it is only recently that something used it in the common code path.

mlir/include/mlir/IR/Operation.h

 //===- Operation.h - MLIR Operation Class -----------------------*- C++ -*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
private:
   }
   size_t numTrailingObjects(OverloadToken<BlockOperand>) const {
     return numSuccs;
   }
   size_t numTrailingObjects(OverloadToken<Region>) const { return numRegions; }
 };
 
 inline raw_ostream &operator<<(raw_ostream &os, Operation &op) {
-  op.print(os);
+  op.print(os, OpPrintingFlags().useLocalScope());
   return os;
 }
 
 } // end namespace mlir
 
 namespace llvm {
 /// Provide isa functionality for operation casts.
 template <typename T> struct isa_impl<T, ::mlir::Operation> {

mlir/lib/IR/AsmPrinter.cpp

 //===- AsmPrinter.cpp - MLIR Assembly Printer Implementation --------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
void Value::printAsOperand(raw_ostream &os, AsmState &state) {
   // 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, OpPrintingFlags flags) {
-  // Handle top-level operations or local printing.
-  if (!getParent() || flags.shouldUseLocalScope()) {
-    AsmState state(this);
-    OperationPrinter(os, flags, state.getImpl()).print(this);
-    return;
-  }
+  // Find the operation to number from based upon the provided flags.
+  Operation *printedOp = this;
+  bool shouldUseLocalScope = flags.shouldUseLocalScope();
+  do {
+    // If we are printing local scope, stop at the first operation that is
+    // isolated from above.
+    if (shouldUseLocalScope && printedOp->isKnownIsolatedFromAbove())
+      break;
 
-  Operation *parentOp = getParentOp();
-  if (!parentOp) {
-    os << "<<UNLINKED OPERATION>>\n";
-    return;
-  }
-  // Get the top-level op.
-  while (auto *nextOp = parentOp->getParentOp())
-    parentOp = nextOp;
+    // Otherwise, traverse up to the next parent.
+    Operation *parentOp = printedOp->getParentOp();
+    if (!parentOp)
+      break;
+    printedOp = parentOp;
+  } while (true);
 
-  AsmState state(parentOp);
+  AsmState state(printedOp);
   print(os, state, flags);
 }
 void Operation::print(raw_ostream &os, AsmState &state, OpPrintingFlags flags) {
   OperationPrinter(os, flags, state.getImpl()).print(this);
 }
 
 void Operation::dump() {
   print(llvm::errs(), OpPrintingFlags().useLocalScope());