Legate Zarr

legate/benchmarks/zarr_read.py

@@ -0,0 +1,195 @@
+# Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved.
+# See file LICENSE for terms.
+
+import argparse
+import contextlib
+import functools
+import pathlib
+import tempfile
+from time import perf_counter as clock
+from typing import ContextManager
+
+import numpy as np
+import zarr
+from kvikio.zarr import GDSStore
+from zarr.errors import ArrayNotFoundError
+
+
+def try_open_zarr_array(dirpath, shape, chunks, dtype):
+    try:
+        a = zarr.open_array(dirpath, mode="r")
+        chunks = chunks or a.chunks
+        if a.shape == shape and a.chunks == chunks and a.dtype == dtype:
+            return a
+    except ArrayNotFoundError:
+        pass
+    return None
+
+
+def create_zarr_array(dirpath, shape, chunks=None, dtype=np.float64) -> None:
+    ret = try_open_zarr_array(dirpath, shape, chunks, dtype)
+    if ret is None:
+        ret = zarr.open_array(
+            dirpath,
+            shape=shape,
+            dtype=dtype,
+            mode="w",
+            chunks=chunks,
+            compressor=None,
+        )
+        ret[:] = np.random.random(shape)
+
+    print(
+        f"Zarr '{ret.store.path}': shape: {ret.shape}, "
+        f"chunks: {ret.chunks}, dtype: {ret.dtype}"
+    )
+
+
+@contextlib.contextmanager
+def run_dask(args, *, use_cupy):
+    from dask import array as da
+    from dask_cuda import LocalCUDACluster
+    from distributed import Client
+
+    def f():
+        t0 = clock()
+        if use_cupy:
+            import cupy
+
+            az = zarr.open_array(GDSStore(args.dir / "A"), meta_array=cupy.empty(()))
+            bz = zarr.open_array(GDSStore(args.dir / "B"), meta_array=cupy.empty(()))
+        else:
+            az = args.dir / "A"
+            bz = args.dir / "B"
+
+        a = da.from_zarr(az)
+        b = da.from_zarr(bz)
+        c = args.op(da, a, b)
+
+        int(c.sum().compute())
+        t1 = clock()
+        return t1 - t0
+
+    with LocalCUDACluster(n_workers=args.n_workers) as cluster:
+        with Client(cluster):
+            yield f
+
+
+@contextlib.contextmanager
+def run_legate(args):
+    import cunumeric as num
+    from legate_kvikio.zarr import read_array
+
+    from legate.core import get_legate_runtime
+
+    def f():
+        get_legate_runtime().issue_execution_fence(block=True)
+        t0 = clock()
+        a = read_array(args.dir / "A")
+        b = read_array(args.dir / "B")
+        c = args.op(num, a, b)
+        int(c.sum())
+        t1 = clock()
+        return t1 - t0
+
+    yield f
+
+
+API = {
+    "dask-cpu": functools.partial(run_dask, use_cupy=False),
+    "dask-gpu": functools.partial(run_dask, use_cupy=True),
+    "legate": run_legate,
+}
+
+OP = {"add": lambda xp, a, b: a + b, "matmul": lambda xp, a, b: xp.matmul(a, b)}
+
+
+def main(args):
+    create_zarr_array(args.dir / "A", shape=(args.m, args.m))
+    create_zarr_array(args.dir / "B", shape=(args.m, args.m))
+
+    timings = []
+    with API[args.api](args) as f:
+        for _ in range(args.n_warmup_runs):
+            elapsed = f()
+            print("elapsed[warmup]: ", elapsed)
+        for i in range(args.nruns):
+            elapsed = f()
+            print(f"elapsed[run #{i}]: ", elapsed)
+            timings.append(elapsed)
+    print(f"elapsed mean: {np.mean(timings):.5}s (std: {np.std(timings):.5}s)")
+
+
+if __name__ == "__main__":
+
+    def parse_directory(x):
+        if x is None:
+            return x
+        else:
+            p = pathlib.Path(x)
+            if not p.is_dir():
+                raise argparse.ArgumentTypeError("Must be a directory")
+            return p
+
+    parser = argparse.ArgumentParser(description="Matrix operation on two Zarr files")
+    parser.add_argument(
+        "-m",
+        default=100,
+        type=int,
+        help="Dimension of the two squired input matrix (MxM) (default: %(default)s).",
+    )
+    parser.add_argument(
+        "-d",
+        "--dir",
+        metavar="PATH",
+        default=None,
+        type=parse_directory,
+        help="Path to the directory to r/w from (default: tempfile.TemporaryDirectory)",
+    )
+    parser.add_argument(
+        "--nruns",
+        metavar="RUNS",
+        default=1,
+        type=int,
+        help="Number of runs (default: %(default)s).",
+    )
+    parser.add_argument(
+        "--api",
+        metavar="API",
+        default=tuple(API.keys())[0],
+        choices=tuple(API.keys()),
+        help="API to use {%(choices)s}",
+    )
+    parser.add_argument(
+        "--n-workers",
+        default=1,
+        type=int,
+        help="Number of workers (default: %(default)s).",
+    )
+    parser.add_argument(
+        "--op",
+        metavar="OP",
+        default=tuple(OP.keys())[0],
+        choices=tuple(OP.keys()),
+        help="Operation to run {%(choices)s}",
+    )
+    parser.add_argument(
+        "--n-warmup-runs",
+        default=1,
+        type=int,
+        help="Number of warmup runs (default: %(default)s).",
+    )
+
+    args = parser.parse_args()
+    args.op = OP[args.op]  # Parse the operation argument
+
+    # Create a temporary directory if user didn't specify a directory
+    temp_dir: tempfile.TemporaryDirectory | ContextManager
+    if args.dir is None:
+        temp_dir = tempfile.TemporaryDirectory()
+        args.dir = pathlib.Path(temp_dir.name)
+    else:
+        temp_dir = contextlib.nullcontext()
+
+    with temp_dir:
+        main(args)

legate/cpp/legate_mapping.cpp

@@ -14,8 +14,10 @@
  * limitations under the License.
  */
 
-#include "legate_mapping.hpp"
+#include <vector>
+
 #include "core/mapping/mapping.h"
+#include "legate_mapping.hpp"
 #include "task_opcodes.hpp"
 
 namespace legate_kvikio {

legate/cpp/task_opcodes.hpp

@@ -18,5 +18,7 @@
 enum TaskOpCode {
   OP_WRITE,
   OP_READ,
+  OP_TILE_WRITE,
+  OP_TILE_READ,
   OP_NUM_TASK_IDS,  // must be last
 };

legate/cpp/tile_io.cpp

@@ -0,0 +1,180 @@
+/*
+ * Copyright (c) 2023, NVIDIA CORPORATION.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#include <array>
+#include <cassert>
+#include <cstddef>
+#include <cstdint>
+#include <filesystem>
+#include <iostream>
+#include <utility>
+
+#include "legate_mapping.hpp"
+#include "task_opcodes.hpp"
+
+#include <kvikio/file_handle.hpp>
+
+namespace {
+
+/**
+ * @brief Get the tile coordinate based on a task index
+ *
+ * @param task_index Task index
+ * @param tile_start The start tile coordinate
+ * @return Tile coordinate
+ */
+legate::DomainPoint get_tile_coord(legate::DomainPoint task_index,
+                                   legate::Span<const uint64_t>& tile_start)
+{
+  for (uint32_t i = 0; i < task_index.dim; ++i) {
+    task_index[i] += tile_start[i];
+  }
+  return task_index;
+}
+
+/**
+ * @brief Get the file path of a tile
+ *
+ * @param dirpath The path to the root directory of the Zarr file
+ * @param tile_coord The coordinate of the tile
+ * @param delimiter The delimiter
+ * @return Path to the file representing the requested tile
+ */
+std::filesystem::path get_file_path(const std::string& dirpath,
+                                    const legate::DomainPoint& tile_coord,
+                                    const std::string& delimiter = ".")
+{
+  std::stringstream ss;
+  for (int32_t idx = 0; idx < tile_coord.dim; ++idx) {
+    if (idx != 0) { ss << delimiter; }
+    ss << tile_coord[idx];
+  }
+  return std::filesystem::path(dirpath) / ss.str();
+}
+
+/**
+ * @brief Functor for tiling read or write Legate store to or from disk using KvikIO
+ *
+ * @tparam IsReadOperation Whether the operation is a read or a write operation
+ * @param context The Legate task context
+ * @param store The Legate store ti read or write
+ */
+template <bool IsReadOperation>
+struct tile_read_write_fn {
+  template <legate::LegateTypeCode CODE, int32_t DIM>
+  void operator()(legate::TaskContext& context, legate::Store& store)
+  {
+    using DTYPE                             = legate::legate_type_of<CODE>;
+    const auto task_index                   = context.get_task_index();
+    const std::string path                  = context.scalars().at(0).value<std::string>();
+    legate::Span<const uint64_t> tile_shape = context.scalars().at(1).values<uint64_t>();
+    legate::Span<const uint64_t> tile_start = context.scalars().at(2).values<uint64_t>();
+    const auto tile_coord                   = get_tile_coord(task_index, tile_start);
+    const auto filepath                     = get_file_path(path, tile_coord);
+
+    auto shape        = store.shape<DIM>();
+    auto shape_volume = shape.volume();
+    if (shape_volume == 0) { return; }
+    size_t nbytes = shape_volume * sizeof(DTYPE);
+    std::array<size_t, DIM> strides{};
+
+    // We know that the accessor is contiguous because we set `policy.exact = true`
+    // in `Mapper::store_mappings()`.
+    if constexpr (IsReadOperation) {
+      kvikio::FileHandle f(filepath, "r");
+      auto* data = store.write_accessor<DTYPE, DIM>().ptr(shape, strides.data());
+      f.pread(data, nbytes).get();
+    } else {
+      kvikio::FileHandle f(filepath, "w");
+      const auto* data = store.read_accessor<DTYPE, DIM>().ptr(shape, strides.data());
+      f.pwrite(data, nbytes).get();
+    }
+  }
+};
+
+}  // namespace
+
+namespace legate_kvikio {
+
+/**
+ * @brief Write a tiled Legate store to disk using KvikIO
+ * Task signature:
+ *   - scalars:
+ *     - path: std::string
+ *     - tile_shape: tuple of int64_t
+ *     - tile_start: tuple of int64_t
+ *   - inputs:
+ *     - buffer: store (any dtype)
+ *
+ * NB: the store must be contigues. To make Legate in force this,
+ *     set `policy.exact = true` in `Mapper::store_mappings()`.
+ *
+ */
+class TileWriteTask : public Task<TileWriteTask, TaskOpCode::OP_TILE_WRITE> {
+ public:
+  static void cpu_variant(legate::TaskContext& context)
+  {
+    legate::Store& store = context.inputs().at(0);
+    legate::double_dispatch(store.dim(), store.code(), tile_read_write_fn<false>{}, context, store);
+  }
+
+  static void gpu_variant(legate::TaskContext& context)
+  {
+    // Since KvikIO supports both GPU and CPU memory seamlessly, we reuse the CPU variant.
+    cpu_variant(context);
+  }
+};
+
+/**
+ * @brief Read a tiled Legate store to disk using KvikIO
+ * Task signature:
+ *   - scalars:
+ *     - path: std::string
+ *     - tile_shape: tuple of int64_t
+ *     - tile_start: tuple of int64_t
+ *   - outputs:
+ *     - buffer: store (any dtype)
+ *
+ * NB: the store must be contigues. To make Legate in force this,
+ *     set `policy.exact = true` in `Mapper::store_mappings()`.
+ *
+ */
+class TileReadTask : public Task<TileReadTask, TaskOpCode::OP_TILE_READ> {
+ public:
+  static void cpu_variant(legate::TaskContext& context)
+  {
+    legate::Store& store = context.outputs().at(0);
+    legate::double_dispatch(store.dim(), store.code(), tile_read_write_fn<true>{}, context, store);
+  }
+
+  static void gpu_variant(legate::TaskContext& context)
+  {
+    // Since KvikIO supports both GPU and CPU memory seamlessly, we reuse the CPU variant.
+    cpu_variant(context);
+  }
+};
+
+}  // namespace legate_kvikio
+
+namespace {
+
+void __attribute__((constructor)) register_tasks()
+{
+  legate_kvikio::TileWriteTask::register_variants();
+  legate_kvikio::TileReadTask::register_variants();
+}
+
+}  // namespace