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/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

examples/motif_3/miniapp_parallel.py

@@ -0,0 +1,86 @@
+import argparse, asyncio, yaml, random, logging
+import time
+
+from radical.asyncflow import WorkflowEngine
+from radical.asyncflow import DragonExecutionBackend
+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):
+    logger = logging.getLogger(__name__)
+    init_default_logger(logging.DEBUG)
+    backend = await DragonExecutionBackend()
+    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")