temp

Signed-off-by: YdrMaster <ydrml@hotmail.com>
InfiniTensor · Nov 17, 2023 · 06fe900 · 06fe900
1 parent db27623
commit 06fe900
Show file tree

Hide file tree

Showing 5 changed files with 40 additions and 12 deletions.
diff --git a/src/03runtime/include/runtime/stream.h b/src/03runtime/include/runtime/stream.h
@@ -17,7 +17,7 @@ namespace refactor::runtime {
     struct Address {
         std::variant<size_t, mem_manager::SharedForeignBlob> value;
 
-        void *operator()(void *stack);
+        void *operator()(void *stack) const;
 
         bool isBlob() const noexcept;
         bool isOffset() const noexcept;
@@ -49,6 +49,7 @@ namespace refactor::runtime {
         auto prepare() -> std::vector<count_t>;
         void run();
         auto bench(void (*sync)()) -> std::vector<std::chrono::nanoseconds>;
+        auto trace(void record(void const **, void **));
     };
 
 }// namespace refactor::runtime

diff --git a/src/03runtime/src/stream.cc b/src/03runtime/src/stream.cc
@@ -6,7 +6,7 @@ namespace refactor::runtime {
 
     void emptyRoutine(runtime::Resources &, void const **, void **) {}
 
-    void *Address::operator()(void *stack) {
+    void *Address::operator()(void *stack) const {
         if (isBlob()) {
             auto blob = std::get<mem_manager::SharedForeignBlob>(value);
             return blob ? (void *) *blob : nullptr;
@@ -82,15 +82,25 @@ namespace refactor::runtime {
         return unknownInputs;
     }
 
+    template<class I, class O>
+    std::pair<void const **, void **> collectAddress(
+        void *stack,
+        std::vector<Address> const &addresses,
+        std::vector<void *> &buffer,
+        I i, O o) {
+        buffer.resize(i.size() + o.size());
+        auto inputs = buffer.data(),
+             outputs = std::transform(i.begin(), i.end(), inputs, [&](auto i) { return addresses[i](stack); });
+        /* alignnnnn*/ std::transform(o.begin(), o.end(), outputs, [&](auto i) { return addresses[i](stack); });
+        return {const_cast<void const **>(inputs), outputs};
+    }
+
     void Stream::run() {
         auto map = [this](auto i) { return _internal.edges[i](*_stack); };
         std::vector<void *> buffer(16);
         for (auto const [nodeIdx, i, o] : _internal.topology) {
-            buffer.resize(i.size() + o.size());
-            auto inputs_ = buffer.data(),
-                 outputs_ = std::transform(i.begin(), i.end(), inputs_, map);
-            /* alignment */ std::transform(o.begin(), o.end(), outputs_, map);
-            _internal.nodes[nodeIdx](_resources, const_cast<void const **>(inputs_), outputs_);
+            auto [inputs, outputs] = collectAddress(*_stack, _internal.edges, buffer, i, o);
+            _internal.nodes[nodeIdx](_resources, inputs, outputs);
         }
     }
 
@@ -99,17 +109,24 @@ namespace refactor::runtime {
         std::vector<void *> buffer(16);
         std::vector<std::chrono::nanoseconds> ans(_internal.nodes.size());
         for (auto const [nodeIdx, i, o] : _internal.topology) {
-            buffer.resize(i.size() + o.size());
-            auto inputs_ = buffer.data(),
-                 outputs_ = std::transform(i.begin(), i.end(), inputs_, map);
-            /* alignment */ std::transform(o.begin(), o.end(), outputs_, map);
+            auto [inputs, outputs] = collectAddress(*_stack, _internal.edges, buffer, i, o);
             auto t0 = std::chrono::high_resolution_clock::now();
-            _internal.nodes[nodeIdx](_resources, const_cast<void const **>(inputs_), outputs_);
+            _internal.nodes[nodeIdx](_resources, inputs, outputs);
             if (sync) { sync(); }
             auto t1 = std::chrono::high_resolution_clock::now();
             ans[nodeIdx] = t1 - t0;
         }
         return ans;
     }
 
+    auto Stream::trace(void record(void const **, void **)) {
+        auto map = [this](auto i) { return _internal.edges[i](*_stack); };
+        std::vector<void *> buffer(16);
+        for (auto const [nodeIdx, i, o] : _internal.topology) {
+            auto [inputs, outputs] = collectAddress(*_stack, _internal.edges, buffer, i, o);
+            _internal.nodes[nodeIdx](_resources, inputs, outputs);
+            record(inputs, outputs);
+        }
+    }
+
 }// namespace refactor::runtime
diff --git a/src/09python_ffi/src/executor.cc b/src/09python_ffi/src/executor.cc
@@ -1,4 +1,5 @@
 #include "executor.h"
+#include <filesystem>
 
 #ifdef USE_CUDA
 #include "kernel/cuda/functions.cuh"
@@ -47,6 +48,13 @@ namespace refactor::python_ffi {
         }
     }
 
+    void trace(std::string path_) {
+        namespace fs = std::filesystem;
+        auto path = fs::path(std::move(path_));
+        fs::create_directories(path);
+        ASSERT(fs::is_directory(path), "Failed to create \"{}\"", path.c_str());
+    }
+
     void Executor::debugInfo() const noexcept {
         auto const &nodes = _graph.internal().contiguous().nodes;
         for (auto i : range0_(nodes.size())) {

diff --git a/src/09python_ffi/src/executor.h b/src/09python_ffi/src/executor.h
@@ -18,6 +18,7 @@ namespace refactor::python_ffi {
         auto prepare() -> std::vector<count_t>;
         void run();
         void bench(bool sync);
+        void trace(std::string);
         void debugInfo() const noexcept;
     };
 

diff --git a/src/09python_ffi/src/main.cpp b/src/09python_ffi/src/main.cpp
@@ -40,6 +40,7 @@ namespace refactor::python_ffi {
             .def("prepare"         , &Executor::prepare          , return_::move      )
             .def("run"             , &Executor::run              , return_::automatic )
             .def("bench"           , &Executor::bench            , return_::automatic )
+            .def("trace"           , &Executor::trace            , return_::automatic )
             .def("dbg"             , &Executor::debugInfo        , return_::automatic );
 
         // clang-format on