[mlir][vector] Add experimental memref cast operation.

[ThomasRaoux]

Summary

Add an experimental op to be able to cast static memref to a memref of the same size with different shape/element type. This would lower to a pointer bitcast.
This is needed for vector to GPU path as we need to be able to load/store from memref with different granularity.

[mehdi_amini]

Does this op really belong to the vector dialect?

It seems like you're plugging a hole somewhere, but it isn't clear what or why this is desirable. Can you explain this in a thread on Discourse?

[aartbik]

In D85058#2188725, @mehdi_amini wrote:

Does this op really belong to the vector dialect?

It seems like you're plugging a hole somewhere, but it isn't clear what or why this is desirable. Can you explain this in a thread on Discourse?

I agree that understanding the need better would help. We already have reshape/type_cast/shape_cast ops with somewhat specific behaviors, and I want to make sure we design the right operations (preferable just a few).

[ThomasRaoux]

In D85058#2189117, @aartbik wrote:

In D85058#2188725, @mehdi_amini wrote:

Does this op really belong to the vector dialect?

It seems like you're plugging a hole somewhere, but it isn't clear what or why this is desirable. Can you explain this in a thread on Discourse?

I agree that understanding the need better would help. We already have reshape/type_cast/shape_cast ops with somewhat specific behaviors, and I want to make sure we design the right operations (preferable just a few).

Sounds good, I started a discussion here: https://llvm.discourse.group/t/memref-cast/1514
The reason why I added it to the vector dialect was because it seemed analogue to the existing vector.type_cast and it is meant to be used for vectorization. However I understand this may be confusing and I don't have a strong opinion on this.

[ThomasRaoux]

I'm abandoning this change as it is not needed after https://reviews.llvm.org/D85244.

mlir/include/mlir/Dialect/Vector/VectorOps.td

let description = [{
     %1 = vector.flat_transpose %0 { rows = 4: i32, columns = 4: i32 }
        : (vector<16xf32>) -> vector<16xf32>
     ```
   }];
   let verifier = ?;
   let assemblyFormat = "$matrix attr-dict `:` type($matrix) `->` type($res)";
 }
 
+// Experimental op used by the path vector to GPU. This allows generating IR
+// with correct type for cases where we need to load/store from a memref with
+// different granularity.
+def Vector_ReinterpretCastOp : Vector_Op<"reinterpret_cast", [NoSideEffect]>,
+    Arguments<(ins StaticShapeMemRefOf<[AnyType]>:$memref)>,
+    Results<(outs StaticShapeMemRefOf<[AnyType]>:$result)> {
+  let summary = "reinterpret_cast converts a memref into an equivalent memref.";
+
+  let description = [{
+    Performs a conversion from a memref to an memref other memref with different
+    shape and or element type. This is needed to be able to access a memref with
+    different granularity without using subviews. This is analogue to a pointer
+    cast.
+
+    This is an experimental op and not all lowering will be supported.
+
+    Syntax:
+
+    ```
+    operation ::= `vector.reinterpret_cast` ssa-use : memref-type to memref-type
+    ```
+
+    Example:
+
+    ```mlir
+    %A = alloc() : memref<128x2xvector<4xi32>>
+    %RA = vector.reinterpret_cast %A : memref<128x2xvector<4xi32>> to memref<128x32xi8>
+    ```
+  }];
+
+  let extraClassDeclaration = [{
+    MemRefType getMemRefType() {
+      return memref().getType().cast<MemRefType>();
+    }
+    MemRefType getResultMemRefType() {
+      return getResult().getType().cast<MemRefType>();
+    }
+  }];
+
+  let assemblyFormat = [{
+    $memref attr-dict `:` type($memref) `to` type($result)
+  }];
+}
+
 #endif // VECTOR_OPS

mlir/lib/Dialect/Vector/VectorOps.cpp

static LogicalResult verify(CreateMaskOp op) {
   // Verify that an operand was specified for each result vector each dimension.
   if (op.getNumOperands() !=
       op.getResult().getType().cast<VectorType>().getRank())
     return op.emitOpError(
         "must specify an operand for each result vector dimension");
   return success();
 }
 
+//===----------------------------------------------------------------------===//
+// ReinterpretCastOp
+//===----------------------------------------------------------------------===//
+
+static LogicalResult verify(ReinterpretCastOp op) {
+  MemRefType canonicalType = canonicalizeStridedLayout(op.getMemRefType());
+  if (!canonicalType.getAffineMaps().empty())
+    return op.emitOpError("expects operand to be a memref with no layout");
+  if (!op.getResultMemRefType().getAffineMaps().empty())
+    return op.emitOpError("expects result to be a memref with no layout");
+  if (op.getResultMemRefType().getMemorySpace() !=
+      op.getMemRefType().getMemorySpace())
+    return op.emitOpError("expects result in same memory space");
+  if (op.getResultMemRefType().getSizeInBits() !=
+      op.getMemRefType().getSizeInBits())
+    return op.emitOpError(
+        "expects result type to have the same size as operand");
+  return success();
+}
+
 namespace {
 
 // Pattern to rewrite a CreateMaskOp with a ConstantMaskOp.
 class CreateMaskFolder final : public OpRewritePattern<CreateMaskOp> {
 public:
   using OpRewritePattern<CreateMaskOp>::OpRewritePattern;
 
   LogicalResult matchAndRewrite(CreateMaskOp createMaskOp,

mlir/lib/IR/StandardTypes.cpp

 int64_t ShapedType::getSizeInBits() const {
   assert(hasStaticShape() &&
          "cannot get the bit size of an aggregate with a dynamic shape");
 
   auto elementType = getElementType();
   if (elementType.isIntOrFloat())
     return elementType.getIntOrFloatBitWidth() * getNumElements();
 
-  // Tensors can have vectors and other tensors as elements, other shaped types
-  // cannot.
-  assert(isa<TensorType>() && "unsupported element type");
-  assert((elementType.isa<VectorType, TensorType>()) &&
+  // Tensors and Memref can have vectors and other tensors as elements, other
+  // shaped types cannot.
+  assert((isa<TensorType>() || isa<MemRefType>()) &&
+         "unsupported element type");
+  assert((elementType.isa<VectorType, TensorType, MemRefType>()) &&
          "unsupported tensor element type");
   return getNumElements() * elementType.cast<ShapedType>().getSizeInBits();
 }
 
 ArrayRef<int64_t> ShapedType::getShape() const {
   switch (getKind()) {
   case StandardTypes::Vector:
     return cast<VectorType>().getShape();

mlir/test/Dialect/Vector/invalid.mlir

 }
 
 // -----
 
 func @scatter_dim_mask_mismatch(%base: memref<?xf32>, %indices: vector<16xi32>, %mask: vector<17xi1>, %value: vector<16xf32>) {
   // expected-error@+1 {{'vector.scatter' op expected value dim to match mask dim}}
   vector.scatter %base, %indices, %mask, %value : vector<16xi32>, vector<17xi1>, vector<16xf32> into memref<?xf32>
 }
+
+// -----
+
+func @reinterpret_cast_size_mismatch(%base: memref<7x24xi32>) -> memref<7x10xvector<4xi32>> {
+  // expected-error@+1 {{'vector.reinterpret_cast' op expects result type to have the same size as operand}}
+  %0 = vector.reinterpret_cast %base : memref<7x24xi32> to memref<7x10xvector<4xi32>>
+  return %0 : memref<7x10xvector<4xi32>>
+}

mlir/test/Dialect/Vector/ops.mlir

func @gather_and_scatter(%base: memref<?xf32>, %indices: vector<16xi32>, %mask: vector<16xi1>) {
   // CHECK: %[[X:.*]] = vector.gather %{{.*}}, %{{.*}}, %{{.*}} : (memref<?xf32>, vector<16xi32>, vector<16xi1>) -> vector<16xf32>
   %0 = vector.gather %base, %indices, %mask : (memref<?xf32>, vector<16xi32>, vector<16xi1>) -> vector<16xf32>
   // CHECK: %[[Y:.*]] = vector.gather %{{.*}}, %{{.*}}, %{{.*}}, %[[X]] : (memref<?xf32>, vector<16xi32>, vector<16xi1>, vector<16xf32>) -> vector<16xf32>
   %1 = vector.gather %base, %indices, %mask, %0 : (memref<?xf32>, vector<16xi32>, vector<16xi1>, vector<16xf32>) -> vector<16xf32>
   // CHECK: vector.scatter %{{.*}}, %{{.*}}, %{{.*}}, %[[Y]] : vector<16xi32>, vector<16xi1>, vector<16xf32> into memref<?xf32>
   vector.scatter %base, %indices, %mask, %1 : vector<16xi32>, vector<16xi1>, vector<16xf32> into memref<?xf32>
   return
 }
+
+// CHECK-LABEL: @reinterpret_cast
+func @reinterpret_cast(%base: memref<7x12xi8>) -> memref<7x3xi32> {
+  // CHECK: %{{.*}} = vector.reinterpret_cast %{{.*}} : memref<7x12xi8> to memref<7x3xi32>
+  %0 = vector.reinterpret_cast %base : memref<7x12xi8> to memref<7x3xi32>
+  return %0 : memref<7x3xi32>
+}
+
+// CHECK-LABEL: @reinterpret_cast_vector
+func @reinterpret_cast_vector(%base: memref<7x24xi32>) -> memref<7x6xvector<4xi32>> {
+  // CHECK: %{{.*}} = vector.reinterpret_cast %{{.*}} : memref<7x24xi32> to memref<7x6xvector<4xi32>>
+  %0 = vector.reinterpret_cast %base : memref<7x24xi32> to memref<7x6xvector<4xi32>>
+  return %0 : memref<7x6xvector<4xi32>>
+}