From 5eb941cdde685dccee9a8f755200c8374ded35e4 Mon Sep 17 00:00:00 2001 From: Nader Rahhal <107228500+Nader-Rahhal@users.noreply.github.com> Date: Mon, 18 May 2026 12:32:11 -0500 Subject: [PATCH 1/4] partitioning wrapper and setup --- lib/legate_jl_wrapper/src/module.cpp | 5 ++++- src/api/data.jl | 21 +++++++++++++++++++++ src/api/tasks.jl | 5 +++-- src/api/types.jl | 10 ++++++++++ 4 files changed, 38 insertions(+), 3 deletions(-) diff --git a/lib/legate_jl_wrapper/src/module.cpp b/lib/legate_jl_wrapper/src/module.cpp index 2e216c5f..ae1c8612 100644 --- a/lib/legate_jl_wrapper/src/module.cpp +++ b/lib/legate_jl_wrapper/src/module.cpp @@ -157,7 +157,10 @@ JLCXX_MODULE define_julia_module(jlcxx::Module& mod) { .method("promote", &LogicalStore::promote) .method("slice", &LogicalStore::slice) .method("get_physical_store", &LogicalStore::get_physical_store) - .method("equal_storage", &LogicalStore::equal_storage); + .method("equal_storage", &LogicalStore::equal_storage) + .method("partition_by_tiling_", &LogicalStore::partition_by_tiling_); + + mod.add_type("LogicalStorePartitionImpl"); mod.add_type("PhysicalArray") .method("dim", &PhysicalArray::dim) diff --git a/src/api/data.jl b/src/api/data.jl index 092b002c..12ae92c5 100644 --- a/src/api/data.jl +++ b/src/api/data.jl @@ -284,3 +284,24 @@ function get_ptr(arr::PhysicalStore) # PhysicalStore -> Ptr return _get_ptr(CxxWrap.CxxPtr(arr)) # cxxwrap call end + +""" + partition_by_tiling(store::LogicalStore, tile_shape) -> LogicalStorePartition + partition_by_tiling(store::LogicalStore, tile_shape, color_shape) -> LogicalStorePartition + +Create a tiled partition of the store. + +# Arguments +- `store`: The logical store to partition. +- `tile_shape`: Shape of each tile. +- `color_shape`: Optional color shape to force the tiling into. +""" +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 diff --git a/src/api/tasks.jl b/src/api/tasks.jl index 12425642..e7537539 100644 --- a/src/api/tasks.jl +++ b/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) diff --git a/src/api/types.jl b/src/api/types.jl index 2bdc65e5..a480d374 100644 --- a/src/api/types.jl +++ b/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::LogicalStorePartitionImpl +end \ No newline at end of file From 4e7e91253947a5adde8aa7bf8cd0e036dad36793 Mon Sep 17 00:00:00 2001 From: Nader Rahhal <107228500+Nader-Rahhal@users.noreply.github.com> Date: Mon, 18 May 2026 12:38:42 -0500 Subject: [PATCH 2/4] fix typo --- lib/legate_jl_wrapper/src/module.cpp | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/lib/legate_jl_wrapper/src/module.cpp b/lib/legate_jl_wrapper/src/module.cpp index ae1c8612..55da6b0b 100644 --- a/lib/legate_jl_wrapper/src/module.cpp +++ b/lib/legate_jl_wrapper/src/module.cpp @@ -158,7 +158,7 @@ JLCXX_MODULE define_julia_module(jlcxx::Module& mod) { .method("slice", &LogicalStore::slice) .method("get_physical_store", &LogicalStore::get_physical_store) .method("equal_storage", &LogicalStore::equal_storage) - .method("partition_by_tiling_", &LogicalStore::partition_by_tiling_); + .method("partition_by_tiling", &LogicalStore::partition_by_tiling); mod.add_type("LogicalStorePartitionImpl"); From 22a9fb8a388415101f98bcbde51c358157225514 Mon Sep 17 00:00:00 2001 From: Nader Rahhal <107228500+Nader-Rahhal@users.noreply.github.com> Date: Mon, 18 May 2026 14:30:07 -0500 Subject: [PATCH 3/4] get all working --- lib/legate_jl_wrapper/include/wrapper.inl | 12 +++++++ lib/legate_jl_wrapper/src/module.cpp | 38 +++++++++++++++++------ src/api/data.jl | 13 +------- src/api/types.jl | 2 +- 4 files changed, 42 insertions(+), 23 deletions(-) diff --git a/lib/legate_jl_wrapper/include/wrapper.inl b/lib/legate_jl_wrapper/include/wrapper.inl index f1c9610d..ebb30b14 100644 --- a/lib/legate_jl_wrapper/include/wrapper.inl +++ b/lib/legate_jl_wrapper/include/wrapper.inl @@ -310,6 +310,17 @@ inline void* get_ptr(legate::PhysicalStore* store) { return legate::double_dispatch(dim, code, GetPtrFunctor{}, store); } +inline std::shared_ptr partition_by_tiling(LogicalStore& store, std::vector tile_shape) { + return std::make_shared( + store.partition_by_tiling(tile_shape)); +} + +inline std::shared_ptr partition_by_tiling(LogicalStore& store, std::vector tile_shape, + std::vector color_shape) { + return std::make_shared( + store.partition_by_tiling(tile_shape, color_shape)); +} + } // namespace data namespace time { @@ -330,4 +341,5 @@ inline uint64_t time_nanoseconds() { return legate::timing::measure_nanoseconds().value(); } } // namespace time + } // namespace legate_wrapper diff --git a/lib/legate_jl_wrapper/src/module.cpp b/lib/legate_jl_wrapper/src/module.cpp index 55da6b0b..29687067 100644 --- a/lib/legate_jl_wrapper/src/module.cpp +++ b/lib/legate_jl_wrapper/src/module.cpp @@ -150,18 +150,35 @@ JLCXX_MODULE define_julia_module(jlcxx::Module& mod) { .method("get_obj_ptr", [](PhysicalStore& s) { return static_cast(&s); }); - mod.add_type("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) - .method("partition_by_tiling", &LogicalStore::partition_by_tiling); - + mod.add_type("LogicalStoreImpl"); mod.add_type("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 tile_shape) { + return legate_wrapper::data::partition_by_tiling(store, tile_shape); + }); + + mod.method("partition_by_tiling", + [](LogicalStore& store, std::vector tile_shape, + std::vector color_shape) { + return legate_wrapper::data::partition_by_tiling(store, tile_shape, color_shape); + }); + mod.method("color_shape", [](std::shared_ptr p) { + auto s = p->color_shape(); + std::vector result(s.begin(), s.end()); + return result; + }); + mod.method("store", [](std::shared_ptr p) { + return p->store(); + }); + mod.add_type("PhysicalArray") .method("dim", &PhysicalArray::dim) .method("type", &PhysicalArray::type) @@ -218,6 +235,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); diff --git a/src/api/data.jl b/src/api/data.jl index 12ae92c5..7a2383b1 100644 --- a/src/api/data.jl +++ b/src/api/data.jl @@ -285,17 +285,6 @@ function get_ptr(arr::PhysicalStore) return _get_ptr(CxxWrap.CxxPtr(arr)) # cxxwrap call end -""" - partition_by_tiling(store::LogicalStore, tile_shape) -> LogicalStorePartition - partition_by_tiling(store::LogicalStore, tile_shape, color_shape) -> LogicalStorePartition - -Create a tiled partition of the store. - -# Arguments -- `store`: The logical store to partition. -- `tile_shape`: Shape of each tile. -- `color_shape`: Optional color shape to force the tiling into. -""" 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) @@ -304,4 +293,4 @@ 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 +end \ No newline at end of file diff --git a/src/api/types.jl b/src/api/types.jl index a480d374..b5f70493 100644 --- a/src/api/types.jl +++ b/src/api/types.jl @@ -131,5 +131,5 @@ Represents a tiled partition of a `LogicalStore`. Created via `partition_by_tili Wraps the underlying C++ `LogicalStorePartitionImpl`. """ struct LogicalStorePartition{T,N} - handle::LogicalStorePartitionImpl + handle::CxxWrap.StdLib.SharedPtr{LogicalStorePartitionImpl} end \ No newline at end of file From a842808568c61b6ab00fdde527b2dfa534adda21 Mon Sep 17 00:00:00 2001 From: Nader Rahhal <107228500+Nader-Rahhal@users.noreply.github.com> Date: Wed, 20 May 2026 01:49:58 -0500 Subject: [PATCH 4/4] add for solve --- lib/legate_jl_wrapper/include/wrapper.inl | 13 ++++++++----- lib/legate_jl_wrapper/src/module.cpp | 20 ++++++++++++-------- 2 files changed, 20 insertions(+), 13 deletions(-) diff --git a/lib/legate_jl_wrapper/include/wrapper.inl b/lib/legate_jl_wrapper/include/wrapper.inl index ebb30b14..9570310c 100644 --- a/lib/legate_jl_wrapper/include/wrapper.inl +++ b/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 { @@ -311,14 +316,12 @@ inline void* get_ptr(legate::PhysicalStore* store) { } inline std::shared_ptr partition_by_tiling(LogicalStore& store, std::vector tile_shape) { - return std::make_shared( - store.partition_by_tiling(tile_shape)); + return std::make_shared(store.partition_by_tiling(tile_shape)); } -inline std::shared_ptr partition_by_tiling(LogicalStore& store, std::vector tile_shape, +inline std::shared_ptr partition_by_tiling(LogicalStore& store, std::vector tile_shape, std::vector color_shape) { - return std::make_shared( - store.partition_by_tiling(tile_shape, color_shape)); + return std::make_shared(store.partition_by_tiling(tile_shape, color_shape)); } } // namespace data diff --git a/lib/legate_jl_wrapper/src/module.cpp b/lib/legate_jl_wrapper/src/module.cpp index 29687067..8c030a04 100644 --- a/lib/legate_jl_wrapper/src/module.cpp +++ b/lib/legate_jl_wrapper/src/module.cpp @@ -208,14 +208,16 @@ JLCXX_MODULE define_julia_module(jlcxx::Module& mod) { [](AutoTask& t) { return static_cast(&t); }); mod.add_type("ManualTask") - .method("add_input", static_cast( - &ManualTask::add_input)) - .method("add_output", static_cast( - &ManualTask::add_output)) - .method("add_scalar", static_cast( - &ManualTask::add_scalar_arg)) - .method("get_obj_ptr", - [](ManualTask& t) { return static_cast(&t); }); + .method("add_input", static_cast(&ManualTask::add_input)) + .method("add_output", static_cast(&ManualTask::add_output)) + .method("add_input", [](ManualTask& t, std::shared_ptr p) { + t.add_input(*p); + }) + .method("add_output", [](ManualTask& t, std::shared_ptr p) { + t.add_output(*p); + }) + .method("add_scalar", static_cast(&ManualTask::add_scalar_arg)) + .method("get_obj_ptr", [](ManualTask& t) { return static_cast(&t); }); /* runtime */ mod.add_type("Runtime").method( @@ -255,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); }