[MLIR][Standard] Fix cast materialization for incompatible types

[frgossen]

Summary

Only create llvm.mlir.cast when types are cast compatible and fail otherwise.

[herhut]

mlir/lib/Conversion/StandardToLLVM/StandardToLLVM.cpp

174

This is addressing the FIXME below partially. A better way could be to expose the verifyMLIRCastType lambda from LLVMDialect::DialectCastOp::verify as a static function on LLVMDialect::DialectCastOp and call that here? Then the FIXME could disappear.

[silvas]

Is there anything about this specific to index type? Can this issue happen with any user-defined type?

[rriddle]

In D85634#2207661, @silvas wrote:

Is there anything about this specific to index type? Can this issue happen with any user-defined type?

The solution in this revision looks like a hack. Returning the input when a materialization is requested defeats the whole purpose of materializations, and only works because the assert isn't turned on yet. Realistically, the underlying issue can happen with any type right now. I have a proper fix in the pipeline that should solve this case, where the conversion infra removes casts that become redundant.

[rriddle]

mlir/lib/Conversion/StandardToLLVM/StandardToLLVM.cpp

174

Agreed, ideally DialectCastOp should expose a static bool areCastCompatible(Type a, Type b); method like many of the other cast operations do. We could call that from here which would remove the TODO, but if that fails this method should fail to materialize.

[mehdi_amini]

The simultaneous application of patterns from SCF to standard and from standard to LLVM rewrites the following to invalid IR.

Do we actually need a simultaneous application?

[herhut]

In D85634#2208550, @mehdi_amini wrote:

The simultaneous application of patterns from SCF to standard and from standard to LLVM rewrites the following to invalid IR.

Do we actually need a simultaneous application?

We can also do this in steps but combining patterns should work. Or at least it should fail in a reasonable way and not produce invalid IR.

[herhut]

In D85634#2207669, @rriddle wrote:

In D85634#2207661, @silvas wrote:

Is there anything about this specific to index type? Can this issue happen with any user-defined type?

The solution in this revision looks like a hack. Returning the input when a materialization is requested defeats the whole purpose of materializations, and only works because the assert isn't turned on yet. Realistically, the underlying issue can happen with any type right now. I have a proper fix in the pipeline that should solve this case, where the conversion infra removes casts that become redundant.

Would the right thing be to return llvm::None to make this fail during conversion? The problem here is not that the cast is redundant but that the cast cannot be materialized as it is invalid.

Or should we extend the special cast operation such that it supports all casts as long as its operand and result types can be converted using the type converter?

[rriddle]

In D85634#2209388, @herhut wrote:

In D85634#2207669, @rriddle wrote:

In D85634#2207661, @silvas wrote:

Is there anything about this specific to index type? Can this issue happen with any user-defined type?

The solution in this revision looks like a hack. Returning the input when a materialization is requested defeats the whole purpose of materializations, and only works because the assert isn't turned on yet. Realistically, the underlying issue can happen with any type right now. I have a proper fix in the pipeline that should solve this case, where the conversion infra removes casts that become redundant.

Would the right thing be to return llvm::None to make this fail during conversion? The problem here is not that the cast is redundant but that the cast cannot be materialized as it is invalid.

Yes, if a conversion can't be materialized this should return llvm::None.

Or should we extend the special cast operation such that it supports all casts as long as its operand and result types can be converted using the type converter?

Realistically, we probably want to be able to materialize casts for a majority(if not all) of the conversions that the default LLVM type converter provides. This makes it much easier to have progressive lowering, which is something we should strive to enable.

[ftynse]

In D85634#2209497, @rriddle wrote:

Realistically, we probably want to be able to materialize casts for a majority(if not all) of the conversions that the default LLVM type converter provides. This makes it much easier to have progressive lowering, which is something we should strive to enable.

Yes, and this should let us disentangle the existing whatever-to-LLVM from std-to-LLVM conversions.

[frgossen]

Address comments

[frgossen]

Fully agree.

The mlir-cuda-runner tests are still failing because they materialized invalid casts. They did not fail because the verifier is not invoked before the lowering to LLVM (?) and the the lowering for llvm.mlir.cast is actually able to deal with the index case at least.

@rriddle, has the fix you mentioned landed yet? I could not find it.

[ftynse]

mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp

1021

Let's add a comment in addCastCompatible that asks to update the message in the verifier when more cases are added.

[frgossen]

Add comment

[frgossen]

@rriddle, the mlir-cuda-runner tests are still failing because they materialize invalid casts. They did not fail earlier because there the verifier is not invoked when lowering to LLVM (?) and the the lowering for llvm.mlir.cast is actually able to deal with the index, it seems.

Has the fix you mentioned landed yet? I could not find it.

[rriddle]

In D85634#2238441, @frgossen wrote:

@rriddle, the mlir-cuda-runner tests are still failing because they materialize invalid casts. They did not fail earlier because there the verifier is not invoked when lowering to LLVM (?) and the the lowering for llvm.mlir.cast is actually able to deal with the index, it seems.

Has the fix you mentioned landed yet? I could not find it.

Unfortunately my attention has been pulled in another direction, and I likely won't have time to return to this at least for a few weeks.

[frgossen]

@rriddle, fair enough. Can you let me know once it lands? I'm almost sure that fixes the failing tests here.
I will leave this CL until then to avoid duplicate work.

[herhut]

Is there an update on this? Is this still relevant?

[frgossen]

I believe this is still relevant as the FIXME is still present. The mlir-cuda-runner tests still fail.
@rriddle

mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td

let description = [{
 
     Notice that index type is not supported, as it's Standard-specific.
 
     Example:
       llvm.mlir.cast %v : f16 to llvm.half
       llvm.mlir.cast %v : llvm.float to f32
       llvm.mlir.cast %v : !llvm<"<2 x float>"> to vector<2xf32>
   }];
+  let extraClassDeclaration = [{
+    /// Return true if `a` and `b` are valid operand and result types for the
+    /// operation.
+    static bool areCastCompatible(Type a, Type b);
+  }];
   let assemblyFormat = "$in attr-dict `:` type($in) `to` type($res)";
   let verifier = "return ::verify(*this);";
 }
 
 // Operations that correspond to LLVM intrinsics. With MLIR operation set being
 // extendable, there is no reason to introduce a hard boundary between "core"
 // operations and intrinsics. However, we systematically prefix them with
 // "intr." to avoid potential name clashes.

mlir/lib/Conversion/StandardToLLVM/StandardToLLVM.cpp

LLVMTypeConverter::LLVMTypeConverter(MLIRContext *ctx,
   });
   // Add generic source and target materializations to handle cases where
   // non-LLVM types persist after an LLVM conversion.
   addSourceMaterialization([&](OpBuilder &builder, Type resultType,
                                ValueRange inputs,
                                Location loc) -> Optional<Value> {
     if (inputs.size() != 1)
       return llvm::None;
-    // FIXME: These should check LLVM::DialectCastOp can actually be constructed
-    // from the input and result.
+    if (!LLVM::DialectCastOp::areCastCompatible(inputs[0].getType(),
+                                                resultType))
+      return llvm::None;
     return builder.create<LLVM::DialectCastOp>(loc, resultType, inputs[0])
         .getResult();
   });
   addTargetMaterialization([&](OpBuilder &builder, Type resultType,
                                ValueRange inputs,
                                Location loc) -> Optional<Value> {
     if (inputs.size() != 1)
       return llvm::None;
-    // FIXME: These should check LLVM::DialectCastOp can actually be constructed
-    // from the input and result.
+    if (!LLVM::DialectCastOp::areCastCompatible(inputs[0].getType(),
+                                                resultType))
+      return llvm::None;
     return builder.create<LLVM::DialectCastOp>(loc, resultType, inputs[0])
         .getResult();
   });
 }
 
 /// Returns the MLIR context.
 MLIRContext &LLVMTypeConverter::getContext() {
   return *getDialect()->getContext();

mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp

static void printGlobalOp(OpAsmPrinter &p, GlobalOp op) {
   if (!initializer.empty())
     p.printRegion(initializer, /*printEntryBlockArgs=*/false);
 }
 
 //===----------------------------------------------------------------------===//
 // Verifier for LLVM::DialectCastOp.
 //===----------------------------------------------------------------------===//
 
-static LogicalResult verify(DialectCastOp op) {
-  auto verifyMLIRCastType = [&op](Type type) -> LogicalResult {
+/// Please update the error message in `verify(DialectCastOp)` when changing
+/// this function.
+bool mlir::LLVM::DialectCastOp::areCastCompatible(Type a, Type b) {
+  auto verifyMLIRCastType = [&](Type type) -> bool {
     if (auto llvmType = type.dyn_cast<LLVM::LLVMType>()) {
       if (llvmType.isVectorTy())
         llvmType = llvmType.getVectorElementType();
       if (llvmType.isIntegerTy() || llvmType.isBFloatTy() ||
           llvmType.isHalfTy() || llvmType.isFloatTy() ||
           llvmType.isDoubleTy()) {
-        return success();
+        return true;
       }
-      return op.emitOpError("type must be non-index integer types, float "
-                            "types, or vector of mentioned types.");
+      return false;
     }
     if (auto vectorType = type.dyn_cast<VectorType>()) {
       if (vectorType.getShape().size() > 1)
-        return op.emitOpError("only 1-d vector is allowed");
+        return false;
       type = vectorType.getElementType();
     }
     if (type.isSignlessIntOrFloat())
-      return success();
+      return true;
     // Note that memrefs are not supported. We currently don't have a use case
     // for it, but even if we do, there are challenges:
     // * if we allow memrefs to cast from/to memref descriptors, then the
     // semantics of the cast op depends on the implementation detail of the
     // descriptor.
     // * if we allow memrefs to cast from/to bare pointers, some users might
     // alternatively want metadata that only present in the descriptor.
     //
     // TODO: re-evaluate the memref cast design when it's needed.
-    return op.emitOpError("type must be non-index integer types, float types, "
-                          "or vector of mentioned types.");
+    return false;
   };
-  return failure(failed(verifyMLIRCastType(op.in().getType())) ||
-                 failed(verifyMLIRCastType(op.getType())));
+  return verifyMLIRCastType(a) && verifyMLIRCastType(b);
+}
+
+static LogicalResult verify(DialectCastOp op) {
+  if (!DialectCastOp::areCastCompatible(op.in().getType(), op.getType()))
+    return op.emitOpError("types must be non-index integer types, float types, "
+                          "or 1-d vector of mentioned types.");
+  return success();
 }
 
 // Parses one of the keywords provided in the list `keywords` and returns the
 // position of the parsed keyword in the list. If none of the keywords from the
 // list is parsed, returns -1.
 static int parseOptionalKeywordAlternative(OpAsmParser &parser,
                                            ArrayRef<StringRef> keywords) {
   for (auto en : llvm::enumerate(keywords)) {

mlir/test/Conversion/StandardToLLVM/invalid.mlir

 // RUN: mlir-opt %s -convert-std-to-llvm -verify-diagnostics -split-input-file
 
 func @mlir_cast_to_llvm(%0 : index) -> !llvm.i64 {
-  // expected-error@+1 {{'llvm.mlir.cast' op type must be non-index integer types, float types, or vector of mentioned types}}
+  // expected-error@+1 {{'llvm.mlir.cast' op types must be non-index integer types, float types, or 1-d vector of mentioned types}}
   %1 = llvm.mlir.cast %0 : index to !llvm.i64
   return %1 : !llvm.i64
 }
 
 // -----
 
 func @mlir_cast_from_llvm(%0 : !llvm.i64) -> index {
-  // expected-error@+1 {{'llvm.mlir.cast' op type must be non-index integer types, float types, or vector of mentioned types}}
+  // expected-error@+1 {{'llvm.mlir.cast' op types must be non-index integer types, float types, or 1-d vector of mentioned types}}
   %1 = llvm.mlir.cast %0 : !llvm.i64 to index
   return %1 : index
 }
 
 // -----
 
 func @mlir_cast_to_llvm_int(%0 : i32) -> !llvm.i64 {
   // expected-error@+1 {{failed to legalize operation 'llvm.mlir.cast' that was explicitly marked illegal}}
   %1 = llvm.mlir.cast %0 : i32 to !llvm.i64
   return %1 : !llvm.i64
 }
 
 // -----
 
 func @mlir_cast_to_llvm_vec(%0 : vector<1x1xf32>) -> !llvm.vec<1 x float> {
-  // expected-error@+1 {{'llvm.mlir.cast' op only 1-d vector is allowed}}
+  // expected-error@+1 {{'llvm.mlir.cast' op types must be non-index integer types, float types, or 1-d vector of mentioned types}}
   %1 = llvm.mlir.cast %0 : vector<1x1xf32> to !llvm.vec<1 x float>
   return %1 : !llvm.vec<1 x float>
 }