Add buffer cacheline size metric (pytorch#4228)

Summary:
Pull Request resolved: pytorch#4228

{F1753540374}

Differential Revision: D59649561

Author: Esteban Padilla Cerdio

backends/vulkan/tools/gpuinfo/glsl/buf_cacheline_size.glsl

@@ -0,0 +1,39 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ * All rights reserved.
+ *
+ * This source code is licensed under the BSD-style license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+#version 450 core
+
+#define PRECISION ${PRECISION}
+
+layout(std430) buffer;
+
+layout(set = 0, binding = 0) buffer PRECISION restrict readonly InBuffer {
+  float data[];
+}
+source;
+
+layout(set = 0, binding = 1) buffer PRECISION restrict writeonly OutBuffer {
+  float data[];
+}
+destination;
+
+layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in;
+
+layout(constant_id = 3) const int niter = 1;
+layout(constant_id = 4) const int stride = 1;
+layout(constant_id = 5) const int pitch = 1;
+
+
+void main() {
+  float c = 0;
+  for (int i = 0; i < niter; ++i) {
+    c += source.data[stride * 0 + pitch * gl_GlobalInvocationID[0]];
+    c += source.data[stride * 1 + pitch * gl_GlobalInvocationID[0]];
+  }
+  destination.data[0] = c;
+}

backends/vulkan/tools/gpuinfo/glsl/buf_cacheline_size.yaml

@@ -0,0 +1,12 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+buf_cacheline_size:
+  parameter_names_with_default_values:
+    DTYPE: float
+    STORAGE: buffer
+  shader_variants:
+    - NAME: buf_cacheline_size

backends/vulkan/tools/gpuinfo/src/app.cpp

@@ -44,6 +44,7 @@ class App {
     const uint32_t NREG_MAX = 512;
     const uint32_t NREG_STEP = 1;
 
+    // TODO: Make these values configurable
     const double COMPENSATE = 0.01;
     const double THRESHOLD = 3;
 
@@ -146,11 +147,78 @@ class App {
               << std::endl;
     std::cout << "Register type," << reg_ty << std::endl;
   }
+
+  void buf_cacheline_size() {
+    std::cout << "\n------ Buffer Cacheline Size ------" << std::endl;
+
+    // TODO: Make these values configurable
+    const double COMPENSATE = 0.01;
+    const double THRESHOLD = 2;
+
+    const uint32_t PITCH = buf_cache_size_ * 2 / nthread_logic_;
+    const uint32_t BUF_SIZE = PITCH * nthread_logic_;
+    const uint32_t MAX_STRIDE = PITCH / 2;
+
+    uint32_t NITER;
+
+    auto bench = [&](int stride) {
+      size_t len = sizeof(float);
+      StorageBuffer in_buf(context(), vkapi::kFloat, BUF_SIZE);
+      StorageBuffer out_buf(context(), vkapi::kFloat, len);
+      vkapi::PipelineBarrier pipeline_barrier{};
+
+      auto shader_name = "buf_cacheline_size";
+
+      uint32_t stride_div = stride / sizeof(float);
+      uint32_t pitch_div = PITCH / sizeof(float);
+
+      auto time = benchmark_on_gpu(shader_name, 100, [&]() {
+        context()->submit_compute_job(
+            VK_KERNEL_FROM_STR(shader_name),
+            pipeline_barrier,
+            {nthread_logic_, 1, 1},
+            {nthread_logic_, 1, 1},
+            {SV(NITER), SV(stride_div), SV(pitch_div)},
+            VK_NULL_HANDLE,
+            0,
+            in_buf.buffer(),
+            out_buf.buffer());
+      });
+      return time;
+    };
+
+    ensure_min_niter(1000, NITER, [&]() { return bench(sizeof(float)); });
+
+    uint32_t cacheline_size;
+
+    DtJumpFinder<5> dj(COMPENSATE, THRESHOLD);
+    uint32_t stride = sizeof(float);
+    for (; stride <= MAX_STRIDE; stride += sizeof(float)) {
+      double time = bench(stride);
+      std::cout << "Testing stride=" << stride << ", time=\t" << time
+                << std::endl;
+
+      if (dj.push(time)) {
+        cacheline_size = stride;
+        break;
+      }
+    }
+    if (stride >= MAX_STRIDE) {
+      std::cout
+          << "Reached end of memory, line size must be equal to global memory cache size available across all threads"
+          << std::endl;
+      cacheline_size = MAX_STRIDE;
+    }
+
+    std::cout << "BufTopLevelCachelineSize," << cacheline_size << std::endl;
+  }
 };
 
 int main(int argc, const char** argv) {
   App app;
 
-  app.reg_count();
+  // app.reg_count();
+  app.buf_cacheline_size();
+
   return 0;
 }