Motif

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()

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()
+

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()

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).

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