BEGIN_PUBLIC [mlir] Add support for unranked case for `tensor_store` and `tensor_load` ops. END_PUBLIC

[herhut]

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

2946

Add the inverse to getTensorTypeFromMemRefType?

2951–2956

Accidental?

[pifon2a]

Address the comments.

[ftynse]

Please drop additional tags (BEGIN_PUBLIC / END_PUBLIC) from the commit message.

[herhut]

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

2952

Can this ever anything else than TensorType? Both ranked and unranked are TensorType, right?

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

2993

Turns out this is not actually needed, right? Sorry, I was confused.

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

let description = [{
     Example:
 
     ```mlir
     // Produces a value of tensor<4x?xf32> type.
     %12 = tensor_load %10 : memref<4x?xf32, #layout, memspace0>
     ```
   }];
 
-  let arguments = (ins Arg<AnyMemRef, "the reference to load from",
-                           [MemRead]>:$memref);
+  let arguments = (ins Arg<AnyRankedOrUnrankedMemRef,
+                       "the reference to load from", [MemRead]>:$memref);
   let results = (outs AnyTensor:$result);
   // TensorLoadOp is fully verified by traits.
   let verifier = ?;
 
   let builders = [OpBuilder<
     "OpBuilder &builder, OperationState &result, Value memref", [{
-      auto memrefType = memref.getType().cast<MemRefType>();
-      auto resultType = RankedTensorType::get(memrefType.getShape(),
-                                              memrefType.getElementType());
       result.addOperands(memref);
-      result.addTypes(resultType);
+      result.addTypes(getTensorTypeFromMemRefType(memref.getType()));
   }]>];
 
   let extraClassDeclaration = [{
     /// The result of a tensor_load is always a tensor.
-    TensorType getType() { return getResult().getType().cast<TensorType>(); }
+    TensorType getType() { 
+      Type resultType = getResult().getType();
+      if (resultType.isa<TensorType>())
+        return resultType.cast<TensorType>();
+      return {};
+    }
   }];
 
   let assemblyFormat = "$memref attr-dict `:` type($memref)";
 }
 
 //===----------------------------------------------------------------------===//
 // TensorStoreOp
 //===----------------------------------------------------------------------===//
let description = [{
 
     ```mlir
     %9 = dim %8, 1 : tensor<4x?xf32>
     %10 = alloc(%9) : memref<4x?xf32, #layout, memspace0>
     tensor_store %8, %10 : memref<4x?xf32, #layout, memspace0>
     ```
   }];
 
-  let arguments = (ins AnyTensor:$tensor,
-                       Arg<AnyMemRef, "the reference to store to",
-                           [MemWrite]>:$memref);
+  let arguments = (ins AnyTensor:$tensor, Arg<AnyRankedOrUnrankedMemRef,
+                       "the reference to store to", [MemWrite]>:$memref);
   // TensorStoreOp is fully verified by traits.
   let verifier = ?;
 
   let assemblyFormat = "$tensor `,` $memref attr-dict `:` type($memref)";
 }
 
 //===----------------------------------------------------------------------===//
 // TruncateIOp

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

 
 //===----------------------------------------------------------------------===//
 // Helpers for Tensor[Load|Store]Op
 //===----------------------------------------------------------------------===//
 
 static Type getTensorTypeFromMemRefType(Type type) {
   if (auto memref = type.dyn_cast<MemRefType>())
     return RankedTensorType::get(memref.getShape(), memref.getElementType());
+  if (auto memref = type.dyn_cast<UnrankedMemRefType>())
+    return UnrankedTensorType::get(memref.getElementType());
+  return NoneType::get(type.getContext());
+}
+
+static Type getMemRefTypeFromTensorType(Type type) {
+  if (auto tensor = type.dyn_cast<MemRefType>())
+    return MemRefType::get(tensor.getShape(), tensor.getElementType());
+  if (auto tensor = type.dyn_cast<UnrankedMemRefType>())
+    return UnrankedMemRefType::get(tensor.getElementType(),
+                                   tensor.getMemorySpace());
   return NoneType::get(type.getContext());
 }
 
 //===----------------------------------------------------------------------===//
 // TruncateIOp
 //===----------------------------------------------------------------------===//
 
 static LogicalResult verify(TruncateIOp op) {

mlir/test/IR/core-ops.mlir

 func @tensor_load_store(%0 : memref<4x4xi32>) {
   // CHECK: %[[TENSOR:.*]] = tensor_load %[[MEMREF:.*]] : memref<4x4xi32>
   %1 = tensor_load %0 : memref<4x4xi32>
   // CHECK: tensor_store %[[TENSOR]], %[[MEMREF]] : memref<4x4xi32>
   tensor_store %1, %0 : memref<4x4xi32>
   return
 }
 
+// CHECK-LABEL: func @unranked_tensor_load_store
+func @unranked_tensor_load_store(%0 : memref<*xi32>) {
+  // CHECK: %[[TENSOR:.*]] = tensor_load %[[MEMREF:.*]] : memref<*xi32>
+  %1 = tensor_load %0 : memref<*xi32>
+  // CHECK: tensor_store %[[TENSOR]], %[[MEMREF]] : memref<*xi32>
+  tensor_store %1, %0 : memref<*xi32>
+  return
+}
+
 // CHECK-LABEL: func @atomic_rmw
 // CHECK-SAME: ([[BUF:%.*]]: memref<10xf32>, [[VAL:%.*]]: f32, [[I:%.*]]: index)
 func @atomic_rmw(%I: memref<10xf32>, %val: f32, %i : index) {
   %x = atomic_rmw "addf" %val, %I[%i] : (f32, memref<10xf32>) -> f32
   // CHECK: atomic_rmw "addf" [[VAL]], [[BUF]]{{\[}}[[I]]]
   return
 }