Prototype of realm subgraph compilation approach

src/realm/cuda/cuda_internal.h

@@ -191,15 +191,6 @@ namespace Realm {
       GPUProcessor *proc = nullptr; // TODO(cperry): delete me
     };
 
-    // an interface for receiving completion notification for a GPU operation
-    //  (right now, just copies)
-    class GPUCompletionNotification {
-    public:
-      virtual ~GPUCompletionNotification(void) {}
-
-      virtual void request_completed(void) = 0;
-    };
-
     class GPUWorkFence : public Realm::Operation::AsyncWorkItem {
     public:
       GPUWorkFence(GPU *gpu, Realm::Operation *op);
@@ -258,7 +249,7 @@ namespace Realm {
       void add_notification(GPUCompletionNotification *notification);
       void add_event(CUevent event, GPUWorkFence *fence,
                      GPUCompletionNotification *notification = NULL,
-                     GPUWorkStart *start = NULL);
+                     GPUWorkStart *start = NULL, bool return_event = true);
       void wait_on_streams(const std::set<GPUStream *> &other_streams);
 
       // atomically checks rate limit counters and returns true if 'bytes'
@@ -280,6 +271,7 @@ namespace Realm {
       Mutex mutex;
       struct PendingEvent {
         CUevent event;
+        bool return_event = true;
         GPUWorkFence *fence;
         GPUWorkStart *start;
         GPUCompletionNotification *notification;
@@ -345,6 +337,7 @@ namespace Realm {
 
       CUevent get_event(bool external = false);
       void return_event(CUevent e, bool external = false);
+      void return_events(CUevent* events, size_t num_events);
 
     protected:
       Mutex mutex;
@@ -545,6 +538,14 @@ namespace Realm {
 
       virtual void shutdown(void);
 
+      virtual void push_subgraph_replay_context() override;
+      virtual void pop_subgraph_replay_context() override;
+
+      virtual void push_subgraph_task_replay_context(SubgraphOperationProfilingInfo* prof) override;
+      virtual void pop_subgraph_task_replay_context(void** token, void* trigger, SubgraphOperationProfilingInfo* prof) override;
+      virtual void sync_task_async_effect(void* token) override;
+      virtual void return_subgraph_async_tokens(const std::vector<void*>& tokens) override;
+
     protected:
       virtual void execute_task(Processor::TaskFuncID func_id,
                                 const ByteArrayRef &task_args);
@@ -773,6 +774,18 @@ namespace Realm {
       size_t width_in_bytes, height, depth;
     };
 
+    // Utility methods for asynchronous XferDes objects to
+    // interact with subgraphs. Must be called from within
+    // a CUDA context.
+
+    // Synchronize with any previous subgraph work launched on the given stream.
+    void sync_subgraph_incoming_deps(ProcSubgraphReplayState* all_proc_states, unsigned index, GPUStream* stream);
+    // Add a GPUCompletionNotification to the given stream, which tracks
+    // when work launched by a transfer is done.
+    void add_transfer_completion_notification(ProcSubgraphReplayState* all_proc_states, GPUStream* stream, GPUStream* depstream, GPUCompletionNotification* completion);
+    // Notify the subgraph that the control side of GPU subgraph operation has completed.
+    void notify_subgraph_control_completion(ProcSubgraphReplayState* all_proc_states, unsigned subgraph_index, unsigned xd_index, GPUStream* depstream);
+
     class GPUChannel;
 
     class GPUXferDes : public XferDes {
@@ -785,6 +798,10 @@ namespace Realm {
 
       bool progress_xd(GPUChannel *channel, TimeLimit work_until);
 
+      bool launches_async_work_locally() override { return true; }
+      void on_subgraph_control_completion() override;
+      LocalTaskProcessor* get_async_event_proc() override;
+
     private:
       std::vector<GPU *> src_gpus, dst_gpus;
       std::vector<bool> dst_is_ipc;
@@ -802,6 +819,10 @@ namespace Realm {
       long get_requests(Request **requests, long nr);
       bool progress_xd(GPUIndirectChannel *channel, TimeLimit work_until);
 
+      bool launches_async_work_locally() override { return true; }
+      void on_subgraph_control_completion() override;
+      LocalTaskProcessor* get_async_event_proc() override;
+
     protected:
       std::vector<GPU *> src_gpus, dst_gpus;
       std::vector<bool> dst_is_ipc;
@@ -811,7 +832,6 @@ namespace Realm {
       : public SingleXDQChannel<GPUIndirectChannel, GPUIndirectXferDes> {
     public:
       GPUIndirectChannel(GPU *_src_gpu, XferDesKind _kind, BackgroundWorkManager *bgwork);
-      ~GPUIndirectChannel();
 
       // multi-threading of cuda copies for a given device is disabled by
       //  default (can be re-enabled with -cuda:mtdma 1)
@@ -841,10 +861,12 @@ namespace Realm {
 
       long submit(Request **requests, long nr);
       GPU *get_gpu() const { return src_gpu; }
+      void shutdown() override;
 
     protected:
       friend class GPUIndirectXferDes;
       GPU *src_gpu;
+      GPUStream* subgraph_stream;
     };
 
     class GPUIndirectRemoteChannelInfo : public SimpleRemoteChannelInfo {
@@ -887,7 +909,6 @@ namespace Realm {
     class GPUChannel : public SingleXDQChannel<GPUChannel, GPUXferDes> {
     public:
       GPUChannel(GPU *_src_gpu, XferDesKind _kind, BackgroundWorkManager *bgwork);
-      ~GPUChannel();
 
       // multi-threading of cuda copies for a given device is disabled by
       //  default (can be re-enabled with -cuda:mtdma 1)
@@ -903,10 +924,13 @@ namespace Realm {
 
       long submit(Request **requests, long nr);
       GPU *get_gpu() const { return src_gpu; }
+      void shutdown() override;
 
     private:
+      friend class GPUXferDes;
       GPU *src_gpu;
-      // std::deque<Request*> pending_copies;
+      // Stream to manage interaction with potential subgraph replays.
+      GPUStream* subgraph_stream;
     };
 
     class GPUfillChannel;
@@ -918,12 +942,19 @@ namespace Realm {
                      const std::vector<XferDesPortInfo> &outputs_info, int _priority,
                      const void *_fill_data, size_t _fill_size, size_t _fill_total);
 
+      void reset(const std::vector<off_t>& ib_offsets) override;
+
       long get_requests(Request **requests, long nr);
 
       bool progress_xd(GPUfillChannel *channel, TimeLimit work_until);
 
+      bool launches_async_work_locally() override { return true; }
+      void on_subgraph_control_completion() override;
+      LocalTaskProcessor* get_async_event_proc() override;
+
     protected:
       size_t reduced_fill_size;
+      size_t fill_total;
     };
 
     class GPUfillChannel : public SingleXDQChannel<GPUfillChannel, GPUfillXferDes> {
@@ -942,11 +973,13 @@ namespace Realm {
                                        size_t fill_total);
 
       long submit(Request **requests, long nr);
+      void shutdown() override;
 
     protected:
       friend class GPUfillXferDes;
 
       GPU *gpu;
+      GPUStream* subgraph_stream;
     };
 
     class GPUreduceChannel;
@@ -962,6 +995,10 @@ namespace Realm {
 
       bool progress_xd(GPUreduceChannel *channel, TimeLimit work_until);
 
+      bool launches_async_work_locally() override { return true; }
+      void on_subgraph_control_completion() override;
+      LocalTaskProcessor* get_async_event_proc() override;
+
     protected:
       XferDesRedopInfo redop_info;
       const ReductionOpUntyped *redop;
@@ -992,11 +1029,13 @@ namespace Realm {
                                size_t fill_total) override;
 
       long submit(Request **requests, long nr) override;
+      void shutdown() override;
 
     protected:
       friend class GPUreduceXferDes;
 
       GPU *gpu;
+      GPUStream* subgraph_stream;
     };
 
     class GPUreduceRemoteChannelInfo : public SimpleRemoteChannelInfo {

src/realm/cuda/cuda_module.cc

@@ -230,7 +230,7 @@ namespace Realm {
 
     void GPUStream::add_event(CUevent event, GPUWorkFence *fence,
                               GPUCompletionNotification *notification,
-                              GPUWorkStart *start)
+                              GPUWorkStart *start, bool return_event)
     {
       bool add_to_worker = false;
       {
@@ -245,6 +245,7 @@ namespace Realm {
         e.fence = fence;
         e.start = start;
         e.notification = notification;
+        e.return_event = return_event;
 
         pending_events.push_back(e);
       }
@@ -320,14 +321,15 @@ namespace Realm {
         log_stream.debug() << "CUDA event " << event << " triggered on stream " << stream
                            << " (GPU " << gpu << ")";
 
-        // give event back to GPU for reuse
-        gpu->event_pool.return_event(event);
-
         // this event has triggered, so figure out the fence/notification to trigger
         //  and also peek at the next event
         GPUWorkFence *fence = 0;
         GPUWorkStart *start = 0;
         GPUCompletionNotification *notification = 0;
+        bool return_event = false;
+        // The logic below may change the value stored in
+        // `event`, so save it before manipulating the queue.
+        CUevent event_to_return = event;
 
         {
           AutoLock<> al(mutex);
@@ -337,6 +339,7 @@ namespace Realm {
           fence = e.fence;
           start = e.start;
           notification = e.notification;
+          return_event = e.return_event;
           pending_events.pop_front();
 
           if(pending_events.empty()) {
@@ -346,6 +349,10 @@ namespace Realm {
             event = pending_events.front().event;
         }
 
+        if (return_event) {
+          gpu->event_pool.return_event(event_to_return);
+        }
+
         if(start) {
           start->mark_gpu_work_start();
         }
@@ -585,6 +592,16 @@ namespace Realm {
       available_events[current_size++] = e;
     }
 
+    void GPUEventPool::return_events(CUevent* events, size_t num_events) {
+      AutoLock<> al(mutex);
+
+      assert(current_size < total_size);
+
+      for (size_t i = 0; i < num_events; i++) {
+        available_events[current_size++] = events[i];
+      }
+    }
+
     ////////////////////////////////////////////////////////////////////////
     //
     // class ContextSynchronizer
@@ -990,9 +1007,9 @@ namespace Realm {
         }
       }
       params.set_num_cores(1);
-      params.set_alu_usage(params.CORE_USAGE_SHARED);
-      params.set_fpu_usage(params.CORE_USAGE_SHARED);
-      params.set_ldst_usage(params.CORE_USAGE_SHARED);
+      params.set_alu_usage(params.CORE_USAGE_EXCLUSIVE);
+      params.set_fpu_usage(params.CORE_USAGE_EXCLUSIVE);
+      params.set_ldst_usage(params.CORE_USAGE_EXCLUSIVE);
       params.set_max_stack_size(_stack_size);
 
       std::string name = stringbuilder() << "GPU proc " << _me;
@@ -1013,6 +1030,85 @@ namespace Realm {
 
     GPUProcessor::~GPUProcessor(void) { delete core_rsrv; }
 
+   void GPUProcessor::push_subgraph_replay_context() {
+     gpu->push_context();
+   }
+
+   void GPUProcessor::pop_subgraph_replay_context() {
+     gpu->pop_context();
+   }
+
+   void GPUProcessor::push_subgraph_task_replay_context(SubgraphOperationProfilingInfo* prof) {
+     // TODO (rohany): A good chunk of this is ripped from the
+     //  GPUContextManager...
+     assert(ThreadLocal::current_gpu_stream == nullptr);
+     GPUStream *s = gpu->get_next_task_stream();
+     ThreadLocal::current_gpu_stream = s;
+     assert(!ThreadLocal::created_gpu_streams);
+
+     // a task can force context sync on task completion either on or off during
+     //  execution, so use -1 as a "no preference" value
+     ThreadLocal::context_sync_required = -1;
+
+     // Push profiling information, if necessary. This completion operation
+     // will set just the GPU start timeline.
+     if (prof && prof->wants_gpu_timeline) {
+       auto e = gpu->event_pool.get_event();
+       CHECK_CU(CUDA_DRIVER_FNPTR(cuEventRecord)(e, s->get_stream()));
+       auto completion = new GPUProfInfoTrigger(prof, true /* start */);
+       s->add_event(e, nullptr, completion, nullptr, true /* return_event */);
+     }
+  }
+
+   void GPUProcessor::pop_subgraph_task_replay_context(void** token, void* trigger, SubgraphOperationProfilingInfo* prof) {
+    GPUStream *s = ThreadLocal::current_gpu_stream;
+    assert(!ThreadLocal::created_gpu_streams);
+
+    assert(!gpu->module->config->cfg_task_legacy_sync);
+    assert(ThreadLocal::context_sync_required == 0);
+
+    // Again, handle post-task GPU timeline information. This must happen
+    // before we enqueue the trigger, otherwise the profiling information
+    // may not get recorded by the time the subgraph finishes and it is
+    // sent out during subgraph cleanup. It has to handle several different
+    // kinds of profiling requests on operation finish.
+    if (prof && (prof->wants_gpu_timeline || prof->wants_timeline || prof->wants_fevent)) {
+      auto e = gpu->event_pool.get_event();
+      CHECK_CU(CUDA_DRIVER_FNPTR(cuEventRecord)(e, s->get_stream()));
+      auto completion = new GPUProfInfoTrigger(prof, false /* start */);
+      s->add_event(e, nullptr, completion, nullptr, true /* return_event */);
+    }
+
+    // Record an event for the completion of kernels launched by the task.
+    auto e = gpu->event_pool.get_event();
+    // Return the event to the caller. The caller is responsible
+    // for cleaning up this
+    *(CUevent_st**)(token) = e;
+    // Actually record the event.
+    CHECK_CU(CUDA_DRIVER_FNPTR(cuEventRecord)(e, s->get_stream()));
+
+    // Enqueue a poller for the completion of the launched kernels.
+    if (trigger) {
+      auto completion = static_cast<GPUCompletionNotification*>(trigger);
+      // The subgraph cleanup will be responsible for cleaning up
+      // and returning events to the pool.
+      s->add_event(e, nullptr, completion, nullptr, false /* return_event */);
+    }
+
+    ThreadLocal::current_gpu_stream = nullptr;
+  }
+
+   void GPUProcessor::sync_task_async_effect(void *token) {
+     // The GPUProcessor will only ever return CUevent_st* to the scheduler.
+     CUevent_st* e = (CUevent_st*)(token);
+     GPUStream *s = ThreadLocal::current_gpu_stream;
+     CHECK_CU(CUDA_DRIVER_FNPTR(cuStreamWaitEvent)(s->get_stream(), e, 0));
+   }
+
+   void GPUProcessor::return_subgraph_async_tokens(const std::vector<void*>& tokens) {
+     gpu->event_pool.return_events((CUevent_st**)tokens.data(), tokens.size());
+   }
+
     GPUStream *GPU::find_stream(CUstream stream) const
     {
       for(std::vector<GPUStream *>::const_iterator it = task_streams.begin();

src/realm/cuda/cuda_module.h

@@ -106,6 +106,16 @@ namespace Realm {
     class GPUZCMemory;
     class GPUReplHeapListener;
 
+    // TODO (rohany): Moving this out of the CUDA module (for now).
+    // an interface for receiving completion notification for a GPU operation
+    //  (right now, just copies)
+    class GPUCompletionNotification {
+      public:
+        virtual ~GPUCompletionNotification(void) {}
+
+        virtual void request_completed(void) = 0;
+    };
+
     class CudaModuleConfig : public ModuleConfig {
       friend class CudaModule;