[MLIR][Standard] Lower `std.dim` with dynamic dimension operand to LLVM

[frgossen]

Summary

Implement the missing lowering from std.dim to the LLVM dialect in case of a
dynamic dimension.

[pifon2a]

mlir/test/Conversion/StandardToLLVM/convert-dynamic-memref-ops.mlir

416

please, capture !llvm<"{ float*, float*, i64, [3 x i64], [3 x i64] }"> as [[DESCR_TY]] and use everywhere below.

[herhut]

mlir/test/Conversion/StandardToLLVM/convert-dynamic-memref-ops.mlir

427

Does this work when translated to LLVM proper? (You can check with mlir-translate). I think the argument to a GEP needs to be a pointer but in this case it is a value.

[frgossen]

Apply suggestion

[frgossen]

Thanks

[ftynse]

Have you checked that bare-pointer calling convention also works?

mlir/lib/Conversion/StandardToLLVM/StandardToLLVM.cpp

574

Value is a descriptor (LLVM struct type), you need to extractvalue from it before doing GEP.

mlir/test/Conversion/StandardToLLVM/convert-dynamic-memref-ops.mlir

411

Nit: I think "dynamic-memref-ops" was intended for ops working on dynamic memrefs, i.e. those with dynamic dimensions. This is not the case here.

416

Cool trick!

[frgossen]

Fix getelementptr issue.

mlir/lib/Conversion/StandardToLLVM/StandardToLLVM.cpp

574

I see. Unfortunately not even the sizes in the MemRefDescriptor are a pointer here :/

mlir/test/Conversion/StandardToLLVM/convert-dynamic-memref-ops.mlir

427

The new version does. It copies everything to stack-allocated memory. Is this what you meant earlier today? I checked this manually with mlir-translate. Should that be reflected in this test, i.e. should I add the corresponding //RUN: ... statement?

[frgossen]

In D81834#2092607, @ftynse wrote:

Have you checked that bare-pointer calling convention also works?

How can I do that?

[nicolasvasilache]

mlir/test/Conversion/StandardToLLVM/convert-dynamic-memref-ops.mlir

416

Facepalm .. Thanks much @pifon2a MUCH better!

[herhut]

In the long run we need a better way to do this so that we only materialize the descriptor once if needed. Until then, this will work (with some nits).

We can land this without restoring the stack for now.

mlir/lib/Conversion/StandardToLLVM/StandardToLLVM.cpp

580

Does LLVM allow to store the entire array at once? I think store supports writing array-typed values.

One concern with this implementation is that it could fill up the stack quickly. Assume this gets called in a loop, then each iteration would put a new copy on the stack. We could use @llvm.stackrestore to limit the lifetime of the alloca'd memory. Not sure what that does to LLVMs optimizations, though.

@ftynse Should stackrestore be an operation in the llvm dialect? Seems handy to have it there.

[ftynse]

In D81834#2092607, @ftynse wrote:
Have you checked that bare-pointer calling convention also works?

How can I do that?

mlir-opt -convert-std-to-llvmir="use-bare-ptr-memref-call-conv=1". There are also some tests for this flag.

Approving now, we can revise to the stacksave/stackrestore version in a separate commit. (It can be handy for dynamically-ranked memrefs as well).

mlir/lib/Conversion/StandardToLLVM/StandardToLLVM.cpp

580

At some point, we were storing the entire descriptor on stack, and it did lead to stack explosion when called in a loop. I spent quite some time cleaning that up and now I am wary of alloca'ed memory for memrefs. I don't see another option here though.

The common trick I've seen is to put stack allocations in the entry block of the function. One cannot branch to that block so there will be no loops. And since this memory is only used in the immediately succeeding instructions and then discarded, we don't need any lifetime analysis to do the trick. This does not play well with pattern-based rewrites though...

@ftynse Should stackrestore be an operation in the llvm dialect? Seems handy to have it there.

Sure, it's an intrinsic, so it should be really easy to add.

[frgossen]

Address comments

[frgossen]

Address comments.

[frgossen]

Remove test with bare-pointer calling convention

[frgossen]

@ftynse, the new test case fails with bare-pointer calling convention because its argument is a dynamically shaped memref.
All the other tests with a dynamically shaped memref fail in the same way.
I tried the new test case with a statically shaped memref and it all works just fine.

Should I move the test case to another file, make the argument statically shaped, and test with bare-pointer calling convention?

mlir/lib/Conversion/StandardToLLVM/StandardToLLVM.cpp

580

@herhut, can you check if this is what you had in mind?

[ftynse]

Should I move the test case to another file, make the argument statically shaped, and test with bare-pointer calling convention?

No, just asking to double check. Thanks!

mlir/include/mlir/Conversion/StandardToLLVM/ConvertStandardToLLVM.h

public:
   Value offset(OpBuilder &builder, Location loc);
 
   /// Builds IR inserting the offset into the descriptor.
   void setOffset(OpBuilder &builder, Location loc, Value offset);
   void setConstantOffset(OpBuilder &builder, Location loc, uint64_t offset);
 
   /// Builds IR extracting the pos-th size from the descriptor.
   Value size(OpBuilder &builder, Location loc, unsigned pos);
+  Value size(OpBuilder &builder, Location loc, Value pos, int64_t rank);
 
   /// Builds IR inserting the pos-th size into the descriptor
   void setSize(OpBuilder &builder, Location loc, unsigned pos, Value size);
   void setConstantSize(OpBuilder &builder, Location loc, unsigned pos,
                        uint64_t size);
 
   /// Builds IR extracting the pos-th size from the descriptor.
   Value stride(OpBuilder &builder, Location loc, unsigned pos);

mlir/lib/Conversion/StandardToLLVM/StandardToLLVM.cpp

 
 /// Builds IR extracting the pos-th size from the descriptor.
 Value MemRefDescriptor::size(OpBuilder &builder, Location loc, unsigned pos) {
   return builder.create<LLVM::ExtractValueOp>(
       loc, indexType, value,
       builder.getI64ArrayAttr({kSizePosInMemRefDescriptor, pos}));
 }
 
+Value MemRefDescriptor::size(OpBuilder &builder, Location loc, Value pos,
+                             int64_t rank) {
+  auto indexTy = indexType.cast<LLVM::LLVMType>();
+  auto indexPtrTy = indexTy.getPointerTo();
+  auto arrayTy = LLVM::LLVMType::getArrayTy(indexTy, rank);
+  auto arrayPtrTy = arrayTy.getPointerTo();
+
+  // Copy size values to stack-allocated memory.
+  auto zero = createIndexAttrConstant(builder, loc, indexType, 0);
+  auto one = createIndexAttrConstant(builder, loc, indexType, 1);
+  auto sizes = builder.create<LLVM::ExtractValueOp>(
+      loc, arrayTy, value,
+      builder.getI64ArrayAttr({kSizePosInMemRefDescriptor}));
+  auto sizesPtr =
+      builder.create<LLVM::AllocaOp>(loc, arrayPtrTy, one, /*alignment=*/0);
+  builder.create<LLVM::StoreOp>(loc, sizes, sizesPtr);
+
+  // Load an return size value of interest.
+  auto resultPtr = builder.create<LLVM::GEPOp>(loc, indexPtrTy, sizesPtr,
+                                               ValueRange({zero, pos}));
+  return builder.create<LLVM::LoadOp>(loc, resultPtr);
+}
+
 /// Builds IR inserting the pos-th size into the descriptor
 void MemRefDescriptor::setSize(OpBuilder &builder, Location loc, unsigned pos,
                                Value size) {
   value = builder.create<LLVM::InsertValueOp>(
       loc, structType, value, size,
       builder.getI64ArrayAttr({kSizePosInMemRefDescriptor, pos}));
 }
 
-/// Builds IR inserting the pos-th size into the descriptor
 void MemRefDescriptor::setConstantSize(OpBuilder &builder, Location loc,
                                        unsigned pos, uint64_t size) {
   setSize(builder, loc, pos,
           createIndexAttrConstant(builder, loc, indexType, size));
 }
 
 /// Builds IR extracting the pos-th stride from the descriptor.
 Value MemRefDescriptor::stride(OpBuilder &builder, Location loc, unsigned pos) {
   return builder.create<LLVM::ExtractValueOp>(
       loc, indexType, value,
       builder.getI64ArrayAttr({kStridePosInMemRefDescriptor, pos}));
 }
 
 /// Builds IR inserting the pos-th stride into the descriptor
 void MemRefDescriptor::setStride(OpBuilder &builder, Location loc, unsigned pos,
                                  Value stride) {
   value = builder.create<LLVM::InsertValueOp>(
       loc, structType, value, stride,
       builder.getI64ArrayAttr({kStridePosInMemRefDescriptor, pos}));
 }
 
-/// Builds IR inserting the pos-th stride into the descriptor
 void MemRefDescriptor::setConstantStride(OpBuilder &builder, Location loc,
                                          unsigned pos, uint64_t stride) {
   setStride(builder, loc, pos,
             createIndexAttrConstant(builder, loc, indexType, stride));
 }
 
 LLVM::LLVMType MemRefDescriptor::getElementType() {
   return value.getType().cast<LLVM::LLVMType>().getStructElementType(

 // size stored in the memref descriptor for dynamic sizes.
 struct DimOpLowering : public ConvertOpToLLVMPattern<DimOp> {
   using ConvertOpToLLVMPattern<DimOp>::ConvertOpToLLVMPattern;
 
   LogicalResult
   matchAndRewrite(Operation *op, ArrayRef<Value> operands,
                   ConversionPatternRewriter &rewriter) const override {
     auto dimOp = cast<DimOp>(op);
+    auto loc = op->getLoc();
     DimOp::Adaptor transformed(operands);
-    MemRefType type = dimOp.memrefOrTensor().getType().cast<MemRefType>();
-
-    Optional<int64_t> index = dimOp.getConstantIndex();
-    if (!index.hasValue()) {
-      // TODO: Implement this lowering.
-      return failure();
-    }
 
+    // Take advantage if index is constant.
+    MemRefType memRefType = dimOp.memrefOrTensor().getType().cast<MemRefType>();
+    if (Optional<int64_t> index = dimOp.getConstantIndex()) {
       int64_t i = index.getValue();
+      if (memRefType.isDynamicDim(i)) {
         // Extract dynamic size from the memref descriptor.
-    if (type.isDynamicDim(i))
-      rewriter.replaceOp(op, {MemRefDescriptor(transformed.memrefOrTensor())
-                                  .size(rewriter, op->getLoc(), i)});
-    else
+        MemRefDescriptor descriptor(transformed.memrefOrTensor());
+        rewriter.replaceOp(op, {descriptor.size(rewriter, loc, i)});
+      } else {
         // Use constant for static size.
-      rewriter.replaceOp(
-          op, createIndexConstant(rewriter, op->getLoc(), type.getDimSize(i)));
+        int64_t dimSize = memRefType.getDimSize(i);
+        rewriter.replaceOp(op, createIndexConstant(rewriter, loc, dimSize));
+      }
+      return success();
+    }
 
+    Value index = dimOp.index();
+    int64_t rank = memRefType.getRank();
+    MemRefDescriptor memrefDescriptor(transformed.memrefOrTensor());
+    rewriter.replaceOp(op, {memrefDescriptor.size(rewriter, loc, index, rank)});
     return success();
   }
 };
 
 // Common base for load and store operations on MemRefs.  Restricts the match
 // to supported MemRef types.  Provides functionality to emit code accessing a
 // specific element of the underlying data buffer.
 template <typename Derived>

mlir/lib/Dialect/StandardOps/IR/Ops.cpp

static LogicalResult verify(DimOp op) {
   } else {
     llvm_unreachable("expected operand with tensor or memref type");
   }
 
   return success();
 }
 
 OpFoldResult DimOp::fold(ArrayRef<Attribute> operands) {
-  auto index = operands[1].dyn_cast<IntegerAttr>();
+  auto index = operands[1].dyn_cast_or_null<IntegerAttr>();
 
   // All forms of folding require a known index.
   if (!index)
     return {};
 
   // Fold if the shape extent along the given index is known.
   auto argTy = memrefOrTensor().getType();
   if (auto shapedTy = argTy.dyn_cast<ShapedType>()) {

mlir/test/Conversion/StandardToLLVM/convert-dynamic-memref-ops.mlir

 // CHECK: llvm.mlir.constant(13 : index) : !llvm.i64
   %c3 = constant 3 : index
   %3 = dim %mixed, %c3 : memref<42x?x?x13x?xf32>
 // CHECK: llvm.extractvalue %[[ld]][3, 4] : !llvm<"{ float*, float*, i64, [5 x i64], [5 x i64] }">
   %c4 = constant 4 : index
   %4 = dim %mixed, %c4 : memref<42x?x?x13x?xf32>
   return
 }
+
+// CHECK-LABEL: @memref_dim_with_dyn_index
+// CHECK-SAME: %[[ALLOC_PTR:.*]]: !llvm<"float*">, %[[ALIGN_PTR:.*]]: !llvm<"float*">, %[[OFFSET:.*]]: !llvm.i64, %[[SIZE0:.*]]: !llvm.i64, %[[SIZE1:.*]]: !llvm.i64, %[[STRIDE0:.*]]: !llvm.i64, %[[STRIDE1:.*]]: !llvm.i64, %[[IDX:.*]]: !llvm.i64) -> !llvm.i64
+func @memref_dim_with_dyn_index(%arg : memref<3x?xf32>, %idx : index) -> index {
+  // CHECK-NEXT: %[[DESCR0:.*]] = llvm.mlir.undef : [[DESCR_TY:!llvm<"{ float\*, float\*, i64, \[2 x i64\], \[2 x i64\] }">]]
+  // CHECK-NEXT: %[[DESCR1:.*]] = llvm.insertvalue %[[ALLOC_PTR]], %[[DESCR0]][0] : [[DESCR_TY]]
+  // CHECK-NEXT: %[[DESCR2:.*]] = llvm.insertvalue %[[ALIGN_PTR]], %[[DESCR1]][1] : [[DESCR_TY]]
+  // CHECK-NEXT: %[[DESCR3:.*]] = llvm.insertvalue %[[OFFSET]],    %[[DESCR2]][2] : [[DESCR_TY]]
+  // CHECK-NEXT: %[[DESCR4:.*]] = llvm.insertvalue %[[SIZE0]],     %[[DESCR3]][3, 0] : [[DESCR_TY]]
+  // CHECK-NEXT: %[[DESCR5:.*]] = llvm.insertvalue %[[STRIDE0]],   %[[DESCR4]][4, 0] : [[DESCR_TY]]
+  // CHECK-NEXT: %[[DESCR6:.*]] = llvm.insertvalue %[[SIZE1]],     %[[DESCR5]][3, 1] : [[DESCR_TY]]
+  // CHECK-NEXT: %[[DESCR7:.*]] = llvm.insertvalue %[[STRIDE1]],   %[[DESCR6]][4, 1] : [[DESCR_TY]]
+  // CHECK-DAG: %[[C0:.*]] = llvm.mlir.constant(0 : index) : !llvm.i64
+  // CHECK-DAG: %[[C1:.*]] = llvm.mlir.constant(1 : index) : !llvm.i64
+  // CHECK-DAG: %[[SIZES:.*]] = llvm.extractvalue %[[DESCR7]][3] : [[DESCR_TY]]
+  // CHECK-DAG: %[[SIZES_PTR:.*]] = llvm.alloca %[[C1]] x !llvm<"[2 x i64]"> : (!llvm.i64) -> !llvm<"[2 x i64]*">
+  // CHECK-DAG: llvm.store %[[SIZES]], %[[SIZES_PTR]] : !llvm<"[2 x i64]*">
+  // CHECK-DAG: %[[RESULT_PTR:.*]] = llvm.getelementptr %[[SIZES_PTR]][%[[C0]], %[[IDX]]] : (!llvm<"[2 x i64]*">, !llvm.i64, !llvm.i64) -> !llvm<"i64*">
+  // CHECK-DAG: %[[RESULT:.*]] = llvm.load %[[RESULT_PTR]] : !llvm<"i64*">
+  // CHECK-DAG: llvm.return %[[RESULT]] : !llvm.i64
+  %result = dim %arg, %idx : memref<3x?xf32>
+  return %result : index
+}