[Mlir-commits] [mlir] [MLIR][XeVM] Set the required sub-group size on every kernel (PR #232844)

[Md Abdullah Shahneous Bari]

https://github.com/mshahneo updated https://github.com/llvm/llvm-project/pull/232844

>From 955aab1baab6478eb57fd34ad718bc6abc48f309 Mon Sep 17 00:00:00 2001
From: "Shahneous Bari, Md Abdullah" <md.abdullah.shahneous.bari at intel.com>
Date: Fri, 16 Oct 2026 16:15:30 +0000
Subject: [PATCH 1/2] [MLIR][XeVM] Set the required sub-group size on every
 kernel

When it serializes a gpu.module to SPIR-V, the XeVM target sets
intel_reqd_sub_group_size = 16 on the first kernel only: the walk stops
after the first kernel it finds. The other kernels of the module get no
required size, so the GPU compiler may pick a different sub-group size
for them, for example 32. Code that the XeGPU and XeVM lowering emits
assumes 16 lanes, so those kernels can compute wrong results.

Set the required sub-group size on every kernel of the module.

Co-Authored-By: Claude Opus 5.5 <noreply at anthropic.com>
---
 mlir/lib/Target/LLVM/XeVM/Target.cpp          |  6 +----
 mlir/test/Dialect/GPU/lit.local.cfg           |  7 ++++++
 .../Dialect/GPU/module-to-binary-xevm.mlir    | 25 +++++++++++++++++++
 3 files changed, 33 insertions(+), 5 deletions(-)
 create mode 100644 mlir/test/Dialect/GPU/lit.local.cfg
 create mode 100644 mlir/test/Dialect/GPU/module-to-binary-xevm.mlir

diff --git a/mlir/lib/Target/LLVM/XeVM/Target.cpp b/mlir/lib/Target/LLVM/XeVM/Target.cpp
index 731ac4bd677f90..d271408302cd14 100644
--- a/mlir/lib/Target/LLVM/XeVM/Target.cpp
+++ b/mlir/lib/Target/LLVM/XeVM/Target.cpp
@@ -619,12 +619,8 @@ XeVMTargetAttrImpl::serializeToObject(Attribute attribute, Operation *module,
   auto xeTarget = cast<XeVMTargetAttr>(attribute);
   if (xeTarget.getTriple().starts_with("spirv")) {
     gpuMod.walk([&](LLVM::LLVMFuncOp funcOp) {
-      if (funcOp->hasDiscardableAttr(
-              gpu::GPUDialect::getKernelFuncAttrName())) {
+      if (funcOp->hasDiscardableAttr(gpu::GPUDialect::getKernelFuncAttrName()))
         funcOp.setIntelReqdSubGroupSize(16);
-        return WalkResult::interrupt();
-      }
-      return WalkResult::advance();
     });
 
     SPIRVSerializer serializer(*module, cast<XeVMTargetAttr>(attribute),
diff --git a/mlir/test/Dialect/GPU/lit.local.cfg b/mlir/test/Dialect/GPU/lit.local.cfg
new file mode 100644
index 00000000000000..df8e228268bf67
--- /dev/null
+++ b/mlir/test/Dialect/GPU/lit.local.cfg
@@ -0,0 +1,7 @@
+xevm_tests = [
+    "module-to-binary-xevm.mlir",
+]
+
+# Exclude the XeVM tests if the SPIRV target is disabled.
+if not config.run_xevm_tests:
+    config.excludes.update(xevm_tests)
diff --git a/mlir/test/Dialect/GPU/module-to-binary-xevm.mlir b/mlir/test/Dialect/GPU/module-to-binary-xevm.mlir
new file mode 100644
index 00000000000000..682fd409ec63d0
--- /dev/null
+++ b/mlir/test/Dialect/GPU/module-to-binary-xevm.mlir
@@ -0,0 +1,25 @@
+// REQUIRES: asserts
+// RUN: mlir-opt %s --gpu-module-to-binary="format=llvm" \
+// RUN:   -debug-only=serialize-to-llvm -o /dev/null 2>&1 | FileCheck %s
+
+// Every kernel of the module gets the required sub-group size, not only the
+// first one. Other functions get none.
+
+module attributes {gpu.container_module} {
+  // CHECK-LABEL: LLVM IR for module: "kernels"
+  gpu.module @kernels [#xevm.target] {
+    // CHECK: define spir_kernel void @kernel1() !intel_reqd_sub_group_size ![[SG:[0-9]+]] {
+    llvm.func spir_kernelcc @kernel1() attributes {gpu.kernel} {
+      llvm.return
+    }
+    // CHECK: define spir_kernel void @kernel2() !intel_reqd_sub_group_size ![[SG]] {
+    llvm.func spir_kernelcc @kernel2() attributes {gpu.kernel} {
+      llvm.return
+    }
+    // CHECK: define spir_func void @helper() {
+    llvm.func spir_funccc @helper() {
+      llvm.return
+    }
+  }
+  // CHECK: ![[SG]] = !{i32 16}
+}

>From e17c2694daebf7d1b7dcd7f1b2a5d6b2f66bc397 Mon Sep 17 00:00:00 2001
From: "Shahneous Bari, Md Abdullah" <md.abdullah.shahneous.bari at intel.com>
Date: Fri, 16 Oct 2026 17:50:12 +0000
Subject: [PATCH 2/2] [MLIR][XeVM] Gate the new test on a lit feature

Replace Dialect/GPU/lit.local.cfg with a `host-supports-spirv` lit
feature, set when the SPIRV target is built. This matches the NVVM and
ROCDL module-to-binary tests in the same directory, which use
`host-supports-nvptx` and `host-supports-amdgpu`.

Co-Authored-By: Claude Opus 5.5 <noreply at anthropic.com>
---
 mlir/test/Dialect/GPU/lit.local.cfg              | 7 -------
 mlir/test/Dialect/GPU/module-to-binary-xevm.mlir | 2 +-
 mlir/test/lit.cfg.py                             | 3 +++
 3 files changed, 4 insertions(+), 8 deletions(-)
 delete mode 100644 mlir/test/Dialect/GPU/lit.local.cfg

diff --git a/mlir/test/Dialect/GPU/lit.local.cfg b/mlir/test/Dialect/GPU/lit.local.cfg
deleted file mode 100644
index df8e228268bf67..00000000000000
--- a/mlir/test/Dialect/GPU/lit.local.cfg
+++ /dev/null
@@ -1,7 +0,0 @@
-xevm_tests = [
-    "module-to-binary-xevm.mlir",
-]
-
-# Exclude the XeVM tests if the SPIRV target is disabled.
-if not config.run_xevm_tests:
-    config.excludes.update(xevm_tests)
diff --git a/mlir/test/Dialect/GPU/module-to-binary-xevm.mlir b/mlir/test/Dialect/GPU/module-to-binary-xevm.mlir
index 682fd409ec63d0..3fc8bf22347ec6 100644
--- a/mlir/test/Dialect/GPU/module-to-binary-xevm.mlir
+++ b/mlir/test/Dialect/GPU/module-to-binary-xevm.mlir
@@ -1,4 +1,4 @@
-// REQUIRES: asserts
+// REQUIRES: asserts, host-supports-spirv
 // RUN: mlir-opt %s --gpu-module-to-binary="format=llvm" \
 // RUN:   -debug-only=serialize-to-llvm -o /dev/null 2>&1 | FileCheck %s
 
diff --git a/mlir/test/lit.cfg.py b/mlir/test/lit.cfg.py
index dbd18f8e8b0f2f..8b732213668a45 100644
--- a/mlir/test/lit.cfg.py
+++ b/mlir/test/lit.cfg.py
@@ -398,6 +398,9 @@ def have_host_jit_feature_support(feature_name):
 if config.run_rocm_tests:
     config.available_features.add("host-supports-amdgpu")
 
+if config.run_xevm_tests:
+    config.available_features.add("host-supports-spirv")
+
 if config.arm_emulator_executable:
     config.available_features.add("arm-emulator")