[mlir] Fix bufferization.alloc_tensor canonicalization crash

[joker-eph]

This make sure that an invalid negative dimension is ignored and stays dynamic instead of crashing the compiler.

Fixes #70887

[llvmbot]

@llvm/pr-subscribers-mlir

@llvm/pr-subscribers-mlir-bufferization

Author: Mehdi Amini (joker-eph)

Changes

This make sure that an invalid negative dimension is ignored and stays dynamic instead of crashing the compiler.

Fixes #70887

---

Full diff: https://github.com/llvm/llvm-project/pull/70891.diff

2 Files Affected:

• (modified) mlir/lib/Dialect/Bufferization/IR/BufferizationOps.cpp (+5-1)
• (modified) mlir/test/Dialect/Bufferization/canonicalize.mlir (+13)

diff --git a/mlir/lib/Dialect/Bufferization/IR/BufferizationOps.cpp b/mlir/lib/Dialect/Bufferization/IR/BufferizationOps.cpp
index 8f19245efdba6c8..033eeac1939fc49 100644
--- a/mlir/lib/Dialect/Bufferization/IR/BufferizationOps.cpp
+++ b/mlir/lib/Dialect/Bufferization/IR/BufferizationOps.cpp
@@ -314,7 +314,11 @@ struct ReplaceStaticShapeDims : OpRewritePattern<AllocTensorOp> {
       Value value = op.getDynamicSizes()[dynValCounter++];
       APInt intVal;
       if (matchPattern(value, m_ConstantInt(&intVal))) {
-        newShape[i] = intVal.getSExtValue();
+        int64_t dim = intVal.getSExtValue();
+        if (dim >= 0)
+          newShape[i] = intVal.getSExtValue();
+        else
+          newDynamicSizes.push_back(value);
       } else {
         newDynamicSizes.push_back(value);
       }
diff --git a/mlir/test/Dialect/Bufferization/canonicalize.mlir b/mlir/test/Dialect/Bufferization/canonicalize.mlir
index 12f13743febb73d..3ba283928a83f0e 100644
--- a/mlir/test/Dialect/Bufferization/canonicalize.mlir
+++ b/mlir/test/Dialect/Bufferization/canonicalize.mlir
@@ -351,3 +351,16 @@ func.func @dealloc_base_memref_extract_of_alloc(%arg0: memref<2xi32>) {
 //  CHECK-SAME:([[ARG0:%.+]]: memref<2xi32>)
 //   CHECK-NOT: memref.alloc(
 //       CHECK: bufferization.dealloc ([[ARG0]] : memref<2xi32>) if (%true
+
+// -----
+
+// CHECK-LABEL: func @negative_input
+func.func @negative_input() -> tensor<?x?x?xf16> {
+  %idx27 = index.constant 27
+  %idx-3 = index.constant -3  // negative integer?
+  %c10 = arith.constant 10 : index
+// CHECK: bufferization.alloc_tensor
+// CHECK-SAME: tensor<10x?x27xf16>
+  %11 = bufferization.alloc_tensor(%c10, %idx-3, %idx27) : tensor<?x?x?xf16>
+  return %11 : tensor<?x?x?xf16>
+}