Partitioning for lin alg ops

lib/legate_jl_wrapper/include/wrapper.inl

@@ -63,6 +63,11 @@ inline bool has_started() { return legate::has_started(); }
  * @brief Check whether the Legate runtime has finished.
  */
 inline bool has_finished() { return legate::has_finished(); }
+
+inline int32_t num_procs() {
+  return legate::Runtime::get_runtime()->get_machine().count();
+}
+
 }  // namespace runtime
 
 namespace tasking {
@@ -310,6 +315,15 @@ inline void* get_ptr(legate::PhysicalStore* store) {
   return legate::double_dispatch(dim, code, GetPtrFunctor{}, store);
 }
 
+inline std::shared_ptr<LogicalStorePartition> partition_by_tiling(LogicalStore& store, std::vector<uint64_t> tile_shape) {
+  return std::make_shared<LogicalStorePartition>(store.partition_by_tiling(tile_shape));
+}
+
+inline std::shared_ptr<LogicalStorePartition> partition_by_tiling(LogicalStore& store, std::vector<uint64_t> tile_shape,
+  std::vector<uint64_t> color_shape) {
+  return std::make_shared<LogicalStorePartition>(store.partition_by_tiling(tile_shape, color_shape));
+}
+
 }  // namespace data
 
 namespace time {
@@ -330,4 +344,5 @@ inline uint64_t time_nanoseconds() {
   return legate::timing::measure_nanoseconds().value();
 }
 }  // namespace time
+
 }  // namespace legate_wrapper

lib/legate_jl_wrapper/src/module.cpp

@@ -150,14 +150,34 @@ JLCXX_MODULE define_julia_module(jlcxx::Module& mod) {
       .method("get_obj_ptr",
               [](PhysicalStore& s) { return static_cast<void*>(&s); });
 
-  mod.add_type<LogicalStore>("LogicalStoreImpl")
-      .method("dim", &LogicalStore::dim)
-      .method("type", &LogicalStore::type)
-      .method("reinterpret_as", &LogicalStore::reinterpret_as)
-      .method("promote", &LogicalStore::promote)
-      .method("slice", &LogicalStore::slice)
-      .method("get_physical_store", &LogicalStore::get_physical_store)
-      .method("equal_storage", &LogicalStore::equal_storage);
+  mod.add_type<LogicalStore>("LogicalStoreImpl");
+  mod.add_type<LogicalStorePartition>("LogicalStorePartitionImpl");
+
+  mod.method("dim", [](LogicalStore& s) { return s.dim(); });
+  mod.method("type", [](LogicalStore& s) { return s.type(); });
+  mod.method("reinterpret_as", [](LogicalStore& s, legate::Type t) { return s.reinterpret_as(t); });
+  mod.method("promote", [](LogicalStore& s, int32_t extra_dim, size_t dim_size) { return s.promote(extra_dim, dim_size); });
+  mod.method("slice", [](LogicalStore& s, int32_t dim, legate::Slice sl) { return s.slice(dim, sl); });
+  mod.method("get_physical_store", [](LogicalStore& s) { return s.get_physical_store(); });
+  mod.method("equal_storage", [](LogicalStore& s, LogicalStore& other) { return s.equal_storage(other); });
+  mod.method("partition_by_tiling",
+    [](LogicalStore& store, std::vector<uint64_t> tile_shape) {
+        return legate_wrapper::data::partition_by_tiling(store, tile_shape);
+  });
+
+  mod.method("partition_by_tiling",
+    [](LogicalStore& store, std::vector<uint64_t> tile_shape,
+                            std::vector<uint64_t> color_shape) {
+        return legate_wrapper::data::partition_by_tiling(store, tile_shape, color_shape);
+  });
+  mod.method("color_shape", [](std::shared_ptr<LogicalStorePartition> p) {
+    auto s = p->color_shape();
+    std::vector<uint64_t> result(s.begin(), s.end());
+    return result;
+  });
+  mod.method("store", [](std::shared_ptr<LogicalStorePartition> p) {
+    return p->store();
+  });
 
   mod.add_type<PhysicalArray>("PhysicalArray")
       .method("dim", &PhysicalArray::dim)
@@ -188,14 +208,16 @@ JLCXX_MODULE define_julia_module(jlcxx::Module& mod) {
               [](AutoTask& t) { return static_cast<void*>(&t); });
 
   mod.add_type<ManualTask>("ManualTask")
-      .method("add_input", static_cast<void (ManualTask::*)(LogicalStore)>(
-                               &ManualTask::add_input))
-      .method("add_output", static_cast<void (ManualTask::*)(LogicalStore)>(
-                                &ManualTask::add_output))
-      .method("add_scalar", static_cast<void (ManualTask::*)(const Scalar&)>(
-                                &ManualTask::add_scalar_arg))
-      .method("get_obj_ptr",
-              [](ManualTask& t) { return static_cast<void*>(&t); });
+    .method("add_input", static_cast<void (ManualTask::*)(LogicalStore)>(&ManualTask::add_input))
+    .method("add_output", static_cast<void (ManualTask::*)(LogicalStore)>(&ManualTask::add_output))
+    .method("add_input", [](ManualTask& t, std::shared_ptr<LogicalStorePartition> p) {
+        t.add_input(*p);
+    })
+    .method("add_output", [](ManualTask& t, std::shared_ptr<LogicalStorePartition> p) {
+        t.add_output(*p);
+    })
+    .method("add_scalar", static_cast<void (ManualTask::*)(const Scalar&)>(&ManualTask::add_scalar_arg))
+    .method("get_obj_ptr", [](ManualTask& t) { return static_cast<void*>(&t); });
 
   /* runtime */
   mod.add_type<Runtime>("Runtime").method(
@@ -215,6 +237,7 @@ JLCXX_MODULE define_julia_module(jlcxx::Module& mod) {
   mod.method("submit_auto_task", &legate_wrapper::tasking::submit_auto_task);
   mod.method("submit_manual_task",
              &legate_wrapper::tasking::submit_manual_task);
+
   /* array management */
   mod.method("create_unbound_array",
              &legate_wrapper::data::create_unbound_array);
@@ -234,5 +257,7 @@ JLCXX_MODULE define_julia_module(jlcxx::Module& mod) {
   mod.method("time_microseconds", &legate_wrapper::time::time_microseconds);
   mod.method("time_nanoseconds", &legate_wrapper::time::time_nanoseconds);
 
+  mod.method("num_procs", &legate_wrapper::runtime::num_procs);
+
   wrap_ufi(mod);
 }

src/api/data.jl

@@ -284,3 +284,13 @@ function get_ptr(arr::PhysicalStore)
     # PhysicalStore -> Ptr
     return _get_ptr(CxxWrap.CxxPtr(arr)) # cxxwrap call
 end
+
+function partition_by_tiling(store::LogicalStore{T,N}, tile_shape) where {T,N}
+    impl = partition_by_tiling(store.handle, to_cxx_vector(tile_shape)) # cxxwrap call
+    return LogicalStorePartition{T,N}(impl)
+end
+
+function partition_by_tiling(store::LogicalStore{T,N}, tile_shape, color_shape) where {T,N}
+    impl = partition_by_tiling(store.handle, to_cxx_vector(tile_shape), to_cxx_vector(color_shape)) # cxxwrap call
+    return LogicalStorePartition{T,N}(impl)
+end

src/api/tasks.jl

@@ -85,7 +85,7 @@ Add a logical array/store as an input to the task.
 """
 function add_input(
     task::Union{AutoTask,ManualTask},
-    item::Union{LogicalArray,LogicalStore},
+    item::Union{LogicalArray,LogicalStore,LogicalStorePartition},
 )
     add_input(task, item.handle)
 end
@@ -98,11 +98,12 @@ Add a logical array/store as an output of the task.
 """
 function add_output(
     task::Union{AutoTask,ManualTask},
-    item::Union{LogicalArray,LogicalStore},
+    item::Union{LogicalArray,LogicalStore,LogicalStorePartition},
 )
     add_output(task, item.handle)
 end
 
+
 """
     add_scalar(AutoTask, scalar::Scalar)
     add_scalar(ManualTask, scalar::Scalar)

src/api/types.jl

@@ -123,3 +123,13 @@ Base.size(a::LogicalArray, i::Integer) = size(a)[i]
 Datatype of object within Legate. See `Legate.supported_types()` to see supported types.
 """
 LegateType
+
+
+"""
+    LogicalStorePartition{T,N}
+Represents a tiled partition of a `LogicalStore`. Created via `partition_by_tiling`.
+Wraps the underlying C++ `LogicalStorePartitionImpl`.
+"""
+struct LogicalStorePartition{T,N}
+    handle::CxxWrap.StdLib.SharedPtr{LogicalStorePartitionImpl}
+end