[mlir][LLVM] Add memref dialect cast

[kdobros]

Summary

This change adds support for casting memref from/to memref descriptor
struct via llvm.mlir.dialect_cast op.
To provide verification for this operation rewriting verify method was
needed, which strengthens previous checks by additionally checking
input-output types compatibility.

Such cast is useful when interfacing with code that doesn't need all
information that memref contains, for example requiring only bitcasted
pointer to data.

[ftynse]

Thanks!

Have you considered leveraging the TypeConverter instead of re-implementing half of the conversion rules in the verifier? It sounds like, in most cases, it should be sufficient to verify that LLVM::TypeConverter(...).convert(stdType) == llvmType, index type being the notable exception.

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

990–991

Please don't use auto unless it improves readability (e.g., the RHS has dyn_cast or is a constructor/factory).

990–991

Nit: drop trivial braces

993

Please terminate sentences with dots in all comments.

1002

Please provide tests for all user-visible error messages.

1011–1013

This looks like it supports casting f16 to !llvm.double, which isn't desirable. Consider using stdType.isF16()-style functions for stricter checks.

1014

Should we also check the bitwidth of the integer?

1029

Don't use else after return.

1047

Does this handle the special case of function calling conventions? The default convention expands the memref descriptor into multiple types, and we also have the "bare-pointer" convention that uses a plain pointer intsead.

1076

Consider inverting the condition and returning early to decrease horizontal indentation.

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

1 ↗

(On Diff #283270)

I don't understand the purpose of these tests. The code only tests the verifier of an op, so I would expect "roundtrip" tests that parse and print back accepted versions of the op and "invalid" tests that check for errors produced by the verifier.

[kdobros]

Addressed review comments.
Rewrote tests to check only the verifier - roundtrip.mlir and invalid.mlir.
Reduced number of error messages.

[kdobros]

Thanks for the review.

In D85323#2197144, @ftynse wrote:

Have you considered leveraging the TypeConverter instead of re-implementing half of the conversion rules in the verifier? It sounds like, in most cases, it should be sufficient to verify that LLVM::TypeConverter(...).convert(stdType) == llvmType, index type being the notable exception.

I have thought about this, but I'm not a fan of making verifier depend on conversion, and also working around index type would be tricky,
I think best way would be to move some type conversions, that don't depend on convention, into dialect and reuse them in both StandardToLLVM and in verifier.
This however seems like a little tricky and somewhat unrelated, so I've added TODO mentioning this.

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

1014

Seems reasonable if we're checking FloatType compatibility.

1047

Sorry, I'm not aware where calling convention could interact with dialect cast.
From what I understand irrespective of calling convention all function arguments are immediately packed back into memref descriptor struct.
Only place I guess pointer could be casted to memref is when moving from LLVM to standard and from LLVM comes only this pointer.
For now such case can be only supported with manually creating correct descriptor.
Automatically converting something like this would require adding rewrite pattern (or maybe legalization), so verifier rejects it, at least for now.
Could you maybe explain, if I'm misunderstanding something?

[ftynse]

I have thought about this, but I'm not a fan of making verifier depend on conversion, and also working around index type would be tricky,

And I am not a fan of duplicating the logic :) Working around index (and calling convention) should be easier when we have a data layout attached to the module, it will be possible to just get the proper index size from the layout and use it.

I think best way would be to move some type conversions, that don't depend on convention, into dialect and reuse them in both StandardToLLVM and in verifier.

I disagree. The reason why we have lib/Conversion/AToB is that conversions depend on both A and B and we did not want to have dialect B depend on A. Putting type conversions in the dialect is the exact opposite of this goal. (Well, to be honest, standard types aren't a part of the standard dialect, but that's just a quirk of standard stuff). There aren't many other cases where type casts are inserted, and the LLVM dialect has always been the trailblazer in the area, but we need to find a cleaner way. We put the "cast" op into the LLVM because we did not want to have it in standard as it could be mistaken for some magic "cast any type to any other type" op with no verification. Maybe we should consider a separate "cast" dialect with pluggable rules, or a separate dialect that can depend on both LLVM and Standard, but not making LLVM depend on Standard (or any other dialect that may lower into LLVM, e.g. Linalg also has type conversions).

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

995

It's valid but useless, should we add a canonicalization pattern that drops such casts? (not for this commit)

1047

I don't remember offhand how the casts and calling convention interact... The best thing to do is arguably try to run the newly introduced tests with both calling conventions and check that they still pass.

1082–1084

can we do something like

op.emitOpError("incorrect cast").attachNote("casting is supported for integers and floats, vectors and memrefs thereof");

that is more concise and decouples the error message from the additional explanation?

[kdobros]

Split emitted error into error and note

[kdobros]

In D85323#2202875, @ftynse wrote:

I think best way would be to move some type conversions, that don't depend on convention, into dialect and reuse them in both StandardToLLVM and in verifier.

I disagree. The reason why we have lib/Conversion/AToB is that conversions depend on both A and B and we did not want to have dialect B depend on A. Putting type conversions in the dialect is the exact opposite of this goal. (Well, to be honest, standard types aren't a part of the standard dialect, but that's just a quirk of standard stuff). There aren't many other cases where type casts are inserted, and the LLVM dialect has always been the trailblazer in the area, but we need to find a cleaner way. We put the "cast" op into the LLVM because we did not want to have it in standard as it could be mistaken for some magic "cast any type to any other type" op with no verification. Maybe we should consider a separate "cast" dialect with pluggable rules, or a separate dialect that can depend on both LLVM and Standard, but not making LLVM depend on Standard (or any other dialect that may lower into LLVM, e.g. Linalg also has type conversions).

Sure, I think it's still not clear to me where core IR ends and where Standard dialect begins, especially for types.
To me making verifier of dialect depend on conversion logic also breaks this separation goal, just in more roundabout way.
IMHO separate cast dialect with pluggable rules makes most sense. Biggest potential advantage here is that it could allow for casts between dialects that don't know about each other by looking for chain of cast rules that would satisfy input/output types. It could also allow to limit explosion of rules by using some pivot dialect like Standard. However implementation seems well above my current familiarity with MLIR.

---

I think I have addressed all comments.
If you think the change still looks good, could you maybe commit on my behalf? I don't have commit access.
My name and mail: Konrad Dobros <[e-mail]>

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

995

Sure, I actually though there already was such canonicalization since I saw that dialect casts with matching types are removed in lowering from standard.

1047

I looked at that first version where separate tests for Standard to LLVM lowering were added.
I saw no difference between calling conventions and memrefs being packed to descriptor struct immediately.

1082–1084

Sure, changed.
I haven't noticed that there is this attachNote method.

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 {
-    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 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: 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())));
+  Type inType = op.in().getType();
+  Type outType = op.getType();
+
+  // Casting between same types is valid.
+  if (inType == outType)
+
+  auto inLLVMType = inType.dyn_cast<LLVM::LLVMType>();
+  auto outLLVMType = outType.dyn_cast<LLVM::LLVMType>();
+
+  // If types are different then exactly one of them must be LLVM type.
+  if (!inLLVMType && !outLLVMType)
+    return op.emitOpError("incorrect cast between two non-LLVM types");
+  if (inLLVMType && outLLVMType)
+    return op.emitOpError("incorrect cast between two different LLVM types");
+
+  LLVMType llvmType = inLLVMType ? inLLVMType : outLLVMType;
+  Type otherType = inLLVMType ? outType : inType;
+
+  // TODO: Reduce duplication between below code and LLVMTypeConverter.
+  auto verifyScalarCastTypes = [](LLVMType llvmType,
+                                  Type stdType) -> LogicalResult {
+    if (stdType.isa<FloatType>())
+      return success((llvmType.isBFloatTy() && stdType.isBF16()) ||
+                     (llvmType.isDoubleTy() && stdType.isF64()) ||
+                     (llvmType.isFloatTy() && stdType.isF32()) ||
+                     (llvmType.isHalfTy() && stdType.isF16()));
+
+    if (auto integerType = stdType.dyn_cast<IntegerType>())
+      return success(llvmType.isIntegerTy() &&
+                     llvmType.getIntegerBitWidth() == integerType.getWidth());
+
+    return failure();
+  };
+
+  // Check valid casts in form of allow-list.
+  if (succeeded(verifyScalarCastTypes(llvmType, otherType)))
+    return success();
+
+  if (auto vectorType = otherType.dyn_cast<VectorType>()) {
+    if (!llvmType.isVectorTy())
+      return op.emitOpError("incorrect cast between vector and "
+                            "non-vector types");
+
+    LLVMType llvmElementType = llvmType.getVectorElementType();
+    Type vectorElementType = vectorType.getElementType();
+
+    if (failed(verifyScalarCastTypes(llvmElementType, vectorElementType)))
+      return op.emitOpError("vector elements must have matching "
+                            "non-index integer types or float types");
+
+    return success();
+
+  if (auto memRefType = otherType.dyn_cast<MemRefType>()) {
+    Type memRefElementType = memRefType.getElementType();
+
+    if (!llvmType.isStructTy())
+      return op.emitOpError(
+          "memref can be cast only from/to memref decriptor struct type");
+
+    if ((memRefType.getRank() > 0 && llvmType.getStructNumElements() != 5) ||
+        (memRefType.getRank() == 0 && llvmType.getStructNumElements() != 3))
+      return op.emitOpError(
+          "incorrect number of elements in memref descriptor struct");
+
+    LLVMType allocatedPtrType = llvmType.getStructElementType(0);
+    LLVMType alignedPtrType = llvmType.getStructElementType(1);
+
+    if (!allocatedPtrType.isPointerTy() || !alignedPtrType.isPointerTy())
+      return op.emitOpError("first and second element of memref descriptor "
+                            "must have pointer type");
+
+    LLVMType allocatedElementType = allocatedPtrType.getPointerElementTy();
+    LLVMType alignedElementType = alignedPtrType.getPointerElementTy();
+
+    if (allocatedElementType != alignedElementType ||
+        failed(verifyScalarCastTypes(allocatedElementType, memRefElementType)))
+      return op.emitOpError(
+          "first and second element of memref descriptor must have same types, "
+          "that is pointer to type convertible to memref's element type");
+
+    // For now skip checking other fields of memref descriptor as they rely on
+    // index lowering.
+
+  }
+
+  return op.emitOpError("incorrect cast").attachNote()
+         << "casting is supported for non-index integers and floats, "
+            "vectors and memrefs of those";
 }
 
 // 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())))

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@+2 {{'llvm.mlir.cast' op incorrect cast}}
+  // expected-note@+1  {{casting is supported for}}
   %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@+2 {{'llvm.mlir.cast' op incorrect cast}}
+  // expected-note@+1  {{casting is supported for}}
   %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}}
+  // expected-error@+2 {{'llvm.mlir.cast' op incorrect cast}}
+  // expected-note@+1  {{casting is supported for}}
   %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>">
 }

mlir/test/Dialect/LLVMIR/invalid.mlir

 }
 
 // -----
 
 func @invalid_ordering_in_fence() {
   // expected-error @+1 {{can be given only acquire, release, acq_rel, and seq_cst orderings}}
   llvm.fence syncscope("agent") monotonic
 }
+
+// -----
+
+// Check dialect cast operation verifier.
+
+func @invaild_dialect_cast() {
+  %0 = "op"() : () -> (index)
+  // expected-error @+2 {{incorrect cast}}
+  // expected-note @+1  {{casting is supported for}}
+  %1 = llvm.mlir.cast %0 : index to !llvm.i64
+}
+
+// -----
+
+func @invaild_dialect_cast() {
+  %0 = "op"() : () -> (f32)
+  // expected-error @+1 {{incorrect cast between two non-LLVM types}}
+  %1 = llvm.mlir.cast %0 : f32 to f16
+}
+
+// -----
+
+func @invaild_dialect_cast() {
+  %0 = "op"() : () -> (!llvm.i32)
+  // expected-error @+1 {{incorrect cast between two different LLVM types}}
+  %1 = llvm.mlir.cast %0 : !llvm.i32 to !llvm.float
+}
+
+// -----
+
+func @invaild_dialect_cast() {
+  %0 = "op"() : () -> (f32)
+  // expected-error @+2 {{incorrect cast}}
+  // expected-note @+1  {{casting is supported for}}
+  %1 = llvm.mlir.cast %0 : f32 to !llvm.half
+}
+
+// -----
+
+func @invaild_dialect_cast() {
+  %0 = "op"() : () -> (vector<2x3xf32>)
+  // expected-error @+1 {{only 1-d vector is allowed}}
+  %1 = llvm.mlir.cast %0 : vector<2x3xf32> to !llvm<"[2 x <3 x float>]">
+}
+
+// -----
+
+func @invaild_dialect_cast() {
+  %0 = "op"() : () -> (vector<2xf32>)
+  // expected-error @+1 {{incorrect cast between vector and non-vector types}}
+  %1 = llvm.mlir.cast %0 : vector<2xf32> to !llvm.float
+}
+
+// -----
+
+func @invaild_dialect_cast() {
+  %0 = "op"() : () -> (vector<2xf32>)
+  // expected-error @+1 {{vector elements must have matching non-index integer types or float types}}
+  %1 = llvm.mlir.cast %0 : vector<2xf32> to !llvm<"<2 x i32>">
+}
+
+// -----
+
+func @invaild_dialect_cast() {
+  %0 = "op"() : () -> (memref<2xf32>)
+  // expected-error @+1 {{memref can be cast only from/to memref decriptor struct type}}
+  %1 = llvm.mlir.cast %0 : memref<2xf32> to !llvm.float
+}
+
+// -----
+
+func @invaild_dialect_cast() {
+  %0 = "op"() : () -> (memref<2xf32>)
+  // expected-error @+1 {{incorrect number of elements in memref descriptor struct}}
+  %1 = llvm.mlir.cast %0 : memref<2xf32> to !llvm<"{float*, float*, i64}">
+}
+
+// -----
+
+func @invaild_dialect_cast() {
+  %0 = "op"() : () -> (memref<2xf32>)
+  // expected-error @+1 {{first and second element of memref descriptor must have pointer type}}
+  %1 = llvm.mlir.cast %0 : memref<2xf32> to !llvm<"{float, float*, i64, [2 x i64], [2 x i64]}">
+}
+
+// -----
+
+func @invaild_dialect_cast() {
+  %0 = "op"() : () -> (memref<2xf32>)
+  // expected-error @+1 {{first and second element of memref descriptor must have same types, that is pointer to type convertible to memref's element type}}
+  %1 = llvm.mlir.cast %0 : memref<2xf32> to !llvm<"{half*, half*, i64, [2 x i64], [2 x i64]}">
+}

mlir/test/Dialect/LLVMIR/roundtrip.mlir

func @useFenceInst() {
   // CHECK:  syncscope("agent") seq_cst
   llvm.fence syncscope("agent") seq_cst
   // CHECK:  seq_cst
   llvm.fence syncscope("") seq_cst
   // CHECK:  release
   llvm.fence release
   return
 }
+
+// Dialect cast operation.
+
+// CHECK-LABEL: @castI32
+//  CHECK-SAME:     %[[ARG0:.*]]: !llvm.i32
+//       CHECK:   %[[CAST:.*]] = llvm.mlir.cast %[[ARG0]] : !llvm.i32 to i32
+//       CHECK:   llvm.mlir.cast %[[CAST]] : i32 to !llvm.i32
+func @castI32(%arg0: !llvm.i32) {
+  %0 = llvm.mlir.cast %arg0 : !llvm.i32 to i32
+  %1 = llvm.mlir.cast %0 : i32 to !llvm.i32
+  return
+}
+
+// CHECK-LABEL: @castFloat
+//  CHECK-SAME:     %[[ARG0:.*]]: !llvm.float
+//       CHECK:   %[[CAST:.*]] = llvm.mlir.cast %[[ARG0]] : !llvm.float to f32
+//       CHECK:   llvm.mlir.cast %[[CAST]] : f32 to !llvm.float
+func @castFloat(%arg0: !llvm.float) {
+  %0 = llvm.mlir.cast %arg0 : !llvm.float to f32
+  %1 = llvm.mlir.cast %0 : f32 to !llvm.float
+  return
+}
+
+// CHECK-LABEL: @castVector
+//  CHECK-SAME:     %[[ARG0:.*]]: !llvm<"<3 x float>">
+//       CHECK:   %[[CAST:.*]] = llvm.mlir.cast %[[ARG0]] : !llvm<"<3 x float>"> to vector<3xf32>
+//       CHECK:   llvm.mlir.cast %[[CAST]] : vector<3xf32> to !llvm<"<3 x float>">
+func @castVector(%arg0: !llvm<"<3 x float>">) {
+  %0 = llvm.mlir.cast %arg0 : !llvm<"<3 x float>"> to vector<3xf32>
+  %1 = llvm.mlir.cast %0 : vector<3xf32> to !llvm<"<3 x float>">
+  return
+}
+
+// CHECK-LABEL: @castMemref
+//  CHECK-SAME:     %[[ARG0:.*]]: !llvm<"{ float*, float*, i64, [2 x i64], [2 x i64] }">
+//       CHECK:   %[[CAST:.*]] = llvm.mlir.cast %[[ARG0]] : !llvm<"{ float*, float*, i64, [2 x i64], [2 x i64] }"> to memref<2x3xf32>
+//       CHECK:   llvm.mlir.cast %[[CAST]] : memref<2x3xf32> to !llvm<"{ float*, float*, i64, [2 x i64], [2 x i64] }">
+func @castMemref(%arg0: !llvm<"{ float*, float*, i64, [2 x i64], [2 x i64] }">) {
+  %0 = llvm.mlir.cast %arg0 : !llvm<"{ float*, float*, i64, [2 x i64], [2 x i64] }"> to memref<2x3xf32>
+  %1 = llvm.mlir.cast %0 : memref<2x3xf32> to !llvm<"{ float*, float*, i64, [2 x i64], [2 x i64] }">
+  return
+}
+
+// CHECK-LABEL: @castMemref0D
+//  CHECK-SAME:     %[[ARG0:.*]]: !llvm<"{ float*, float*, i64 }">
+//       CHECK:   %[[CAST:.*]] = llvm.mlir.cast %[[ARG0]] : !llvm<"{ float*, float*, i64 }"> to memref<f32>
+//       CHECK:   llvm.mlir.cast %[[CAST]] : memref<f32> to !llvm<"{ float*, float*, i64 }">
+func @castMemref0D(%arg0: !llvm<"{ float*, float*, i64 }">) {
+  %0 = llvm.mlir.cast %arg0 : !llvm<"{ float*, float*, i64 }"> to memref<f32>
+  %1 = llvm.mlir.cast %0 : memref<f32> to !llvm<"{ float*, float*, i64 }">
+  return
+}