[mlir] translate memref.reshape ops that have static shapes

[ashay-github]

Summary

This patch references code for translating memref.reinterpret_cast ops
to add translation rules for memref.reshape ops that have a static shape
argument. Since reshape ops don't have offsets, sizes, or strides, this
patch simply sets the allocated and aligned pointers of the MemRef
descriptor.

[ftynse]

Since reshape ops don't have offsets, sizes, or strides, this
patch simply sets the allocated and aligned pointers of the MemRef
descriptor.

The convention is that the descriptor always contains the correct size and stride information even if the lowerings may chose to ignore it and use the information statically available from types. The reason for this is the descriptor being also used to interface with C code that does not have access to the information containing in the memref type. Leaving size/stride information uninitialized will lead to UB in such code.

[ashay-github]

set stride, offset, and size attributes of MemRefDescriptor

The code now uses getStridesAndOffset() on the result MemRefType to infer the
stride, offset, and size attributes.

[ashay-github]

Thanks for the suggestion. I've updated the code and test to set the offset, stride, and size. PTAL.

[cathyzhyi]

I left some comments but I am also pretty new to this. Hope someone who's expert on this help to take a look.

mlir/lib/Conversion/MemRefToLLVM/MemRefToLLVM.cpp

1122

nit: could remove the auto successStrides and check if (failed(getStridesAndOffset...)) directly

1125

My understanding is that it's still possible to get non static offset and strides for static shape memref so need to check all static before calling desc.setConstantXXX.

[ashay-github]

check for static strides and offsets, otherwise return failure

[ashay-github]

Thanks @cathyzhyi! I've updated the code accordingly, but please let me know if I still got it wrong. Thanks!

[cathyzhyi]

mlir/lib/Conversion/MemRefToLLVM/MemRefToLLVM.cpp

1128

Need to check for constant offset as well.

[ashay-github]

check for dynamic offset

[ashay-github]

Thanks @cathyzhyi! I added a check for !isStaticStrideOrOffset(offset) as well.

[cathyzhyi]

LGTM! @ftynse could you help to double check because I am not super confident either.

[ftynse]

mlir/lib/Conversion/MemRefToLLVM/MemRefToLLVM.cpp

1119

Nit: do not specify the number of stack elements in SmallVector unless you have a strong reason to use a specific number.

[ashay-github]

dropping element count in SmallVector usage

mlir/lib/Conversion/MemRefToLLVM/MemRefToLLVM.cpp

struct MemRefReshapeOpLowering
   }
 
 private:
   LogicalResult
   convertSourceMemRefToDescriptor(ConversionPatternRewriter &rewriter,
                                   Type srcType, memref::ReshapeOp reshapeOp,
                                   memref::ReshapeOp::Adaptor adaptor,
                                   Value *descriptor) const {
-    // Conversion for statically-known shape args is performed via
-    // `memref_reinterpret_cast`.
     auto shapeMemRefType = reshapeOp.shape().getType().cast<MemRefType>();
-    if (shapeMemRefType.hasStaticShape())
+    if (shapeMemRefType.hasStaticShape()) {
+      MemRefType targetMemRefType =
+          reshapeOp.getResult().getType().cast<MemRefType>();
+      auto llvmTargetDescriptorTy =
+          typeConverter->convertType(targetMemRefType)
+              .dyn_cast_or_null<LLVM::LLVMStructType>();
+      if (!llvmTargetDescriptorTy)
         return failure();
 
+      // Create descriptor.
+      Location loc = reshapeOp.getLoc();
+      auto desc =
+          MemRefDescriptor::undef(rewriter, loc, llvmTargetDescriptorTy);
+
+      // Set allocated and aligned pointers.
+      Value allocatedPtr, alignedPtr;
+      extractPointersAndOffset(loc, rewriter, *getTypeConverter(),
+                               reshapeOp.source(), adaptor.source(),
+                               &allocatedPtr, &alignedPtr);
+      desc.setAllocatedPtr(rewriter, loc, allocatedPtr);
+      desc.setAlignedPtr(rewriter, loc, alignedPtr);
+
+      *descriptor = desc;
+      return success();
+    }
+
     // The shape is a rank-1 tensor with unknown length.
     Location loc = reshapeOp.getLoc();
     MemRefDescriptor shapeDesc(adaptor.shape());
     Value resultRank = shapeDesc.size(rewriter, loc, 0);
 
     // Extract address space and element type.
     auto targetType =
         reshapeOp.getResult().getType().cast<UnrankedMemRefType>();
     unsigned addressSpace = targetType.getMemorySpaceAsInt();
     Type elementType = targetType.getElementType();
 
     // Create the unranked memref descriptor that holds the ranked one. The
     // inner descriptor is allocated on stack.
     auto targetDesc = UnrankedMemRefDescriptor::undef(
         rewriter, loc, typeConverter->convertType(targetType));
     targetDesc.setRank(rewriter, loc, resultRank);

mlir/test/Conversion/MemRefToLLVM/convert-static-memref-ops.mlir

func.func @address() {
     // CHECK: llvm.add {{.*}} : i32
     // CHECK: llvm.urem {{.*}} : i32
     // CHECK: llvm.sub {{.*}} : i32
     // CHECK: llvm.inttoptr %{{.*}} : i32 to !llvm.ptr
     return
   }
 }
 
+// -----
+
+memref.global "private" constant @__constant_3xi64 : memref<3xi64> = dense<[2, 6, 20]>
+func @forward(%arg0: memref<4x5x6xf32>) -> memref<2x6x20xf32> {
+  %0 = memref.get_global @__constant_3xi64 : memref<3xi64>
+
+  // CHECK: %[[undef:.*]] = llvm.mlir.undef : !llvm.struct<(ptr<f32>, ptr<f32>, i64, array<3 x i64>, array<3 x i64>)>
+  // CHECK: %[[elem0:.*]] = llvm.extractvalue %0[0] : !llvm.struct<(ptr<f32>, ptr<f32>, i64, array<3 x i64>, array<3 x i64>)>
+  // CHECK: %[[elem1:.*]] = llvm.extractvalue %0[1] : !llvm.struct<(ptr<f32>, ptr<f32>, i64, array<3 x i64>, array<3 x i64>)>
+  // CHECK: %[[first:.*]] = llvm.insertvalue %[[elem0]], %[[undef]][0] : !llvm.struct<(ptr<f32>, ptr<f32>, i64, array<3 x i64>, array<3 x i64>)>
+  // CHECK: %[[second:.*]] = llvm.insertvalue %[[elem1]], %[[first:.*]][1] : !llvm.struct<(ptr<f32>, ptr<f32>, i64, array<3 x i64>, array<3 x i64>)>
+  // CHECK: %[[cast:.*]] = builtin.unrealized_conversion_cast %[[second]] : !llvm.struct<(ptr<f32>, ptr<f32>, i64, array<3 x i64>, array<3 x i64>)> to memref<2x6x20xf32>
+  %1 = memref.reshape %arg0(%0) : (memref<4x5x6xf32>, memref<3xi64>) -> memref<2x6x20xf32>
+
+  // CHECK: return %[[cast]] : memref<2x6x20xf32>
+  return %1 : memref<2x6x20xf32>
+}