MetaTensor/MetaTensorBuilder.h at master · LogDS/MetaTensor · 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
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
//
// Created by gyankos on 12/10/26.
//

#ifndef METATENSOR_METATENSORBUILDER_H
#define METATENSOR_METATENSORBUILDER_H
#include <optional>

#include "MetaTensorBaseVisitor.h"
#include <yaucl/structures/eq_class.h>
#include <yaucl/structures/LinkedHashMultimap.h>

#include "logds/language/MetaTensorLanguage.h"
#include "yaucl/structures/gamma_context/HierarchicalListView.h"


namespace logds {
    namespace metatensor {
        class MetaTensorBuilder : public  MetaTensorBaseVisitor {
            equivalence_classv2<std::string> var_eq_class_dimensions;
            std::set<std::string> registered_dimensions;
            std::stack<std::unique_ptr<yaucl::gamma_context::HierarchicalListView<std::string, AnyType>>> Gamma;
            std::stack<std::vector<Statement>> pre_visit_declarations;
            uint64_t fresh_var_id = 0;

            std::string generateFreshId() {
                auto val = std::string("_")+std::to_string(fresh_var_id++);
                return val;
            }

            void postProcessFunctions(std::vector<Function> & vector) {
                // TODO: for instance, apply the equality class for each instance of the
            }
        public:

            std::any visitProgram(MetaTensorParser::ProgramContext *context) override;
            std::any visitDim_decl(MetaTensorParser::Dim_declContext *context) override;
            std::any visitFf(MetaTensorParser::FfContext *context) override;
            std::any visitTt(MetaTensorParser::TtContext *context) override;
            std::any visitAll(MetaTensorParser::AllContext *context) override;
            std::any visitAny(MetaTensorParser::AnyContext *context) override;
            std::any visitLambda_float(MetaTensorParser::Lambda_floatContext *context) override;
            std::any visitBand(MetaTensorParser::BandContext *context) override;
            std::any visitBor(MetaTensorParser::BorContext *context) override;
            std::any visitB_xor(MetaTensorParser::B_xorContext *context) override;
            std::any visitB_neg(MetaTensorParser::B_negContext *context) override;
            std::any visitBpar(MetaTensorParser::BparContext *context) override {
                if (!context) {
                    return std::nullopt;
                }
                return visit(context->bool_());
            }

            std::any visitLambda_id(MetaTensorParser::Lambda_idContext *context) override {
                if (!context) {
                    return std::nullopt;
                }
                auto var_name = context->ID()->getText();
                assert(pre_visit_declarations.size()>0); // I should have at least one element within the stack
                assert(Gamma.size()>0);
                auto type = Gamma.top()->find_node(var_name);
                assert(type); // The context should recognize the variable as having a specific type
                if (std::holds_alternative<TLTypes>(type->value)) {
                    Float val;
                    if (std::get<0>(type->value) == TLTypes::T_BOOL) {
                        auto tmp_var_id = generateFreshId();
                        Gamma.top()->emplace(tmp_var_id, TLTypes::T_FLOAT64);
                        pre_visit_declarations.top().emplace_back(tmp_var_id, var_name, TLTypes::T_FLOAT64);
                        val = tmp_var_id;
                    } else if (std::get<0>(type->value) != TLTypes::T_GEN_TENSOR) {
                        val = var_name;
                    } else {
                        assert(false);
                    }
                    return val;
                }
            }

            std::any visitBoolT(MetaTensorParser::BoolTContext *context) override { return TLTypes::T_BOOL;    }
            std::any visitFloatT(MetaTensorParser::FloatTContext *ctx) override   { return TLTypes::T_FLOAT32; }
            std::any visitDoubleT(MetaTensorParser::DoubleTContext *ctx) override { return TLTypes::T_FLOAT64; }
            std::any visitInt32T(MetaTensorParser::Int32TContext *ctx) override   { return TLTypes::T_INT32;   }
            std::any visitInt64T(MetaTensorParser::Int64TContext *ctx) override   { return TLTypes::T_INT64;   }
            std::any visitUint32(MetaTensorParser::Uint32Context *ctx) override   { return TLTypes::T_UINT32;  };
        };
    } // metatensor
} // logds

#endif //METATENSOR_METATENSORBUILDER_H