[mlir] Lower DimOp to LLVM for unranked memrefs.

[pifon2a]

lint

[herhut]

mlir/lib/Conversion/StandardToLLVM/StandardToLLVM.cpp

2400

Should this be a helper in UnrankedMemRefDescriptor? This will likely get used in other places (we already to it in some hlo cast operation I think). Something like getOffsetPtr getSizesPtr and getStridesPtr which return the pointer to index?

2439

Just fall through to the implementation below? If it is dynamic, i being constant does not help.

2451

Why does this need rank explicit and the above does not?

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

1345–1346

Nit: You could cast to ShapedType and then query hasRank. We have a cast to shaped type below already.

1346

Remove extra {}

[pifon2a]

Address the comments and fix mlir-cpu-runner tests in cmake.

mlir/lib/Conversion/StandardToLLVM/StandardToLLVM.cpp

2400

Definitely, I was planning to do that when lowering memrefcast ops. It is also not trivial, because this code uses TypeConverter and addressSpace that are not a part of UnrankedMemRefDescriptor. I don't think it should be a member of the descriptor. But it can be a useful utility. Not in this CL.

2439

i left the old code pretty much untouched. Thanks for noticing. These two different size methods are emit different code. See the comment below.

2451

the one with the rank emits GEPs and they use rank to construct array type with shapes (it is not really needed though).
the one without the rank emits ExtractElementOp .

[herhut]

mlir/lib/Conversion/StandardToLLVM/StandardToLLVM.cpp

2439

So you could do if (!memRefType.isDynamicDim(i)) and otherwise fall through to the below case, removing the extra rank. Otherwise we have the same code essentially twice here.

2451

So can we just drop it here?

[herhut]

Thanks for clarifying!

mlir/lib/Conversion/StandardToLLVM/StandardToLLVM.cpp

2439

So, discussed this offline. The difference here is that using this form of size creates better code and hence is used here. Having a static index enables this.

Maybe it would be nice if size could do this internally but that is not for this PR. @ftynse WDYT

mlir/lib/Conversion/StandardToLLVM/StandardToLLVM.cpp

 // 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();
+    Type operandType = dimOp.memrefOrTensor().getType();
+    if (operandType.isa<UnrankedMemRefType>()) {
+      rewriter.replaceOp(op, {extractSizeOfUnrankedMemRef(operandType, dimOp,
+                                                          operands, rewriter)});
+
+      return success();
+    }
+    if (operandType.isa<MemRefType>()) {
+      rewriter.replaceOp(op, {extractSizeOfRankedMemRef(operandType, dimOp,
+                                                        operands, rewriter)});
+      return success();
+    }
+    return failure();
+  }
+
+private:
+  Value extractSizeOfUnrankedMemRef(Type operandType, DimOp dimOp,
+                                    ArrayRef<Value> operands,
+                                    ConversionPatternRewriter &rewriter) const {
+    Location loc = dimOp.getLoc();
     DimOp::Adaptor transformed(operands);
 
+    auto unrankedMemRefType = operandType.cast<UnrankedMemRefType>();
+    auto scalarMemRefType =
+        MemRefType::get({}, unrankedMemRefType.getElementType());
+    unsigned addressSpace = unrankedMemRefType.getMemorySpace();
+
+    // Extract pointer to the underlying ranked descriptor and bitcast it to a
+    // memref<element_type> descriptor pointer to minimize the number of GEP
+    // operations.
+    UnrankedMemRefDescriptor unrankedDesc(transformed.memrefOrTensor());
+    Value underlyingRankedDesc = unrankedDesc.memRefDescPtr(rewriter, loc);
+    Value scalarMemRefDescPtr = rewriter.create<LLVM::BitcastOp>(
+        loc,
+        typeConverter.convertType(scalarMemRefType)
+            .cast<LLVM::LLVMType>()
+            .getPointerTo(addressSpace),
+        underlyingRankedDesc);
+
+    // Get pointer to offset field of memref<element_type> descriptor.
+    Type indexPtrTy = typeConverter.getIndexType().getPointerTo(addressSpace);
+    Value two = rewriter.create<LLVM::ConstantOp>(
+        loc, typeConverter.convertType(rewriter.getI32Type()),
+        rewriter.getI32IntegerAttr(2));
+    Value offsetPtr = rewriter.create<LLVM::GEPOp>(
+        loc, indexPtrTy, scalarMemRefDescPtr,
+        ValueRange({createIndexConstant(rewriter, loc, 0), two}));
+
+    // The size value that we have to extract can be obtained using GEPop with
+    // `dimOp.index() + 1` index argument.
+    Value idxPlusOne = rewriter.create<LLVM::AddOp>(
+        loc, createIndexConstant(rewriter, loc, 1), transformed.index());
+    Value sizePtr = rewriter.create<LLVM::GEPOp>(loc, indexPtrTy, offsetPtr,
+                                                 ValueRange({idxPlusOne}));
+    return rewriter.create<LLVM::LoadOp>(loc, sizePtr);
+  }
+
+  Value extractSizeOfRankedMemRef(Type operandType, DimOp dimOp,
+                                  ArrayRef<Value> operands,
+                                  ConversionPatternRewriter &rewriter) const {
+    Location loc = dimOp.getLoc();
+    DimOp::Adaptor transformed(operands);
     // Take advantage if index is constant.
-    MemRefType memRefType = dimOp.memrefOrTensor().getType().cast<MemRefType>();
+    MemRefType memRefType = operandType.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.
+        // extract dynamic size from the memref descriptor.
         MemRefDescriptor descriptor(transformed.memrefOrTensor());
-        rewriter.replaceOp(op, {descriptor.size(rewriter, loc, i)});
-      } else {
-        // Use constant for static size.
-        int64_t dimSize = memRefType.getDimSize(i);
-        rewriter.replaceOp(op, createIndexConstant(rewriter, loc, dimSize));
-      }
-      return success();
+        return descriptor.size(rewriter, loc, i);
+      }
+      return createIndexConstant(rewriter, loc, dimSize);
     }
-
     Value index = dimOp.index();
     int64_t rank = memRefType.getRank();
     MemRefDescriptor memrefDescriptor(transformed.memrefOrTensor());
-    rewriter.replaceOp(op, {memrefDescriptor.size(rewriter, loc, index, rank)});
-    return success();
+    return memrefDescriptor.size(rewriter, loc, index, rank);
   }
 };
 
 struct RankOpLowering : public ConvertOpToLLVMPattern<RankOp> {
   using ConvertOpToLLVMPattern<RankOp>::ConvertOpToLLVMPattern;
 
   LogicalResult
   matchAndRewrite(Operation *op, ArrayRef<Value> operands,

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

static LogicalResult verify(DimOp op) {
   // Check that constant index is not knowingly out of range.
   auto type = op.memrefOrTensor().getType();
   if (auto tensorType = type.dyn_cast<RankedTensorType>()) {
     if (index.getValue() >= tensorType.getRank())
       return op.emitOpError("index is out of range");
   } else if (auto memrefType = type.dyn_cast<MemRefType>()) {
     if (index.getValue() >= memrefType.getRank())
       return op.emitOpError("index is out of range");
-  } else if (type.isa<UnrankedTensorType>()) {
+  } else if (type.isa<UnrankedTensorType>() || type.isa<UnrankedMemRefType>()) {
     // Assume index to be in range.
   } else {
     llvm_unreachable("expected operand with tensor or memref type");
   }
 
   return success();
 }
 
 OpFoldResult DimOp::fold(ArrayRef<Attribute> operands) {
   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();
+  // Fold if the shape extent along the given index is known.
   if (auto shapedTy = argTy.dyn_cast<ShapedType>()) {
+    // Folding for unranked types (UnrankedMemRefType, UnrankedTensorType) is
+    // not supported.
+    if (!shapedTy.hasRank())
+      return {};
     if (!shapedTy.isDynamicDim(index.getInt())) {
       Builder builder(getContext());
       return builder.getIndexAttr(shapedTy.getShape()[index.getInt()]);
     }
   }
 
   // Fold dim to the size argument for an `AllocOp`, `ViewOp`, or `SubViewOp`.
   auto memrefType = argTy.dyn_cast<MemRefType>();
   if (!memrefType)
     return {};
 
   // The size at the given index is now known to be a dynamic size of a memref.
-  auto memref = memrefOrTensor().getDefiningOp();
+  auto *memref = memrefOrTensor().getDefiningOp();
   unsigned unsignedIndex = index.getValue().getZExtValue();
   if (auto alloc = dyn_cast_or_null<AllocOp>(memref))
     return *(alloc.getDynamicSizes().begin() +
              memrefType.getDynamicDimIndex(unsignedIndex));
 
   if (auto view = dyn_cast_or_null<ViewOp>(memref))
     return *(view.getDynamicSizes().begin() +
              memrefType.getDynamicDimIndex(unsignedIndex));

mlir/test/Conversion/StandardToLLVM/convert-to-llvmir.mlir

 // CHECK-LABEL: func @rank_of_ranked
 // CHECK32-LABEL: func @rank_of_ranked
 func @rank_of_ranked(%ranked: memref<?xi32>) {
   %rank = rank %ranked : memref<?xi32>
   return
 }
 // CHECK: llvm.mlir.constant(1 : index) : !llvm.i64
 // CHECK32: llvm.mlir.constant(1 : index) : !llvm.i32
+
+// -----
+
+// CHECK-LABEL: func @dim_of_unranked
+// CHECK32-LABEL: func @dim_of_unranked
+func @dim_of_unranked(%unranked: memref<*xi32>) -> index {
+  %c0 = constant 0 : index
+  %dim = dim %unranked, %c0 : memref<*xi32>
+  return %dim : index
+}
+// CHECK-NEXT: llvm.mlir.undef : !llvm.struct<(i64, ptr<i8>)>
+// CHECK-NEXT: llvm.insertvalue 
+// CHECK-NEXT: %[[UNRANKED_DESC:.*]] = llvm.insertvalue
+// CHECK-NEXT: %[[C0:.*]] = llvm.mlir.constant(0 : index) : !llvm.i64
+
+// CHECK-NEXT: %[[RANKED_DESC:.*]] = llvm.extractvalue %[[UNRANKED_DESC]][1]
+// CHECK-SAME:   : !llvm.struct<(i64, ptr<i8>)>
+
+// CHECK-NEXT: %[[ZERO_D_DESC:.*]] = llvm.bitcast %[[RANKED_DESC]]
+// CHECK-SAME:   : !llvm.ptr<i8> to !llvm.ptr<struct<(ptr<i32>, ptr<i32>, i64)>>
+
+// CHECK-NEXT: %[[C2_i32:.*]] = llvm.mlir.constant(2 : i32) : !llvm.i32
+// CHECK-NEXT: %[[C0_:.*]] = llvm.mlir.constant(0 : index) : !llvm.i64
+
+// CHECK-NEXT: %[[OFFSET_PTR:.*]] = llvm.getelementptr %[[ZERO_D_DESC]]{{\[}}
+// CHECK-SAME:   %[[C0_]], %[[C2_i32]]] : (!llvm.ptr<struct<(ptr<i32>, ptr<i32>,
+// CHECK-SAME:   i64)>>, !llvm.i64, !llvm.i32) -> !llvm.ptr<i64>
+
+// CHECK-NEXT: %[[C1:.*]] = llvm.mlir.constant(1 : index) : !llvm.i64
+// CHECK-NEXT: %[[INDEX_INC:.*]] = llvm.add %[[C1]], %[[C0]] : !llvm.i64
+
+// CHECK-NEXT: %[[SIZE_PTR:.*]] = llvm.getelementptr %[[OFFSET_PTR]]{{\[}}
+// CHECK-SAME:   %[[INDEX_INC]]] : (!llvm.ptr<i64>, !llvm.i64) -> !llvm.ptr<i64>
+
+// CHECK-NEXT: %[[SIZE:.*]] = llvm.load %[[SIZE_PTR]] : !llvm.ptr<i64>
+// CHECK-NEXT: llvm.return %[[SIZE]] : !llvm.i64
+
+// CHECK32: %[[SIZE:.*]] = llvm.load %{{.*}} : !llvm.ptr<i32>
+// CHECK32-NEXT: llvm.return %[[SIZE]] : !llvm.i32

mlir/test/mlir-cpu-runner/unranked_memref.mlir

-// RUN: mlir-opt %s -convert-linalg-to-loops -convert-linalg-to-llvm -convert-std-to-llvm | mlir-cpu-runner -e main -entry-point-result=void -shared-libs=%linalg_test_lib_dir/libmlir_runner_utils%shlibext,%linalg_test_lib_dir/libmlir_test_cblas%shlibext,%linalg_test_lib_dir/libmlir_test_cblas_interface%shlibext | FileCheck %s
+// RUN: mlir-opt %s -convert-linalg-to-loops -convert-linalg-to-llvm -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
 
 // CHECK: rank = 2
 // CHECK-SAME: sizes = [10, 3]
 // CHECK-SAME: strides = [3, 1]
 // CHECK-COUNT-10: [10, 10, 10]
 //
 // CHECK: rank = 2
 // CHECK-SAME: sizes = [10, 3]
func @main() -> () {
     store %i8_z, %I8[]: memref<i8>
     %U4 = memref_cast %I8 : memref<i8> to memref<*xi8>
     call @print_memref_i8(%U4) : (memref<*xi8>) -> ()
 
     dealloc %A : memref<10x3xf32, 0>
 
     call @return_var_memref_caller() : () -> ()
     call @return_two_var_memref_caller() : () -> ()
+    call @dim_op_of_unranked() : () -> ()
     return
 }
 
 func @print_memref_i8(memref<*xi8>) attributes { llvm.emit_c_interface }
 func @print_memref_f32(memref<*xf32>) attributes { llvm.emit_c_interface }
 
 func @return_two_var_memref_caller() {
   %0 = alloca() : memref<4x3xf32>
func @return_var_memref_caller() {
   call @print_memref_f32(%1) : (memref<*xf32>) -> ()
   return
 }
 
 func @return_var_memref(%arg0: memref<4x3xf32>) -> memref<*xf32> {
   %0 = memref_cast %arg0: memref<4x3xf32> to memref<*xf32>
   return %0 : memref<*xf32>
 }
+
+func @print_i64(index) -> ()
+func @print_newline() -> ()
+
+func @dim_op_of_unranked() {
+  %ranked = alloc() : memref<4x3xf32>
+  %unranked = memref_cast %ranked: memref<4x3xf32> to memref<*xf32>
+
+  %c0 = constant 0 : index
+  %dim_0 = dim %unranked, %c0 : memref<*xf32>
+  call @print_i64(%dim_0) : (index) -> ()
+  call @print_newline() : () -> ()
+  // CHECK: 4
+
+  %c1 = constant 1 : index
+  %dim_1 = dim %unranked, %c1 : memref<*xf32>
+  call @print_i64(%dim_1) : (index) -> ()
+  call @print_newline() : () -> ()
+  // CHECK: 3
+
+  return
+}