[mlir] Allow memref.cast to change the signed/unsigned/signless element type

[ezhulenev]

Revert accidental change

[herhut]

Can you explain why this is needed? In usecases I have seen, the entire IR typically transitions from signed to signless types but I have not seen the need for a mixed form, yet.

[ezhulenev]

gpu.memset operations expects both arguments to have the same type:

%0 = arith.constant 0 : ui32
%1 = memef:alloca: memref<ui32>
gpu.memset %1, %0 : memref<ui32>, ui32

But there are no constant operations in "standard" dialects that can define unsigned constant, as a work around this memset can be defined with i32 types if we can cast the memref argument.

[nicolasvasilache]

This seems very undesirable.

memref.reinterpret_cast could be a better place to have this type of behavior as you can also change the sizes, strides etc and would allow richer elemental type changes.

[herhut]

In D133911#3792705, @ezhulenev wrote:

gpu.memset operations expects both arguments to have the same type:

%0 = arith.constant 0 : ui32
%1 = memef:alloca: memref<ui32>
gpu.memset %1, %0 : memref<ui32>, ui32

Would it also be possible to turn the target of the memset into a signless int? Is the fact that the memref is unsigned load bearing here or just an artifact?

I agree on cast being the wrong place, as it is concerned with shapes. I wonder why arith.bitcast rules out signed types. That seems the natural home for this, somehow.

mlir/include/mlir/Dialect/MemRef/IR/MemRefOps.td

let description = [{
     ```
     operation ::= ssa-id `=` `memref.cast` ssa-use `:` type `to` type
     ```
 
     The `memref.cast` operation converts a memref from one type to an equivalent
     type with a compatible shape. The source and destination types are
     compatible if:
 
-    a. Both are ranked memref types with the same element type, address space,
-    and rank and:
+    a. Both memrefs have compatible element types: element type is the same or
+    differs only it the signed/signless/unsigned bit.
+
+    b. Both are ranked memref types with compatible element types, address
+    space, and rank and:
       1. Both have the same layout or both have compatible strided layouts.
       2. The individual sizes (resp. offset and strides in the case of strided
          memrefs) may convert constant dimensions to dynamic dimensions and
          vice-versa.
 
     If the cast converts any dimensions from an unknown to a known size, then it
     acts as an assertion that fails at runtime if the dynamic dimensions
     disagree with resultant destination size.
let description = [{
     %4 = memref.cast %1 : memref<12x4xf32, strided<[?, ?], offset: ?>> to
                           memref<12x4xf32, strided<[4, 1], offset: 5>>
     // Erase static offset and stride information, replacing it with
     // dynamic information.
     %5 = memref.cast %1 : memref<12x4xf32, strided<[4, 1], offset: 5>> to
                           memref<12x4xf32, strided<[?, ?], offset: ?>>
     ```
 
-    b. Either or both memref types are unranked with the same element type, and
-    address space.
+    c. Either or both memref types are unranked with compatible element type,
+    and the same address space.
 
     Example:
 
     ```mlir
     Cast to concrete shape.
         %4 = memref.cast %1 : memref<*xf32> to memref<4x?xf32>
 
     Erase rank information.
let description = [{
     of transformations and foldings to be added as lowering progresses from
     higher-level dialect to lower-level dialects such as the LLVM dialect.
 
     The op requires a strided memref source operand. If the source operand is not
     a strided memref, then verification fails.
 
     This operation is also useful for completeness to the existing memref.dim op.
     While accessing strides, offsets and the base pointer independently is not
-    available, this is useful for composing with its natural complement op: 
+    available, this is useful for composing with its natural complement op:
     `memref.reinterpret_cast`.
 
     Intended Use Cases:
 
     The main use case is to expose the logic for manipulate memref metadata at a
-    higher level than the LLVM dialect. 
+    higher level than the LLVM dialect.
     This makes lowering more progressive and brings the following benefits:
       - not all users of MLIR want to lower to LLVM and the information to e.g.
         lower to library calls---like libxsmm---or to SPIR-V was not available.
-      - foldings and canonicalizations can happen at a higher level in MLIR: 
-        before this op existed, lowering to LLVM would create large amounts of 
+      - foldings and canonicalizations can happen at a higher level in MLIR:
+        before this op existed, lowering to LLVM would create large amounts of
         LLVMIR. Even when LLVM does a good job at folding the low-level IR from
         a performance perspective, it is unnecessarily opaque and inefficient to
         send unkempt IR to LLVM.
 
     Example:
 
     ```mlir
-      %base, %offset, %sizes:2, %strides:2 = 
-        memref.extract_strided_metadata %memref : 
+      %base, %offset, %sizes:2, %strides:2 =
+        memref.extract_strided_metadata %memref :
           memref<10x?xf32>, index, index, index, index, index
 
-      // After folding, the type of %m2 can be memref<10x?xf32> and further 
+      // After folding, the type of %m2 can be memref<10x?xf32> and further
       // folded to %memref.
       %m2 = memref.reinterpret_cast %base to
           offset: [%offset],
           sizes: [%sizes#0, %sizes#1],
           strides: [%strides#0, %strides#1]
         : memref<f32> to memref<?x?xf32, offset: ?, strides: [?, ?]>
     ```
   }];

mlir/lib/Dialect/GPU/IR/GPUDialect.cpp

LogicalResult SubgroupMmaComputeOp::verify() {
   if (aShape[1] != bShape[0] || aShape[0] != cShape[0] ||
       bShape[1] != cShape[1])
     return emitError("operand shapes do not satisfy matmul constraints");
 
   return success();
 }
 
 /// This is a common class used for patterns of the form
-/// "someop(memrefcast) -> someop".  It folds the source of any memref.cast
-/// into the root operation directly.
+/// "someop(memrefcast) -> someop".  It folds the source of a memref.cast
+/// into the root operation directly if source and user have the same memref
+/// element type.
 static LogicalResult foldMemRefCast(Operation *op) {
   bool folded = false;
+
+  // Returns true iff memref types have the same element type.
+  auto sameEltType = [](Value a, Value b) -> bool {
+    auto aEltType = a.getType().cast<MemRefType>().getElementType();
+    auto bEltType = b.getType().cast<MemRefType>().getElementType();
+    return aEltType == bEltType;
+  };
+
   for (OpOperand &operand : op->getOpOperands()) {
     auto cast = operand.get().getDefiningOp<mlir::memref::CastOp>();
-    if (cast) {
-      operand.set(cast.getOperand());
-      folded = true;
-    }
-  }
+    if (!cast || !sameEltType(operand.get(), cast.getOperand()))
+      continue;
+
   return success(folded);
 }
 
 LogicalResult MemcpyOp::fold(ArrayRef<Attribute> operands,
                              SmallVectorImpl<::mlir::OpFoldResult> &results) {
   return foldMemRefCast(*this);
 }
 

mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp

auto resultCast =
         rewriter.create<CastOp>(alloc.getLoc(), alloc.getType(), newAlloc);
 
     rewriter.replaceOp(alloc, {resultCast});
     return success();
   }
 };
 
 /// Fold alloc operations with no users or only store and dealloc uses.
-template <typename T>
-struct SimplifyDeadAlloc : public OpRewritePattern<T> {
+template <typename T> struct SimplifyDeadAlloc : public OpRewritePattern<T> {
   using OpRewritePattern<T>::OpRewritePattern;
 
   LogicalResult matchAndRewrite(T alloc,
                                 PatternRewriter &rewriter) const override {
     if (llvm::any_of(alloc->getUsers(), [&](Operation *op) {
           if (auto storeOp = dyn_cast<StoreOp>(op))
             return storeOp.getValue() == alloc;
           return !isa<DeallocOp>(op);
if (inputs.size() != 1 || outputs.size() != 1)
     return false;
   Type a = inputs.front(), b = outputs.front();
   auto aT = a.dyn_cast<MemRefType>();
   auto bT = b.dyn_cast<MemRefType>();
 
   auto uaT = a.dyn_cast<UnrankedMemRefType>();
   auto ubT = b.dyn_cast<UnrankedMemRefType>();
 
+  // Strips signed/unsigned bit from integer types.
+  auto stripSign = [](Type type) -> Type {
+    if (auto integer = type.dyn_cast<IntegerType>())
+      return IntegerType::get(type.getContext(), integer.getWidth());
+    return type;
+  };
+
   if (aT && bT) {
-    if (aT.getElementType() != bT.getElementType())
+    if (stripSign(aT.getElementType()) != stripSign(bT.getElementType()))
       return false;
     if (aT.getLayout() != bT.getLayout()) {
       int64_t aOffset, bOffset;
       SmallVector<int64_t, 4> aStrides, bStrides;
       if (failed(getStridesAndOffset(aT, aStrides, aOffset)) ||
           failed(getStridesAndOffset(bT, bStrides, bOffset)) ||
           aStrides.size() != bStrides.size())
         return false;
if (aT && bT) {
     if (!bT && !ubT)
       return false;
     // Unranked to unranked casting is unsupported
     if (uaT && ubT)
       return false;
 
     auto aEltType = (aT) ? aT.getElementType() : uaT.getElementType();
     auto bEltType = (bT) ? bT.getElementType() : ubT.getElementType();
-    if (aEltType != bEltType)
+    if (stripSign(aEltType) != stripSign(bEltType))
       return false;
 
     auto aMemSpace = (aT) ? aT.getMemorySpace() : uaT.getMemorySpace();
     auto bMemSpace = (bT) ? bT.getMemorySpace() : ubT.getMemorySpace();
     return aMemSpace == bMemSpace;
   }
 
   return false;

mlir/test/Dialect/MemRef/ops.mlir

 // CHECK-LABEL: func @memref_dealloc
 func.func @memref_dealloc() {
   %0 = memref.alloc() : memref<2xf32>
   %1 = memref.cast %0 : memref<2xf32> to memref<*xf32>
   memref.dealloc %1 : memref<*xf32>
   return
 }
 
+// CHECK-LABEL: func @memref_cast
+func.func @memref_cast() {
+  %0 = memref.alloc() : memref<2xi32>
+  %1 = memref.cast %0 : memref<2xi32> to memref<2xui32>
+  %2 = memref.cast %1 : memref<2xui32> to memref<2xsi32>
+  %3 = memref.cast %2 : memref<2xsi32> to memref<*xi32>
+  return
+}
 
 // CHECK-LABEL: func @memref_alloca_scope
 func.func @memref_alloca_scope() {
   memref.alloca_scope {
     memref.alloca_scope.return
   }
   return
 }
func.func @expand_collapse_shape_transposed_layout(
   %r0 = memref.expand_shape %m0 [[0], [1, 2]] :
     memref<?x?xf32, strided<[1, 10], offset: 0>> into
     memref<?x?x5xf32, strided<[1, 50, 10], offset: 0>>
   %rr0 = memref.collapse_shape %r0 [[0], [1, 2]] :
     memref<?x?x5xf32, strided<[1, 50, 10], offset: 0>> into
     memref<?x?xf32, strided<[1, 10], offset: 0>>
 
   %r1 = memref.expand_shape %m1 [[0, 1], [2], [3, 4]] :
-    memref<4x5x6xf32, strided<[1, ?, 1000], offset: 0>> into 
+    memref<4x5x6xf32, strided<[1, ?, 1000], offset: 0>> into
     memref<2x2x5x2x3xf32, strided<[2, 1, ?, 3000, 1000], offset: 0>>
   %rr1 = memref.collapse_shape %r1 [[0, 1], [2], [3, 4]] :
     memref<2x2x5x2x3xf32, strided<[2, 1, ?, 3000, 1000], offset: 0>> into
     memref<4x5x6xf32, strided<[1, ?, 1000], offset: 0>>
   return
 }
 
 // -----
^bb0(%old_value : f32):
       memref.atomic_yield %out : f32
   // CHECK: index_attr = 8 : index
   } { index_attr = 8 : index }
   return
 }
 
 // -----
 
-func.func @extract_strided_metadata(%memref : memref<10x?xf32>) 
+func.func @extract_strided_metadata(%memref : memref<10x?xf32>)
   -> memref<?x?xf32, strided<[?, ?], offset: ?>> {
 
   %base, %offset, %sizes:2, %strides:2 = memref.extract_strided_metadata %memref
     : memref<10x?xf32> -> memref<f32>, index, index, index, index, index
 
   %m2 = memref.reinterpret_cast %base to
       offset: [%offset],
       sizes: [%sizes#0, %sizes#1],
       strides: [%strides#0, %strides#1]
     : memref<f32> to memref<?x?xf32, strided<[?, ?], offset: ?>>
 
   return %m2: memref<?x?xf32, strided<[?, ?], offset: ?>>
 }