[Mlir-commits] [mlir] 2257ebd - [mlir][arith][bufferize] Fix crash when bufferizing elementwise arith.select

[Matthias Springer]

Author: Matthias Springer
Date: 2023-08-23T17:40:53+02:00
New Revision: 2257ebd33ad1f3b7912cfab18edc3bbcfa8fa046

URL: https://github.com/llvm/llvm-project/commit/2257ebd33ad1f3b7912cfab18edc3bbcfa8fa046
DIFF: https://github.com/llvm/llvm-project/commit/2257ebd33ad1f3b7912cfab18edc3bbcfa8fa046.diff

LOG: [mlir][arith][bufferize] Fix crash when bufferizing elementwise arith.select

Elementwise arith.select are currently not supported. Emit an error message instead of crashing. This fixes #61707.

Differential Revision: https://reviews.llvm.org/D158617

Added: 
    

Modified: 
    mlir/lib/Dialect/Arith/Transforms/BufferizableOpInterfaceImpl.cpp
    mlir/test/Dialect/Arith/bufferize.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Dialect/Arith/Transforms/BufferizableOpInterfaceImpl.cpp b/mlir/lib/Dialect/Arith/Transforms/BufferizableOpInterfaceImpl.cpp
index 1bfb0c7e102cd3..f69b2557eec922 100644
--- a/mlir/lib/Dialect/Arith/Transforms/BufferizableOpInterfaceImpl.cpp
+++ b/mlir/lib/Dialect/Arith/Transforms/BufferizableOpInterfaceImpl.cpp
@@ -134,6 +134,13 @@ struct SelectOpInterface
     auto selectOp = cast<arith::SelectOp>(op);
     Location loc = selectOp.getLoc();
 
+    // Elementwise conditions are not supported yet. To bufferize such an op,
+    // it could be lowered to an elementwise "linalg.generic" with a new
+    // "tensor.empty" out tensor, followed by "empty tensor elimination". Such
+    // IR will bufferize.
+    if (!selectOp.getCondition().getType().isInteger(1))
+      return op->emitOpError("only i1 condition values are supported");
+
     // TODO: It would be more efficient to copy the result of the `select` op
     // instead of its OpOperands. In the worst case, 2 copies are inserted at
     // the moment (one for each tensor). When copying the op result, only one

diff  --git a/mlir/test/Dialect/Arith/bufferize.mlir b/mlir/test/Dialect/Arith/bufferize.mlir
index 9fcf466f2649cb..944954e9e4edd8 100644
--- a/mlir/test/Dialect/Arith/bufferize.mlir
+++ b/mlir/test/Dialect/Arith/bufferize.mlir
@@ -1,5 +1,5 @@
-// RUN: mlir-opt %s -arith-bufferize -split-input-file | FileCheck %s
-// RUN: mlir-opt %s -arith-bufferize=alignment=64 -split-input-file | FileCheck --check-prefix=ALIGNED %s
+// RUN: mlir-opt %s -arith-bufferize -split-input-file -verify-diagnostics | FileCheck %s
+// RUN: mlir-opt %s -arith-bufferize=alignment=64 -split-input-file -verify-diagnostics | FileCheck --check-prefix=ALIGNED %s
 
 // CHECK-LABEL:   func @index_cast(
 // CHECK-SAME:  %[[TENSOR:.*]]: tensor<i32>, %[[SCALAR:.*]]: i32
@@ -96,3 +96,12 @@ func.func @select(%arg0: i1, %arg1: tensor<f32>, %arg2: tensor<f32>) -> tensor<f
   %0 = arith.select %arg0, %arg1, %arg2 : tensor<f32>
   return %0 : tensor<f32>
 }
+
+// -----
+
+func.func @elementwise_select(%arg0: tensor<5xi1>, %arg1: tensor<5xi32>, %arg2: tensor<5xi32>) -> tensor<5xi32> {
+  // expected-error @below{{only i1 condition values are supported}}
+  // expected-error @below{{failed to bufferize op}}
+  %0 = arith.select %arg0, %arg1, %arg2 : tensor<5xi1>, tensor<5xi32>
+  return %0 : tensor<5xi32>
+}