[Mlir-commits] [mlir] 5a71f7a - [mlir] Fix bufferization.alloc_tensor canonicalization crash (#70891)

[llvmlistbot at llvm.org]

Author: Mehdi Amini
Date: 2023-10-31T20:58:29-07:00
New Revision: 5a71f7a4b725db3b9bf1457e2ca6be2db6aa3996

URL: https://github.com/llvm/llvm-project/commit/5a71f7a4b725db3b9bf1457e2ca6be2db6aa3996
DIFF: https://github.com/llvm/llvm-project/commit/5a71f7a4b725db3b9bf1457e2ca6be2db6aa3996.diff

LOG: [mlir] Fix bufferization.alloc_tensor canonicalization crash (#70891)

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

Fixes #70887

Added: 
    

Modified: 
    mlir/lib/Dialect/Bufferization/IR/BufferizationOps.cpp
    mlir/test/Dialect/Bufferization/canonicalize.mlir

Removed: 
    


################################################################################
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>
+}