[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

1126

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

1129

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

1132

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

1123

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

 #include "mlir/IR/AffineMap.h"
 #include "mlir/IR/BlockAndValueMapping.h"
 #include "llvm/ADT/SmallBitVector.h"
 
 using namespace mlir;
 
 namespace {
 
+bool isStaticStride(int64_t stride) {
+  return !ShapedType::isDynamicStrideOrOffset(stride);
+}
+
 struct AllocOpLowering : public AllocLikeOpLLVMLowering {
   AllocOpLowering(LLVMTypeConverter &converter)
       : AllocLikeOpLLVMLowering(memref::AllocOp::getOperationName(),
                                 converter) {}
 
   std::tuple<Value, Value> allocateBuffer(ConversionPatternRewriter &rewriter,
                                           Location loc, Value sizeBytes,
                                           Operation *op) const override {
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);
+
+      // Extract the offset and strides from the type.
+      int64_t offset;
+      SmallVector<int64_t, 4> strides;
+      if (failed(getStridesAndOffset(targetMemRefType, strides, offset)))
+        return rewriter.notifyMatchFailure(
+            reshapeOp, "failed to get stride and offset exprs");
+
+      if (!llvm::all_of(strides, isStaticStride))
+        return rewriter.notifyMatchFailure(reshapeOp,
+                                           "dynamic strides are unsupported");
+
+      desc.setConstantOffset(rewriter, loc, offset);
+      for (unsigned i = 0, e = targetMemRefType.getRank(); i < e; ++i) {
+        desc.setConstantSize(rewriter, loc, i, targetMemRefType.getDimSize(i));
+        desc.setConstantStride(rewriter, loc, i, strides[i]);
+      }
+
+      *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>();
matchAndRewrite(ReshapeOp reshapeOp, typename ReshapeOp::Adaptor adaptor,
     Type llvmIndexType =
         this->typeConverter->convertType(rewriter.getIndexType());
     SmallVector<Value> dstShape = getDynamicOutputShape(
         rewriter, loc, llvmIndexType, reshapeOp.getReassociationIndices(),
         srcStaticShape, srcDesc, dstStaticShape);
     for (auto &en : llvm::enumerate(dstShape))
       dstDesc.setSize(rewriter, loc, en.index(), en.value());
 
-    auto isStaticStride = [](int64_t stride) {
-      return !ShapedType::isDynamicStrideOrOffset(stride);
-    };
     if (llvm::all_of(strides, isStaticStride)) {
       for (auto &en : llvm::enumerate(strides))
         dstDesc.setConstantStride(rewriter, loc, en.index(), en.value());
     } else if (srcType.getLayout().isIdentity() &&
                dstType.getLayout().isIdentity()) {
       Value c1 = rewriter.create<LLVM::ConstantOp>(loc, llvmIndexType,
                                                    rewriter.getIndexAttr(1));
       Value stride = c1;

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]>
+
+// CHECK-LABEL: func @memref.reshape
+// CHECK-SAME:    %[[arg0:.*]]: memref<4x5x6xf32>) -> memref<2x6x20xf32>
+func.func @memref.reshape(%arg0: memref<4x5x6xf32>) -> memref<2x6x20xf32> {
+  // CHECK: %[[cast0:.*]] = builtin.unrealized_conversion_cast %arg0 : memref<4x5x6xf32> to !llvm.struct<(ptr<f32>, ptr<f32>, i64, array<3 x i64>, array<3 x i64>)>
+  %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 %[[cast0]][0] : !llvm.struct<(ptr<f32>, ptr<f32>, i64, array<3 x i64>, array<3 x i64>)>
+  // CHECK: %[[elem1:.*]] = llvm.extractvalue %[[cast0]][1] : !llvm.struct<(ptr<f32>, ptr<f32>, i64, array<3 x i64>, array<3 x i64>)>
+  // CHECK: %[[insert0:.*]] = llvm.insertvalue %[[elem0]], %[[undef]][0] : !llvm.struct<(ptr<f32>, ptr<f32>, i64, array<3 x i64>, array<3 x i64>)>
+  // CHECK: %[[insert1:.*]] = llvm.insertvalue %[[elem1]], %[[insert0:.*]][1] : !llvm.struct<(ptr<f32>, ptr<f32>, i64, array<3 x i64>, array<3 x i64>)>
+  // CHECK: %[[zero:.*]] = llvm.mlir.constant(0 : index) : i64
+  // CHECK: %[[insert2:.*]] = llvm.insertvalue %[[zero]], %[[insert1]][2] : !llvm.struct<(ptr<f32>, ptr<f32>, i64, array<3 x i64>, array<3 x i64>)>
+  // CHECK: %[[two:.*]] = llvm.mlir.constant(2 : index) : i64
+  // CHECK: %[[insert3:.*]] = llvm.insertvalue %[[two]], %[[insert2]][3, 0] : !llvm.struct<(ptr<f32>, ptr<f32>, i64, array<3 x i64>, array<3 x i64>)>
+  // CHECK: %[[hundred_and_twenty:.*]] = llvm.mlir.constant(120 : index) : i64
+  // CHECK: %[[insert4:.*]] = llvm.insertvalue %[[hundred_and_twenty]], %[[insert3]][4, 0] : !llvm.struct<(ptr<f32>, ptr<f32>, i64, array<3 x i64>, array<3 x i64>)>
+  // CHECK: %[[six:.*]] = llvm.mlir.constant(6 : index) : i64
+  // CHECK: %[[insert5:.*]] = llvm.insertvalue %[[six]], %[[insert4]][3, 1] : !llvm.struct<(ptr<f32>, ptr<f32>, i64, array<3 x i64>, array<3 x i64>)>
+  // CHECK: %[[twenty0:.*]] = llvm.mlir.constant(20 : index) : i64
+  // CHECK: %[[insert6:.*]] = llvm.insertvalue %[[twenty0]], %[[insert5]][4, 1] : !llvm.struct<(ptr<f32>, ptr<f32>, i64, array<3 x i64>, array<3 x i64>)>
+  // CHECK: %[[twenty1:.*]] = llvm.mlir.constant(20 : index) : i64
+  // CHECK: %[[insert7:.*]] = llvm.insertvalue %[[twenty1]], %[[insert6]][3, 2] : !llvm.struct<(ptr<f32>, ptr<f32>, i64, array<3 x i64>, array<3 x i64>)>
+  // CHECK: %[[one:.*]] = llvm.mlir.constant(1 : index) : i64
+  // CHECK: %[[insert8:.*]] = llvm.insertvalue %[[one]], %[[insert7]][4, 2] : !llvm.struct<(ptr<f32>, ptr<f32>, i64, array<3 x i64>, array<3 x i64>)>
+  // CHECK: %[[cast1:.*]] = builtin.unrealized_conversion_cast %[[insert8]] : !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 %[[cast1]] : memref<2x6x20xf32>
+  return %1 : memref<2x6x20xf32>
+}