wip

microsoft · Apr 5, 2024 · 473288c · 473288c
1 parent 20840bb
commit 473288c
Show file tree

Hide file tree

Showing 152 changed files with 4,053 additions and 1,659 deletions.
diff --git a/.vscode/settings.json b/.vscode/settings.json
@@ -10,4 +10,102 @@
     "cmake.ctestArgs": [
         "--verbose"
     ],
+    "files.associations": {
+        "ostream": "cpp",
+        "stdexcept": "cpp",
+        "string": "cpp",
+        "iosfwd": "cpp",
+        "memory": "cpp",
+        "stack": "cpp",
+        "any": "cpp",
+        "array": "cpp",
+        "atomic": "cpp",
+        "bit": "cpp",
+        "*.tcc": "cpp",
+        "bitset": "cpp",
+        "cctype": "cpp",
+        "cfenv": "cpp",
+        "chrono": "cpp",
+        "cinttypes": "cpp",
+        "clocale": "cpp",
+        "cmath": "cpp",
+        "codecvt": "cpp",
+        "compare": "cpp",
+        "complex": "cpp",
+        "concepts": "cpp",
+        "condition_variable": "cpp",
+        "csignal": "cpp",
+        "cstdarg": "cpp",
+        "cstddef": "cpp",
+        "cstdint": "cpp",
+        "cstdio": "cpp",
+        "cstdlib": "cpp",
+        "cstring": "cpp",
+        "ctime": "cpp",
+        "cwchar": "cpp",
+        "cwctype": "cpp",
+        "deque": "cpp",
+        "forward_list": "cpp",
+        "list": "cpp",
+        "map": "cpp",
+        "set": "cpp",
+        "unordered_map": "cpp",
+        "unordered_set": "cpp",
+        "vector": "cpp",
+        "exception": "cpp",
+        "algorithm": "cpp",
+        "functional": "cpp",
+        "iterator": "cpp",
+        "memory_resource": "cpp",
+        "numeric": "cpp",
+        "optional": "cpp",
+        "random": "cpp",
+        "ratio": "cpp",
+        "regex": "cpp",
+        "string_view": "cpp",
+        "system_error": "cpp",
+        "tuple": "cpp",
+        "type_traits": "cpp",
+        "utility": "cpp",
+        "fstream": "cpp",
+        "future": "cpp",
+        "initializer_list": "cpp",
+        "iomanip": "cpp",
+        "iostream": "cpp",
+        "istream": "cpp",
+        "limits": "cpp",
+        "mutex": "cpp",
+        "new": "cpp",
+        "numbers": "cpp",
+        "ranges": "cpp",
+        "semaphore": "cpp",
+        "span": "cpp",
+        "sstream": "cpp",
+        "stop_token": "cpp",
+        "streambuf": "cpp",
+        "thread": "cpp",
+        "typeinfo": "cpp",
+        "valarray": "cpp",
+        "variant": "cpp",
+        "__nullptr": "cpp",
+        "__hash_table": "cpp",
+        "__split_buffer": "cpp",
+        "__tree": "cpp",
+        "queue": "cpp",
+        "__locale": "cpp",
+        "*.ipp": "cpp",
+        "strstream": "cpp",
+        "typeindex": "cpp",
+        "locale": "cpp",
+        "__node_handle": "cpp",
+        "__threading_support": "cpp",
+        "__functional_03": "cpp",
+        "filesystem": "cpp",
+        "__bit_reference": "cpp",
+        "__config": "cpp",
+        "__debug": "cpp",
+        "version": "cpp",
+        "__functional_base": "cpp",
+        "__memory": "cpp"
+    },
 }
diff --git a/ark/CMakeLists.txt b/ark/CMakeLists.txt
@@ -1,9 +1,9 @@
 # Copyright (c) Microsoft Corporation.
 # Licensed under the MIT license.
 
-file(GLOB_RECURSE SOURCES CONFIGURE_DEPENDS *.cc)
-file(GLOB_RECURSE UT_SOURCES CONFIGURE_DEPENDS *_test.cc *_test.cu)
-file(GLOB_RECURSE UT_COMMON_SOURCES CONFIGURE_DEPENDS ${CMAKE_CURRENT_SOURCE_DIR}/unittest/*.cc)
+file(GLOB_RECURSE SOURCES CONFIGURE_DEPENDS *.cpp)
+file(GLOB_RECURSE UT_SOURCES CONFIGURE_DEPENDS *_test.cpp)
+file(GLOB_RECURSE UT_COMMON_SOURCES CONFIGURE_DEPENDS ${CMAKE_CURRENT_SOURCE_DIR}/unittest/*.cpp)
 list(REMOVE_ITEM SOURCES ${UT_SOURCES} ${UT_COMMON_SOURCES})
 
 if(USE_ROCM)

diff --git a/ark/cpu_timer.cc → ark/cpu_timer.cpp b/ark/cpu_timer.cc → ark/cpu_timer.cpp
diff --git a/ark/dims.cc → ark/dims.cpp b/ark/dims.cc → ark/dims.cpp
@@ -1,9 +1,11 @@
 // Copyright (c) Microsoft Corporation.
 // Licensed under the MIT license.
 
+#include "ark/dims.hpp"
+
 #include <vector>
 
-#include "include/ark.h"
+#include "error.hpp"
 #include "logging.h"
 
 namespace ark {
@@ -101,6 +103,19 @@ Dims Dims::dims4() const {
     return ret;
 }
 
+// Return true if all valid dimensions are zero.
+bool Dims::is_zeros() const {
+    if (this->is_invalid()) {
+        return false;
+    }
+    const DimType *v = this->data;
+    for (int i = 0; i < DIMS_LEN; ++i) {
+        if (v[i] == NO_DIM) break;
+        if (v[i] != 0) return false;
+    }
+    return true;
+}
+
 // Return true if the dimensions are empty.
 bool Dims::is_no_dim() const {
     const DimType *v = this->data;

diff --git a/ark/env.cc → ark/env.cpp b/ark/env.cc → ark/env.cpp
diff --git a/ark/error.hpp b/ark/error.hpp
@@ -0,0 +1,29 @@
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT license.
+
+#ifndef ARK_ERROR_HPP_
+#define ARK_ERROR_HPP_
+
+#include <stdexcept>
+#include <string>
+
+namespace ark {
+
+#define REGISTER_ERROR_TYPE(_name)                                 \
+    class _name : public std::runtime_error {                      \
+       public:                                                     \
+        _name(const std::string &msg) : std::runtime_error(msg) {} \
+    };
+
+REGISTER_ERROR_TYPE(InvalidUsageError)
+REGISTER_ERROR_TYPE(ModelError)
+REGISTER_ERROR_TYPE(SchedulerError)
+REGISTER_ERROR_TYPE(ExecutorError)
+REGISTER_ERROR_TYPE(SystemError)
+REGISTER_ERROR_TYPE(GpuError)
+REGISTER_ERROR_TYPE(RuntimeError)
+REGISTER_ERROR_TYPE(UnitTestError)
+
+}  // namespace ark
+
+#endif  // ARK_ERROR_HPP_
diff --git a/ark/executor/executor.cc b/ark/executor/executor.cc
@@ -0,0 +1,87 @@
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT license.
+
+#include "ark/executor.hpp"
+
+#include <algorithm>
+#include <memory>
+#include <string>
+
+#include "env.h"
+#include "gpu/gpu_loop_kernel.h"
+#include "logging.h"
+#include "sched/sched.h"
+
+namespace ark {
+
+class Executor::Impl {
+   public:
+    Impl(int rank, int world_size, const Model &model, const Schedule &schedule, const std::string &name,
+         int num_warps_per_sm);
+    ~Impl() = default;
+
+    void compile();
+    void launch();
+    void run(int iter);
+    void wait();
+    float stop();
+
+   private:
+    const int rank_;
+    const int world_size_;
+    int gpu_id_;
+
+    std::shared_ptr<GpuContext> ctx_;
+    std::unique_ptr<GpuLoopKernel> glk_;
+    std::shared_ptr<GpuStream> stream_;
+};
+
+Executor::Impl::Impl(int rank, int world_size, const Model &model, const Schedule &schedule, 
+                     const std::string &name, int num_warps_per_sm)
+    : rank_(rank), world_size_(world_size) {
+    gpu_id_ = rank_ % get_env().num_ranks_per_host;
+    sched_.reset(static_cast<BaseScheduler *>(new DefaultScheduler{
+        model, gpu_id_, rank_, world_size_, num_warps_per_sm}));
+
+    ctx_ = sched_->create_context();
+    const GpuManager::Info &ginfo = ctx_->get_gpu_manager()->info();
+    stream_ = ctx_->get_gpu_manager()->create_stream();
+    glk_ = std::make_unique<GpuLoopKernel>(
+        ctx_, name, sched_->gen_code(), ginfo.num_sm, num_warps_per_sm,
+        (unsigned int)ginfo.smem_block_total);
+}
+
+void Executor::Impl::compile() { glk_->compile(); }
+
+void Executor::Impl::launch() {
+    glk_->load();
+    glk_->launch(stream_, false);
+}
+
+void Executor::Impl::run(int iter) { glk_->run(iter); }
+
+void Executor::Impl::wait() { glk_->wait(); }
+
+float Executor::Impl::stop() {
+    glk_->stop();
+    return glk_->get_elapsed_msec();
+}
+
+Executor::Executor(int rank, int world_size, const Model &model, const Schedule &schedule, 
+                   const std::string &name, int num_warps_per_sm)
+    : impl_{std::make_unique<Executor::Impl>(rank, world_size, model, schedule, name,
+                                             num_warps_per_sm)} {}
+
+Executor::~Executor() = default;
+
+void Executor::compile() { impl_->compile(); }
+
+void Executor::launch() { impl_->launch(); }
+
+void Executor::run(int iter) { impl_->run(iter); }
+
+void Executor::wait() { impl_->wait(); }
+
+float Executor::stop() { return impl_->stop(); }
+
+}  // namespace ark
diff --git a/ark/file_io.cc → ark/file_io.cpp b/ark/file_io.cc → ark/file_io.cpp
@@ -7,7 +7,6 @@
 #include <fstream>
 #include <sstream>
 
-#include "include/ark.h"
 #include "logging.h"
 
 namespace fs = std::filesystem;

diff --git a/ark/gpu/gpu_buffer.cc → ark/gpu/gpu_buffer.cpp b/ark/gpu/gpu_buffer.cc → ark/gpu/gpu_buffer.cpp
diff --git a/ark/gpu/gpu_comm_sw.cc → ark/gpu/gpu_comm_sw.cpp b/ark/gpu/gpu_comm_sw.cc → ark/gpu/gpu_comm_sw.cpp
@@ -13,7 +13,6 @@
 #include "env.h"
 #include "gpu/gpu_logging.h"
 #include "gpu/gpu_manager.h"
-#include "include/ark.h"
 #include "ipc/ipc_hosts.h"
 #include "ipc/ipc_socket.h"
 

diff --git a/ark/gpu/gpu_compile.cc → ark/gpu/gpu_compile.cpp b/ark/gpu/gpu_compile.cc → ark/gpu/gpu_compile.cpp
@@ -22,8 +22,7 @@
 #include "env.h"
 #include "file_io.h"
 #include "gpu/gpu_logging.h"
-#include "include/ark.h"
-#include "random.h"
+#include "ark/random.hpp"
 
 #define ARK_DEBUG_KERNEL 0
 

diff --git a/ark/gpu/gpu_context.cc → ark/gpu/gpu_context.cpp b/ark/gpu/gpu_context.cc → ark/gpu/gpu_context.cpp
diff --git a/ark/gpu/gpu_context_test.cc → ark/gpu/gpu_context_test.cpp b/ark/gpu/gpu_context_test.cc → ark/gpu/gpu_context_test.cpp
@@ -5,7 +5,6 @@
 
 #include <numeric>
 
-#include "include/ark.h"
 #include "unittest/unittest_utils.h"
 
 // Test initializing and destroying GpuContext
@@ -178,7 +177,6 @@ ark::unittest::State test_gpu_context_remote() {
 }
 
 int main() {
-    ark::init();
     UNITTEST(test_gpu_context_basic);
     UNITTEST(test_gpu_context_buffer_free);
     UNITTEST(test_gpu_context_buffer_alloc);

diff --git a/ark/gpu/gpu_event.cc → ark/gpu/gpu_event.cpp b/ark/gpu/gpu_event.cc → ark/gpu/gpu_event.cpp
diff --git a/ark/gpu/gpu_kernel.cc → ark/gpu/gpu_kernel.cpp b/ark/gpu/gpu_kernel.cc → ark/gpu/gpu_kernel.cpp
diff --git a/ark/gpu/gpu_kernel_test.cc → ark/gpu/gpu_kernel_test.cpp b/ark/gpu/gpu_kernel_test.cc → ark/gpu/gpu_kernel_test.cpp
@@ -4,7 +4,6 @@
 #include "gpu/gpu_kernel.h"
 
 #include "gpu/gpu_loop_kernel.h"
-#include "include/ark.h"
 #include "unittest/unittest_utils.h"
 
 const std::string void_kernel = "extern \"C\" __global__ void kernel() {}";

diff --git a/ark/gpu/gpu_logging.h b/ark/gpu/gpu_logging.h
@@ -5,7 +5,6 @@
 #define ARK_GPU_LOGGING_H_
 
 #include "gpu/gpu.h"
-#include "include/ark.h"
 #include "logging.h"
 
 #define GLOG(cmd)                                           \

diff --git a/ark/gpu/gpu_loop_kernel.cc → ark/gpu/gpu_loop_kernel.cpp b/ark/gpu/gpu_loop_kernel.cc → ark/gpu/gpu_loop_kernel.cpp
diff --git a/ark/gpu/gpu_manager.cc → ark/gpu/gpu_manager.cpp b/ark/gpu/gpu_manager.cc → ark/gpu/gpu_manager.cpp
diff --git a/ark/gpu/gpu_memory.cc → ark/gpu/gpu_memory.cpp b/ark/gpu/gpu_memory.cc → ark/gpu/gpu_memory.cpp
diff --git a/ark/gpu/gpu_stream.cc → ark/gpu/gpu_stream.cpp b/ark/gpu/gpu_stream.cc → ark/gpu/gpu_stream.cpp
diff --git a/ark/include/ark.h → ark/include/ark.hpp b/ark/include/ark.h → ark/include/ark.hpp
@@ -11,10 +11,10 @@
 #define ARK_PATCH 0
 #define ARK_VERSION (ARK_MAJOR * 10000 + ARK_MINOR * 100 + ARK_PATCH)
 
-#include "ark/dims.h"
-#include "ark/error.h"
-#include "ark/executor.h"
-#include "ark/model.h"
+#include "ark/dims.hpp"
+#include "ark/error.hpp"
+// #include "ark/executor.hpp"
+#include "ark/model.hpp"
 
 namespace ark {
 

diff --git a/ark/include/ark/dims.h → ark/include/ark/dims.hpp b/ark/include/ark/dims.h → ark/include/ark/dims.hpp
@@ -1,8 +1,8 @@
 // Copyright (c) Microsoft Corporation.
 // Licensed under the MIT license.
 
-#ifndef ARK_DIMS_H
-#define ARK_DIMS_H
+#ifndef ARK_DIMS_HPP
+#define ARK_DIMS_HPP
 
 #include <ostream>
 #include <string>
@@ -11,11 +11,12 @@
 namespace ark {
 
 // Data type for dimension.
-typedef long long int DimType;
+typedef int64_t DimType;
 
 // DIMS_LEN is the maximum number of dimensions of a tensor. If a tensor
 // has less than DIMS_LEN dimensions, the remaining dimensions will be NO_DIM.
-enum { DIMS_LEN = 4, NO_DIM = -1 };
+constexpr DimType NO_DIM = -1;
+constexpr DimType DIMS_LEN = 4;
 
 // Up-to-`DIMS_LEN`-dimensional vector.
 class Dims {
@@ -39,6 +40,8 @@ class Dims {
     int ndims() const;
     // Return a new Dims object with 4 valid dimensions by prepending 1s.
     Dims dims4() const;
+    // Return true if all valid dimensions are zero.
+    bool is_zeros() const;
     // Return true if the dimensions are empty.
     bool is_no_dim() const;
     // Return true if the dimensions are invalid.
@@ -64,4 +67,4 @@ class Dims {
 
 }  // namespace ark
 
-#endif  // ARK_DIMS_H
+#endif  // ARK_DIMS_HPP