[mlir] Add MemRefReinterpretCastOp definition to Standard.

[pifon2a]

Summary

Reuse most code for printing/parsing/verification from SubViewOp.

https://llvm.discourse.group/t/rfc-standard-memref-cast-ops/1454/15

[pifon2a]

Rebase

mlir/lib/Conversion/StandardToLLVM/StandardToLLVM.cpp

if (srcType.isa<MemRefType>() && dstType.isa<UnrankedMemRefType>()) {
       auto loadOp = rewriter.create<LLVM::LoadOp>(loc, castPtr);
       rewriter.replaceOp(op, loadOp.getResult());
     } else {
       llvm_unreachable("Unsupported unranked memref to unranked memref cast");
     }
   }
 };
 
+struct MemRefReinterpretCastOpLowering
+    : public ConvertOpToLLVMPattern<MemRefReinterpretCastOp> {
+  using ConvertOpToLLVMPattern<MemRefReinterpretCastOp>::ConvertOpToLLVMPattern;
+
+  LogicalResult
+  matchAndRewrite(Operation *op, ArrayRef<Value> operands,
+                  ConversionPatternRewriter &rewriter) const override {
+    auto castOp = cast<MemRefReinterpretCastOp>(op);
+    MemRefReinterpretCastOp::Adaptor adaptor(operands, op->getAttrDictionary());
+    Type srcType = castOp.source().getType();
+
+    Value descriptor;
+    if (failed(ConvertSourceMemRefToDescriptor(rewriter, srcType, castOp,
+                                               adaptor, &descriptor)))
+      return failure();
+    rewriter.replaceOp(op, {descriptor});
+    return success();
+  }
+
+private:
+  LogicalResult
+  ConvertSourceMemRefToDescriptor(ConversionPatternRewriter &rewriter,
+                                  Type srcType, MemRefReinterpretCastOp castOp,
+                                  MemRefReinterpretCastOp::Adaptor adaptor,
+                                  Value *descriptor) const {
+    MemRefType targetMemRefType =
+        castOp.getResult().getType().cast<MemRefType>();
+    auto llvmTargetDescriptorTy = typeConverter.convertType(targetMemRefType)
+                                      .dyn_cast_or_null<LLVM::LLVMType>();
+    if (!llvmTargetDescriptorTy || !llvmTargetDescriptorTy.isStructTy())
+      return failure();
+
+    // Create descriptor.
+    Location loc = castOp.getLoc();
+    auto desc = MemRefDescriptor::undef(rewriter, loc, llvmTargetDescriptorTy);
+
+    // Set allocated and aligned pointers.
+    Value allocatedPtr, alignedPtr;
+    ExtractPointers(loc, rewriter, castOp.source(), adaptor.source(),
+                    &allocatedPtr, &alignedPtr);
+    desc.setAllocatedPtr(rewriter, loc, allocatedPtr);
+    desc.setAlignedPtr(rewriter, loc, alignedPtr);
+
+    // Set offset.
+    if (castOp.isDynamicOffset(0)) {
+      desc.setOffset(rewriter, loc, adaptor.offsets()[0]);
+    } else {
+      desc.setConstantOffset(rewriter, loc, castOp.getStaticOffset(0));
+    }
+
+    // Set sizes and strides.
+    unsigned dynSizeId = 0;
+    unsigned dynStrideId = 0;
+    for (unsigned i = 0, e = targetMemRefType.getRank(); i < e; ++i) {
+      if (castOp.isDynamicSize(i)) {
+        desc.setSize(rewriter, loc, i, adaptor.sizes()[dynSizeId++]);
+      } else {
+        desc.setConstantSize(rewriter, loc, i, castOp.getStaticSize(i));
+      }
+      if (castOp.isDynamicStride(i)) {
+        desc.setStride(rewriter, loc, i, adaptor.strides()[dynStrideId++]);
+      } else {
+        desc.setConstantStride(rewriter, loc, i, castOp.getStaticStride(i));
+      }
+    }
+    *descriptor = desc;
+    return success();
+  }
+
+  void ExtractPointers(Location loc, ConversionPatternRewriter &rewriter,
+                       Value originalOperand, Value convertedOperand,
+                       Value *allocatedPtr, Value *alignedPtr) const {
+    Type operandType = originalOperand.getType();
+    if (operandType.isa<MemRefType>()) {
+      MemRefDescriptor desc(convertedOperand);
+      *allocatedPtr = desc.allocatedPtr(rewriter, loc);
+      *alignedPtr = desc.alignedPtr(rewriter, loc);
+      return;
+    }
+
+    unsigned memorySpace =
+        operandType.cast<UnrankedMemRefType>().getMemorySpace();
+    LLVM::LLVMType elementType =
+        typeConverter
+            .convertType(
+                operandType.cast<UnrankedMemRefType>().getElementType())
+            .cast<LLVM::LLVMType>();
+    LLVM::LLVMType elementPtrPtrType =
+        elementType.getPointerTo(memorySpace).getPointerTo(memorySpace);
+
+    // Extract pointer to the underlying ranked memref descriptor and cast it to
+    // ElemType**.
+    UnrankedMemRefDescriptor unrankedDesc(convertedOperand);
+    Value underlyingDescPtr = unrankedDesc.memRefDescPtr(rewriter, loc);
+    Value elementPtrPtr = rewriter.create<LLVM::BitcastOp>(
+        loc, elementPtrPtrType, underlyingDescPtr);
+
+    LLVM::LLVMType int32Type =
+        typeConverter.convertType(rewriter.getI32Type()).cast<LLVM::LLVMType>();
+
+    // Extract and set allocated pointer.
+    Value zero = rewriter.create<LLVM::ConstantOp>(
+        loc, int32Type, rewriter.getI32IntegerAttr(memorySpace));
+    Value base_gep = rewriter.create<LLVM::GEPOp>(
+        loc, elementPtrPtrType, elementPtrPtr, ValueRange({zero}));
+    *allocatedPtr = rewriter.create<LLVM::LoadOp>(loc, base_gep);
+
+    // Extract and set aligned pointer.
+    Value one = rewriter.create<LLVM::ConstantOp>(
+        loc, int32Type, rewriter.getI32IntegerAttr(1));
+    Value aligned_gep = rewriter.create<LLVM::GEPOp>(
+        loc, elementPtrPtrType, elementPtrPtr, ValueRange({one}));
+    *alignedPtr = rewriter.create<LLVM::LoadOp>(loc, aligned_gep);
+  }
+};
+
 struct DialectCastOpLowering
     : public ConvertOpToLLVMPattern<LLVM::DialectCastOp> {
   using ConvertOpToLLVMPattern<LLVM::DialectCastOp>::ConvertOpToLLVMPattern;
 
   LogicalResult
   matchAndRewrite(Operation *op, ArrayRef<Value> operands,
                   ConversionPatternRewriter &rewriter) const override {
     auto castOp = cast<LLVM::DialectCastOp>(op);

mlir/test/mlir-cpu-runner/memref_reinterpret_cast.mlir

+// RUN: mlir-opt %s -convert-scf-to-std  -convert-std-to-llvm | mlir-cpu-runner -e main -entry-point-result=void -shared-libs=%mlir_runner_utils_dir/libmlir_runner_utils%shlibext,%mlir_runner_utils_dir/libmlir_c_runner_utils%shlibext | FileCheck %s
+
+func @print_memref_f32(memref<*xf32>) attributes { llvm.emit_c_interface }
+
+func @main() -> () {
+  %c0 = constant 0 : index
+  %c1 = constant 1 : index
+
+  // Initialize input.
+  %input = alloc() : memref<2x3xf32>
+  %dim_x = dim %input, %c0 : memref<2x3xf32>
+  %dim_y = dim %input, %c1 : memref<2x3xf32>
+  scf.parallel (%i, %j) = (%c0, %c0) to (%dim_x, %dim_y) step (%c1, %c1) {
+    %prod = muli %i,  %dim_y : index
+    %val = addi %prod, %j : index
+    %val_i64 = index_cast %val : index to i64
+    %val_f32 = sitofp %val_i64 : i64 to f32
+    store %val_f32, %input[%i, %j] : memref<2x3xf32>
+  }
+  %unranked_input = memref_cast %input : memref<2x3xf32> to memref<*xf32>
+  call @print_memref_f32(%unranked_input) : (memref<*xf32>) -> ()
+  // CHECK: rank = 2 offset = 0 sizes = [2, 3] strides = [3, 1]
+  // CHECK-NEXT: [0,   1,   2]
+  // CHECK-NEXT: [3,   4,   5]
+
+  // Test cases.
+  call @cast_ranked_memref_to_static_shape(%input) : (memref<2x3xf32>) -> ()
+  call @cast_ranked_memref_to_dynamic_shape(%input) : (memref<2x3xf32>) -> ()
+  call @cast_unranked_memref_to_static_shape(%input) : (memref<2x3xf32>) -> ()
+  call @cast_unranked_memref_to_dynamic_shape(%input) : (memref<2x3xf32>) -> ()
+  return
+}
+
+func @cast_ranked_memref_to_static_shape(%input : memref<2x3xf32>) {
+  %output = memref_reinterpret_cast %input to
+           offset: [0], sizes: [6, 1], strides: [1, 1]
+           : memref<2x3xf32> to memref<6x1xf32>
+
+  %unranked_output = memref_cast %output
+      : memref<6x1xf32> to memref<*xf32>
+  call @print_memref_f32(%unranked_output) : (memref<*xf32>) -> ()
+  // CHECK: rank = 2 offset = 0 sizes = [6, 1] strides = [1, 1] data =
+  // CHECK-NEXT: [0],
+  // CHECK-NEXT: [1],
+  // CHECK-NEXT: [2],
+  // CHECK-NEXT: [3],
+  // CHECK-NEXT: [4],
+  // CHECK-NEXT: [5]
+  return
+}
+
+func @cast_ranked_memref_to_dynamic_shape(%input : memref<2x3xf32>) {
+  %c0 = constant 0 : index
+  %c1 = constant 1 : index
+  %c6 = constant 6 : index
+  %output = memref_reinterpret_cast %input to
+           offset: [%c0], sizes: [%c1, %c6], strides: [%c6, %c1]
+           : memref<2x3xf32> to memref<?x?xf32, offset: ?, strides: [?, ?]>
+
+  %unranked_output = memref_cast %output
+      : memref<?x?xf32, offset: ?, strides: [?, ?]> to memref<*xf32>
+  call @print_memref_f32(%unranked_output) : (memref<*xf32>) -> ()
+  // CHECK: rank = 2 offset = 0 sizes = [1, 6] strides = [6, 1] data =
+  // CHECK-NEXT: [0,   1,   2,   3,   4,   5]
+  return
+}
+
+func @cast_unranked_memref_to_static_shape(%input : memref<2x3xf32>) {
+  %unranked_input = memref_cast %input : memref<2x3xf32> to memref<*xf32>
+  %output = memref_reinterpret_cast %unranked_input to
+           offset: [0], sizes: [6, 1], strides: [1, 1]
+           : memref<*xf32> to memref<6x1xf32>
+
+  %unranked_output = memref_cast %output
+      : memref<6x1xf32> to memref<*xf32>
+  call @print_memref_f32(%unranked_output) : (memref<*xf32>) -> ()
+  // CHECK: rank = 2 offset = 0 sizes = [6, 1] strides = [1, 1] data =
+  // CHECK-NEXT: [0],
+  // CHECK-NEXT: [1],
+  // CHECK-NEXT: [2],
+  // CHECK-NEXT: [3],
+  // CHECK-NEXT: [4],
+  // CHECK-NEXT: [5]
+  return
+}
+
+func @cast_unranked_memref_to_dynamic_shape(%input : memref<2x3xf32>) {
+  %unranked_input = memref_cast %input : memref<2x3xf32> to memref<*xf32>
+  %c0 = constant 0 : index
+  %c1 = constant 1 : index
+  %c6 = constant 6 : index
+  %output = memref_reinterpret_cast %unranked_input to
+           offset: [%c0], sizes: [%c1, %c6], strides: [%c6, %c1]
+           : memref<*xf32> to memref<?x?xf32, offset: ?, strides: [?, ?]>
+
+  %unranked_output = memref_cast %output
+      : memref<?x?xf32, offset: ?, strides: [?, ?]> to memref<*xf32>
+  call @print_memref_f32(%unranked_output) : (memref<*xf32>) -> ()
+  // CHECK: rank = 2 offset = 0 sizes = [1, 6] strides = [6, 1] data =
+  // CHECK-NEXT: [0,   1,   2,   3,   4,   5]
+  return
+}