issue/497 - support tensor from/to files

Author: PanZezhong1725

include/infinicore/ops.hpp

@@ -1,5 +1,8 @@
 #pragma once
 
-#include "op/matmul.hpp"
-#include "op/ones.hpp"
-#include "op/rearrange.hpp"
+#include "ops/add.hpp"
+#include "ops/attention.hpp"
+#include "ops/matmul.hpp"
+#include "ops/ones.hpp"
+#include "ops/rearrange.hpp"
+#include "ops/rms_norm.hpp"

include/infinicore/ops/add.hpp

@@ -0,0 +1,17 @@
+#pragma once
+
+#include "../device.hpp"
+#include "common/op.hpp"
+
+namespace infinicore::op {
+class Add {
+public:
+    using schema = void (*)(Tensor, Tensor, Tensor);
+    static void execute(Tensor c, Tensor a, Tensor b);
+    static common::OpDispatcher<schema> &dispatcher();
+};
+
+Tensor add(Tensor a, Tensor b);
+void add_(Tensor c, Tensor a, Tensor b);
+Tensor operator+(Tensor a, Tensor b);
+} // namespace infinicore::op

include/infinicore/ops/attention.hpp

@@ -0,0 +1,16 @@
+#pragma once
+
+#include "../device.hpp"
+#include "common/op.hpp"
+
+namespace infinicore::op {
+class Attention {
+public:
+    using schema = void (*)(Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, size_t);
+    static void execute(Tensor out, Tensor q, Tensor k, Tensor v, Tensor k_cache, Tensor v_cache, size_t pos);
+    static common::OpDispatcher<schema> &dispatcher();
+};
+
+Tensor attention(Tensor q, Tensor k, Tensor v, Tensor k_cache, Tensor v_cache, size_t pos);
+void attention_(Tensor out, Tensor q, Tensor k, Tensor v, Tensor k_cache, Tensor v_cache, size_t pos);
+} // namespace infinicore::op

include/infinicore/op/common/cache.hpp renamed to include/infinicore/ops/common/cache.hpp

@@ -8,19 +8,19 @@ namespace infinicore::op::common {
 template <typename Fn>
 class OpDispatcher {
 public:
-    void registerDevice(Device::Type device_type, Fn fn, bool override_existing=true) {
-        if (table_[(size_t)device_type] == nullptr || override_existing){
+    void registerDevice(Device::Type device_type, Fn fn, bool override_existing = true) {
+        if (table_[(size_t)device_type] == nullptr || override_existing) {
             table_[(size_t)device_type] = fn;
         }
     }
 
-    void registerDevice(std::initializer_list<Device::Type> device_types, Fn fn, bool override_existing=true) {
+    void registerDevice(std::initializer_list<Device::Type> device_types, Fn fn, bool override_existing = true) {
         for (auto device_type : device_types) {
             registerDevice(device_type, fn, override_existing);
         }
     }
 
-    void registerAll(Fn fn, bool override_existing=true) {
+    void registerAll(Fn fn, bool override_existing = true) {
         for (size_t device_type = 0; device_type < static_cast<size_t>(Device::Type::COUNT); ++device_type) {
             registerDevice((Device::Type)device_type, fn, override_existing);
         }

include/infinicore/op/common/dispatcher.hpp renamed to include/infinicore/ops/common/dispatcher.hpp

@@ -0,0 +1,16 @@
+#pragma once
+
+#include "../device.hpp"
+#include "common/op.hpp"
+
+namespace infinicore::op {
+class RMSNorm {
+public:
+    using schema = void (*)(Tensor, Tensor, Tensor, float);
+    static void execute(Tensor y, Tensor x, Tensor weight, float epsilon = 1e-5f);
+    static common::OpDispatcher<schema> &dispatcher();
+};
+
+Tensor rms_norm(Tensor x, Tensor weight, float epsilon = 1e-5f);
+void rms_norm_(Tensor y, Tensor x, Tensor weight, float epsilon = 1e-5f);
+} // namespace infinicore::op