Merge pull request #56 from InfiniTensor/dev

Dev

Author: YdrMaster

CMakeLists.txt

@@ -25,6 +25,8 @@ if(USE_CUDA)
     endif()
     message(STATUS "CMAKE_CUDA_HOST_COMPILER set to " ${CMAKE_CUDA_HOST_COMPILER})
     message(STATUS "CMAKE_CUDA_ARCHITECTURES set to " ${CMAKE_CUDA_ARCHITECTURES})
+    message(STATUS "CUDA_PATH set to " ${CMAKE_CUDA_TOOLKIT_INCLUDE_DIRECTORIES})
+    add_compile_definitions(CUDA_INCLUDE_PATH="${CMAKE_CUDA_TOOLKIT_INCLUDE_DIRECTORIES}")
 endif()
 
 if(USE_KUNLUN)

src/00common/include/common/data_type.h

@@ -45,6 +45,7 @@ namespace refactor {
         bool isFloat() const noexcept;
         bool isSignedLarge() const noexcept;
         bool isSigned() const noexcept;
+        bool isUnsigned() const noexcept;
         bool isNumberic() const noexcept;
         bool isCpuNumberic() const noexcept;
         bool isBool() const noexcept;

src/00common/include/common/natural.h

@@ -7,7 +7,7 @@ namespace refactor {
 
     template<class t = size_t>
     class natural_t : public std::iterator<std::input_iterator_tag, t> {
-        size_t _i;
+        t _i;
 
     public:
         natural_t(t val) noexcept : _i(val) {}
@@ -33,7 +33,7 @@ namespace refactor {
 
     template<class t = size_t>
     class rev_natural_t : public std::iterator<std::input_iterator_tag, t> {
-        size_t _i;
+        t _i;
 
     public:
         rev_natural_t(t val) noexcept : _i(val - 1) {}

src/00common/include/common/range.h

@@ -16,7 +16,7 @@ namespace refactor {
 
         bool empty() const noexcept { return end_ == begin_; }
         size_t size() const noexcept { return end_ - begin_; }
-        t at(size_t i) const noexcept {
+        t at(size_t i) const {
             ASSERT(i < size(), "Index out of range");
             return operator[](i);
         }

src/00common/src/data_type.cc

@@ -82,6 +82,11 @@ namespace refactor {
             DT::I8, DT::I16, DT::I32, DT::I64};
         return set.contains(internal);
     }
+    bool DT::isUnsigned() const noexcept {
+        static const std::unordered_set<Enum> set{
+            DT::U8, DT::U16, DT::U32, DT::U64};
+        return set.contains(internal);
+    }
     bool DT::isNumberic() const noexcept {
         static const std::unordered_set<Enum> set{
             DT::F32, DT::U8, DT::I8, DT::U16, DT::I16,

src/01graph_topo/include/graph_topo/linked_graph.hpp

@@ -37,7 +37,7 @@ namespace refactor::graph_topo {
         Rc<Node> pushNode(TN, std::vector<Rc<Edge>>);
         void eraseNode(count_t);
         void eraseNode(Rc<Node>);
-        size_t cleanup(bool useless(TE const &) = nullptr);
+        size_t cleanup(bool useful(TE const &) = nullptr);
         bool sort();
     };
 
@@ -86,6 +86,7 @@ namespace refactor::graph_topo {
             node.disconnect(i);
         }
         for (auto const &out : node._outputs) {
+            ASSERT(out->_targets.empty(), "Output edge should not have targets");
             out->_source = nullptr;
         }
     }
@@ -180,23 +181,22 @@ namespace refactor::graph_topo {
         _nodes.erase(it);
     }
 
-    LINKED_GRAPH_FN cleanup(bool useless(TE const &))->size_t {
+    LINKED_GRAPH_FN cleanup(bool useful(TE const &))->size_t {
         std::unordered_set<Edge *> outputs;
         outputs.reserve(_outputs.size());
         std::transform(_outputs.begin(), _outputs.end(), std::inserter(outputs, outputs.end()), [](auto const &e) { return e.get(); });
-        auto useful = [&](Rc<Edge> const &e) {
-            return !e->_targets.empty() ||     // 还有节点连接到这个边
-                   outputs.contains(e.get()) ||// 这个边是全图输出
-                   !useless ||                 // 不需要其他判断
-                   !useless(e->_info);         // 这个边其他原因有用
+        auto useful_ = [&](Rc<Edge> const &e) {
+            return !e->_targets.empty() ||      // 还有节点连接到这个边
+                   outputs.contains(e.get()) || // 这个边是全图输出
+                   (useful && useful(e->_info));// 这个边其他原因有用
         };
 
         auto before = _nodes.size();
         while (true) {
             auto endit = std::remove_if(
                 _nodes.begin(), _nodes.end(),
                 [&, this](auto &n) {
-                    auto useless_ = std::none_of(n->_outputs.begin(), n->_outputs.end(), useful);
+                    auto useless_ = std::none_of(n->_outputs.begin(), n->_outputs.end(), useful_);
                     if (useless_) { _cleanupNode(*n); }
                     return useless_;
                 });

src/03runtime/include/runtime/stream.h

@@ -24,6 +24,7 @@ namespace refactor::runtime {
     struct Edge {
         Arc<hardware::Device::Blob> blob;
         size_t stackOffset;
+        std::string name;
     };
 
     class Stream {
@@ -39,6 +40,8 @@ namespace refactor::runtime {
                std::vector<Node>,
                std::vector<Edge>,
                decltype(_device));
+
+        decltype(_graph) const &graph() const noexcept { return _graph; }
         void setData(count_t, void const *, size_t);
         void setData(count_t, Arc<hardware::Device::Blob>);
         bool getData(count_t, void *, size_t) const;

src/04kernel/cuda/include/kernel/cuda/functions.cuh

@@ -3,6 +3,8 @@
 
 namespace refactor::kernel::cuda {
 
+    int currentDevice();
+
     void sync();
 
     void copyOut(void *dst, const void *src, size_t size);

src/04kernel/cuda/include/kernel/cuda/scatter_nd.cuh

@@ -0,0 +1,21 @@
+#ifndef KERNEL_CUDA_SCATTER_ND_CUH
+#define KERNEL_CUDA_SCATTER_ND_CUH
+
+#include "threads_distributer.cuh"
+
+namespace refactor::kernel::cuda {
+
+    void launchScatterND(
+        KernelLaunchParameters const &,
+        void const *data,
+        void const *indices,
+        void const *updates,
+        void *output,
+        unsigned int const *strides,
+        size_t rank,
+        unsigned int blockCount,
+        size_t blockSize);
+
+}// namespace refactor::kernel::cuda
+
+#endif// KERNEL_CUDA_SCATTER_ND_CUH

src/04kernel/cuda/src/functions.cu

@@ -1,19 +1,22 @@
 #include "kernel/cuda/functions.cuh"
+#include "macro.cuh"
 #include <cstdio>
 
 namespace refactor::kernel::cuda {
 
+    int currentDevice() {
+        int device;
+        CUDA_ASSERT(cudaGetDevice(&device));
+        return device;
+    }
+
     void sync() {
-        auto state = cudaDeviceSynchronize();
-        if (state != cudaSuccess) {
-            printf("cudaDeviceSynchronize failed: %s\n", cudaGetErrorString(state));
-            exit(1);
-        }
+        CUDA_ASSERT(cudaDeviceSynchronize());
     }
 
     void copyOut(void *dst, const void *src, size_t size) {
         sync();
-        cudaMemcpy(dst, src, size, cudaMemcpyDeviceToHost);
+        CUDA_ASSERT(cudaMemcpy(dst, src, size, cudaMemcpyDeviceToHost));
     }
 
 }// namespace refactor::kernel::cuda