[mlir] translate memref.reshape with static shapes but dynamic dims

[ashay-github]

Summary

Prior to this patch, the lowering of memref.reshape operations to the
LLVM dialect failed if the shape argument had a static shape with
dynamic dimensions. This patch adds the necessary support so that when
the shape argument has dynamic values, the lowering probes the dimension
at runtime to set the size in the MemRefDescriptor type. This patch
also computes the stride for dynamic dimensions by deriving it from the
sizes of the inner dimensions.

[ashay-github]

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

237

Since the dimensions are now iterated in reverse order, this patch only _reorders_ statements (instead of adding new statements) in the existing test.

[ftynse]

mlir/lib/Conversion/MemRefToLLVM/MemRefToLLVM.cpp

1138

If the last (i.e., the first in iteration order) stride is dynamic, strideValue will be set to nullptr, and it will be used below to create IR. This looks broken.

1146–1152

Consider using a regular if conditional instead, it will likely be shorter and more readable.

1151

Please avoid calling code that might create IR from within rewriter.create arguments. This may lead to the IR being emitted in unspecified and thus untestable order (I know that it will not in this specific case, but it is easier to enforce the rule of never doing this rather than to debug potential future changes in function calls).

[ashay-github]

Now using regular if instead of a ternary statement for better readability

[ashay-github]

mlir/lib/Conversion/MemRefToLLVM/MemRefToLLVM.cpp

1138

I thought the same, but (if I am correct) the ViewOpLowering::getStride() method at line 1911 in the same file also seems to assume that the innermost dimension's stride will not be dynamic.

If the innermost dimension's stride is indeed dynamic, is it safe to say that it will be 1?

1146–1152

Makes sense, fixed.

1151

Thanks, I hadn't considered the potential problems. I've fixed it in the most recent version.

[ftynse]

mlir/lib/Conversion/MemRefToLLVM/MemRefToLLVM.cpp

1138

I suppose you mean this https://github.com/llvm/llvm-project/blob/5a2e640eb794fffeba6afa259058115d62e7b32e/mlir/lib/Conversion/MemRefToLLVM/MemRefToLLVM.cpp#L1811 (line 1811, not 1911). It has enough checks to never attempt constructing the IR from a null Value, the code here does not. The lowering for ViewOp also leverages the semantic restriction of the op itself that requires the target memref to be contiguous, i.e. to have the innermost stride of 1. Here, there is actually a stronger restriction -- the target memref must have the identity layout, which subsumes innermost stride 1 and the following strides being products of sizes.

1142

Nit: I understand the desire to indicate default-initialization with braces, but we don't usually do this here.

1152

Nit: you shouldn't need .getResult(), single-result ops are implicitly convertible to Value.

[ashay-github]

Ensure that null stride value is never used

mlir/lib/Conversion/MemRefToLLVM/MemRefToLLVM.cpp

1138

Thanks for the explanation. I had missed the fact that reshape ops have to have identity layout maps for both source and destination memref types.

I've placed an assertion at the top of the loop, in addition to ensuring that the IR isn't constructed with a null stride values. Please let me know if I didn't get that right. Thanks!

[ftynse]

mlir/lib/Conversion/MemRefToLLVM/MemRefToLLVM.cpp

1134–1135

Nit: put the comment inside the assert - assert(condition && "message");

[ashay-github]

Place comment in the assert message

mlir/lib/Conversion/MemRefToLLVM/MemRefToLLVM.cpp

if (shapeMemRefType.hasStaticShape()) {
       SmallVector<int64_t> strides;
       if (failed(getStridesAndOffset(targetMemRefType, strides, offset)))
         return rewriter.notifyMatchFailure(
             reshapeOp, "failed to get stride and offset exprs");
 
       if (!isStaticStrideOrOffset(offset))
         return rewriter.notifyMatchFailure(reshapeOp,
                                            "dynamic offset is unsupported");
-      if (!llvm::all_of(strides, isStaticStrideOrOffset))
-        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]);
+
+      assert(targetMemRefType.getLayout().isIdentity() &&
+             "Identity layout map is a precondition of a valid reshape op");
+
+      Value stride = nullptr;
+      int64_t targetRank = targetMemRefType.getRank();
+      for (auto i : llvm::reverse(llvm::seq<int64_t>(0, targetRank))) {
+        if (!ShapedType::isDynamicStrideOrOffset(strides[i])) {
+          // If the stride for this dimension is dynamic, then use the product
+          // of the sizes of the inner dimensions.
+          stride = createIndexConstant(rewriter, loc, strides[i]);
+        } else if (!stride) {
+          // `stride` is null only in the first iteration of the loop.  However,
+          // since the target memref has an identity layout, we can safely set
+          // the innermost stride to 1.
+          stride = createIndexConstant(rewriter, loc, 1);
+        }
+
+        Value dimSize;
+        int64_t size = targetMemRefType.getDimSize(i);
+        // If the size of this dimension is dynamic, then load it at runtime
+        // from the shape operand.
+        if (!ShapedType::isDynamic(size)) {
+          dimSize = createIndexConstant(rewriter, loc, size);
+        } else {
+          Value shapeOp = reshapeOp.shape();
+          Value index = createIndexConstant(rewriter, loc, i);
+          dimSize = rewriter.create<memref::LoadOp>(loc, shapeOp, index);
+        }
+
+        desc.setSize(rewriter, loc, i, dimSize);
+        desc.setStride(rewriter, loc, i, stride);
+
+        // Prepare the stride value for the next dimension.
+        stride = rewriter.create<LLVM::MulOp>(loc, stride, dimSize);
       }
 
       *descriptor = desc;
       return success();
     }
 
     // The shape is a rank-1 tensor with unknown length.
     Location loc = reshapeOp.getLoc();

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

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: %[[one:.*]] = llvm.mlir.constant(1 : index) : 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: %[[insert3:.*]] = llvm.insertvalue %[[twenty0]], %[[insert2]][3, 2] : !llvm.struct<(ptr<f32>, ptr<f32>, i64, array<3 x i64>, array<3 x i64>)>
+  // CHECK: %[[insert4:.*]] = llvm.insertvalue %[[one]], %[[insert3]][4, 2] : !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: %[[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: %[[insert6:.*]] = llvm.insertvalue %[[twenty1]], %[[insert5]][4, 1] : !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: %[[two:.*]] = llvm.mlir.constant(2 : index) : i64
+  // CHECK: %[[insert7:.*]] = llvm.insertvalue %[[two]], %[[insert6]][3, 0] : !llvm.struct<(ptr<f32>, ptr<f32>, i64, array<3 x i64>, array<3 x i64>)>
+  // CHECK: %[[insert8:.*]] = llvm.insertvalue %[[hundred_and_twenty]], %[[insert7]][4, 0] : !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>
 }
+
+// -----
+
+// CHECK-LABEL: func @memref.reshape.dynamic.dim
+// CHECK-SAME:    %[[arg:.*]]: memref<?x?x?xf32>, %[[shape:.*]]: memref<4xi64>) -> memref<?x?x12x32xf32>
+func.func @memref.reshape.dynamic.dim(%arg: memref<?x?x?xf32>, %shape: memref<4xi64>) -> memref<?x?x12x32xf32> {
+  // CHECK: %[[arg_cast:.*]] = builtin.unrealized_conversion_cast %[[arg]] : memref<?x?x?xf32> to !llvm.struct<(ptr<f32>, ptr<f32>, i64, array<3 x i64>, array<3 x i64>)>
+  // CHECK: %[[shape_cast:.*]] = builtin.unrealized_conversion_cast %[[shape]] : memref<4xi64> to !llvm.struct<(ptr<i64>, ptr<i64>, i64, array<1 x i64>, array<1 x i64>)>
+  // CHECK: %[[undef:.*]] = llvm.mlir.undef : !llvm.struct<(ptr<f32>, ptr<f32>, i64, array<4 x i64>, array<4 x i64>)>
+  // CHECK: %[[alloc_ptr:.*]] = llvm.extractvalue %[[arg_cast]][0] : !llvm.struct<(ptr<f32>, ptr<f32>, i64, array<3 x i64>, array<3 x i64>)>
+  // CHECK: %[[align_ptr:.*]] = llvm.extractvalue %[[arg_cast]][1] : !llvm.struct<(ptr<f32>, ptr<f32>, i64, array<3 x i64>, array<3 x i64>)>
+  // CHECK: %[[insert0:.*]] = llvm.insertvalue %[[alloc_ptr]], %[[undef]][0] : !llvm.struct<(ptr<f32>, ptr<f32>, i64, array<4 x i64>, array<4 x i64>)>
+  // CHECK: %[[insert1:.*]] = llvm.insertvalue %[[align_ptr]], %[[insert0]][1] : !llvm.struct<(ptr<f32>, ptr<f32>, i64, array<4 x i64>, array<4 x i64>)>
+
+  // CHECK: %[[zero0:.*]] = llvm.mlir.constant(0 : index) : i64
+  // CHECK: %[[insert2:.*]] = llvm.insertvalue %[[zero0]], %[[insert1]][2] : !llvm.struct<(ptr<f32>, ptr<f32>, i64, array<4 x i64>, array<4 x i64>)>
+
+  // CHECK: %[[one0:.*]] = llvm.mlir.constant(1 : index) : i64
+  // CHECK: %[[thirty_two0:.*]] = llvm.mlir.constant(32 : index) : i64
+  // CHECK: %[[insert3:.*]] = llvm.insertvalue %[[thirty_two0]], %[[insert2]][3, 3] : !llvm.struct<(ptr<f32>, ptr<f32>, i64, array<4 x i64>, array<4 x i64>)>
+  // CHECK: %[[insert4:.*]] = llvm.insertvalue %[[one0]], %[[insert3]][4, 3] : !llvm.struct<(ptr<f32>, ptr<f32>, i64, array<4 x i64>, array<4 x i64>)>
+
+  // CHECK: %[[thirty_two1:.*]] = llvm.mlir.constant(32 : index) : i64
+  // CHECK: %[[twelve:.*]] = llvm.mlir.constant(12 : index) : i64
+  // CHECK: %[[insert5:.*]] = llvm.insertvalue %[[twelve]], %[[insert4]][3, 2] : !llvm.struct<(ptr<f32>, ptr<f32>, i64, array<4 x i64>, array<4 x i64>)>
+  // CHECK: %[[insert6:.*]] = llvm.insertvalue %[[thirty_two1]], %[[insert5]][4, 2] : !llvm.struct<(ptr<f32>, ptr<f32>, i64, array<4 x i64>, array<4 x i64>)>
+
+  // CHECK: %[[three_hundred_and_eighty_four:.*]] = llvm.mlir.constant(384 : index) : i64
+  // CHECK: %[[one1:.*]] = llvm.mlir.constant(1 : index) : i64
+  // CHECK: %[[shape_ptr0:.*]] = llvm.extractvalue %[[shape_cast]][1] : !llvm.struct<(ptr<i64>, ptr<i64>, i64, array<1 x i64>, array<1 x i64>)>
+  // CHECK: %[[shape_gep0:.*]] = llvm.getelementptr %[[shape_ptr0]][%[[one1]]] : (!llvm.ptr<i64>, i64) -> !llvm.ptr<i64>
+  // CHECK: %[[shape_load0:.*]] = llvm.load %[[shape_gep0]] : !llvm.ptr<i64>
+  // CHECK: %[[insert7:.*]] = llvm.insertvalue %[[shape_load0]], %[[insert6]][3, 1] : !llvm.struct<(ptr<f32>, ptr<f32>, i64, array<4 x i64>, array<4 x i64>)>
+  // CHECK: %[[insert8:.*]] = llvm.insertvalue %[[three_hundred_and_eighty_four]], %[[insert7]][4, 1] : !llvm.struct<(ptr<f32>, ptr<f32>, i64, array<4 x i64>, array<4 x i64>)>
+
+  // CHECK: %[[mul:.*]] = llvm.mul %19, %23  : i64
+  // CHECK: %[[zero1:.*]] = llvm.mlir.constant(0 : index) : i64
+  // CHECK: %[[shape_ptr1:.*]] = llvm.extractvalue %[[shape_cast]][1] : !llvm.struct<(ptr<i64>, ptr<i64>, i64, array<1 x i64>, array<1 x i64>)>
+  // CHECK: %[[shape_gep1:.*]] = llvm.getelementptr %[[shape_ptr1]][%[[zero1]]] : (!llvm.ptr<i64>, i64) -> !llvm.ptr<i64>
+  // CHECK: %[[shape_load1:.*]] = llvm.load %[[shape_gep1]] : !llvm.ptr<i64>
+  // CHECK: %[[insert9:.*]] = llvm.insertvalue %[[shape_load1]], %[[insert8]][3, 0] : !llvm.struct<(ptr<f32>, ptr<f32>, i64, array<4 x i64>, array<4 x i64>)>
+  // CHECK: %[[insert10:.*]] = llvm.insertvalue %[[mul]], %[[insert9]][4, 0] : !llvm.struct<(ptr<f32>, ptr<f32>, i64, array<4 x i64>, array<4 x i64>)>
+
+  // CHECK: %[[result_cast:.*]] = builtin.unrealized_conversion_cast %[[insert10]] : !llvm.struct<(ptr<f32>, ptr<f32>, i64, array<4 x i64>, array<4 x i64>)> to memref<?x?x12x32xf32>
+  %0 = memref.reshape %arg(%shape) : (memref<?x?x?xf32>, memref<4xi64>) -> memref<?x?x12x32xf32>
+
+  return %0 : memref<?x?x12x32xf32>
+  // CHECK: return %[[result_cast]] : memref<?x?x12x32xf32>
+}