From 6bcc300e00cd3c03d57de233e9eaece61c6842b2 Mon Sep 17 00:00:00 2001
From: Karol Herbst <kherbst@redhat.com>
Date: Mon, 18 Dec 2023 23:06:11 +0100
Subject: [PATCH] rusticl: add RUSTICL_MAX_WORK_GROUPS

Part-of: <https://gitlab.freedesktop.org/mesa/mesa/-/merge_requests/27666>
---
 docs/envvars.rst                               | 5 +++++
 src/gallium/frontends/rusticl/core/device.rs   | 9 +++++++--
 src/gallium/frontends/rusticl/core/platform.rs | 7 +++++++
 3 files changed, 19 insertions(+), 2 deletions(-)

diff --git a/docs/envvars.rst b/docs/envvars.rst
index a68483ae352..1c3f19d88de 100644
--- a/docs/envvars.rst
+++ b/docs/envvars.rst
@@ -1085,6 +1085,11 @@ Rusticl environment variables
    - ``sync`` waits on the GPU to complete after every event
    - ``validate`` validates any internally generated SPIR-Vs, e.g. through compiling OpenCL C code
 
+.. envvar:: RUSTICL_MAX_WORK_GROUPS
+
+   Limits the amount of threads per dimension in a work-group. Useful for splitting up long running
+   tasks to increase responsiveness or to simulate the lowering of huge global sizes for testing.
+
 .. _clc-env-var:
 
 clc environment variables
diff --git a/src/gallium/frontends/rusticl/core/device.rs b/src/gallium/frontends/rusticl/core/device.rs
index e4662c25d14..0e5c0601171 100644
--- a/src/gallium/frontends/rusticl/core/device.rs
+++ b/src/gallium/frontends/rusticl/core/device.rs
@@ -893,8 +893,13 @@ impl Device {
     }
 
     pub fn max_grid_size(&self) -> Vec<u64> {
-        self.screen
-            .compute_param(pipe_compute_cap::PIPE_COMPUTE_CAP_MAX_GRID_SIZE)
+        let v: Vec<u64> = self
+            .screen
+            .compute_param(pipe_compute_cap::PIPE_COMPUTE_CAP_MAX_GRID_SIZE);
+
+        v.into_iter()
+            .map(|a| min(a, Platform::dbg().max_grid_size))
+            .collect()
     }
 
     pub fn max_clock_freq(&self) -> cl_uint {
diff --git a/src/gallium/frontends/rusticl/core/platform.rs b/src/gallium/frontends/rusticl/core/platform.rs
index 40f18f2ca51..305600a70cf 100644
--- a/src/gallium/frontends/rusticl/core/platform.rs
+++ b/src/gallium/frontends/rusticl/core/platform.rs
@@ -21,6 +21,7 @@ pub struct PlatformDebug {
     pub allow_invalid_spirv: bool,
     pub clc: bool,
     pub program: bool,
+    pub max_grid_size: u64,
     pub sync_every_event: bool,
     pub validate_spirv: bool,
 }
@@ -66,6 +67,7 @@ static mut PLATFORM_DBG: PlatformDebug = PlatformDebug {
     allow_invalid_spirv: false,
     clc: false,
     program: false,
+    max_grid_size: 0,
     sync_every_event: false,
     validate_spirv: false,
 };
@@ -91,6 +93,11 @@ fn load_env() {
         }
     }
 
+    debug.max_grid_size = env::var("RUSTICL_MAX_WORK_GROUPS")
+        .ok()
+        .and_then(|s| s.parse().ok())
+        .unwrap_or(u64::MAX);
+
     // SAFETY: no other references exist at this point
     let features = unsafe { &mut *addr_of_mut!(PLATFORM_FEATURES) };
     if let Ok(feature_flags) = env::var("RUSTICL_FEATURES") {