[mlir] Lower RankOp to LLVM for unranked memrefs.

[herhut]

Thanks!

mlir/lib/Conversion/StandardToLLVM/StandardToLLVM.cpp

2412

Maybe rephrase into a positive: RankOp should be folded away for ranked operands.

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

1301

Maybe return rank to make it used?

[ftynse]

mlir/lib/Conversion/StandardToLLVM/StandardToLLVM.cpp

2412

Can't we just emit a constant instead?

[pifon2a]

Address the comments.

[pifon2a]

Test

[pifon2a]

Lint

[ftynse]

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

1307

Irrelevant to this patch, but shouldn't we use i32 for the rank if bitwidth of index is 32?

[pifon2a]

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

1301

Does it change anything?

[pifon2a]

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

1307

https://reviews.llvm.org/D85300

mlir/lib/Conversion/StandardToLLVM/StandardToLLVM.cpp

matchAndRewrite(Operation *op, ArrayRef<Value> operands,
     Value index = dimOp.index();
     int64_t rank = memRefType.getRank();
     MemRefDescriptor memrefDescriptor(transformed.memrefOrTensor());
     rewriter.replaceOp(op, {memrefDescriptor.size(rewriter, loc, index, rank)});
     return success();
   }
 };
 
+struct RankOpLowering : public ConvertOpToLLVMPattern<RankOp> {
+  using ConvertOpToLLVMPattern<RankOp>::ConvertOpToLLVMPattern;
+
+  LogicalResult
+  matchAndRewrite(Operation *op, ArrayRef<Value> operands,
+                  ConversionPatternRewriter &rewriter) const override {
+    Location loc = op->getLoc();
+    Type operandType = cast<RankOp>(op).memrefOrTensor().getType();
+    if (auto unrankedMemRefType = operandType.dyn_cast<UnrankedMemRefType>()) {
+      UnrankedMemRefDescriptor desc(RankOp::Adaptor(operands).memrefOrTensor());
+      rewriter.replaceOp(op, {desc.rank(rewriter, loc)});
+      return success();
+    }
+    if (auto rankedMemRefType = operandType.dyn_cast<MemRefType>()) {
+      rewriter.replaceOp(
+          op, {createIndexConstant(rewriter, loc, rankedMemRefType.getRank())});
+      return success();
+    }
+    return failure();
+  }
+};
+
 // 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>
 struct LoadStoreOpLowering : public ConvertOpToLLVMPattern<Derived> {
   using ConvertOpToLLVMPattern<Derived>::ConvertOpToLLVMPattern;
   using Base = LoadStoreOpLowering<Derived>;
 
void mlir::populateStdToLLVMMemoryConversionPatterns(
     LLVMTypeConverter &converter, OwningRewritePatternList &patterns) {
   // clang-format off
   patterns.insert<
       AssumeAlignmentOpLowering,
       DeallocOpLowering,
       DimOpLowering,
       LoadOpLowering,
       MemRefCastOpLowering,
+      RankOpLowering,
       StoreOpLowering,
       SubViewOpLowering,
       ViewOpLowering,
       AllocOpLowering>(converter);
   // clang-format on
 }
 
 void mlir::populateStdToLLVMFuncOpConversionPattern(

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

 // CHECK-SAME: -> !llvm.struct<(ptr<i64>, ptr<i64>, i64, array<1 x i64>, array<1 x i64>)>
 // CHECK32-LABEL: func @memref_index
 // CHECK32-SAME: %arg0: !llvm.ptr<i32>, %arg1: !llvm.ptr<i32>,
 // CHECK32-SAME: %arg2: !llvm.i32, %arg3: !llvm.i32, %arg4: !llvm.i32)
 // CHECK32-SAME: -> !llvm.struct<(ptr<i32>, ptr<i32>, i32, array<1 x i32>, array<1 x i32>)>
 func @memref_index(%arg0: memref<32xindex>) -> memref<32xindex> {
   return %arg0 : memref<32xindex>
 }
+
+// -----
+
+// CHECK-LABEL: func @rank_of_unranked
+// CHECK32-LABEL: func @rank_of_unranked
+func @rank_of_unranked(%unranked: memref<*xi32>) {
+  %rank = rank %unranked : memref<*xi32>
+  return
+}
+// CHECK-NEXT: llvm.mlir.undef
+// CHECK-NEXT: llvm.insertvalue
+// CHECK-NEXT: llvm.insertvalue
+// CHECK-NEXT: llvm.extractvalue %{{.*}}[0] : !llvm<"{ i64, i8* }">
+// CHECK32: llvm.extractvalue %{{.*}}[0] : !llvm<"{ i64, i8* }">
+
+// 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