[mlir] add utilites for DiagnosedSilenceableFailure

[ftynse]

Summary

This class adds helper functions similar to emitError for the
DiagnosedSilenceableFailure class in both the silenceable and definite
failure cases. These helpers simplify the use of said class and make
tranfsorm op application code idiomatic.

[springerm]

Nice, this will make the error handling much easier to use!

mlir/include/mlir/Dialect/Transform/IR/TransformInterfaces.h

186

Optional: To be consistent with op->emitError, this could take an Operation * or Value instead of Location.

266

nit: Without looking at the impl of DiagnosedSilenceableFailure: I expected something like DiagnosedSilenceableFailure(loc, message);, similar to emitDefiniteFailure.

[ftynse]

Address review.

mlir/include/mlir/Dialect/Transform/IR/TransformInterfaces.h

266

There's no such constructor, it would be unclear in which state the silenceable failure is created.

mlir/include/mlir/Dialect/Transform/IR/TransformInterfaces.h

#if LLVM_ENABLE_ABI_BREAKING_CHECKS
   /// Whether the associated diagnostics have been reported.
   /// Diagnostics reporting consumes the diagnostics, so we need a mechanism to
   /// differentiate reported diagnostics from a state where it was never
   /// created.
   bool reported = false;
 #endif // LLVM_ENABLE_ABI_BREAKING_CHECKS
 };
 
+class DiagnosedDefiniteFailure;
+
+DiagnosedDefiniteFailure emitDefiniteFailure(Location loc,
+                                             const Twine &message = {});
+
+/// A compatibility class connecting `InFlightDiagnostic` to
+/// `DiagnosedSilenceableFailure` while providing an interface similar to the
+/// former. Implicitly convertible to `DiagnosticSilenceableFailure` in definite
+/// failure state and to `LogicalResult` failure. Reports the error on
+/// conversion or on destruction. Instances of this class can be created by
+/// `emitDefiniteFailure()`.
+class DiagnosedDefiniteFailure {
+  friend DiagnosedDefiniteFailure emitDefiniteFailure(Location loc,
+                                                      const Twine &message);
+
+public:
+  /// Only move-constructible because it carries an in-flight diagnostic.
+  DiagnosedDefiniteFailure(DiagnosedDefiniteFailure &&) = default;
+
+  /// Forward the message to the diagnostic.
+  template <typename T>
+  DiagnosedDefiniteFailure &operator<<(T &&value) & {
+    diag << std::forward<T>(value);
+    return *this;
+  }
+  template <typename T>
+  DiagnosedDefiniteFailure &&operator<<(T &&value) && {
+    return std::move(this->operator<<(std::forward<T>(value)));
+  }
+
+  /// Attaches a note to the error.
+  Diagnostic &attachNote(Optional<Location> loc = llvm::None) {
+    return diag.attachNote(loc);
+  }
+
+  /// Implicit conversion to DiagnosedSilenceableFailure in the definite failure
+  /// state. Reports the error.
+  operator DiagnosedSilenceableFailure() {
+    diag.report();
+    return DiagnosedSilenceableFailure::definiteFailure();
+  }
+
+  /// Implicit conversion to LogicalResult in the failure state. Reports the
+  /// error.
+  operator LogicalResult() {
+    diag.report();
+    return failure();
+  }
+
+private:
+  /// Constructs a definite failure at the given location with the given
+  /// message.
+  explicit DiagnosedDefiniteFailure(Location loc, const Twine &message)
+      : diag(emitError(loc, message)) {}
+
+  /// Copy-construction and any assignment is disallowed to prevent repeated
+  /// error reporting.
+  DiagnosedDefiniteFailure(const DiagnosedDefiniteFailure &) = delete;
+  DiagnosedDefiniteFailure &
+  operator=(const DiagnosedDefiniteFailure &) = delete;
+  DiagnosedDefiniteFailure &operator=(DiagnosedDefiniteFailure &&) = delete;
+
+  /// The error message.
+  InFlightDiagnostic diag;
+};
+
+/// Emits a definite failure with the given message. The returned object allows
+/// for last-minute modification to the error message, such as attaching notes
+/// and completing the message. It will be reported when the object is
+/// destructed or converted.
+inline DiagnosedDefiniteFailure emitDefiniteFailure(Location loc,
+                                                    const Twine &message) {
+  return DiagnosedDefiniteFailure(loc, message);
+}
+inline DiagnosedDefiniteFailure emitDefiniteFailure(Operation *op,
+                                                    const Twine &message = {}) {
+  return emitDefiniteFailure(op->getLoc(), message);
+}
+
+/// Emits a silenceable failure with the given message. A silenceable failure
+/// must be either suppressed or converted into a definite failure and reported
+/// to the user.
+inline DiagnosedSilenceableFailure
+emitSilenceableFailure(Location loc, const Twine &message = {}) {
+  Diagnostic diag(loc, DiagnosticSeverity::Error);
+  diag << message;
+  return DiagnosedSilenceableFailure::silenceableFailure(std::move(diag));
+}
+inline DiagnosedSilenceableFailure
+emitSilenceableFailure(Operation *op, const Twine &message = {}) {
+  return emitSilenceableFailure(op->getLoc(), message);
+}
+
 namespace transform {
 
 class TransformOpInterface;
 
 /// Options controlling the application of transform operations by the
 /// TransformState.
 class TransformOptions {
 public:

mlir/include/mlir/Dialect/Transform/IR/TransformInterfaces.td

let extraSharedClassDeclaration = [{
     ::mlir::LogicalResult reportUnknownTransformError(
         ::mlir::Operation *target) {
       ::mlir::InFlightDiagnostic diag = $_op->emitError() << "failed to apply";
       diag.attachNote(target->getLoc()) << "attempted to apply to this op";
       return diag;
     }
 
     /// Creates the silenceable failure object with a diagnostic located at the
-    /// current operation.
-    DiagnosedSilenceableFailure emitSilenceableError() {
-      Diagnostic diag($_op->getLoc(), DiagnosticSeverity::Error);
-      return DiagnosedSilenceableFailure::silenceableFailure(std::move(diag));
+    /// current operation. Silenceable failure must be suppressed or reported
+    /// explicitly at some later time.
+    DiagnosedSilenceableFailure
+    emitSilenceableError(const ::llvm::Twine &message = {}) {
+      return ::mlir::emitSilenceableFailure($_op);
+    }
+
+    /// Creates the definite failure object with a diagnostic located at the
+    /// current operation. Definite failure will be reported when the object
+    /// is destroyed or converted.
+    DiagnosedDefiniteFailure
+    emitDefiniteFailure(const ::llvm::Twine &message = {}) {
+      return ::mlir::emitDefiniteFailure($_op, message);
     }
 
     /// Creates the default silenceable failure for a transform op that failed
     /// to properly apply to a target.
     DiagnosedSilenceableFailure emitDefaultSilenceableFailure(
         Operation *target) {
-      Diagnostic diag($_op->getLoc(), DiagnosticSeverity::Error);
+      DiagnosedSilenceableFailure diag = emitSilenceableFailure($_op->getLoc());
       diag << $_op->getName() << " failed to apply";
       diag.attachNote(target->getLoc()) << "when applied to this op";
-      return DiagnosedSilenceableFailure::silenceableFailure(std::move(diag));
+      return diag;
     }
   }];
 }
 
 def TransformTypeInterface : TypeInterface<"TransformTypeInterface"> {
   let description = [{
     Types that can be used for Transform dialect handle values. Such types
     define the properties of Payload IR operations associated with the handle.

mlir/lib/Dialect/GPU/TransformOps/GPUTransformOps.cpp

static DiagnosedSilenceableFailure rewriteOneForeachThreadToGpuThreads(
     RewriterBase &rewriter, scf::ForeachThreadOp foreachThreadOp,
     const SmallVectorImpl<int64_t> &globalBlockDims, bool syncAfterDistribute,
     llvm::Optional<TransformOpInterface> transformOp) {
   auto failureHelper =
       [&](const Twine &message) -> DiagnosedSilenceableFailure {
     if (transformOp.has_value()) {
       return transformOp->emitSilenceableError() << message;
     }
-    foreachThreadOp->emitError() << message;
-    return DiagnosedSilenceableFailure::definiteFailure();
+    return emitDefiniteFailure(foreachThreadOp, message);
   };
 
   if (foreachThreadOp.getNumResults() > 0)
     return failureHelper(
         "only bufferized scf.foreach_thread lowers to gpu.thread_id");
 
   if (foreachThreadOp.getNumThreads().size() > 3)
     return failureHelper(

mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp

transform::FuseIntoContainingOp::apply(transform::TransformResults &results,
   // If nothing to fuse, propagate success.
   if (producerOps.empty()) {
     results.set(getFusedOp().cast<OpResult>(),
                 SmallVector<mlir::Operation *>{});
     return DiagnosedSilenceableFailure::success();
   }
   ArrayRef<Operation *> containingOps = state.getPayloadOps(getContainingOp());
   if (containingOps.size() != 1) {
-    // Definite failure.
-    return DiagnosedSilenceableFailure(
-        this->emitOpError("requires exactly one containing_op handle (got ")
-        << containingOps.size() << ")");
+    return emitDefiniteFailure()
+           << "requires exactly one containing_op handle (got "
+           << containingOps.size() << ")";
   }
   Operation *containingOp = containingOps.front();
 
   // Helper function to find the next producer that should be fused. Take any
   // producer that has a use inside the containing op.
   SmallVector<Operation *> remainingProducers(producerOps.begin(),
                                               producerOps.end());
   auto getNextProducer = [&]() -> FailureOr<Operation *> {
splitPoints = llvm::to_vector(llvm::map_range(
         }));
     if (diag.isSilenceableFailure()) {
       results.set(getFirst().cast<OpResult>(), {});
       results.set(getSecond().cast<OpResult>(), {});
       return diag;
     }
 
     if (splitPoints.size() != payload.size()) {
-      emitError() << "expected the dynamic split point handle to point to as "
-                     "many operations ("
-                  << splitPoints.size() << ") as the target handle ("
-                  << payload.size() << ")";
-      return DiagnosedSilenceableFailure::definiteFailure();
+      return emitDefiniteFailure()
+             << "expected the dynamic split point handle to point to as "
     }
   } else {
     splitPoints.resize(payload.size(),
                        rewriter.getIndexAttr(getStaticSplitPoint()));
   }
 
   // Split each target operation.
   SmallVector<Operation *> first, second;

mlir/lib/Dialect/Transform/IR/TransformOps.cpp

transform::AlternativesOp::apply(transform::TransformResults &results,
   SmallVector<Operation *> originals;
   if (Value scopeHandle = getScope())
     llvm::append_range(originals, state.getPayloadOps(scopeHandle));
   else
     originals.push_back(state.getTopLevel());
 
   for (Operation *original : originals) {
     if (original->isAncestor(getOperation())) {
-      InFlightDiagnostic diag =
-          emitError() << "scope must not contain the transforms being applied";
+      auto diag = emitDefiniteFailure()
+                  << "scope must not contain the transforms being applied";
       diag.attachNote(original->getLoc()) << "scope";
-      return DiagnosedSilenceableFailure::definiteFailure();
+      return diag;
     }
     if (!original->hasTrait<OpTrait::IsIsolatedFromAbove>()) {
-      InFlightDiagnostic diag =
-          emitError()
+      auto diag = emitDefiniteFailure()
                   << "only isolated-from-above ops can be alternative scopes";
       diag.attachNote(original->getLoc()) << "scope";
-      return DiagnosedSilenceableFailure(std::move(diag));
+      return diag;
     }
   }
 
   for (Region &reg : getAlternatives()) {
     // Clone the scope operations and make the transforms in this alternative
     // region apply to them by virtue of mapping the block argument (the only
     // visible handle) to the cloned scope operations. This effectively prevents
     // the transformation from accessing any IR outside the scope.
transform::PDLMatchOp::apply(transform::TransformResults &results,
                              transform::TransformState &state) {
   auto *extension = state.getExtension<PatternApplicatorExtension>();
   assert(extension &&
          "expected PatternApplicatorExtension to be attached by the parent op");
   SmallVector<Operation *> targets;
   for (Operation *root : state.getPayloadOps(getRoot())) {
     if (failed(extension->findAllMatches(
             getPatternName().getLeafReference().getValue(), root, targets))) {
-      emitOpError() << "could not find pattern '" << getPatternName() << "'";
-      return DiagnosedSilenceableFailure::definiteFailure();
+      emitDefiniteFailure()
+          << "could not find pattern '" << getPatternName() << "'";
     }
   }
   results.set(getResult().cast<OpResult>(), targets);
   return DiagnosedSilenceableFailure::success();
 }
 
 void transform::PDLMatchOp::getEffects(
     SmallVectorImpl<MemoryEffects::EffectInstance> &effects) {

mlir/test/Dialect/Transform/transform-state-extension.mlir

^bb0(%arg0: !pdl.operation):
     // expected-remark @below {{extension present, A}}
     test_check_if_test_extension_present %arg0
     // expected-note @below {{associated payload op}}
     test_remap_operand_to_self %arg0
     // expected-remark @below {{extension present, A}}
     test_check_if_test_extension_present %arg0
   }
 }
+
+// -----
+
+module {
+  transform.sequence failures(suppress) {
+  ^bb0(%arg0: !pdl.operation):
+    // expected-error @below {{TestTransformStateExtension missing}}
+    test_remap_operand_to_self %arg0
+  }
+}

mlir/test/lib/Dialect/Transform/TestTransformDialectExtension.cpp

#endif // NDEBUG
   }
 
   return DiagnosedSilenceableFailure::success();
 }
 
 DiagnosedSilenceableFailure mlir::test::TestRemapOperandPayloadToSelfOp::apply(
     transform::TransformResults &results, transform::TransformState &state) {
   auto *extension = state.getExtension<TestTransformStateExtension>();
-  if (!extension) {
-    emitError() << "TestTransformStateExtension missing";
-    return DiagnosedSilenceableFailure::definiteFailure();
-  }
+  if (!extension)
+    return emitDefiniteFailure("TestTransformStateExtension missing");
 
   if (failed(extension->updateMapping(state.getPayloadOps(getOperand()).front(),
                                       getOperation())))
     return DiagnosedSilenceableFailure::definiteFailure();
   return DiagnosedSilenceableFailure::success();
 }
 
 DiagnosedSilenceableFailure mlir::test::TestRemoveTestExtensionOp::apply(