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

[llvmlistbot at llvm.org]

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

>From 6342d393feaf3a6573b19bf64a8af77da377e91b Mon Sep 17 00:00:00 2001
From: mugiwaraluffy56 <myakampuneeth at gmail.com>
Date: Sat, 31 Jan 2026 23:08:50 +0530
Subject: [PATCH] [mlir][linalg] Fix linalg.select crash with index type
 operands

The buildTernaryFn function in RegionBuilderHelper had overly restrictive
type checking that only allowed integer and floating point types, causing
an UNREACHABLE when linalg.select was used with index type operands.

This patch simplifies the logic by removing the type restrictions entirely.
The select operation only requires an i1 condition, and the values can be
any element type that tensors support. The arith::SelectOp verifier will
handle type validation.

Fixes #179046
---
 mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp |  9 ++-------
 mlir/test/Dialect/Linalg/named-ops.mlir  | 13 +++++++++++++
 2 files changed, 15 insertions(+), 7 deletions(-)

diff --git a/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp b/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp
index 0f0e308bba78e..d930721592562 100644
--- a/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp
+++ b/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp
@@ -612,17 +612,11 @@ class RegionBuilderHelper {
   // Build the ternary functions defined by OpDSL.
   Value buildTernaryFn(TernaryFn ternaryFn, Value arg0, Value arg1, Value arg2,
                        function_ref<InFlightDiagnostic()> emitError = {}) {
-    bool headBool =
-        isInteger(arg0) && arg0.getType().getIntOrFloatBitWidth() == 1;
-    bool tailFloatingPoint =
-        isFloatingPoint(arg0) && isFloatingPoint(arg1) && isFloatingPoint(arg2);
-    bool tailInteger = isInteger(arg0) && isInteger(arg1) && isInteger(arg2);
     OpBuilder::InsertionGuard g(builder);
     builder.setInsertionPointToEnd(&block);
     switch (ternaryFn) {
     case TernaryFn::select:
-      if (!headBool && !(tailFloatingPoint || tailInteger))
-        llvm_unreachable("unsupported non numeric type");
+      // select requires i1 condition; values can be any type tensors support
       return arith::SelectOp::create(builder, arg0.getLoc(), arg0, arg1, arg2);
     }
     if (emitError) {
@@ -705,6 +699,7 @@ 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..ee4128570d3c0 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 values (condition must be i1).
+// CHECK-LABEL: func @select_index
+func.func @select_index(%arg0: tensor<4x8x16xi1>, %arg1: tensor<4x8x16xindex>, %arg2: tensor<4x8x16xindex>) -> tensor<4x8x16xindex> {
+  %0 = tensor.empty() : tensor<4x8x16xindex>
+  // CHECK: linalg.select
+  // CHECK-SAME: ins(%{{.+}}, %{{.+}}, %{{.+}} : tensor<4x8x16xi1>, tensor<4x8x16xindex>, tensor<4x8x16xindex>)
+  // CHECK-SAME: outs(%{{.+}} : tensor<4x8x16xindex>)
+  %1 = linalg.select ins(%arg0, %arg1, %arg2 : tensor<4x8x16xi1>, tensor<4x8x16xindex>, tensor<4x8x16xindex>) outs(%0: tensor<4x8x16xindex>) -> tensor<4x8x16xindex>
+  return %1 : tensor<4x8x16xindex>
+}
+
 //===----------------------------------------------------------------------===//
 // linalg.pack + linalg.unpack
 //===----------------------------------------------------------------------===//