Fold certain ops during dialect conversion

[bondhugula]

Summary

Ops like builtin.unrealized_cast are meant to be folded during dialect
conversion. Update dialect conversion driver to fold ops with trait
FoldOnDialectConversion irrespective of whether they are marked illegal.

This revision also moves the check on "ignored ops" to earlier which has
functionally no impact but potentially an efficiency win.

[bondhugula]

Tweak comment.

[bondhugula]

Drop unrelated change.

[GMNGeoffrey]

IMO it would be better to have a way for the person specifying the conversion target to say "get rid of these ops if you can". Unrealized conversion cast is a bit of an odd duck, but it seems that "get rid of this during dialect conversion if possible" is more often a property of the conversion than a property of the op. I guess it could be either or both (where unrealized conversion is the relatively special case where it makes sense to have it on the op). It also feels like an arbitrary limitation that such an op can only be removed with folding and not with canonicalization or some other pattern.

[mehdi_amini]

it seems that "get rid of this during dialect conversion if possible" is more often a property of the conversion than a property of the op

That would be my feeling as well on this.

[bondhugula]

In D116661#3222768, @GMNGeoffrey wrote:

I guess it could be either or both (where unrealized conversion is the relatively special case where it makes sense to have it on the op). It also feels like an arbitrary limitation that such an op
can only be removed with folding and not with canonicalization or some other pattern.

Right - we may need both, i.e., special ops (with a specific trait) or ops that the user wants to remove. The latter would need a change to the API itself or a new "target" class like addLegalButOptimize which isn't the goal of this revision. Besides, if there is a use case to canonicalize and fold-out a specific set of ops or op types, you can always do that with another call to applyPatternsAndFoldGreedily or applyOpPatternsAndFold with just those ops' canonicalization patterns added. It's arguable as to whether you'd want dialect conversion to take on that concern. Separately, it's easy for me to make a change to run legalizeWithPattern as well on these FoldDuringDialectConversion ops so that you get both folding and canonicalization for them -- the trait will then have to be renamed SimplifyDuringDialectConversion.

[ftynse]

Does this actually remove the need to run the reconciliation pass in our test suite?

[bondhugula]

In D116661#3230841, @ftynse wrote:

Does this actually remove the need to run the reconciliation pass in our test suite?

I can check this but the goal is to remove the need to run this reconciliation for all downstream users as well -- I wouldn't want to merge this in otherwise. Besides LLVM, I can check IREE as well.

mlir/include/mlir/IR/BuiltinOps.td

def ModuleOp : Builtin_Op<"module", [
   let skipDefaultBuilders = 1;
 }
 
 //===----------------------------------------------------------------------===//
 // UnrealizedConversionCastOp
 //===----------------------------------------------------------------------===//
 
 def UnrealizedConversionCastOp : Builtin_Op<"unrealized_conversion_cast", [
-    DeclareOpInterfaceMethods<CastOpInterface>, NoSideEffect
+    DeclareOpInterfaceMethods<CastOpInterface>, FoldOnDialectConversion,
+    NoSideEffect
   ]> {
   let summary = "An unrealized conversion from one set of types to another";
   let description = [{
     An `unrealized_conversion_cast` operation represents an unrealized
     conversion from one set of types to another, that is used to enable the
     inter-mixing of different type systems. This operation should not be
     attributed any special representational or execution semantics, and is
     generally only intended to be used to satisfy the temporary intermixing of

mlir/include/mlir/IR/OpBase.td

 // X op Y == Y op X
 def Commutative  : NativeOpTrait<"IsCommutative">;
 // op op X == op X (unary) / X op X == X (binary)
 def Idempotent  : NativeOpTrait<"IsIdempotent">;
 // op op X == X
 def Involution  : NativeOpTrait<"IsInvolution">;
 // Op behaves like a constant.
 def ConstantLike : NativeOpTrait<"ConstantLike">;
+// A fold is always attempted on `op` during dialect conversion.
+def FoldOnDialectConversion : NativeOpTrait<"FoldOnDialectConversion">;
 // Op behaves like a function.
 def FunctionLike : NativeOpTrait<"FunctionLike">;
 // Op is isolated from above.
 def IsolatedFromAbove : NativeOpTrait<"IsIsolatedFromAbove">;
 // Op results are float or vectors/tensors thereof.
 def ResultsAreFloatLike : NativeOpTrait<"ResultsAreFloatLike">;
 // Op has the same operand type.
 def SameTypeOperands : NativeOpTrait<"SameTypeOperands">;

mlir/include/mlir/IR/OpDefinition.h

 ///   an `Elementwise` op.
 template <typename ConcreteType>
 struct Elementwise : public TraitBase<ConcreteType, Elementwise> {
   static LogicalResult verifyTrait(Operation *op) {
     return ::mlir::OpTrait::impl::verifyElementwise(op);
   }
 };
 
-/// This trait tags `Elementwise` operatons that can be systematically
+/// A trait for operations that are meant to be folded during dialect
+/// conversion even if they have been marked legal. This trait suits certain
+/// operations that are meant to be folded by design and not outlive the
+/// conversion pass whenever possible.
+template <typename ConcreteType>
+class FoldOnDialectConversion
+    : public TraitBase<ConcreteType, FoldOnDialectConversion> {
+public:
+  static LogicalResult verifyTrait(Operation *op) { return success(); }
+};
+
+/// This trait tags `Elementwise` operations that can be systematically
 /// scalarized. All vector/tensor operands and results are then replaced by
 /// scalars of the respective element type. Semantically, this is the operation
 /// on a single element of the vector/tensor.
 ///
 /// Rationale:
 /// Allow to define the vector/tensor semantics of elementwise operations based
 /// on the same op's behavior on scalars. This provides a constructive procedure
 /// for IR transformations to, e.g., create scalar loop bodies from tensor ops.

mlir/lib/Transforms/Utils/DialectConversion.cpp

LLVM_DEBUG({
 
     // If the operation has no regions, just print it here.
     if (op->getNumRegions() == 0) {
       op->print(logger.startLine(), OpPrintingFlags().printGenericOpForm());
       logger.getOStream() << "\n\n";
     }
   });
 
+  // Check to see if the operation is ignored and doesn't need to be converted.
+  if (rewriter.getImpl().isOpIgnored(op)) {
+    LLVM_DEBUG({
+      logSuccess(logger, "operation marked 'ignored' during conversion");
+      logger.startLine() << logLineComment;
+    });
+    return success();
+  }
+
+  // Ops with this trait are always folded even if marked legal.
+  if (op->hasTrait<OpTrait::FoldOnDialectConversion>() &&
+      succeeded(legalizeWithFold(op, rewriter))) {
+    LLVM_DEBUG({
+      logSuccess(logger,
+                 "operation with trait 'FoldDuringDialectConversion' folded");
+      logger.startLine() << logLineComment;
+    });
+    return success();
+  }
+
   // Check if this operation is legal on the target.
   if (auto legalityInfo = target.isLegal(op)) {
     LLVM_DEBUG({
       logSuccess(
           logger, "operation marked legal by the target{0}",
           legalityInfo->isRecursivelyLegal
               ? "; NOTE: operation is recursively legal; skipping internals"
               : "");
       logger.startLine() << logLineComment;
     });
 
     // If this operation is recursively legal, mark its children as ignored so
     // that we don't consider them for legalization.
     if (legalityInfo->isRecursivelyLegal)
       rewriter.getImpl().markNestedOpsIgnored(op);
     return success();
   }
 
-  // Check to see if the operation is ignored and doesn't need to be converted.
-  if (rewriter.getImpl().isOpIgnored(op)) {
-    LLVM_DEBUG({
-      logSuccess(logger, "operation marked 'ignored' during conversion");
-      logger.startLine() << logLineComment;
-    });
-    return success();
-  }
-
   // If the operation isn't legal, try to fold it in-place.
   // TODO: Should we always try to do this, even if the op is
   // already legal?
   if (succeeded(legalizeWithFold(op, rewriter))) {
     LLVM_DEBUG({
       logSuccess(logger, "operation was folded");
       logger.startLine() << logLineComment;
     });