[Mlir-commits] [mlir] [mlir][SPIR-V] Add lowering for gpu.lane_id op (PR #90873)

[Sang Ik Lee]

https://github.com/silee2 created https://github.com/llvm/llvm-project/pull/90873

Add gpu.lane_id op lower for convert-gpu-to-spirv pass

>From e3c080db019042a6670a96006845e39a68d5a42a Mon Sep 17 00:00:00 2001
From: "Lee, Sang Ik" <sang.ik.lee at intel.com>
Date: Thu, 2 May 2024 16:28:42 +0000
Subject: [PATCH] [mlir][SPIR-V] Add lowering for gpu.lane_id op

---
 mlir/lib/Conversion/GPUToSPIRV/GPUToSPIRV.cpp |  2 ++
 .../SPIRV/Transforms/SPIRVConversion.cpp      |  3 ++-
 .../GPUToSPIRV/builtins-opencl.mlir           | 22 +++++++++++++++++++
 3 files changed, 26 insertions(+), 1 deletion(-)

diff --git a/mlir/lib/Conversion/GPUToSPIRV/GPUToSPIRV.cpp b/mlir/lib/Conversion/GPUToSPIRV/GPUToSPIRV.cpp
index d7885e0359592d..1560b3360577d3 100644
--- a/mlir/lib/Conversion/GPUToSPIRV/GPUToSPIRV.cpp
+++ b/mlir/lib/Conversion/GPUToSPIRV/GPUToSPIRV.cpp
@@ -629,6 +629,8 @@ void mlir::populateGPUToSPIRVPatterns(SPIRVTypeConverter &typeConverter,
                                       spirv::BuiltIn::NumSubgroups>,
       SingleDimLaunchConfigConversion<gpu::SubgroupSizeOp,
                                       spirv::BuiltIn::SubgroupSize>,
+      SingleDimLaunchConfigConversion<
+          gpu::LaneIdOp, spirv::BuiltIn::SubgroupLocalInvocationId>,
       WorkGroupSizeConversion, GPUAllReduceConversion,
       GPUSubgroupReduceConversion>(typeConverter, patterns.getContext());
 }
diff --git a/mlir/lib/Dialect/SPIRV/Transforms/SPIRVConversion.cpp b/mlir/lib/Dialect/SPIRV/Transforms/SPIRVConversion.cpp
index 4072608dc8f873..eba773d23773e6 100644
--- a/mlir/lib/Dialect/SPIRV/Transforms/SPIRVConversion.cpp
+++ b/mlir/lib/Dialect/SPIRV/Transforms/SPIRVConversion.cpp
@@ -867,7 +867,8 @@ getOrInsertBuiltinVariable(Block &body, Location loc, spirv::BuiltIn builtin,
   }
   case spirv::BuiltIn::SubgroupId:
   case spirv::BuiltIn::NumSubgroups:
-  case spirv::BuiltIn::SubgroupSize: {
+  case spirv::BuiltIn::SubgroupSize:
+  case spirv::BuiltIn::SubgroupLocalInvocationId: {
     auto ptrType =
         spirv::PointerType::get(integerType, spirv::StorageClass::Input);
     std::string name = getBuiltinVarName(builtin, prefix, suffix);
diff --git a/mlir/test/Conversion/GPUToSPIRV/builtins-opencl.mlir b/mlir/test/Conversion/GPUToSPIRV/builtins-opencl.mlir
index 8990d066e4e277..d4fe618b9df29c 100644
--- a/mlir/test/Conversion/GPUToSPIRV/builtins-opencl.mlir
+++ b/mlir/test/Conversion/GPUToSPIRV/builtins-opencl.mlir
@@ -50,3 +50,25 @@ module attributes {
     }
   }
 }
+
+// -----
+
+module attributes {
+  gpu.container_module,
+  spirv.target_env = #spirv.target_env<#spirv.vce<v1.0, [Addresses, Kernel, Int64], []>, #spirv.resource_limits<>>
+} {
+  // INDEX32-LABEL:  spirv.module @{{.*}} Physical32 OpenCL
+  // INDEX32: spirv.GlobalVariable [[LANEID:@.*]] built_in("SubgroupLocalInvocationId") : !spirv.ptr<i32, Input>
+  // INDEX64-LABEL:  spirv.module @{{.*}} Physical64 OpenCL
+  // INDEX64: spirv.GlobalVariable [[LANEID:@.*]] built_in("SubgroupLocalInvocationId") : !spirv.ptr<i32, Input>
+  gpu.module @kernels {
+    gpu.func @builtin_laneid() kernel
+      attributes {spirv.entry_point_abi = #spirv.entry_point_abi<workgroup_size = [16, 1, 1]>} {
+      // INDEX32: [[ADDRESS:%.*]] = spirv.mlir.addressof [[LANEID]]
+      // INDEX32-NEXT: {{%.*}} = spirv.Load "Input" [[ADDRESS]]
+      // INDEX64: spirv.UConvert %{{.+}} : i32 to i64
+      %0 = gpu.lane_id
+      gpu.return
+    }
+  }
+}