Merge pull request #3 from wsttiger/master

Initial PR for RTGLib (CPU operators and tests)

src/include/rtg/operators.hpp

 #ifndef RTG_GUARD_OPERATORS_HPP
 #define RTG_GUARD_OPERATORS_HPP
 
+#include <array>
 #include <rtg/operation.hpp>
 #include <rtg/stringutils.hpp>
 #include <rtg/streamutils.hpp>
@@ -55,6 +56,13 @@ struct check_shapes
         return *this;
     }
 
+    const check_shapes& same_ndims() const
+    {
+        if(!this->same([](const shape& s) { return s.lens().size(); }))
+            RTG_THROW("Dimensions do not match");
+        return *this;
+    }
+
     template <class F>
     bool same(F f) const
     {
@@ -93,7 +101,7 @@ struct convolution
     std::string name() const { return "convolution"; }
     shape compute_shape(std::vector<shape> inputs) const
     {
-        check_shapes{inputs}.has(2).same_type().same_dims().only_dims(4);
+        check_shapes{inputs}.has(2).same_type().same_ndims().only_dims(4);
 
         const shape& input   = inputs.at(0);
         const shape& weights = inputs.at(1);
@@ -253,6 +261,141 @@ struct reshape
     }
 };
 
+struct gemm
+{
+    std::string name() const { return "gemm"; }
+    shape compute_shape(std::vector<shape> inputs) const
+    {
+        check_shapes{inputs}.has(2).same_type();
+        const shape& a = inputs.at(0);
+        const shape& b = inputs.at(1);
+        auto t         = a.type();
+
+        if(a.lens()[1] != b.lens()[0])
+            RTG_THROW("Inner dimensions do not match");
+        return {t, {a.lens()[0], b.lens()[1]}};
+    }
+
+    argument compute(shape, std::vector<argument>) const { RTG_THROW("not computable"); }
+
+    friend std::ostream& operator<<(std::ostream& os, const gemm& op)
+    {
+        os << op.name() << "[";
+        os << "]";
+        return os;
+    }
+};
+
+struct unary
+{
+    shape compute_shape(std::vector<shape> inputs) const
+    {
+        check_shapes{inputs}.has(1);
+        return inputs.at(0);
+    }
+    argument compute(shape, std::vector<argument>) const { RTG_THROW("not computable"); }
+};
+
+struct identity : unary
+{
+    std::string name() const { return "identity"; }
+};
+
+struct abs : unary
+{
+    std::string name() const { return "abs"; }
+};
+
+struct exp : unary
+{
+    std::string name() const { return "exp"; }
+};
+
+struct sin : unary
+{
+    std::string name() const { return "sin"; }
+};
+
+struct cos : unary
+{
+    std::string name() const { return "cos"; }
+};
+
+struct tan : unary
+{
+    std::string name() const { return "tan"; }
+};
+
+struct asin : unary
+{
+    std::string name() const { return "asin"; }
+};
+
+struct acos : unary
+{
+    std::string name() const { return "acos"; }
+};
+
+struct atan : unary
+{
+    std::string name() const { return "atan"; }
+};
+
+struct softmax : unary
+{
+    std::string name() const { return "softmax"; }
+};
+
+struct tanh : unary
+{
+    std::string name() const { return "tanh"; }
+};
+
+struct sigmoid : unary
+{
+    std::string name() const { return "sigmoid"; }
+};
+
+struct neg : unary
+{
+    std::string name() const { return "neg"; }
+};
+
+struct flatten
+{
+    std::string name() const { return "flatten"; }
+};
+
+struct binary
+{
+    shape compute_shape(std::vector<shape> inputs) const
+    {
+        // TODO(wsttiger@gmail.com) Check this for numpy-style broadcasting operations
+        check_shapes{inputs}.has(2).same_type().same_dims();
+        return inputs.at(0);
+    }
+};
+
+struct add : binary
+{
+    std::string name() const { return "add"; }
+};
+
+struct sub : binary
+{
+    std::string name() const { return "sub"; }
+};
+
+struct mul : binary
+{
+    std::string name() const { return "mul"; }
+};
+
+struct div : binary
+{
+    std::string name() const { return "div"; }
+};
+
 struct outline
 {
     shape s;

src/include/rtg/tensor_view.hpp

@@ -11,6 +11,7 @@ namespace rtg {
 template <class T>
 struct tensor_view
 {
+    using value_type = T;
     tensor_view() : m_data(nullptr) {}
     tensor_view(shape s, T* d) : m_data(d), m_shape(s) {}
 

src/onnx/read_onnx.cpp

@@ -6,6 +6,7 @@
 #include <fstream>
 #include <unordered_map>
 #include <functional>
+#include <array>
 
 #include <rtg/fallthrough.hpp>
 #include <rtg/program.hpp>

src/targets/cpu/cpu_target.cpp

@@ -7,6 +7,12 @@
 namespace rtg {
 namespace cpu {
 
+template <typename T>
+T zero(const T&)
+{
+    return T(0);
+}
+
 struct cpu_convolution
 {
     convolution op;
@@ -48,25 +54,285 @@ struct cpu_convolution
     }
 };
 
-struct relu
+struct cpu_reshape
+{
+    reshape op;
+    std::string name() const { return "cpu::reshape"; }
+    shape compute_shape(std::vector<shape> inputs) const { return op.compute_shape(inputs); }
+
+    argument compute(shape output_shape, std::vector<argument> args) const
+    {
+        return {output_shape, std::move(args.front().data)};
+    }
+};
+
+struct cpu_gemm
+{
+    gemm op;
+    std::string name() const { return "cpu::gemm"; }
+    shape compute_shape(std::vector<shape> inputs) const { return op.compute_shape(inputs); }
+
+    argument compute(shape output_shape, std::vector<argument> args) const
+    {
+        argument result{output_shape};
+        visit_all(result, args[0], args[1])([&](auto cmat, auto amat, auto bmat) {
+            auto m = amat.get_shape().lens()[0];
+            auto n = bmat.get_shape().lens()[1];
+            auto k = bmat.get_shape().lens()[0];
+
+            auto a = amat.data();
+            auto b = bmat.data();
+            auto c = cmat.data();
+            for(int ii = 0; ii < m; ii++)
+            {
+                for(int jj = 0; jj < n; jj++)
+                {
+                    c[ii * n + jj] = 0;
+                }
+            }
+            for(int ii = 0; ii < m; ii++)
+            {
+                for(int kk = 0; kk < k; kk++)
+                {
+                    auto aik  = a[ii * k + kk];
+                    auto* bkj = &b[kk * n];
+                    auto* cij = &c[ii * n];
+                    for(int jj = 0; jj < n; jj++, cij++, bkj++)
+                    {
+                        *cij += aik * (*bkj);
+                    }
+                }
+            }
+        });
+        return result;
+    }
+};
+
+struct identity_op
+{
+    std::string name() const { return "cpu::identity"; }
+    auto fcn() const
+    {
+        return [](auto x) { return x; };
+    }
+};
+
+struct abs_op
+{
+    std::string name() const { return "cpu::abs"; }
+    auto fcn() const
+    {
+        return [](auto x) { return std::abs(x); };
+    }
+};
+
+struct exp_op
+{
+    std::string name() const { return "cpu::exp"; }
+    auto fcn() const
+    {
+        return [](auto x) { return std::exp(x); };
+    }
+};
+
+struct sin_op
+{
+    std::string name() const { return "cpu::sin"; }
+    auto fcn() const
+    {
+        return [](auto x) { return std::sin(x); };
+    }
+};
+
+struct cos_op
+{
+    std::string name() const { return "cpu::cos"; }
+    auto fcn() const
+    {
+        return [](auto x) { return std::cos(x); };
+    }
+};
+
+struct tan_op
+{
+    std::string name() const { return "cpu::tan"; }
+    auto fcn() const
+    {
+        return [](auto x) { return std::tan(x); };
+    }
+};
+
+struct asin_op
+{
+    std::string name() const { return "cpu::asin"; }
+    auto fcn() const
+    {
+        return [](auto x) { return std::asin(x); };
+    }
+};
+
+struct acos_op
+{
+    std::string name() const { return "cpu::acos"; }
+    auto fcn() const
+    {
+        return [](auto x) { return std::acos(x); };
+    }
+};
+
+struct atan_op
+{
+    std::string name() const { return "cpu::atan"; }
+    auto fcn() const
+    {
+        return [](auto x) { return std::atan(x); };
+    }
+};
+
+struct tanh_op
+{
+    std::string name() const { return "cpu::tanh"; }
+    auto fcn() const
+    {
+        return [](auto x) { return std::tanh(x); };
+    }
+};
+
+struct sigmoid_op
+{
+    std::string name() const { return "cpu::sigmoid"; }
+    auto fcn() const
+    {
+        return [](auto x) { return 1.f / (1.f + std::exp(-x)); };
+    }
+};
+
+struct neg_op
+{
+    std::string name() const { return "cpu::neg"; }
+    auto fcn() const
+    {
+        return [](auto x) { return -x; };
+    }
+};
+
+struct relu_op
 {
     std::string name() const { return "cpu::relu"; }
-    shape compute_shape(std::vector<shape> inputs) const { return inputs.front(); }
+    auto fcn() const
+    {
+        return [](auto x) { return x > 0 ? x : 0; };
+    }
+};
 
+template <typename Op>
+struct cpu_unary
+{
+    Op op;
+    std::string name() const { return op.name(); }
+    shape compute_shape(std::vector<shape> inputs) const { return inputs.front(); }
     argument compute(shape output_shape, std::vector<argument> args) const
     {
         argument result{output_shape};
         result.visit([&](auto output) {
             args[0].visit([&](auto input) {
-                std::transform(input.begin(), input.end(), output.begin(), [](auto x) {
-                    return x > 0 ? x : 0;
-                });
+                std::transform(input.begin(), input.end(), output.begin(), op.fcn());
             });
         });
         return result;
     }
 };
 
+struct softmax2d
+{
+    std::string name() const { return "cpu::softmax2d"; }
+    shape compute_shape(std::vector<shape> inputs) const { return inputs.front(); }
+    argument compute(shape output_shape, std::vector<argument> args) const
+    {
+        argument result{output_shape};
+        visit_all(result, args[0])([&](auto output, auto input) {
+            using value_type = typename decltype(input)::value_type;
+            auto nb          = input.get_shape().lens()[0];
+            auto nc          = input.get_shape().lens()[1];
+            auto nh          = input.get_shape().lens()[2];
+            auto nw          = input.get_shape().lens()[3];
+            dfor(nb, nh, nw)([&](std::size_t b, std::size_t i, std::size_t j) {
+                value_type cmax = std::numeric_limits<value_type>::lowest();
+                for(int c = 0; c < nc; c++)
+                {
+                    cmax = std::max(cmax, input(b, c, i, j));
+                }
+                for(int c = 0; c < nc; c++)
+                {
+                    output(b, c, i, j) = std::exp(input(b, c, i, j) - cmax);
+                }
+                value_type sum = value_type(0);
+                for(int c = 0; c < nc; c++)
+                {
+                    sum += output(b, c, i, j);
+                }
+                for(int c = 0; c < nc; c++)
+                {
+                    output(b, c, i, j) = output(b, c, i, j) / sum;
+                }
+            });
+        });
+        return result;
+    }
+};
+
+struct add_op
+{
+    std::string name() const { return "add"; }
+    auto fcn() const
+    {
+        return [](auto x, auto y) { return x + y; };
+    }
+};
+
+struct sub_op
+{
+    std::string name() const { return "sub"; }
+    auto fcn() const
+    {
+        return [](auto x, auto y) { return x - y; };
+    }
+};
+
+struct mul_op
+{
+    std::string name() const { return "mul"; }
+    auto fcn() const
+    {
+        return [](auto x, auto y) { return x * y; };
+    }
+};
+
+struct div_op
+{
+    std::string name() const { return "div"; }
+    auto fcn() const
+    {
+        return [](auto x, auto y) { return x / y; };
+    }
+};
+
+template <typename Op>
+struct cpu_binary
+{
+    Op op;
+    std::string name() const { return op.name(); }
+    shape compute_shape(std::vector<shape> inputs) const { return inputs.front(); }
+    argument compute(shape output_shape, std::vector<argument> args) const
+    {
+        argument result{output_shape};
+        visit_all(result, args[0], args[1])([&](auto output, auto input1, auto input2) {
+            std::transform(input1.begin(), input1.end(), input2.begin(), output.begin(), op.fcn());
+        });
+        return result;
+    }
+};
+
 struct cpu_apply
 {
     program* prog;
@@ -79,10 +345,54 @@ struct cpu_apply
             {
                 apply_convolution(it);
             }
+            else if(it->op.name() == "gemm")
+            {
+                apply_gemm(it);
+            }
+            else if(it->op.name() == "reshape")
+            {
+                apply_reshape(it);
+            }
             else if(it->op.name() == "activation")
             {
                 apply_activation(it);
             }
+            else if(it->op.name() == "identity")
+            {
+                apply_identity(it);
+            }
+            else if(it->op.name() == "softmax")
+            {
+                apply_softmax(it);
+            }
+            else if(it->op.name() == "tanh")
+            {
+                apply_tanh(it);
+            }
+            else if(it->op.name() == "sigmoid")
+            {
+                apply_sigmoid(it);
+            }
+            else if(it->op.name() == "exp")
+            {
+                apply_exp(it);
+            }
+            else if(it->op.name() == "neg")
+            {
+                apply_neg(it);
+            }
+            else if(it->op.name() == "sin")
+            {
+                apply_sin(it);
+            }
+            else if(it->op.name() == "cos")
+            {
+                apply_cos(it);
+            }
+            else if(it->op.name() == "tan")
+            {
+                apply_tan(it);
+            }
         }
     }
 
@@ -92,11 +402,68 @@ struct cpu_apply
         prog->replace_instruction(ins, cpu_convolution{op}, ins->arguments);
     }
 
+    void apply_gemm(instruction_ref ins)
+    {
+        auto&& op = any_cast<gemm>(ins->op);
+        prog->replace_instruction(ins, cpu_gemm{op}, ins->arguments);
+    }
+
+    void apply_reshape(instruction_ref ins)
+    {
+        auto&& op = any_cast<reshape>(ins->op);
+        prog->replace_instruction(ins, cpu_reshape{op}, ins->arguments);
+    }
+
     void apply_activation(instruction_ref ins)
     {
         auto&& op = any_cast<activation>(ins->op);
         if(op.mode == "relu")
-            prog->replace_instruction(ins, relu{}, ins->arguments);
+            prog->replace_instruction(ins, cpu_unary<relu_op>{}, ins->arguments);
+    }
+
+    void apply_identity(instruction_ref ins)
+    {
+        prog->replace_instruction(ins, cpu_unary<identity_op>{}, ins->arguments);
+    }
+
+    void apply_softmax(instruction_ref ins)
+    {
+        prog->replace_instruction(ins, softmax2d{}, ins->arguments);
+    }
+
+    void apply_tanh(instruction_ref ins)
+    {
+        prog->replace_instruction(ins, cpu_unary<tanh_op>{}, ins->arguments);
+    }
+
+    void apply_sigmoid(instruction_ref ins)
+    {
+        prog->replace_instruction(ins, cpu_unary<sigmoid_op>{}, ins->arguments);
+    }
+
+    void apply_exp(instruction_ref ins)
+    {
+        prog->replace_instruction(ins, cpu_unary<exp_op>{}, ins->arguments);
+    }
+
+    void apply_neg(instruction_ref ins)
+    {
+        prog->replace_instruction(ins, cpu_unary<neg_op>{}, ins->arguments);
+    }
+
+    void apply_sin(instruction_ref ins)
+    {
+        prog->replace_instruction(ins, cpu_unary<sin_op>{}, ins->arguments);
+    }
+
+    void apply_cos(instruction_ref ins)
+    {
+        prog->replace_instruction(ins, cpu_unary<cos_op>{}, ins->arguments);
+    }
+
+    void apply_tan(instruction_ref ins)
+    {
+        prog->replace_instruction(ins, cpu_unary<tan_op>{}, ins->arguments);
     }
 };