From 3302d2d534f3023d13abeea4b9a1d857687bc8ff Mon Sep 17 00:00:00 2001
From: Adam Siemieniuk <adam.siemieniuk@intel.com>
Date: Wed, 22 May 2024 17:22:06 +0200
Subject: [PATCH] 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>
---
 lib/TPP/GPU/SetSPIRVAbiAttribute.cpp        | 10 +++++++++-
 lib/TPP/GPU/SetSPIRVCapabilities.cpp        |  7 +++++--
 test/GPU/Vulkan/set-spirv-capabilities.mlir |  2 +-
 3 files changed, 15 insertions(+), 4 deletions(-)

diff --git a/lib/TPP/GPU/SetSPIRVAbiAttribute.cpp b/lib/TPP/GPU/SetSPIRVAbiAttribute.cpp
index 9d601f0f4..4908a99af 100644
--- 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
index 8730c11e4..843c0c080 100644
--- 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
index ff53af88d..8450d0711 100644
--- 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<>>} {