feat(ptodsl): implement simt_allreduce_sum for SIMT cross-workitem all-reduce

ptodsl/ptodsl/pto.py

@@ -143,6 +143,9 @@
     LoopHandle, BranchHandle,
 )
 
+# ── All-reduce ─────────────────────────────────────────────────────────────────
+from ._allreduce import simt_allreduce_sum, simt_allreduce_max, simt_allreduce_min  # noqa: F401
+
 # ── Decorator ─────────────────────────────────────────────────────────────────
 from ._jit import jit, KernelHandle, merge_jit_modules      # noqa: F401
 from ._subkernels import cube, simd, simt     # noqa: F401

test/vpto/cases/micro-op/simt/allreduce_cross_max/compare.py

@@ -0,0 +1,15 @@
+#!/usr/bin/env python3
+import sys, numpy as np
+
+def main():
+    golden = np.fromfile("golden_out.bin", dtype=np.float32)
+    out = np.fromfile("out.bin", dtype=np.float32)
+    if golden.shape != out.shape or not np.allclose(golden, out, rtol=1e-5, atol=1e-5):
+        mismatches = np.nonzero(~np.isclose(golden, out, rtol=1e-5, atol=1e-5))[0]
+        idx = int(mismatches[0]) if mismatches.size else 0
+        print(f"[ERROR] mismatch at idx={{idx}}, golden={{golden[idx]:.6f}}, out={{out[idx]:.6f}}")
+        sys.exit(2)
+    print("[INFO] compare passed")
+
+if __name__ == "__main__":
+    main()

test/vpto/cases/micro-op/simt/allreduce_cross_max/golden.py

@@ -0,0 +1,22 @@
+#!/usr/bin/env python3
+import argparse, numpy as np
+from pathlib import Path
+
+NLANES = 128
+EXPECTED = 1.0
+
+def generate(output_dir: Path) -> None:
+    output_dir.mkdir(parents=True, exist_ok=True)
+    out = np.zeros(NLANES, dtype=np.float32)
+    out.tofile(output_dir / "out.bin")
+    golden = np.full(NLANES, EXPECTED, dtype=np.float32)
+    golden.tofile(output_dir / "golden_out.bin")
+
+def main():
+    p = argparse.ArgumentParser()
+    p.add_argument("--output-dir", type=Path, default=Path("."))
+    a = p.parse_args()
+    generate(a.output_dir)
+
+if __name__ == "__main__":
+    main()

test/vpto/cases/micro-op/simt/allreduce_cross_max/kernel.pto

@@ -0,0 +1,65 @@
+module attributes {pto.kernel_kind = #pto.kernel_kind<vector>, pto.mode = "auto", pto.target_arch = "a5"} {
+  func.func @_kernel(%arg0: !pto.ptr<f32, gm>) attributes {pto.aicore} {
+    %c0_i64 = arith.constant 0 : i64
+    %0 = builtin.unrealized_conversion_cast %c0_i64 : i64 to ui64
+    %1 = pto.castptr %0 : ui64 -> !pto.ptr<f32, ub>
+    %c128_i32 = arith.constant 128 : i32
+    %c1_i32 = arith.constant 1 : i32
+    %c1_i32_0 = arith.constant 1 : i32
+    pto.simt_launch @_body__simt_0<<<%c128_i32, %c1_i32, %c1_i32_0>>>(%1, %arg0) : (!pto.ptr<f32, ub>, !pto.ptr<f32, gm>) -> ()
+    pto.barrier <PIPE_ALL>
+    return
+  }
+  func.func @_body__simt_0(%arg0: !pto.ptr<f32, ub>, %arg1: !pto.ptr<f32, gm>) attributes {pto.simt_entry} {
+    %0 = pto.get_tid_x : i32
+    %cst = arith.constant 1.000000e+00 : f32
+    %c5_i32 = arith.constant 5 : i32
+    %c31_i32 = arith.constant 31 : i32
+    %c32_i32 = arith.constant 32 : i32
+    %c1_i32 = arith.constant 1 : i32
+    %c4_i32 = arith.constant 4 : i32
+    %c0_i32 = arith.constant 0 : i32
+    %c0 = arith.constant 0 : index
+    %cst_0 = arith.constant 0xFF800000 : f32
+    %1 = pto.get_tid_x : i32
+    %2 = arith.shrui %1, %c5_i32 : i32
+    %3 = pto.get_laneid : i32
+    %4 = pto.redux_max %cst : f32 -> f32
+    %5 = arith.cmpi ult, %3, %c1_i32 : i32
+    scf.if %5 {
+      %13 = arith.muli %2, %c1_i32 : i32
+      %14 = arith.addi %13, %3 : i32
+      %15 = arith.index_cast %14 : i32 to index
+      pto.store %4, %arg0[%15] : !pto.ptr<f32, ub>, f32
+    }
+    pto.syncthreads
+    %6 = arith.cmpi ult, %1, %c32_i32 : i32
+    %7 = scf.if %6 -> (f32) {
+      %13 = arith.cmpi ult, %3, %c4_i32 : i32
+      %14 = scf.if %13 -> (f32) {
+        %16 = arith.index_cast %3 : i32 to index
+        %17 = pto.load %arg0[%16] : !pto.ptr<f32, ub> -> f32
+        scf.yield %17 : f32
+      } else {
+        scf.yield %cst_0 : f32
+      }
+      %15 = pto.redux_max %14 : f32 -> f32
+      scf.yield %15 : f32
+    } else {
+      scf.yield %cst_0 : f32
+    }
+    %8 = arith.cmpi ult, %1, %c1_i32 : i32
+    scf.if %8 {
+      %13 = arith.index_cast %1 : i32 to index
+      pto.store %7, %arg0[%13] : !pto.ptr<f32, ub>, f32
+    }
+    pto.syncthreads
+    %9 = arith.remui %1, %c1_i32 : i32
+    %10 = arith.index_cast %9 : i32 to index
+    %11 = pto.load %arg0[%10] : !pto.ptr<f32, ub> -> f32
+    pto.syncthreads
+    %12 = arith.index_cast %0 : i32 to index
+    pto.store %11, %arg1[%12] : !pto.ptr<f32, gm>, f32
+    return
+  }
+}

test/vpto/cases/micro-op/simt/allreduce_cross_max/launch.cpp

@@ -0,0 +1,11 @@
+#ifndef __VEC_SCOPE__
+#define __VEC_SCOPE__
+#endif
+#include <stdint.h>
+#ifndef __CPU_SIM
+#include "acl/acl.h"
+#endif
+extern "C" __global__ [aicore] void _kernel(__gm__ float *out);
+void Launch_kernel(float *out, void *stream) {
+  _kernel<<<1, nullptr, stream>>>((__gm__ float *)out);
+}

test/vpto/cases/micro-op/simt/allreduce_cross_max/main.cpp

@@ -0,0 +1,43 @@
+#include "test_common.h"
+#include "acl/acl.h"
+#include <cstdio>
+#include <cstdlib>
+#include <cstring>
+
+using namespace PtoTestCommon;
+
+#define ACL_CHECK(expr) do { const aclError _ret = (expr); if (_ret != ACL_SUCCESS) { std::fprintf(stderr, "[ERROR] %s failed: %d (%s:%d)\n", #expr, (int)_ret, __FILE__, __LINE__); rc = 1; goto cleanup; } } while (0)
+
+void Launch_kernel(float *out, void *stream);
+
+int main() {
+  size_t elemCount = 128;
+  size_t fileSize = elemCount * sizeof(float);
+  float *outHost = nullptr;
+  float *outDevice = nullptr;
+  int rc = 0;
+  bool aclInited = false, deviceSet = false;
+  int deviceId = 0;
+  aclrtStream stream = nullptr;
+
+  ACL_CHECK(aclInit(nullptr)); aclInited = true;
+  if (const char *e = std::getenv("ACL_DEVICE_ID")) deviceId = std::atoi(e);
+  ACL_CHECK(aclrtSetDevice(deviceId)); deviceSet = true;
+  ACL_CHECK(aclrtCreateStream(&stream));
+  ACL_CHECK(aclrtMallocHost((void **)(&outHost), fileSize));
+  ACL_CHECK(aclrtMalloc((void **)&outDevice, fileSize, ACL_MEM_MALLOC_HUGE_FIRST));
+  // Input: zero-initialize output buffer (kernel writes results)
+  std::memset(outHost, 0, fileSize);
+  ACL_CHECK(aclrtMemcpy(outDevice, fileSize, outHost, fileSize, ACL_MEMCPY_HOST_TO_DEVICE));
+  Launch_kernel(outDevice, stream);
+  ACL_CHECK(aclrtSynchronizeStream(stream));
+  ACL_CHECK(aclrtMemcpy(outHost, fileSize, outDevice, fileSize, ACL_MEMCPY_DEVICE_TO_HOST));
+  WriteFile("./out.bin", outHost, fileSize);
+
+cleanup:
+  aclrtFree(outDevice); aclrtFreeHost(outHost);
+  if (stream) aclrtDestroyStream(stream);
+  if (deviceSet) aclrtResetDevice(deviceId);
+  if (aclInited) aclFinalize();
+  return rc;
+}

test/vpto/cases/micro-op/simt/allreduce_cross_min/compare.py

@@ -0,0 +1,15 @@
+#!/usr/bin/env python3
+import sys, numpy as np
+
+def main():
+    golden = np.fromfile("golden_out.bin", dtype=np.float32)
+    out = np.fromfile("out.bin", dtype=np.float32)
+    if golden.shape != out.shape or not np.allclose(golden, out, rtol=1e-5, atol=1e-5):
+        mismatches = np.nonzero(~np.isclose(golden, out, rtol=1e-5, atol=1e-5))[0]
+        idx = int(mismatches[0]) if mismatches.size else 0
+        print(f"[ERROR] mismatch at idx={{idx}}, golden={{golden[idx]:.6f}}, out={{out[idx]:.6f}}")
+        sys.exit(2)
+    print("[INFO] compare passed")
+
+if __name__ == "__main__":
+    main()

test/vpto/cases/micro-op/simt/allreduce_cross_min/golden.py

@@ -0,0 +1,22 @@
+#!/usr/bin/env python3
+import argparse, numpy as np
+from pathlib import Path
+
+NLANES = 128
+EXPECTED = 1.0
+
+def generate(output_dir: Path) -> None:
+    output_dir.mkdir(parents=True, exist_ok=True)
+    out = np.zeros(NLANES, dtype=np.float32)
+    out.tofile(output_dir / "out.bin")
+    golden = np.full(NLANES, EXPECTED, dtype=np.float32)
+    golden.tofile(output_dir / "golden_out.bin")
+
+def main():
+    p = argparse.ArgumentParser()
+    p.add_argument("--output-dir", type=Path, default=Path("."))
+    a = p.parse_args()
+    generate(a.output_dir)
+
+if __name__ == "__main__":
+    main()

test/vpto/cases/micro-op/simt/allreduce_cross_min/kernel.pto

@@ -0,0 +1,65 @@
+module attributes {pto.kernel_kind = #pto.kernel_kind<vector>, pto.mode = "auto", pto.target_arch = "a5"} {
+  func.func @_kernel(%arg0: !pto.ptr<f32, gm>) attributes {pto.aicore} {
+    %c0_i64 = arith.constant 0 : i64
+    %0 = builtin.unrealized_conversion_cast %c0_i64 : i64 to ui64
+    %1 = pto.castptr %0 : ui64 -> !pto.ptr<f32, ub>
+    %c128_i32 = arith.constant 128 : i32
+    %c1_i32 = arith.constant 1 : i32
+    %c1_i32_0 = arith.constant 1 : i32
+    pto.simt_launch @_body__simt_0<<<%c128_i32, %c1_i32, %c1_i32_0>>>(%1, %arg0) : (!pto.ptr<f32, ub>, !pto.ptr<f32, gm>) -> ()
+    pto.barrier <PIPE_ALL>
+    return
+  }
+  func.func @_body__simt_0(%arg0: !pto.ptr<f32, ub>, %arg1: !pto.ptr<f32, gm>) attributes {pto.simt_entry} {
+    %0 = pto.get_tid_x : i32
+    %cst = arith.constant 1.000000e+00 : f32
+    %c5_i32 = arith.constant 5 : i32
+    %c31_i32 = arith.constant 31 : i32
+    %c32_i32 = arith.constant 32 : i32
+    %c1_i32 = arith.constant 1 : i32
+    %c4_i32 = arith.constant 4 : i32
+    %c0_i32 = arith.constant 0 : i32
+    %c0 = arith.constant 0 : index
+    %cst_0 = arith.constant 0x7F800000 : f32
+    %1 = pto.get_tid_x : i32
+    %2 = arith.shrui %1, %c5_i32 : i32
+    %3 = pto.get_laneid : i32
+    %4 = pto.redux_min %cst : f32 -> f32
+    %5 = arith.cmpi ult, %3, %c1_i32 : i32
+    scf.if %5 {
+      %13 = arith.muli %2, %c1_i32 : i32
+      %14 = arith.addi %13, %3 : i32
+      %15 = arith.index_cast %14 : i32 to index
+      pto.store %4, %arg0[%15] : !pto.ptr<f32, ub>, f32
+    }
+    pto.syncthreads
+    %6 = arith.cmpi ult, %1, %c32_i32 : i32
+    %7 = scf.if %6 -> (f32) {
+      %13 = arith.cmpi ult, %3, %c4_i32 : i32
+      %14 = scf.if %13 -> (f32) {
+        %16 = arith.index_cast %3 : i32 to index
+        %17 = pto.load %arg0[%16] : !pto.ptr<f32, ub> -> f32
+        scf.yield %17 : f32
+      } else {
+        scf.yield %cst_0 : f32
+      }
+      %15 = pto.redux_min %14 : f32 -> f32
+      scf.yield %15 : f32
+    } else {
+      scf.yield %cst_0 : f32
+    }
+    %8 = arith.cmpi ult, %1, %c1_i32 : i32
+    scf.if %8 {
+      %13 = arith.index_cast %1 : i32 to index
+      pto.store %7, %arg0[%13] : !pto.ptr<f32, ub>, f32
+    }
+    pto.syncthreads
+    %9 = arith.remui %1, %c1_i32 : i32
+    %10 = arith.index_cast %9 : i32 to index
+    %11 = pto.load %arg0[%10] : !pto.ptr<f32, ub> -> f32
+    pto.syncthreads
+    %12 = arith.index_cast %0 : i32 to index
+    pto.store %11, %arg1[%12] : !pto.ptr<f32, gm>, f32
+    return
+  }
+}

test/vpto/cases/micro-op/simt/allreduce_cross_min/launch.cpp

@@ -0,0 +1,11 @@
+#ifndef __VEC_SCOPE__
+#define __VEC_SCOPE__
+#endif
+#include <stdint.h>
+#ifndef __CPU_SIM
+#include "acl/acl.h"
+#endif
+extern "C" __global__ [aicore] void _kernel(__gm__ float *out);
+void Launch_kernel(float *out, void *stream) {
+  _kernel<<<1, nullptr, stream>>>((__gm__ float *)out);
+}

test/vpto/cases/micro-op/simt/allreduce_cross_min/main.cpp

@@ -0,0 +1,43 @@
+#include "test_common.h"
+#include "acl/acl.h"
+#include <cstdio>
+#include <cstdlib>
+#include <cstring>
+
+using namespace PtoTestCommon;
+
+#define ACL_CHECK(expr) do { const aclError _ret = (expr); if (_ret != ACL_SUCCESS) { std::fprintf(stderr, "[ERROR] %s failed: %d (%s:%d)\n", #expr, (int)_ret, __FILE__, __LINE__); rc = 1; goto cleanup; } } while (0)
+
+void Launch_kernel(float *out, void *stream);
+
+int main() {
+  size_t elemCount = 128;
+  size_t fileSize = elemCount * sizeof(float);
+  float *outHost = nullptr;
+  float *outDevice = nullptr;
+  int rc = 0;
+  bool aclInited = false, deviceSet = false;
+  int deviceId = 0;
+  aclrtStream stream = nullptr;
+
+  ACL_CHECK(aclInit(nullptr)); aclInited = true;
+  if (const char *e = std::getenv("ACL_DEVICE_ID")) deviceId = std::atoi(e);
+  ACL_CHECK(aclrtSetDevice(deviceId)); deviceSet = true;
+  ACL_CHECK(aclrtCreateStream(&stream));
+  ACL_CHECK(aclrtMallocHost((void **)(&outHost), fileSize));
+  ACL_CHECK(aclrtMalloc((void **)&outDevice, fileSize, ACL_MEM_MALLOC_HUGE_FIRST));
+  // Input: zero-initialize output buffer (kernel writes results)
+  std::memset(outHost, 0, fileSize);
+  ACL_CHECK(aclrtMemcpy(outDevice, fileSize, outHost, fileSize, ACL_MEMCPY_HOST_TO_DEVICE));
+  Launch_kernel(outDevice, stream);
+  ACL_CHECK(aclrtSynchronizeStream(stream));
+  ACL_CHECK(aclrtMemcpy(outHost, fileSize, outDevice, fileSize, ACL_MEMCPY_DEVICE_TO_HOST));
+  WriteFile("./out.bin", outHost, fileSize);
+
+cleanup:
+  aclrtFree(outDevice); aclrtFreeHost(outHost);
+  if (stream) aclrtDestroyStream(stream);
+  if (deviceSet) aclrtResetDevice(deviceId);
+  if (aclInited) aclFinalize();
+  return rc;
+}

test/vpto/cases/micro-op/simt/allreduce_cross_sum/compare.py

@@ -0,0 +1,15 @@
+#!/usr/bin/env python3
+import sys, numpy as np
+
+def main():
+    golden = np.fromfile("golden_out.bin", dtype=np.float32)
+    out = np.fromfile("out.bin", dtype=np.float32)
+    if golden.shape != out.shape or not np.allclose(golden, out, rtol=1e-5, atol=1e-5):
+        mismatches = np.nonzero(~np.isclose(golden, out, rtol=1e-5, atol=1e-5))[0]
+        idx = int(mismatches[0]) if mismatches.size else 0
+        print(f"[ERROR] mismatch at idx={{idx}}, golden={{golden[idx]:.6f}}, out={{out[idx]:.6f}}")
+        sys.exit(2)
+    print("[INFO] compare passed")
+
+if __name__ == "__main__":
+    main()

test/vpto/cases/micro-op/simt/allreduce_cross_sum/golden.py

@@ -0,0 +1,22 @@
+#!/usr/bin/env python3
+import argparse, numpy as np
+from pathlib import Path
+
+NLANES = 128
+EXPECTED = 128.0
+
+def generate(output_dir: Path) -> None:
+    output_dir.mkdir(parents=True, exist_ok=True)
+    out = np.zeros(NLANES, dtype=np.float32)
+    out.tofile(output_dir / "out.bin")
+    golden = np.full(NLANES, EXPECTED, dtype=np.float32)
+    golden.tofile(output_dir / "golden_out.bin")
+
+def main():
+    p = argparse.ArgumentParser()
+    p.add_argument("--output-dir", type=Path, default=Path("."))
+    a = p.parse_args()
+    generate(a.output_dir)
+
+if __name__ == "__main__":
+    main()