From 9c2acfe5b224d617e572806f73724d1e28c483e9 Mon Sep 17 00:00:00 2001
From: Bartosz Dunajski <bartosz.dunajski@intel.com>
Date: Wed, 24 Jul 2024 12:03:48 +0000
Subject: [PATCH] refactor: pass WG count to helper method

Signed-off-by: Bartosz Dunajski <bartosz.dunajski@intel.com>
---
 .../command_container/walker_partition_xehp_and_later.h   | 8 ++++----
 .../command_container/walker_partition_xehp_and_later.inl | 2 +-
 .../source/xe2_hpg_core/implicit_scaling_xe2_hpg_core.cpp | 2 +-
 .../source/xe_hpc_core/implicit_scaling_xe_hpc_core.cpp   | 2 +-
 .../source/xe_hpg_core/implicit_scaling_xe_hpg_core.cpp   | 2 +-
 5 files changed, 8 insertions(+), 8 deletions(-)

diff --git a/shared/source/command_container/walker_partition_xehp_and_later.h b/shared/source/command_container/walker_partition_xehp_and_later.h
index 9474c533b5..c88b09adc2 100644
--- a/shared/source/command_container/walker_partition_xehp_and_later.h
+++ b/shared/source/command_container/walker_partition_xehp_and_later.h
@@ -55,7 +55,7 @@ template <typename GfxFamily>
 using POST_SYNC_OPERATION = typename PIPE_CONTROL<GfxFamily>::POST_SYNC_OPERATION;
 
 template <typename GfxFamily, typename WalkerType>
-void appendWalkerFields(WalkerType &walkerCmd, uint32_t tileCount);
+void appendWalkerFields(WalkerType &walkerCmd, uint32_t tileCount, uint32_t workgroupCount);
 
 template <typename Command>
 Command *putCommand(void *&inputAddress, uint32_t &totalBytesProgrammed) {
@@ -509,7 +509,7 @@ void *programPartitionedWalker(void *&inputAddress, uint32_t &totalBytesProgramm
 
         inputWalker->setWorkloadPartitionEnable(true);
 
-        auto workgroupCount = 0u;
+        uint32_t workgroupCount = 0;
         if (partitionType == WalkerType::PARTITION_TYPE::PARTITION_TYPE_X) {
             workgroupCount = inputWalker->getThreadGroupIdXDimension();
         } else if (partitionType == WalkerType::PARTITION_TYPE::PARTITION_TYPE_Y) {
@@ -523,9 +523,9 @@ void *programPartitionedWalker(void *&inputAddress, uint32_t &totalBytesProgramm
         } else {
             inputWalker->setPartitionSize(Math::divideAndRoundUp(workgroupCount, partitionCount));
         }
-    }
 
-    appendWalkerFields<GfxFamily, WalkerType>(*inputWalker, tileCount);
+        appendWalkerFields<GfxFamily, WalkerType>(*inputWalker, tileCount, workgroupCount);
+    }
 
     *computeWalker = *inputWalker;
 
diff --git a/shared/source/command_container/walker_partition_xehp_and_later.inl b/shared/source/command_container/walker_partition_xehp_and_later.inl
index 138da6a461..387bcd0f02 100644
--- a/shared/source/command_container/walker_partition_xehp_and_later.inl
+++ b/shared/source/command_container/walker_partition_xehp_and_later.inl
@@ -9,6 +9,6 @@
 
 namespace WalkerPartition {
 template <typename GfxFamily, typename WalkerType>
-void appendWalkerFields(WalkerType &walkerCmd, uint32_t tileCount) {
+void appendWalkerFields(WalkerType &walkerCmd, uint32_t tileCount, uint32_t workgroupCount) {
 }
 } // namespace WalkerPartition
diff --git a/shared/source/xe2_hpg_core/implicit_scaling_xe2_hpg_core.cpp b/shared/source/xe2_hpg_core/implicit_scaling_xe2_hpg_core.cpp
index e32725fe5c..9921b34e78 100644
--- a/shared/source/xe2_hpg_core/implicit_scaling_xe2_hpg_core.cpp
+++ b/shared/source/xe2_hpg_core/implicit_scaling_xe2_hpg_core.cpp
@@ -26,4 +26,4 @@ template void ImplicitScalingDispatch<Family>::dispatchCommands<DefaultWalkerTyp
 template size_t ImplicitScalingDispatch<Family>::getSize<DefaultWalkerType>(bool apiSelfCleanup, bool preferStaticPartitioning, const DeviceBitfield &devices, const Vec3<size_t> &groupStart, const Vec3<size_t> &groupCount);
 } // namespace NEO
 
-template void WalkerPartition::appendWalkerFields<NEO::Family, NEO::DefaultWalkerType>(NEO::DefaultWalkerType &walkerCmd, uint32_t tileCount);
+template void WalkerPartition::appendWalkerFields<NEO::Family, NEO::DefaultWalkerType>(NEO::DefaultWalkerType &walkerCmd, uint32_t tileCount, uint32_t workgroupCount);
diff --git a/shared/source/xe_hpc_core/implicit_scaling_xe_hpc_core.cpp b/shared/source/xe_hpc_core/implicit_scaling_xe_hpc_core.cpp
index 05f4cc2168..f3f646e0cf 100644
--- a/shared/source/xe_hpc_core/implicit_scaling_xe_hpc_core.cpp
+++ b/shared/source/xe_hpc_core/implicit_scaling_xe_hpc_core.cpp
@@ -39,4 +39,4 @@ template void ImplicitScalingDispatch<Family>::dispatchCommands<DefaultWalkerTyp
 template size_t ImplicitScalingDispatch<Family>::getSize<DefaultWalkerType>(bool apiSelfCleanup, bool preferStaticPartitioning, const DeviceBitfield &devices, const Vec3<size_t> &groupStart, const Vec3<size_t> &groupCount);
 } // namespace NEO
 
-template void WalkerPartition::appendWalkerFields<NEO::Family, NEO::DefaultWalkerType>(NEO::DefaultWalkerType &walkerCmd, uint32_t tileCount);
+template void WalkerPartition::appendWalkerFields<NEO::Family, NEO::DefaultWalkerType>(NEO::DefaultWalkerType &walkerCmd, uint32_t tileCount, uint32_t workgroupCount);
diff --git a/shared/source/xe_hpg_core/implicit_scaling_xe_hpg_core.cpp b/shared/source/xe_hpg_core/implicit_scaling_xe_hpg_core.cpp
index 13e9c3210b..2c42370f1d 100644
--- a/shared/source/xe_hpg_core/implicit_scaling_xe_hpg_core.cpp
+++ b/shared/source/xe_hpg_core/implicit_scaling_xe_hpg_core.cpp
@@ -26,4 +26,4 @@ template void ImplicitScalingDispatch<Family>::dispatchCommands<DefaultWalkerTyp
 template size_t ImplicitScalingDispatch<Family>::getSize<DefaultWalkerType>(bool apiSelfCleanup, bool preferStaticPartitioning, const DeviceBitfield &devices, const Vec3<size_t> &groupStart, const Vec3<size_t> &groupCount);
 } // namespace NEO
 
-template void WalkerPartition::appendWalkerFields<NEO::Family, NEO::DefaultWalkerType>(NEO::DefaultWalkerType &walkerCmd, uint32_t tileCount);
+template void WalkerPartition::appendWalkerFields<NEO::Family, NEO::DefaultWalkerType>(NEO::DefaultWalkerType &walkerCmd, uint32_t tileCount, uint32_t workgroupCount);