[mlir][SPIR-V] Add lowering for gpu.lane_id op

[silee2]

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

[llvmbot]

@llvm/pr-subscribers-mlir
@llvm/pr-subscribers-mlir-gpu

@llvm/pr-subscribers-mlir-spirv

Author: Sang Ik Lee (silee2)

Changes

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

---

Full diff: https://github.com/llvm/llvm-project/pull/90873.diff

3 Files Affected:

• (modified) mlir/lib/Conversion/GPUToSPIRV/GPUToSPIRV.cpp (+2)
• (modified) mlir/lib/Dialect/SPIRV/Transforms/SPIRVConversion.cpp (+2-1)
• (modified) mlir/test/Conversion/GPUToSPIRV/builtins-opencl.mlir (+22)

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
+    }
+  }
+}