[mlir] Support unranked types in func signature conversion in BufferPlacement.

[pifon2a]

Summary

Currently, only ranked tensor args and results can be converted to memref types.

[herhut]

Nice!

mlir/lib/Transforms/BufferPlacement.cpp

718

Would isa<BaseMemRefType> work?

[pifon2a]

mlir/lib/Transforms/BufferPlacement.cpp

718

yep, it does. Thanks!

[pifon2a]

Addressed the comment

mlir/lib/Transforms/BufferPlacement.cpp

 //===----------------------------------------------------------------------===//
 // BufferAssignmentTypeConverter
 //===----------------------------------------------------------------------===//
 
 /// Registers conversions into BufferAssignmentTypeConverter
 BufferAssignmentTypeConverter::BufferAssignmentTypeConverter() {
   // Keep all types unchanged.
   addConversion([](Type type) { return type; });
-  // A type conversion that converts ranked-tensor type to memref type.
+  // Convert RankedTensorType to MemRefType.
   addConversion([](RankedTensorType type) {
     return (Type)MemRefType::get(type.getShape(), type.getElementType());
   });
+  // Convert UnrankedTensorType to UnrankedMemRefType.
+  addConversion([](UnrankedTensorType type) {
+    return (Type)UnrankedMemRefType::get(type.getElementType(), 0);
+  });
 }
 
 /// Checks if `type` has been converted from non-memref type to memref.
 bool BufferAssignmentTypeConverter::isConvertedMemref(Type type, Type before) {
-  return type.isa<MemRefType>() && !before.isa<MemRefType>();
+  auto isMemRefType = [](Type type) -> bool {
+    return type.isa<MemRefType>() || type.isa<UnrankedMemRefType>();
+  };
+  return isMemRefType(type) && !isMemRefType(before);
 }
 
 //===----------------------------------------------------------------------===//
 // BufferPlacementPass construction
 //===----------------------------------------------------------------------===//
 
 std::unique_ptr<Pass> mlir::createBufferPlacementPass() {
   return std::make_unique<BufferPlacementPass>();
 }

mlir/test/Transforms/buffer-placement-preparation-allowed-memref-results.mlir

 func @simple_signature_conversion(%arg0: tensor<4x8xf32>) -> tensor<4x8xf32> {
   return %arg0 : tensor<4x8xf32>
 }
 //      CHECK: (%[[ARG0:.*]]: [[TYPE:.*]]<[[RANK:.*]]>) -> [[TYPE]]<[[RANK]]>
 // CHECK-NEXT: return %[[ARG0]]
 
 // -----
 
+// CHECK-LABEL: func @func_with_unranked_arg_and_result
+func @func_with_unranked_arg_and_result(%arg0: tensor<*xf32>) -> tensor<*xf32> {
+  return %arg0 : tensor<*xf32>
+}
+// CHECK-SAME: ([[ARG:%.*]]: memref<*xf32>) -> memref<*xf32>
+// CHECK-NEXT: return [[ARG]] : memref<*xf32>
+
+// -----
+
 // CHECK-LABEL: func @func_and_block_signature_conversion
 func @func_and_block_signature_conversion(%arg0 : tensor<2xf32>, %cond : i1, %arg1: tensor<4x4xf32>) -> tensor<4x4xf32>{
     cond_br %cond, ^bb1, ^bb2
   ^bb1:
     br ^exit(%arg0 : tensor<2xf32>)
   ^bb2:
     br ^exit(%arg0 : tensor<2xf32>)
   ^exit(%arg2: tensor<2xf32>):

mlir/test/Transforms/buffer-placement-preparation.mlir

 }
 // CHECK: (%[[CALLER_ARG:.*]]: memref<5xf32>, %[[CALLER_RESULT:.*]]: memref<5xf32>)
 // CHECK: %[[X0:.*]] = alloc()
 // CHECK: %[[X1:.*]] = call @callee(%[[CALLER_ARG]], %[[X0]])
 // CHECK: %[[Y0:.*]] = alloc()
 // CHECK: %[[Y1:.*]] = call @callee(%[[X0]], %[[Y0]])
 // CHECK: linalg.copy(%[[Y0]], %[[CALLER_RESULT]])
 // CHECK: return
+
+// CHECK-LABEL: func @func_with_unranked_arg
+func @func_with_unranked_arg(%arg0: tensor<*xf32>) {
+  return
+}
+// CHECK-SAME: ([[ARG:%.*]]: memref<*xf32>)