[mlir][Standard] Allow unranked memrefs as operands to dim and rank

[herhut]

Summary

std.dim currently only accepts ranked memrefs and std.rank is limited to tensors.

[ftynse]

Please fix the broken test.

mlir/include/mlir/Dialect/StandardOps/IR/Ops.td

1412

Irrelevant to this commit, but I find it confusing that AnyRankedOrUnrankedMemRef is less strict than AnyMemRef.

[herhut]

Fix test

mlir/include/mlir/Dialect/StandardOps/IR/Ops.td

1412

I agree, I found this confusing and got it wrong. I thought about cleaning this up but the problem is that if I modify the type constraints than downstream users will see changes in their dialect and likely won't notice, as we widen the meaning of AnyMemRef. Any ideas on a good process to clean this up?

[ftynse]

mlir/include/mlir/Dialect/StandardOps/IR/Ops.td

1412

I'd post an announcement on Discourse saying you intend to do this by <date> and give people some time to react. IMO, a week should be okay.

mlir/include/mlir/Dialect/StandardOps/IR/Ops.td

let description = [{
     %y = dim %A, %c1 : tensor<4 x ? x f32>
 
     // Equivalent generic form:
     %x = "std.dim"(%A, %c0) : (tensor<4 x ? x f32>, index) -> index
     %y = "std.dim"(%A, %c1) : (tensor<4 x ? x f32>, index) -> index
     ```
   }];
 
-  let arguments = (ins AnyTypeOf<[AnyMemRef, AnyTensor],
+  let arguments = (ins AnyTypeOf<[AnyRankedOrUnrankedMemRef, AnyTensor],
                                  "any tensor or memref type">:$memrefOrTensor,
                        Index:$index);
   let results = (outs Index:$result);
 
   let assemblyFormat = [{
     attr-dict $memrefOrTensor `,` $index `:` type($memrefOrTensor)
   }];
 

 
 //===----------------------------------------------------------------------===//
 // RankOp
 //===----------------------------------------------------------------------===//
 
 def RankOp : Std_Op<"rank", [NoSideEffect]> {
   let summary = "rank operation";
   let description = [{
-    The `rank` operation takes a tensor operand and returns its rank.
+    The `rank` operation takes a memref/tensor operand and returns its rank.
 
     Example:
 
     ```mlir
-    %1 = rank %0 : tensor<*xf32>
+    %1 = rank %arg0 : tensor<*xf32>
+    %2 = rank %arg1 : memref<*xf32>
     ```
   }];
 
-  let arguments = (ins AnyTensor);
+  let arguments = (ins AnyTypeOf<[AnyRankedOrUnrankedMemRef, AnyTensor],
+                                 "any tensor or memref type">:$memrefOrTensor);
   let results = (outs Index);
   let verifier = ?;
 
   let builders = [OpBuilder<
     "OpBuilder &builder, OperationState &result, Value tensor", [{
       auto indexType = builder.getIndexType();
       build(builder, result, indexType, tensor);
     }]>];
 
   let hasFolder = 1;
-  let assemblyFormat = "operands attr-dict `:` type(operands)";
+  let assemblyFormat = "$memrefOrTensor attr-dict `:` type($memrefOrTensor)";
 }
 
 //===----------------------------------------------------------------------===//
 // ReOp
 //===----------------------------------------------------------------------===//
 
 def ReOp : Std_Op<"re",
     [NoSideEffect,

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

LogicalResult PrefetchOp::fold(ArrayRef<Attribute> cstOperands,
   return foldMemRefCast(*this);
 }
 
 //===----------------------------------------------------------------------===//
 // RankOp
 //===----------------------------------------------------------------------===//
 
 OpFoldResult RankOp::fold(ArrayRef<Attribute> operands) {
-  // Constant fold rank when the rank of the tensor is known.
+  // Constant fold rank when the rank of the operand is known.
   auto type = getOperand().getType();
-  if (auto tensorType = type.dyn_cast<RankedTensorType>())
-    return IntegerAttr::get(IndexType::get(getContext()), tensorType.getRank());
+  if (auto shapedType = type.dyn_cast<ShapedType>())
+    if (shapedType.hasRank())
+      return IntegerAttr::get(IndexType::get(getContext()),
+                              shapedType.getRank());
   return IntegerAttr();
 }
 
 //===----------------------------------------------------------------------===//
 // ReturnOp
 //===----------------------------------------------------------------------===//
 
 static LogicalResult verify(ReturnOp op) {

mlir/test/Transforms/constant-fold.mlir

func @fold_rank() -> (index) {
   %rank_0 = rank %const_0 : tensor<2x1x4xi32>
 
   // CHECK-NEXT: return [[C3]]
   return %rank_0 : index
 }
 
 // -----
 
+// CHECK-LABEL: func @fold_rank_memref
+func @fold_rank_memref(%arg0 : memref<?x?xf32>) -> (index) {
+  // Fold a rank into a constant
+  // CHECK-NEXT: [[C2:%.+]] = constant 2 : index
+  %rank_0 = rank %arg0 : memref<?x?xf32>
+
+  // CHECK-NEXT: return [[C2]]
+  return %rank_0 : index
+}
+
+// -----
+
 // CHECK-LABEL: func @nested_isolated_region
 func @nested_isolated_region() {
   // CHECK-NEXT: func @isolated_op
   // CHECK-NEXT: constant 2
   func @isolated_op() {
     %0 = constant 1 : i32
     %2 = addi %0, %0 : i32
     "foo.yield"(%2) : (i32) -> ()