diff --git a/lib/legate_jl_wrapper/include/wrapper.inl b/lib/legate_jl_wrapper/include/wrapper.inl index f1c9610d..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 { @@ -310,6 +315,15 @@ 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 +344,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 2e216c5f..8c030a04 100644 --- a/lib/legate_jl_wrapper/src/module.cpp +++ b/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(&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); + 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) @@ -188,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( @@ -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); } diff --git a/src/api/data.jl b/src/api/data.jl index 092b002c..7a2383b1 100644 --- a/src/api/data.jl +++ b/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 \ No newline at end of file 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..b5f70493 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::CxxWrap.StdLib.SharedPtr{LogicalStorePartitionImpl} +end \ No newline at end of file