add convolve3D and complex trace attributes to scalability experiment

Author: SerodioJ

run_scalable_experiment_on_scope.py

@@ -98,7 +98,9 @@ def run_experiments(args):
 
     for attr in attr_classes[shape_map[scope["data_size"]]]:
         if run_specs["CPU"]:
-            for spec in ["TVM", "JAX", "TORCH_C", "TORCH_N"]:
+            for spec in ["Baseline", "TVM", "JAX", "TORCH_C", "TORCH_N"]:
+                if spec == "Baseline" and "glcm" in attr:
+                    continue
                 if run_specs[spec]:
                     output_path = os.path.join(args.output_path, f"{attr}+{spec}+cpu+{exp_id}.zarr")
                     rm_aux_path = Path(output_path)

src/run_scalable.py

@@ -10,7 +10,7 @@
 
 
 
-from scalable_integration.utils import get_glcm_chunksize_overlap
+from scalable_integration.utils import get_glcm_chunksize_overlap, get_complex_trace_chunksize_overlap, get_conv_chunksize_overlap
 from scalable_integration.custom_worker import DaskOperatorWorker
 
 
@@ -40,6 +40,10 @@ def run_exp(attr, rt, device, input_path, output_path, base_path, dtype, shape):
 
     if "glcm" in attr:
         chunksize, overlap = get_glcm_chunksize_overlap(rt=rt, exp_shape=shape)
+    elif "conv" in attr:
+        chunksize, overlap = get_conv_chunksize_overlap(rt=rt, exp_shape=shape)
+    else:
+        chunksize, overlap = get_complex_trace_chunksize_overlap(rt=rt, exp_shape=shape)
 
 
     input_data = zarr.open(input_path)
@@ -51,16 +55,22 @@ def run_exp(attr, rt, device, input_path, output_path, base_path, dtype, shape):
     )
 
     if "glcm" in attr:
-        from scalable_integration.texture import glcm_base
-        tasks = glcm_base(
-            rt=rt,
-            input_data=input_data,
-            output_data=output_data,
-            chunksize=chunksize,
-            overlap=overlap,
-            dtype=dtype,
-            device=device,
-        )
+        from scalable_integration.texture import glcm_base as get_tasks
+    elif "conv" in attr:
+        from scalable_integration.signal import conv_base as get_tasks
+    else:
+        from scalable_integration.complex_trace import complex_trace_base as get_tasks
+    tasks = get_tasks(
+        rt=rt,
+        input_data=input_data,
+        output_data=output_data,
+        chunksize=chunksize,
+        overlap=overlap,
+        dtype=dtype,
+        device=device,
+    )
+        
+
 
     client.compute(tasks, sync=True)
     end = perf_counter()
@@ -93,6 +103,7 @@ def run_exp(attr, rt, device, input_path, output_path, base_path, dtype, shape):
             "relative-amplitude-change",
             "amplitude-acceleration",
             "inst-bandwidth",
+            "convolve3d",
             "glcm-asm",
             "glcm-contrast",
             "glcm-variance",

src/scalable_integration/complex_trace.py

@@ -0,0 +1,283 @@
+import dask
+import numpy as np
+import cupy as cp
+import tvm
+import jax
+import torch
+from scalable_integration.utils import get_chunks
+from scalable_integration.custom_worker import get_operator
+jax.config.update("jax_enable_x64", True)
+
+
+def complex_trace_base(
+    rt,
+    input_data,
+    output_data,
+    chunksize,
+    overlap,
+    dtype,
+    device,
+):
+
+    in_ind, out_ind, padding = get_chunks(
+        data_shape=input_data.shape,
+        chunksize=chunksize,
+        overlap=overlap
+    )
+
+    task_funcs = {
+        "tvm": complex_trace_tvm,
+        "baseline": complex_trace_baseline,
+        "torch_c": complex_trace_torch_c,
+        "torch_n": complex_trace_torch_n,
+        "jax": complex_trace_jax,
+    }
+
+    task_func = task_funcs[rt]
+
+    tasks = [
+        task_func(
+            input_data=input_data,
+            output_data=output_data,
+            indx=i,
+            out_indx=out_i,
+            chunksize=chunksize,
+            pad_width=p,
+            overlap=overlap,
+            dtype=dtype,
+            device=device,
+        )
+        for i, out_i, p in zip(in_ind, out_ind, padding)
+    ]
+
+    return tasks, chunksize
+
+    
+
+@dask.delayed
+def complex_trace_tvm(
+    input_data,
+    output_data,
+    indx,
+    out_indx,
+    chunksize,
+    pad_width,
+    overlap,
+    dtype,
+    device
+):
+    sli = tuple(
+        slice(i, i + c + o[0] + o[1] - p[0] - p[1])
+        for i, c, p, o in zip(indx, chunksize, pad_width, overlap)
+    )
+    chunk = input_data[sli].astype(dtype)
+    chunk = np.pad(
+        chunk, pad_width=pad_width, mode="constant", constant_values=0
+    )
+
+    operator = get_operator()
+
+    data_tvm = tvm.nd.array(chunk, device=operator._dev)
+    res = tvm.nd.empty(data_tvm.shape, dtype=data_tvm.dtype, device=operator._dev)
+    operator.transform(data_tvm, res)
+
+    res = res.numpy()
+
+
+    useful_slice = [
+        min(c, o - i)
+        for i, c, o in zip(out_indx, chunksize, output_data.shape)
+    ]
+    out_sli = tuple(
+        slice(i, i + u)
+        for i, u in zip(out_indx, useful_slice)
+    )
+
+    res_sli = tuple(
+        slice(o[0], o[0] + u)
+        for u, o in zip(useful_slice, overlap)
+    )
+
+    output_data[out_sli] = res[res_sli]
+
+@dask.delayed
+def complex_trace_baseline(
+    input_data,
+    output_data,
+    indx,
+    out_indx,
+    chunksize,
+    pad_width,
+    overlap,
+    dtype,
+    device
+):
+    sli = tuple(
+        slice(i, i + c + o[0] + o[1] - p[0] - p[1])
+        for i, c, p, o in zip(indx, chunksize, pad_width, overlap)
+    )
+    chunk = input_data[sli].astype(dtype)
+    chunk = np.pad(
+        chunk, pad_width=pad_width, mode="constant", constant_values=0
+    )
+    operator = get_operator()
+
+    if device == "cpu":
+        res = operator._transform_cpu(chunk)
+    else:
+        chunk = cp.asarray(chunk)
+        res = operator._transform_gpu(chunk).get()
+
+
+    useful_slice = [
+        min(c, o - i)
+        for i, c, o in zip(out_indx, chunksize, output_data.shape)
+    ]
+    out_sli = tuple(
+        slice(i, i + u)
+        for i, u in zip(out_indx, useful_slice)
+    )
+
+    res_sli = tuple(
+        slice(o[0], o[0] + u)
+        for u, o in zip(useful_slice, overlap)
+    )
+
+    output_data[out_sli] = res[res_sli]
+
+@dask.delayed
+def complex_trace_jax(
+    input_data,
+    output_data,
+    indx,
+    out_indx,
+    chunksize,
+    pad_width,
+    overlap,
+    dtype,
+    device
+):
+    sli = tuple(
+        slice(i, i + c + o[0] + o[1] - p[0] - p[1])
+        for i, c, p, o in zip(indx, chunksize, pad_width, overlap)
+    )
+    chunk = input_data[sli].astype(dtype)
+    chunk = np.pad(
+        chunk, pad_width=pad_width, mode="constant", constant_values=0
+    )
+    chunk = jax.device_put(chunk, device=jax.devices(device)[0])
+
+    operator = get_operator()
+    if device == "cpu":
+        res = operator._transform_cpu(chunk)
+    else:
+        res = operator._transform_gpu(chunk)
+    res = np.asarray(res)
+
+
+    useful_slice = [
+        min(c, o - i)
+        for i, c, o in zip(out_indx, chunksize, output_data.shape)
+    ]
+    out_sli = tuple(
+        slice(i, i + u)
+        for i, u in zip(out_indx, useful_slice)
+    )
+
+    res_sli = tuple(
+        slice(o[0], o[0] + u)
+        for u, o in zip(useful_slice, overlap)
+    )
+
+    output_data[out_sli] = res[res_sli]
+
+@dask.delayed
+def complex_trace_torch_c(
+    input_data,
+    output_data,
+    indx,
+    out_indx,
+    chunksize,
+    pad_width,
+    overlap,
+    dtype,
+    device
+):
+    sli = tuple(
+        slice(i, i + c + o[0] + o[1] - p[0] - p[1])
+        for i, c, p, o in zip(indx, chunksize, pad_width, overlap)
+    )
+    chunk = input_data[sli].astype(dtype)
+    chunk = np.pad(
+        chunk, pad_width=pad_width, mode="constant", constant_values=0
+    )
+    chunk = torch.from_numpy(chunk).to(torch.device("cpu" if device == "cpu" else "cuda"))
+    operator = get_operator()
+    if device == "cpu":
+        res = operator._transform_cpu(chunk)
+    else:
+        res = operator._transform_gpu(chunk).cpu()
+    res = res.numpy()
+
+
+    useful_slice = [
+        min(c, o - i)
+        for i, c, o in zip(out_indx, chunksize, output_data.shape)
+    ]
+    out_sli = tuple(
+        slice(i, i + u)
+        for i, u in zip(out_indx, useful_slice)
+    )
+
+    res_sli = tuple(
+        slice(o[0], o[0] + u)
+        for u, o in zip(useful_slice, overlap)
+    )
+
+    output_data[out_sli] = res[res_sli]
+
+
+@dask.delayed
+def complex_trace_torch_n(
+    input_data,
+    output_data,
+    indx,
+    out_indx,
+    chunksize,
+    pad_width,
+    overlap,
+    dtype,
+    device
+):
+    sli = tuple(
+        slice(i, i + c + o[0] + o[1] - p[0] - p[1])
+        for i, c, p, o in zip(indx, chunksize, pad_width, overlap)
+    )
+    chunk = input_data[sli].astype(dtype)
+    chunk = np.pad(
+        chunk, pad_width=pad_width, mode="constant", constant_values=0
+    )
+    chunk = torch.from_numpy(chunk).to(torch.device("cpu" if device == "cpu" else "cuda"))
+    operator = get_operator()
+    if device == "cpu":
+        res = operator._nocompile_cpu(chunk)
+    else:
+        res = operator._nocompile_gpu(chunk).cpu()
+    res = res.numpy()
+
+
+    useful_slice = [
+        min(c, o - i)
+        for i, c, o in zip(out_indx, chunksize, output_data.shape)
+    ]
+    out_sli = tuple(
+        slice(i, i + u)
+        for i, u in zip(out_indx, useful_slice)
+    )
+
+    res_sli = tuple(
+        slice(o[0], o[0] + u)
+        for u, o in zip(useful_slice, overlap)
+    )
+
+    output_data[out_sli] = res[res_sli]