From b8ee8c7ebfdad68418899cfe43d77b3f735c5106 Mon Sep 17 00:00:00 2001 From: GKNB Date: Thu, 10 Jul 2025 19:05:24 +0000 Subject: [PATCH 01/46] Add features in wfMiniAPI for listing supported functions, add a separate timing example, add a stupid MD sim example, add a CUDA version of matmal for performance check --- wfMiniAPI/python/wfMiniAPI/kernel.py | 23 +++++- wfMiniAPI/python/wfMiniAPI/matmul.cu | 104 +++++++++++++++++++++++++ wfMiniAPI/python/wfMiniAPI/registry.py | 37 +++++++++ wfMiniAPI/python/wfMiniAPI/sim.py | 48 ++++++++++++ wfMiniAPI/python/wfMiniAPI/timing.py | 42 ++++++++++ 5 files changed, 252 insertions(+), 2 deletions(-) create mode 100644 wfMiniAPI/python/wfMiniAPI/matmul.cu create mode 100644 wfMiniAPI/python/wfMiniAPI/registry.py create mode 100644 wfMiniAPI/python/wfMiniAPI/sim.py create mode 100644 wfMiniAPI/python/wfMiniAPI/timing.py diff --git a/wfMiniAPI/python/wfMiniAPI/kernel.py b/wfMiniAPI/python/wfMiniAPI/kernel.py index 1d25be0..96f714b 100644 --- a/wfMiniAPI/python/wfMiniAPI/kernel.py +++ b/wfMiniAPI/python/wfMiniAPI/kernel.py @@ -3,6 +3,9 @@ import os import sys +from registry import annotate_kernel, list_kernels, kernel_params, run_kernel + + print("Python executable location:", sys.executable) print("NumPy version:", np.__version__) print("NumPy location:", np.__file__) @@ -31,9 +34,11 @@ #misc ################# +@annotate_kernel def sleep(seconds): time.sleep(seconds) +@annotate_kernel def get_device_module(device): if device == "gpu": if not CUPY_AVAILABLE: @@ -47,6 +52,7 @@ def get_device_module(device): #io ################# +@annotate_kernel def writeSingleRank(num_bytes, data_root_dir): if not MPI4PY_AVAILABLE: raise ImportError("mpi4py is not installed. Install mpi4py to use multi-process read/write.") @@ -65,7 +71,7 @@ def writeSingleRank(num_bytes, data_root_dir): with h5py.File(filename, 'w') as f: dset = f.create_dataset("data", data = data) - +@annotate_kernel def writeNonMPI(num_bytes, data_root_dir, filename_suffix=None): if not MPI4PY_AVAILABLE: raise ImportError("mpi4py is not installed. Install mpi4py to use multi-process read/write.") @@ -87,6 +93,7 @@ def writeNonMPI(num_bytes, data_root_dir, filename_suffix=None): with h5py.File(filename, 'w') as f: dset = f.create_dataset("data", data = data) +@annotate_kernel def writeWithMPI(num_bytes, data_root_dir, filename_suffix=None): if not MPI4PY_AVAILABLE: raise ImportError("mpi4py is not installed. Install mpi4py to use multi-process read/write.") @@ -112,6 +119,7 @@ def writeWithMPI(num_bytes, data_root_dir, filename_suffix=None): offset = rank * num_elem dset[offset:offset+num_elem] = data +@annotate_kernel def readNonMPI(num_bytes, data_root_dir, filename_suffix=None): if not MPI4PY_AVAILABLE: raise ImportError("mpi4py is not installed. Install mpi4py to use multi-process read/write.") @@ -132,6 +140,7 @@ def readNonMPI(num_bytes, data_root_dir, filename_suffix=None): with h5py.File(filename, 'r') as f: data = f['data'][0:num_elem] +@annotate_kernel def readWithMPI(num_bytes, data_root_dir, filename_suffix=None): if not MPI4PY_AVAILABLE: raise ImportError("mpi4py is not installed. Install mpi4py to use multi-process read/write.") @@ -162,6 +171,7 @@ def readWithMPI(num_bytes, data_root_dir, filename_suffix=None): #comm ################# +@annotate_kernel def MPIallReduce(device, data_size): xp = get_device_module(device) if not MPI4PY_AVAILABLE: @@ -183,6 +193,7 @@ def MPIallReduce(device, data_size): comm_nccl.allReduce(sendbuf.data.ptr, recvbuf.data.ptr, data_size, nccl.NCCL_FLOAT32, nccl.NCCL_SUM, cp.cuda.Stream.null) cp.cuda.Stream.null.synchronize() +@annotate_kernel def MPIallGather(device, data_size): xp = get_device_module(device) if not MPI4PY_AVAILABLE: @@ -209,6 +220,7 @@ def MPIallGather(device, data_size): #data movement ################# +@annotate_kernel def dataCopyH2D(data_size): if not CUPY_AVAILABLE: raise ImportError("CuPy is not installed. Install CuPy to use GPU capabilities.") @@ -216,6 +228,7 @@ def dataCopyH2D(data_size): data_h = np.empty(data_size, dtype=np.float32) data_d = cp.asarray(data_h) +@annotate_kernel def dataCopyD2H(data_size): if not CUPY_AVAILABLE: raise ImportError("CuPy is not installed. Install CuPy to use GPU capabilities.") @@ -228,18 +241,21 @@ def dataCopyD2H(data_size): #computation ################# +@annotate_kernel def matMulSimple2D(device, size): xp = get_device_module(device) matrix_a = xp.empty((size, size), dtype=xp.float32) matrix_b = xp.empty((size, size), dtype=xp.float32) matrix_c = xp.matmul(matrix_a, matrix_b) +@annotate_kernel def matMulGeneral(device, size_a, size_b, axis): xp = get_device_module(device) matrix_a = xp.empty(tuple(size_a), dtype=xp.float32) matrix_b = xp.empty(tuple(size_b), dtype=xp.float32) matrix_c = xp.tensordot(matrix_a, matrix_b, axis) +@annotate_kernel def fft(device, data_size, type_in, transform_dim): xp = get_device_module(device) if type_in == "float": @@ -255,13 +271,14 @@ def fft(device, data_size, type_in, transform_dim): out = xp.fft.fft(data_in, axis=transform_dim) - +@annotate_kernel def axpy(device, size): xp = get_device_module(device) x = xp.empty(size, dtype=xp.float32) y = xp.empty(size, dtype=xp.float32) y += 1.01 * x +@annotate_kernel def implaceCompute(device, size, num_op, op): xp = get_device_module(device) x = xp.empty(size, dtype=xp.float32) @@ -277,10 +294,12 @@ def implaceCompute(device, size, num_op, op): for _ in range(num_op): x = func(x) +@annotate_kernel def generateRandomNumber(device, size): xp = get_device_module(device) x = xp.random.rand(size) +@annotate_kernel def scatterAdd(device, x_size, y_size): xp = get_device_module(device) y = xp.empty(y_size, dtype=xp.float32) diff --git a/wfMiniAPI/python/wfMiniAPI/matmul.cu b/wfMiniAPI/python/wfMiniAPI/matmul.cu new file mode 100644 index 0000000..2d10079 --- /dev/null +++ b/wfMiniAPI/python/wfMiniAPI/matmul.cu @@ -0,0 +1,104 @@ +// matmul_cublas.cu +#include +#include +#include +#include + +#define CHECK_CUDA(call) \ + do { \ + cudaError_t err = (call); \ + if (err != cudaSuccess) { \ + std::cerr << "CUDA error at " << __FILE__ << ":" << __LINE__ \ + << " code=" << err << " \"" << cudaGetErrorString(err) \ + << "\"" << std::endl; \ + std::exit(EXIT_FAILURE); \ + } \ + } while (0) + +#define CHECK_CUBLAS(call) \ + do { \ + cublasStatus_t stat = (call); \ + if (stat != CUBLAS_STATUS_SUCCESS) { \ + std::cerr << "cuBLAS error at " << __FILE__ << ":" << __LINE__ \ + << " code=" << stat << std::endl; \ + std::exit(EXIT_FAILURE); \ + } \ + } while (0) + +int main(int argc, char* argv[]) { + if (argc != 3) { + std::cerr << "Usage: " << argv[0] << " \n"; + return EXIT_FAILURE; + } + + int dev = std::atoi(argv[1]); + int N = std::atoi(argv[2]); + + // Select GPU + CHECK_CUDA(cudaSetDevice(dev)); + + // cuBLAS handle + cublasHandle_t handle; + CHECK_CUBLAS(cublasCreate(&handle)); + + // Allocate device matrices A, B, C (uninitialized, like cupy.empty) + size_t bytes = size_t(N) * N * sizeof(float); + float *d_A, *d_B, *d_C; + CHECK_CUDA(cudaMalloc(&d_A, bytes)); + CHECK_CUDA(cudaMalloc(&d_B, bytes)); + CHECK_CUDA(cudaMalloc(&d_C, bytes)); + + // Warm-up: one matmul to initialize kernels / cache + const float alpha = 1.0f, beta = 0.0f; + CHECK_CUBLAS( + cublasSgemm(handle, + CUBLAS_OP_N, CUBLAS_OP_N, + N, N, N, + &alpha, + d_A, N, + d_B, N, + &beta, + d_C, N) + ); + // Ensure warm-up completed + CHECK_CUDA(cudaDeviceSynchronize()); + + // Create CUDA events for timing + cudaEvent_t start, stop; + CHECK_CUDA(cudaEventCreate(&start)); + CHECK_CUDA(cudaEventCreate(&stop)); + + // Record start, run 20 matmuls, record stop + CHECK_CUDA(cudaEventRecord(start)); + for (int i = 0; i < 20; ++i) { + CHECK_CUBLAS( + cublasSgemm(handle, + CUBLAS_OP_N, CUBLAS_OP_N, + N, N, N, + &alpha, + d_A, N, + d_B, N, + &beta, + d_C, N) + ); + } + CHECK_CUDA(cudaEventRecord(stop)); + CHECK_CUDA(cudaEventSynchronize(stop)); + + float ms = 0; + CHECK_CUDA(cudaEventElapsedTime(&ms, start, stop)); + + std::cout << "Matrix size: " << N << "×" << N << "\n" + << "Total time for 20 calls: " << ms << " ms\n" + << "Average per call: " << (ms / 20.0f) << " ms\n"; + + // Cleanup + cudaEventDestroy(start); + cudaEventDestroy(stop); + cudaFree(d_A); + cudaFree(d_B); + cudaFree(d_C); + cublasDestroy(handle); + + return 0; +} diff --git a/wfMiniAPI/python/wfMiniAPI/registry.py b/wfMiniAPI/python/wfMiniAPI/registry.py new file mode 100644 index 0000000..b0ca5ef --- /dev/null +++ b/wfMiniAPI/python/wfMiniAPI/registry.py @@ -0,0 +1,37 @@ +import inspect +from collections import OrderedDict + +_KERNELS = OrderedDict() + +class KernelSpec: + def __init__(self, func): + self.func = func + self.name = func.__name__ + self.sig = inspect.signature(func) + self.doc = inspect.getdoc(func) or "" + + def params(self): + out = {} + for name, param in self.sig.parameters.items(): + default = param.default if param.default is not inspect._empty else None + ann = param.annotation if param.annotation is not inspect._empty else None + out[name] = { + 'default': default, + 'annotation': ann, + } + return out + +def annotate_kernel(func): + _KERNELS[func.__name__] = KernelSpec(func) + return func + +def list_kernels(): + return list(_KERNELS) + +def kernel_params(name): + return _KERNELS[name].params() + +def run_kernel(name, **kwargs): + func = _KERNELS[name].func + out = func(**kwargs) + return out diff --git a/wfMiniAPI/python/wfMiniAPI/sim.py b/wfMiniAPI/python/wfMiniAPI/sim.py new file mode 100644 index 0000000..d7db36e --- /dev/null +++ b/wfMiniAPI/python/wfMiniAPI/sim.py @@ -0,0 +1,48 @@ +def triad_kernel(N_atoms, avg_neighbors, flops_per_pair, alpha=1.0000001, beta=0.0000001): + """ + A batched AXPY/Triad on a 2D array: + arr <- arr * alpha + beta + Parameters: + - N_atoms: number of "particles" + - avg_neighbors: length of each vector slice + - flops_per_pair: number of times to repeat the triad + """ + arr = xp.ones((N_atoms, avg_neighbors), dtype=xp.float32) + for _ in range(flops_per_pair): + arr = arr * alpha + beta + return arr + +def reduction_kernel(arr): + """ + Global reduction (sum) over all elements. + """ + return arr.sum() + +def run_miniapp(nsteps, + N_atoms=10000, + avg_neighbors=50, + flops_per_pair=20, + build_freq=1): + """ + Run the mini-app for `nsteps`, calling the two kernels each step. + Tunable parameters: + - N_atoms, avg_neighbors, flops_per_pair control the work per kernel. + - build_freq controls how often triad_kernel is called (can model neighbor-list rebuild frequency). + """ + for step in range(nsteps): + # Optionally skip triad kernel on some steps + if step % build_freq == 0: + arr = triad_kernel(N_atoms, avg_neighbors, flops_per_pair) + # Always do reduction to mimic energy sum + energy = reduction_kernel(arr) + if step % max(nsteps // 5, 1) == 0: + print(f"Step {step}: energy={energy:.3e}") + # Return final energy (and free GPU memory if used) + if gpu_enabled: + xp.get_default_memory_pool().free_all_blocks() + return energy + +if __name__ == "__main__": + # Example: 10 steps, 5k atoms, 60 neighbors, 30 flops per pair, triad every step + run_miniapp(nsteps=10, N_atoms=5000, avg_neighbors=60, flops_per_pair=30, build_freq=1) + diff --git a/wfMiniAPI/python/wfMiniAPI/timing.py b/wfMiniAPI/python/wfMiniAPI/timing.py new file mode 100644 index 0000000..6dd41c1 --- /dev/null +++ b/wfMiniAPI/python/wfMiniAPI/timing.py @@ -0,0 +1,42 @@ +import kernel +import time + + +print(kernel.get_device_module("cpu")) +print(kernel.get_device_module("gpu")) + +print(kernel.list_kernels()) +print(kernel.kernel_params("matMulGeneral")) +print(kernel.kernel_params("writeWithMPI")) + +n_repeat = 20 +kernel.run_kernel("matMulSimple2D", device="cpu", size=8192) +t0 = time.time() +for _ in range(n_repeat): + kernel.matMulSimple2D(device="cpu", size=8192) +# kernel.run_kernel("matMulSimple2D", device="cpu", size=8192) +print("took", time.time() - t0) +total_ms = (time.time() - t0) * 1000 +avg_ms = total_ms / n_repeat +print(f"CPU: Total time for {n_repeat} runs: {total_ms} ms") +print(f"CPU: Average per run: {avg_ms} ms") + + +import cupy as cp +kernel.matMulSimple2D(device="gpu", size=8192) +cp.cuda.Stream.null.synchronize() + +start = cp.cuda.Event() +end = cp.cuda.Event() +start.record() + +for _ in range(n_repeat): + kernel.matMulSimple2D(device="gpu", size=8192) +# kernel.run_kernel("matMulSimple2D", device="gpu", size=8192) +end.record() +end.synchronize() + +total_ms = cp.cuda.get_elapsed_time(start, end) +avg_ms = total_ms / n_repeat +print(f"GPU: Total time for {n_repeat} runs: {total_ms} ms") +print(f"GPU: Average per run: {avg_ms} ms") From b991d112ab080f02bb33865a0b50cfc470aa4457 Mon Sep 17 00:00:00 2001 From: GKNB Date: Thu, 24 Jul 2025 13:43:18 +0000 Subject: [PATCH 02/46] Fix installation --- wfMiniAPI/MANIFEST.in | 3 + wfMiniAPI/{README => README.md} | 0 wfMiniAPI/VERSION | 1 + wfMiniAPI/requirements.txt | 4 + wfMiniAPI/setup.py | 132 +++++++++++++++--- wfMiniAPI/{python => src}/wfMiniAPI/kernel.py | 0 wfMiniAPI/{python => src}/wfMiniAPI/matmul.cu | 0 .../{python => src}/wfMiniAPI/registry.py | 0 wfMiniAPI/{python => src}/wfMiniAPI/sim.py | 0 wfMiniAPI/{python => src}/wfMiniAPI/timing.py | 0 10 files changed, 120 insertions(+), 20 deletions(-) create mode 100644 wfMiniAPI/MANIFEST.in rename wfMiniAPI/{README => README.md} (100%) create mode 100644 wfMiniAPI/VERSION create mode 100644 wfMiniAPI/requirements.txt rename wfMiniAPI/{python => src}/wfMiniAPI/kernel.py (100%) rename wfMiniAPI/{python => src}/wfMiniAPI/matmul.cu (100%) rename wfMiniAPI/{python => src}/wfMiniAPI/registry.py (100%) rename wfMiniAPI/{python => src}/wfMiniAPI/sim.py (100%) rename wfMiniAPI/{python => src}/wfMiniAPI/timing.py (100%) diff --git a/wfMiniAPI/MANIFEST.in b/wfMiniAPI/MANIFEST.in new file mode 100644 index 0000000..00be847 --- /dev/null +++ b/wfMiniAPI/MANIFEST.in @@ -0,0 +1,3 @@ + +include *.py *.json *.md *.sh *.txt +include MANIFEST.in README.md VERSION diff --git a/wfMiniAPI/README b/wfMiniAPI/README.md similarity index 100% rename from wfMiniAPI/README rename to wfMiniAPI/README.md diff --git a/wfMiniAPI/VERSION b/wfMiniAPI/VERSION new file mode 100644 index 0000000..49d5957 --- /dev/null +++ b/wfMiniAPI/VERSION @@ -0,0 +1 @@ +0.1 diff --git a/wfMiniAPI/requirements.txt b/wfMiniAPI/requirements.txt new file mode 100644 index 0000000..9bda694 --- /dev/null +++ b/wfMiniAPI/requirements.txt @@ -0,0 +1,4 @@ + +numpy +h5py +mpi4py diff --git a/wfMiniAPI/setup.py b/wfMiniAPI/setup.py index 147e75c..bf3f66b 100644 --- a/wfMiniAPI/setup.py +++ b/wfMiniAPI/setup.py @@ -1,21 +1,113 @@ -from setuptools import setup, find_packages - -setup( - name='wfminiAPI', - version='0.1', - packages=find_packages(), - install_requires=[ - 'numpy', - 'h5py', - 'mpi4py' - ], - extras_require={ - 'gpu': ['cupy'] +#!/usr/bin/env python3 + +__author__ = 'RADICAL-Cybertools Team, Tianle Wang, Ozgur Kilic' +__email__ = 'info@radical-cybertools.org' +__copyright__ = 'Copyright 2022-25, The RADICAL-Cybertools Team' +__license__ = 'MIT' + + +''' Setup script, only usable via pip. ''' + +import os + +import subprocess as sp + +from glob import glob +from setuptools import setup, Command, find_packages + + +# ------------------------------------------------------------------------------ +# +repo = 'workflow-mini-apps' +name = 'wfMiniAPI' +mod_root = 'src/%s/' % name + +root = os.path.dirname(__file__) or '.' +readme = open('%s/README.md' % root, encoding='utf-8').read() +descr = 'An open source library that is used to make implementing emulated' \ + ' task in workflow mini-app simple. It support both Python and C++' \ + ' (OpenMP) backend and is targetting various different' \ + ' architecture including CPU, NVIDIA GPU, AMD GPU and Intel GPU' +keywords = ['radical', 'cybertools', 'mini-app'] + + +# ------------------------------------------------------------------------------ +# get version info +version = open('%s/VERSION' % root).read().strip() + + +# ------------------------------------------------------------------------------ +# +class RunTwine(Command): + user_options = [] + def initialize_options(self): pass + def finalize_options(self): pass + def run(self): + _, _, _ret = sh_callout('python3 setup.py sdist upload -r pypi') + raise SystemExit(_ret) + + +# ------------------------------------------------------------------------------ +# +with open('%s/requirements.txt' % root, encoding='utf-8') as freq: + requirements = freq.readlines() + + + +# ------------------------------------------------------------------------------ +# +setup_args = { + 'name' : name, + 'version' : version, + 'description' : descr, + 'long_description' : readme, + 'long_description_content_type' : 'text/markdown', + 'author' : __author__, + 'author_email' : __email__, + 'maintainer' : 'The RADICAL Group', + 'maintainer_email' : 'radical@rutgers.edu', + 'url' : 'http://radical-cybertools.github.io/%s/' % repo, + 'project_urls' : { + 'Documentation': 'https://%s.readthedocs.io/en/latest/' % name, + 'Source' : 'https://github.com/radical-cybertools/%s/' % repo, + 'Issues' : 'https://github.com/radical-cybertools/%s/issues' % repo, }, - author='Tianle Wang, Ozgur Kilic', - author_email='twang3@bnl.gov', - description='An open source library that is used to make implementing emulated task in workflow mini-app simple. It support both Python and C++ (OpenMP) backend and is targetting various different types of task', - license='MIT', - keywords='example keywords', - url='', -) + 'license' : 'MIT', + 'keywords' : keywords, + 'python_requires' : '>=3.8', + 'classifiers' : [ + 'Development Status :: 5 - Production/Stable', + 'Intended Audience :: Developers', + 'Environment :: Console', + 'License :: OSI Approved :: MIT License', + 'Programming Language :: Python', + 'Programming Language :: Python :: 3', + 'Programming Language :: Python :: 3.8', + 'Topic :: Utilities', + 'Topic :: System :: Distributed Computing', + 'Topic :: Scientific/Engineering', + 'Operating System :: POSIX', + 'Operating System :: Unix' + ], + 'packages' : find_packages('src'), + 'package_dir' : {'': 'src'}, + 'package_data' : {'': ['*.txt', '*.sh', '*.json', '*.gz', '*.c', + '*.md', 'VERSION']}, + 'install_requires' : requirements, + 'extras_require' : {'gpu': ['cupy-cuda12x']}, + 'zip_safe' : False, + 'cmdclass' : {'upload': RunTwine}, +} + + +# ------------------------------------------------------------------------------ +# +setup(**setup_args) + + +# ------------------------------------------------------------------------------ +# clean temporary files from source tree +os.system('rm -vrf src/%s.egg-info' % name) + + +# ------------------------------------------------------------------------------ diff --git a/wfMiniAPI/python/wfMiniAPI/kernel.py b/wfMiniAPI/src/wfMiniAPI/kernel.py similarity index 100% rename from wfMiniAPI/python/wfMiniAPI/kernel.py rename to wfMiniAPI/src/wfMiniAPI/kernel.py diff --git a/wfMiniAPI/python/wfMiniAPI/matmul.cu b/wfMiniAPI/src/wfMiniAPI/matmul.cu similarity index 100% rename from wfMiniAPI/python/wfMiniAPI/matmul.cu rename to wfMiniAPI/src/wfMiniAPI/matmul.cu diff --git a/wfMiniAPI/python/wfMiniAPI/registry.py b/wfMiniAPI/src/wfMiniAPI/registry.py similarity index 100% rename from wfMiniAPI/python/wfMiniAPI/registry.py rename to wfMiniAPI/src/wfMiniAPI/registry.py diff --git a/wfMiniAPI/python/wfMiniAPI/sim.py b/wfMiniAPI/src/wfMiniAPI/sim.py similarity index 100% rename from wfMiniAPI/python/wfMiniAPI/sim.py rename to wfMiniAPI/src/wfMiniAPI/sim.py diff --git a/wfMiniAPI/python/wfMiniAPI/timing.py b/wfMiniAPI/src/wfMiniAPI/timing.py similarity index 100% rename from wfMiniAPI/python/wfMiniAPI/timing.py rename to wfMiniAPI/src/wfMiniAPI/timing.py From 2cd77994af0118bb615cb53979c1aed8ff544777 Mon Sep 17 00:00:00 2001 From: GKNB Date: Sun, 27 Jul 2025 16:50:56 +0000 Subject: [PATCH 03/46] Add dependency for rose --- wfMiniAPI/requirements-gpu.txt | 6 ++++++ wfMiniAPI/requirements.txt | 3 ++- 2 files changed, 8 insertions(+), 1 deletion(-) create mode 100644 wfMiniAPI/requirements-gpu.txt diff --git a/wfMiniAPI/requirements-gpu.txt b/wfMiniAPI/requirements-gpu.txt new file mode 100644 index 0000000..ce00423 --- /dev/null +++ b/wfMiniAPI/requirements-gpu.txt @@ -0,0 +1,6 @@ +numpy +h5py +mpi4py +cupy +radical-pilot +git+https://github.com/radical-cybertools/ROSE/tree/main diff --git a/wfMiniAPI/requirements.txt b/wfMiniAPI/requirements.txt index 9bda694..041476a 100644 --- a/wfMiniAPI/requirements.txt +++ b/wfMiniAPI/requirements.txt @@ -1,4 +1,5 @@ - numpy h5py mpi4py +radical-pilot +git+https://github.com/radical-cybertools/ROSE/tree/main From aa9525925fa277333440d7bd0c6715b25f30e86a Mon Sep 17 00:00:00 2001 From: GKNB Date: Sun, 27 Jul 2025 16:53:36 +0000 Subject: [PATCH 04/46] Fix requirement --- wfMiniAPI/requirements-gpu.txt | 2 +- wfMiniAPI/requirements.txt | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/wfMiniAPI/requirements-gpu.txt b/wfMiniAPI/requirements-gpu.txt index ce00423..3c8be1b 100644 --- a/wfMiniAPI/requirements-gpu.txt +++ b/wfMiniAPI/requirements-gpu.txt @@ -3,4 +3,4 @@ h5py mpi4py cupy radical-pilot -git+https://github.com/radical-cybertools/ROSE/tree/main +git+https://github.com/radical-cybertools/ROSE.git diff --git a/wfMiniAPI/requirements.txt b/wfMiniAPI/requirements.txt index 041476a..aefbf98 100644 --- a/wfMiniAPI/requirements.txt +++ b/wfMiniAPI/requirements.txt @@ -2,4 +2,4 @@ numpy h5py mpi4py radical-pilot -git+https://github.com/radical-cybertools/ROSE/tree/main +git+https://github.com/radical-cybertools/ROSE.git From 5f4855a16608af0c422b452eebdd9836244a1a71 Mon Sep 17 00:00:00 2001 From: GKNB Date: Mon, 28 Jul 2025 02:35:54 +0000 Subject: [PATCH 05/46] Add timing --- wfMiniAPI/requirements-gpu.txt | 2 +- wfMiniAPI/requirements.txt | 2 +- wfMiniAPI/src/wfMiniAPI/kernel.py | 18 ++----------- wfMiniAPI/src/wfMiniAPI/registry.py | 39 +++++++++++++++++++++++++++++ 4 files changed, 43 insertions(+), 18 deletions(-) diff --git a/wfMiniAPI/requirements-gpu.txt b/wfMiniAPI/requirements-gpu.txt index 3c8be1b..c2348ec 100644 --- a/wfMiniAPI/requirements-gpu.txt +++ b/wfMiniAPI/requirements-gpu.txt @@ -3,4 +3,4 @@ h5py mpi4py cupy radical-pilot -git+https://github.com/radical-cybertools/ROSE.git +ROSE @ git+https://github.com/radical-cybertools/ROSE.git diff --git a/wfMiniAPI/requirements.txt b/wfMiniAPI/requirements.txt index aefbf98..4b84da6 100644 --- a/wfMiniAPI/requirements.txt +++ b/wfMiniAPI/requirements.txt @@ -2,4 +2,4 @@ numpy h5py mpi4py radical-pilot -git+https://github.com/radical-cybertools/ROSE.git +ROSE @ git+https://github.com/radical-cybertools/ROSE.git diff --git a/wfMiniAPI/src/wfMiniAPI/kernel.py b/wfMiniAPI/src/wfMiniAPI/kernel.py index 96f714b..2d9bcc0 100644 --- a/wfMiniAPI/src/wfMiniAPI/kernel.py +++ b/wfMiniAPI/src/wfMiniAPI/kernel.py @@ -3,7 +3,7 @@ import os import sys -from registry import annotate_kernel, list_kernels, kernel_params, run_kernel +from .registry import annotate_kernel, list_kernels, kernel_params, run_kernel print("Python executable location:", sys.executable) @@ -308,18 +308,4 @@ def scatterAdd(device, x_size, y_size): if xp == np: y += x[idx] elif xp == cp: - scatter_add_kernel = cp.RawKernel(r''' - extern "C" __global__ - void my_scatter_add_kernel(const float *x, const float *y, const int *idx) - { - int tid = blockDim.x * blockIdx.x + threadIdx.x; - - } - ''', 'my_scatter_add_kernel') - - - -#for the tutorial, three things: -#exalearn (CPU + GPU v1), ddmd v1, how to build wk-miniapp -#show installation script + run script, in installation script, show how to install assuming we are working in a brand new env (container for example) - + cp.add.at(y, idx, x) diff --git a/wfMiniAPI/src/wfMiniAPI/registry.py b/wfMiniAPI/src/wfMiniAPI/registry.py index b0ca5ef..c008116 100644 --- a/wfMiniAPI/src/wfMiniAPI/registry.py +++ b/wfMiniAPI/src/wfMiniAPI/registry.py @@ -35,3 +35,42 @@ def run_kernel(name, **kwargs): func = _KERNELS[name].func out = func(**kwargs) return out + +def time_kernel(name, device, n_warmup=1, n_repeat=20, **kwargs): + xp = get_device_module(device) + + # CPU timing + if device.lower() == "cpu": + for _ in range(n_warmup): + run_kernel(name, device=device, **kwargs) + t0 = time.time() + for _ in range(n_repeat): + run_kernel(name, device=device, **kwargs) + t1 = time.time() + + total_ms = (t1 - t0) * 1000 + avg_ms = total_ms / n_repeat + print(f"CPU: Total {n_repeat} runs: {total_ms:.2f} ms") + print(f"CPU: Avg per run: {avg_ms:.2f} ms") + return total_ms, avg_ms + + # GPU timing + elif device.lower() == "gpu": + # warmup + sync + for _ in range(n_warmup): + run_kernel(name, device=device, **kwargs) + cp.cuda.Stream.null.synchronize() + + start = cp.cuda.Event() + end = cp.cuda.Event() + start.record() + for _ in range(n_repeat): + run_kernel(name, device=device, **kwargs) + end.record() + end.synchronize() + + total_ms = cp.cuda.get_elapsed_time(start, end) + avg_ms = total_ms / n_repeat + print(f"GPU: Total {n_repeat} runs: {total_ms:.2f} ms") + print(f"GPU: Avg per run: {avg_ms:.2f} ms") + return total_ms, avg_ms From de2be37feadcec8d4498c321139ce32b35ee6690 Mon Sep 17 00:00:00 2001 From: GKNB Date: Mon, 28 Jul 2025 02:46:52 +0000 Subject: [PATCH 06/46] timing --- wfMiniAPI/src/wfMiniAPI/kernel.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/wfMiniAPI/src/wfMiniAPI/kernel.py b/wfMiniAPI/src/wfMiniAPI/kernel.py index 2d9bcc0..875012c 100644 --- a/wfMiniAPI/src/wfMiniAPI/kernel.py +++ b/wfMiniAPI/src/wfMiniAPI/kernel.py @@ -3,7 +3,7 @@ import os import sys -from .registry import annotate_kernel, list_kernels, kernel_params, run_kernel +from .registry import annotate_kernel, list_kernels, kernel_params, run_kernel, time_kernel print("Python executable location:", sys.executable) From 10278e2a2b2837336465483d8e4b39b224714b46 Mon Sep 17 00:00:00 2001 From: GKNB Date: Mon, 28 Jul 2025 02:53:21 +0000 Subject: [PATCH 07/46] fix --- wfMiniAPI/src/wfMiniAPI/registry.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/wfMiniAPI/src/wfMiniAPI/registry.py b/wfMiniAPI/src/wfMiniAPI/registry.py index c008116..b0b8519 100644 --- a/wfMiniAPI/src/wfMiniAPI/registry.py +++ b/wfMiniAPI/src/wfMiniAPI/registry.py @@ -37,8 +37,6 @@ def run_kernel(name, **kwargs): return out def time_kernel(name, device, n_warmup=1, n_repeat=20, **kwargs): - xp = get_device_module(device) - # CPU timing if device.lower() == "cpu": for _ in range(n_warmup): From bbf244da00b931ca86ea9825fd90f5010f2ab237 Mon Sep 17 00:00:00 2001 From: GKNB Date: Mon, 28 Jul 2025 03:32:31 +0000 Subject: [PATCH 08/46] fix time --- wfMiniAPI/src/wfMiniAPI/registry.py | 1 + 1 file changed, 1 insertion(+) diff --git a/wfMiniAPI/src/wfMiniAPI/registry.py b/wfMiniAPI/src/wfMiniAPI/registry.py index b0b8519..04b90ba 100644 --- a/wfMiniAPI/src/wfMiniAPI/registry.py +++ b/wfMiniAPI/src/wfMiniAPI/registry.py @@ -1,5 +1,6 @@ import inspect from collections import OrderedDict +import time _KERNELS = OrderedDict() From 886df95b787f1cda363fece9e808ece8a889c144 Mon Sep 17 00:00:00 2001 From: GKNB Date: Mon, 28 Jul 2025 03:46:37 +0000 Subject: [PATCH 09/46] cupy --- wfMiniAPI/src/wfMiniAPI/registry.py | 1 + 1 file changed, 1 insertion(+) diff --git a/wfMiniAPI/src/wfMiniAPI/registry.py b/wfMiniAPI/src/wfMiniAPI/registry.py index 04b90ba..7c833bc 100644 --- a/wfMiniAPI/src/wfMiniAPI/registry.py +++ b/wfMiniAPI/src/wfMiniAPI/registry.py @@ -55,6 +55,7 @@ def time_kernel(name, device, n_warmup=1, n_repeat=20, **kwargs): # GPU timing elif device.lower() == "gpu": + import cupy as cp # warmup + sync for _ in range(n_warmup): run_kernel(name, device=device, **kwargs) From c671d716e9d3682879faf690012c11837f1fc6b7 Mon Sep 17 00:00:00 2001 From: GKNB Date: Mon, 28 Jul 2025 03:58:48 +0000 Subject: [PATCH 10/46] benchmark for axpy --- wfMiniAPI/benchmark/axpy_benchmark.cu | 91 +++++++++++++++++++++++++++ 1 file changed, 91 insertions(+) create mode 100644 wfMiniAPI/benchmark/axpy_benchmark.cu diff --git a/wfMiniAPI/benchmark/axpy_benchmark.cu b/wfMiniAPI/benchmark/axpy_benchmark.cu new file mode 100644 index 0000000..de29171 --- /dev/null +++ b/wfMiniAPI/benchmark/axpy_benchmark.cu @@ -0,0 +1,91 @@ +#include +#include +#include +#include + +#define CHECK_CUDA(call) \ + do { \ + cudaError_t err = call; \ + if (err != cudaSuccess) { \ + fprintf(stderr, "CUDA Error %s:%d: %s\n", __FILE__, __LINE__, \ + cudaGetErrorString(err)); \ + std::exit(EXIT_FAILURE); \ + } \ + } while (0) + +#define CHECK_CUBLAS(call) \ + do { \ + cublasStatus_t st = call; \ + if (st != CUBLAS_STATUS_SUCCESS) { \ + fprintf(stderr, "cuBLAS Error %s:%d: %d\n", __FILE__, __LINE__, st); \ + std::exit(EXIT_FAILURE); \ + } \ + } while (0) + +int main() { + const int N = 1024 * 1024; + const float alpha = 2.5f; + const int n_warmup = 3; + const int n_repeat = 50; + + // Host allocations + float *h_x = (float*)malloc(N * sizeof(float)); + float *h_y = (float*)malloc(N * sizeof(float)); + for (int i = 0; i < N; ++i) { + h_x[i] = 1.0f; // or whatever initialization + h_y[i] = 0.0f; + } + + // Device allocations + float *d_x, *d_y; + CHECK_CUDA(cudaMalloc((void**)&d_x, N * sizeof(float))); + CHECK_CUDA(cudaMalloc((void**)&d_y, N * sizeof(float))); + + // Copy data to device + CHECK_CUDA(cudaMemcpy(d_x, h_x, N * sizeof(float), cudaMemcpyHostToDevice)); + CHECK_CUDA(cudaMemcpy(d_y, h_y, N * sizeof(float), cudaMemcpyHostToDevice)); + + // Create cuBLAS handle + cublasHandle_t handle; + CHECK_CUBLAS(cublasCreate(&handle)); + + // Timing events + cudaEvent_t start, stop; + CHECK_CUDA(cudaEventCreate(&start)); + CHECK_CUDA(cudaEventCreate(&stop)); + + // Warm‐up + for (int i = 0; i < n_warmup; ++i) { + CHECK_CUBLAS(cublasSaxpy(handle, N, &alpha, d_x, 1, d_y, 1)); + } + CHECK_CUDA(cudaDeviceSynchronize()); + + // Timed repeats + float total_ms = 0.0f; + for (int i = 0; i < n_repeat; ++i) { + CHECK_CUDA(cudaEventRecord(start, 0)); + CHECK_CUBLAS(cublasSaxpy(handle, N, &alpha, d_x, 1, d_y, 1)); + CHECK_CUDA(cudaEventRecord(stop, 0)); + CHECK_CUDA(cudaEventSynchronize(stop)); + float iter_ms = 0.0f; + CHECK_CUDA(cudaEventElapsedTime(&iter_ms, start, stop)); + total_ms += iter_ms; + } + + float avg_ms = total_ms / n_repeat; + double bytes = double(N) * sizeof(float) * 2; // read x and read/write y + double bandwidth = bytes / (avg_ms / 1e3) / 1e9; // GB/s + + printf("AXPY (N=%d) average over %d runs: %f ms, bandwidth = %f GB/s\n", + N, n_repeat, avg_ms, bandwidth); + + // Cleanup + cublasDestroy(handle); + cudaFree(d_x); + cudaFree(d_y); + free(h_x); + free(h_y); + + return 0; +} + From 2c59af0c30cad46b161a34ab9cb1bc5c6d8b600b Mon Sep 17 00:00:00 2001 From: GKNB Date: Mon, 28 Jul 2025 04:06:24 +0000 Subject: [PATCH 11/46] matmul --- wfMiniAPI/benchmark/matmul_benchmark.cu | 120 ++++++++++++++++++++++++ 1 file changed, 120 insertions(+) create mode 100644 wfMiniAPI/benchmark/matmul_benchmark.cu diff --git a/wfMiniAPI/benchmark/matmul_benchmark.cu b/wfMiniAPI/benchmark/matmul_benchmark.cu new file mode 100644 index 0000000..a0f5585 --- /dev/null +++ b/wfMiniAPI/benchmark/matmul_benchmark.cu @@ -0,0 +1,120 @@ +#include +#include +#include +#include + +#define CHECK_CUDA(call) \ + do { \ + cudaError_t err = call; \ + if (err != cudaSuccess) { \ + fprintf(stderr, "CUDA Error %s:%d: %s\n", __FILE__, __LINE__, \ + cudaGetErrorString(err)); \ + std::exit(EXIT_FAILURE); \ + } \ + } while (0) + +#define CHECK_CUBLAS(call) \ + do { \ + cublasStatus_t st = call; \ + if (st != CUBLAS_STATUS_SUCCESS) { \ + fprintf(stderr, "cuBLAS Error %s:%d: %d\n", __FILE__, __LINE__, st); \ + std::exit(EXIT_FAILURE); \ + } \ + } while (0) + +int main() { + const int N = 1024; + const size_t bytes = size_t(N) * N * sizeof(float); + const float alpha = 1.0f; + const float beta = 0.0f; + const int n_warmup = 3; + const int n_repeat = 50; + + // Host allocations & initialize + float *h_A = (float*)malloc(bytes); + float *h_B = (float*)malloc(bytes); + float *h_C = (float*)malloc(bytes); + for (int i = 0; i < N*N; ++i) { + h_A[i] = 1.0f; + h_B[i] = 1.0f; + h_C[i] = 0.0f; + } + + // Device allocations + float *d_A, *d_B, *d_C; + CHECK_CUDA(cudaMalloc(&d_A, bytes)); + CHECK_CUDA(cudaMalloc(&d_B, bytes)); + CHECK_CUDA(cudaMalloc(&d_C, bytes)); + + // Copy data to device + CHECK_CUDA(cudaMemcpy(d_A, h_A, bytes, cudaMemcpyHostToDevice)); + CHECK_CUDA(cudaMemcpy(d_B, h_B, bytes, cudaMemcpyHostToDevice)); + CHECK_CUDA(cudaMemcpy(d_C, h_C, bytes, cudaMemcpyHostToDevice)); + + // cuBLAS handle + cublasHandle_t handle; + CHECK_CUBLAS(cublasCreate(&handle)); + + // Create CUDA events for timing + cudaEvent_t start, stop; + CHECK_CUDA(cudaEventCreate(&start)); + CHECK_CUDA(cudaEventCreate(&stop)); + + // Warm‑up SGEMM calls (not timed) + for (int i = 0; i < n_warmup; ++i) { + CHECK_CUBLAS(cublasSgemm( + handle, + CUBLAS_OP_N, CUBLAS_OP_N, + N, N, N, + &alpha, + d_A, N, + d_B, N, + &beta, + d_C, N + )); + } + CHECK_CUDA(cudaDeviceSynchronize()); + + // Timed repeats + float total_ms = 0.0f; + for (int i = 0; i < n_repeat; ++i) { + CHECK_CUDA(cudaEventRecord(start, 0)); + CHECK_CUBLAS(cublasSgemm( + handle, + CUBLAS_OP_N, CUBLAS_OP_N, + N, N, N, + &alpha, + d_A, N, + d_B, N, + &beta, + d_C, N + )); + CHECK_CUDA(cudaEventRecord(stop, 0)); + CHECK_CUDA(cudaEventSynchronize(stop)); + + float iter_ms = 0.0f; + CHECK_CUDA(cudaEventElapsedTime(&iter_ms, start, stop)); + total_ms += iter_ms; + } + + float avg_ms = total_ms / n_repeat; + double flops = 2.0 * double(N) * N * N; // 2·N³ operations for SGEMM + double gflops = (flops / (avg_ms / 1e3)) / 1e9; + + printf("cuBLAS SGEMM (N=%d) average over %d runs: %f ms → %f GFLOPS\n", + N, n_repeat, avg_ms, gflops); + + // Cleanup + CHECK_CUBLAS(cublasDestroy(handle)); + CHECK_CUDA(cudaFree(d_A)); + CHECK_CUDA(cudaFree(d_B)); + CHECK_CUDA(cudaFree(d_C)); + CHECK_CUDA(cudaEventDestroy(start)); + CHECK_CUDA(cudaEventDestroy(stop)); + free(h_A); + free(h_B); + free(h_C); + + return 0; +} + From 8e3962f2290d0e70402870e2bee621a333673bae Mon Sep 17 00:00:00 2001 From: GKNB Date: Mon, 28 Jul 2025 04:15:04 +0000 Subject: [PATCH 12/46] fft_benchmark --- wfMiniAPI/benchmark/fft_benchmark.cu | 103 +++++++++++++++++++++++++++ 1 file changed, 103 insertions(+) create mode 100644 wfMiniAPI/benchmark/fft_benchmark.cu diff --git a/wfMiniAPI/benchmark/fft_benchmark.cu b/wfMiniAPI/benchmark/fft_benchmark.cu new file mode 100644 index 0000000..d3d7c5b --- /dev/null +++ b/wfMiniAPI/benchmark/fft_benchmark.cu @@ -0,0 +1,103 @@ +#include +#include +#include +#include +#include + +#define CHECK_CUDA(call) \ + do { \ + cudaError_t err = call; \ + if (err != cudaSuccess) { \ + fprintf(stderr, "CUDA Error %s:%d: %s\n", __FILE__, __LINE__, \ + cudaGetErrorString(err)); \ + std::exit(EXIT_FAILURE); \ + } \ + } while (0) + +#define CHECK_CUFFT(call) \ + do { \ + cufftResult err = call; \ + if (err != CUFFT_SUCCESS) { \ + fprintf(stderr, "cuFFT Error %s:%d: %d\n", __FILE__, __LINE__, err); \ + std::exit(EXIT_FAILURE); \ + } \ + } while (0) + +int main() { + const int N = 1024; // length of each FFT + const int batch = 1024; // number of transforms (one per column) + const int n_warmup = 3; + const int n_repeat = 50; + + // host-side real input (simulate data_in); we'll pack into complex + size_t real_elems = size_t(N) * batch; + float *h_real = (float*)malloc(real_elems * sizeof(float)); + for (size_t i = 0; i < real_elems; ++i) { + h_real[i] = 1.0f; // or load your actual data here + } + + // Device buffers: complex in/out + cufftComplex *d_in = nullptr, *d_out = nullptr; + size_t complex_bytes = real_elems * sizeof(cufftComplex); + CHECK_CUDA(cudaMalloc(&d_in, complex_bytes)); + CHECK_CUDA(cudaMalloc(&d_out, complex_bytes)); + + // Pack real→complex on host and copy once + cufftComplex *h_pack = (cufftComplex*)malloc(complex_bytes); + for (size_t i = 0; i < real_elems; ++i) { + h_pack[i].x = h_real[i]; + h_pack[i].y = 0.0f; + } + CHECK_CUDA(cudaMemcpy(d_in, h_pack, complex_bytes, cudaMemcpyHostToDevice)); + free(h_pack); + free(h_real); + + // Create a batched 1D C2C plan + cufftHandle plan; + CHECK_CUFFT(cufftPlan1d(&plan, N, CUFFT_C2C, batch)); + + // Create CUDA events for timing + cudaEvent_t start, stop; + CHECK_CUDA(cudaEventCreate(&start)); + CHECK_CUDA(cudaEventCreate(&stop)); + + // Warm‑up (not timed) + for (int i = 0; i < n_warmup; ++i) { + CHECK_CUFFT(cufftExecC2C(plan, d_in, d_out, CUFFT_FORWARD)); + } + CHECK_CUDA(cudaDeviceSynchronize()); + + // Timed repeats + float total_ms = 0.0f; + for (int i = 0; i < n_repeat; ++i) { + CHECK_CUDA(cudaEventRecord(start, 0)); + CHECK_CUFFT(cufftExecC2C(plan, d_in, d_out, CUFFT_FORWARD)); + CHECK_CUDA(cudaEventRecord(stop, 0)); + CHECK_CUDA(cudaEventSynchronize(stop)); + + float iter_ms = 0.0f; + CHECK_CUDA(cudaEventElapsedTime(&iter_ms, start, stop)); + total_ms += iter_ms; + } + + float avg_ms = total_ms / n_repeat; + + // Approximate FLOPS: ~5·N·log2(N) per transform + double flops_per_transform = 5.0 * N * std::log2(double(N)); + double total_flops = flops_per_transform * batch; + double gflops = total_flops / (avg_ms/1e3) / 1e9; + + printf("cuFFT C2C 1D FFT (N=%d, batch=%d)\n", N, batch); + printf(" Average over %d runs: %f ms\n", n_repeat, avg_ms); + printf(" Approx. throughput: %f GFLOPS\n", gflops); + + // Cleanup + CHECK_CUFFT(cufftDestroy(plan)); + CHECK_CUDA(cudaFree(d_in)); + CHECK_CUDA(cudaFree(d_out)); + CHECK_CUDA(cudaEventDestroy(start)); + CHECK_CUDA(cudaEventDestroy(stop)); + + return 0; +} + From 9c111541fac569e9f1a302f83b713eeaa2d32dfe Mon Sep 17 00:00:00 2001 From: GKNB Date: Mon, 28 Jul 2025 04:22:08 +0000 Subject: [PATCH 13/46] benchmark --- wfMiniAPI/benchmark/axpy_benchmark.cu | 19 ++++--------------- wfMiniAPI/benchmark/fft_benchmark.cu | 22 ++++------------------ wfMiniAPI/benchmark/matmul_benchmark.cu | 15 ++------------- 3 files changed, 10 insertions(+), 46 deletions(-) diff --git a/wfMiniAPI/benchmark/axpy_benchmark.cu b/wfMiniAPI/benchmark/axpy_benchmark.cu index de29171..1ea1ca9 100644 --- a/wfMiniAPI/benchmark/axpy_benchmark.cu +++ b/wfMiniAPI/benchmark/axpy_benchmark.cu @@ -24,43 +24,36 @@ int main() { const int N = 1024 * 1024; - const float alpha = 2.5f; + const float alpha = 1.1f; const int n_warmup = 3; const int n_repeat = 50; - // Host allocations float *h_x = (float*)malloc(N * sizeof(float)); float *h_y = (float*)malloc(N * sizeof(float)); for (int i = 0; i < N; ++i) { - h_x[i] = 1.0f; // or whatever initialization + h_x[i] = 1.0f; h_y[i] = 0.0f; } - // Device allocations float *d_x, *d_y; CHECK_CUDA(cudaMalloc((void**)&d_x, N * sizeof(float))); CHECK_CUDA(cudaMalloc((void**)&d_y, N * sizeof(float))); - // Copy data to device CHECK_CUDA(cudaMemcpy(d_x, h_x, N * sizeof(float), cudaMemcpyHostToDevice)); CHECK_CUDA(cudaMemcpy(d_y, h_y, N * sizeof(float), cudaMemcpyHostToDevice)); - // Create cuBLAS handle cublasHandle_t handle; CHECK_CUBLAS(cublasCreate(&handle)); - // Timing events cudaEvent_t start, stop; CHECK_CUDA(cudaEventCreate(&start)); CHECK_CUDA(cudaEventCreate(&stop)); - // Warm‐up for (int i = 0; i < n_warmup; ++i) { CHECK_CUBLAS(cublasSaxpy(handle, N, &alpha, d_x, 1, d_y, 1)); } CHECK_CUDA(cudaDeviceSynchronize()); - // Timed repeats float total_ms = 0.0f; for (int i = 0; i < n_repeat; ++i) { CHECK_CUDA(cudaEventRecord(start, 0)); @@ -73,13 +66,9 @@ int main() { } float avg_ms = total_ms / n_repeat; - double bytes = double(N) * sizeof(float) * 2; // read x and read/write y - double bandwidth = bytes / (avg_ms / 1e3) / 1e9; // GB/s + printf("AXPY (N=%d) average over %d runs: %f ms\n", + N, n_repeat, avg_ms); - printf("AXPY (N=%d) average over %d runs: %f ms, bandwidth = %f GB/s\n", - N, n_repeat, avg_ms, bandwidth); - - // Cleanup cublasDestroy(handle); cudaFree(d_x); cudaFree(d_y); diff --git a/wfMiniAPI/benchmark/fft_benchmark.cu b/wfMiniAPI/benchmark/fft_benchmark.cu index d3d7c5b..94a0e08 100644 --- a/wfMiniAPI/benchmark/fft_benchmark.cu +++ b/wfMiniAPI/benchmark/fft_benchmark.cu @@ -24,25 +24,22 @@ } while (0) int main() { - const int N = 1024; // length of each FFT - const int batch = 1024; // number of transforms (one per column) + const int N = 1024; + const int batch = 1024; const int n_warmup = 3; const int n_repeat = 50; - // host-side real input (simulate data_in); we'll pack into complex size_t real_elems = size_t(N) * batch; float *h_real = (float*)malloc(real_elems * sizeof(float)); for (size_t i = 0; i < real_elems; ++i) { - h_real[i] = 1.0f; // or load your actual data here + h_real[i] = 1.0f; } - // Device buffers: complex in/out cufftComplex *d_in = nullptr, *d_out = nullptr; size_t complex_bytes = real_elems * sizeof(cufftComplex); CHECK_CUDA(cudaMalloc(&d_in, complex_bytes)); CHECK_CUDA(cudaMalloc(&d_out, complex_bytes)); - // Pack real→complex on host and copy once cufftComplex *h_pack = (cufftComplex*)malloc(complex_bytes); for (size_t i = 0; i < real_elems; ++i) { h_pack[i].x = h_real[i]; @@ -52,22 +49,18 @@ int main() { free(h_pack); free(h_real); - // Create a batched 1D C2C plan cufftHandle plan; CHECK_CUFFT(cufftPlan1d(&plan, N, CUFFT_C2C, batch)); - // Create CUDA events for timing cudaEvent_t start, stop; CHECK_CUDA(cudaEventCreate(&start)); CHECK_CUDA(cudaEventCreate(&stop)); - // Warm‑up (not timed) for (int i = 0; i < n_warmup; ++i) { CHECK_CUFFT(cufftExecC2C(plan, d_in, d_out, CUFFT_FORWARD)); } CHECK_CUDA(cudaDeviceSynchronize()); - // Timed repeats float total_ms = 0.0f; for (int i = 0; i < n_repeat; ++i) { CHECK_CUDA(cudaEventRecord(start, 0)); @@ -82,16 +75,9 @@ int main() { float avg_ms = total_ms / n_repeat; - // Approximate FLOPS: ~5·N·log2(N) per transform - double flops_per_transform = 5.0 * N * std::log2(double(N)); - double total_flops = flops_per_transform * batch; - double gflops = total_flops / (avg_ms/1e3) / 1e9; - - printf("cuFFT C2C 1D FFT (N=%d, batch=%d)\n", N, batch); + printf("cuFFT C2C 1D FFT (N=%d, batch=%d)", N, batch); printf(" Average over %d runs: %f ms\n", n_repeat, avg_ms); - printf(" Approx. throughput: %f GFLOPS\n", gflops); - // Cleanup CHECK_CUFFT(cufftDestroy(plan)); CHECK_CUDA(cudaFree(d_in)); CHECK_CUDA(cudaFree(d_out)); diff --git a/wfMiniAPI/benchmark/matmul_benchmark.cu b/wfMiniAPI/benchmark/matmul_benchmark.cu index a0f5585..bbb9c67 100644 --- a/wfMiniAPI/benchmark/matmul_benchmark.cu +++ b/wfMiniAPI/benchmark/matmul_benchmark.cu @@ -30,7 +30,6 @@ int main() { const int n_warmup = 3; const int n_repeat = 50; - // Host allocations & initialize float *h_A = (float*)malloc(bytes); float *h_B = (float*)malloc(bytes); float *h_C = (float*)malloc(bytes); @@ -40,27 +39,22 @@ int main() { h_C[i] = 0.0f; } - // Device allocations float *d_A, *d_B, *d_C; CHECK_CUDA(cudaMalloc(&d_A, bytes)); CHECK_CUDA(cudaMalloc(&d_B, bytes)); CHECK_CUDA(cudaMalloc(&d_C, bytes)); - // Copy data to device CHECK_CUDA(cudaMemcpy(d_A, h_A, bytes, cudaMemcpyHostToDevice)); CHECK_CUDA(cudaMemcpy(d_B, h_B, bytes, cudaMemcpyHostToDevice)); CHECK_CUDA(cudaMemcpy(d_C, h_C, bytes, cudaMemcpyHostToDevice)); - // cuBLAS handle cublasHandle_t handle; CHECK_CUBLAS(cublasCreate(&handle)); - // Create CUDA events for timing cudaEvent_t start, stop; CHECK_CUDA(cudaEventCreate(&start)); CHECK_CUDA(cudaEventCreate(&stop)); - // Warm‑up SGEMM calls (not timed) for (int i = 0; i < n_warmup; ++i) { CHECK_CUBLAS(cublasSgemm( handle, @@ -75,7 +69,6 @@ int main() { } CHECK_CUDA(cudaDeviceSynchronize()); - // Timed repeats float total_ms = 0.0f; for (int i = 0; i < n_repeat; ++i) { CHECK_CUDA(cudaEventRecord(start, 0)); @@ -98,13 +91,10 @@ int main() { } float avg_ms = total_ms / n_repeat; - double flops = 2.0 * double(N) * N * N; // 2·N³ operations for SGEMM - double gflops = (flops / (avg_ms / 1e3)) / 1e9; - printf("cuBLAS SGEMM (N=%d) average over %d runs: %f ms → %f GFLOPS\n", - N, n_repeat, avg_ms, gflops); + printf("cuBLAS SGEMM (N=%d) average over %d runs: %f ms\n", + N, n_repeat, avg_ms); - // Cleanup CHECK_CUBLAS(cublasDestroy(handle)); CHECK_CUDA(cudaFree(d_A)); CHECK_CUDA(cudaFree(d_B)); @@ -117,4 +107,3 @@ int main() { return 0; } - From 19f09a7ebffd09310375488cd8bc37a6b8c336f4 Mon Sep 17 00:00:00 2001 From: GKNB Date: Mon, 28 Jul 2025 04:44:53 +0000 Subject: [PATCH 14/46] axpy improve --- wfMiniAPI/benchmark/axpy_benchmark.cu | 2 +- wfMiniAPI/benchmark/axpy_naive_benchmark.cu | 78 +++++++++++++++ wfMiniAPI/src/wfMiniAPI/kernel.py | 20 +++- wfMiniAPI/src/wfMiniAPI/matmul.cu | 104 -------------------- 4 files changed, 98 insertions(+), 106 deletions(-) create mode 100644 wfMiniAPI/benchmark/axpy_naive_benchmark.cu delete mode 100644 wfMiniAPI/src/wfMiniAPI/matmul.cu diff --git a/wfMiniAPI/benchmark/axpy_benchmark.cu b/wfMiniAPI/benchmark/axpy_benchmark.cu index 1ea1ca9..57c640a 100644 --- a/wfMiniAPI/benchmark/axpy_benchmark.cu +++ b/wfMiniAPI/benchmark/axpy_benchmark.cu @@ -23,7 +23,7 @@ } while (0) int main() { - const int N = 1024 * 1024; + const int N = 1024 * 1024 * 32; const float alpha = 1.1f; const int n_warmup = 3; const int n_repeat = 50; diff --git a/wfMiniAPI/benchmark/axpy_naive_benchmark.cu b/wfMiniAPI/benchmark/axpy_naive_benchmark.cu new file mode 100644 index 0000000..ea01a23 --- /dev/null +++ b/wfMiniAPI/benchmark/axpy_naive_benchmark.cu @@ -0,0 +1,78 @@ +#include +#include +#include + +#define CHECK_CUDA(call) \ + do { \ + cudaError_t err = call; \ + if (err != cudaSuccess) { \ + fprintf(stderr, "CUDA Error %s:%d: %s\n", __FILE__, __LINE__, \ + cudaGetErrorString(err)); \ + std::exit(EXIT_FAILURE); \ + } \ + } while (0) + +__global__ void axpy_kernel(int N, float alpha, const float* x, float* y) { + int idx = blockIdx.x * blockDim.x + threadIdx.x; + if (idx < N) { + y[idx] += alpha * x[idx]; + } +} + +int main() { + const int N = 1024 * 1024 * 32; + const float alpha = 1.1f; + const int n_warmup = 3; + const int n_repeat = 50; + + float *h_x = (float*)malloc(N * sizeof(float)); + float *h_y = (float*)malloc(N * sizeof(float)); + for (int i = 0; i < N; ++i) { + h_x[i] = 1.0f; + h_y[i] = 0.0f; + } + + float *d_x, *d_y; + CHECK_CUDA(cudaMalloc((void**)&d_x, N * sizeof(float))); + CHECK_CUDA(cudaMalloc((void**)&d_y, N * sizeof(float))); + + CHECK_CUDA(cudaMemcpy(d_x, h_x, N * sizeof(float), cudaMemcpyHostToDevice)); + CHECK_CUDA(cudaMemcpy(d_y, h_y, N * sizeof(float), cudaMemcpyHostToDevice)); + + const int TPB = 256; + int blocks = (N + TPB - 1) / TPB; + + cudaEvent_t start, stop; + CHECK_CUDA(cudaEventCreate(&start)); + CHECK_CUDA(cudaEventCreate(&stop)); + + for (int i = 0; i < n_warmup; ++i) { + axpy_kernel<<>>(N, alpha, d_x, d_y); + } + CHECK_CUDA(cudaDeviceSynchronize()); + + float total_ms = 0.0f; + for (int i = 0; i < n_repeat; ++i) { + CHECK_CUDA(cudaEventRecord(start, 0)); + axpy_kernel<<>>(N, alpha, d_x, d_y); + CHECK_CUDA(cudaEventRecord(stop, 0)); + CHECK_CUDA(cudaEventSynchronize(stop)); + float iter_ms = 0.0f; + CHECK_CUDA(cudaEventElapsedTime(&iter_ms, start, stop)); + total_ms += iter_ms; + } + + float avg_ms = total_ms / n_repeat; + printf("AXPY kernel (N=%d) average over %d runs: %f ms\n", + N, n_repeat, avg_ms); + + CHECK_CUDA(cudaEventDestroy(start)); + CHECK_CUDA(cudaEventDestroy(stop)); + CHECK_CUDA(cudaFree(d_x)); + CHECK_CUDA(cudaFree(d_y)); + free(h_x); + free(h_y); + + return 0; +} + diff --git a/wfMiniAPI/src/wfMiniAPI/kernel.py b/wfMiniAPI/src/wfMiniAPI/kernel.py index 875012c..4ca3697 100644 --- a/wfMiniAPI/src/wfMiniAPI/kernel.py +++ b/wfMiniAPI/src/wfMiniAPI/kernel.py @@ -272,12 +272,30 @@ def fft(device, data_size, type_in, transform_dim): out = xp.fft.fft(data_in, axis=transform_dim) @annotate_kernel -def axpy(device, size): +def axpy_slow(device, size): xp = get_device_module(device) x = xp.empty(size, dtype=xp.float32) y = xp.empty(size, dtype=xp.float32) y += 1.01 * x + +_axpy = cp.ElementwiseKernel( + 'float32 alpha, float32 x, float32 y', + 'float32 y', + 'y += alpha * x', + 'axpy_kernel' +) + +@annotate_kernel +def axpy(device, size): + xp = get_device_module(device) + x = xp.empty(size, dtype=xp.float32) + y = xp.empty(size, dtype=xp.float32) + if xp == np: + y += 1.01 * x + elif xp == cp: + _axpy(1.01, x, y, out=y) + @annotate_kernel def implaceCompute(device, size, num_op, op): xp = get_device_module(device) diff --git a/wfMiniAPI/src/wfMiniAPI/matmul.cu b/wfMiniAPI/src/wfMiniAPI/matmul.cu deleted file mode 100644 index 2d10079..0000000 --- a/wfMiniAPI/src/wfMiniAPI/matmul.cu +++ /dev/null @@ -1,104 +0,0 @@ -// matmul_cublas.cu -#include -#include -#include -#include - -#define CHECK_CUDA(call) \ - do { \ - cudaError_t err = (call); \ - if (err != cudaSuccess) { \ - std::cerr << "CUDA error at " << __FILE__ << ":" << __LINE__ \ - << " code=" << err << " \"" << cudaGetErrorString(err) \ - << "\"" << std::endl; \ - std::exit(EXIT_FAILURE); \ - } \ - } while (0) - -#define CHECK_CUBLAS(call) \ - do { \ - cublasStatus_t stat = (call); \ - if (stat != CUBLAS_STATUS_SUCCESS) { \ - std::cerr << "cuBLAS error at " << __FILE__ << ":" << __LINE__ \ - << " code=" << stat << std::endl; \ - std::exit(EXIT_FAILURE); \ - } \ - } while (0) - -int main(int argc, char* argv[]) { - if (argc != 3) { - std::cerr << "Usage: " << argv[0] << " \n"; - return EXIT_FAILURE; - } - - int dev = std::atoi(argv[1]); - int N = std::atoi(argv[2]); - - // Select GPU - CHECK_CUDA(cudaSetDevice(dev)); - - // cuBLAS handle - cublasHandle_t handle; - CHECK_CUBLAS(cublasCreate(&handle)); - - // Allocate device matrices A, B, C (uninitialized, like cupy.empty) - size_t bytes = size_t(N) * N * sizeof(float); - float *d_A, *d_B, *d_C; - CHECK_CUDA(cudaMalloc(&d_A, bytes)); - CHECK_CUDA(cudaMalloc(&d_B, bytes)); - CHECK_CUDA(cudaMalloc(&d_C, bytes)); - - // Warm-up: one matmul to initialize kernels / cache - const float alpha = 1.0f, beta = 0.0f; - CHECK_CUBLAS( - cublasSgemm(handle, - CUBLAS_OP_N, CUBLAS_OP_N, - N, N, N, - &alpha, - d_A, N, - d_B, N, - &beta, - d_C, N) - ); - // Ensure warm-up completed - CHECK_CUDA(cudaDeviceSynchronize()); - - // Create CUDA events for timing - cudaEvent_t start, stop; - CHECK_CUDA(cudaEventCreate(&start)); - CHECK_CUDA(cudaEventCreate(&stop)); - - // Record start, run 20 matmuls, record stop - CHECK_CUDA(cudaEventRecord(start)); - for (int i = 0; i < 20; ++i) { - CHECK_CUBLAS( - cublasSgemm(handle, - CUBLAS_OP_N, CUBLAS_OP_N, - N, N, N, - &alpha, - d_A, N, - d_B, N, - &beta, - d_C, N) - ); - } - CHECK_CUDA(cudaEventRecord(stop)); - CHECK_CUDA(cudaEventSynchronize(stop)); - - float ms = 0; - CHECK_CUDA(cudaEventElapsedTime(&ms, start, stop)); - - std::cout << "Matrix size: " << N << "×" << N << "\n" - << "Total time for 20 calls: " << ms << " ms\n" - << "Average per call: " << (ms / 20.0f) << " ms\n"; - - // Cleanup - cudaEventDestroy(start); - cudaEventDestroy(stop); - cudaFree(d_A); - cudaFree(d_B); - cudaFree(d_C); - cublasDestroy(handle); - - return 0; -} From 3c200a29e3bf40e18e5ece1f511a43ff64b2fcc0 Mon Sep 17 00:00:00 2001 From: GKNB Date: Mon, 28 Jul 2025 04:47:39 +0000 Subject: [PATCH 15/46] matmul --- wfMiniAPI/benchmark/matmul_benchmark.cu | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/wfMiniAPI/benchmark/matmul_benchmark.cu b/wfMiniAPI/benchmark/matmul_benchmark.cu index bbb9c67..65477ef 100644 --- a/wfMiniAPI/benchmark/matmul_benchmark.cu +++ b/wfMiniAPI/benchmark/matmul_benchmark.cu @@ -23,7 +23,7 @@ } while (0) int main() { - const int N = 1024; + const int N = 4096; const size_t bytes = size_t(N) * N * sizeof(float); const float alpha = 1.0f; const float beta = 0.0f; From fa48be47da8f6692a11602ac43b34c4b9b4a1c4f Mon Sep 17 00:00:00 2001 From: GKNB Date: Mon, 28 Jul 2025 05:00:17 +0000 Subject: [PATCH 16/46] New axpy --- wfMiniAPI/src/wfMiniAPI/kernel.py | 14 +++++++------- 1 file changed, 7 insertions(+), 7 deletions(-) diff --git a/wfMiniAPI/src/wfMiniAPI/kernel.py b/wfMiniAPI/src/wfMiniAPI/kernel.py index 4ca3697..44b5a83 100644 --- a/wfMiniAPI/src/wfMiniAPI/kernel.py +++ b/wfMiniAPI/src/wfMiniAPI/kernel.py @@ -278,12 +278,12 @@ def axpy_slow(device, size): y = xp.empty(size, dtype=xp.float32) y += 1.01 * x - -_axpy = cp.ElementwiseKernel( - 'float32 alpha, float32 x, float32 y', - 'float32 y', - 'y += alpha * x', - 'axpy_kernel' +_axpy_inplace = cp.ElementwiseKernel( + 'float32 alpha, float32 x, raw float32 y', + '', + 'y += alpha * x', + 'axpy_inplace_kernel', + no_return=True ) @annotate_kernel @@ -294,7 +294,7 @@ def axpy(device, size): if xp == np: y += 1.01 * x elif xp == cp: - _axpy(1.01, x, y, out=y) + _axpy_inplace(1.01, x, y) @annotate_kernel def implaceCompute(device, size, num_op, op): From b8129f74266f93fa093af52058bfc2b69c3f0439 Mon Sep 17 00:00:00 2001 From: GKNB Date: Mon, 28 Jul 2025 05:11:14 +0000 Subject: [PATCH 17/46] axpy --- wfMiniAPI/src/wfMiniAPI/kernel.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/wfMiniAPI/src/wfMiniAPI/kernel.py b/wfMiniAPI/src/wfMiniAPI/kernel.py index 44b5a83..6e41ceb 100644 --- a/wfMiniAPI/src/wfMiniAPI/kernel.py +++ b/wfMiniAPI/src/wfMiniAPI/kernel.py @@ -279,9 +279,9 @@ def axpy_slow(device, size): y += 1.01 * x _axpy_inplace = cp.ElementwiseKernel( - 'float32 alpha, float32 x, raw float32 y', + 'float32 alpha, raw float32 x, raw float32 y', '', - 'y += alpha * x', + 'y[i] += alpha * x[i]', 'axpy_inplace_kernel', no_return=True ) From a85f597351fc550cb0fdc79ff505b9e39e7683a1 Mon Sep 17 00:00:00 2001 From: GKNB Date: Mon, 28 Jul 2025 05:15:24 +0000 Subject: [PATCH 18/46] axpy --- wfMiniAPI/src/wfMiniAPI/kernel.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/wfMiniAPI/src/wfMiniAPI/kernel.py b/wfMiniAPI/src/wfMiniAPI/kernel.py index 6e41ceb..b1b33cc 100644 --- a/wfMiniAPI/src/wfMiniAPI/kernel.py +++ b/wfMiniAPI/src/wfMiniAPI/kernel.py @@ -294,7 +294,7 @@ def axpy(device, size): if xp == np: y += 1.01 * x elif xp == cp: - _axpy_inplace(1.01, x, y) + _axpy_inplace(1.01, x, y, size=size) @annotate_kernel def implaceCompute(device, size, num_op, op): From a61e5503132f426e0923b4b459ff57305ebdc5b5 Mon Sep 17 00:00:00 2001 From: GKNB Date: Mon, 28 Jul 2025 05:19:43 +0000 Subject: [PATCH 19/46] axpy --- wfMiniAPI/src/wfMiniAPI/kernel.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/wfMiniAPI/src/wfMiniAPI/kernel.py b/wfMiniAPI/src/wfMiniAPI/kernel.py index b1b33cc..ec3e655 100644 --- a/wfMiniAPI/src/wfMiniAPI/kernel.py +++ b/wfMiniAPI/src/wfMiniAPI/kernel.py @@ -279,11 +279,11 @@ def axpy_slow(device, size): y += 1.01 * x _axpy_inplace = cp.ElementwiseKernel( - 'float32 alpha, raw float32 x, raw float32 y', - '', - 'y[i] += alpha * x[i]', + 'float32 alpha, raw float32 x', + 'raw float32 y', + 'y[i] += alpha * x[i]', 'axpy_inplace_kernel', - no_return=True + no_return=True ) @annotate_kernel From ece6702d2cad3702b6599c78dde7f132bd654244 Mon Sep 17 00:00:00 2001 From: GKNB Date: Mon, 28 Jul 2025 05:23:35 +0000 Subject: [PATCH 20/46] axpy --- wfMiniAPI/src/wfMiniAPI/kernel.py | 36 +++++++++++++++---------------- 1 file changed, 18 insertions(+), 18 deletions(-) diff --git a/wfMiniAPI/src/wfMiniAPI/kernel.py b/wfMiniAPI/src/wfMiniAPI/kernel.py index ec3e655..95a0bbf 100644 --- a/wfMiniAPI/src/wfMiniAPI/kernel.py +++ b/wfMiniAPI/src/wfMiniAPI/kernel.py @@ -272,29 +272,29 @@ def fft(device, data_size, type_in, transform_dim): out = xp.fft.fft(data_in, axis=transform_dim) @annotate_kernel -def axpy_slow(device, size): +def axpy(device, size): xp = get_device_module(device) x = xp.empty(size, dtype=xp.float32) y = xp.empty(size, dtype=xp.float32) y += 1.01 * x -_axpy_inplace = cp.ElementwiseKernel( - 'float32 alpha, raw float32 x', - 'raw float32 y', - 'y[i] += alpha * x[i]', - 'axpy_inplace_kernel', - no_return=True -) - -@annotate_kernel -def axpy(device, size): - xp = get_device_module(device) - x = xp.empty(size, dtype=xp.float32) - y = xp.empty(size, dtype=xp.float32) - if xp == np: - y += 1.01 * x - elif xp == cp: - _axpy_inplace(1.01, x, y, size=size) +#_axpy_inplace = cp.ElementwiseKernel( +# 'float32 alpha, raw float32 x', +# 'raw float32 y', +# 'y[i] += alpha * x[i]', +# 'axpy_inplace_kernel', +# no_return=True +#) +# +#@annotate_kernel +#def axpy(device, size): +# xp = get_device_module(device) +# x = xp.empty(size, dtype=xp.float32) +# y = xp.empty(size, dtype=xp.float32) +# if xp == np: +# y += 1.01 * x +# elif xp == cp: +# _axpy_inplace(1.01, x, y, size=size) @annotate_kernel def implaceCompute(device, size, num_op, op): From 650db7da19b7138da6e02bae55739be0e8745818 Mon Sep 17 00:00:00 2001 From: GKNB Date: Mon, 28 Jul 2025 07:35:03 +0000 Subject: [PATCH 21/46] nwarmup --- wfMiniAPI/src/wfMiniAPI/registry.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/wfMiniAPI/src/wfMiniAPI/registry.py b/wfMiniAPI/src/wfMiniAPI/registry.py index 7c833bc..e7c01e7 100644 --- a/wfMiniAPI/src/wfMiniAPI/registry.py +++ b/wfMiniAPI/src/wfMiniAPI/registry.py @@ -37,7 +37,7 @@ def run_kernel(name, **kwargs): out = func(**kwargs) return out -def time_kernel(name, device, n_warmup=1, n_repeat=20, **kwargs): +def time_kernel(name, device, n_warmup=3, n_repeat=20, **kwargs): # CPU timing if device.lower() == "cpu": for _ in range(n_warmup): From dae57b1db8677df164c894ce9c3e0dd152412282 Mon Sep 17 00:00:00 2001 From: GKNB Date: Mon, 28 Jul 2025 08:14:49 +0000 Subject: [PATCH 22/46] timing --- wfMiniAPI/src/wfMiniAPI/registry.py | 38 +++++++++++++++++------------ 1 file changed, 22 insertions(+), 16 deletions(-) diff --git a/wfMiniAPI/src/wfMiniAPI/registry.py b/wfMiniAPI/src/wfMiniAPI/registry.py index e7c01e7..80f7517 100644 --- a/wfMiniAPI/src/wfMiniAPI/registry.py +++ b/wfMiniAPI/src/wfMiniAPI/registry.py @@ -1,6 +1,7 @@ import inspect from collections import OrderedDict import time +import numpy as np _KERNELS = OrderedDict() @@ -42,35 +43,40 @@ def time_kernel(name, device, n_warmup=3, n_repeat=20, **kwargs): if device.lower() == "cpu": for _ in range(n_warmup): run_kernel(name, device=device, **kwargs) - t0 = time.time() + run_times = [] for _ in range(n_repeat): + t0 = time.time() run_kernel(name, device=device, **kwargs) - t1 = time.time() + t1 = time.time() + run_times.append((t1 - t0) * 1000) - total_ms = (t1 - t0) * 1000 + total_ms = sum(run_times) avg_ms = total_ms / n_repeat + std_ms = np.std(run_times) print(f"CPU: Total {n_repeat} runs: {total_ms:.2f} ms") - print(f"CPU: Avg per run: {avg_ms:.2f} ms") - return total_ms, avg_ms + print(f"CPU: Avg per run: {avg_ms:.2f} \u00B1 {std_ms:.4f} ms") + return total_ms, avg_ms, std_ms # GPU timing elif device.lower() == "gpu": import cupy as cp - # warmup + sync for _ in range(n_warmup): run_kernel(name, device=device, **kwargs) cp.cuda.Stream.null.synchronize() - start = cp.cuda.Event() - end = cp.cuda.Event() - start.record() + run_times = [] for _ in range(n_repeat): + start = cp.cuda.Event() + end = cp.cuda.Event() + start.record() run_kernel(name, device=device, **kwargs) - end.record() - end.synchronize() - - total_ms = cp.cuda.get_elapsed_time(start, end) - avg_ms = total_ms / n_repeat + end.record() + end.synchronize() + elapsed_time = cp.cuda.get_elapsed_time(start, end) # in ms + run_times.append(elapsed_time) + total_ms = sum(run_times) + avg_ms = total_ms / n_repeat + std_ms = np.std(run_times) print(f"GPU: Total {n_repeat} runs: {total_ms:.2f} ms") - print(f"GPU: Avg per run: {avg_ms:.2f} ms") - return total_ms, avg_ms + print(f"GPU: Avg per run: {avg_ms:.2f} \u00B1 {std_ms:.4f} ms") + return total_ms, avg_ms, std_ms From 00376acdc5eecb22d98218cd8602a3682f144065 Mon Sep 17 00:00:00 2001 From: GKNB Date: Mon, 28 Jul 2025 15:11:04 +0000 Subject: [PATCH 23/46] timing, fft --- wfMiniAPI/benchmark/fft_benchmark.cu | 2 +- wfMiniAPI/benchmark/fft_benchmark_3d.cu | 91 +++++++++++++++++++++++++ wfMiniAPI/src/wfMiniAPI/registry.py | 8 +++ 3 files changed, 100 insertions(+), 1 deletion(-) create mode 100644 wfMiniAPI/benchmark/fft_benchmark_3d.cu diff --git a/wfMiniAPI/benchmark/fft_benchmark.cu b/wfMiniAPI/benchmark/fft_benchmark.cu index 94a0e08..2dcd79d 100644 --- a/wfMiniAPI/benchmark/fft_benchmark.cu +++ b/wfMiniAPI/benchmark/fft_benchmark.cu @@ -35,7 +35,7 @@ int main() { h_real[i] = 1.0f; } - cufftComplex *d_in = nullptr, *d_out = nullptr; + cufftComplex *d_in, *d_out; size_t complex_bytes = real_elems * sizeof(cufftComplex); CHECK_CUDA(cudaMalloc(&d_in, complex_bytes)); CHECK_CUDA(cudaMalloc(&d_out, complex_bytes)); diff --git a/wfMiniAPI/benchmark/fft_benchmark_3d.cu b/wfMiniAPI/benchmark/fft_benchmark_3d.cu new file mode 100644 index 0000000..2fa5df6 --- /dev/null +++ b/wfMiniAPI/benchmark/fft_benchmark_3d.cu @@ -0,0 +1,91 @@ +#include +#include +#include +#include +#include + +#define CHECK_CUDA(call) \ + do { \ + cudaError_t err = call; \ + if (err != cudaSuccess) { \ + fprintf(stderr, "CUDA Error %s:%d: %s\n", __FILE__, __LINE__, \ + cudaGetErrorString(err)); \ + std::exit(EXIT_FAILURE); \ + } \ + } while (0) + +#define CHECK_CUFFT(call) \ + do { \ + cufftResult err = call; \ + if (err != CUFFT_SUCCESS) { \ + fprintf(stderr, "cuFFT Error %s:%d: %d\n", __FILE__, __LINE__, err); \ + std::exit(EXIT_FAILURE); \ + } \ + } while (0) + +int main() { + const int Nx = 256; + const int Ny = 256; + const int Nz = 256; + const int n_warmup = 3; + const int n_repeat = 50; + + size_t real_elems = size_t(Nx) * Ny * Nz; + + float *h_real = (float*)malloc(real_elems * sizeof(float)); + for (size_t i = 0; i < real_elems; ++i) { + h_real[i] = 1.0f; + } + + cufftComplex *d_in, *d_out; + size_t complex_bytes = real_elems * sizeof(cufftComplex); + CHECK_CUDA(cudaMalloc(&d_in, complex_bytes)); + CHECK_CUDA(cudaMalloc(&d_out, complex_bytes)); + + cufftComplex *h_pack = (cufftComplex*)malloc(complex_bytes); + for (size_t i = 0; i < real_elems; ++i) { + h_pack[i].x = h_real[i]; + h_pack[i].y = 0.0f; + } + CHECK_CUDA(cudaMemcpy(d_in, h_pack, complex_bytes, cudaMemcpyHostToDevice)); + free(h_pack); + free(h_real); + + cufftHandle plan; + CHECK_CUFFT(cufftPlan3d(&plan, Nx, Ny, Nz, CUFFT_C2C)); + + cudaEvent_t start, stop; + CHECK_CUDA(cudaEventCreate(&start)); + CHECK_CUDA(cudaEventCreate(&stop)); + + for (int i = 0; i < n_warmup; ++i) { + CHECK_CUFFT(cufftExecC2C(plan, d_in, d_out, CUFFT_FORWARD)); + } + CHECK_CUDA(cudaDeviceSynchronize()); + + float total_ms = 0.0f; + for (int i = 0; i < n_repeat; ++i) { + CHECK_CUDA(cudaEventRecord(start, 0)); + CHECK_CUFFT(cufftExecC2C(plan, d_in, d_out, CUFFT_FORWARD)); + CHECK_CUDA(cudaEventRecord(stop, 0)); + CHECK_CUDA(cudaEventSynchronize(stop)); + + float iter_ms = 0.0f; + CHECK_CUDA(cudaEventElapsedTime(&iter_ms, start, stop)); + total_ms += iter_ms; + } + + float avg_ms = total_ms / n_repeat; + + printf("cuFFT C2C 3D FFT (Nx=%d, Ny=%d, Nz=%d)", Nx, Ny, Nz); + printf(" Average over %d runs: %f ms\n", n_repeat, avg_ms); + + CHECK_CUFFT(cufftDestroy(plan)); + CHECK_CUDA(cudaFree(d_in)); + CHECK_CUDA(cudaFree(d_out)); + CHECK_CUDA(cudaEventDestroy(start)); + CHECK_CUDA(cudaEventDestroy(stop)); + + return 0; +} + diff --git a/wfMiniAPI/src/wfMiniAPI/registry.py b/wfMiniAPI/src/wfMiniAPI/registry.py index 80f7517..e1fdd09 100644 --- a/wfMiniAPI/src/wfMiniAPI/registry.py +++ b/wfMiniAPI/src/wfMiniAPI/registry.py @@ -64,6 +64,14 @@ def time_kernel(name, device, n_warmup=3, n_repeat=20, **kwargs): run_kernel(name, device=device, **kwargs) cp.cuda.Stream.null.synchronize() + for _ in range(n_warmup): + start = cp.cuda.Event() + end = cp.cuda.Event() + start.record() + run_kernel(name, device=device, **kwargs) + end.record() + end.synchronize() + run_times = [] for _ in range(n_repeat): start = cp.cuda.Event() From c42a91f7d0be6f5226d5e47a39fd8053544b1506 Mon Sep 17 00:00:00 2001 From: GKNB Date: Mon, 28 Jul 2025 15:16:16 +0000 Subject: [PATCH 24/46] fft --- wfMiniAPI/src/wfMiniAPI/kernel.py | 16 ++++++++++++++++ 1 file changed, 16 insertions(+) diff --git a/wfMiniAPI/src/wfMiniAPI/kernel.py b/wfMiniAPI/src/wfMiniAPI/kernel.py index 95a0bbf..370b1c4 100644 --- a/wfMiniAPI/src/wfMiniAPI/kernel.py +++ b/wfMiniAPI/src/wfMiniAPI/kernel.py @@ -271,6 +271,22 @@ def fft(device, data_size, type_in, transform_dim): out = xp.fft.fft(data_in, axis=transform_dim) +@annotate_kernel +def fftn(device, data_size, type_in, transform_dim): + xp = get_device_module(device) + if type_in == "float": + data_in = xp.empty(tuple(data_size), dtype=xp.float32) + elif type_in == "double": + data_in = xp.empty(tuple(data_size), dtype=xp.float64) + elif type_in == "complexF": + data_in = xp.empty(tuple(data_size), dtype=xp.complex64) + elif type_in == "complexD": + data_in = xp.empty(tuple(data_size), dtype=xp.complex128) + else: + raise TypeError("In fftn call, type_in must be one of the following: [float, double, complexF, complexD]") + + out = xp.fft.fftn(data_in, axis=transform_dim) + @annotate_kernel def axpy(device, size): xp = get_device_module(device) From ce75af9d14462642ad15ca5d1c9d4a34f0ed3fda Mon Sep 17 00:00:00 2001 From: GKNB Date: Mon, 28 Jul 2025 15:22:26 +0000 Subject: [PATCH 25/46] fix 3d fft --- wfMiniAPI/src/wfMiniAPI/kernel.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/wfMiniAPI/src/wfMiniAPI/kernel.py b/wfMiniAPI/src/wfMiniAPI/kernel.py index 370b1c4..af28c65 100644 --- a/wfMiniAPI/src/wfMiniAPI/kernel.py +++ b/wfMiniAPI/src/wfMiniAPI/kernel.py @@ -285,7 +285,7 @@ def fftn(device, data_size, type_in, transform_dim): else: raise TypeError("In fftn call, type_in must be one of the following: [float, double, complexF, complexD]") - out = xp.fft.fftn(data_in, axis=transform_dim) + out = xp.fft.fftn(data_in, axes=transform_dim) @annotate_kernel def axpy(device, size): From a7090731e67ea750568a0c74c1c708c17ea3871c Mon Sep 17 00:00:00 2001 From: GKNB Date: Tue, 29 Jul 2025 09:31:05 +0000 Subject: [PATCH 26/46] add top_k and emulated task --- examples/Tutorial/tasks/active_learn.py | 66 ++++++++++++++++++++ examples/Tutorial/tasks/simulation.py | 81 +++++++++++++++++++++++++ examples/Tutorial/tasks/training.py | 77 +++++++++++++++++++++++ wfMiniAPI/src/wfMiniAPI/kernel.py | 34 +++++++++-- 4 files changed, 254 insertions(+), 4 deletions(-) create mode 100644 examples/Tutorial/tasks/active_learn.py create mode 100644 examples/Tutorial/tasks/simulation.py create mode 100644 examples/Tutorial/tasks/training.py diff --git a/examples/Tutorial/tasks/active_learn.py b/examples/Tutorial/tasks/active_learn.py new file mode 100644 index 0000000..51e9a8b --- /dev/null +++ b/examples/Tutorial/tasks/active_learn.py @@ -0,0 +1,66 @@ +#!/usr/bin/env python + +import io, os, sys, socket +import time +import argparse +import wfMiniAPI.kernel as kernel + +def parse_args(): + parser = argparse.ArgumentParser(description='Active Learning') + + # AL parameters + parser.add_argument('--num_sample', type=int, default=65536, + help='number of samples to evaluate uncertainty (default: 65536)') + parser.add_argument('--batch_size', type=int, default=64, + help='batch size in training') + parser.add_argument('--device', default='gpu', + help='Wheter this is running on cpu or gpu') + parser.add_argument('--dense_dim_in', type=int, default=2048, + help='dim for most heavy dense layer, input') + parser.add_argument('--dense_dim_out', type=int, default=512, + help='dim for most heavy dense layer, output') + parser.add_argument('--top_k', type=int, default=4, + help='the number of points return with biggest uncertainty') + + # Tuning knobs + parser.add_argument('--num_mult', type=int, default=5, + help='number of matrix mult to perform') + + # Task related parameters + parser.add_argument('--experiment_dir', required=True, + help='the root dir of gsas output data') + parser.add_argument('--task_index', required=True, + help='the index of task to prevent colliding') + + args = parser.parse_args() + + return args + +def main(): + + start_time = time.time() + + args = parse_args() + print(args) + + root_path = args.experiment_dir + '/{}'.format(args.task_index) + '/' + print("root_path for data = ", root_path) + + num_batch = args.num_sample // args.batch_size + for _ in range(num_batch): + kernel.dataCopyH2D(args.batch_size * args.dense_dim_in) + for ii in range(args.num_mult): + kernel.matMulGeneral(args.device, [args.batch_size, args.dense_dim_in], [args.dense_dim_in, args.dense_dim_out], ([1], [0])) + kernel.axpy(args.device, args.dense_dim_in * args.dense_dim_out) + kernel.top_k(args.device, args.num_sample, args.top_k) + if args.device == 'gpu': + kernel.dataCopyH2D(args.top_k * 2) + + os.makedirs(root_path, exist_ok=True) + file_name = os.path.join(root_path, "result.txt") + kernel.writeSingleRank(args.top_k * 2, file_name) + end_time = time.time() + print("Total running time is {} seconds".format(end_time - start_time)) + +if __name__ == '__main__': + main() diff --git a/examples/Tutorial/tasks/simulation.py b/examples/Tutorial/tasks/simulation.py new file mode 100644 index 0000000..4c2f2b2 --- /dev/null +++ b/examples/Tutorial/tasks/simulation.py @@ -0,0 +1,81 @@ +#!/usr/bin/env python + +import io, os, sys, socket +import time +import argparse +import wfMiniAPI.kernel as kernel + +def parse_args(): + + parser = argparse.ArgumentParser(description="Molecular Dynamics Simulation Configuration") + + # Simulation parameters + parser.add_argument('--N_atoms', type=int, default=10000, + help='Number of particles (default: 10000)') + parser.add_argument('--grid_size', type=int, default=64, + help='PME grid size (default: 64)') + parser.add_argument('--neighbor_freq', type=int, default=10, + help='Steps between neighboring list updates (default: 10)') + parser.add_argument('--log_freq', type=int, default=20, + help='Steps between logging outputs (default: 20)') + parser.add_argument('--n_steps', type=int, default=100, + help='Total MD steps to emulate (default: 100)') + parser.add_argument('--device', type=str, default='cpu', choices=['cpu', 'gpu'], + help='Device to run the simulation on (default: cpu)') + + # Tuning knobs + parser.add_argument('--n_force', type=int, default=100, + help='Short-range AXPY (default: 100)') + parser.add_argument('--n_fft', type=int, default=2, + help='Number of forward+inverse FFTs (default: 2)') + parser.add_argument('--n_int', type=int, default=20, + help='Integration AXPY count (default: 20)') + parser.add_argument('--bytes_per_atom', type=int, default=144, + help='Number of bytes per atom (default: 144)') + + # Task related parameters + parser.add_argument('--experiment_dir', required=True, + help='the root dir of gsas output data') + parser.add_argument('--task_index', required=True, + help='the index of task to prevent colliding') + + args = parser.parse_args() + args.io_bytes = args.N_atoms * args.bytes_per_atom + + return args + +def main(): + + start_time = time.time() + + args = parse_args() + print(args) + + root_path = args.experiment_dir + '/{}'.format(args.task_index) + '/' + print("root_path for data = ", root_path) + + kernel.generateRandomNumber(args.device, args.N_atoms * 3 * 3) + print("After Initialization takes ", time.time() - start_time) + + for i in range(args.n_steps): + if i % args.neighbor_freq == 0: + kernel.matMulGeneral(args.device, size_a=(args.N_atoms,3), size_b=(3, args.N_atoms), axis=1) + for j in range(args.n_force): + kernel.axpy_fast(args.device, 3*args.N_atoms) + for j in range(args.n_fft): + kernel.fftn(args.device, (args.grid_size, args.grid_size, args.grid_size), 'complexF', (0,1,2)) + for j in range(args.n_int): + kernel.axpy_fast(args.device, 3*args.N_atoms) + if i % log_freq == 0: + dir_name = os.path.join(root_path, f"./log_step_{i}") + os.makedirs(dir_name, exist_ok=True) + kernel.writeSingleRank(io_bytes, dir_name) + print("After main loop takes ", time.time() - start_time) + + if args.device == 'gpu': + kernel.dataCopyD2H(args.N_atoms * 3 * 3) + print("Total running time is {} seconds".format(end_time - start_time)) + +if __name__ == '__main__': + main() + diff --git a/examples/Tutorial/tasks/training.py b/examples/Tutorial/tasks/training.py new file mode 100644 index 0000000..b5f81f6 --- /dev/null +++ b/examples/Tutorial/tasks/training.py @@ -0,0 +1,77 @@ +#!/usr/bin/env python + +import io, os, sys, socket +import time +import argparse +import wfMiniAPI.kernel as kernel + +def parse_args(): + parser = argparse.ArgumentParser(description='Bayasian ML Training') + + # Training parameters + parser.add_argument('--num_epochs', type=int, default=200, + help='number of epochs to train (default: 200)') + parser.add_argument('--num_sample', type=int, default=512, + help='num of samples in matrix mult') + parser.add_argument('--batch_size', type=int, default=64, + help='batch size in training') + parser.add_argument('--device', default='gpu', + help='Wheter this is running on cpu or gpu') + parser.add_argument('--dense_dim_in', type=int, default=2048, + help='dim for most heavy dense layer, input') + parser.add_argument('--dense_dim_out', type=int, default=512, + help='dim for most heavy dense layer, output') + parser.add_argument('--log_freq', type=int, default=20, + help='epochs between logging outputs (default: 20)') + + # Tuning knobs + parser.add_argument('--write_size', type=int, default=0, + help='size of bytes written to disk') + parser.add_argument('--num_mult', type=int, default=10, + help='number of matrix mult to perform') + + # Task related parameters + parser.add_argument('--experiment_dir', required=True, + help='the root dir of gsas output data') + parser.add_argument('--task_index', required=True, + help='the index of task to prevent colliding') + + args = parser.parse_args() + + return args + + +def main(): + + start_time = time.time() + + args = parse_args() + print(args) + + root_path = args.experiment_dir + '/{}'.format(args.task_index) + '/' + print("root_path for data = ", root_path) + + kernel.generateRandomNumber(args.device, args.dense_dim_in * args.dense_dim_out) + if args.device == 'gpu': + kernel.dataCopyH2D(args.dense_dim_in * args.dense_dim_out) + + for epoch in range(args.num_epochs): + num_batch = args.num_sample // args.batch_size + for _ in range(num_batch): + kernel.dataCopyH2D(args.batch_size * args.dense_dim_in) + for ii in range(args.num_mult): + kernel.matMulGeneral(args.device, [args.batch_size, args.dense_dim_in], [args.dense_dim_in, args.dense_dim_out], ([1], [0])) + kernel.axpy(args.device, args.dense_dim_in * args.dense_dim_out) + if epoch % args.log_freq == 0: + dir_name = os.path.join(root_path, f"./epoch_{epoch}") + os.makedirs(dir_name, exist_ok=True) + kernel.writeSingleRank(args.write_size, dir_name) + + if args.device == 'gpu': + kernel.dataCopyD2H(args.dense_dim_in * args.dense_dim_out) + + end_time = time.time() + print("Total running time is {} seconds".format(end_time - start_time)) + +if __name__ == '__main__': + main() diff --git a/wfMiniAPI/src/wfMiniAPI/kernel.py b/wfMiniAPI/src/wfMiniAPI/kernel.py index af28c65..bb9d0d1 100644 --- a/wfMiniAPI/src/wfMiniAPI/kernel.py +++ b/wfMiniAPI/src/wfMiniAPI/kernel.py @@ -294,11 +294,11 @@ def axpy(device, size): y = xp.empty(size, dtype=xp.float32) y += 1.01 * x -#_axpy_inplace = cp.ElementwiseKernel( +#_axpy_fuse = cp.ElementwiseKernel( # 'float32 alpha, raw float32 x', # 'raw float32 y', # 'y[i] += alpha * x[i]', -# 'axpy_inplace_kernel', +# 'axpy_fuse_kernel', # no_return=True #) # @@ -310,10 +310,28 @@ def axpy(device, size): # if xp == np: # y += 1.01 * x # elif xp == cp: -# _axpy_inplace(1.01, x, y, size=size) +# _axpy_fuse(1.01, x, y, size=size) + +_axpy_fuse_fast = cp.ElementwiseKernel( + 'float32 alpha, raw float32 x', + 'raw float32 y', + 'y[i] += alpha * x[i]', + 'axpy_fuse_kernel', + no_return=True +) @annotate_kernel -def implaceCompute(device, size, num_op, op): +def axpy_fast(device, size): + xp = get_device_module(device) + x = xp.empty(size, dtype=xp.float32) + y = xp.empty(size, dtype=xp.float32) + if xp == np: + y += 1.01 * x + elif xp == cp: + _axpy_fuse_fast(1.01, x, y, size=size) + +@annotate_kernel +def inplaceCompute(device, size, num_op, op): xp = get_device_module(device) x = xp.empty(size, dtype=xp.float32) if isinstance(op, str): @@ -343,3 +361,11 @@ def scatterAdd(device, x_size, y_size): y += x[idx] elif xp == cp: cp.add.at(y, idx, x) + +@annotate_kernel +def top_k(device, size, k): + xp = get_device_module(device) + arr = xp.empty(size, dtype=xp.float32) + indices = xp.argpartition(-arr, k)[:k] + sorted_indices = indices[xp.argsort(-arr[indices])] + top_values = arr[sorted_indices] From 7ebf01872d0428f1952d9fcd2f6b572ff08901f9 Mon Sep 17 00:00:00 2001 From: GKNB Date: Tue, 29 Jul 2025 10:56:35 +0000 Subject: [PATCH 27/46] fast axpy --- examples/Tutorial/tasks/active_learn.py | 8 ++++---- examples/Tutorial/tasks/simulation.py | 6 ++++-- examples/Tutorial/tasks/training.py | 2 +- 3 files changed, 9 insertions(+), 7 deletions(-) diff --git a/examples/Tutorial/tasks/active_learn.py b/examples/Tutorial/tasks/active_learn.py index 51e9a8b..6d7f359 100644 --- a/examples/Tutorial/tasks/active_learn.py +++ b/examples/Tutorial/tasks/active_learn.py @@ -51,14 +51,14 @@ def main(): kernel.dataCopyH2D(args.batch_size * args.dense_dim_in) for ii in range(args.num_mult): kernel.matMulGeneral(args.device, [args.batch_size, args.dense_dim_in], [args.dense_dim_in, args.dense_dim_out], ([1], [0])) - kernel.axpy(args.device, args.dense_dim_in * args.dense_dim_out) + kernel.axpy_fast(args.device, args.dense_dim_in * args.dense_dim_out) kernel.top_k(args.device, args.num_sample, args.top_k) if args.device == 'gpu': kernel.dataCopyH2D(args.top_k * 2) - os.makedirs(root_path, exist_ok=True) - file_name = os.path.join(root_path, "result.txt") - kernel.writeSingleRank(args.top_k * 2, file_name) + dir_name = os.path.join(root_path, "result") + os.makedirs(dir_name, exist_ok=True) + kernel.writeSingleRank(args.top_k * 2, dir_name) end_time = time.time() print("Total running time is {} seconds".format(end_time - start_time)) diff --git a/examples/Tutorial/tasks/simulation.py b/examples/Tutorial/tasks/simulation.py index 4c2f2b2..c3fb78e 100644 --- a/examples/Tutorial/tasks/simulation.py +++ b/examples/Tutorial/tasks/simulation.py @@ -66,14 +66,16 @@ def main(): kernel.fftn(args.device, (args.grid_size, args.grid_size, args.grid_size), 'complexF', (0,1,2)) for j in range(args.n_int): kernel.axpy_fast(args.device, 3*args.N_atoms) - if i % log_freq == 0: + if i % args.log_freq == 0: dir_name = os.path.join(root_path, f"./log_step_{i}") os.makedirs(dir_name, exist_ok=True) - kernel.writeSingleRank(io_bytes, dir_name) + kernel.writeSingleRank(args.io_bytes, dir_name) print("After main loop takes ", time.time() - start_time) if args.device == 'gpu': kernel.dataCopyD2H(args.N_atoms * 3 * 3) + + end_time = time.time() print("Total running time is {} seconds".format(end_time - start_time)) if __name__ == '__main__': diff --git a/examples/Tutorial/tasks/training.py b/examples/Tutorial/tasks/training.py index b5f81f6..729d0e0 100644 --- a/examples/Tutorial/tasks/training.py +++ b/examples/Tutorial/tasks/training.py @@ -61,7 +61,7 @@ def main(): kernel.dataCopyH2D(args.batch_size * args.dense_dim_in) for ii in range(args.num_mult): kernel.matMulGeneral(args.device, [args.batch_size, args.dense_dim_in], [args.dense_dim_in, args.dense_dim_out], ([1], [0])) - kernel.axpy(args.device, args.dense_dim_in * args.dense_dim_out) + kernel.axpy_fast(args.device, args.dense_dim_in * args.dense_dim_out) if epoch % args.log_freq == 0: dir_name = os.path.join(root_path, f"./epoch_{epoch}") os.makedirs(dir_name, exist_ok=True) From 6625ada1f441e0e13c0ec45876bf5d03ce464205 Mon Sep 17 00:00:00 2001 From: iznoanygod Date: Fri, 5 Sep 2025 03:34:13 -0400 Subject: [PATCH 28/46] staging --- examples/motif_4/config.yaml | 36 ++++++++++++++ examples/motif_4/miniapp.py | 53 +++++++++++++++++++++ examples/motif_4/stages/assimilate_stage.py | 35 ++++++++++++++ examples/motif_4/stages/control_stage.py | 25 ++++++++++ examples/motif_4/stages/infer_stage.py | 43 +++++++++++++++++ examples/motif_4/stages/sim_stage.py | 48 +++++++++++++++++++ examples/motif_4/stages/train_stage.py | 36 ++++++++++++++ examples/motif_4/utils/metrics.py | 8 ++++ 8 files changed, 284 insertions(+) create mode 100644 examples/motif_4/config.yaml create mode 100644 examples/motif_4/miniapp.py create mode 100644 examples/motif_4/stages/assimilate_stage.py create mode 100644 examples/motif_4/stages/control_stage.py create mode 100644 examples/motif_4/stages/infer_stage.py create mode 100644 examples/motif_4/stages/sim_stage.py create mode 100644 examples/motif_4/stages/train_stage.py create mode 100644 examples/motif_4/utils/metrics.py diff --git a/examples/motif_4/config.yaml b/examples/motif_4/config.yaml new file mode 100644 index 0000000..dc26d93 --- /dev/null +++ b/examples/motif_4/config.yaml @@ -0,0 +1,36 @@ +simulation: + steps: 100 + state_dim: 1024 + read_size_bytes: 1048576 # 1 MiB pre-step read + write_size_bytes: 1048576 # 1 MiB post-step write + device: "cpu" # "cpu" or "gpu" + sleep_ms: 0 + +assimilation: + window: 4 + device: "cpu" + emit_every: 1 + +training: + device: "cpu" + train_every: 5 + model_dim: 512 + use_collective: false + collective_bytes: 4194304 # emulate AllReduce volume + +inference: + device: "cpu" + copy_bytes: 2097152 + vec_len: 1000000 + +control: + device: "cpu" + policy_cost: 65536 # ~256x256 matmul + +queues: + sim_to_assim: 8 + assim_to_train: 8 + assim_to_infer: 8 + train_to_infer: 4 + infer_to_control: 8 + control_to_sim: 4 \ No newline at end of file diff --git a/examples/motif_4/miniapp.py b/examples/motif_4/miniapp.py new file mode 100644 index 0000000..209366a --- /dev/null +++ b/examples/motif_4/miniapp.py @@ -0,0 +1,53 @@ +import argparse, asyncio, os, time, json +from pathlib import Path + +from stages.sim_stage import SimulationStage +from stages.assimilate_stage import AssimilationStage +from stages.train_stage import TrainingStage +from stages.infer_stage import InferenceStage +from stages.control_stage import ControlStage +from utils.metrics import Timer, log +import yaml + +async def run_pipeline(cfg): + q_sim_to_assim = asyncio.Queue(maxsize=cfg["queues"]["sim_to_assim"]) + q_assim_to_train = asyncio.Queue(maxsize=cfg["queues"]["assim_to_train"]) + q_assim_to_infer = asyncio.Queue(maxsize=cfg["queues"]["assim_to_infer"]) + q_train_to_infer = asyncio.Queue(maxsize=cfg["queues"]["train_to_infer"]) + q_infer_to_control = asyncio.Queue(maxsize=cfg["queues"]["infer_to_control"]) + q_control_to_sim = asyncio.Queue(maxsize=cfg["queues"]["control_to_sim"]) + + # Stages + sim = SimulationStage(cfg) + assim = AssimilationStage(cfg) + train = TrainingStage(cfg) + infer = InferenceStage(cfg) + control = ControlStage(cfg) + + # Tasks (concurrent, real-time style) + tasks = [ + asyncio.create_task(sim.run(q_out=q_sim_to_assim, q_in=q_control_to_sim)), + asyncio.create_task(assim.run(q_in=q_sim_to_assim, + q_out_train=q_assim_to_train, + q_out_infer=q_assim_to_infer)), + asyncio.create_task(train.run(q_in=q_assim_to_train, q_out=q_train_to_infer)), + asyncio.create_task(infer.run(q_state=q_assim_to_infer, + q_model=q_train_to_infer, + q_out=q_infer_to_control)), + asyncio.create_task(control.run(q_in=q_infer_to_control, q_out=q_control_to_sim)), + ] + + await asyncio.gather(*tasks) + +def load_cfg(path): + with open(path, "r") as f: + return yaml.safe_load(f) + +if __name__ == "__main__": + ap = argparse.ArgumentParser(description="Digital Twin (Motif 4) Mini-App") + ap.add_argument("--config", type=str, default="config.yaml") + args = ap.parse_args() + cfg = load_cfg(args.config) + t = Timer().start() + asyncio.run(run_pipeline(cfg)) + log(f"DONE in {t.stop():.3f}s") diff --git a/examples/motif_4/stages/assimilate_stage.py b/examples/motif_4/stages/assimilate_stage.py new file mode 100644 index 0000000..94d03c4 --- /dev/null +++ b/examples/motif_4/stages/assimilate_stage.py @@ -0,0 +1,35 @@ +import asyncio, collections +from wfMiniAPI import kernel as kern +from utils.metrics import log + +class AssimilationStage: + """ + Windowed assimilation (e.g., reductions / filtering). + Emits to both training (for model updates) and inference (for online predictions). + """ + def __init__(self, cfg): + a = cfg["assimilation"] + self.window = a["window"] + self.device = a["device"] + self.buf = collections.deque(maxlen=self.window) + self.emit_every = a.get("emit_every", 1) + + async def run(self, q_in, q_out_train, q_out_infer): + n = 0 + while True: + msg = await q_in.get() + if msg is None: + await q_out_train.put(None) + await q_out_infer.put(None) + log("AssimilationStage: complete") + return + + self.buf.append(msg["state_dim"]) + data_size = max(1, sum(self.buf)) # scale only + kern.reduction(device=self.device, data_size=data_size) + + if n % self.emit_every == 0: + state_summary = {"t": msg["t"], "feat": data_size} + await q_out_train.put(state_summary) + await q_out_infer.put(state_summary) + n += 1 diff --git a/examples/motif_4/stages/control_stage.py b/examples/motif_4/stages/control_stage.py new file mode 100644 index 0000000..2bf4c53 --- /dev/null +++ b/examples/motif_4/stages/control_stage.py @@ -0,0 +1,25 @@ +import asyncio, math +from wfMiniAPI import kernel as kern +from utils.metrics import log + +class ControlStage: + """ + Takes predictions and computes control signals to steer the sim. + """ + def __init__(self, cfg): + c = cfg["control"] + self.device = c["device"] + self.policy_cost = c.get("policy_cost", 256*256) + + async def run(self, q_in, q_out): + while True: + pred = await q_in.get() + if pred is None: + await q_out.put(None) + log("ControlStage: complete") + return + # emulate tiny policy compute + dim = int(math.sqrt(self.policy_cost)) + kern.matMulSimple2D(device=self.device, dim=dim) + ctrl = {"gain": 0.01, "model_ver": pred["model_ver"], "t": pred["t"]} + await q_out.put(ctrl) diff --git a/examples/motif_4/stages/infer_stage.py b/examples/motif_4/stages/infer_stage.py new file mode 100644 index 0000000..32c5aaa --- /dev/null +++ b/examples/motif_4/stages/infer_stage.py @@ -0,0 +1,43 @@ +import asyncio +from wfMiniAPI import kernel as kern +from utils.metrics import log + +class InferenceStage: + """ + Online inference using the surrogate. + Models memory copies + axpy-like compute. + """ + def __init__(self, cfg): + inf = cfg["inference"] + self.device = inf["device"] + self.copy_bytes = inf.get("copy_bytes", 2 * 1024 * 1024) + self.vec_len = inf.get("vec_len", 1_000_000) + + async def run(self, q_state, q_model, q_out): + current_model = {"version": -1, "dim": 1} + + async def model_updater(): + nonlocal current_model + while True: + m = await q_model.get() + if m is None: + return + current_model = m + + updater_task = asyncio.create_task(model_updater()) + + while True: + s = await q_state.get() + if s is None: + await q_out.put(None) + updater_task.cancel() + log("InferenceStage: complete") + return + + # emulate H2D/D2H copies around inference + kern.dataCopyH2D(data_size=self.copy_bytes) + kern.axpy(device=self.device, data_size=self.vec_len) # y = a*x + y (emulated) + kern.dataCopyD2H(data_size=self.copy_bytes) + + pred = {"t": s["t"], "model_ver": current_model["version"], "health": 1.0} + await q_out.put(pred) diff --git a/examples/motif_4/stages/sim_stage.py b/examples/motif_4/stages/sim_stage.py new file mode 100644 index 0000000..adc482c --- /dev/null +++ b/examples/motif_4/stages/sim_stage.py @@ -0,0 +1,48 @@ +import asyncio, time, math, random +from wfMiniAPI import kernel as kern +from utils.metrics import log + +class SimulationStage: + """ + Emulates a time-stepping simulation producing state snapshots. + Uses matMulSimple2D + read/write emulation to capture CPU/GPU + I/O. + Accepts control updates (e.g., parameter steering) from the controller. + """ + def __init__(self, cfg): + sim = cfg["simulation"] + self.steps = sim["steps"] + self.state_dim = sim["state_dim"] + self.read_size = sim["read_size_bytes"] + self.write_size = sim["write_size_bytes"] + self.device = sim["device"] + self.sleep_ms = sim.get("sleep_ms", 0) + + async def run(self, q_out, q_in): + for step in range(self.steps): + # apply pending control update (non-blocking) + ctrl = None + try: + ctrl = q_in.get_nowait() + except asyncio.QueueEmpty: + pass + + # read/input emulate + kern.readNonMPI(nbytes=self.read_size) + + # compute: simulate main step + kern.RNG(device=self.device, data_size=self.state_dim) + kern.matMulSimple2D(device=self.device, dim=self.state_dim) + + # write/output emulate + kern.writeNonMPI(nbytes=self.write_size) + + # emit "state" token (size metadata; payload omitted for lightness) + msg = {"t": step, "state_dim": self.state_dim, "ctrl": ctrl} + await q_out.put(msg) + + if self.sleep_ms > 0: + await asyncio.sleep(self.sleep_ms / 1000.0) + + # signal completion + await q_out.put(None) + log("SimulationStage: complete") diff --git a/examples/motif_4/stages/train_stage.py b/examples/motif_4/stages/train_stage.py new file mode 100644 index 0000000..c94a22b --- /dev/null +++ b/examples/motif_4/stages/train_stage.py @@ -0,0 +1,36 @@ +import asyncio +from wfMiniAPI import kernel as kern +from utils.metrics import log + +class TrainingStage: + """ + Periodic surrogate training step. + Uses matMulGeneral + (optional) MPIallReduce to emulate data-parallel training. + """ + def __init__(self, cfg): + tr = cfg["training"] + self.device = tr["device"] + self.train_every = tr["train_every"] + self.model_dim = tr["model_dim"] + self.use_collective = tr.get("use_collective", False) + self.collective_bytes = tr.get("collective_bytes", 4 * 1024 * 1024) + + async def run(self, q_in, q_out): + k = 0 + while True: + msg = await q_in.get() + if msg is None: + await q_out.put(None) + log("TrainingStage: complete") + return + if k % self.train_every == 0: + # emulate backprop compute + kern.matMulGeneral(device=self.device, dim_list=[self.model_dim, + self.model_dim, + self.model_dim]) + # emulate gradient allreduce + if self.use_collective: + kern.MPIallReduce(device=self.device, data_size=self.collective_bytes) + model = {"version": k, "dim": self.model_dim} + await q_out.put(model) + k += 1 diff --git a/examples/motif_4/utils/metrics.py b/examples/motif_4/utils/metrics.py new file mode 100644 index 0000000..2d91aaf --- /dev/null +++ b/examples/motif_4/utils/metrics.py @@ -0,0 +1,8 @@ +import time, sys + +class Timer: + def __init__(self): self.t0 = None + def start(self): self.t0 = time.time(); return self + def stop(self): return time.time() - self.t0 + +def log(msg): print(f"[dt] {msg}", file=sys.stdout, flush=True) \ No newline at end of file From e2bc5f47546355053fdfb535444ade71f3128f42 Mon Sep 17 00:00:00 2001 From: iznoanygod Date: Fri, 5 Sep 2025 11:57:05 -0400 Subject: [PATCH 29/46] inverse design --- examples/motif_3/config.yaml | 32 ++++++++++++ examples/motif_3/miniapp.py | 43 ++++++++++++++++ examples/motif_3/stages/evaluate_stage.py | 48 ++++++++++++++++++ examples/motif_3/stages/propose_stage.py | 47 +++++++++++++++++ examples/motif_3/stages/select_stage.py | 62 +++++++++++++++++++++++ examples/motif_3/stages/train_stage.py | 38 ++++++++++++++ examples/motif_3/utils/metrics.py | 6 +++ 7 files changed, 276 insertions(+) create mode 100644 examples/motif_3/config.yaml create mode 100644 examples/motif_3/miniapp.py create mode 100644 examples/motif_3/stages/evaluate_stage.py create mode 100644 examples/motif_3/stages/propose_stage.py create mode 100644 examples/motif_3/stages/select_stage.py create mode 100644 examples/motif_3/stages/train_stage.py create mode 100644 examples/motif_3/utils/metrics.py diff --git a/examples/motif_3/config.yaml b/examples/motif_3/config.yaml new file mode 100644 index 0000000..564d77b --- /dev/null +++ b/examples/motif_3/config.yaml @@ -0,0 +1,32 @@ +propose: + rounds: 10 + batch_size: 16 + device: "cpu" # "cpu" or "gpu" + vec_len: 1000000 # emulate vectorized param proposals + copy_bytes: 2097152 # 2 MiB H2D/D2H around proposals + +evaluate: + device: "cpu" + read_size_bytes: 1048576 + write_size_bytes: 1048576 + matmul_dim: 512 + workers: 4 # parallel evaluators + +training: + device: "cpu" + model_dim: 512 + train_every: 8 + use_collective: false + collective_bytes: 4194304 + +select: + device: "cpu" + top_k: 5 + emit_every: 8 + policy_cost: 65536 # ~256x256 matmul for re-scoring + +queues: + candidates: 64 + results: 128 + models: 16 + feedback: 32 \ No newline at end of file diff --git a/examples/motif_3/miniapp.py b/examples/motif_3/miniapp.py new file mode 100644 index 0000000..9986211 --- /dev/null +++ b/examples/motif_3/miniapp.py @@ -0,0 +1,43 @@ +import argparse, asyncio, yaml +from stages.propose_stage import ProposeStage +from stages.evaluate_stage import EvaluateStage +from stages.train_stage import TrainingStage +from stages.select_stage import SelectStage +from utils.metrics import Timer, log + +async def run_pipeline(cfg): + # Queues wire the inverse-design loop + q_candidates = asyncio.Queue(maxsize=cfg["queues"]["candidates"]) + q_results_train = asyncio.Queue(maxsize=cfg["queues"]["results"]) + q_results_select = asyncio.Queue(maxsize=cfg["queues"]["results"]) + q_model_updates = asyncio.Queue(maxsize=cfg["queues"]["models"]) + q_feedback = asyncio.Queue(maxsize=cfg["queues"]["feedback"]) + + propose = ProposeStage(cfg) + evaluate = EvaluateStage(cfg) + train = TrainingStage(cfg) + select = SelectStage(cfg) + + tasks = [ + asyncio.create_task(propose.run(q_out=q_candidates, q_fb=q_feedback)), + asyncio.create_task(evaluate.run(q_in=q_candidates, + q_out_train=q_results_train, + q_out_select=q_results_select)), + asyncio.create_task(train.run(q_in=q_results_train, q_out=q_model_updates)), + asyncio.create_task(select.run(q_in_results=q_results_select, + q_in_model=q_model_updates, + q_out_fb=q_feedback)), + ] + await asyncio.gather(*tasks) + +def load_cfg(path): + with open(path, "r") as f: return yaml.safe_load(f) + +if __name__ == "__main__": + ap = argparse.ArgumentParser(description="Inverse Design (Motif 3) Mini-App") + ap.add_argument("--config", type=str, default="config.yaml") + args = ap.parse_args() + cfg = load_cfg(args.config) + t = Timer().start() + asyncio.run(run_pipeline(cfg)) + log(f"DONE in {t.stop():.3f}s") \ No newline at end of file diff --git a/examples/motif_3/stages/evaluate_stage.py b/examples/motif_3/stages/evaluate_stage.py new file mode 100644 index 0000000..ddc982c --- /dev/null +++ b/examples/motif_3/stages/evaluate_stage.py @@ -0,0 +1,48 @@ +import asyncio, math, random +from wfMiniAPI import kernel as kern +from utils.metrics import log + +class EvaluateStage: + """ + Embarrassingly-parallel evaluation pool: + - read/write emulation (I/O) + - main compute via matMulSimple2D (dominant) + Produces result records sent to both training and selection branches. + """ + def __init__(self, cfg): + e = cfg["evaluate"] + self.device = e["device"] + self.read_size = e["read_size_bytes"] + self.write_size = e["write_size_bytes"] + self.matmul_dim = e["matmul_dim"] + self.workers = e.get("workers", 4) + + async def worker(self, name, q_in, q_out_train, q_out_select): + while True: + cand = await q_in.get() + if cand is None: + # propagate sentinel to other workers and downstream + await q_in.put(None) + await q_out_train.put(None) + await q_out_select.put(None) + log(f"EvaluateStage[{name}]: complete") + return + + # emulate evaluation + kern.readNonMPI(nbytes=self.read_size) + kern.RNG(device=self.device, data_size=cand["vec_len"]) + kern.matMulSimple2D(device=self.device, dim=self.matmul_dim) + kern.writeNonMPI(nbytes=self.write_size) + + # fake objective: lower is better + obj = 1.0 / (1.0 + self.matmul_dim) + random.random() * 0.01 + res = {"rid": cand["rid"], "cid": cand["cid"], "objective": obj} + + await q_out_train.put(res) + await q_out_select.put(res) + + async def run(self, q_in, q_out_train, q_out_select): + tasks = [asyncio.create_task( + self.worker(f"W{i}", q_in, q_out_train, q_out_select)) + for i in range(self.workers)] + await asyncio.gather(*tasks) \ No newline at end of file diff --git a/examples/motif_3/stages/propose_stage.py b/examples/motif_3/stages/propose_stage.py new file mode 100644 index 0000000..a6a8844 --- /dev/null +++ b/examples/motif_3/stages/propose_stage.py @@ -0,0 +1,47 @@ +import asyncio, random +from wfMiniAPI import kernel as kern +from utils.metrics import log + +class ProposeStage: + """ + Proposes candidate designs in rounds; consumes feedback to bias proposals. + Emulates proposal compute & data staging via RNG + H2D/D2H copies. + """ + def __init__(self, cfg): + p = cfg["propose"] + self.rounds = p["rounds"] + self.batch = p["batch_size"] + self.device = p["device"] + self.vec_len = p.get("vec_len", 1_000_000) + self.copy_bytes = p.get("copy_bytes", 2 * 1024 * 1024) + + async def run(self, q_out, q_fb): + for r in range(self.rounds): + # consume *latest* feedback if available (non-blocking) + fb_latest = None + try: + while True: + fb_latest = q_fb.get_nowait() + if fb_latest is None: + # early termination signal + await q_out.put(None) + log("ProposeStage: early stop") + return + except asyncio.QueueEmpty: + pass + + for i in range(self.batch): + # emulate cheap generator compute+copies + kern.RNG(device=self.device, data_size=self.vec_len) + kern.dataCopyH2D(data_size=self.copy_bytes) + kern.dataCopyD2H(data_size=self.copy_bytes) + cand = { + "rid": r, "cid": i, + "params": {"seed": random.random()}, + "vec_len": self.vec_len + } + await q_out.put(cand) + + # finished proposing + await q_out.put(None) + log("ProposeStage: complete") \ No newline at end of file diff --git a/examples/motif_3/stages/select_stage.py b/examples/motif_3/stages/select_stage.py new file mode 100644 index 0000000..48d1ef1 --- /dev/null +++ b/examples/motif_3/stages/select_stage.py @@ -0,0 +1,62 @@ +import asyncio, heapq +from wfMiniAPI import kernel as kern +from utils.metrics import log + +class SelectStage: + """ + Maintains top-K designs and periodically sends feedback to the proposer. + Consumes (a) evaluation results and (b) model updates (optional). + """ + def __init__(self, cfg): + s = cfg["select"] + self.top_k = s["top_k"] + self.emit_every = s.get("emit_every", 1) + self.device = s["device"] + self.policy_cost = s.get("policy_cost", 256 * 256) + + # min-heap storing (-score, rid, cid) so we keep best (lowest objective) + self.heap = [] + + def _push(self, res): + score = res["objective"] + item = (score, res["rid"], res["cid"]) + if len(self.heap) < self.top_k: + heapq.heappush(self.heap, (-score, item)) + else: + worst = self.heap[0] + if -score > worst[0]: + heapq.heapreplace(self.heap, (-score, item)) + + async def run(self, q_in_results, q_in_model, q_out_fb): + # model updates are optional; listen asynchronously + async def _model_listener(): + while True: + m = await q_in_model.get() + if m is None: return + # tiny compute to emulate re-scoring policy + kern.matMulSimple2D(device=self.device, + dim=int(self.policy_cost ** 0.5)) + + listener = asyncio.create_task(_model_listener()) + + n = 0 + while True: + res = await q_in_results.get() + if res is None: + await q_out_fb.put(None) # tell proposer we’re done + listener.cancel() + log("SelectStage: complete") + return + + self._push(res) + + if n % self.emit_every == 0 and self.heap: + # Emit feedback with current best + best_neg, (score, rid, cid) = max(self.heap), None + # heap stores (-score, item); invert again: + best = max(self.heap, key=lambda x: x[0]) + _, (top_score, top_rid, top_cid) = best + fb = {"best": {"rid": top_rid, "cid": top_cid, + "objective": top_score}} + await q_out_fb.put(fb) + n += 1 \ No newline at end of file diff --git a/examples/motif_3/stages/train_stage.py b/examples/motif_3/stages/train_stage.py new file mode 100644 index 0000000..a759226 --- /dev/null +++ b/examples/motif_3/stages/train_stage.py @@ -0,0 +1,38 @@ +import asyncio +from wfMiniAPI import kernel as kern +from utils.metrics import log + +class TrainingStage: + """ + Optional surrogate training on a stream of evaluation results. + - Compute via matMulGeneral + - Optional MPIallReduce to emulate DP sync + Emits lightweight model-version updates. + """ + def __init__(self, cfg): + tr = cfg["training"] + self.device = tr["device"] + self.model_dim = tr["model_dim"] + self.train_every = tr["train_every"] + self.use_collective = tr.get("use_collective", False) + self.collective_bytes = tr.get("collective_bytes", 4 * 1024 * 1024) + + async def run(self, q_in, q_out): + k = 0 + while True: + msg = await q_in.get() + if msg is None: + await q_out.put(None) + log("TrainingStage: complete") + return + + if k % self.train_every == 0: + kern.matMulGeneral(device=self.device, + dim_list=[self.model_dim, + self.model_dim, + self.model_dim]) + if self.use_collective: + kern.MPIallReduce(device=self.device, + data_size=self.collective_bytes) + await q_out.put({"version": k, "dim": self.model_dim}) + k += 1 \ No newline at end of file diff --git a/examples/motif_3/utils/metrics.py b/examples/motif_3/utils/metrics.py new file mode 100644 index 0000000..03184e5 --- /dev/null +++ b/examples/motif_3/utils/metrics.py @@ -0,0 +1,6 @@ +import time, sys +class Timer: + def __init__(self): self.t0 = None + def start(self): self.t0 = time.time(); return self + def stop(self): return time.time() - self.t0 +def log(msg): print(f"[id] {msg}", file=sys.stdout, flush=True) \ No newline at end of file From c41b86303fa5cd1d5f1a91a23f7e2fcab43b34d7 Mon Sep 17 00:00:00 2001 From: iznoanygod Date: Fri, 26 Sep 2025 11:54:18 -0400 Subject: [PATCH 30/46] migrating --- examples/motif_3/miniapp_async.py | 89 +++++++++++++++++++++++++++++++ 1 file changed, 89 insertions(+) create mode 100644 examples/motif_3/miniapp_async.py diff --git a/examples/motif_3/miniapp_async.py b/examples/motif_3/miniapp_async.py new file mode 100644 index 0000000..c7c97b6 --- /dev/null +++ b/examples/motif_3/miniapp_async.py @@ -0,0 +1,89 @@ +import argparse, asyncio, yaml, random +from utils.metrics import Timer, log + +from radical.asyncflow import WorkflowEngine +from radical.asyncflow import ConcurrentExecutionBackend + +from wfMiniAPI import kernel as kern + +from stages.sim_stage import SimulationStage +from stages.assimilate_stage import AssimilationStage +from stages.train_stage import TrainingStage +from stages.infer_stage import InferenceStage +from stages.control_stage import ControlStage + +from concurrent.futures import ThreadPoolExecutor + +def load_cfg(path): + with open(path, "r") as f: return yaml.safe_load(f) + +async def workflow(cfg): + backend = await ConcurrentExecutionBackend(ThreadPoolExecutor()) + flow = await WorkflowEngine.create(backend=backend) + + @flow.function_task + async def propose(): + p = cfg["propose"] + device = p["device"] + vec_len = p.get("vec_len", 1_000_000) + copy_bytes = p.get("copy_bytes", 2 * 1024 * 1024) + + kern.RNG(device=device, data_size=vec_len) + kern.dataCopyH2D(data_size=copy_bytes) + kern.dataCopyD2H(data_size=copy_bytes) + cand = { + "params": {"seed": random.random()}, + "vec_len": vec_len + } + return cand + + @flow.function_task + async def evaluate(cand): + e = cfg["evaluate"] + device = e["device"] + read_size = e["read_size_bytes"] + write_size = e["write_size_bytes"] + matmul_dim = e["matmul_dim"] + kern.readNonMPI(nbytes=read_size) + kern.RNG(device=device, data_size=cand["vec_len"]) + kern.matMulSimple2D(device=device, dim=matmul_dim) + kern.writeNonMPI(nbytes=write_size) + + # fake objective: lower is better + obj = 1.0 / (1.0 + matmul_dim) + random.random() * 0.01 + res = {"objective": obj} + return res + + @flow.function_task + async def select(): + return None + + @flow.function_task + async def train(res, k): + tr = cfg["training"] + device = tr["device"] + model_dim = tr["model_dim"] + use_collective = tr.get("use_collective", False) + collective_bytes = tr.get("collective_bytes", 4 * 1024 * 1024) + + kern.matMulGeneral(device=device, + dim_list=[model_dim, model_dim, model_dim]) + if use_collective: + kern.MPIallReduce(device=device, data_size=collective_bytes) + model_update = {"version": k, "dim": model_dim} + return model_update + + propose_t = propose() + evaluate_t = evaluate(propose_t) + train_t = await train(evaluate_t, k=1) + + await flow.shutdown() + +if __name__ == "__main__": + ap = argparse.ArgumentParser(description="Inverse Design (Motif 3) Mini-App") + ap.add_argument("--config", type=str, default="config.yaml") + args = ap.parse_args() + cfg = load_cfg(args.config) + t = Timer().start() + asyncio.run(workflow(cfg)) + log(f"DONE in {t.stop():.3f}s") \ No newline at end of file From 96720b9ae3f4b6b98187006df830be8206151950 Mon Sep 17 00:00:00 2001 From: iznoanygod Date: Fri, 26 Sep 2025 14:32:10 -0400 Subject: [PATCH 31/46] asynflow miniapps in development --- examples/motif_3/miniapp_async.py | 2 +- examples/motif_4/miniapp_async.py | 108 ++++++++++++++++++++++++++++++ 2 files changed, 109 insertions(+), 1 deletion(-) create mode 100644 examples/motif_4/miniapp_async.py diff --git a/examples/motif_3/miniapp_async.py b/examples/motif_3/miniapp_async.py index c7c97b6..61964ed 100644 --- a/examples/motif_3/miniapp_async.py +++ b/examples/motif_3/miniapp_async.py @@ -49,13 +49,13 @@ async def evaluate(cand): kern.matMulSimple2D(device=device, dim=matmul_dim) kern.writeNonMPI(nbytes=write_size) - # fake objective: lower is better obj = 1.0 / (1.0 + matmul_dim) + random.random() * 0.01 res = {"objective": obj} return res @flow.function_task async def select(): + # TODO write selection task return None @flow.function_task diff --git a/examples/motif_4/miniapp_async.py b/examples/motif_4/miniapp_async.py new file mode 100644 index 0000000..95a0514 --- /dev/null +++ b/examples/motif_4/miniapp_async.py @@ -0,0 +1,108 @@ +import argparse, asyncio, yaml, random +from utils.metrics import Timer, log + +from radical.asyncflow import WorkflowEngine +from radical.asyncflow import ConcurrentExecutionBackend + +from wfMiniAPI import kernel as kern + +from concurrent.futures import ThreadPoolExecutor + +def load_cfg(path): + with open(path, "r") as f: return yaml.safe_load(f) + +async def workflow(cfg): + backend = await ConcurrentExecutionBackend(ThreadPoolExecutor()) + flow = await WorkflowEngine.create(backend=backend) + + @flow.function_task + async def simulation(): + sim = cfg["simulation"] + steps = sim["steps"] + state_dim = sim["state_dim"] + read_size = sim["read_size_bytes"] + write_size = sim["write_size_bytes"] + device = sim["device"] + sleep_ms = sim.get("sleep_ms", 0) + kern.readNonMPI(nbytes=read_size) + + kern.RNG(device=device, data_size=state_dim) + kern.matMulSimple2D(device=device, dim=state_dim) + + kern.writeNonMPI(nbytes=write_size) + + msg = {"t": step, "state_dim": state_dim, "ctrl": ctrl} + + if sleep_ms > 0: + await asyncio.sleep(sleep_ms / 1000.0) + return msg + + @flow.function_task + async def assimilation(state): + assim = cfg["assimilation"] + read_size = assim["read_size_bytes"] + write_size = assim["write_size_bytes"] + device = assim["device"] + kern.readNonMPI(nbytes=read_size) + kern.RNG(device=device, data_size=state["state_dim"]) + kern.matMulSimple2D(device=device, dim=state["state_dim"]) + kern.writeNonMPI(nbytes=write_size) + msg = {"t": state["t"], "state_dim": state["state_dim"]} + return msg + + @flow.function_task + async def training(state): + train = cfg["training"] + read_size = train["read_size_bytes"] + write_size = train["write_size_bytes"] + device = train["device"] + kern.readNonMPI(nbytes=read_size) + kern.RNG(device=device, data_size=state["state_dim"]) + kern.matMulSimple2D(device=device, dim=state["state_dim"]) + kern.writeNonMPI(nbytes=write_size) + msg = {"t": state["t"], "model_dim": state["state_dim"] // 2} + return msg + + @flow.function_task + async def inference(state, model): + infer = cfg["inference"] + read_size = infer["read_size_bytes"] + write_size = infer["write_size_bytes"] + device = infer["device"] + kern.readNonMPI(nbytes=read_size) + kern.RNG(device=device, data_size=state["state_dim"]) + kern.matMulSimple2D(device=device, dim=state["state_dim"]) + kern.matMulSimple2D(device=device, dim=model["model_dim"]) + kern.writeNonMPI(nbytes=write_size) + msg = {"t": state["t"], "state_dim": state["state_dim"]} + return msg + + @flow.function_task + async def control(state): + ctrl = cfg["control"] + read_size = ctrl["read_size_bytes"] + write_size = ctrl["write_size_bytes"] + device = ctrl["device"] + kern.readNonMPI(nbytes=read_size) + kern.RNG(device=device, data_size=state["state_dim"]) + kern.matMulSimple2D(device=device, dim=state["state_dim"]) + kern.writeNonMPI(nbytes=write_size) + msg = {"t": state["t"], "ctrl_param": random.random()} + return msg + + sim_t = simulation() + assim_t = assimilation(sim_t) + train_t = training(assim_t) + infer_t = inference(assim_t, train_t) + control_t = await control(infer_t) + + await flow.shutdown() + +if __name__ == "__main__": + ap = argparse.ArgumentParser(description="Digital Twin (Motif 4) Mini-App") + ap.add_argument("--config", type=str, default="config.yaml") + args = ap.parse_args() + cfg = load_cfg(args.config) + t = Timer().start() + asyncio.run(workflow(cfg)) + log(f"DONE in {t.stop():.3f}s") From cc3730355b6955d40f41254fe2baf5e07fe20995 Mon Sep 17 00:00:00 2001 From: iznoanygod Date: Thu, 2 Oct 2025 19:25:48 -0400 Subject: [PATCH 32/46] dragon --- examples/motif_3/miniapp_async.py | 6 --- examples/motif_3/miniapp_dragon.py | 82 ++++++++++++++++++++++++++++++ examples/motif_4/miniapp_async.py | 2 +- 3 files changed, 83 insertions(+), 7 deletions(-) create mode 100644 examples/motif_3/miniapp_dragon.py diff --git a/examples/motif_3/miniapp_async.py b/examples/motif_3/miniapp_async.py index 61964ed..be47c4c 100644 --- a/examples/motif_3/miniapp_async.py +++ b/examples/motif_3/miniapp_async.py @@ -27,7 +27,6 @@ async def propose(): device = p["device"] vec_len = p.get("vec_len", 1_000_000) copy_bytes = p.get("copy_bytes", 2 * 1024 * 1024) - kern.RNG(device=device, data_size=vec_len) kern.dataCopyH2D(data_size=copy_bytes) kern.dataCopyD2H(data_size=copy_bytes) @@ -53,11 +52,6 @@ async def evaluate(cand): res = {"objective": obj} return res - @flow.function_task - async def select(): - # TODO write selection task - return None - @flow.function_task async def train(res, k): tr = cfg["training"] diff --git a/examples/motif_3/miniapp_dragon.py b/examples/motif_3/miniapp_dragon.py new file mode 100644 index 0000000..6057150 --- /dev/null +++ b/examples/motif_3/miniapp_dragon.py @@ -0,0 +1,82 @@ +import argparse, asyncio, yaml, random, logging +from utils.metrics import Timer, log + +from radical.asyncflow import WorkflowEngine +from radical.asyncflow import DragonExecutionBackend +from radical.asyncflow.logging import init_default_logger + +from wfMiniAPI import kernel as kern + + +from concurrent.futures import ThreadPoolExecutor + +def load_cfg(path): + with open(path, "r") as f: return yaml.safe_load(f) + +async def workflow(cfg): + logger = logging.getLogger(__name__) + init_default_logger(logging.DEBUG) + backend = await DragonExecutionBackend() + init_default_logger(logging.DEBUG) + flow = await WorkflowEngine.create(backend=backend) + + @flow.function_task + async def propose(): + p = cfg["propose"] + device = p["device"] + vec_len = p.get("vec_len", 1_000_000) + copy_bytes = p.get("copy_bytes", 2 * 1024 * 1024) + kern.RNG(device=device, data_size=vec_len) + kern.dataCopyH2D(data_size=copy_bytes) + kern.dataCopyD2H(data_size=copy_bytes) + cand = { + "params": {"seed": random.random()}, + "vec_len": vec_len + } + return cand + + @flow.function_task + async def evaluate(cand): + e = cfg["evaluate"] + device = e["device"] + read_size = e["read_size_bytes"] + write_size = e["write_size_bytes"] + matmul_dim = e["matmul_dim"] + kern.readNonMPI(nbytes=read_size) + kern.RNG(device=device, data_size=cand["vec_len"]) + kern.matMulSimple2D(device=device, dim=matmul_dim) + kern.writeNonMPI(nbytes=write_size) + + obj = 1.0 / (1.0 + matmul_dim) + random.random() * 0.01 + res = {"objective": obj} + return res + + @flow.function_task + async def train(res, k): + tr = cfg["training"] + device = tr["device"] + model_dim = tr["model_dim"] + use_collective = tr.get("use_collective", False) + collective_bytes = tr.get("collective_bytes", 4 * 1024 * 1024) + + kern.matMulGeneral(device=device, + dim_list=[model_dim, model_dim, model_dim]) + if use_collective: + kern.MPIallReduce(device=device, data_size=collective_bytes) + model_update = {"version": k, "dim": model_dim} + return model_update + + propose_t = propose() + evaluate_t = evaluate(propose_t) + train_t = await train(evaluate_t, k=1) + + await flow.shutdown() + +if __name__ == "__main__": + ap = argparse.ArgumentParser(description="Inverse Design (Motif 3) Mini-App") + ap.add_argument("--config", type=str, default="config.yaml") + args = ap.parse_args() + cfg = load_cfg(args.config) + t = Timer().start() + asyncio.run(workflow(cfg)) + log(f"DONE in {t.stop():.3f}s") \ No newline at end of file diff --git a/examples/motif_4/miniapp_async.py b/examples/motif_4/miniapp_async.py index 95a0514..fd4ec16 100644 --- a/examples/motif_4/miniapp_async.py +++ b/examples/motif_4/miniapp_async.py @@ -31,7 +31,7 @@ async def simulation(): kern.writeNonMPI(nbytes=write_size) - msg = {"t": step, "state_dim": state_dim, "ctrl": ctrl} + msg = {"state_dim": state_dim} if sleep_ms > 0: await asyncio.sleep(sleep_ms / 1000.0) From afa497c5d25dadfcae5bd4113fbc95d4c139d3d7 Mon Sep 17 00:00:00 2001 From: iznoanygod Date: Thu, 9 Oct 2025 16:28:48 -0400 Subject: [PATCH 33/46] dr --- examples/motif_3/miniapp_dragon.py | 24 +++++++++++++++++------- 1 file changed, 17 insertions(+), 7 deletions(-) diff --git a/examples/motif_3/miniapp_dragon.py b/examples/motif_3/miniapp_dragon.py index 6057150..eba4136 100644 --- a/examples/motif_3/miniapp_dragon.py +++ b/examples/motif_3/miniapp_dragon.py @@ -7,7 +7,6 @@ from wfMiniAPI import kernel as kern - from concurrent.futures import ThreadPoolExecutor def load_cfg(path): @@ -28,6 +27,7 @@ async def propose(): copy_bytes = p.get("copy_bytes", 2 * 1024 * 1024) kern.RNG(device=device, data_size=vec_len) kern.dataCopyH2D(data_size=copy_bytes) + kern.matMulSimple2D(device=device, dim=vec_len) kern.dataCopyD2H(data_size=copy_bytes) cand = { "params": {"seed": random.random()}, @@ -42,25 +42,34 @@ async def evaluate(cand): read_size = e["read_size_bytes"] write_size = e["write_size_bytes"] matmul_dim = e["matmul_dim"] - kern.readNonMPI(nbytes=read_size) + kern.readNonMPI(num_bytes=read_size, data_root_dir="./") kern.RNG(device=device, data_size=cand["vec_len"]) kern.matMulSimple2D(device=device, dim=matmul_dim) - kern.writeNonMPI(nbytes=write_size) + kern.writeNonMPI(num_bytes=write_size, data_root_dir="./") obj = 1.0 / (1.0 + matmul_dim) + random.random() * 0.01 res = {"objective": obj} return res @flow.function_task - async def train(res, k): + async def select(res): + s = cfg["select"] + device = s["device"] + policy_cost = s.get("policy_cost", 256 * 256) + kern.matMulSimple2D(device=device, size=int(policy_cost ** 0.5)) + best = {"rid": 0, "cid": 0, "objective": res["objective"]} + return {"best": best} + + @flow.function_task + async def train(sel, k): tr = cfg["training"] device = tr["device"] model_dim = tr["model_dim"] use_collective = tr.get("use_collective", False) collective_bytes = tr.get("collective_bytes", 4 * 1024 * 1024) - kern.matMulGeneral(device=device, - dim_list=[model_dim, model_dim, model_dim]) + kern.matMulGeneral(device=device, size_a=model_dim, size_b=model_dim, axis=0) + if use_collective: kern.MPIallReduce(device=device, data_size=collective_bytes) model_update = {"version": k, "dim": model_dim} @@ -68,7 +77,8 @@ async def train(res, k): propose_t = propose() evaluate_t = evaluate(propose_t) - train_t = await train(evaluate_t, k=1) + select_t = select(evaluate_t) + train_t = await train(select_t, k=1) await flow.shutdown() From 3dc2645880229a4850617abbe1b421f380b6d1db Mon Sep 17 00:00:00 2001 From: Andrew Park Date: Tue, 10 Mar 2026 17:59:50 -0400 Subject: [PATCH 34/46] Cleaning up --- examples/motif_3/config.yaml | 36 +---- examples/motif_3/miniapp.py | 43 ----- examples/motif_3/miniapp_async.py | 83 ---------- examples/motif_3/miniapp_dragon.py | 92 ----------- examples/motif_3/miniapp_parallel.py | 81 ++++++++++ examples/motif_3/miniapp_simple.py | 75 +++++++++ examples/motif_3/stages/evaluate_stage.py | 48 ------ examples/motif_3/stages/propose_stage.py | 47 ------ examples/motif_3/stages/select_stage.py | 62 ------- examples/motif_3/stages/train_stage.py | 38 ----- examples/motif_4/config.yaml | 49 +++--- examples/motif_4/miniapp.py | 53 ------ examples/motif_4/miniapp_async.py | 171 +++++++++++--------- examples/motif_4/stages/assimilate_stage.py | 35 ---- examples/motif_4/stages/control_stage.py | 25 --- examples/motif_4/stages/infer_stage.py | 43 ----- examples/motif_4/stages/sim_stage.py | 48 ------ examples/motif_4/stages/train_stage.py | 36 ----- 18 files changed, 277 insertions(+), 788 deletions(-) delete mode 100644 examples/motif_3/miniapp.py delete mode 100644 examples/motif_3/miniapp_async.py delete mode 100644 examples/motif_3/miniapp_dragon.py create mode 100644 examples/motif_3/miniapp_parallel.py create mode 100644 examples/motif_3/miniapp_simple.py delete mode 100644 examples/motif_3/stages/evaluate_stage.py delete mode 100644 examples/motif_3/stages/propose_stage.py delete mode 100644 examples/motif_3/stages/select_stage.py delete mode 100644 examples/motif_3/stages/train_stage.py delete mode 100644 examples/motif_4/miniapp.py delete mode 100644 examples/motif_4/stages/assimilate_stage.py delete mode 100644 examples/motif_4/stages/control_stage.py delete mode 100644 examples/motif_4/stages/infer_stage.py delete mode 100644 examples/motif_4/stages/sim_stage.py delete mode 100644 examples/motif_4/stages/train_stage.py diff --git a/examples/motif_3/config.yaml b/examples/motif_3/config.yaml index 564d77b..299c2e6 100644 --- a/examples/motif_3/config.yaml +++ b/examples/motif_3/config.yaml @@ -1,32 +1,12 @@ -propose: - rounds: 10 - batch_size: 16 - device: "cpu" # "cpu" or "gpu" - vec_len: 1000000 # emulate vectorized param proposals - copy_bytes: 2097152 # 2 MiB H2D/D2H around proposals - -evaluate: +simulate: device: "cpu" - read_size_bytes: 1048576 - write_size_bytes: 1048576 - matmul_dim: 512 - workers: 4 # parallel evaluators + read_size_bytes: 4294967296 + write_size_bytes: 1073741824 + matmul_dim: 8192 training: device: "cpu" - model_dim: 512 - train_every: 8 - use_collective: false - collective_bytes: 4194304 - -select: - device: "cpu" - top_k: 5 - emit_every: 8 - policy_cost: 65536 # ~256x256 matmul for re-scoring - -queues: - candidates: 64 - results: 128 - models: 16 - feedback: 32 \ No newline at end of file + read_size_bytes: 1342177280 + write_size_bytes: 512 + data_copy_size_bytes: 2097152 + matmul_dim: 8192 diff --git a/examples/motif_3/miniapp.py b/examples/motif_3/miniapp.py deleted file mode 100644 index 9986211..0000000 --- a/examples/motif_3/miniapp.py +++ /dev/null @@ -1,43 +0,0 @@ -import argparse, asyncio, yaml -from stages.propose_stage import ProposeStage -from stages.evaluate_stage import EvaluateStage -from stages.train_stage import TrainingStage -from stages.select_stage import SelectStage -from utils.metrics import Timer, log - -async def run_pipeline(cfg): - # Queues wire the inverse-design loop - q_candidates = asyncio.Queue(maxsize=cfg["queues"]["candidates"]) - q_results_train = asyncio.Queue(maxsize=cfg["queues"]["results"]) - q_results_select = asyncio.Queue(maxsize=cfg["queues"]["results"]) - q_model_updates = asyncio.Queue(maxsize=cfg["queues"]["models"]) - q_feedback = asyncio.Queue(maxsize=cfg["queues"]["feedback"]) - - propose = ProposeStage(cfg) - evaluate = EvaluateStage(cfg) - train = TrainingStage(cfg) - select = SelectStage(cfg) - - tasks = [ - asyncio.create_task(propose.run(q_out=q_candidates, q_fb=q_feedback)), - asyncio.create_task(evaluate.run(q_in=q_candidates, - q_out_train=q_results_train, - q_out_select=q_results_select)), - asyncio.create_task(train.run(q_in=q_results_train, q_out=q_model_updates)), - asyncio.create_task(select.run(q_in_results=q_results_select, - q_in_model=q_model_updates, - q_out_fb=q_feedback)), - ] - await asyncio.gather(*tasks) - -def load_cfg(path): - with open(path, "r") as f: return yaml.safe_load(f) - -if __name__ == "__main__": - ap = argparse.ArgumentParser(description="Inverse Design (Motif 3) Mini-App") - ap.add_argument("--config", type=str, default="config.yaml") - args = ap.parse_args() - cfg = load_cfg(args.config) - t = Timer().start() - asyncio.run(run_pipeline(cfg)) - log(f"DONE in {t.stop():.3f}s") \ No newline at end of file diff --git a/examples/motif_3/miniapp_async.py b/examples/motif_3/miniapp_async.py deleted file mode 100644 index be47c4c..0000000 --- a/examples/motif_3/miniapp_async.py +++ /dev/null @@ -1,83 +0,0 @@ -import argparse, asyncio, yaml, random -from utils.metrics import Timer, log - -from radical.asyncflow import WorkflowEngine -from radical.asyncflow import ConcurrentExecutionBackend - -from wfMiniAPI import kernel as kern - -from stages.sim_stage import SimulationStage -from stages.assimilate_stage import AssimilationStage -from stages.train_stage import TrainingStage -from stages.infer_stage import InferenceStage -from stages.control_stage import ControlStage - -from concurrent.futures import ThreadPoolExecutor - -def load_cfg(path): - with open(path, "r") as f: return yaml.safe_load(f) - -async def workflow(cfg): - backend = await ConcurrentExecutionBackend(ThreadPoolExecutor()) - flow = await WorkflowEngine.create(backend=backend) - - @flow.function_task - async def propose(): - p = cfg["propose"] - device = p["device"] - vec_len = p.get("vec_len", 1_000_000) - copy_bytes = p.get("copy_bytes", 2 * 1024 * 1024) - kern.RNG(device=device, data_size=vec_len) - kern.dataCopyH2D(data_size=copy_bytes) - kern.dataCopyD2H(data_size=copy_bytes) - cand = { - "params": {"seed": random.random()}, - "vec_len": vec_len - } - return cand - - @flow.function_task - async def evaluate(cand): - e = cfg["evaluate"] - device = e["device"] - read_size = e["read_size_bytes"] - write_size = e["write_size_bytes"] - matmul_dim = e["matmul_dim"] - kern.readNonMPI(nbytes=read_size) - kern.RNG(device=device, data_size=cand["vec_len"]) - kern.matMulSimple2D(device=device, dim=matmul_dim) - kern.writeNonMPI(nbytes=write_size) - - obj = 1.0 / (1.0 + matmul_dim) + random.random() * 0.01 - res = {"objective": obj} - return res - - @flow.function_task - async def train(res, k): - tr = cfg["training"] - device = tr["device"] - model_dim = tr["model_dim"] - use_collective = tr.get("use_collective", False) - collective_bytes = tr.get("collective_bytes", 4 * 1024 * 1024) - - kern.matMulGeneral(device=device, - dim_list=[model_dim, model_dim, model_dim]) - if use_collective: - kern.MPIallReduce(device=device, data_size=collective_bytes) - model_update = {"version": k, "dim": model_dim} - return model_update - - propose_t = propose() - evaluate_t = evaluate(propose_t) - train_t = await train(evaluate_t, k=1) - - await flow.shutdown() - -if __name__ == "__main__": - ap = argparse.ArgumentParser(description="Inverse Design (Motif 3) Mini-App") - ap.add_argument("--config", type=str, default="config.yaml") - args = ap.parse_args() - cfg = load_cfg(args.config) - t = Timer().start() - asyncio.run(workflow(cfg)) - log(f"DONE in {t.stop():.3f}s") \ No newline at end of file diff --git a/examples/motif_3/miniapp_dragon.py b/examples/motif_3/miniapp_dragon.py deleted file mode 100644 index eba4136..0000000 --- a/examples/motif_3/miniapp_dragon.py +++ /dev/null @@ -1,92 +0,0 @@ -import argparse, asyncio, yaml, random, logging -from utils.metrics import Timer, log - -from radical.asyncflow import WorkflowEngine -from radical.asyncflow import DragonExecutionBackend -from radical.asyncflow.logging import init_default_logger - -from wfMiniAPI import kernel as kern - -from concurrent.futures import ThreadPoolExecutor - -def load_cfg(path): - with open(path, "r") as f: return yaml.safe_load(f) - -async def workflow(cfg): - logger = logging.getLogger(__name__) - init_default_logger(logging.DEBUG) - backend = await DragonExecutionBackend() - init_default_logger(logging.DEBUG) - flow = await WorkflowEngine.create(backend=backend) - - @flow.function_task - async def propose(): - p = cfg["propose"] - device = p["device"] - vec_len = p.get("vec_len", 1_000_000) - copy_bytes = p.get("copy_bytes", 2 * 1024 * 1024) - kern.RNG(device=device, data_size=vec_len) - kern.dataCopyH2D(data_size=copy_bytes) - kern.matMulSimple2D(device=device, dim=vec_len) - kern.dataCopyD2H(data_size=copy_bytes) - cand = { - "params": {"seed": random.random()}, - "vec_len": vec_len - } - return cand - - @flow.function_task - async def evaluate(cand): - e = cfg["evaluate"] - device = e["device"] - read_size = e["read_size_bytes"] - write_size = e["write_size_bytes"] - matmul_dim = e["matmul_dim"] - kern.readNonMPI(num_bytes=read_size, data_root_dir="./") - kern.RNG(device=device, data_size=cand["vec_len"]) - kern.matMulSimple2D(device=device, dim=matmul_dim) - kern.writeNonMPI(num_bytes=write_size, data_root_dir="./") - - obj = 1.0 / (1.0 + matmul_dim) + random.random() * 0.01 - res = {"objective": obj} - return res - - @flow.function_task - async def select(res): - s = cfg["select"] - device = s["device"] - policy_cost = s.get("policy_cost", 256 * 256) - kern.matMulSimple2D(device=device, size=int(policy_cost ** 0.5)) - best = {"rid": 0, "cid": 0, "objective": res["objective"]} - return {"best": best} - - @flow.function_task - async def train(sel, k): - tr = cfg["training"] - device = tr["device"] - model_dim = tr["model_dim"] - use_collective = tr.get("use_collective", False) - collective_bytes = tr.get("collective_bytes", 4 * 1024 * 1024) - - kern.matMulGeneral(device=device, size_a=model_dim, size_b=model_dim, axis=0) - - if use_collective: - kern.MPIallReduce(device=device, data_size=collective_bytes) - model_update = {"version": k, "dim": model_dim} - return model_update - - propose_t = propose() - evaluate_t = evaluate(propose_t) - select_t = select(evaluate_t) - train_t = await train(select_t, k=1) - - await flow.shutdown() - -if __name__ == "__main__": - ap = argparse.ArgumentParser(description="Inverse Design (Motif 3) Mini-App") - ap.add_argument("--config", type=str, default="config.yaml") - args = ap.parse_args() - cfg = load_cfg(args.config) - t = Timer().start() - asyncio.run(workflow(cfg)) - log(f"DONE in {t.stop():.3f}s") \ No newline at end of file diff --git a/examples/motif_3/miniapp_parallel.py b/examples/motif_3/miniapp_parallel.py new file mode 100644 index 0000000..b3f8017 --- /dev/null +++ b/examples/motif_3/miniapp_parallel.py @@ -0,0 +1,81 @@ +import argparse, asyncio, yaml, random, logging +import time +from utils.metrics import Timer, log + +from radical.asyncflow import WorkflowEngine +from radical.asyncflow import DragonExecutionBackend +from radical.asyncflow.logging import init_default_logger + +from wfMiniAPI import kernel as kern + +def load_cfg(path): + with open(path, "r") as f: return yaml.safe_load(f) + +async def workflow(cfg): + logger = logging.getLogger(__name__) + init_default_logger(logging.DEBUG) + backend = await DragonExecutionBackend() + init_default_logger(logging.DEBUG) + flow = await WorkflowEngine.create(backend=backend) + + @flow.function_task + async def simulate(): + e = cfg["simulate"] + device = e["device"] + read_size = e["read_size_bytes"] + write_size = e["write_size_bytes"] + matmul_dim = e["matmul_dim"] + + log(f"Simulating candidate with params...") + log(time.strftime("%H:%M:%S", time.localtime())) + kern.generateRandomNumber(device=device, size=matmul_dim) + for i in range(8): + kern.matMulSimple2D(device=device, size=matmul_dim) + kern.writeNonMPI(num_bytes=write_size, data_root_dir="./") + kern.readNonMPI(num_bytes=read_size, data_root_dir="./") + log(f"Finished simulating candidate with params...") + log(time.strftime("%H:%M:%S", time.localtime())) + return random.random() + + @flow.function_task + async def train(res): + e = cfg["training"] + device = e["device"] + read_size = e["read_size_bytes"] + write_size = e["write_size_bytes"] + copy_size = e["data_copy_size_bytes"] + matmul_dim = e["matmul_dim"] + + log(f"Training model with evaluation results...") + log(time.strftime("%H:%M:%S", time.localtime())) + kern.readNonMPI(num_bytes=read_size, data_root_dir="./") + kern.dataCopyH2D(data_size=copy_size) + for i in range(6): + kern.matMulSimple2D(device="gpu", size=matmul_dim) + kern.matMulSimple2D(device="cpu", size=matmul_dim) + kern.dataCopyH2D(data_size=copy_size) + kern.writeNonMPI(num_bytes=write_size, data_root_dir="./") + log(f"Finished training model with evaluation results...") + log(time.strftime("%H:%M:%S", time.localtime())) + return random.random() + + simulate_t0= simulate() + + train_t1 = train(simulate_t0, None) + simulate_t1 = simulate(simulate_t0) + + train_t2 = train(simulate_t1, train_t1) + simulate_t2 = simulate(train_t1) + + train_t3 = await train(simulate_t2, train_t2) + + await flow.shutdown() + +if __name__ == "__main__": + ap = argparse.ArgumentParser(description="Inverse Design (Motif 3) Mini-App") + ap.add_argument("--config", type=str, default="config.yaml") + args = ap.parse_args() + cfg = load_cfg(args.config) + t = Timer().start() + asyncio.run(workflow(cfg)) + log(f"DONE in {t.stop():.3f}s") \ No newline at end of file diff --git a/examples/motif_3/miniapp_simple.py b/examples/motif_3/miniapp_simple.py new file mode 100644 index 0000000..95bb97b --- /dev/null +++ b/examples/motif_3/miniapp_simple.py @@ -0,0 +1,75 @@ +import argparse, asyncio, yaml, random, logging +import time +from utils.metrics import Timer, log + +from radical.asyncflow import WorkflowEngine +from radical.asyncflow import DragonExecutionBackend +from radical.asyncflow.logging import init_default_logger + +from wfMiniAPI import kernel as kern + +def load_cfg(path): + with open(path, "r") as f: return yaml.safe_load(f) + +async def workflow(cfg): + logger = logging.getLogger(__name__) + init_default_logger(logging.DEBUG) + backend = await DragonExecutionBackend() + init_default_logger(logging.DEBUG) + flow = await WorkflowEngine.create(backend=backend) + + @flow.function_task + async def simulate(): + e = cfg["simulate"] + device = e["device"] + read_size = e["read_size_bytes"] + write_size = e["write_size_bytes"] + matmul_dim = e["matmul_dim"] + + log(f"Simulating candidate with params...") + log(time.strftime("%H:%M:%S", time.localtime())) + kern.generateRandomNumber(device=device, size=matmul_dim) + for i in range(8): + kern.matMulSimple2D(device=device, size=matmul_dim) + kern.writeNonMPI(num_bytes=write_size, data_root_dir="./") + kern.readNonMPI(num_bytes=read_size, data_root_dir="./") + log(f"Finished simulating candidate with params...") + log(time.strftime("%H:%M:%S", time.localtime())) + return random.random() + + @flow.function_task + async def train(res): + e = cfg["training"] + device = e["device"] + read_size = e["read_size_bytes"] + write_size = e["write_size_bytes"] + copy_size = e["data_copy_size_bytes"] + matmul_dim = e["matmul_dim"] + + log(f"Training model with evaluation results...") + log(time.strftime("%H:%M:%S", time.localtime())) + kern.readNonMPI(num_bytes=read_size, data_root_dir="./") + kern.dataCopyH2D(data_size=copy_size) + for i in range(6): + kern.matMulSimple2D(device="gpu", size=matmul_dim) + kern.matMulSimple2D(device="cpu", size=matmul_dim) + kern.dataCopyH2D(data_size=copy_size) + kern.writeNonMPI(num_bytes=write_size, data_root_dir="./") + log(f"Finished training model with evaluation results...") + log(time.strftime("%H:%M:%S", time.localtime())) + return random.random() + + for i in range(3): + simulate_t= simulate() + train_t = await train(simulate_t) + + await flow.shutdown() + +if __name__ == "__main__": + ap = argparse.ArgumentParser(description="Inverse Design (Motif 3) Mini-App") + ap.add_argument("--config", type=str, default="config.yaml") + args = ap.parse_args() + cfg = load_cfg(args.config) + t = Timer().start() + asyncio.run(workflow(cfg)) + log(f"DONE in {t.stop():.3f}s") \ No newline at end of file diff --git a/examples/motif_3/stages/evaluate_stage.py b/examples/motif_3/stages/evaluate_stage.py deleted file mode 100644 index ddc982c..0000000 --- a/examples/motif_3/stages/evaluate_stage.py +++ /dev/null @@ -1,48 +0,0 @@ -import asyncio, math, random -from wfMiniAPI import kernel as kern -from utils.metrics import log - -class EvaluateStage: - """ - Embarrassingly-parallel evaluation pool: - - read/write emulation (I/O) - - main compute via matMulSimple2D (dominant) - Produces result records sent to both training and selection branches. - """ - def __init__(self, cfg): - e = cfg["evaluate"] - self.device = e["device"] - self.read_size = e["read_size_bytes"] - self.write_size = e["write_size_bytes"] - self.matmul_dim = e["matmul_dim"] - self.workers = e.get("workers", 4) - - async def worker(self, name, q_in, q_out_train, q_out_select): - while True: - cand = await q_in.get() - if cand is None: - # propagate sentinel to other workers and downstream - await q_in.put(None) - await q_out_train.put(None) - await q_out_select.put(None) - log(f"EvaluateStage[{name}]: complete") - return - - # emulate evaluation - kern.readNonMPI(nbytes=self.read_size) - kern.RNG(device=self.device, data_size=cand["vec_len"]) - kern.matMulSimple2D(device=self.device, dim=self.matmul_dim) - kern.writeNonMPI(nbytes=self.write_size) - - # fake objective: lower is better - obj = 1.0 / (1.0 + self.matmul_dim) + random.random() * 0.01 - res = {"rid": cand["rid"], "cid": cand["cid"], "objective": obj} - - await q_out_train.put(res) - await q_out_select.put(res) - - async def run(self, q_in, q_out_train, q_out_select): - tasks = [asyncio.create_task( - self.worker(f"W{i}", q_in, q_out_train, q_out_select)) - for i in range(self.workers)] - await asyncio.gather(*tasks) \ No newline at end of file diff --git a/examples/motif_3/stages/propose_stage.py b/examples/motif_3/stages/propose_stage.py deleted file mode 100644 index a6a8844..0000000 --- a/examples/motif_3/stages/propose_stage.py +++ /dev/null @@ -1,47 +0,0 @@ -import asyncio, random -from wfMiniAPI import kernel as kern -from utils.metrics import log - -class ProposeStage: - """ - Proposes candidate designs in rounds; consumes feedback to bias proposals. - Emulates proposal compute & data staging via RNG + H2D/D2H copies. - """ - def __init__(self, cfg): - p = cfg["propose"] - self.rounds = p["rounds"] - self.batch = p["batch_size"] - self.device = p["device"] - self.vec_len = p.get("vec_len", 1_000_000) - self.copy_bytes = p.get("copy_bytes", 2 * 1024 * 1024) - - async def run(self, q_out, q_fb): - for r in range(self.rounds): - # consume *latest* feedback if available (non-blocking) - fb_latest = None - try: - while True: - fb_latest = q_fb.get_nowait() - if fb_latest is None: - # early termination signal - await q_out.put(None) - log("ProposeStage: early stop") - return - except asyncio.QueueEmpty: - pass - - for i in range(self.batch): - # emulate cheap generator compute+copies - kern.RNG(device=self.device, data_size=self.vec_len) - kern.dataCopyH2D(data_size=self.copy_bytes) - kern.dataCopyD2H(data_size=self.copy_bytes) - cand = { - "rid": r, "cid": i, - "params": {"seed": random.random()}, - "vec_len": self.vec_len - } - await q_out.put(cand) - - # finished proposing - await q_out.put(None) - log("ProposeStage: complete") \ No newline at end of file diff --git a/examples/motif_3/stages/select_stage.py b/examples/motif_3/stages/select_stage.py deleted file mode 100644 index 48d1ef1..0000000 --- a/examples/motif_3/stages/select_stage.py +++ /dev/null @@ -1,62 +0,0 @@ -import asyncio, heapq -from wfMiniAPI import kernel as kern -from utils.metrics import log - -class SelectStage: - """ - Maintains top-K designs and periodically sends feedback to the proposer. - Consumes (a) evaluation results and (b) model updates (optional). - """ - def __init__(self, cfg): - s = cfg["select"] - self.top_k = s["top_k"] - self.emit_every = s.get("emit_every", 1) - self.device = s["device"] - self.policy_cost = s.get("policy_cost", 256 * 256) - - # min-heap storing (-score, rid, cid) so we keep best (lowest objective) - self.heap = [] - - def _push(self, res): - score = res["objective"] - item = (score, res["rid"], res["cid"]) - if len(self.heap) < self.top_k: - heapq.heappush(self.heap, (-score, item)) - else: - worst = self.heap[0] - if -score > worst[0]: - heapq.heapreplace(self.heap, (-score, item)) - - async def run(self, q_in_results, q_in_model, q_out_fb): - # model updates are optional; listen asynchronously - async def _model_listener(): - while True: - m = await q_in_model.get() - if m is None: return - # tiny compute to emulate re-scoring policy - kern.matMulSimple2D(device=self.device, - dim=int(self.policy_cost ** 0.5)) - - listener = asyncio.create_task(_model_listener()) - - n = 0 - while True: - res = await q_in_results.get() - if res is None: - await q_out_fb.put(None) # tell proposer we’re done - listener.cancel() - log("SelectStage: complete") - return - - self._push(res) - - if n % self.emit_every == 0 and self.heap: - # Emit feedback with current best - best_neg, (score, rid, cid) = max(self.heap), None - # heap stores (-score, item); invert again: - best = max(self.heap, key=lambda x: x[0]) - _, (top_score, top_rid, top_cid) = best - fb = {"best": {"rid": top_rid, "cid": top_cid, - "objective": top_score}} - await q_out_fb.put(fb) - n += 1 \ No newline at end of file diff --git a/examples/motif_3/stages/train_stage.py b/examples/motif_3/stages/train_stage.py deleted file mode 100644 index a759226..0000000 --- a/examples/motif_3/stages/train_stage.py +++ /dev/null @@ -1,38 +0,0 @@ -import asyncio -from wfMiniAPI import kernel as kern -from utils.metrics import log - -class TrainingStage: - """ - Optional surrogate training on a stream of evaluation results. - - Compute via matMulGeneral - - Optional MPIallReduce to emulate DP sync - Emits lightweight model-version updates. - """ - def __init__(self, cfg): - tr = cfg["training"] - self.device = tr["device"] - self.model_dim = tr["model_dim"] - self.train_every = tr["train_every"] - self.use_collective = tr.get("use_collective", False) - self.collective_bytes = tr.get("collective_bytes", 4 * 1024 * 1024) - - async def run(self, q_in, q_out): - k = 0 - while True: - msg = await q_in.get() - if msg is None: - await q_out.put(None) - log("TrainingStage: complete") - return - - if k % self.train_every == 0: - kern.matMulGeneral(device=self.device, - dim_list=[self.model_dim, - self.model_dim, - self.model_dim]) - if self.use_collective: - kern.MPIallReduce(device=self.device, - data_size=self.collective_bytes) - await q_out.put({"version": k, "dim": self.model_dim}) - k += 1 \ No newline at end of file diff --git a/examples/motif_4/config.yaml b/examples/motif_4/config.yaml index dc26d93..f35526d 100644 --- a/examples/motif_4/config.yaml +++ b/examples/motif_4/config.yaml @@ -1,36 +1,27 @@ +experiment: + steps: 8 + read_size_bytes: 1024 # 1 KiB pre-step read + write_size_bytes: 8589934592 # 8 GiB post-step write + device: "cpu" # "cpu" or "gpu" + matmul_dim: 8192 # dimension for square matrix multiplication + simulation: - steps: 100 - state_dim: 1024 + steps: 16 read_size_bytes: 1048576 # 1 MiB pre-step read - write_size_bytes: 1048576 # 1 MiB post-step write + write_size_bytes: 4294967296 # 4 GiB post-step write device: "cpu" # "cpu" or "gpu" - sleep_ms: 0 - -assimilation: - window: 4 - device: "cpu" - emit_every: 1 + matmul_dim: 2048 # dimension for square matrix multiplication training: - device: "cpu" - train_every: 5 - model_dim: 512 - use_collective: false - collective_bytes: 4194304 # emulate AllReduce volume + steps: 6 + read_size_bytes: 2147483648 # 2 GiB pre-step read + write_size_bytes: 512 # 512 B post-step write + data_copy_size_bytes: 2097152 # 2 MiB data copy + matmul_dim: 8192 # dimension for square matrix multiplication inference: - device: "cpu" - copy_bytes: 2097152 - vec_len: 1000000 - -control: - device: "cpu" - policy_cost: 65536 # ~256x256 matmul - -queues: - sim_to_assim: 8 - assim_to_train: 8 - assim_to_infer: 8 - train_to_infer: 4 - infer_to_control: 8 - control_to_sim: 4 \ No newline at end of file + steps: 4 + read_size_bytes: 1342177280 # 1.25 GiB pre + write_size_bytes: 2147483648 # 2 GiB post-step write + data_copy_size_bytes: 2097152 # 2 MiB data copy + matmul_dim: 8192 # dimension for square matrix multiplication \ No newline at end of file diff --git a/examples/motif_4/miniapp.py b/examples/motif_4/miniapp.py deleted file mode 100644 index 209366a..0000000 --- a/examples/motif_4/miniapp.py +++ /dev/null @@ -1,53 +0,0 @@ -import argparse, asyncio, os, time, json -from pathlib import Path - -from stages.sim_stage import SimulationStage -from stages.assimilate_stage import AssimilationStage -from stages.train_stage import TrainingStage -from stages.infer_stage import InferenceStage -from stages.control_stage import ControlStage -from utils.metrics import Timer, log -import yaml - -async def run_pipeline(cfg): - q_sim_to_assim = asyncio.Queue(maxsize=cfg["queues"]["sim_to_assim"]) - q_assim_to_train = asyncio.Queue(maxsize=cfg["queues"]["assim_to_train"]) - q_assim_to_infer = asyncio.Queue(maxsize=cfg["queues"]["assim_to_infer"]) - q_train_to_infer = asyncio.Queue(maxsize=cfg["queues"]["train_to_infer"]) - q_infer_to_control = asyncio.Queue(maxsize=cfg["queues"]["infer_to_control"]) - q_control_to_sim = asyncio.Queue(maxsize=cfg["queues"]["control_to_sim"]) - - # Stages - sim = SimulationStage(cfg) - assim = AssimilationStage(cfg) - train = TrainingStage(cfg) - infer = InferenceStage(cfg) - control = ControlStage(cfg) - - # Tasks (concurrent, real-time style) - tasks = [ - asyncio.create_task(sim.run(q_out=q_sim_to_assim, q_in=q_control_to_sim)), - asyncio.create_task(assim.run(q_in=q_sim_to_assim, - q_out_train=q_assim_to_train, - q_out_infer=q_assim_to_infer)), - asyncio.create_task(train.run(q_in=q_assim_to_train, q_out=q_train_to_infer)), - asyncio.create_task(infer.run(q_state=q_assim_to_infer, - q_model=q_train_to_infer, - q_out=q_infer_to_control)), - asyncio.create_task(control.run(q_in=q_infer_to_control, q_out=q_control_to_sim)), - ] - - await asyncio.gather(*tasks) - -def load_cfg(path): - with open(path, "r") as f: - return yaml.safe_load(f) - -if __name__ == "__main__": - ap = argparse.ArgumentParser(description="Digital Twin (Motif 4) Mini-App") - ap.add_argument("--config", type=str, default="config.yaml") - args = ap.parse_args() - cfg = load_cfg(args.config) - t = Timer().start() - asyncio.run(run_pipeline(cfg)) - log(f"DONE in {t.stop():.3f}s") diff --git a/examples/motif_4/miniapp_async.py b/examples/motif_4/miniapp_async.py index fd4ec16..1dc9cb2 100644 --- a/examples/motif_4/miniapp_async.py +++ b/examples/motif_4/miniapp_async.py @@ -1,100 +1,115 @@ -import argparse, asyncio, yaml, random +import argparse, asyncio, yaml, random, logging +import time from utils.metrics import Timer, log from radical.asyncflow import WorkflowEngine -from radical.asyncflow import ConcurrentExecutionBackend +from radical.asyncflow import DragonExecutionBackend +from radical.asyncflow.logging import init_default_logger from wfMiniAPI import kernel as kern -from concurrent.futures import ThreadPoolExecutor - def load_cfg(path): with open(path, "r") as f: return yaml.safe_load(f) async def workflow(cfg): - backend = await ConcurrentExecutionBackend(ThreadPoolExecutor()) + logger = logging.getLogger(__name__) + init_default_logger(logging.DEBUG) + backend = await DragonExecutionBackend() + init_default_logger(logging.DEBUG) flow = await WorkflowEngine.create(backend=backend) @flow.function_task - async def simulation(): - sim = cfg["simulation"] - steps = sim["steps"] - state_dim = sim["state_dim"] - read_size = sim["read_size_bytes"] - write_size = sim["write_size_bytes"] - device = sim["device"] - sleep_ms = sim.get("sleep_ms", 0) - kern.readNonMPI(nbytes=read_size) - - kern.RNG(device=device, data_size=state_dim) - kern.matMulSimple2D(device=device, dim=state_dim) - - kern.writeNonMPI(nbytes=write_size) - - msg = {"state_dim": state_dim} - - if sleep_ms > 0: - await asyncio.sleep(sleep_ms / 1000.0) - return msg - - @flow.function_task - async def assimilation(state): - assim = cfg["assimilation"] - read_size = assim["read_size_bytes"] - write_size = assim["write_size_bytes"] - device = assim["device"] - kern.readNonMPI(nbytes=read_size) - kern.RNG(device=device, data_size=state["state_dim"]) - kern.matMulSimple2D(device=device, dim=state["state_dim"]) - kern.writeNonMPI(nbytes=write_size) - msg = {"t": state["t"], "state_dim": state["state_dim"]} - return msg + async def experiment(): + e = cfg["experiment"] + steps = e["steps"] + read_size = e["read_size_bytes"] + write_size = e["write_size_bytes"] + device = e["device"] + matmul_dim = e["matmul_dim"] + + log(f"Experiment Data...") + log(time.strftime("%H:%M:%S", time.localtime())) + kern.generateRandomNumber(device=device, size=matmul_dim) + for i in range(steps): + kern.matMulSimple2D(device=device, size=matmul_dim) + kern.writeNonMPI(num_bytes=write_size, data_root_dir="./") + kern.readNonMPI(num_bytes=read_size, data_root_dir="./") + log(f"Finished simulating...") + log(time.strftime("%H:%M:%S", time.localtime())) + return random.random() @flow.function_task - async def training(state): - train = cfg["training"] - read_size = train["read_size_bytes"] - write_size = train["write_size_bytes"] - device = train["device"] - kern.readNonMPI(nbytes=read_size) - kern.RNG(device=device, data_size=state["state_dim"]) - kern.matMulSimple2D(device=device, dim=state["state_dim"]) - kern.writeNonMPI(nbytes=write_size) - msg = {"t": state["t"], "model_dim": state["state_dim"] // 2} - return msg + async def simulation(i): + e = cfg["simulation"] + steps = e["steps"] + read_size = e["read_size_bytes"] + write_size = e["write_size_bytes"] + device = e["device"] + matmul_dim = e["matmul_dim"] + log(f"Simulating {i}...") + log(time.strftime("%H:%M:%S", time.localtime())) + kern.readNonMPI(num_bytes=read_size, data_root_dir="./") + kern.generateRandomNumber(device=device, size=matmul_dim) + for i in range(steps): + kern.matMulSimple2D(device=device, size=matmul_dim) + if i == 0: + kern.writeNonMPI(num_bytes=write_size, data_root_dir="./tmp") + log(f"Finished simulating...") + log(time.strftime("%H:%M:%S", time.localtime())) + return random.random() + @flow.function_task - async def inference(state, model): - infer = cfg["inference"] - read_size = infer["read_size_bytes"] - write_size = infer["write_size_bytes"] - device = infer["device"] - kern.readNonMPI(nbytes=read_size) - kern.RNG(device=device, data_size=state["state_dim"]) - kern.matMulSimple2D(device=device, dim=state["state_dim"]) - kern.matMulSimple2D(device=device, dim=model["model_dim"]) - kern.writeNonMPI(nbytes=write_size) - msg = {"t": state["t"], "state_dim": state["state_dim"]} - return msg + async def training(simulation, experiment): + e = cfg["training"] + read_size = e["read_size_bytes"] + write_size = e["write_size_bytes"] + copy_size = e["data_copy_size_bytes"] + matmul_dim = e["matmul_dim"] + + log(f"Training model with simulation results...") + log(time.strftime("%H:%M:%S", time.localtime())) + kern.readNonMPI(num_bytes=read_size, data_root_dir="./") + kern.dataCopyH2D(data_size=copy_size) + for i in range(6): + kern.matMulSimple2D(device="gpu", size=matmul_dim) + kern.matMulSimple2D(device="cpu", size=matmul_dim) + kern.dataCopyH2D(data_size=copy_size) + kern.writeNonMPI(num_bytes=write_size, data_root_dir="./") + log(f"Finished training model with evaluation results...") + log(time.strftime("%H:%M:%S", time.localtime())) + return random.random() @flow.function_task - async def control(state): - ctrl = cfg["control"] - read_size = ctrl["read_size_bytes"] - write_size = ctrl["write_size_bytes"] - device = ctrl["device"] - kern.readNonMPI(nbytes=read_size) - kern.RNG(device=device, data_size=state["state_dim"]) - kern.matMulSimple2D(device=device, dim=state["state_dim"]) - kern.writeNonMPI(nbytes=write_size) - msg = {"t": state["t"], "ctrl_param": random.random()} - return msg - - sim_t = simulation() - assim_t = assimilation(sim_t) - train_t = training(assim_t) - infer_t = inference(assim_t, train_t) - control_t = await control(infer_t) + async def inference(training): + e = cfg["inference"] + steps = e["steps"] + read_size = e["read_size_bytes"] + write_size = e["write_size_bytes"] + copy_size = e["data_copy_size_bytes"] + matmul_dim = e["matmul_dim"] + + log(f"Inferencing model with evaluation results...") + log(time.strftime("%H:%M:%S", time.localtime())) + kern.readNonMPI(num_bytes=read_size, data_root_dir="./") + kern.dataCopyH2D(data_size=copy_size) + for i in range(steps): + kern.matMulSimple2D(device="gpu", size=matmul_dim) + kern.matMulSimple2D(device="cpu", size=matmul_dim) + kern.dataCopyH2D(data_size=copy_size) + kern.writeNonMPI(num_bytes=write_size, data_root_dir="./") + log(f"Finished training model with evaluation results...") + log(time.strftime("%H:%M:%S", time.localtime())) + return random.random() + + for i in range(3): + experiment_t = await experiment() + sim_t = [] + for i in range(32): + sim_t.append(simulation(i)) + await asyncio.gather(*sim_t) + train_t = training(sim_t, experiment_t) + infer_t = await inference(train_t) await flow.shutdown() diff --git a/examples/motif_4/stages/assimilate_stage.py b/examples/motif_4/stages/assimilate_stage.py deleted file mode 100644 index 94d03c4..0000000 --- a/examples/motif_4/stages/assimilate_stage.py +++ /dev/null @@ -1,35 +0,0 @@ -import asyncio, collections -from wfMiniAPI import kernel as kern -from utils.metrics import log - -class AssimilationStage: - """ - Windowed assimilation (e.g., reductions / filtering). - Emits to both training (for model updates) and inference (for online predictions). - """ - def __init__(self, cfg): - a = cfg["assimilation"] - self.window = a["window"] - self.device = a["device"] - self.buf = collections.deque(maxlen=self.window) - self.emit_every = a.get("emit_every", 1) - - async def run(self, q_in, q_out_train, q_out_infer): - n = 0 - while True: - msg = await q_in.get() - if msg is None: - await q_out_train.put(None) - await q_out_infer.put(None) - log("AssimilationStage: complete") - return - - self.buf.append(msg["state_dim"]) - data_size = max(1, sum(self.buf)) # scale only - kern.reduction(device=self.device, data_size=data_size) - - if n % self.emit_every == 0: - state_summary = {"t": msg["t"], "feat": data_size} - await q_out_train.put(state_summary) - await q_out_infer.put(state_summary) - n += 1 diff --git a/examples/motif_4/stages/control_stage.py b/examples/motif_4/stages/control_stage.py deleted file mode 100644 index 2bf4c53..0000000 --- a/examples/motif_4/stages/control_stage.py +++ /dev/null @@ -1,25 +0,0 @@ -import asyncio, math -from wfMiniAPI import kernel as kern -from utils.metrics import log - -class ControlStage: - """ - Takes predictions and computes control signals to steer the sim. - """ - def __init__(self, cfg): - c = cfg["control"] - self.device = c["device"] - self.policy_cost = c.get("policy_cost", 256*256) - - async def run(self, q_in, q_out): - while True: - pred = await q_in.get() - if pred is None: - await q_out.put(None) - log("ControlStage: complete") - return - # emulate tiny policy compute - dim = int(math.sqrt(self.policy_cost)) - kern.matMulSimple2D(device=self.device, dim=dim) - ctrl = {"gain": 0.01, "model_ver": pred["model_ver"], "t": pred["t"]} - await q_out.put(ctrl) diff --git a/examples/motif_4/stages/infer_stage.py b/examples/motif_4/stages/infer_stage.py deleted file mode 100644 index 32c5aaa..0000000 --- a/examples/motif_4/stages/infer_stage.py +++ /dev/null @@ -1,43 +0,0 @@ -import asyncio -from wfMiniAPI import kernel as kern -from utils.metrics import log - -class InferenceStage: - """ - Online inference using the surrogate. - Models memory copies + axpy-like compute. - """ - def __init__(self, cfg): - inf = cfg["inference"] - self.device = inf["device"] - self.copy_bytes = inf.get("copy_bytes", 2 * 1024 * 1024) - self.vec_len = inf.get("vec_len", 1_000_000) - - async def run(self, q_state, q_model, q_out): - current_model = {"version": -1, "dim": 1} - - async def model_updater(): - nonlocal current_model - while True: - m = await q_model.get() - if m is None: - return - current_model = m - - updater_task = asyncio.create_task(model_updater()) - - while True: - s = await q_state.get() - if s is None: - await q_out.put(None) - updater_task.cancel() - log("InferenceStage: complete") - return - - # emulate H2D/D2H copies around inference - kern.dataCopyH2D(data_size=self.copy_bytes) - kern.axpy(device=self.device, data_size=self.vec_len) # y = a*x + y (emulated) - kern.dataCopyD2H(data_size=self.copy_bytes) - - pred = {"t": s["t"], "model_ver": current_model["version"], "health": 1.0} - await q_out.put(pred) diff --git a/examples/motif_4/stages/sim_stage.py b/examples/motif_4/stages/sim_stage.py deleted file mode 100644 index adc482c..0000000 --- a/examples/motif_4/stages/sim_stage.py +++ /dev/null @@ -1,48 +0,0 @@ -import asyncio, time, math, random -from wfMiniAPI import kernel as kern -from utils.metrics import log - -class SimulationStage: - """ - Emulates a time-stepping simulation producing state snapshots. - Uses matMulSimple2D + read/write emulation to capture CPU/GPU + I/O. - Accepts control updates (e.g., parameter steering) from the controller. - """ - def __init__(self, cfg): - sim = cfg["simulation"] - self.steps = sim["steps"] - self.state_dim = sim["state_dim"] - self.read_size = sim["read_size_bytes"] - self.write_size = sim["write_size_bytes"] - self.device = sim["device"] - self.sleep_ms = sim.get("sleep_ms", 0) - - async def run(self, q_out, q_in): - for step in range(self.steps): - # apply pending control update (non-blocking) - ctrl = None - try: - ctrl = q_in.get_nowait() - except asyncio.QueueEmpty: - pass - - # read/input emulate - kern.readNonMPI(nbytes=self.read_size) - - # compute: simulate main step - kern.RNG(device=self.device, data_size=self.state_dim) - kern.matMulSimple2D(device=self.device, dim=self.state_dim) - - # write/output emulate - kern.writeNonMPI(nbytes=self.write_size) - - # emit "state" token (size metadata; payload omitted for lightness) - msg = {"t": step, "state_dim": self.state_dim, "ctrl": ctrl} - await q_out.put(msg) - - if self.sleep_ms > 0: - await asyncio.sleep(self.sleep_ms / 1000.0) - - # signal completion - await q_out.put(None) - log("SimulationStage: complete") diff --git a/examples/motif_4/stages/train_stage.py b/examples/motif_4/stages/train_stage.py deleted file mode 100644 index c94a22b..0000000 --- a/examples/motif_4/stages/train_stage.py +++ /dev/null @@ -1,36 +0,0 @@ -import asyncio -from wfMiniAPI import kernel as kern -from utils.metrics import log - -class TrainingStage: - """ - Periodic surrogate training step. - Uses matMulGeneral + (optional) MPIallReduce to emulate data-parallel training. - """ - def __init__(self, cfg): - tr = cfg["training"] - self.device = tr["device"] - self.train_every = tr["train_every"] - self.model_dim = tr["model_dim"] - self.use_collective = tr.get("use_collective", False) - self.collective_bytes = tr.get("collective_bytes", 4 * 1024 * 1024) - - async def run(self, q_in, q_out): - k = 0 - while True: - msg = await q_in.get() - if msg is None: - await q_out.put(None) - log("TrainingStage: complete") - return - if k % self.train_every == 0: - # emulate backprop compute - kern.matMulGeneral(device=self.device, dim_list=[self.model_dim, - self.model_dim, - self.model_dim]) - # emulate gradient allreduce - if self.use_collective: - kern.MPIallReduce(device=self.device, data_size=self.collective_bytes) - model = {"version": k, "dim": self.model_dim} - await q_out.put(model) - k += 1 From 061dc852c5ebb32192f3d6392a2d2808d9742f65 Mon Sep 17 00:00:00 2001 From: Andrew Park Date: Tue, 10 Mar 2026 18:07:13 -0400 Subject: [PATCH 35/46] Minifix --- examples/motif_3/config.yaml | 2 ++ examples/motif_3/miniapp_parallel.py | 6 ++++-- examples/motif_3/miniapp_simple.py | 6 ++++-- 3 files changed, 10 insertions(+), 4 deletions(-) diff --git a/examples/motif_3/config.yaml b/examples/motif_3/config.yaml index 299c2e6..74ae1c6 100644 --- a/examples/motif_3/config.yaml +++ b/examples/motif_3/config.yaml @@ -1,10 +1,12 @@ simulate: + steps: 8 device: "cpu" read_size_bytes: 4294967296 write_size_bytes: 1073741824 matmul_dim: 8192 training: + steps: 6 device: "cpu" read_size_bytes: 1342177280 write_size_bytes: 512 diff --git a/examples/motif_3/miniapp_parallel.py b/examples/motif_3/miniapp_parallel.py index b3f8017..dd4f6f2 100644 --- a/examples/motif_3/miniapp_parallel.py +++ b/examples/motif_3/miniapp_parallel.py @@ -21,6 +21,7 @@ async def workflow(cfg): @flow.function_task async def simulate(): e = cfg["simulate"] + steps = e["steps"] device = e["device"] read_size = e["read_size_bytes"] write_size = e["write_size_bytes"] @@ -29,7 +30,7 @@ async def simulate(): log(f"Simulating candidate with params...") log(time.strftime("%H:%M:%S", time.localtime())) kern.generateRandomNumber(device=device, size=matmul_dim) - for i in range(8): + for i in range(steps): kern.matMulSimple2D(device=device, size=matmul_dim) kern.writeNonMPI(num_bytes=write_size, data_root_dir="./") kern.readNonMPI(num_bytes=read_size, data_root_dir="./") @@ -40,6 +41,7 @@ async def simulate(): @flow.function_task async def train(res): e = cfg["training"] + steps = e["steps"] device = e["device"] read_size = e["read_size_bytes"] write_size = e["write_size_bytes"] @@ -50,7 +52,7 @@ async def train(res): log(time.strftime("%H:%M:%S", time.localtime())) kern.readNonMPI(num_bytes=read_size, data_root_dir="./") kern.dataCopyH2D(data_size=copy_size) - for i in range(6): + for i in range(steps): kern.matMulSimple2D(device="gpu", size=matmul_dim) kern.matMulSimple2D(device="cpu", size=matmul_dim) kern.dataCopyH2D(data_size=copy_size) diff --git a/examples/motif_3/miniapp_simple.py b/examples/motif_3/miniapp_simple.py index 95bb97b..b5d023a 100644 --- a/examples/motif_3/miniapp_simple.py +++ b/examples/motif_3/miniapp_simple.py @@ -21,6 +21,7 @@ async def workflow(cfg): @flow.function_task async def simulate(): e = cfg["simulate"] + steps = e["steps"] device = e["device"] read_size = e["read_size_bytes"] write_size = e["write_size_bytes"] @@ -29,7 +30,7 @@ async def simulate(): log(f"Simulating candidate with params...") log(time.strftime("%H:%M:%S", time.localtime())) kern.generateRandomNumber(device=device, size=matmul_dim) - for i in range(8): + for i in range(steps): kern.matMulSimple2D(device=device, size=matmul_dim) kern.writeNonMPI(num_bytes=write_size, data_root_dir="./") kern.readNonMPI(num_bytes=read_size, data_root_dir="./") @@ -40,6 +41,7 @@ async def simulate(): @flow.function_task async def train(res): e = cfg["training"] + steps = e["steps"] device = e["device"] read_size = e["read_size_bytes"] write_size = e["write_size_bytes"] @@ -50,7 +52,7 @@ async def train(res): log(time.strftime("%H:%M:%S", time.localtime())) kern.readNonMPI(num_bytes=read_size, data_root_dir="./") kern.dataCopyH2D(data_size=copy_size) - for i in range(6): + for i in range(steps): kern.matMulSimple2D(device="gpu", size=matmul_dim) kern.matMulSimple2D(device="cpu", size=matmul_dim) kern.dataCopyH2D(data_size=copy_size) From 8e3fe524bcfb76995168b2b034eaba1b5388ac42 Mon Sep 17 00:00:00 2001 From: Andrew Park Date: Tue, 10 Mar 2026 18:16:51 -0400 Subject: [PATCH 36/46] One last try --- examples/motif_4/miniapp_async.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/examples/motif_4/miniapp_async.py b/examples/motif_4/miniapp_async.py index 1dc9cb2..4966428 100644 --- a/examples/motif_4/miniapp_async.py +++ b/examples/motif_4/miniapp_async.py @@ -51,7 +51,7 @@ async def simulation(i): log(time.strftime("%H:%M:%S", time.localtime())) kern.readNonMPI(num_bytes=read_size, data_root_dir="./") kern.generateRandomNumber(device=device, size=matmul_dim) - for i in range(steps): + for _ in range(steps): kern.matMulSimple2D(device=device, size=matmul_dim) if i == 0: kern.writeNonMPI(num_bytes=write_size, data_root_dir="./tmp") @@ -62,6 +62,7 @@ async def simulation(i): @flow.function_task async def training(simulation, experiment): e = cfg["training"] + steps = e["steps"] read_size = e["read_size_bytes"] write_size = e["write_size_bytes"] copy_size = e["data_copy_size_bytes"] @@ -71,7 +72,7 @@ async def training(simulation, experiment): log(time.strftime("%H:%M:%S", time.localtime())) kern.readNonMPI(num_bytes=read_size, data_root_dir="./") kern.dataCopyH2D(data_size=copy_size) - for i in range(6): + for i in range(steps): kern.matMulSimple2D(device="gpu", size=matmul_dim) kern.matMulSimple2D(device="cpu", size=matmul_dim) kern.dataCopyH2D(data_size=copy_size) @@ -98,7 +99,7 @@ async def inference(training): kern.matMulSimple2D(device="cpu", size=matmul_dim) kern.dataCopyH2D(data_size=copy_size) kern.writeNonMPI(num_bytes=write_size, data_root_dir="./") - log(f"Finished training model with evaluation results...") + log(f"Finished inferencing model with evaluation results...") log(time.strftime("%H:%M:%S", time.localtime())) return random.random() From a278bbac2f890ee695b39a6b099e661f943365fd Mon Sep 17 00:00:00 2001 From: Andrew Park Date: Tue, 10 Mar 2026 18:52:28 -0400 Subject: [PATCH 37/46] _ --- examples/Tutorial/tasks/active_learn.py | 2 +- examples/Tutorial/tasks/training.py | 2 +- examples/motif_3/miniapp_parallel.py | 31 +++++++++-------- examples/motif_3/miniapp_simple.py | 27 ++++++++------- examples/motif_3/utils/metrics.py | 6 ---- examples/motif_4/miniapp_async.py | 45 +++++++++++++------------ examples/motif_4/utils/metrics.py | 8 ----- 7 files changed, 58 insertions(+), 63 deletions(-) delete mode 100644 examples/motif_3/utils/metrics.py delete mode 100644 examples/motif_4/utils/metrics.py diff --git a/examples/Tutorial/tasks/active_learn.py b/examples/Tutorial/tasks/active_learn.py index 6d7f359..9e85304 100644 --- a/examples/Tutorial/tasks/active_learn.py +++ b/examples/Tutorial/tasks/active_learn.py @@ -14,7 +14,7 @@ def parse_args(): parser.add_argument('--batch_size', type=int, default=64, help='batch size in training') parser.add_argument('--device', default='gpu', - help='Wheter this is running on cpu or gpu') + help='Whether this is running on cpu or gpu') parser.add_argument('--dense_dim_in', type=int, default=2048, help='dim for most heavy dense layer, input') parser.add_argument('--dense_dim_out', type=int, default=512, diff --git a/examples/Tutorial/tasks/training.py b/examples/Tutorial/tasks/training.py index 729d0e0..1ef1f9a 100644 --- a/examples/Tutorial/tasks/training.py +++ b/examples/Tutorial/tasks/training.py @@ -16,7 +16,7 @@ def parse_args(): parser.add_argument('--batch_size', type=int, default=64, help='batch size in training') parser.add_argument('--device', default='gpu', - help='Wheter this is running on cpu or gpu') + help='Whether this is running on cpu or gpu') parser.add_argument('--dense_dim_in', type=int, default=2048, help='dim for most heavy dense layer, input') parser.add_argument('--dense_dim_out', type=int, default=512, diff --git a/examples/motif_3/miniapp_parallel.py b/examples/motif_3/miniapp_parallel.py index dd4f6f2..16d670d 100644 --- a/examples/motif_3/miniapp_parallel.py +++ b/examples/motif_3/miniapp_parallel.py @@ -1,6 +1,5 @@ import argparse, asyncio, yaml, random, logging import time -from utils.metrics import Timer, log from radical.asyncflow import WorkflowEngine from radical.asyncflow import DragonExecutionBackend @@ -8,6 +7,11 @@ from wfMiniAPI import kernel as kern +class Timer: + def __init__(self): self.t0 = None + def start(self): self.t0 = time.time(); return self + def stop(self): return time.time() - self.t0 + def load_cfg(path): with open(path, "r") as f: return yaml.safe_load(f) @@ -15,11 +19,10 @@ async def workflow(cfg): logger = logging.getLogger(__name__) init_default_logger(logging.DEBUG) backend = await DragonExecutionBackend() - init_default_logger(logging.DEBUG) flow = await WorkflowEngine.create(backend=backend) @flow.function_task - async def simulate(): + async def simulate(_train=None): e = cfg["simulate"] steps = e["steps"] device = e["device"] @@ -27,19 +30,19 @@ async def simulate(): write_size = e["write_size_bytes"] matmul_dim = e["matmul_dim"] - log(f"Simulating candidate with params...") - log(time.strftime("%H:%M:%S", time.localtime())) + logger.info(f"Simulating with params...") + logger.info(time.strftime("%H:%M:%S", time.localtime())) kern.generateRandomNumber(device=device, size=matmul_dim) for i in range(steps): kern.matMulSimple2D(device=device, size=matmul_dim) kern.writeNonMPI(num_bytes=write_size, data_root_dir="./") kern.readNonMPI(num_bytes=read_size, data_root_dir="./") - log(f"Finished simulating candidate with params...") - log(time.strftime("%H:%M:%S", time.localtime())) + logger.info(f"Finished simulating with params...") + logger.info(time.strftime("%H:%M:%S", time.localtime())) return random.random() @flow.function_task - async def train(res): + async def train(_sim, _train=None): e = cfg["training"] steps = e["steps"] device = e["device"] @@ -48,8 +51,8 @@ async def train(res): copy_size = e["data_copy_size_bytes"] matmul_dim = e["matmul_dim"] - log(f"Training model with evaluation results...") - log(time.strftime("%H:%M:%S", time.localtime())) + logger.info(f"Training model with evaluation results...") + logger.info(time.strftime("%H:%M:%S", time.localtime())) kern.readNonMPI(num_bytes=read_size, data_root_dir="./") kern.dataCopyH2D(data_size=copy_size) for i in range(steps): @@ -57,13 +60,13 @@ async def train(res): kern.matMulSimple2D(device="cpu", size=matmul_dim) kern.dataCopyH2D(data_size=copy_size) kern.writeNonMPI(num_bytes=write_size, data_root_dir="./") - log(f"Finished training model with evaluation results...") - log(time.strftime("%H:%M:%S", time.localtime())) + logger.info(f"Finished training model with evaluation results...") + logger.info(time.strftime("%H:%M:%S", time.localtime())) return random.random() simulate_t0= simulate() - train_t1 = train(simulate_t0, None) + train_t1 = train(simulate_t0) simulate_t1 = simulate(simulate_t0) train_t2 = train(simulate_t1, train_t1) @@ -80,4 +83,4 @@ async def train(res): cfg = load_cfg(args.config) t = Timer().start() asyncio.run(workflow(cfg)) - log(f"DONE in {t.stop():.3f}s") \ No newline at end of file + print(f"DONE in {t.stop():.3f}s") \ No newline at end of file diff --git a/examples/motif_3/miniapp_simple.py b/examples/motif_3/miniapp_simple.py index b5d023a..b42331a 100644 --- a/examples/motif_3/miniapp_simple.py +++ b/examples/motif_3/miniapp_simple.py @@ -1,6 +1,5 @@ import argparse, asyncio, yaml, random, logging import time -from utils.metrics import Timer, log from radical.asyncflow import WorkflowEngine from radical.asyncflow import DragonExecutionBackend @@ -8,6 +7,11 @@ from wfMiniAPI import kernel as kern +class Timer: + def __init__(self): self.t0 = None + def start(self): self.t0 = time.time(); return self + def stop(self): return time.time() - self.t0 + def load_cfg(path): with open(path, "r") as f: return yaml.safe_load(f) @@ -15,7 +19,6 @@ async def workflow(cfg): logger = logging.getLogger(__name__) init_default_logger(logging.DEBUG) backend = await DragonExecutionBackend() - init_default_logger(logging.DEBUG) flow = await WorkflowEngine.create(backend=backend) @flow.function_task @@ -27,19 +30,19 @@ async def simulate(): write_size = e["write_size_bytes"] matmul_dim = e["matmul_dim"] - log(f"Simulating candidate with params...") - log(time.strftime("%H:%M:%S", time.localtime())) + logger.info(f"Simulating with params...") + logger.info(time.strftime("%H:%M:%S", time.localtime())) kern.generateRandomNumber(device=device, size=matmul_dim) for i in range(steps): kern.matMulSimple2D(device=device, size=matmul_dim) kern.writeNonMPI(num_bytes=write_size, data_root_dir="./") kern.readNonMPI(num_bytes=read_size, data_root_dir="./") - log(f"Finished simulating candidate with params...") - log(time.strftime("%H:%M:%S", time.localtime())) + logger.info(f"Finished simulating with params...") + logger.info(time.strftime("%H:%M:%S", time.localtime())) return random.random() @flow.function_task - async def train(res): + async def train(_sim=None): e = cfg["training"] steps = e["steps"] device = e["device"] @@ -48,8 +51,8 @@ async def train(res): copy_size = e["data_copy_size_bytes"] matmul_dim = e["matmul_dim"] - log(f"Training model with evaluation results...") - log(time.strftime("%H:%M:%S", time.localtime())) + logger.info(f"Training model with evaluation results...") + logger.info(time.strftime("%H:%M:%S", time.localtime())) kern.readNonMPI(num_bytes=read_size, data_root_dir="./") kern.dataCopyH2D(data_size=copy_size) for i in range(steps): @@ -57,8 +60,8 @@ async def train(res): kern.matMulSimple2D(device="cpu", size=matmul_dim) kern.dataCopyH2D(data_size=copy_size) kern.writeNonMPI(num_bytes=write_size, data_root_dir="./") - log(f"Finished training model with evaluation results...") - log(time.strftime("%H:%M:%S", time.localtime())) + logger.info(f"Finished training model with evaluation results...") + logger.info(time.strftime("%H:%M:%S", time.localtime())) return random.random() for i in range(3): @@ -74,4 +77,4 @@ async def train(res): cfg = load_cfg(args.config) t = Timer().start() asyncio.run(workflow(cfg)) - log(f"DONE in {t.stop():.3f}s") \ No newline at end of file + print(f"DONE in {t.stop():.3f}s") \ No newline at end of file diff --git a/examples/motif_3/utils/metrics.py b/examples/motif_3/utils/metrics.py deleted file mode 100644 index 03184e5..0000000 --- a/examples/motif_3/utils/metrics.py +++ /dev/null @@ -1,6 +0,0 @@ -import time, sys -class Timer: - def __init__(self): self.t0 = None - def start(self): self.t0 = time.time(); return self - def stop(self): return time.time() - self.t0 -def log(msg): print(f"[id] {msg}", file=sys.stdout, flush=True) \ No newline at end of file diff --git a/examples/motif_4/miniapp_async.py b/examples/motif_4/miniapp_async.py index 4966428..5dba85f 100644 --- a/examples/motif_4/miniapp_async.py +++ b/examples/motif_4/miniapp_async.py @@ -1,6 +1,5 @@ import argparse, asyncio, yaml, random, logging import time -from utils.metrics import Timer, log from radical.asyncflow import WorkflowEngine from radical.asyncflow import DragonExecutionBackend @@ -8,6 +7,11 @@ from wfMiniAPI import kernel as kern +class Timer: + def __init__(self): self.t0 = None + def start(self): self.t0 = time.time(); return self + def stop(self): return time.time() - self.t0 + def load_cfg(path): with open(path, "r") as f: return yaml.safe_load(f) @@ -15,7 +19,6 @@ async def workflow(cfg): logger = logging.getLogger(__name__) init_default_logger(logging.DEBUG) backend = await DragonExecutionBackend() - init_default_logger(logging.DEBUG) flow = await WorkflowEngine.create(backend=backend) @flow.function_task @@ -27,15 +30,15 @@ async def experiment(): device = e["device"] matmul_dim = e["matmul_dim"] - log(f"Experiment Data...") - log(time.strftime("%H:%M:%S", time.localtime())) + logger.info(f"Experiment Data...") + logger.info(time.strftime("%H:%M:%S", time.localtime())) kern.generateRandomNumber(device=device, size=matmul_dim) for i in range(steps): kern.matMulSimple2D(device=device, size=matmul_dim) kern.writeNonMPI(num_bytes=write_size, data_root_dir="./") kern.readNonMPI(num_bytes=read_size, data_root_dir="./") - log(f"Finished simulating...") - log(time.strftime("%H:%M:%S", time.localtime())) + logger.info(f"Finished simulating...") + logger.info(time.strftime("%H:%M:%S", time.localtime())) return random.random() @flow.function_task @@ -47,20 +50,20 @@ async def simulation(i): device = e["device"] matmul_dim = e["matmul_dim"] - log(f"Simulating {i}...") - log(time.strftime("%H:%M:%S", time.localtime())) + logger.info(f"Simulating {i}...") + logger.info(time.strftime("%H:%M:%S", time.localtime())) kern.readNonMPI(num_bytes=read_size, data_root_dir="./") kern.generateRandomNumber(device=device, size=matmul_dim) for _ in range(steps): kern.matMulSimple2D(device=device, size=matmul_dim) if i == 0: kern.writeNonMPI(num_bytes=write_size, data_root_dir="./tmp") - log(f"Finished simulating...") - log(time.strftime("%H:%M:%S", time.localtime())) + logger.info(f"Finished simulating...") + logger.info(time.strftime("%H:%M:%S", time.localtime())) return random.random() @flow.function_task - async def training(simulation, experiment): + async def training(_simulation, _experiment): e = cfg["training"] steps = e["steps"] read_size = e["read_size_bytes"] @@ -68,8 +71,8 @@ async def training(simulation, experiment): copy_size = e["data_copy_size_bytes"] matmul_dim = e["matmul_dim"] - log(f"Training model with simulation results...") - log(time.strftime("%H:%M:%S", time.localtime())) + logger.info(f"Training model with simulation results...") + logger.info(time.strftime("%H:%M:%S", time.localtime())) kern.readNonMPI(num_bytes=read_size, data_root_dir="./") kern.dataCopyH2D(data_size=copy_size) for i in range(steps): @@ -77,12 +80,12 @@ async def training(simulation, experiment): kern.matMulSimple2D(device="cpu", size=matmul_dim) kern.dataCopyH2D(data_size=copy_size) kern.writeNonMPI(num_bytes=write_size, data_root_dir="./") - log(f"Finished training model with evaluation results...") - log(time.strftime("%H:%M:%S", time.localtime())) + logger.info(f"Finished training model with evaluation results...") + logger.info(time.strftime("%H:%M:%S", time.localtime())) return random.random() @flow.function_task - async def inference(training): + async def inference(_training=None): e = cfg["inference"] steps = e["steps"] read_size = e["read_size_bytes"] @@ -90,8 +93,8 @@ async def inference(training): copy_size = e["data_copy_size_bytes"] matmul_dim = e["matmul_dim"] - log(f"Inferencing model with evaluation results...") - log(time.strftime("%H:%M:%S", time.localtime())) + logger.info(f"Inferencing model with evaluation results...") + logger.info(time.strftime("%H:%M:%S", time.localtime())) kern.readNonMPI(num_bytes=read_size, data_root_dir="./") kern.dataCopyH2D(data_size=copy_size) for i in range(steps): @@ -99,8 +102,8 @@ async def inference(training): kern.matMulSimple2D(device="cpu", size=matmul_dim) kern.dataCopyH2D(data_size=copy_size) kern.writeNonMPI(num_bytes=write_size, data_root_dir="./") - log(f"Finished inferencing model with evaluation results...") - log(time.strftime("%H:%M:%S", time.localtime())) + logger.info(f"Finished inferencing model with evaluation results...") + logger.info(time.strftime("%H:%M:%S", time.localtime())) return random.random() for i in range(3): @@ -121,4 +124,4 @@ async def inference(training): cfg = load_cfg(args.config) t = Timer().start() asyncio.run(workflow(cfg)) - log(f"DONE in {t.stop():.3f}s") + print(f"DONE in {t.stop():.3f}s") diff --git a/examples/motif_4/utils/metrics.py b/examples/motif_4/utils/metrics.py deleted file mode 100644 index 2d91aaf..0000000 --- a/examples/motif_4/utils/metrics.py +++ /dev/null @@ -1,8 +0,0 @@ -import time, sys - -class Timer: - def __init__(self): self.t0 = None - def start(self): self.t0 = time.time(); return self - def stop(self): return time.time() - self.t0 - -def log(msg): print(f"[dt] {msg}", file=sys.stdout, flush=True) \ No newline at end of file From f709c97ec95856f69a4f40f562a0b746b817b5b9 Mon Sep 17 00:00:00 2001 From: Andrew Park Date: Wed, 11 Mar 2026 16:34:13 -0400 Subject: [PATCH 38/46] Addressing issues --- examples/motif_3/miniapp_parallel.py | 14 +++++++++++--- examples/motif_3/miniapp_simple.py | 14 +++++++++++--- examples/motif_4/config.yaml | 2 +- examples/motif_4/miniapp_async.py | 12 ++++++++++-- wfMiniAPI/requirements.txt | 4 +++- 5 files changed, 36 insertions(+), 10 deletions(-) diff --git a/examples/motif_3/miniapp_parallel.py b/examples/motif_3/miniapp_parallel.py index 16d670d..5c2f839 100644 --- a/examples/motif_3/miniapp_parallel.py +++ b/examples/motif_3/miniapp_parallel.py @@ -1,8 +1,9 @@ import argparse, asyncio, yaml, random, logging import time +from rhapsody.backends import DragonExecutionBackendV3 + from radical.asyncflow import WorkflowEngine -from radical.asyncflow import DragonExecutionBackend from radical.asyncflow.logging import init_default_logger from wfMiniAPI import kernel as kern @@ -16,9 +17,16 @@ def load_cfg(path): with open(path, "r") as f: return yaml.safe_load(f) async def workflow(cfg): + import multiprocessing as mp logger = logging.getLogger(__name__) init_default_logger(logging.DEBUG) - backend = await DragonExecutionBackend() + + # Create Dragon Batch backend (1 nodes with 32 workers) + nodes = 1 + backend = await DragonExecutionBackendV3( + num_workers=nodes * mp.cpu_count(), + disable_background_batching=False, + ) flow = await WorkflowEngine.create(backend=backend) @flow.function_task @@ -83,4 +91,4 @@ async def train(_sim, _train=None): cfg = load_cfg(args.config) t = Timer().start() asyncio.run(workflow(cfg)) - print(f"DONE in {t.stop():.3f}s") \ No newline at end of file + print(f"DONE in {t.stop():.3f}s") diff --git a/examples/motif_3/miniapp_simple.py b/examples/motif_3/miniapp_simple.py index b42331a..1ad2ecb 100644 --- a/examples/motif_3/miniapp_simple.py +++ b/examples/motif_3/miniapp_simple.py @@ -1,8 +1,9 @@ import argparse, asyncio, yaml, random, logging import time +from rhapsody.backends import DragonExecutionBackendV3 + from radical.asyncflow import WorkflowEngine -from radical.asyncflow import DragonExecutionBackend from radical.asyncflow.logging import init_default_logger from wfMiniAPI import kernel as kern @@ -16,9 +17,16 @@ def load_cfg(path): with open(path, "r") as f: return yaml.safe_load(f) async def workflow(cfg): + import multiprocessing as mp logger = logging.getLogger(__name__) init_default_logger(logging.DEBUG) - backend = await DragonExecutionBackend() + + # Create Dragon Batch backend (1 nodes with 32 workers) + nodes = 1 + backend = await DragonExecutionBackendV3( + num_workers=nodes * mp.cpu_count(), + disable_background_batching=False, + ) flow = await WorkflowEngine.create(backend=backend) @flow.function_task @@ -77,4 +85,4 @@ async def train(_sim=None): cfg = load_cfg(args.config) t = Timer().start() asyncio.run(workflow(cfg)) - print(f"DONE in {t.stop():.3f}s") \ No newline at end of file + print(f"DONE in {t.stop():.3f}s") diff --git a/examples/motif_4/config.yaml b/examples/motif_4/config.yaml index f35526d..68c3e61 100644 --- a/examples/motif_4/config.yaml +++ b/examples/motif_4/config.yaml @@ -24,4 +24,4 @@ inference: read_size_bytes: 1342177280 # 1.25 GiB pre write_size_bytes: 2147483648 # 2 GiB post-step write data_copy_size_bytes: 2097152 # 2 MiB data copy - matmul_dim: 8192 # dimension for square matrix multiplication \ No newline at end of file + matmul_dim: 8192 # dimension for square matrix multiplication diff --git a/examples/motif_4/miniapp_async.py b/examples/motif_4/miniapp_async.py index 5dba85f..f923e3a 100644 --- a/examples/motif_4/miniapp_async.py +++ b/examples/motif_4/miniapp_async.py @@ -1,8 +1,9 @@ import argparse, asyncio, yaml, random, logging import time +from rhapsody.backends import DragonExecutionBackendV3 + from radical.asyncflow import WorkflowEngine -from radical.asyncflow import DragonExecutionBackend from radical.asyncflow.logging import init_default_logger from wfMiniAPI import kernel as kern @@ -16,9 +17,16 @@ def load_cfg(path): with open(path, "r") as f: return yaml.safe_load(f) async def workflow(cfg): + import multiprocessing as mp logger = logging.getLogger(__name__) init_default_logger(logging.DEBUG) - backend = await DragonExecutionBackend() + + # Create Dragon Batch backend (1 nodes with 32 workers) + nodes = 1 + backend = await DragonExecutionBackendV3( + num_workers=nodes * mp.cpu_count(), + disable_background_batching=False, + ) flow = await WorkflowEngine.create(backend=backend) @flow.function_task diff --git a/wfMiniAPI/requirements.txt b/wfMiniAPI/requirements.txt index 4b84da6..4fa0150 100644 --- a/wfMiniAPI/requirements.txt +++ b/wfMiniAPI/requirements.txt @@ -2,4 +2,6 @@ numpy h5py mpi4py radical-pilot -ROSE @ git+https://github.com/radical-cybertools/ROSE.git +radical-asyncflow +dragonhpc +rhapsody-py[all] From c2cc61cf093aaf2f24a2e59effe1457f2a6fab1e Mon Sep 17 00:00:00 2001 From: Andrew Park Date: Wed, 11 Mar 2026 16:36:47 -0400 Subject: [PATCH 39/46] Forgot this one --- wfMiniAPI/requirements-gpu.txt | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/wfMiniAPI/requirements-gpu.txt b/wfMiniAPI/requirements-gpu.txt index c2348ec..5479d79 100644 --- a/wfMiniAPI/requirements-gpu.txt +++ b/wfMiniAPI/requirements-gpu.txt @@ -3,4 +3,6 @@ h5py mpi4py cupy radical-pilot -ROSE @ git+https://github.com/radical-cybertools/ROSE.git +radical-asyncflow +dragonhpc +rhapsody-py[all] From 096de7e93d73cdceec1a4ba8a289f8156d9bca61 Mon Sep 17 00:00:00 2001 From: Andrew Park Date: Wed, 11 Mar 2026 20:01:02 -0400 Subject: [PATCH 40/46] fix dragon --- examples/motif_3/miniapp_parallel.py | 3 ++- examples/motif_3/miniapp_simple.py | 1 + examples/motif_4/miniapp_async.py | 1 + 3 files changed, 4 insertions(+), 1 deletion(-) diff --git a/examples/motif_3/miniapp_parallel.py b/examples/motif_3/miniapp_parallel.py index 5c2f839..d4ac600 100644 --- a/examples/motif_3/miniapp_parallel.py +++ b/examples/motif_3/miniapp_parallel.py @@ -20,9 +20,10 @@ async def workflow(cfg): import multiprocessing as mp logger = logging.getLogger(__name__) init_default_logger(logging.DEBUG) - + # Create Dragon Batch backend (1 nodes with 32 workers) nodes = 1 + mp.set_start_method("dragon") backend = await DragonExecutionBackendV3( num_workers=nodes * mp.cpu_count(), disable_background_batching=False, diff --git a/examples/motif_3/miniapp_simple.py b/examples/motif_3/miniapp_simple.py index 1ad2ecb..df6aff4 100644 --- a/examples/motif_3/miniapp_simple.py +++ b/examples/motif_3/miniapp_simple.py @@ -23,6 +23,7 @@ async def workflow(cfg): # Create Dragon Batch backend (1 nodes with 32 workers) nodes = 1 + mp.set_start_method("dragon") backend = await DragonExecutionBackendV3( num_workers=nodes * mp.cpu_count(), disable_background_batching=False, diff --git a/examples/motif_4/miniapp_async.py b/examples/motif_4/miniapp_async.py index f923e3a..2442b46 100644 --- a/examples/motif_4/miniapp_async.py +++ b/examples/motif_4/miniapp_async.py @@ -23,6 +23,7 @@ async def workflow(cfg): # Create Dragon Batch backend (1 nodes with 32 workers) nodes = 1 + mp.set_start_method("dragon") backend = await DragonExecutionBackendV3( num_workers=nodes * mp.cpu_count(), disable_background_batching=False, From 1d11b4395d58e405216247623fb894dc0fe7caea Mon Sep 17 00:00:00 2001 From: Andrew Park Date: Fri, 13 Mar 2026 16:30:42 -0500 Subject: [PATCH 41/46] Add init file to fix setup --- wfMiniAPI/src/wfMiniAPI/__init__.py | 7 +++++++ 1 file changed, 7 insertions(+) create mode 100644 wfMiniAPI/src/wfMiniAPI/__init__.py diff --git a/wfMiniAPI/src/wfMiniAPI/__init__.py b/wfMiniAPI/src/wfMiniAPI/__init__.py new file mode 100644 index 0000000..d6ca307 --- /dev/null +++ b/wfMiniAPI/src/wfMiniAPI/__init__.py @@ -0,0 +1,7 @@ +from . import kernel, registry, sim + +__all__ = [ + 'kernel', + 'registry', + 'sim', +] \ No newline at end of file From dfd3b0e99efa40870bef93f2e9580e58bc1e2d2e Mon Sep 17 00:00:00 2001 From: Andrew Park Date: Mon, 16 Mar 2026 05:08:39 -0500 Subject: [PATCH 42/46] Running with dragon --- examples/motif_4/miniapp_async.py | 20 ++++++++++---------- 1 file changed, 10 insertions(+), 10 deletions(-) diff --git a/examples/motif_4/miniapp_async.py b/examples/motif_4/miniapp_async.py index 2442b46..b55bb4a 100644 --- a/examples/motif_4/miniapp_async.py +++ b/examples/motif_4/miniapp_async.py @@ -42,7 +42,7 @@ async def experiment(): logger.info(f"Experiment Data...") logger.info(time.strftime("%H:%M:%S", time.localtime())) kern.generateRandomNumber(device=device, size=matmul_dim) - for i in range(steps): + for j in range(steps): kern.matMulSimple2D(device=device, size=matmul_dim) kern.writeNonMPI(num_bytes=write_size, data_root_dir="./") kern.readNonMPI(num_bytes=read_size, data_root_dir="./") @@ -51,7 +51,7 @@ async def experiment(): return random.random() @flow.function_task - async def simulation(i): + async def simulation(i, _experiment): e = cfg["simulation"] steps = e["steps"] read_size = e["read_size_bytes"] @@ -66,13 +66,13 @@ async def simulation(i): for _ in range(steps): kern.matMulSimple2D(device=device, size=matmul_dim) if i == 0: - kern.writeNonMPI(num_bytes=write_size, data_root_dir="./tmp") + kern.writeNonMPI(num_bytes=write_size, data_root_dir="./") logger.info(f"Finished simulating...") logger.info(time.strftime("%H:%M:%S", time.localtime())) return random.random() @flow.function_task - async def training(_simulation, _experiment): + async def training(): e = cfg["training"] steps = e["steps"] read_size = e["read_size_bytes"] @@ -82,19 +82,19 @@ async def training(_simulation, _experiment): logger.info(f"Training model with simulation results...") logger.info(time.strftime("%H:%M:%S", time.localtime())) + kern.writeNonMPI(num_bytes=write_size, data_root_dir="./") kern.readNonMPI(num_bytes=read_size, data_root_dir="./") kern.dataCopyH2D(data_size=copy_size) for i in range(steps): kern.matMulSimple2D(device="gpu", size=matmul_dim) kern.matMulSimple2D(device="cpu", size=matmul_dim) kern.dataCopyH2D(data_size=copy_size) - kern.writeNonMPI(num_bytes=write_size, data_root_dir="./") logger.info(f"Finished training model with evaluation results...") logger.info(time.strftime("%H:%M:%S", time.localtime())) return random.random() @flow.function_task - async def inference(_training=None): + async def inference(_training): e = cfg["inference"] steps = e["steps"] read_size = e["read_size_bytes"] @@ -115,13 +115,13 @@ async def inference(_training=None): logger.info(time.strftime("%H:%M:%S", time.localtime())) return random.random() - for i in range(3): - experiment_t = await experiment() + for j in range(3): sim_t = [] + experiment_t = experiment() for i in range(32): - sim_t.append(simulation(i)) + sim_t.append(simulation(i, experiment_t)) await asyncio.gather(*sim_t) - train_t = training(sim_t, experiment_t) + train_t = training() infer_t = await inference(train_t) await flow.shutdown() From 1c345804020a86408c42ec93f793a85ed41908f0 Mon Sep 17 00:00:00 2001 From: Andrew Park Date: Mon, 16 Mar 2026 14:32:45 -0500 Subject: [PATCH 43/46] Add instructions for how to install and run in README --- examples/motif_3/README.md | 63 ++++++++++++++++++++++++++++++++++++++ examples/motif_4/README.md | 61 ++++++++++++++++++++++++++++++++++++ 2 files changed, 124 insertions(+) create mode 100644 examples/motif_3/README.md create mode 100644 examples/motif_4/README.md diff --git a/examples/motif_3/README.md b/examples/motif_3/README.md new file mode 100644 index 0000000..f668caa --- /dev/null +++ b/examples/motif_3/README.md @@ -0,0 +1,63 @@ +# Inverse Design Motif Workflow Mini-app +Workflow mini-app build based on an Inverse Design workflow. + +To execute the mini-app please follow these steps + +### Load your modules +First load any modules you need for your environment. +We did our experiments on Delta, so your needed modules may change. +Generally + +- `$ module load cudatoolkit/25.3_12.8` +- `$ module load python/3.13.5-gcc13.3.1` + +### If using a venv, load the venv now + +- `$ source path_to_venv/bin/activate` + +If you have already installed wfMiniAPI, you can skip ahead to running the workflow. +### Installing wfMiniAPI and its dependencies +First clone the wfMiniAPI repository +- `$ git clone git@github.com:radical-cybertools/workflow-mini-apps.git` +- `$ cd workflow-mini-apps/` + +Now before installing, we must patch the kernel.py to use the correct axpy kernel. +Edit `wfMiniAPI/src/wfMiniAPI/kernel.py` by uncommenting lines `305-313` and commenting lines `315-321`. +It should look like this: +```python +@annotate_kernel +def axpy_fuse(device, size): + xp = get_device_module(device) + x = xp.empty(size, dtype=xp.float32) + y = xp.empty(size, dtype=xp.float32) + if xp == np: + y += 1.01 * x + elif xp == cp: + _axpy_fuse(1.01, x, y, size=size) + +#_axpy_fuse_fast = cp.ElementwiseKernel( +# 'float32 alpha, raw float32 x', +# 'raw float32 y', +# 'y[i] += alpha * x[i]', +# 'axpy_fuse_kernel', +# no_return=True +#) +``` +There are also some dependencies which are not installed by `pip`. +You will need `cupy`, the specific version you need is determined by the CUDA version. +In our case, on DELTA we are using CUDA 12.8 +- `$ pip install cupy-cuda12x` + +Now we can install the wfMiniAPI +- `$ cd wfMiniAPI` +- `$ pip install .` + +### Running the workflow +Navigate to the example motif and launch the workflow mini-app using Dragon +- `$ cd workflow-mini-apps/examples/motif_3` +- `$ dragon -s workflow_simple.py` +or + `$ dragon -s workflow_parallel.py` + +This will launch the workflows using Dragon with a single node. +Remove the `-s` to run the workflow with more than one node. \ No newline at end of file diff --git a/examples/motif_4/README.md b/examples/motif_4/README.md new file mode 100644 index 0000000..e6c9727 --- /dev/null +++ b/examples/motif_4/README.md @@ -0,0 +1,61 @@ +# Digital Replica Motif Workflow Mini-app +Workflow mini-app build based on an Digital Replica workflow. + +To execute the mini-app please follow these steps + +### Load your modules +First load any modules you need for your environment. +We did our experiments on Delta, so your needed modules may change. +Generally + +- `$ module load cudatoolkit/25.3_12.8` +- `$ module load python/3.13.5-gcc13.3.1` + +### If using a venv, load the venv now + +- `$ source path_to_venv/bin/activate` + +If you have already installed wfMiniAPI, you can skip ahead to running the workflow. +### Installing wfMiniAPI and its dependencies +First clone the wfMiniAPI repository +- `$ git clone git@github.com:radical-cybertools/workflow-mini-apps.git` +- `$ cd workflow-mini-apps/` + +Now before installing, we must patch the kernel.py to use the correct axpy kernel. +Edit `wfMiniAPI/src/wfMiniAPI/kernel.py` by uncommenting lines `305-313` and commenting lines `315-321`. +It should look like this: +```python +@annotate_kernel +def axpy_fuse(device, size): + xp = get_device_module(device) + x = xp.empty(size, dtype=xp.float32) + y = xp.empty(size, dtype=xp.float32) + if xp == np: + y += 1.01 * x + elif xp == cp: + _axpy_fuse(1.01, x, y, size=size) + +#_axpy_fuse_fast = cp.ElementwiseKernel( +# 'float32 alpha, raw float32 x', +# 'raw float32 y', +# 'y[i] += alpha * x[i]', +# 'axpy_fuse_kernel', +# no_return=True +#) +``` +There are also some dependencies which are not installed by `pip`. +You will need `cupy`, the specific version you need is determined by the CUDA version. +In our case, on DELTA we are using CUDA 12.8 +- `$ pip install cupy-cuda12x` + +Now we can install the wfMiniAPI +- `$ cd wfMiniAPI` +- `$ pip install .` + +### Running the workflow +Navigate to the example motif and launch the workflow mini-app using Dragon +- `$ cd workflow-mini-apps/examples/motif_4` +- `$ dragon -s workflow_async.py` + +This will launch the workflows using Dragon with a single node. +Remove the `-s` to run the workflow with more than one node. \ No newline at end of file From 177a117e4d04dd598f7f43d720233cc3bdfe974c Mon Sep 17 00:00:00 2001 From: Andrew Park Date: Wed, 25 Mar 2026 12:10:53 -0500 Subject: [PATCH 44/46] dragon mpi motif --- examples/dragon-mpi/README.md | 20 ++++ examples/dragon-mpi/config.yaml | 31 ++++++ examples/dragon-mpi/miniapp_mpi.py | 164 +++++++++++++++++++++++++++++ 3 files changed, 215 insertions(+) create mode 100644 examples/dragon-mpi/README.md create mode 100644 examples/dragon-mpi/config.yaml create mode 100644 examples/dragon-mpi/miniapp_mpi.py diff --git a/examples/dragon-mpi/README.md b/examples/dragon-mpi/README.md new file mode 100644 index 0000000..edb2df9 --- /dev/null +++ b/examples/dragon-mpi/README.md @@ -0,0 +1,20 @@ +# Workflow mini-app using MPI and Rhapsody with Dragon backend +This example goes over launching MPI tasks +## First load the proper modules + +### Load cray-mpich-abi +`$ module load cray-mpich-abi/8.1.32` + +### Load cuda toolkit +`$ module load cudatoolkit/25.3_12.8` + +### If using a venv, load the venv now +`$ source path_to_venv/bin/activate` + +### Now you can launch using dragon +`$ dragon launch.py` + +> [!NOTE] +> DRAGON assumes it is launched via SLURM and will use SLURM environment variables. +> Trying to run DRAGON without these set will cause a crash or hang. +> This can be manually fixed by running `export SLURM_JOB_NUM_NODES=1` (or however many nodes you would like). \ No newline at end of file diff --git a/examples/dragon-mpi/config.yaml b/examples/dragon-mpi/config.yaml new file mode 100644 index 0000000..c63bcba --- /dev/null +++ b/examples/dragon-mpi/config.yaml @@ -0,0 +1,31 @@ +stage1: + ranks: 2 + steps: 2 + read_size_bytes: 1024 # 1 KiB pre-step read + write_size_bytes: 1342177280 # 1.25 GiB post-step write + device: "cpu" # "cpu" or "gpu" + matmul_dim: 8192 # dimension for square matrix multiplication + +stage2: + ranks: 4 + steps: 2 + read_size_bytes: 1048576 # 1 MiB pre-step read + write_size_bytes: 2147483648 # 2 GiB post-step write + device: "cpu" # "cpu" or "gpu" + matmul_dim: 2048 # dimension for square matrix multiplication + +stage3: + ranks: 8 + steps: 2 + read_size_bytes: 2147483648 # 2 GiB pre-step read + write_size_bytes: 512 # 512 B post-step write + data_copy_size_bytes: 2097152 # 2 MiB data copy + matmul_dim: 8192 # dimension for square matrix multiplication + +stage4: + ranks: 16 + steps: 2 + read_size_bytes: 1342177280 # 1.25 GiB pre + write_size_bytes: 2147483648 # 2 GiB post-step write + data_copy_size_bytes: 2097152 # 2 MiB data copy + matmul_dim: 8192 # dimension for square matrix multiplication diff --git a/examples/dragon-mpi/miniapp_mpi.py b/examples/dragon-mpi/miniapp_mpi.py new file mode 100644 index 0000000..59e6b46 --- /dev/null +++ b/examples/dragon-mpi/miniapp_mpi.py @@ -0,0 +1,164 @@ +import argparse, asyncio, yaml, random, logging, os +import time + +from rhapsody.backends import DragonExecutionBackendV3 + +from radical.asyncflow import WorkflowEngine +from radical.asyncflow.logging import init_default_logger + +from wfMiniAPI import kernel as kern + +class Timer: + def __init__(self): self.t0 = None + def start(self): self.t0 = time.time(); return self + def stop(self): return time.time() - self.t0 + +def load_cfg(path): + with open(path, "r") as f: return yaml.safe_load(f) + +async def workflow(cfg): + import multiprocessing as mp + import mpi4py + from mpi4py import MPI + logger = logging.getLogger(__name__) + init_default_logger(logging.DEBUG) + + # Create Dragon Batch backend (1 nodes with 32 workers) + nodes = 1 + mp.set_start_method("dragon") + backend = await DragonExecutionBackendV3( + num_workers=nodes * mp.cpu_count(), + disable_background_batching=False, + ) + flow = await WorkflowEngine.create(backend=backend) + + #MPI.Init() + os.makedirs("./input", exist_ok=True) + os.makedirs("./output", exist_ok=True) + kern.writeNonMPI(num_bytes=64, data_root_dir="./input") + + s1 = cfg["stage1"] + @flow.function_task + async def stage1(task_description={'ranks': s1['ranks'], 'type': 'mpi'}, *args): + logger.info(time.strftime("%H:%M:%S", time.localtime())) + import mpi4py + from mpi4py import MPI + s1 = cfg["stage1"] + comm = MPI.COMM_WORLD + rank = comm.Get_rank() + size = comm.Get_size() + print(f"Rank {rank} of {size} says: Hello from Stage 1!", flush=True) + steps = s1["steps"] + read_size = s1["read_size_bytes"] + write_size = s1["write_size_bytes"] + device = s1["device"] + matmul_dim = s1["matmul_dim"] + + kern.readWithMPI(num_bytes=read_size, data_root_dir="./input") + kern.generateRandomNumber(device=device, size=matmul_dim) + for j in range(steps): + kern.matMulSimple2D(device=device, size=matmul_dim) + kern.writeWithMPI(num_bytes=write_size, data_root_dir="./output") + logger.info(f"Finished stage 1...") + logger.info(time.strftime("%H:%M:%S", time.localtime())) + return random.random() + return rank + + s2 = cfg["stage2"] + @flow.function_task + async def stage2(task_description={'ranks': s2['ranks'], 'type': 'mpi'}, *args): + logger.info(time.strftime("%H:%M:%S", time.localtime())) + import mpi4py + from mpi4py import MPI + comm = MPI.COMM_WORLD + rank = comm.Get_rank() + size = comm.Get_size() + print(f"Rank {rank} of {size} says: Hello from Stage 2!", flush=True) + s2 = cfg["stage2"] + steps = s2["steps"] + read_size = s2["read_size_bytes"] + write_size = s2["write_size_bytes"] + device = s2["device"] + matmul_dim = s2["matmul_dim"] + + kern.readWithMPI(num_bytes=read_size, data_root_dir="./input") + kern.generateRandomNumber(device=device, size=matmul_dim) + for _ in range(steps): + kern.matMulSimple2D(device=device, size=matmul_dim) + kern.writeWithMPI(num_bytes=write_size, data_root_dir="./output") + logger.info(f"Finished stage 2...") + logger.info(time.strftime("%H:%M:%S", time.localtime())) + return random.random() + + s3 = cfg["stage3"] + @flow.function_task + async def stage3(task_description={'ranks': s3['ranks'], 'type': 'mpi'}, *args): + logger.info(time.strftime("%H:%M:%S", time.localtime())) + import mpi4py + from mpi4py import MPI + comm = MPI.COMM_WORLD + rank = comm.Get_rank() + size = comm.Get_size() + print(f"Rank {rank} of {size} says: Hello from Stage 3!", flush=True) + s3 = cfg["stage3"] + steps = s3["steps"] + read_size = s3["read_size_bytes"] + write_size = s3["write_size_bytes"] + copy_size = s3["data_copy_size_bytes"] + matmul_dim = s3["matmul_dim"] + + kern.readWithMPI(num_bytes=read_size, data_root_dir="./input") + kern.dataCopyH2D(data_size=copy_size) + for i in range(steps): + kern.matMulSimple2D(device="gpu", size=matmul_dim) + kern.matMulSimple2D(device="cpu", size=matmul_dim) + kern.dataCopyH2D(data_size=copy_size) + kern.writeWithMPI(num_bytes=write_size, data_root_dir="./output") + logger.info(f"Finished stage 3...") + logger.info(time.strftime("%H:%M:%S", time.localtime())) + return random.random() + + s4 = cfg["stage4"] + @flow.function_task + async def stage4(task_description={'ranks': s4['ranks'], 'type': 'mpi'}, *args): + logger.info(time.strftime("%H:%M:%S", time.localtime())) + import mpi4py + from mpi4py import MPI + comm = MPI.COMM_WORLD + rank = comm.Get_rank() + size = comm.Get_size() + print(f"Rank {rank} of {size} says: Hello from Stage 4!", flush=True) + s4 = cfg["stage4"] + steps = s4["steps"] + read_size = s4["read_size_bytes"] + write_size = s4["write_size_bytes"] + copy_size = s4["data_copy_size_bytes"] + matmul_dim = s4["matmul_dim"] + + logger.info(time.strftime("%H:%M:%S", time.localtime())) + kern.readWithMPI(num_bytes=read_size, data_root_dir="./input") + kern.dataCopyH2D(data_size=copy_size) + for i in range(steps): + kern.matMulSimple2D(device="gpu", size=matmul_dim) + kern.matMulSimple2D(device="cpu", size=matmul_dim) + kern.dataCopyH2D(data_size=copy_size) + kern.writeWithMPI(num_bytes=write_size, data_root_dir="./output") + logger.info(f"Finished stage 4...") + logger.info(time.strftime("%H:%M:%S", time.localtime())) + return random.random() + + stage1_t = stage1() + stage2_t = stage2(stage1_t) + stage3_t = stage3(stage1_t) + stage4_t = await stage4(stage2_t, stage3_t) + + await flow.shutdown() + +if __name__ == "__main__": + ap = argparse.ArgumentParser(description="MPI with Dragon Mini-App") + ap.add_argument("--config", type=str, default="config.yaml") + args = ap.parse_args() + cfg = load_cfg(args.config) + t = Timer().start() + asyncio.run(workflow(cfg)) + print(f"DONE in {t.stop():.3f}s") From 08cd0374872a1ae52aa69fdd8457d467cb47f585 Mon Sep 17 00:00:00 2001 From: Andrew Park Date: Wed, 25 Mar 2026 13:02:17 -0500 Subject: [PATCH 45/46] README includes sectino on h5py --- examples/dragon-mpi/README.md | 9 ++++++--- .../{miniapp_mpi.py => miniapp_rhapsody.py} | 20 ++++++++++--------- 2 files changed, 17 insertions(+), 12 deletions(-) rename examples/dragon-mpi/{miniapp_mpi.py => miniapp_rhapsody.py} (88%) diff --git a/examples/dragon-mpi/README.md b/examples/dragon-mpi/README.md index edb2df9..731eb5a 100644 --- a/examples/dragon-mpi/README.md +++ b/examples/dragon-mpi/README.md @@ -3,16 +3,19 @@ This example goes over launching MPI tasks ## First load the proper modules ### Load cray-mpich-abi -`$ module load cray-mpich-abi/8.1.32` +`$ module load cray-mpich-abi` ### Load cuda toolkit -`$ module load cudatoolkit/25.3_12.8` +`$ module load cudatoolkit` + +### Load h5py with MPI support +`$ module load cray-hdf5` ### If using a venv, load the venv now `$ source path_to_venv/bin/activate` ### Now you can launch using dragon -`$ dragon launch.py` +`$ dragon miniapp_mpi.py` > [!NOTE] > DRAGON assumes it is launched via SLURM and will use SLURM environment variables. diff --git a/examples/dragon-mpi/miniapp_mpi.py b/examples/dragon-mpi/miniapp_rhapsody.py similarity index 88% rename from examples/dragon-mpi/miniapp_mpi.py rename to examples/dragon-mpi/miniapp_rhapsody.py index 59e6b46..1b98800 100644 --- a/examples/dragon-mpi/miniapp_mpi.py +++ b/examples/dragon-mpi/miniapp_rhapsody.py @@ -39,7 +39,9 @@ async def workflow(cfg): s1 = cfg["stage1"] @flow.function_task - async def stage1(task_description={'ranks': s1['ranks'], 'type': 'mpi'}, *args): + async def stage1(task_backend_specific_kwargs={ + "process_templates": [(2, {}), (2, {})] + }, *args): logger.info(time.strftime("%H:%M:%S", time.localtime())) import mpi4py from mpi4py import MPI @@ -54,11 +56,11 @@ async def stage1(task_description={'ranks': s1['ranks'], 'type': 'mpi'}, *args): device = s1["device"] matmul_dim = s1["matmul_dim"] - kern.readWithMPI(num_bytes=read_size, data_root_dir="./input") + #kern.readWithMPI(num_bytes=read_size, data_root_dir="./input") kern.generateRandomNumber(device=device, size=matmul_dim) for j in range(steps): kern.matMulSimple2D(device=device, size=matmul_dim) - kern.writeWithMPI(num_bytes=write_size, data_root_dir="./output") + #kern.writeWithMPI(num_bytes=write_size, data_root_dir="./output") logger.info(f"Finished stage 1...") logger.info(time.strftime("%H:%M:%S", time.localtime())) return random.random() @@ -81,11 +83,11 @@ async def stage2(task_description={'ranks': s2['ranks'], 'type': 'mpi'}, *args): device = s2["device"] matmul_dim = s2["matmul_dim"] - kern.readWithMPI(num_bytes=read_size, data_root_dir="./input") + #kern.readWithMPI(num_bytes=read_size, data_root_dir="./input") kern.generateRandomNumber(device=device, size=matmul_dim) for _ in range(steps): kern.matMulSimple2D(device=device, size=matmul_dim) - kern.writeWithMPI(num_bytes=write_size, data_root_dir="./output") + #kern.writeWithMPI(num_bytes=write_size, data_root_dir="./output") logger.info(f"Finished stage 2...") logger.info(time.strftime("%H:%M:%S", time.localtime())) return random.random() @@ -107,13 +109,13 @@ async def stage3(task_description={'ranks': s3['ranks'], 'type': 'mpi'}, *args): copy_size = s3["data_copy_size_bytes"] matmul_dim = s3["matmul_dim"] - kern.readWithMPI(num_bytes=read_size, data_root_dir="./input") + #kern.readWithMPI(num_bytes=read_size, data_root_dir="./input") kern.dataCopyH2D(data_size=copy_size) for i in range(steps): kern.matMulSimple2D(device="gpu", size=matmul_dim) kern.matMulSimple2D(device="cpu", size=matmul_dim) kern.dataCopyH2D(data_size=copy_size) - kern.writeWithMPI(num_bytes=write_size, data_root_dir="./output") + #kern.writeWithMPI(num_bytes=write_size, data_root_dir="./output") logger.info(f"Finished stage 3...") logger.info(time.strftime("%H:%M:%S", time.localtime())) return random.random() @@ -136,13 +138,13 @@ async def stage4(task_description={'ranks': s4['ranks'], 'type': 'mpi'}, *args): matmul_dim = s4["matmul_dim"] logger.info(time.strftime("%H:%M:%S", time.localtime())) - kern.readWithMPI(num_bytes=read_size, data_root_dir="./input") + #kern.readWithMPI(num_bytes=read_size, data_root_dir="./input") kern.dataCopyH2D(data_size=copy_size) for i in range(steps): kern.matMulSimple2D(device="gpu", size=matmul_dim) kern.matMulSimple2D(device="cpu", size=matmul_dim) kern.dataCopyH2D(data_size=copy_size) - kern.writeWithMPI(num_bytes=write_size, data_root_dir="./output") + #kern.writeWithMPI(num_bytes=write_size, data_root_dir="./output") logger.info(f"Finished stage 4...") logger.info(time.strftime("%H:%M:%S", time.localtime())) return random.random() From 735ffa9235de66f05299e4a605cf085a82843d2e Mon Sep 17 00:00:00 2001 From: iznoanygod Date: Fri, 19 Jun 2026 14:12:26 -0400 Subject: [PATCH 46/46] fixes to make this work with new version of rhapsody and dragon --- examples/dragon-mpi/miniapp_rhapsody.py | 28 +++++++++---------------- examples/motif_3/miniapp_parallel.py | 7 +------ examples/motif_3/miniapp_simple.py | 7 +------ examples/motif_4/miniapp_async.py | 7 +------ 4 files changed, 13 insertions(+), 36 deletions(-) diff --git a/examples/dragon-mpi/miniapp_rhapsody.py b/examples/dragon-mpi/miniapp_rhapsody.py index 1b98800..1fcbc09 100644 --- a/examples/dragon-mpi/miniapp_rhapsody.py +++ b/examples/dragon-mpi/miniapp_rhapsody.py @@ -23,25 +23,17 @@ async def workflow(cfg): logger = logging.getLogger(__name__) init_default_logger(logging.DEBUG) - # Create Dragon Batch backend (1 nodes with 32 workers) - nodes = 1 mp.set_start_method("dragon") - backend = await DragonExecutionBackendV3( - num_workers=nodes * mp.cpu_count(), - disable_background_batching=False, - ) + backend = await DragonExecutionBackendV3() flow = await WorkflowEngine.create(backend=backend) - #MPI.Init() os.makedirs("./input", exist_ok=True) os.makedirs("./output", exist_ok=True) kern.writeNonMPI(num_bytes=64, data_root_dir="./input") s1 = cfg["stage1"] @flow.function_task - async def stage1(task_backend_specific_kwargs={ - "process_templates": [(2, {}), (2, {})] - }, *args): + async def stage1(task_description={'ranks': s1['ranks'], 'type': 'mpi'}, *args): logger.info(time.strftime("%H:%M:%S", time.localtime())) import mpi4py from mpi4py import MPI @@ -56,11 +48,11 @@ async def stage1(task_backend_specific_kwargs={ device = s1["device"] matmul_dim = s1["matmul_dim"] - #kern.readWithMPI(num_bytes=read_size, data_root_dir="./input") + kern.readWithMPI(num_bytes=read_size, data_root_dir="./input") kern.generateRandomNumber(device=device, size=matmul_dim) for j in range(steps): kern.matMulSimple2D(device=device, size=matmul_dim) - #kern.writeWithMPI(num_bytes=write_size, data_root_dir="./output") + kern.writeWithMPI(num_bytes=write_size, data_root_dir="./output") logger.info(f"Finished stage 1...") logger.info(time.strftime("%H:%M:%S", time.localtime())) return random.random() @@ -83,11 +75,11 @@ async def stage2(task_description={'ranks': s2['ranks'], 'type': 'mpi'}, *args): device = s2["device"] matmul_dim = s2["matmul_dim"] - #kern.readWithMPI(num_bytes=read_size, data_root_dir="./input") + kern.readWithMPI(num_bytes=read_size, data_root_dir="./input") kern.generateRandomNumber(device=device, size=matmul_dim) for _ in range(steps): kern.matMulSimple2D(device=device, size=matmul_dim) - #kern.writeWithMPI(num_bytes=write_size, data_root_dir="./output") + kern.writeWithMPI(num_bytes=write_size, data_root_dir="./output") logger.info(f"Finished stage 2...") logger.info(time.strftime("%H:%M:%S", time.localtime())) return random.random() @@ -109,13 +101,13 @@ async def stage3(task_description={'ranks': s3['ranks'], 'type': 'mpi'}, *args): copy_size = s3["data_copy_size_bytes"] matmul_dim = s3["matmul_dim"] - #kern.readWithMPI(num_bytes=read_size, data_root_dir="./input") + kern.readWithMPI(num_bytes=read_size, data_root_dir="./input") kern.dataCopyH2D(data_size=copy_size) for i in range(steps): kern.matMulSimple2D(device="gpu", size=matmul_dim) kern.matMulSimple2D(device="cpu", size=matmul_dim) kern.dataCopyH2D(data_size=copy_size) - #kern.writeWithMPI(num_bytes=write_size, data_root_dir="./output") + kern.writeWithMPI(num_bytes=write_size, data_root_dir="./output") logger.info(f"Finished stage 3...") logger.info(time.strftime("%H:%M:%S", time.localtime())) return random.random() @@ -138,13 +130,13 @@ async def stage4(task_description={'ranks': s4['ranks'], 'type': 'mpi'}, *args): matmul_dim = s4["matmul_dim"] logger.info(time.strftime("%H:%M:%S", time.localtime())) - #kern.readWithMPI(num_bytes=read_size, data_root_dir="./input") + kern.readWithMPI(num_bytes=read_size, data_root_dir="./input") kern.dataCopyH2D(data_size=copy_size) for i in range(steps): kern.matMulSimple2D(device="gpu", size=matmul_dim) kern.matMulSimple2D(device="cpu", size=matmul_dim) kern.dataCopyH2D(data_size=copy_size) - #kern.writeWithMPI(num_bytes=write_size, data_root_dir="./output") + kern.writeWithMPI(num_bytes=write_size, data_root_dir="./output") logger.info(f"Finished stage 4...") logger.info(time.strftime("%H:%M:%S", time.localtime())) return random.random() diff --git a/examples/motif_3/miniapp_parallel.py b/examples/motif_3/miniapp_parallel.py index d4ac600..cc7a6ca 100644 --- a/examples/motif_3/miniapp_parallel.py +++ b/examples/motif_3/miniapp_parallel.py @@ -21,13 +21,8 @@ async def workflow(cfg): logger = logging.getLogger(__name__) init_default_logger(logging.DEBUG) - # Create Dragon Batch backend (1 nodes with 32 workers) - nodes = 1 mp.set_start_method("dragon") - backend = await DragonExecutionBackendV3( - num_workers=nodes * mp.cpu_count(), - disable_background_batching=False, - ) + backend = await DragonExecutionBackendV3() flow = await WorkflowEngine.create(backend=backend) @flow.function_task diff --git a/examples/motif_3/miniapp_simple.py b/examples/motif_3/miniapp_simple.py index df6aff4..885f994 100644 --- a/examples/motif_3/miniapp_simple.py +++ b/examples/motif_3/miniapp_simple.py @@ -21,13 +21,8 @@ async def workflow(cfg): logger = logging.getLogger(__name__) init_default_logger(logging.DEBUG) - # Create Dragon Batch backend (1 nodes with 32 workers) - nodes = 1 mp.set_start_method("dragon") - backend = await DragonExecutionBackendV3( - num_workers=nodes * mp.cpu_count(), - disable_background_batching=False, - ) + backend = await DragonExecutionBackendV3() flow = await WorkflowEngine.create(backend=backend) @flow.function_task diff --git a/examples/motif_4/miniapp_async.py b/examples/motif_4/miniapp_async.py index b55bb4a..5967016 100644 --- a/examples/motif_4/miniapp_async.py +++ b/examples/motif_4/miniapp_async.py @@ -21,13 +21,8 @@ async def workflow(cfg): logger = logging.getLogger(__name__) init_default_logger(logging.DEBUG) - # Create Dragon Batch backend (1 nodes with 32 workers) - nodes = 1 mp.set_start_method("dragon") - backend = await DragonExecutionBackendV3( - num_workers=nodes * mp.cpu_count(), - disable_background_batching=False, - ) + backend = await DragonExecutionBackendV3() flow = await WorkflowEngine.create(backend=backend) @flow.function_task