From ac6472b105c39b0b8a605d22785202a6f924b30f Mon Sep 17 00:00:00 2001
From: Zbigniew Zdanowicz <zbigniew.zdanowicz@intel.com>
Date: Wed, 2 Oct 2024 09:50:19 +0000
Subject: [PATCH] fix: limit slm size based on hw informations

- add override default values method

Signed-off-by: Zbigniew Zdanowicz <zbigniew.zdanowicz@intel.com>
---
 shared/source/command_container/command_encoder.h      |  3 +++
 .../command_encoder_bdw_and_later.inl                  |  5 +++++
 .../command_container/command_encoder_enablers.inl     |  1 +
 .../command_encoder_xehp_and_later.inl                 | 10 +++++++++-
 4 files changed, 18 insertions(+), 1 deletion(-)

diff --git a/shared/source/command_container/command_encoder.h b/shared/source/command_container/command_encoder.h
index 4845be9094..c75582554f 100644
--- a/shared/source/command_container/command_encoder.h
+++ b/shared/source/command_container/command_encoder.h
@@ -226,6 +226,9 @@ struct EncodeDispatchKernel {
     static uint32_t computeSlmValues(const HardwareInfo &hwInfo, uint32_t slmSize);
 
     static bool singleTileExecImplicitScalingRequired(bool cooperativeKernel);
+
+    template <typename WalkerType, typename InterfaceDescriptorType>
+    static void overrideDefaultValues(WalkerType &walkerCmd, InterfaceDescriptorType &interfaceDescriptor);
 };
 
 template <typename GfxFamily>
diff --git a/shared/source/command_container/command_encoder_bdw_and_later.inl b/shared/source/command_container/command_encoder_bdw_and_later.inl
index 0efba9dbb3..f783de29b6 100644
--- a/shared/source/command_container/command_encoder_bdw_and_later.inl
+++ b/shared/source/command_container/command_encoder_bdw_and_later.inl
@@ -630,4 +630,9 @@ void InOrderPatchCommandHelpers::PatchCmd<Family>::patchComputeWalker(uint64_t a
     UNRECOVERABLE_IF(true);
 }
 
+template <typename Family>
+template <typename WalkerType, typename InterfaceDescriptorType>
+void EncodeDispatchKernel<Family>::overrideDefaultValues(WalkerType &walkerCmd, InterfaceDescriptorType &interfaceDescriptor) {
+}
+
 } // namespace NEO
diff --git a/shared/source/command_container/command_encoder_enablers.inl b/shared/source/command_container/command_encoder_enablers.inl
index f280f1ee6d..31d2e5686a 100644
--- a/shared/source/command_container/command_encoder_enablers.inl
+++ b/shared/source/command_container/command_encoder_enablers.inl
@@ -31,6 +31,7 @@ template void NEO::EncodeDispatchKernel<Family>::patchScratchAddressInImplicitAr
 template void NEO::EncodeDispatchKernel<Family>::forceComputeWalkerPostSyncFlushWithWrite<Family::DefaultWalkerType>(Family::DefaultWalkerType &walkerCmd);
 template void NEO::EncodeDispatchKernel<Family>::setWalkerRegionSettings<Family::DefaultWalkerType>(Family::DefaultWalkerType &walkerCmd, const HardwareInfo &hwInfo, uint32_t partitionCount,
                                                                                                     uint32_t workgroupSize, uint32_t maxWgCountPerTile, bool requiredWalkOrder);
+template void NEO::EncodeDispatchKernel<Family>::overrideDefaultValues<Family::DefaultWalkerType, Family::INTERFACE_DESCRIPTOR_DATA>(Family::DefaultWalkerType &walkerCmd, Family::INTERFACE_DESCRIPTOR_DATA &interfaceDescriptor);
 
 template struct NEO::EncodeStates<Family>;
 template struct NEO::EncodeMath<Family>;
diff --git a/shared/source/command_container/command_encoder_xehp_and_later.inl b/shared/source/command_container/command_encoder_xehp_and_later.inl
index 29a1d382c5..7c3531d8b6 100644
--- a/shared/source/command_container/command_encoder_xehp_and_later.inl
+++ b/shared/source/command_container/command_encoder_xehp_and_later.inl
@@ -412,6 +412,8 @@ void EncodeDispatchKernel<Family>::encode(CommandContainer &container, EncodeDis
         args.device->getDeviceInfo().maxFrontEndThreads};
     EncodeDispatchKernel<Family>::encodeAdditionalWalkerFields(rootDeviceEnvironment, walkerCmd, walkerArgs);
 
+    EncodeDispatchKernel<Family>::overrideDefaultValues(walkerCmd, idd);
+
     uint32_t workgroupSize = args.dispatchInterface->getGroupSize()[0] * args.dispatchInterface->getGroupSize()[1] * args.dispatchInterface->getGroupSize()[2];
     bool isRequiredWorkGroupOrder = args.requiredDispatchWalkOrder != NEO::RequiredDispatchWalkOrder::none;
     if (args.partitionCount > 1 && !args.isInternal) {
@@ -1076,7 +1078,7 @@ void EncodeDispatchKernel<Family>::setupPreferredSlmSize(InterfaceDescriptorType
         break;
     case SlmPolicy::slmPolicyLargeSlm:
     default:
-        slmSize = slmTotalSize * workGroupCountPerDss;
+        slmSize = std::min(slmTotalSize * workGroupCountPerDss, static_cast<uint32_t>(hwInfo.capabilityTable.slmSize * MemoryConstants::kiloByte));
         break;
     }
 
@@ -1111,4 +1113,10 @@ void InOrderPatchCommandHelpers::PatchCmd<Family>::patchComputeWalker(uint64_t a
     auto &postSync = walkerCmd->getPostSync();
     postSync.setImmediateData(baseCounterValue + appendCounterValue);
 }
+
+template <typename Family>
+template <typename WalkerType, typename InterfaceDescriptorType>
+void EncodeDispatchKernel<Family>::overrideDefaultValues(WalkerType &walkerCmd, InterfaceDescriptorType &interfaceDescriptor) {
+}
+
 } // namespace NEO