Refactoring changes (#581)

* Move the Legate task implemention out from the header * Fix the bug in task name extraction introduced by refactoring * Factor out variant detection code from the base task class * Move task details to an inline file * Refactoring to get rid of Legion references in registration * Hide variant registration details with a friend class * Take over reduction operator registrations from Legion * Re-export necessary Legion names and use fully qualified names for all other Legion references * Use the right copyright year
nv-legate · Feb 25, 2023 · cb65d91 · cb65d91
1 parent 8fc4440
commit cb65d91
Show file tree

Hide file tree

Showing 50 changed files with 1,345 additions and 1,006 deletions.
diff --git a/legate_core_cpp.cmake b/legate_core_cpp.cmake
@@ -205,8 +205,10 @@ list(APPEND legate_core_SOURCES
   src/core/runtime/projection.cc
   src/core/runtime/runtime.cc
   src/core/runtime/shard.cc
+  src/core/task/registrar.cc
   src/core/task/return.cc
   src/core/task/task.cc
+  src/core/task/variant.cc
   src/core/utilities/debug.cc
   src/core/utilities/deserializer.cc
   src/core/utilities/machine.cc
@@ -355,13 +357,16 @@ install(
 
 install(
   FILES src/core/runtime/context.h
+        src/core/runtime/context.inl
         src/core/runtime/runtime.h
   DESTINATION ${CMAKE_INSTALL_INCLUDEDIR}/legate/core/runtime)
 
 install(
   FILES src/core/task/exception.h
         src/core/task/return.h
         src/core/task/task.h
+        src/core/task/task.inl
+        src/core/task/variant.h
   DESTINATION ${CMAKE_INSTALL_INCLUDEDIR}/legate/core/task)
 
 install(

diff --git a/src/core/comm/coll.cc b/src/core/comm/coll.cc
@@ -35,7 +35,6 @@ namespace legate {
 namespace comm {
 namespace coll {
 
-using namespace Legion;
 Logger log_coll("coll");
 
 BackendNetwork* backend_network = nullptr;

diff --git a/src/core/comm/comm_cpu.cc b/src/core/comm/comm_cpu.cc
@@ -19,8 +19,6 @@
 
 #include "core/comm/coll.h"
 
-using namespace Legion;
-
 namespace legate {
 namespace comm {
 namespace cpu {
@@ -30,7 +28,7 @@ static int init_cpucoll_mapping(const Legion::Task* task,
                                 Legion::Context context,
                                 Legion::Runtime* runtime)
 {
-  Core::show_progress(task, context, runtime, task->get_task_name());
+  Core::show_progress(task, context, runtime);
   int mpi_rank = 0;
 #if defined(LEGATE_USE_NETWORK)
   if (coll::backend_network->comm_type == coll::CollCommType::CollMPI) {
@@ -46,7 +44,7 @@ static coll::CollComm init_cpucoll(const Legion::Task* task,
                                    Legion::Context context,
                                    Legion::Runtime* runtime)
 {
-  Core::show_progress(task, context, runtime, task->get_task_name());
+  Core::show_progress(task, context, runtime);
 
   const int point = task->index_point[0];
   int num_ranks   = task->index_domain.get_volume();
@@ -80,7 +78,7 @@ static void finalize_cpucoll(const Legion::Task* task,
                              Legion::Context context,
                              Legion::Runtime* runtime)
 {
-  Core::show_progress(task, context, runtime, task->get_task_name());
+  Core::show_progress(task, context, runtime);
 
   assert(task->futures.size() == 1);
   coll::CollComm comm = task->futures[0].get_result<coll::CollComm>();
@@ -95,32 +93,29 @@ void register_tasks(Legion::Machine machine,
                     Legion::Runtime* runtime,
                     const LibraryContext& context)
 {
-  const InputArgs& command_args = Legion::Runtime::get_input_args();
-  int argc                      = command_args.argc;
-  char** argv                   = command_args.argv;
-  coll::collInit(argc, argv);
+  const auto& command_args = Legion::Runtime::get_input_args();
+  coll::collInit(command_args.argc, command_args.argv);
 
-  const TaskID init_cpucoll_mapping_task_id =
-    context.get_task_id(LEGATE_CORE_INIT_CPUCOLL_MAPPING_TASK_ID);
+  auto init_cpucoll_mapping_task_id = context.get_task_id(LEGATE_CORE_INIT_CPUCOLL_MAPPING_TASK_ID);
   const char* init_cpucoll_mapping_task_name = "core::comm::cpu::init_mapping";
   runtime->attach_name(init_cpucoll_mapping_task_id,
                        init_cpucoll_mapping_task_name,
                        false /*mutable*/,
                        true /*local only*/);
 
-  const TaskID init_cpucoll_task_id  = context.get_task_id(LEGATE_CORE_INIT_CPUCOLL_TASK_ID);
+  auto init_cpucoll_task_id          = context.get_task_id(LEGATE_CORE_INIT_CPUCOLL_TASK_ID);
   const char* init_cpucoll_task_name = "core::comm::cpu::init";
   runtime->attach_name(
     init_cpucoll_task_id, init_cpucoll_task_name, false /*mutable*/, true /*local only*/);
 
-  const TaskID finalize_cpucoll_task_id = context.get_task_id(LEGATE_CORE_FINALIZE_CPUCOLL_TASK_ID);
+  auto finalize_cpucoll_task_id = context.get_task_id(LEGATE_CORE_FINALIZE_CPUCOLL_TASK_ID);
   const char* finalize_cpucoll_task_name = "core::comm::cpu::finalize";
   runtime->attach_name(
     finalize_cpucoll_task_id, finalize_cpucoll_task_name, false /*mutable*/, true /*local only*/);
 
   auto make_registrar = [&](auto task_id, auto* task_name, auto proc_kind) {
-    TaskVariantRegistrar registrar(task_id, task_name);
-    registrar.add_constraint(ProcessorConstraint(proc_kind));
+    Legion::TaskVariantRegistrar registrar(task_id, task_name);
+    registrar.add_constraint(Legion::ProcessorConstraint(proc_kind));
     registrar.set_leaf(true);
     registrar.global_registration = false;
     return registrar;

diff --git a/src/core/comm/comm_nccl.cu b/src/core/comm/comm_nccl.cu
@@ -17,14 +17,14 @@
 #include "core/comm/comm_nccl.h"
 #include "core/cuda/cuda_help.h"
 #include "core/cuda/stream_pool.h"
+#include "core/data/buffer.h"
 #include "core/utilities/nvtx_help.h"
+#include "core/utilities/typedefs.h"
 #include "legate.h"
 
 #include <cuda.h>
 #include <nccl.h>
 
-using namespace Legion;
-
 namespace legate {
 namespace comm {
 namespace nccl {
@@ -59,7 +59,7 @@ static ncclUniqueId init_nccl_id(const Legion::Task* task,
 {
   legate::nvtx::Range auto_range("core::comm::nccl::init_id");
 
-  Core::show_progress(task, context, runtime, task->get_task_name());
+  Core::show_progress(task, context, runtime);
 
   ncclUniqueId id;
   CHECK_NCCL(ncclGetUniqueId(&id));
@@ -74,7 +74,7 @@ static ncclComm_t* init_nccl(const Legion::Task* task,
 {
   legate::nvtx::Range auto_range("core::comm::nccl::init");
 
-  Core::show_progress(task, context, runtime, task->get_task_name());
+  Core::show_progress(task, context, runtime);
 
   assert(task->futures.size() == 1);
 
@@ -92,13 +92,8 @@ static ncclComm_t* init_nccl(const Legion::Task* task,
 
   // Perform a warm-up all-to-all
 
-  using namespace Legion;
-
-  DeferredBuffer<_Payload, 1> src_buffer(Memory::GPU_FB_MEM,
-                                         Domain(Rect<1>{Point<1>{0}, Point<1>{num_ranks - 1}}));
-
-  DeferredBuffer<_Payload, 1> tgt_buffer(Memory::GPU_FB_MEM,
-                                         Domain(Rect<1>{Point<1>{0}, Point<1>{num_ranks - 1}}));
+  auto src_buffer = create_buffer<_Payload>(num_ranks, Memory::Kind::GPU_FB_MEM);
+  auto tgt_buffer = create_buffer<_Payload>(num_ranks, Memory::Kind::GPU_FB_MEM);
 
   CHECK_NCCL(ncclGroupStart());
   for (auto idx = 0; idx < num_ranks; ++idx) {
@@ -119,7 +114,7 @@ static void finalize_nccl(const Legion::Task* task,
 {
   legate::nvtx::Range auto_range("core::comm::nccl::finalize");
 
-  Core::show_progress(task, context, runtime, task->get_task_name());
+  Core::show_progress(task, context, runtime);
 
   assert(task->futures.size() == 1);
   auto comm = task->futures[0].get_result<ncclComm_t*>();
@@ -131,24 +126,24 @@ void register_tasks(Legion::Machine machine,
                     Legion::Runtime* runtime,
                     const LibraryContext& context)
 {
-  const TaskID init_nccl_id_task_id  = context.get_task_id(LEGATE_CORE_INIT_NCCL_ID_TASK_ID);
+  auto init_nccl_id_task_id          = context.get_task_id(LEGATE_CORE_INIT_NCCL_ID_TASK_ID);
   const char* init_nccl_id_task_name = "core::comm::nccl::init_id";
   runtime->attach_name(
     init_nccl_id_task_id, init_nccl_id_task_name, false /*mutable*/, true /*local only*/);
 
-  const TaskID init_nccl_task_id  = context.get_task_id(LEGATE_CORE_INIT_NCCL_TASK_ID);
+  auto init_nccl_task_id          = context.get_task_id(LEGATE_CORE_INIT_NCCL_TASK_ID);
   const char* init_nccl_task_name = "core::comm::nccl::init";
   runtime->attach_name(
     init_nccl_task_id, init_nccl_task_name, false /*mutable*/, true /*local only*/);
 
-  const TaskID finalize_nccl_task_id  = context.get_task_id(LEGATE_CORE_FINALIZE_NCCL_TASK_ID);
+  auto finalize_nccl_task_id          = context.get_task_id(LEGATE_CORE_FINALIZE_NCCL_TASK_ID);
   const char* finalize_nccl_task_name = "core::comm::nccl::finalize";
   runtime->attach_name(
     finalize_nccl_task_id, finalize_nccl_task_name, false /*mutable*/, true /*local only*/);
 
   auto make_registrar = [&](auto task_id, auto* task_name, auto proc_kind) {
-    TaskVariantRegistrar registrar(task_id, task_name);
-    registrar.add_constraint(ProcessorConstraint(proc_kind));
+    Legion::TaskVariantRegistrar registrar(task_id, task_name);
+    registrar.add_constraint(Legion::ProcessorConstraint(proc_kind));
     registrar.set_leaf(true);
     registrar.global_registration = false;
     return registrar;

diff --git a/src/core/comm/local_comm.cc b/src/core/comm/local_comm.cc
@@ -27,7 +27,6 @@ namespace legate {
 namespace comm {
 namespace coll {
 
-using namespace Legion;
 extern Logger log_coll;
 
 // public functions start from here
@@ -348,4 +347,4 @@ void LocalNetwork::barrierLocal(CollComm global_comm)
 
 }  // namespace coll
 }  // namespace comm
-}  // namespace legate
+}  // namespace legate
diff --git a/src/core/comm/mpi_comm.cc b/src/core/comm/mpi_comm.cc
@@ -27,7 +27,6 @@ namespace legate {
 namespace comm {
 namespace coll {
 
-using namespace Legion;
 extern Logger log_coll;
 
 enum CollTag : int {
@@ -572,4 +571,4 @@ int MPINetwork::generateGatherTag(int rank, CollComm global_comm)
 
 }  // namespace coll
 }  // namespace comm
-}  // namespace legate
+}  // namespace legate
diff --git a/src/core/data/allocator.cc b/src/core/data/allocator.cc
@@ -19,7 +19,7 @@
 
 namespace legate {
 
-ScopedAllocator::ScopedAllocator(Legion::Memory::Kind kind, bool scoped, size_t alignment)
+ScopedAllocator::ScopedAllocator(Memory::Kind kind, bool scoped, size_t alignment)
   : target_kind_(kind), scoped_(scoped), alignment_(alignment)
 {
 }
@@ -59,4 +59,4 @@ void ScopedAllocator::deallocate(void* ptr)
   buffer.destroy();
 }
 
-}  // namespace legate
+}  // namespace legate
diff --git a/src/core/data/allocator.h b/src/core/data/allocator.h
@@ -31,18 +31,18 @@ class ScopedAllocator {
 
   // Iff 'scoped', all allocations will be released upon destruction.
   // Otherwise this is up to the runtime after the task has finished.
-  ScopedAllocator(Legion::Memory::Kind kind, bool scoped = true, size_t alignment = 16);
+  ScopedAllocator(Memory::Kind kind, bool scoped = true, size_t alignment = 16);
   ~ScopedAllocator();
 
  public:
   void* allocate(size_t bytes);
   void deallocate(void* ptr);
 
  private:
-  Legion::Memory::Kind target_kind_{Legion::Memory::Kind::SYSTEM_MEM};
+  Memory::Kind target_kind_{Memory::Kind::SYSTEM_MEM};
   bool scoped_;
   size_t alignment_;
   std::unordered_map<const void*, ByteBuffer> buffers_{};
 };
 
-}  // namespace legate
+}  // namespace legate
diff --git a/src/core/data/buffer.h b/src/core/data/buffer.h
@@ -19,6 +19,7 @@
 #include "legion.h"
 
 #include "core/utilities/machine.h"
+#include "core/utilities/typedefs.h"
 
 namespace legate {
 
@@ -43,11 +44,10 @@ using Buffer = Legion::DeferredBuffer<VAL, DIM>;
 // after it's technically been deallocated.
 
 template <typename VAL, int32_t DIM>
-Buffer<VAL, DIM> create_buffer(const Legion::Point<DIM>& extents,
-                               Legion::Memory::Kind kind = Legion::Memory::Kind::NO_MEMKIND,
-                               size_t alignment          = 16)
+Buffer<VAL, DIM> create_buffer(const Point<DIM>& extents,
+                               Memory::Kind kind = Memory::Kind::NO_MEMKIND,
+                               size_t alignment  = 16)
 {
-  using namespace Legion;
   if (Memory::Kind::NO_MEMKIND == kind) kind = find_memory_kind_for_executing_processor(false);
   auto hi = extents - Point<DIM>::ONES();
   // We just avoid creating empty buffers, as they cause all sorts of headaches.
@@ -58,10 +58,10 @@ Buffer<VAL, DIM> create_buffer(const Legion::Point<DIM>& extents,
 
 template <typename VAL>
 Buffer<VAL> create_buffer(size_t size,
-                          Legion::Memory::Kind kind = Legion::Memory::Kind::NO_MEMKIND,
-                          size_t alignment          = 16)
+                          Memory::Kind kind = Memory::Kind::NO_MEMKIND,
+                          size_t alignment  = 16)
 {
-  return create_buffer<VAL, 1>(Legion::Point<1>(size), kind, alignment);
+  return create_buffer<VAL, 1>(Point<1>(size), kind, alignment);
 }
 
 }  // namespace legate