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

[springerm]

Summary

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

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

return {{op->getOpResult(0) /*result*/, BufferRelation::Equivalent,
              /*isDefinite=*/false}};
   }
 
   LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
                           const BufferizationOptions &options) const {
     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
     // copy would be needed.
     FailureOr<Value> maybeTrueBuffer =
         getBuffer(rewriter, selectOp.getTrueValue(), options);
     FailureOr<Value> maybeFalseBuffer =
         getBuffer(rewriter, selectOp.getFalseValue(), options);

mlir/test/Dialect/Arith/bufferize.mlir

-// 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
 func.func @index_cast(%tensor: tensor<i32>, %scalar: i32) -> (tensor<index>, index) {
   %index_tensor = arith.index_cast %tensor : tensor<i32> to tensor<index>
   %index_scalar = arith.index_cast %scalar : i32 to index
   return %index_tensor, %index_scalar : tensor<index>, index
 }

 // CHECK-DAG:           %[[FALSE_VAL_MEMREF:.*]] = bufferization.to_memref %[[FALSE_VAL]] : memref<f32>
 // CHECK:           %[[RET_MEMREF:.*]] = arith.select %[[PRED]], %[[TRUE_VAL_MEMREF]], %[[FALSE_VAL_MEMREF]] : memref<f32>
 // CHECK:           %[[RET:.*]] = bufferization.to_tensor %[[RET_MEMREF]] : memref<f32>
 // CHECK:           return %[[RET]] : tensor<f32>
 func.func @select(%arg0: i1, %arg1: tensor<f32>, %arg2: tensor<f32>) -> tensor<f32> {
   %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>
+}