axiom-compute/atlas.cpp at master · privateDuck/axiom-compute · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
// atlas.cpp : This file contains the 'main' function. Program execution begins and ends there.
//

#include <format>
#include <iostream>
#include "../exec/ExecutionEngine.hpp"
#include "src/data/DataIngester.hpp"

using NType = double;

struct AddNode {
	std::vector<NType> compute(const std::vector<NType>& inputs, std::vector<NType>& outputs, size_t num_outs, int) {
		NType sum = 0;
		for (const auto& val : inputs) sum += val;
		outputs[0] = sum;
	}

	void init(NodeId i) {

	}
};

struct MultNode {
	std::vector<NType> compute(const std::vector<NType>& inputs, std::vector<NType>& outputs, size_t num_outs, int) {
		NType product = 1;
		for (const auto& val : inputs) product *= val;
		outputs[0] = product;
	}
	void init(NodeId i) {

	}
};

struct InputNode {
	NType value;
	void setValue(NType v) { value = v; };

	std::vector<NType> compute(const std::vector<NType>& inputs, std::vector<NType>& outputs, size_t num_outs, int) {
		outputs[0] = inputs[0];
	}
	void init(NodeId i) {

	}
};

using NodeVar = std::variant<AddNode, MultNode, InputNode>;

int main() {
	DAGEngine<NodeVar, double, int> graph;
	auto inA = graph.addNode(InputNode{10.0});
	auto inB = graph.addNode(InputNode{10.0});
	auto add = graph.addNode(AddNode{});
	auto mul = graph.addNode(MultNode{});

	graph.addEdge(inA, 0, add, 0);
	graph.addEdge(inB, 0, add, 1);
	graph.addEdge(add, 0, mul, 0);
	graph.addEdge(inB, 0, mul, 1);

	graph.compile();
	graph.execute(1);
	auto result = graph.getOutput(mul, 0);
	std::cout << result << std::endl;
}