Specialize AVX512 Payload for AMX extension

include/firestarter/Environment/X86/Payload/AVX512Payload.hpp

@@ -35,7 +35,7 @@ using __tilecfg = struct __tile_config {
 };
 
 /// This payload is designed for the AVX512 foundation CPU extension.
-class AVX512Payload final : public X86Payload {
+class AVX512Payload : public X86Payload {
 public:
   AVX512Payload() noexcept
       : X86Payload(/*FeatureRequests=*/{asmjit::CpuFeatures::X86::kAVX512_F}, /*Name=*/"AVX512", /*RegisterSize=*/8,
@@ -69,7 +69,7 @@ class AVX512Payload final : public X86Payload {
   /// \returns The compiled payload that provides access to the init and load functions.
   [[nodiscard]] auto compilePayload(const environment::payload::PayloadSettings& Settings, bool DumpRegisters,
                                     bool ErrorDetection) const
-      -> environment::payload::CompiledPayload::UniquePtr override;
+      -> environment::payload::CompiledPayload::UniquePtr final;
 
 private:
   static void create_AMX_config(__tilecfg* tileinfo);

include/firestarter/Environment/X86/Payload/AVX512WithAMXPayload.hpp

@@ -0,0 +1,39 @@
+/******************************************************************************
+ * FIRESTARTER - A Processor Stress Test Utility
+ * Copyright (C) 2024 TU Dresden, Center for Information Services and High
+ * Performance Computing
+ *
+ * This program is free software: you can redistribute it and/or modify
+ * it under the terms of the GNU General Public License as published by
+ * the Free Software Foundation, either version 3 of the License, or
+ * (at your option) any later version.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ * GNU General Public License for more details.
+ *
+ * You should have received a copy of the GNU General Public License
+ * along with this program.  If not, see <http://www.gnu.org/licenses/\>.
+ *
+ * Contact: [email protected]
+ *****************************************************************************/
+
+#pragma once
+
+#include "asmjit/core/cpuinfo.h"
+#include "firestarter/Environment/X86/Payload/AVX512Payload.hpp"
+
+namespace firestarter::environment::x86::payload {
+
+/// This payload is designed for the AVX512 foundation CPU extension specialized for AMX.
+class AVX512WithAMXPayload : public AVX512Payload {
+public:
+  AVX512WithAMXPayload() noexcept {
+    // Enable the AMX instruction in the AVX512 Payload and request AMX_TILE and AMX_BF16 feature.
+    addInstructionFlops("AMX", 512);
+    addFeatureRequest(asmjit::CpuFeatures::X86::kAMX_TILE);
+    addFeatureRequest(asmjit::CpuFeatures::X86::kAMX_BF16);
+  }
+};
+} // namespace firestarter::environment::x86::payload

include/firestarter/Environment/X86/Payload/X86Payload.hpp

@@ -90,6 +90,20 @@ class X86Payload : public environment::payload::Payload {
   };
 
 protected:
+  /// Add another instruction to the InstructionFlops map. This makes shure this instruction is shown in the set of
+  /// available instructions and we can correctly calculate the FLOPS for it. It is required when specializing Payloads
+  /// with new CPU extensions (e.g. AVX512 with newer CPU Extension like AMX).
+  /// \arg InstructionName The name of the instruction.
+  /// \arg Flops The Flops that are computed for this function.
+  void addInstructionFlops(const std::string& InstructionName, unsigned Flops) {
+    InstructionFlops[InstructionName] = Flops;
+  }
+
+  /// Add another feature request of this payload after it has been initialized. This can be used to specialize Payloads
+  /// (e.g. AVX512 with newer CPU Extension like AMX) and add more requested features as needed.
+  /// \arg Request The requested Cpu Feature.
+  void addFeatureRequest(asmjit::CpuFeatures::X86::Id Request) { FeatureRequests.push_back(Request); }
+
   /// Emit the code to dump the xmm, ymm or zmm registers into memory for the dump registers feature.
   /// \tparam Vec the type of the vector register used.
   /// \arg Cb The asmjit code builder that is used to emit the assembler code.

include/firestarter/Environment/X86/Platform/SapphireRapidsConfig.hpp

@@ -21,7 +21,7 @@
 
 #pragma once
 
-#include "firestarter/Environment/X86/Payload/AVX512Payload.hpp"
+#include "firestarter/Environment/X86/Payload/AVX512WithAMXPayload.hpp"
 #include "firestarter/Environment/X86/Platform/X86PlatformConfig.hpp"
 
 namespace firestarter::environment::x86::platform {
@@ -44,6 +44,6 @@ class SapphireRapidsConfig final : public X86PlatformConfig {
                                                                  {"L1_L", 40},
                                                                  {"REG", 140},
                                                                  {"AMX", 1}}),
-                          /*Payload=*/std::make_shared<const payload::AVX512Payload>()) {}
+                          /*Payload=*/std::make_shared<const payload::AVX512WithAMXPayload>()) {}
 };
 } // namespace firestarter::environment::x86::platform

src/firestarter/Environment/X86/Payload/AVX512Payload.cpp

@@ -62,16 +62,6 @@ auto AVX512Payload::compilePayload(const environment::payload::PayloadSettings&
   auto Repetitions =
       environment::payload::PayloadSettings::getNumberOfSequenceRepetitions(Sequence, Settings.linesPerThread());
 
-  // Check if AMX is in instruction mix and supported by CPU
-  // TODO: move this to a wrapper of AVX512Payload that injects the AMX instruction and the required AMX_BF16 feature
-  if (std::find(Sequence.begin(), Sequence.end(), "AMX") != Sequence.end()) {
-    if (this->supportedFeatures().x86().hasAMX_BF16()) {
-      workerLog::trace() << "AMX BF16 operations are supported by this processor.";
-    } else {
-      workerLog::error() << "AMX BF16 operations are not supported by this processor.";
-    }
-  }
-
   // compute count of flops and memory access for performance report
   environment::payload::PayloadStats Stats;