diff --git a/examples/Tutorial/tasks/active_learn.py b/examples/Tutorial/tasks/active_learn.py new file mode 100644 index 0000000..9e85304 --- /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='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, + 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_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) + + 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)) + +if __name__ == '__main__': + main() diff --git a/examples/Tutorial/tasks/simulation.py b/examples/Tutorial/tasks/simulation.py new file mode 100644 index 0000000..c3fb78e --- /dev/null +++ b/examples/Tutorial/tasks/simulation.py @@ -0,0 +1,83 @@ +#!/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 % args.log_freq == 0: + dir_name = os.path.join(root_path, f"./log_step_{i}") + os.makedirs(dir_name, exist_ok=True) + 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__': + main() + diff --git a/examples/Tutorial/tasks/training.py b/examples/Tutorial/tasks/training.py new file mode 100644 index 0000000..1ef1f9a --- /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='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, + 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_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) + 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/examples/dragon-mpi/README.md b/examples/dragon-mpi/README.md new file mode 100644 index 0000000..731eb5a --- /dev/null +++ b/examples/dragon-mpi/README.md @@ -0,0 +1,23 @@ +# 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` + +### Load cuda toolkit +`$ 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 miniapp_mpi.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_rhapsody.py b/examples/dragon-mpi/miniapp_rhapsody.py new file mode 100644 index 0000000..1fcbc09 --- /dev/null +++ b/examples/dragon-mpi/miniapp_rhapsody.py @@ -0,0 +1,158 @@ +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) + + mp.set_start_method("dragon") + backend = await DragonExecutionBackendV3() + flow = await WorkflowEngine.create(backend=backend) + + 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") 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_3/config.yaml b/examples/motif_3/config.yaml new file mode 100644 index 0000000..74ae1c6 --- /dev/null +++ b/examples/motif_3/config.yaml @@ -0,0 +1,14 @@ +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 + data_copy_size_bytes: 2097152 + matmul_dim: 8192 diff --git a/examples/motif_3/miniapp_parallel.py b/examples/motif_3/miniapp_parallel.py new file mode 100644 index 0000000..cc7a6ca --- /dev/null +++ b/examples/motif_3/miniapp_parallel.py @@ -0,0 +1,90 @@ +import argparse, asyncio, yaml, random, logging +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 + logger = logging.getLogger(__name__) + init_default_logger(logging.DEBUG) + + mp.set_start_method("dragon") + backend = await DragonExecutionBackendV3() + flow = await WorkflowEngine.create(backend=backend) + + @flow.function_task + async def simulate(_train=None): + e = cfg["simulate"] + steps = e["steps"] + device = e["device"] + read_size = e["read_size_bytes"] + write_size = e["write_size_bytes"] + matmul_dim = e["matmul_dim"] + + 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="./") + 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(_sim, _train=None): + e = cfg["training"] + steps = e["steps"] + 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"] + + 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): + 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() + + simulate_t0= simulate() + + train_t1 = train(simulate_t0) + 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)) + print(f"DONE in {t.stop():.3f}s") diff --git a/examples/motif_3/miniapp_simple.py b/examples/motif_3/miniapp_simple.py new file mode 100644 index 0000000..885f994 --- /dev/null +++ b/examples/motif_3/miniapp_simple.py @@ -0,0 +1,84 @@ +import argparse, asyncio, yaml, random, logging +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 + logger = logging.getLogger(__name__) + init_default_logger(logging.DEBUG) + + mp.set_start_method("dragon") + backend = await DragonExecutionBackendV3() + flow = await WorkflowEngine.create(backend=backend) + + @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"] + matmul_dim = e["matmul_dim"] + + 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="./") + 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(_sim=None): + e = cfg["training"] + steps = e["steps"] + 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"] + + 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): + 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() + + 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)) + print(f"DONE in {t.stop():.3f}s") 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 diff --git a/examples/motif_4/config.yaml b/examples/motif_4/config.yaml new file mode 100644 index 0000000..68c3e61 --- /dev/null +++ b/examples/motif_4/config.yaml @@ -0,0 +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: 16 + read_size_bytes: 1048576 # 1 MiB pre-step read + write_size_bytes: 4294967296 # 4 GiB post-step write + device: "cpu" # "cpu" or "gpu" + matmul_dim: 2048 # dimension for square matrix multiplication + +training: + 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: + 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 diff --git a/examples/motif_4/miniapp_async.py b/examples/motif_4/miniapp_async.py new file mode 100644 index 0000000..5967016 --- /dev/null +++ b/examples/motif_4/miniapp_async.py @@ -0,0 +1,131 @@ +import argparse, asyncio, yaml, random, logging +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 + logger = logging.getLogger(__name__) + init_default_logger(logging.DEBUG) + + mp.set_start_method("dragon") + backend = await DragonExecutionBackendV3() + flow = await WorkflowEngine.create(backend=backend) + + @flow.function_task + 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"] + + logger.info(f"Experiment Data...") + logger.info(time.strftime("%H:%M:%S", time.localtime())) + kern.generateRandomNumber(device=device, size=matmul_dim) + 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="./") + logger.info(f"Finished simulating...") + logger.info(time.strftime("%H:%M:%S", time.localtime())) + return random.random() + + @flow.function_task + async def simulation(i, _experiment): + 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"] + + 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="./") + logger.info(f"Finished simulating...") + logger.info(time.strftime("%H:%M:%S", time.localtime())) + return random.random() + + @flow.function_task + async def training(): + 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"] + matmul_dim = e["matmul_dim"] + + 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) + 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): + 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"] + + 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): + 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 inferencing model with evaluation results...") + logger.info(time.strftime("%H:%M:%S", time.localtime())) + return random.random() + + for j in range(3): + sim_t = [] + experiment_t = experiment() + for i in range(32): + sim_t.append(simulation(i, experiment_t)) + await asyncio.gather(*sim_t) + train_t = training() + infer_t = await inference(train_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)) + print(f"DONE in {t.stop():.3f}s") 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/benchmark/axpy_benchmark.cu b/wfMiniAPI/benchmark/axpy_benchmark.cu new file mode 100644 index 0000000..57c640a --- /dev/null +++ b/wfMiniAPI/benchmark/axpy_benchmark.cu @@ -0,0 +1,80 @@ +#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 * 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)); + + cublasHandle_t handle; + CHECK_CUBLAS(cublasCreate(&handle)); + + cudaEvent_t start, stop; + CHECK_CUDA(cudaEventCreate(&start)); + CHECK_CUDA(cudaEventCreate(&stop)); + + for (int i = 0; i < n_warmup; ++i) { + CHECK_CUBLAS(cublasSaxpy(handle, N, &alpha, d_x, 1, d_y, 1)); + } + CHECK_CUDA(cudaDeviceSynchronize()); + + 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; + printf("AXPY (N=%d) average over %d runs: %f ms\n", + N, n_repeat, avg_ms); + + cublasDestroy(handle); + cudaFree(d_x); + cudaFree(d_y); + free(h_x); + free(h_y); + + return 0; +} + 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/benchmark/fft_benchmark.cu b/wfMiniAPI/benchmark/fft_benchmark.cu new file mode 100644 index 0000000..2dcd79d --- /dev/null +++ b/wfMiniAPI/benchmark/fft_benchmark.cu @@ -0,0 +1,89 @@ +#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; + const int batch = 1024; + const int n_warmup = 3; + const int n_repeat = 50; + + 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; + } + + 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(cufftPlan1d(&plan, N, CUFFT_C2C, batch)); + + 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 1D FFT (N=%d, batch=%d)", N, batch); + 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/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/benchmark/matmul_benchmark.cu b/wfMiniAPI/benchmark/matmul_benchmark.cu new file mode 100644 index 0000000..65477ef --- /dev/null +++ b/wfMiniAPI/benchmark/matmul_benchmark.cu @@ -0,0 +1,109 @@ +#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 = 4096; + 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; + + 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; + } + + 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)); + + 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)); + + cublasHandle_t handle; + CHECK_CUBLAS(cublasCreate(&handle)); + + cudaEvent_t start, stop; + CHECK_CUDA(cudaEventCreate(&start)); + CHECK_CUDA(cudaEventCreate(&stop)); + + 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()); + + 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; + + printf("cuBLAS SGEMM (N=%d) average over %d runs: %f ms\n", + N, n_repeat, avg_ms); + + 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; +} diff --git a/wfMiniAPI/requirements-gpu.txt b/wfMiniAPI/requirements-gpu.txt new file mode 100644 index 0000000..5479d79 --- /dev/null +++ b/wfMiniAPI/requirements-gpu.txt @@ -0,0 +1,8 @@ +numpy +h5py +mpi4py +cupy +radical-pilot +radical-asyncflow +dragonhpc +rhapsody-py[all] diff --git a/wfMiniAPI/requirements.txt b/wfMiniAPI/requirements.txt new file mode 100644 index 0000000..4fa0150 --- /dev/null +++ b/wfMiniAPI/requirements.txt @@ -0,0 +1,7 @@ +numpy +h5py +mpi4py +radical-pilot +radical-asyncflow +dragonhpc +rhapsody-py[all] 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/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 diff --git a/wfMiniAPI/python/wfMiniAPI/kernel.py b/wfMiniAPI/src/wfMiniAPI/kernel.py similarity index 80% rename from wfMiniAPI/python/wfMiniAPI/kernel.py rename to wfMiniAPI/src/wfMiniAPI/kernel.py index 1d25be0..bb9d0d1 100644 --- a/wfMiniAPI/python/wfMiniAPI/kernel.py +++ b/wfMiniAPI/src/wfMiniAPI/kernel.py @@ -3,6 +3,9 @@ import os import sys +from .registry import annotate_kernel, list_kernels, kernel_params, run_kernel, time_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,14 +271,67 @@ 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, axes=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 -def implaceCompute(device, size, num_op, op): +#_axpy_fuse = cp.ElementwiseKernel( +# 'float32 alpha, raw float32 x', +# 'raw float32 y', +# 'y[i] += alpha * x[i]', +# 'axpy_fuse_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_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 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): @@ -277,10 +346,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) @@ -289,18 +360,12 @@ 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) +@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] diff --git a/wfMiniAPI/src/wfMiniAPI/registry.py b/wfMiniAPI/src/wfMiniAPI/registry.py new file mode 100644 index 0000000..e1fdd09 --- /dev/null +++ b/wfMiniAPI/src/wfMiniAPI/registry.py @@ -0,0 +1,90 @@ +import inspect +from collections import OrderedDict +import time +import numpy as np + +_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 + +def time_kernel(name, device, n_warmup=3, n_repeat=20, **kwargs): + # CPU timing + if device.lower() == "cpu": + for _ in range(n_warmup): + run_kernel(name, device=device, **kwargs) + run_times = [] + for _ in range(n_repeat): + t0 = time.time() + run_kernel(name, device=device, **kwargs) + t1 = time.time() + run_times.append((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} \u00B1 {std_ms:.4f} ms") + return total_ms, avg_ms, std_ms + + # GPU timing + elif device.lower() == "gpu": + import cupy as cp + for _ in range(n_warmup): + 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() + end = cp.cuda.Event() + start.record() + run_kernel(name, device=device, **kwargs) + 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} \u00B1 {std_ms:.4f} ms") + return total_ms, avg_ms, std_ms diff --git a/wfMiniAPI/src/wfMiniAPI/sim.py b/wfMiniAPI/src/wfMiniAPI/sim.py new file mode 100644 index 0000000..d7db36e --- /dev/null +++ b/wfMiniAPI/src/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/src/wfMiniAPI/timing.py b/wfMiniAPI/src/wfMiniAPI/timing.py new file mode 100644 index 0000000..6dd41c1 --- /dev/null +++ b/wfMiniAPI/src/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")