[Mlir-commits] [mlir] [mlir][spirv] Fix spirv.Select min version requirement (PR #72173)

[Lei Zhang]

https://github.com/antiagainst created https://github.com/llvm/llvm-project/pull/72173

Per the spec, "Before version 1.4, results are only computed per component." So using scalar condition to select composite needs SPIR-V v1.4 at least.

>From be56a259ca8c31e3bea4eace1b6315c691c2267e Mon Sep 17 00:00:00 2001
From: Lei Zhang <antiagainst at gmail.com>
Date: Mon, 13 Nov 2023 15:46:58 -0800
Subject: [PATCH] [mlir][spirv] Fix spirv.Select min version requirement

Per the spec, "Before version 1.4, results are only computed per
component." So using scalar condition to select composite needs
SPIR-V v1.4 at least.
---
 .../mlir/Dialect/SPIRV/IR/SPIRVLogicalOps.td  |  4 ++++
 mlir/lib/Dialect/SPIRV/IR/ControlFlowOps.cpp  | 23 +++++++++++++++++++
 .../SPIRV/Transforms/vce-deduction.mlir       | 10 ++++++++
 3 files changed, 37 insertions(+)

diff --git a/mlir/include/mlir/Dialect/SPIRV/IR/SPIRVLogicalOps.td b/mlir/include/mlir/Dialect/SPIRV/IR/SPIRVLogicalOps.td
index d6e90724edc2733..cf38c15d20dc326 100644
--- a/mlir/include/mlir/Dialect/SPIRV/IR/SPIRVLogicalOps.td
+++ b/mlir/include/mlir/Dialect/SPIRV/IR/SPIRVLogicalOps.td
@@ -1002,6 +1002,10 @@ def SPIRV_SelectOp : SPIRV_Op<"Select",
   let assemblyFormat = [{
     operands attr-dict `:` type($condition) `,` type($result)
   }];
+
+  // These ops require dynamic availability specification based on operand and
+  // result types.
+  bit autogenAvailability = 0;
 }
 
 // -----
diff --git a/mlir/lib/Dialect/SPIRV/IR/ControlFlowOps.cpp b/mlir/lib/Dialect/SPIRV/IR/ControlFlowOps.cpp
index 081f8b601f41f07..580782043c81b46 100644
--- a/mlir/lib/Dialect/SPIRV/IR/ControlFlowOps.cpp
+++ b/mlir/lib/Dialect/SPIRV/IR/ControlFlowOps.cpp
@@ -10,7 +10,9 @@
 //
 //===----------------------------------------------------------------------===//
 
+#include "mlir/Dialect/SPIRV/IR/SPIRVEnums.h"
 #include "mlir/Dialect/SPIRV/IR/SPIRVOps.h"
+#include "mlir/Dialect/SPIRV/IR/SPIRVTypes.h"
 #include "mlir/Interfaces/CallInterfaces.h"
 
 #include "SPIRVOpUtils.h"
@@ -429,6 +431,27 @@ LogicalResult SelectOp::verify() {
   return success();
 }
 
+// Custom availability implementation is needed for spirv.Select given the
+// syntax changes starting v1.4.
+SmallVector<ArrayRef<spirv::Extension>, 1> SelectOp::getExtensions() {
+  return {};
+}
+SmallVector<ArrayRef<spirv::Capability>, 1> SelectOp::getCapabilities() {
+  return {};
+}
+std::optional<spirv::Version> SelectOp::getMinVersion() {
+  // Per the spec, "Before version 1.4, results are only computed per
+  // component."
+  if (isa<spirv::ScalarType>(getCondition().getType()) &&
+      isa<spirv::CompositeType>(getType()))
+    return Version::V_1_4;
+
+  return Version::V_1_0;
+}
+std::optional<spirv::Version> SelectOp::getMaxVersion() {
+  return Version::V_1_6;
+}
+
 //===----------------------------------------------------------------------===//
 // spirv.mlir.selection
 //===----------------------------------------------------------------------===//
diff --git a/mlir/test/Dialect/SPIRV/Transforms/vce-deduction.mlir b/mlir/test/Dialect/SPIRV/Transforms/vce-deduction.mlir
index 45e6eaa17b3ad69..4eaa21d2f94ef61 100644
--- a/mlir/test/Dialect/SPIRV/Transforms/vce-deduction.mlir
+++ b/mlir/test/Dialect/SPIRV/Transforms/vce-deduction.mlir
@@ -32,6 +32,16 @@ spirv.module Logical GLSL450 attributes {
   }
 }
 
+// CHECK: requires #spirv.vce<v1.4, [Shader], []>
+spirv.module Logical GLSL450 attributes {
+  spirv.target_env = #spirv.target_env<#spirv.vce<v1.6, [Shader], []>, #spirv.resource_limits<>>
+} {
+  spirv.func @select_with_scalar_condition(%predicate : i1, %a: vector<2xf32>, %b: vector<2xf32>) -> vector<2xf32> "None" {
+    %0 = spirv.Select %predicate, %a, %b : i1, vector<2xf32>
+    spirv.ReturnValue %0: vector<2xf32>
+  }
+}
+
 //===----------------------------------------------------------------------===//
 // Capability
 //===----------------------------------------------------------------------===//