[mlir][bufferization] Prevent crash in one shot bufferization with unranked tensor cast

[Lewuathe]

Summary

One shot bufferization does not support bufferizing the cast between unranked tensors. To prevent the crash, we can check the compatibility of the result type in advance. Reported in https://github.com/llvm/llvm-project/issues/62369.

[Lewuathe]

Add newline.

[springerm]

mlir/lib/Dialect/Tensor/Transforms/BufferizableOpInterfaceImpl.cpp

97–101

This can be bufferized. The crash is caused by the fact that unraked->unranked casting is valid on tensor but not on memrefs for some reason. (I think we should handle it the same way on memrefs/tensors.) In that case, we can directly replaceOpWithBufferizedValues(rewriter, op, *resultBuffer);.

[Lewuathe]

@springerm Thanks for the advice.
Lowering memref.cast op with unranked memref for input and output to another dialect (e.g. LLVM) is not allowed for now.

./bin/mlir-opt cast.mlir
cast.mlir:2:8: error: 'memref.cast' op operand type 'memref<4x?xi32>' and result type 'memref<4x?xf32>' are cast incompatible
  %0 = memref.cast %arg : memref<4x?xi32> to memref<4x?xf32>
       ^
cast.mlir:2:8: note: see current operation: %0 = "memref.cast"(%arg0) : (memref<4x?xi32>) -> memref<4x?xf32>

We can support a bufferization from unranked to unranked somehow, but it may cause runtime errors later after lowering more. Casting the unranked tensor to an unranked tensor looks like just passing the responsibility to the bufferization with unranked memref, so I did not find any good hint from that.

So what we can do is to make the pass return the failure instead of catching in the assertion?

[springerm]

Lowering memref.cast op with unranked memref for input and output to another dialect (e.g. LLVM) is not allowed for now.

Really? I thought the reason why we have unranked memrefs is so that we can call external functions. How does this work if we cannot lower it to LLVM?

Casting unranked tensor to unranked tensor looks like just passing the responsibility to the bufferization with unranked memref so I did not find any good hint with that.

Casting unranked -> unranked is a no-op. Such casts should fold away. This could be done during bufferization to work around the inconsistency that unranked->unranked is allowed on tensors but not on memrefs:

if (resultBuffer->getType() == *resultMemRefType) {
  // This cast is a no-op.
  replaceOpWithBufferizedValues(rewriter, op, *resultBuffer);
  return success();
}

[Lewuathe]

@springerm

Casting unranked -> unranked is a no-op. Such casts should fold away. This could be done during bufferization to work around the inconsistency that unranked->unranked is allowed on tensors but not on memrefs:

Thanks for the advice. I think that should work. I'll try to eliminate the cast as no-op.

[Lewuathe]

Canonicalize unranked to unranked as no-op.

[Lewuathe]

@springerm No-op casting simply seems to be working as you said. Could you review this change when you get a chance?

[Lewuathe]

@springerm Sorry for bothering you from time to time, but could you review this change when you get a good chance?

[springerm]

Sorry for the delay, I lost track of some reviews.

mlir/lib/Dialect/Tensor/Transforms/BufferizableOpInterfaceImpl.cpp

LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
     if (failed(resultBuffer))
       return failure();
 
     // Compute the new type.
     auto resultMemRefType =
         bufferization::getBufferType(castOp.getResult(), options);
     if (failed(resultMemRefType))
       return failure();
+    if (resultBuffer->getType() == *resultMemRefType) {
+      // This cast is a no-op.
+      replaceOpWithBufferizedValues(rewriter, op, *resultBuffer);
+      return success();
+    }
 
     // Replace the op with a memref.cast.
     assert(memref::CastOp::areCastCompatible(resultBuffer->getType(),
                                              *resultMemRefType) &&
            "CallOp::bufferize: cast incompatible");
     replaceOpWithNewBufferizedOp<memref::CastOp>(
         rewriter, op, *resultMemRefType, *resultBuffer);
 

mlir/test/Dialect/Bufferization/Transforms/one-shot-bufferize.mlir

func.func @read_of_alias(%t: tensor<100xf32>, %pos1: index, %pos2: index,
   // CHECK-TOP-DOWN-ANALYSIS: memref.copy
   // CHECK-TOP-DOWN-ANALYSIS: memref.store %{{.*}}, %[[alloc]]
   %0 = tensor.insert %f into %t[%pos1] : tensor<100xf32>
   %1 = tensor.extract_slice %t[%pos2][%sz][1] : tensor<100xf32> to tensor<?xf32>
   %2 = tensor.extract %1[%pos3] : tensor<?xf32>
   %3 = tensor.extract %0[%pos3] : tensor<100xf32>
   return %2, %3 : f32, f32
 }
+
+// -----
+
+// CHECK-LABEL: func @from_unranked_to_unranked
+func.func @from_unranked_to_unranked(%arg0: tensor<*xi32>) -> tensor<*xi32> {
+  // CHECK: return %arg{{.*}} : tensor<*xi32>
+  %0 = tensor.cast %arg0 : tensor<*xi32> to tensor<*xi32>
+  return %0 : tensor<*xi32>
+}