[Mlir-commits] [mlir] [mlir][linalg] Fix linalg.select crash with index type operands (PR #179056)

[llvmlistbot at llvm.org]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir-linalg

Author: puneeth_aditya_5656 (mugiwaraluffy56)

<details>
<summary>Changes</summary>

## Summary
Fix UNREACHABLE crash in `linalg.select` when using index type operands.

The `buildTernaryFn` function in `RegionBuilderHelper` only validated integer and floating point types, causing an assertion failure when `linalg.select` was used with index type operands.

## Changes
- Add `isIndex()` helper function to check for `IndexType`
- Add `tailIndex` check in `buildTernaryFn` for index type operands
- Add regression test `@<!-- -->select_index` in `named-ops.mlir`

## Test plan
- Added `@<!-- -->select_index` test case that reproduces the crash
- Run: `ninja check-mlir-linalg`

Fixes #<!-- -->179046

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


2 Files Affected:

- (modified) mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp (+5-1) 
- (modified) mlir/test/Dialect/Linalg/named-ops.mlir (+13) 


``````````diff
diff --git a/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp b/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp
index 0f0e308bba78e..e23a8af059fac 100644
--- a/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp
+++ b/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp
@@ -617,11 +617,12 @@ class RegionBuilderHelper {
     bool tailFloatingPoint =
         isFloatingPoint(arg0) && isFloatingPoint(arg1) && isFloatingPoint(arg2);
     bool tailInteger = isInteger(arg0) && isInteger(arg1) && isInteger(arg2);
+    bool tailIndex = isIndex(arg0) && isIndex(arg1) && isIndex(arg2);
     OpBuilder::InsertionGuard g(builder);
     builder.setInsertionPointToEnd(&block);
     switch (ternaryFn) {
     case TernaryFn::select:
-      if (!headBool && !(tailFloatingPoint || tailInteger))
+      if (!headBool && !(tailFloatingPoint || tailInteger || tailIndex))
         llvm_unreachable("unsupported non numeric type");
       return arith::SelectOp::create(builder, arg0.getLoc(), arg0, arg1, arg2);
     }
@@ -705,6 +706,9 @@ class RegionBuilderHelper {
   bool isInteger(Value value) {
     return llvm::isa<IntegerType>(value.getType());
   }
+  bool isIndex(Value value) {
+    return llvm::isa<IndexType>(value.getType());
+  }
 
   OpBuilder &builder;
   Block █
diff --git a/mlir/test/Dialect/Linalg/named-ops.mlir b/mlir/test/Dialect/Linalg/named-ops.mlir
index 1e356c8fb4e72..a84519f31d368 100644
--- a/mlir/test/Dialect/Linalg/named-ops.mlir
+++ b/mlir/test/Dialect/Linalg/named-ops.mlir
@@ -2718,6 +2718,19 @@ func.func @select_tensor(%arg0: tensor<4x8x16xi1>, %arg1: tensor<4x8x16xf32>, %a
   return %1 : tensor<4x8x16xf32>
 }
 
+// -----
+
+// GH#179046: Test linalg.select with index type operands.
+// CHECK-LABEL: func @select_index
+func.func @select_index(%arg0: tensor<4x8x16xindex>, %arg1: tensor<4x8x16xindex>, %arg2: tensor<4x8x16xindex>) -> tensor<4x8x16xindex> {
+  %0 = tensor.empty() : tensor<4x8x16xindex>
+  // CHECK: linalg.select
+  // CHECK-SAME: ins(%{{.+}}, %{{.+}}, %{{.+}} : tensor<4x8x16xindex>, tensor<4x8x16xindex>, tensor<4x8x16xindex>)
+  // CHECK-SAME: outs(%{{.+}} : tensor<4x8x16xindex>)
+  %1 = linalg.select ins(%arg0, %arg1, %arg2 : tensor<4x8x16xindex>, tensor<4x8x16xindex>, tensor<4x8x16xindex>) outs(%0: tensor<4x8x16xindex>) -> tensor<4x8x16xindex>
+  return %1 : tensor<4x8x16xindex>
+}
+
 //===----------------------------------------------------------------------===//
 // linalg.pack + linalg.unpack
 //===----------------------------------------------------------------------===//

``````````

</details>


https://github.com/llvm/llvm-project/pull/179056