[mlir][memref] Add `memref.elt_bitcast` op

[Hardcode84]

Summary

This op allows to reinterpret between memrefs with different element types.
See the discussion https://discourse.llvm.org/t/rfc-memref-bitcasting/66395

[bondhugula]

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

2134

This can be marked MemRefsNormalizable as well.

2138

an existing memref ... as a ...

2138

with a different ...

2139

have the same shape, ...

2140

-> The source memref should be allocated with ...

2141

with the proper

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

3167–3168

Can use auto here since the type is on the RHS.

3170–3174

Nit: Messages can all start in lowercase. No need of full stop at the end since there would be a trailing part.

mlir/test/Dialect/MemRef/canonicalize.mlir

852–854

853–854

Nit: Indent the start of the actual MLIR fragment correctly.

[bondhugula]

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

3176–3180

This check looks incomplete - it doesn't cover vector elemental type cases?

[Hardcode84]

review comments, better folder

[Hardcode84]

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

3176–3180

Checking vector and index sizes require DataLayout, I'm not sure if we should do this in verifier.

[bondhugula]

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

3176–3180

Can elt type bitcasts like those from vector<2xi16> to <i16> or to <i64> still be valid for some data layout? (We don't need to worry about index types here. memrefs of index type may not even be a case of interest for such elt type bit casts.)

[Hardcode84]

Can elt type bitcasts like those from vector<2xi16> to <i16> or to <i64> still be valid for some data layout?

More realistic case will be vector<3xi32> <-> vector<2xi64> cast, where vector<3xi32> may or may not be padded to 128 bits depending on layout.
Default data layout implementation assumes it is https://github.com/llvm/llvm-project/blob/8f60eee9144cd4178938d231ecccb65c43f78cde/mlir/lib/Interfaces/DataLayoutInterfaces.cpp#L78, but it can be overridden.
And I don't want to replicate this vector logic in verifier anyway. So the question is, is it ok to access data layout structures in op verifier?

[Hardcode84]

And another question, do we really need such strict check in verifier? IMO, it better to do a trivial int/float check in verifier and then potentially reject it on memref->llvm conversion level where we have all required info about layout and which seems a proper place for this.

[bondhugula]

Dismissing change request.

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

3191–3196

Code comment needed here and/or above.

[Hardcode84]

rebase, add comment

[Hardcode84]

ping

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

let extraClassDeclaration = [{
     MemRefType getMemRefType() {
       return getMemref().getType().cast<MemRefType>();
     }
   }];
   let hasFolder = 1;
   let hasVerifier = 1;
 }
 
+//===----------------------------------------------------------------------===//
+// EltBitcastOp
+//===----------------------------------------------------------------------===//
+
+def EltBitcastOp : MemRef_Op<"elt_bitcast", [
+    Pure, ViewLikeOpInterface, SameOperandsShape, MemRefsNormalizable
+  ]> {
+  let summary = "memref elementwise bitcast op";
+  let description = [{
+    `elt_bitcast` op reinterprets an existing memref as a memref with a
+    different element type. Source and destination memrefs must have the same
+    shape, layout, memory space and element size. The source memref should be
+    allocated with the proper alignment and memory type for accessing
+    destination memref elements. Memref data is not modified or copied by this
+    op.
+
+    Example:
+
+    ```mlir
+    %dst = memref.elt_bitcast %src : memref<?xf64> to memref<?xvector<2xi32>>
+    ```
+  }];
+
+  let arguments = (ins AnyMemRef:$source);
+  let results = (outs AnyMemRef:$result);
+
+  let extraClassDeclaration = [{
+      ::mlir::Value getViewSource() { return getSource(); }
+  }];
+
+  let hasFolder = 1;
+  let hasVerifier = 1;
+
+  let assemblyFormat = "$source attr-dict `:` type($source) `to` type($result)";
+}
+
 #endif // MEMREF_OPS

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

 OpFoldResult AtomicRMWOp::fold(ArrayRef<Attribute> operands) {
   /// atomicrmw(memrefcast) -> atomicrmw
   if (succeeded(foldMemRefCast(*this, getValue())))
     return getResult();
   return OpFoldResult();
 }
 
 //===----------------------------------------------------------------------===//
+// EltBitcastOp
+//===----------------------------------------------------------------------===//
+
+LogicalResult EltBitcastOp::verify() {
+  auto srcType = getSource().getType().cast<MemRefType>();
+  auto dstType = getResult().getType().cast<MemRefType>();
+  if (srcType.getShape() != dstType.getShape())
+    return emitError("memref bitcast shape mismatch");
+  if (srcType.getLayout() != dstType.getLayout())
+    return emitError("memref bitcast layout mismatch");
+  if (srcType.getMemorySpace() != dstType.getMemorySpace())
+    return emitError("memref bitcast memory space mismatch");
+
+  Type srcElem = srcType.getElementType();
+  Type dstElem = dstType.getElementType();
+  if (srcElem.isIntOrFloat() && dstElem.isIntOrFloat() &&
+      srcElem.getIntOrFloatBitWidth() != dstElem.getIntOrFloatBitWidth())
+    return emitError("memref bitcast element size mismatch");
+  return success();
+}
+
+OpFoldResult EltBitcastOp::fold(llvm::ArrayRef<Attribute> /*operands*/) {
+  Value src = getSource();
+  Type dstType = getResult().getType();
+  if (src.getType() == dstType)
+    return src;
+
+  // Search cast chain for the matching type.
+  while (auto parent = src.getDefiningOp<EltBitcastOp>()) {
+    Value parentSource = parent.getSource();
+    if (parentSource.getType() == dstType)
+      return parentSource;
+
+    src = parentSource;
+  }
+
+  return nullptr;
+}
+
+//===----------------------------------------------------------------------===//
 // TableGen'd op method definitions
 //===----------------------------------------------------------------------===//
 
 #define GET_OP_CLASSES
 #include "mlir/Dialect/MemRef/IR/MemRefOps.cpp.inc"

mlir/test/Dialect/MemRef/canonicalize.mlir

 // CHECK-NOT: memref.realloc
 // CHECK: return %[[SRC]]
 func.func @memref_realloc_dead(%src : memref<2xf32>, %v : f32) -> memref<2xf32>{
   %0 = memref.realloc %src : memref<2xf32> to memref<4xf32>
   %i2 = arith.constant 2 : index
   memref.store %v, %0[%i2] : memref<4xf32>
   return %src : memref<2xf32>
 }
+
+// -----
+
+// CHECK-LABEL: func @memref_bitcast_same_type
+//  CHECK-SAME:   (%[[ARG:.*]]: memref<?xf64>)
+//       CHECK:   return %[[ARG]]
+func.func @memref_bitcast_same_type(%src : memref<?xf64>) -> memref<?xf64> {
+  %0 = memref.elt_bitcast %src : memref<?xf64> to memref<?xf64>
+  return %0 : memref<?xf64>
+}
+
+// -----
+
+// CHECK-LABEL: func @memref_bitcast_chain
+//  CHECK-SAME:   (%[[ARG:.*]]: memref<?xf64>)
+//       CHECK:   return %[[ARG]]
+func.func @memref_bitcast_chain(%src : memref<?xf64>) -> memref<?xf64> {
+  %0 = memref.elt_bitcast %src : memref<?xf64> to memref<?xi64>
+  %1 = memref.elt_bitcast %0 : memref<?xi64> to memref<?xvector<2xi32>>
+  %2 = memref.elt_bitcast %1 : memref<?xvector<2xi32>> to memref<?xf64>
+  return %2 : memref<?xf64>
+}

mlir/test/Dialect/MemRef/ops.mlir

func.func @memref_realloc_dd(%src : memref<?xf32>, %d: index)
   return %0 : memref<?xf32>
 }
 
 // CHECK-LABEL: func @memref_extract_aligned_pointer
 func.func @memref_extract_aligned_pointer(%src : memref<?xf32>) -> index {
   %0 = memref.extract_aligned_pointer_as_index %src : memref<?xf32> -> index
   return %0 : index
 }
+
+// -----
+
+// CHECK-LABEL: func @memref_bitcast
+//  CHECK-SAME:   (%[[ARG:.*]]: memref<?xf64>)
+//       CHECK:   %[[RES:.*]] = memref.elt_bitcast %[[ARG]] : memref<?xf64> to memref<?xvector<2xi32>>
+//       CHECK:   return %[[RES]]
+func.func @memref_bitcast(%src : memref<?xf64>) -> memref<?xvector<2xi32>> {
+  %0 = memref.elt_bitcast %src : memref<?xf64> to memref<?xvector<2xi32>>
+  return %0 : memref<?xvector<2xi32>>
+}