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

[llvmlistbot at llvm.org]

Author: Arseniy Obolenskiy
Date: 2026-05-09T21:44:43+02:00
New Revision: 492d774292daa9b22ebfe7e38c4fbe8fbab249d6

URL: https://github.com/llvm/llvm-project/commit/492d774292daa9b22ebfe7e38c4fbe8fbab249d6
DIFF: https://github.com/llvm/llvm-project/commit/492d774292daa9b22ebfe7e38c4fbe8fbab249d6.diff

LOG: [mlir][SPIR-V] Support spirv.selection_control attribute on scf.if (#196510)

Added: 
    

Modified: 
    mlir/include/mlir/Dialect/SPIRV/IR/TargetAndABI.h
    mlir/lib/Conversion/SCFToSPIRV/SCFToSPIRV.cpp
    mlir/lib/Dialect/SPIRV/IR/SPIRVDialect.cpp
    mlir/lib/Dialect/SPIRV/IR/TargetAndABI.cpp
    mlir/test/Conversion/SCFToSPIRV/if.mlir
    mlir/test/Dialect/SPIRV/IR/target-and-abi.mlir

Removed: 
    


################################################################################
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

diff  --git a/mlir/test/Dialect/SPIRV/IR/target-and-abi.mlir b/mlir/test/Dialect/SPIRV/IR/target-and-abi.mlir
index 63dea6af83556..e634186d3b9a8 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 at +1 {{'spirv.selection_control' must be a spirv::SelectionControlAttr}}
+func.func @selection_control_wrong_attr() attributes {
+  spirv.selection_control = 64
+} { return }