[mlir][spirv] Fix missing cases in UpdateVCE pass

[ThomasRaoux]

Summary

Builtin attached to global variable may have capabilities requirements. Some capabilities may also require minimum spirv version.

[antiagainst]

Good catch Thomas! The required capability for builtin is indeed not tightly wired up. We can synthesis SPIR-V operation's availability requirements, even requirements from using a particular enum attribute. But it does require the enum attribute to be statically declared on the op's ODS. But here for global variables the builtin decoration is not specified in global variable's ODS. This is actually a deeper problem to solve becuase I can see other ops can potentially have attributes that require additional capabilities... I need to file an issue to track this and solve it properly later.

For now to support global variable, I think this might not be the preferred way. Note that this is just one place we query the availability info; we also do similar things in SPIRVConversionTarget. So to support it, we should have the op report the capabilities required for the builtin decoration too. I think we can do the following:

1. Add an bit autogenAvailability = 1 field in SPIRV_Op in SPIRVBase.td to control whether we should auto-generate availablity for a particular SPIR-V op.
2. Modify emitAvailabilityImpl in SPIRVUtilsGen.cpp to return directly if autogenAvailability is false. So with this we can turn off the availablity autogen for global variables.
3. Implement the autogen'd availablity methods manually so you can inject the logic to report versions/capabilities properly there. (You can find lots of examples in the autogen'd SPIRVOpAvailabilityImpl.inc.)

This way all places will automatically get the proper support.

mlir/lib/Dialect/SPIRV/Transforms/UpdateVCEPass.cpp

WalkResult walkResult = module.walk([&](Operation *op) -> WalkResult {
 
     // Op capability requirements
     if (auto capabilities = dyn_cast<spirv::QueryCapabilityInterface>(op))
       if (failed(checkAndUpdateCapabilityRequirements(
               op, targetEnv, capabilities.getCapabilities(),
               deducedCapabilities)))
         return WalkResult::interrupt();
 
+    // Builtin attach to the op may have capability requirements.
+    auto builtInName = llvm::convertToSnakeFromCamelCase(
+        stringifyDecoration(spirv::Decoration::BuiltIn));
+    if (auto builtin = op->getAttrOfType<StringAttr>(builtInName)) {
+      if (auto varBuiltIn = spirv::symbolizeBuiltIn(builtin.getValue())) {
+        if (auto cap = spirv::getCapabilities(varBuiltIn.getValue())) {
+          if (failed(checkAndUpdateCapabilityRequirements(
+                  op, targetEnv,
+                  SmallVector<ArrayRef<spirv::Capability>, 1>(1,
+                                                              cap.getValue()),
+                  deducedCapabilities)))
+            return WalkResult::interrupt();
+        }
+      }
+    }
+
     SmallVector<Type, 4> valueTypes;
     valueTypes.append(op->operand_type_begin(), op->operand_type_end());
     valueTypes.append(op->result_type_begin(), op->result_type_end());
 
     // Special treatment for global variables, whose type requirements are
     // conveyed by type attributes.
     if (auto globalVar = dyn_cast<spirv::GlobalVariableOp>(op))
       valueTypes.push_back(globalVar.type());
for (Type valueType : valueTypes) {
 
       typeCapabilities.clear();
       valueType.cast<spirv::SPIRVType>().getCapabilities(typeCapabilities);
       if (failed(checkAndUpdateCapabilityRequirements(
               op, targetEnv, typeCapabilities, deducedCapabilities)))
         return WalkResult::interrupt();
     }
 
+    // Capabilities may require a minimum version as well.
+    for (auto capability : deducedCapabilities) {
+      if (auto capMinVersion = spirv::getMinVersion(capability)) {
+        deducedVersion = std::max(deducedVersion, capMinVersion.getValue());
+        if (deducedVersion > allowedVersion) {
+          return op->emitError("'") << op->getName() << "' requires "
+                                    << spirv::stringifyCapability(capability)
+                                    << " which requires min version "
+                                    << spirv::stringifyVersion(deducedVersion)
+                                    << " but target environment allows up to "
+                                    << spirv::stringifyVersion(allowedVersion);
+        }
+      }
+    }
+
     return WalkResult::advance();
   });
 
   if (walkResult.wasInterrupted())
     return signalPassFailure();
 
   // TODO(antiagainst): verify that the deduced version is consistent with
   // SPIR-V ops' maximal version requirements.

mlir/test/Dialect/SPIRV/Transforms/vce-deduction.mlir

spv.target_env = #spv.target_env<
      max_compute_workgroup_size = dense<[128, 128, 64]> : vector<3xi32>}>
 } {
   spv.func @fadd_function(%val : f16) -> f16 "None" {
     %0 = spv.FAdd %val, %val : f16
     spv.ReturnValue %0: f16
   }
 }
 
+// Test Builtin capabilities as well as capabilities min version.
+// NumSubgroups requires GroupNonUniform which requires v1.3.
+// CHECK: requires #spv.vce<v1.3, [GroupNonUniform, Shader], []>
+spv.module Logical GLSL450 attributes {
+  spv.target_env = #spv.target_env<
+    #spv.vce<v1.5, [Shader, UniformAndStorageBuffer8BitAccess, StorageBuffer16BitAccess, StorageUniform16, Int16, ImageBuffer, StorageImageExtendedFormats, GroupNonUniform], []>,
+    {max_compute_workgroup_invocations = 128 : i32,
+     max_compute_workgroup_size = dense<[128, 128, 64]> : vector<3xi32>}>
+} {
+   spv.globalVariable @__builtin_var_NumSubgroups__ built_in("NumSubgroups") : !spv.ptr<i32, Input>
+}
+
 //===----------------------------------------------------------------------===//
 // Extension
 //===----------------------------------------------------------------------===//
 
 // Test deducing minimal extensions.
 // spv.SubgroupBallotKHR requires the SPV_KHR_shader_ballot extension.
 
 // CHECK: requires #spv.vce<v1.0, [SubgroupBallotKHR, Shader], [SPV_KHR_shader_ballot]>
spv.func @subgroup_ballot(%predicate : i1) -> vector<4xi32> "None" {
     spv.ReturnValue %0: vector<4xi32>
   }
 }
 
 // Test deducing implied extension.
 // Vulkan memory model requires SPV_KHR_vulkan_memory_model, which is enabled
 // implicitly by v1.5.
 
-// CHECK: requires #spv.vce<v1.0, [VulkanMemoryModel], [SPV_KHR_vulkan_memory_model]>
+// CHECK: requires #spv.vce<v1.5, [VulkanMemoryModel], [SPV_KHR_vulkan_memory_model]>
 spv.module Logical Vulkan attributes {
   spv.target_env = #spv.target_env<
     #spv.vce<v1.5, [Shader, VulkanMemoryModel], []>,
     {max_compute_workgroup_invocations = 128 : i32,
      max_compute_workgroup_size = dense<[128, 128, 64]> : vector<3xi32>}>
 } {
   spv.func @iadd(%val : i32) -> i32 "None" {
     %0 = spv.IAdd %val, %val: i32