PyTorch tracer (#243)

Co-authored-by: Noli Gerawork <[email protected]>

Author: chhwang

ark/api/executor.cpp

@@ -68,6 +68,27 @@ Dims Tensor::torch_strides() const {
     return Dims();
 }
 
+void *Tensor::data() const {
+    if (ref_) {
+        return ref_->data();
+    }
+    return nullptr;
+}
+
+void *Tensor::data(void *data) {
+    if (ref_) {
+        return ref_->data(data);
+    }
+    return nullptr;
+}
+
+bool Tensor::is_external() const {
+    if (ref_) {
+        return ref_->is_external();
+    }
+    return false;
+}
+
 std::ostream &operator<<(std::ostream &os, const Tensor &tensor) {
     if (tensor.is_null()) {
         os << "null";

ark/api/tensor.cpp

@@ -0,0 +1,34 @@
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT license.
+
+#include "buffer_registry.hpp"
+
+#include "gpu/gpu_logging.hpp"
+
+namespace ark {
+
+BufferRegistry &BufferRegistry::get_instance() {
+    static BufferRegistry instance;
+    return instance;
+}
+
+void BufferRegistry::set(size_t id, void *data, int device_id,
+                         bool is_external) {
+    if (data != nullptr && device_id < 0) {
+        gpuPointerAttributes attr;
+        GLOG(gpuPointerGetAttributes(&attr, data));
+        device_id = attr.device;
+    }
+    buffers_[id] =
+        std::make_shared<BufferRegistry::Info>(data, device_id, is_external);
+}
+
+std::shared_ptr<BufferRegistry::Info> BufferRegistry::get(size_t id) const {
+    auto it = buffers_.find(id);
+    if (it != buffers_.end()) {
+        return it->second;
+    }
+    return nullptr;
+}
+
+}  // namespace ark

ark/buffer_registry.cpp

@@ -0,0 +1,41 @@
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT license.
+
+#ifndef ARK_BUFFER_REGISTRY_HPP_
+#define ARK_BUFFER_REGISTRY_HPP_
+
+#include <memory>
+#include <unordered_map>
+
+namespace ark {
+
+/// Manages addresses of all allocated buffers including externally managed
+/// buffers.
+class BufferRegistry {
+   public:
+    struct Info {
+        Info(void *data, int device_id, bool is_external)
+            : data(data), device_id(device_id), is_external(is_external) {}
+        void *data;
+        int device_id;
+        bool is_external;
+    };
+
+    ~BufferRegistry() {}
+
+    static BufferRegistry &get_instance();
+
+    void set(size_t id, void *data, int device_id, bool is_external);
+
+    std::shared_ptr<Info> get(size_t id) const;
+
+   private:
+    std::unordered_map<size_t, std::shared_ptr<Info>> buffers_;
+    BufferRegistry() {}
+    BufferRegistry(const BufferRegistry &) = delete;
+    BufferRegistry &operator=(const BufferRegistry &) = delete;
+};
+
+}  // namespace ark
+
+#endif  // ARK_BUFFER_REGISTRY_HPP_

ark/buffer_registry.hpp

@@ -7,8 +7,8 @@
 #include <utility>
 
 #include "ark/data_type.hpp"
+#include "buffer_registry.hpp"
 #include "env.h"
-#include "external_buffer_registry.hpp"
 #include "file_io.h"
 #include "logging.hpp"
 #include "model/model_buffer.hpp"
@@ -56,9 +56,7 @@ class CodeGenerator::Impl {
    public:
     Impl(const PlanJson &plan,
          const std::map<size_t, size_t> &buffer_id_to_offset,
-         const std::map<size_t, std::pair<std::string, void *>>
-             &buffer_id_to_kernel_arg,
-         const std::string &name);
+         const std::set<size_t> &extra_buffer_ids, const std::string &name);
     ~Impl() = default;
 
    private:
@@ -83,7 +81,7 @@ class CodeGenerator::Impl {
     friend class CodeGenerator;
 
     std::map<size_t, size_t> buffer_id_to_offset_;
-    std::map<size_t, std::pair<std::string, void *>> buffer_id_to_kernel_arg_;
+    std::set<size_t> extra_buffer_ids_;
     std::string name_;
     int rank_;
     int world_size_;
@@ -94,11 +92,10 @@ class CodeGenerator::Impl {
 
 CodeGenerator::Impl::Impl(const PlanJson &plan,
                           const std::map<size_t, size_t> &buffer_id_to_offset,
-                          const std::map<size_t, std::pair<std::string, void *>>
-                              &buffer_id_to_kernel_arg,
+                          const std::set<size_t> &extra_buffer_ids,
                           const std::string &name)
     : buffer_id_to_offset_(buffer_id_to_offset),
-      buffer_id_to_kernel_arg_(buffer_id_to_kernel_arg),
+      extra_buffer_ids_(extra_buffer_ids),
       name_(name) {
     rank_ = plan.at("Rank");
     world_size_ = plan.at("WorldSize");
@@ -191,8 +188,8 @@ CodeGenerator::Impl::Impl(const PlanJson &plan,
 
     // Generate the global arguments
     std::stringstream global_args_ss, function_args_ss, arg_types_ss;
-    for (const auto &[buf_id, kernel_arg] : buffer_id_to_kernel_arg_) {
-        const auto &arg_name = kernel_arg.first;
+    for (auto buf_id : extra_buffer_ids_) {
+        std::string arg_name = "_ext_buf_" + std::to_string(buf_id);
         global_args_ss << "void *" << arg_name << ", ";
         function_args_ss << arg_name << ", ";
         arg_types_ss << "void *, ";
@@ -263,6 +260,7 @@ std::string CodeGenerator::Impl::def_task(const Json &task_json) {
     }
     ss << "__device__ void t" << task_json["Id"]
        << "(char *_buf, int _idx, int _spw, @GLOBAL_ARGS@) {\n";
+    auto &buf_reg = BufferRegistry::get_instance();
     op_idx = 0;
     for (auto &op_json : task_json["Ops"]) {
         auto op = ModelOp::deserialize(op_json);
@@ -273,29 +271,36 @@ std::string CodeGenerator::Impl::def_task(const Json &task_json) {
             if (arg.type_name() == "TENSOR") {
                 auto tns = arg.value<ModelTensorRef>();
                 size_t buffer_id = tns->buffer()->id();
-                auto it = buffer_id_to_kernel_arg_.find(buffer_id);
-                if (it == buffer_id_to_kernel_arg_.end()) {
-                    size_t buffer_offset = buffer_id_to_offset_.at(buffer_id);
+                auto it = buffer_id_to_offset_.find(buffer_id);
+                auto buf_info = buf_reg.get(buffer_id);
+                if ((buf_info && buf_info->is_external) ||
+                    (it == buffer_id_to_offset_.end())) {
+                    ss << "(" << tns->data_type()->type_str() << "*)_ext_buf_"
+                       << buffer_id;
+                } else {
+                    size_t buffer_offset;
+                    buffer_offset = it->second;
                     size_t offset = buffer_offset + ModelOffset(tns).value();
                     ss << "(" << tns->data_type()->type_str() << "*)&_buf["
                        << offset << "]";
-                } else {
-                    const auto &name = it->second.first;
-                    ss << "(" << tns->data_type()->type_str() << "*)" << name;
                 }
             } else if (arg.type_name() == "OFFSET") {
                 auto moff = arg.value<ModelOffset>();
                 size_t buffer_id = moff.buffer_id();
-                auto it = buffer_id_to_kernel_arg_.find(buffer_id);
-                if (it == buffer_id_to_kernel_arg_.end()) {
-                    size_t buffer_offset = buffer_id_to_offset_.at(buffer_id);
+                auto buf_info = buf_reg.get(buffer_id);
+                if (buf_info && buf_info->is_external) {
+                    size_t offset = moff.value();
+                    ss << "(uint64_t)((char*)_ext_buf_" << buffer_id << " + "
+                       << offset << ")";
+                } else {
+                    size_t buffer_offset;
+                    auto it = buffer_id_to_offset_.find(buffer_id);
+                    if (it == buffer_id_to_offset_.end()) {
+                        ERR(InternalError, "buffer ID not found: ", buffer_id);
+                    }
+                    buffer_offset = it->second;
                     size_t offset = buffer_offset + moff.value();
                     ss << offset;
-                } else {
-                    const auto &name = it->second.first;
-                    size_t offset = moff.value();
-                    ss << "(uint64_t)((char*)" << name << " + " << offset
-                       << ")";
                 }
             } else {
                 ss << arg.serialize().begin().value();
@@ -496,11 +501,9 @@ std::string CodeGenerator::Impl::sync_process_range(const Range<size_t> &range,
 
 CodeGenerator::CodeGenerator(
     const PlanJson &plan, const std::map<size_t, size_t> &buffer_id_to_offset,
-    const std::map<size_t, std::pair<std::string, void *>>
-        &buffer_id_to_kernel_arg,
-    const std::string &name)
-    : impl_(std::make_shared<Impl>(plan, buffer_id_to_offset,
-                                   buffer_id_to_kernel_arg, name)) {}
+    const std::set<size_t> &extra_buffer_ids, const std::string &name)
+    : impl_(std::make_shared<Impl>(plan, buffer_id_to_offset, extra_buffer_ids,
+                                   name)) {}
 
 std::string CodeGenerator::code() const { return impl_->code_; }
 

ark/codegen.cpp

@@ -6,8 +6,8 @@
 
 #include <map>
 #include <memory>
+#include <set>
 #include <string>
-#include <utility>
 
 #include "model/model_json.hpp"
 
@@ -17,8 +17,7 @@ class CodeGenerator {
    public:
     CodeGenerator(const PlanJson &plan,
                   const std::map<size_t, size_t> &buffer_id_to_offset,
-                  const std::map<size_t, std::pair<std::string, void *>>
-                      &buffer_id_to_kernel_arg,
+                  const std::set<size_t> &extra_buffer_ids,
                   const std::string &name = "ark_kernel");
 
     ~CodeGenerator() = default;

ark/codegen.hpp

@@ -7,21 +7,25 @@
 #include "gpu/gpu_manager.hpp"
 
 namespace ark {
+
 class GpuEvent::Impl {
    public:
-    Impl(bool disable_timing);
+    Impl(int device_id, bool disable_timing);
     ~Impl();
     Impl(const Impl&) = delete;
     Impl& operator=(const Impl&) = delete;
 
+    int device_id() const { return device_id_; }
     void record(gpuStream stream);
     float elapsed_msec(const GpuEvent& other) const;
 
    private:
+    int device_id_;
     gpuEvent event_;
 };
 
-GpuEvent::Impl::Impl(bool disable_timing) {
+GpuEvent::Impl::Impl(int device_id, bool disable_timing)
+    : device_id_(device_id) {
     unsigned int flags = 0;
     if (disable_timing) {
         flags |= gpuEventDisableTiming;
@@ -41,8 +45,10 @@ float GpuEvent::Impl::elapsed_msec(const GpuEvent& other) const {
     return elapsed;
 }
 
-GpuEvent::GpuEvent(bool disable_timing)
-    : pimpl_(std::make_shared<Impl>(disable_timing)) {}
+GpuEvent::GpuEvent(int device_id, bool disable_timing)
+    : pimpl_(std::make_shared<Impl>(device_id, disable_timing)) {}
+
+int GpuEvent::device_id() const { return pimpl_->device_id(); }
 
 void GpuEvent::record(gpuStream stream) { pimpl_->record(stream); }
 

ark/external_buffer_registry.cpp

@@ -19,13 +19,14 @@ class GpuEvent {
     GpuEvent(const GpuEvent &) = delete;
     GpuEvent &operator=(const GpuEvent &) = delete;
 
+    int device_id() const;
     void record(gpuStream stream);
     float elapsed_msec(const GpuEvent &other) const;
 
    protected:
     friend class GpuManager;
 
-    GpuEvent(bool disable_timing = false);
+    GpuEvent(int device_id, bool disable_timing = false);
 
    private:
     class Impl;

ark/external_buffer_registry.hpp

@@ -118,7 +118,8 @@ std::shared_ptr<GpuHostMemory> GpuManager::malloc_host(size_t bytes,
 }
 
 std::shared_ptr<GpuEvent> GpuManager::create_event(bool disable_timing) const {
-    return std::shared_ptr<GpuEvent>(new GpuEvent(disable_timing));
+    return std::shared_ptr<GpuEvent>(
+        new GpuEvent(pimpl_->gpu_id_, disable_timing));
 }
 
 std::shared_ptr<GpuStream> GpuManager::create_stream() const {