Add Intel GPU SPIR-V attributes (#913)

Adds Intel specific SPIR-V extensions enabling additional intrinsics on Intel GPU. The extra capabilities do not interfere with other generic Vulkan targets so, they can be added by default. --------- Co-authored-by: Renato Golin <rengolin@systemcall.eu>
plaidml · May 22, 2024 · 3302d2d · 3302d2d
1 parent 08caaf6
commit 3302d2d
Show file tree

Hide file tree

Showing 3 changed files with 15 additions and 4 deletions.
diff --git a/lib/TPP/GPU/SetSPIRVAbiAttribute.cpp b/lib/TPP/GPU/SetSPIRVAbiAttribute.cpp
@@ -42,13 +42,21 @@ struct SetSPIRVAbiAttribute
     auto gpuModule = getOperation();
     auto *context = &getContext();
     auto attrName = StringAttr::get(context, spirv::getEntryPointABIAttrName());
-    if (clientAPI == "opencl") {
+
+    bool isIntel = (clientAPI == "intel");
+
+    if ((clientAPI == "opencl") || isIntel) {
       auto abi = spirv::getEntryPointABIAttr(context);
       for (const auto &gpuFunc : gpuModule.getOps<gpu::GPUFuncOp>()) {
         if (!gpu::GPUDialect::isKernel(gpuFunc) || gpuFunc->getAttr(attrName))
           continue;
 
         gpuFunc->setAttr(attrName, abi);
+
+        if (isIntel) {
+          gpuFunc->setAttr("VectorComputeFunctionINTEL",
+                           UnitAttr::get(context));
+        }
       }
     } else if (clientAPI == "vulkan") {
       const SmallVector<gpu::Dimension> dims = {

diff --git a/lib/TPP/GPU/SetSPIRVCapabilities.cpp b/lib/TPP/GPU/SetSPIRVCapabilities.cpp
@@ -59,6 +59,8 @@ struct SetSPIRVCapabilities
         spirv::Capability::AtomicFloat32AddEXT,
         spirv::Capability::ExpectAssumeKHR,
         spirv::Capability::StorageBuffer16BitAccess,
+        spirv::Capability::VectorComputeINTEL,
+        spirv::Capability::VectorAnyINTEL,
         // clang-format on
     };
     spirv::Capability caps_vulkan[] = {
@@ -73,12 +75,13 @@ struct SetSPIRVCapabilities
         spirv::Extension::SPV_EXT_shader_atomic_float_add,
         spirv::Extension::SPV_KHR_expect_assume,
         spirv::Extension::SPV_KHR_16bit_storage,
-        spirv::Extension::SPV_NV_cooperative_matrix};
+        spirv::Extension::SPV_NV_cooperative_matrix,
+        spirv::Extension::SPV_INTEL_vector_compute};
     spirv::Extension exts_vulkan[] = {
         spirv::Extension::SPV_KHR_storage_buffer_storage_class,
         spirv::Extension::SPV_KHR_16bit_storage,
         spirv::Extension::SPV_NV_cooperative_matrix};
-    if (clientAPI == "opencl") {
+    if ((clientAPI == "opencl") || (clientAPI == "intel")) {
       auto triple = spirv::VerCapExtAttr::get(
           spirv::Version::V_1_4, caps_opencl, exts_opencl, context);
       auto attr = spirv::TargetEnvAttr::get(

diff --git a/test/GPU/Vulkan/set-spirv-capabilities.mlir b/test/GPU/Vulkan/set-spirv-capabilities.mlir
@@ -33,4 +33,4 @@ module attributes {gpu.container_module} {
 // VULKAN: gpu.module @entry_kernel attributes {spirv.target_env = #spirv.target_env<#spirv.vce<v1.4, [Shader, Float16, StorageBuffer16BitAccess], [SPV_KHR_storage_buffer_storage_class, SPV_KHR_16bit_storage, SPV_NV_cooperative_matrix]>, api=Vulkan, #spirv.resource_limits<>>} {
 
 // OPENCL: module attributes {gpu.container_module} {
-// OPENCL: gpu.module @entry_kernel attributes {spirv.target_env = #spirv.target_env<#spirv.vce<v1.4, [Addresses, Float16Buffer, Int64, Int16, Int8, Bfloat16ConversionINTEL, Kernel, Linkage, Vector16, GenericPointer, Groups, Float16, Float64, AtomicFloat32AddEXT, ExpectAssumeKHR, StorageBuffer16BitAccess], [SPV_INTEL_bfloat16_conversion, SPV_EXT_shader_atomic_float_add, SPV_KHR_expect_assume, SPV_KHR_16bit_storage, SPV_NV_cooperative_matrix]>, api=OpenCL, #spirv.resource_limits<>>} {
+// OPENCL: gpu.module @entry_kernel attributes {spirv.target_env = #spirv.target_env<#spirv.vce<v1.4, [Addresses, Float16Buffer, Int64, Int16, Int8, Bfloat16ConversionINTEL, Kernel, Linkage, Vector16, GenericPointer, Groups, Float16, Float64, AtomicFloat32AddEXT, ExpectAssumeKHR, StorageBuffer16BitAccess, VectorComputeINTEL, VectorAnyINTEL], [SPV_INTEL_bfloat16_conversion, SPV_EXT_shader_atomic_float_add, SPV_KHR_expect_assume, SPV_KHR_16bit_storage, SPV_NV_cooperative_matrix, SPV_INTEL_vector_compute]>, api=OpenCL, #spirv.resource_limits<>>} {