[MLIR] Modify Partial op conversion mode to optionally track all non-legalizable operations.

[lucyrfox]

Summary

There are three op conversion modes: Partial, Full, and Analysis. This change modifies the Partial mode to optionally take a set of non-legalizable ops. If this parameter is specified, all ops that are not legalizable (i.e. would cause full conversion to fail) are tracked throughout the partial legalization.

[lucyrfox]

Took a first pass at this, but I'm open to different ways of approaching this change. At I've written it now, the basic idea of the Analysis conversion is "we're performing a partial conversion and emitting warnings along the way for any ops that would fail to legalize in a full conversion". Success is returned as long as the partial conversion succeeds. I did it this way because returning failure outright means that the partially converted IR doesn't get emitted, which is useful for debugging.

However, an implication of this is that if the legalization did fully succeed, we don't know that based on the Analysis conversion. As a result, I foresee a lot of users running an Analysis conversion immediately followed by a Full conversion, thus incurring an unnecessary pass over the IR. So ideally, this would return success/failure for the full conversion while still outputting the partially converted IR in the failure case.

What do you think?

[lucyrfox]

Minor formatting change.

[rriddle]

In D78788#2001142, @lucyrfox wrote:

Took a first pass at this, but I'm open to different ways of approaching this change. At I've written it now, the basic idea of the Analysis conversion is "we're performing a partial conversion and emitting warnings along the way for any ops that would fail to legalize in a full conversion". Success is returned as long as the partial conversion succeeds. I did it this way because returning failure outright means that the partially converted IR doesn't get emitted, which is useful for debugging.

Before answering, please note that we do not base these decisions based on debugging. What the result is depends on what the user intends to do with the result.

However, an implication of this is that if the legalization did fully succeed, we don't know that based on the Analysis conversion. As a result, I foresee a lot of users running an Analysis conversion immediately followed by a Full conversion, thus incurring an unnecessary pass over the IR. So ideally, this would return success/failure for the full conversion while still outputting the partially converted IR in the failure case.

Seems like the user can easily tell if the conversion was completely successful by checking if the set of unlegalized operations is empty?

What do you think?

We can change analysis conversion to whatever we want, but the main question is what are you doing with the information? What happens after the analysis conversion and what do you want to convey? We may end up with different types of analysis conversions based on the answer to that question. For example, the original envisioning for this was that a user would performs analysis conversions targeting different backends to detail which parts of the subgraph were really supported for a specific target. In that type of mode you don't want to do any conversions, because you want to splice the original graph based on the analysis. It seems that the use case you are targeting is different. For your use case, what is the pipeline that you are envisioning and how does an analysis conversion fit into it?

mlir/lib/Transforms/DialectConversion.cpp

1575

I'm not sure we always want warnings here, seems like more of a DEBUG log. If a user is relying on an analysis conversion to help slice code that is supported by different backends, they may not care that certain operations weren't converted.

[mehdi_amini]

mlir/lib/Transforms/DialectConversion.cpp

1568

Isn't the mode necessarily Partial/Analysis here?

1577

Isn't this something we'd want to track even in "Partial" mode?

It seems like the difference between Partial and Analysis overall is about whether the conversion is performed or not. Anything else should be uniform.
Partial vs Full is that the latter is aborting early and so does not need to accumulate the unlegalizableOps.

1579

Do we need to warn here?
Can we let signaling up to the client since we populate unlegalizableOps?

[lucyrfox]

[Note: I haven't made code changes in this iteration because I want to settle on an approach before refactoring this. For the same reason I'm leaving some comments unresolved for now.]

In D78788#2005904, @rriddle wrote:

In D78788#2001142, @lucyrfox wrote:

Took a first pass at this, but I'm open to different ways of approaching this change. At I've written it now, the basic idea of the Analysis conversion is "we're performing a partial conversion and emitting warnings along the way for any ops that would fail to legalize in a full conversion". Success is returned as long as the partial conversion succeeds. I did it this way because returning failure outright means that the partially converted IR doesn't get emitted, which is useful for debugging.

Before answering, please note that we do not base these decisions based on debugging. What the result is depends on what the user intends to do with the result.

However, an implication of this is that if the legalization did fully succeed, we don't know that based on the Analysis conversion. As a result, I foresee a lot of users running an Analysis conversion immediately followed by a Full conversion, thus incurring an unnecessary pass over the IR. So ideally, this would return success/failure for the full conversion while still outputting the partially converted IR in the failure case.

Seems like the user can easily tell if the conversion was completely successful by checking if the set of unlegalized operations is empty?

Good point, don't know why I didn't think of this :)

What do you think?

We can change analysis conversion to whatever we want, but the main question is what are you doing with the information? What happens after the analysis conversion and what do you want to convey? We may end up with different types of analysis conversions based on the answer to that question. For example, the original envisioning for this was that a user would performs analysis conversions targeting different backends to detail which parts of the subgraph were really supported for a specific target. In that type of mode you don't want to do any conversions, because you want to splice the original graph based on the analysis. It seems that the use case you are targeting is different. For your use case, what is the pipeline that you are envisioning and how does an analysis conversion fit into it?

Thanks for raising these points; this helps me better understand the current design of the analysis conversion mode.

My use case (for the TF XLA bridge) is the following: we want to perform a full conversion and for the result to be fully legalized (in that sense, the same as full legalization today). However, in the case that it fails, we want to (a) know all of the ops that failed to legalize and (b) have a partially legalized graph output. I suppose in that sense, this conversion mode I'm envisioning would act like full legalization in the success case but like partial legalization in the failure case.

After considering the comments here, I see a couple of possibilities (both of which would involve leaving the existing Analysis conversion mode as it is today):

(1) Modify the partial conversion mode to take an (optional?) unlegalizedOps set as input and populate it with all ops that fail to legalize.
-In our use case, this would mean always using the Partial conversion mode, and if what we really want is a Full conversion, checking whether that set of ops is empty. That's totally fine from the TF XLA perspective, but is that a level of indirection that ought to be in Core?
-Cons: API instability if the unlegalizedOps set is not optional.

(2) Add a new (2nd type of analysis) conversion mode that aligns with this use case.

I lean toward option (1), but I want to solicit some more thoughts before making the code change.

mlir/lib/Transforms/DialectConversion.cpp

1577

Isn't this something we'd want to track even in "Partial" mode?

I think this is a good idea and aligns with option (1) I describe in the larger comment on this reply.

It seems like the difference between Partial and Analysis overall is about whether the conversion is performed or not. Anything else should be uniform.

This is true with Analysis mode as it is today, but this patch (in its current state) makes it such that both Partial and Analysis mode perform the conversion. The key difference is whether or not the unlegalizable ops are collected in that set.

Partial vs Full is that the latter is aborting early and so does not need to accumulate the unlegalizableOps.

This is the difference between Analysis and Full. Presently, Partial doesn't accumulate the unlegalizableOps, either.

[rriddle]

I lean toward option (1), but I want to solicit some more thoughts before making the code change.

(1) seems like a fine option to add.

[lucyrfox]

Modify Partial conversion mode (instead of Analysis mode) to take an optional parameter to store non-legalizable ops.

[lucyrfox]

mlir/lib/Transforms/DialectConversion.cpp

1575

Removed these warnings.

1579

Agreed. Changed.

[mehdi_amini]

mlir/include/mlir/Transforms/DialectConversion.h

669 ↗

(On Diff #261123)

Just a nit, but I believe that if you use backticks ` to refer to arguments then doxygen link these.

[lucyrfox]

Small formatting fix.

mlir/lib/Transforms/DialectConversion.cpp

 //===- DialectConversion.cpp - MLIR dialect conversion generic pass -------===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Transforms/DialectConversion.h"

 
 // This class converts operations to a given conversion target via a set of
 // rewrite patterns. The conversion behaves differently depending on the
 // conversion mode.
 struct OperationConverter {
   explicit OperationConverter(ConversionTarget &target,
                               const OwningRewritePatternList &patterns,
                               OpConversionMode mode,
-                              DenseSet<Operation *> *legalizableOps = nullptr)
+                              DenseSet<Operation *> *unlegalizableOps = nullptr)
       : opLegalizer(target, patterns), mode(mode),
-        legalizableOps(legalizableOps) {}
+        unlegalizableOps(unlegalizableOps) {}
 
   /// Converts the given operations to the conversion target.
   LogicalResult convertOperations(ArrayRef<Operation *> ops,
                                   TypeConverter *typeConverter);
 
 private:
   /// Converts an operation with the given rewriter.
   LogicalResult convert(ConversionPatternRewriter &rewriter, Operation *op);
 
   /// Converts the type signatures of the blocks nested within 'op'.
   LogicalResult convertBlockSignatures(ConversionPatternRewriter &rewriter,
                                        Operation *op);
 
   /// The legalizer to use when converting operations.
   OperationLegalizer opLegalizer;
 
   /// The conversion mode to use when legalizing operations.
   OpConversionMode mode;
 
-  /// A set of pre-existing operations that were found to be legalizable to the
-  /// target. This field is only used when mode == OpConversionMode::Analysis.
-  DenseSet<Operation *> *legalizableOps;
+  /// A set of pre-existing operations that were found not to be legalizable to
+  /// the target. This field is only used when mode ==
+  /// OpConversionMode::Analysis.
+  DenseSet<Operation *> *unlegalizableOps;
 };
 } // end anonymous namespace
 
 LogicalResult
 OperationConverter::convertBlockSignatures(ConversionPatternRewriter &rewriter,
                                            Operation *op) {
   // Check to see if type signatures need to be converted.
   if (!rewriter.getImpl().argConverter.typeConverter)
LogicalResult OperationConverter::convert(ConversionPatternRewriter &rewriter,
   if (failed(opLegalizer.legalize(op, rewriter))) {
     // Handle the case of a failed conversion for each of the different modes.
     /// Full conversions expect all operations to be converted.
     if (mode == OpConversionMode::Full)
       return op->emitError()
              << "failed to legalize operation '" << op->getName() << "'";
     /// Partial conversions allow conversions to fail iff the operation was not
     /// explicitly marked as illegal.
-    if (mode == OpConversionMode::Partial && opLegalizer.isIllegal(op))
+    if ((mode == OpConversionMode::Partial ||
+         mode == OpConversionMode::Analysis) &&
+        opLegalizer.isIllegal(op))
       return op->emitError()
              << "failed to legalize operation '" << op->getName()
              << "' that was explicitly marked illegal";
+    if (mode == OpConversionMode::Analysis) {
+      /// An analysis conversion emits any errors that would occur if a full
+      /// conversion were attempted, but should continue on without failure
+      /// (i.e. performing a partial conversion).
+      unlegalizableOps->insert(op);
+      op->emitWarning() << "failed to legalize operation '" << op->getName()
+                      << "'";
+      return success();
+    }
   } else {
-    /// Analysis conversions don't fail if any operations fail to legalize,
-    /// they are only interested in the operations that were successfully
-    /// legalized.
-    if (mode == OpConversionMode::Analysis)
-      legalizableOps->insert(op);
-
     // If legalization succeeded, convert the types any of the blocks within
     // this operation.
     if (failed(convertBlockSignatures(rewriter, op)))
       return failure();
   }
   return success();
 }
 
OperationConverter::convertOperations(ArrayRef<Operation *> ops,
   }
 
   // Convert each operation and discard rewrites on failure.
   ConversionPatternRewriter rewriter(ops.front()->getContext(), typeConverter);
   for (auto *op : toConvert)
     if (failed(convert(rewriter, op)))
       return rewriter.getImpl().discardRewrites(), failure();
 
-  // Otherwise, the body conversion succeeded. Apply rewrites if this is not an
-  // analysis conversion.
-  if (mode == OpConversionMode::Analysis)
-    rewriter.getImpl().discardRewrites();
-  else
-    rewriter.getImpl().applyRewrites();
+  // Apply rewrites.
+
+  // Otherwise, the body conversion succeeded.
   return success();
 }
 
 //===----------------------------------------------------------------------===//
 // Type Conversion
 //===----------------------------------------------------------------------===//
 
 /// Remap an input of the original signature with a new set of types. The

 mlir::applyFullConversion(Operation *op, ConversionTarget &target,
                           const OwningRewritePatternList &patterns,
                           TypeConverter *converter) {
   return applyFullConversion(llvm::makeArrayRef(op), target, patterns,
                              converter);
 }
 
 /// Apply an analysis conversion on the given operations, and all nested
-/// operations. This method analyzes which operations would be successfully
-/// converted to the target if a conversion was applied. All operations that
-/// were found to be legalizable to the given 'target' are placed within the
-/// provided 'convertedOps' set; note that no actual rewrites are applied to the
-/// operations on success and only pre-existing operations are added to the set.
+/// operations. This method applies a partial conversion and analyzes on the way
+/// which operations would fail to convert successfully if a full conversion
+/// were attempted.
+///
+/// All operations that are found not to be legalizable to the given 'target'
+/// are placed within the provided 'unconvertedOps' set.
 LogicalResult mlir::applyAnalysisConversion(
     ArrayRef<Operation *> ops, ConversionTarget &target,
     const OwningRewritePatternList &patterns,
-    DenseSet<Operation *> &convertedOps, TypeConverter *converter) {
+    DenseSet<Operation *> &unconvertedOps, TypeConverter *converter) {
   OperationConverter opConverter(target, patterns, OpConversionMode::Analysis,
-                                 &convertedOps);
+                                 &unconvertedOps);
   return opConverter.convertOperations(ops, converter);
 }
 LogicalResult
 mlir::applyAnalysisConversion(Operation *op, ConversionTarget &target,
                               const OwningRewritePatternList &patterns,
-                              DenseSet<Operation *> &convertedOps,
+                              DenseSet<Operation *> &unconvertedOps,
                               TypeConverter *converter) {
   return applyAnalysisConversion(llvm::makeArrayRef(op), target, patterns,
-                                 convertedOps, converter);
+                                 unconvertedOps, converter);
 }

mlir/test/Transforms/test-legalizer-analysis.mlir

 // RUN: mlir-opt -allow-unregistered-dialect -test-legalize-patterns -verify-diagnostics -test-legalize-mode=analysis %s | FileCheck %s
-// expected-remark@-2 {{op 'module' is legalizable}}
-// expected-remark@-3 {{op 'module_terminator' is legalizable}}
-
-// expected-remark@+1 {{op 'func' is legalizable}}
 func @test(%arg0: f32) {
-  // expected-remark@+1 {{op 'test.illegal_op_a' is legalizable}}
   %result = "test.illegal_op_a"() : () -> (i32)
+  // expected-warning@+2 {{failed to legalize operation 'foo.region'}}
+  // expected-remark@+1 {{op 'foo.region' is not legalizable}}
   "foo.region"() ({
-      // expected-remark@+1 {{op 'test.invalid' is legalizable}}
       "test.invalid"() : () -> ()
   }) : () -> ()
+  // expected-warning@+2 {{failed to legalize operation 'std.return'}}
+  // expected-remark@+1 {{op 'std.return' is not legalizable}}
   return
 }
 
-// Check that none of the legalizable operations were modified.
+// Check that the legalizable operations were modified.
 // CHECK-LABEL: func @test
-// CHECK-NEXT: "test.illegal_op_a"
-// CHECK: "test.invalid"
+// CHECK-NEXT: "test.legal_op_a"
+// CHECK: "test.valid"

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

if (mode == ConversionMode::Full) {
       (void)applyFullConversion(getOperation(), target, patterns, &converter);
       return;
     }
 
     // Otherwise, handle an analysis conversion.
     assert(mode == ConversionMode::Analysis);
 
     // Analyze the convertible operations.
-    DenseSet<Operation *> legalizedOps;
+    DenseSet<Operation *> unlegalizedOps;
     if (failed(applyAnalysisConversion(getOperation(), target, patterns,
-                                       legalizedOps, &converter)))
+                                       unlegalizedOps, &converter)))
       return signalPassFailure();
 
     // Emit remarks for each legalizable operation.
-    for (auto *op : legalizedOps)
-      op->emitRemark() << "op '" << op->getName() << "' is legalizable";
+    for (auto *op : unlegalizedOps)
+      op->emitRemark() << "op '" << op->getName() << "' is not legalizable";
   }
 
   /// The mode of conversion to use.
   ConversionMode mode;
 };
 } // end anonymous namespace
 
 static llvm::cl::opt<TestLegalizePatternDriver::ConversionMode>