[MLIR] Add llvm.mlir.cast op for semantic preserving cast between dialect types.

[timshen]

Summary

See discussion here: https://llvm.discourse.group/t/rfc-dialect-type-cast-op/538/11

[timshen]

Add more tests.

[timshen]

Error message.

[ftynse]

Makes sense to me, mostly nits.

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

689

Nit: drop "MLIR" from the def name, we are not using it in other operations (constant, undef)

704

Nit: the implicit convention for cast operations is to use type($in) 'to' type($res), let's follow it here as well.

mlir/lib/Conversion/StandardToLLVM/ConvertStandardToLLVM.cpp

1820

Nit: I'd say that the llvm.mlir.cast itself isn't legal, not the conversion.

Actually, I'm not convinced we want to report an error from the pattern at all. There may be cases where a partial conversion is sufficient and the caller will deal with the remaining casts properly, yet we would have reported the error to the user by that time....

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

898

Don't you also want to support memrefs?

[timshen]

Update per comments.

[timshen]

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

689

Well it has to be named as "SomethingCast", as LLVM_CastOp already exists. Renamed to "DialectCast".

mlir/lib/Conversion/StandardToLLVM/ConvertStandardToLLVM.cpp

1820

Good point. Removed the error reporting.

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

898

memref is complicated because of the descriptor. Do we allow the conversion between a memref and the aligned base pointer, or a memref and the descriptor? Neither seems trivial to implement to me, nor I really need it at the moment. Added a TODO.

[ftynse]

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

898

I'd say between memref and descriptor. But you are right that it adds a lot of complexity since the descriptor type depends on the parameters of the lowering.

[rriddle]

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

704

nit: use "" for single line strings.

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

900

nit: Please drop all trivial braces.

914

nit: Please use early return instead. That allows for dropping the else.

[timshen]

Stylistic change and add a comment on why not support memref at the moment.

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

 def LLVM_ConstantOp
     : LLVM_OneResultOp<"mlir.constant", [NoSideEffect]>,
       Arguments<(ins AnyAttr:$value)>,
       LLVM_Builder<"$res = getLLVMConstant($_resultType, $value, $_location);">
 {
   let assemblyFormat = "`(` $value `)` attr-dict `:` type($res)";
 }
 
+def LLVM_DialectCastOp : LLVM_Op<"mlir.cast", [NoSideEffect]>,
+                         Results<(outs AnyType:$res)>,
+                         Arguments<(ins AnyType:$in)> {
+  let summary = "Type cast between LLVM dialect and Standard.";
+  let description = [{
+    llvm.mlir.cast op casts between Standard and LLVM dialects. It only changes
+    the dialect, but does not change compile-time or runtime semantics.
+
+    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 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.
 
 class LLVM_UnaryIntrinsicOp<string func, list<OpTrait> traits = []> :
     LLVM_OneResultIntrOp<func, [], [0],
            !listconcat([NoSideEffect, SameOperandsAndResultType], traits)>,

mlir/lib/Conversion/StandardToLLVM/ConvertStandardToLLVM.cpp

if (srcType.isa<MemRefType>() && dstType.isa<MemRefType>()) {
       auto loadOp = rewriter.create<LLVM::LoadOp>(loc, castPtr);
       rewriter.replaceOp(op, loadOp.getResult());
     } else {
       llvm_unreachable("Unsuppored unranked memref to unranked memref cast");
     }
   }
 };
 
+struct DialectCastOpLowering
+    : public LLVMLegalizationPattern<LLVM::DialectCastOp> {
+  using LLVMLegalizationPattern<LLVM::DialectCastOp>::LLVMLegalizationPattern;
+
+  PatternMatchResult
+  matchAndRewrite(Operation *op, ArrayRef<Value> operands,
+                  ConversionPatternRewriter &rewriter) const override {
+    auto castOp = cast<LLVM::DialectCastOp>(op);
+    OperandAdaptor<LLVM::DialectCastOp> transformed(operands);
+    if (transformed.in().getType() !=
+        typeConverter.convertType(castOp.getType())) {
+      return matchFailure();
+    }
+    rewriter.replaceOp(op, transformed.in());
+    return matchSuccess();
+  }
+};
+
 // A `dim` is converted to a constant for static sizes and to an access to the
 // size stored in the memref descriptor for dynamic sizes.
 struct DimOpLowering : public LLVMLegalizationPattern<DimOp> {
   using LLVMLegalizationPattern<DimOp>::LLVMLegalizationPattern;
 
   PatternMatchResult
   matchAndRewrite(Operation *op, ArrayRef<Value> operands,
                   ConversionPatternRewriter &rewriter) const override {
patterns.insert<
       CallOpLowering,
       CeilFOpLowering,
       CmpFOpLowering,
       CmpIOpLowering,
       CondBranchOpLowering,
       CopySignOpLowering,
       CosOpLowering,
       ConstLLVMOpLowering,
+      DialectCastOpLowering,
       DivFOpLowering,
       ExpOpLowering,
       LogOpLowering,
       Log10OpLowering,
       Log2OpLowering,
       FPExtLowering,
       FPTruncLowering,
       IndexCastOpLowering,
Option<bool> emitCWrappers{
       llvm::cl::desc("Emit C-compatible wrapper functions"),
       llvm::cl::init(false)};
 };
 } // end namespace
 
 mlir::LLVMConversionTarget::LLVMConversionTarget(MLIRContext &ctx)
     : ConversionTarget(ctx) {
   this->addLegalDialect<LLVM::LLVMDialect>();
+  this->addIllegalOp<LLVM::DialectCastOp>();
 }
 
 std::unique_ptr<OpPassBase<ModuleOp>>
 mlir::createLowerToLLVMPass(bool useAlloca, bool useBarePtrCallConv,
                             bool emitCWrappers) {
   return std::make_unique<LLVMLoweringPass>(useAlloca, useBarePtrCallConv,
                                             emitCWrappers);
 }
 
 static PassRegistration<LLVMLoweringPass>
     pass(PASS_NAME, "Convert scalar and vector operations from the "
                     "Standard to the LLVM dialect");

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

if (op.getValueOrNull().dyn_cast_or_null<StringAttr>())
     return;
   p << " : " << op.type();
 
   Region &initializer = op.getInitializerRegion();
   if (!initializer.empty())
     p.printRegion(initializer, /*printEntryBlockArgs=*/false);
 }
 
+//===----------------------------------------------------------------------===//
+// Verifier for LLVM::DialectCastOp.
+//===----------------------------------------------------------------------===//
+
+static LogicalResult verify(DialectCastOp op) {
+  auto verifyMLIRCastType = [&op](Type type) -> LogicalResult {
+    if (auto llvmType = type.dyn_cast<LLVM::LLVMType>()) {
+      if (llvmType.isVectorTy())
+        llvmType = llvmType.getVectorElementType();
+      if (llvmType.isIntegerTy() || llvmType.isHalfTy() ||
+          llvmType.isFloatTy() || llvmType.isDoubleTy()) {
+        return success();
+      }
+      return op.emitOpError("type must be non-index integer types, float "
+                            "types, or vector of mentioned types.");
+    }
+    if (auto vectorType = type.dyn_cast<VectorType>()) {
+      if (vectorType.getShape().size() > 1)
+        return op.emitOpError("only 1-d vector is allowed");
+      type = vectorType.getElementType();
+    }
+    if (type.isSignlessIntOrFloat())
+      return success();
+    // 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(timshen): 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 failure(failed(verifyMLIRCastType(op.in().getType())) ||
+                 failed(verifyMLIRCastType(op.getType())));
+}
+
 // 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)) {
     if (succeeded(parser.parseOptionalKeyword(en.value())))
       return en.index();

mlir/test/Conversion/StandardToLLVM/convert-to-llvmir.mlir

func @assume_alignment(%0 : memref<4x4xf16>) {
   // CHECK-NEXT: %[[MASK:.*]] = llvm.mlir.constant(15 : index) : !llvm.i64
   // CHECK-NEXT: %[[INT:.*]] = llvm.ptrtoint %[[PTR]] : !llvm<"half*"> to !llvm.i64
   // CHECK-NEXT: %[[MASKED_PTR:.*]] = llvm.and %[[INT]], %[[MASK:.*]] : !llvm.i64
   // CHECK-NEXT: %[[CONDITION:.*]] = llvm.icmp "eq" %[[MASKED_PTR]], %[[ZERO]] : !llvm.i64
   // CHECK-NEXT: "llvm.intr.assume"(%[[CONDITION]]) : (!llvm.i1) -> ()
   assume_alignment %0, 16 : memref<4x4xf16>
   return
 }
+
+// -----
+
+// CHECK-LABEL: func @mlir_cast_to_llvm
+// CHECK-SAME: %[[ARG:.*]]:
+func @mlir_cast_to_llvm(%0 : vector<2xf16>) -> !llvm<"<2 x half>"> {
+  %1 = llvm.mlir.cast %0 : vector<2xf16> to !llvm<"<2 x half>">
+  // CHECK-NEXT: llvm.return %[[ARG]]
+  return %1 : !llvm<"<2 x half>">
+}
+
+// CHECK-LABEL: func @mlir_cast_from_llvm
+// CHECK-SAME: %[[ARG:.*]]:
+func @mlir_cast_from_llvm(%0 : !llvm<"<2 x half>">) -> vector<2xf16> {
+  %1 = llvm.mlir.cast %0 : !llvm<"<2 x half>"> to vector<2xf16>
+  // CHECK-NEXT: llvm.return %[[ARG]]
+  return %1 : vector<2xf16>
+}
+
+// -----
+
+// CHECK-LABEL: func @mlir_cast_to_llvm
+// CHECK-SAME: %[[ARG:.*]]:
+func @mlir_cast_to_llvm(%0 : f16) -> !llvm.half {
+  %1 = llvm.mlir.cast %0 : f16 to !llvm.half
+  // CHECK-NEXT: llvm.return %[[ARG]]
+  return %1 : !llvm.half
+}
+
+// CHECK-LABEL: func @mlir_cast_from_llvm
+// CHECK-SAME: %[[ARG:.*]]:
+func @mlir_cast_from_llvm(%0 : !llvm.half) -> f16 {
+  %1 = llvm.mlir.cast %0 : !llvm.half to f16
+  // CHECK-NEXT: llvm.return %[[ARG]]
+  return %1 : f16
+}

mlir/test/Conversion/StandardToLLVM/invalid.mlir

-// RUN: mlir-opt %s -verify-diagnostics -split-input-file
+// RUN: mlir-opt %s -convert-std-to-llvm -verify-diagnostics -split-input-file
 
 #map1 = affine_map<(d0, d1)[s0, s1, s2] -> (d0 * s1 + s0 + d1 * s2)>
 
 func @invalid_memref_cast(%arg0: memref<?x?xf64>) {
   %c1 = constant 1 : index
   %c0 = constant 0 : index
-  // expected-error@+1: 'std.memref_cast' op operand #0 must be unranked.memref of any type values or memref of any type values,
+  // expected-error@+1 {{'std.memref_cast' op operand #0 must be unranked.memref of any type values or memref of any type values, but got '!llvm<"{ double*, double*, i64, [2 x i64], [2 x i64] }">'}}
   %5 = memref_cast %arg0 : memref<?x?xf64> to memref<?x?xf64, #map1>
   %25 = std.subview %5[%c0, %c0][%c1, %c1][] : memref<?x?xf64, #map1> to memref<?x?xf64, #map1>
   return
 }
 
+// -----
+
+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}}
+  %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}}
+  %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<"<1 x float>"> {
+  // expected-error@+1 {{'llvm.mlir.cast' op only 1-d vector is allowed}}
+  %1 = llvm.mlir.cast %0 : vector<1x1xf32> to !llvm<"<1 x float>">
+  return %1 : !llvm<"<1 x float>">
+}