[Mlir-commits] [mlir] [mlir][SPIR-V] Support spirv.selection_control attribute on scf.if (PR #196510)

[Arseniy Obolenskiy]

https://github.com/aobolensk updated https://github.com/llvm/llvm-project/pull/196510

>From db2817def623bdc817e99905e9782fbddaa65e87 Mon Sep 17 00:00:00 2001
From: Arseniy Obolenskiy <arseniy.obolenskiy at amd.com>
Date: Fri, 8 May 2026 13:10:37 +0200
Subject: [PATCH 1/2] [mlir][SPIR-V] Support spirv.selection_control attribute
 on scf.if

---
 .../mlir/Dialect/SPIRV/IR/TargetAndABI.h      |  3 +++
 mlir/lib/Conversion/SCFToSPIRV/SCFToSPIRV.cpp |  8 +++++--
 mlir/lib/Dialect/SPIRV/IR/SPIRVDialect.cpp    |  4 ++++
 mlir/lib/Dialect/SPIRV/IR/TargetAndABI.cpp    |  4 ++++
 mlir/test/Conversion/SCFToSPIRV/if.mlir       | 22 +++++++++++++++++++
 5 files changed, 39 insertions(+), 2 deletions(-)

diff --git a/mlir/include/mlir/Dialect/SPIRV/IR/TargetAndABI.h b/mlir/include/mlir/Dialect/SPIRV/IR/TargetAndABI.h
index 7e11eb653c126..6e302542bde35 100644
--- a/mlir/include/mlir/Dialect/SPIRV/IR/TargetAndABI.h
+++ b/mlir/include/mlir/Dialect/SPIRV/IR/TargetAndABI.h
@@ -112,6 +112,9 @@ ResourceLimitsAttr getDefaultResourceLimits(MLIRContext *context);
 /// Returns the attribute name for specifying loop control.
 StringRef getLoopControlAttrName();
 
+/// Returns the attribute name for specifying selection control.
+StringRef getSelectionControlAttrName();
+
 /// Returns the attribute name for specifying SPIR-V target environment.
 StringRef getTargetEnvAttrName();
 
diff --git a/mlir/lib/Conversion/SCFToSPIRV/SCFToSPIRV.cpp b/mlir/lib/Conversion/SCFToSPIRV/SCFToSPIRV.cpp
index a9c6f7db847d3..d5140f3faa6ff 100644
--- a/mlir/lib/Conversion/SCFToSPIRV/SCFToSPIRV.cpp
+++ b/mlir/lib/Conversion/SCFToSPIRV/SCFToSPIRV.cpp
@@ -249,8 +249,12 @@ struct IfOpConversion : SCFToSPIRVPattern<scf::IfOp> {
 
     // Create `spirv.selection` operation, selection header block and merge
     // block.
-    auto selectionOp = spirv::SelectionOp::create(
-        rewriter, loc, spirv::SelectionControl::None);
+    auto selectionControl = spirv::SelectionControl::None;
+    if (auto attr = ifOp->getAttrOfType<spirv::SelectionControlAttr>(
+            spirv::getSelectionControlAttrName()))
+      selectionControl = attr.getValue();
+    auto selectionOp =
+        spirv::SelectionOp::create(rewriter, loc, selectionControl);
     auto *mergeBlock = rewriter.createBlock(&selectionOp.getBody(),
                                             selectionOp.getBody().end());
     spirv::MergeOp::create(rewriter, loc);
diff --git a/mlir/lib/Dialect/SPIRV/IR/SPIRVDialect.cpp b/mlir/lib/Dialect/SPIRV/IR/SPIRVDialect.cpp
index 2b5e7a571f42d..5821391b426cb 100644
--- a/mlir/lib/Dialect/SPIRV/IR/SPIRVDialect.cpp
+++ b/mlir/lib/Dialect/SPIRV/IR/SPIRVDialect.cpp
@@ -1054,6 +1054,10 @@ LogicalResult SPIRVDialect::verifyOperationAttribute(Operation *op,
     if (!isa<spirv::LoopControlAttr>(attr))
       return op->emitError("'")
              << symbol << "' must be a spirv::LoopControlAttr";
+  } else if (symbol == spirv::getSelectionControlAttrName()) {
+    if (!isa<spirv::SelectionControlAttr>(attr))
+      return op->emitError("'")
+             << symbol << "' must be a spirv::SelectionControlAttr";
   } else {
     return op->emitError("found unsupported '")
            << symbol << "' attribute on operation";
diff --git a/mlir/lib/Dialect/SPIRV/IR/TargetAndABI.cpp b/mlir/lib/Dialect/SPIRV/IR/TargetAndABI.cpp
index 270cb6df20415..c604fd087ba46 100644
--- a/mlir/lib/Dialect/SPIRV/IR/TargetAndABI.cpp
+++ b/mlir/lib/Dialect/SPIRV/IR/TargetAndABI.cpp
@@ -167,6 +167,10 @@ spirv::getDefaultResourceLimits(MLIRContext *context) {
 
 StringRef spirv::getLoopControlAttrName() { return "spirv.loop_control"; }
 
+StringRef spirv::getSelectionControlAttrName() {
+  return "spirv.selection_control";
+}
+
 StringRef spirv::getTargetEnvAttrName() { return "spirv.target_env"; }
 
 spirv::TargetEnvAttr spirv::getDefaultTargetEnv(MLIRContext *context) {
diff --git a/mlir/test/Conversion/SCFToSPIRV/if.mlir b/mlir/test/Conversion/SCFToSPIRV/if.mlir
index 2c18da41dc021..0b3df9a533302 100644
--- a/mlir/test/Conversion/SCFToSPIRV/if.mlir
+++ b/mlir/test/Conversion/SCFToSPIRV/if.mlir
@@ -167,4 +167,26 @@ func.func @unsupported_yield_type(%arg0 : memref<8xi32>, %arg1 : memref<8xi32>,
   return
 }
 
+// CHECK-LABEL: @selection_flatten
+func.func @selection_flatten(%arg2 : memref<10xf32, #spirv.storage_class<StorageBuffer>>, %arg3 : i1) {
+  %value = arith.constant 0.0 : f32
+  %i = arith.constant 0 : index
+  // CHECK: spirv.mlir.selection control(Flatten) {
+  scf.if %arg3 {
+    memref.store %value, %arg2[%i] : memref<10xf32, #spirv.storage_class<StorageBuffer>>
+  } {spirv.selection_control = #spirv.selection_control<Flatten>}
+  return
+}
+
+// CHECK-LABEL: @selection_dont_flatten
+func.func @selection_dont_flatten(%arg2 : memref<10xf32, #spirv.storage_class<StorageBuffer>>, %arg3 : i1) {
+  %value = arith.constant 0.0 : f32
+  %i = arith.constant 0 : index
+  // CHECK: spirv.mlir.selection control(DontFlatten) {
+  scf.if %arg3 {
+    memref.store %value, %arg2[%i] : memref<10xf32, #spirv.storage_class<StorageBuffer>>
+  } {spirv.selection_control = #spirv.selection_control<DontFlatten>}
+  return
+}
+
 } // end module

>From 63888bc31d678c24f0c26354be21e6716d1a1049 Mon Sep 17 00:00:00 2001
From: Arseniy Obolenskiy <arseniy.obolenskiy at amd.com>
Date: Sat, 9 May 2026 09:54:05 +0200
Subject: [PATCH 2/2] Add negative test

---
 mlir/test/Dialect/SPIRV/IR/target-and-abi.mlir | 11 +++++++++++
 1 file changed, 11 insertions(+)

diff --git a/mlir/test/Dialect/SPIRV/IR/target-and-abi.mlir b/mlir/test/Dialect/SPIRV/IR/target-and-abi.mlir
index 63dea6af83556..7ecb8fb363f76 100644
--- a/mlir/test/Dialect/SPIRV/IR/target-and-abi.mlir
+++ b/mlir/test/Dialect/SPIRV/IR/target-and-abi.mlir
@@ -346,3 +346,14 @@ func.func @vce() attributes {
   // CHECK: #spirv.vce<v1.0, [Shader], [SPV_KHR_storage_buffer_storage_class]>
   vce = #spirv.vce<v1.0, [Shader], [SPV_KHR_storage_buffer_storage_class]>
 } { return }
+
+// -----
+
+//===----------------------------------------------------------------------===//
+// spirv.selection_control
+//===----------------------------------------------------------------------===//
+
+// expected-error @+1 {{'spirv.selection_control' must be a spirv::SelectionControlAttr}}
+func.func @selection_control_wrong_attr() attributes {
+  spirv.selection_control = 64
+} { return }