Align rocmlir-tuning-driver benchmarking with Triton do_bench

mlir/include/mlir/Dialect/Rock/IR/AmdArchDb.h

@@ -119,6 +119,10 @@ AmdArchInfo lookupArchInfo(StringRef arch);
 bool isDirectToLDSSupported(GemmFeatures features);
 bool isGlobalPrefetchSupported(StringRef arch);
 bool isAsyncDirectToLDSSupported(StringRef arch);
+
+/// Get the size in bytes of the last-level cache for this architecture (the
+/// AMD Infinity Cache where present, otherwise the L2).
+int64_t getLastLevelCacheSize(StringRef arch);
 } // namespace rock
 } // namespace mlir
 

mlir/lib/Dialect/Rock/IR/AmdArchDb.cpp

@@ -414,6 +414,47 @@ AmdArchInfo mlir::rock::lookupArchInfo(StringRef arch) {
   llvm_unreachable(msg.c_str());
 }
 
+int64_t mlir::rock::getLastLevelCacheSize(StringRef arch) {
+  constexpr int64_t kMiB = 1024 * 1024;
+
+  // We cannot rely on hipDeviceProp_t::l2CacheSize for last-level sizing: it
+  // reports the small per-XCD L2 (~4 MiB on CDNA3/CDNA4), not the last-level
+  // AMD Infinity Cache that actually needs to be evicted between timed runs.
+  // Classify by chip the same way lookupArchInfo does, but distinguish the
+  // generations whose Infinity Cache differs (e.g. gfx101x vs gfx103x).
+  auto [chip, deviceId] = parseArchString(arch);
+  (void)deviceId;
+  StringRef minor = chip.take_back(2);
+  StringRef major = chip.slice(0, chip.size() - 2);
+
+  if (major == "gfx9") {
+    // CDNA3 (gfx942) / CDNA4 (gfx950) carry a large last-level Infinity Cache.
+    if (minor == "42" || minor == "50")
+      return 256 * kMiB;
+    // CDNA1 (gfx908) / CDNA2 (gfx90a) top out at a per-GCD L2.
+    if (minor == "08" || minor == "0a")
+      return 8 * kMiB;
+    // GCN5 / gfx906: L2 is the last level.
+    return 4 * kMiB;
+  }
+  if (major == "gfx10") {
+    // gfx103x (RDNA2) introduced the Infinity Cache; gfx101x (RDNA1) did not.
+    if (minor.starts_with("3"))
+      return 128 * kMiB;
+    return 4 * kMiB;
+  }
+  if (major == "gfx11") // RDNA3
+    return 96 * kMiB;
+  if (major == "gfx12") {
+    // gfx1250 assumed Infinity-Cache-class; TODO confirm once AMD publishes it.
+    if (minor == "50")
+      return 256 * kMiB;
+    return 64 * kMiB; // RDNA4
+  }
+  // Unknown arch: assume an Infinity-Cache-class last-level cache.
+  return 256 * kMiB;
+}
+
 GemmFeatures mlir::rock::AmdArchInfo::getDefaultFeatures(Type dataType) {
   GemmFeatures theseFeatures = defaultFeatures;
   bool isWmma = bitEnumContainsAll(theseFeatures, GemmFeatures::wmma);

mlir/tools/rocmlir-tuning-driver/CacheFlush.cpp

@@ -172,12 +172,10 @@ class CacheFlushState {
     }
   }
 
-  LogicalResult flushL2Cache(hipStream_t stream) {
+  LogicalResult flushCache(hipStream_t stream, bool useLastLevelCacheSize) {
     std::lock_guard<std::mutex> lock(stateMutex);
-    if (failed(allocL2CacheFlushBuffer()))
+    if (failed(allocCacheFlushBuffer(useLastLevelCacheSize)))
       return failure();
-    if (skipL2Flush)
-      return success();
     CHECK_HIP(hipMemsetAsync(flushBuffer.get(), 0, flushSize, stream));
     return success();
   }
@@ -205,7 +203,6 @@ class CacheFlushState {
       result = failure();
     }
     flushSize = 0;
-    skipL2Flush = false;
 #if defined(__HIP_PLATFORM_AMD__)
     if (failed(icacheKernel.cleanup()))
       result = failure();
@@ -216,17 +213,24 @@ class CacheFlushState {
   }
 
 private:
-  LogicalResult allocL2CacheFlushBuffer() {
-    if (flushBuffer || skipL2Flush)
-      return success();
-    size_t l2Size = deviceProps.l2CacheSize;
-    if (l2Size == 0) {
-      llvm::errs() << "Device '" << deviceProps.name
-                   << "' reported zero-sized L2 cache; skipping L2 flush.\n";
-      skipL2Flush = true;
+  LogicalResult allocCacheFlushBuffer(bool useLastLevelCacheSize) {
+    if (flushBuffer)
       return success();
+    if (useLastLevelCacheSize) {
+      // Size the flush buffer to the architecture's last-level cache. We cannot
+      // use hipDeviceProp_t::l2CacheSize because it only reports the small
+      // per-XCD L2 (~4 MiB on CDNA3/CDNA4), not the last-level AMD Infinity
+      // Cache that actually needs to be evicted between timed runs (256 MiB on
+      // MI300X/MI325X/MI350X). rock::getLastLevelCacheSize returns the right
+      // last-level size per arch (Infinity Cache where present, else L2).
+      flushSize = static_cast<size_t>(
+          rock::getLastLevelCacheSize(deviceProps.gcnArchName));
+    } else {
+      // Default: size the flush buffer to the L2 cache reported by the HIP
+      // runtime, plus a 20% margin.
+      size_t l2Size = static_cast<size_t>(deviceProps.l2CacheSize);
+      flushSize = l2Size + (l2Size / 5); // 20% margin
     }
-    flushSize = l2Size + (l2Size / 5); // 20% margin
     void *rawBuffer = nullptr;
     CHECK_HIP(hipMalloc(&rawBuffer, flushSize));
     flushBuffer.reset(rawBuffer);
@@ -265,7 +269,6 @@ class CacheFlushState {
   hipDeviceProp_t deviceProps = {};
   size_t flushSize = 0;
   HipDeviceBuffer flushBuffer;
-  bool skipL2Flush = false;
 #if defined(__HIP_PLATFORM_AMD__)
   static constexpr int32_t kDefaultWaveSize = 64;
   // https://github.com/ROCm/composable_kernel/blob/develop/include/ck_tile/host/flush_icache.hpp
@@ -308,8 +311,8 @@ CacheFlushState &getState() {
 
 } // namespace
 
-LogicalResult flushL2Cache(hipStream_t stream) {
-  return getState().flushL2Cache(stream);
+LogicalResult flushCache(hipStream_t stream, bool useLastLevelCacheSize) {
+  return getState().flushCache(stream, useLastLevelCacheSize);
 }
 
 LogicalResult flushInstructionCache(hipStream_t stream) {

mlir/tools/rocmlir-tuning-driver/CacheFlush.h

@@ -15,10 +15,14 @@
 
 namespace rocmlir::tuningdriver {
 
-/// \brief Flushes the L2 cache by performing a memory write operation.
+/// \brief Flushes the cache by performing a memory write operation.
 /// \param stream The HIP stream to use for the flush operation.
+/// \param useLastLevelCacheSize When true, size the flush buffer to the
+/// architecture's last-level cache (e.g. AMD Infinity Cache) instead of the
+/// per-XCD L2 cache size reported by the HIP runtime.
 /// \return success() if the flush succeeds, failure() otherwise.
-mlir::LogicalResult flushL2Cache(hipStream_t stream);
+mlir::LogicalResult flushCache(hipStream_t stream,
+                               bool useLastLevelCacheSize = false);
 
 /// \brief Flushes the instruction cache to ensure that any modified code is
 /// visible to the device.