tensor_node.h

Go to the documentation of this file.

 
#pragma once
 
#include "sdfg/data_flow/library_nodes/math/math_node.h"
#include "sdfg/types/type.h"
 
namespace sdfg {
namespace math {
namespace tensor {
class TensorLayout;
 
typedef types::PrimitiveType QuantizationType;
constexpr QuantizationType QUANTIZATION_MATCH_INPUTS = types::PrimitiveType::Void;

struct ContractionTypes {
    types::PrimitiveType input = types::PrimitiveType::Void;
    types::PrimitiveType accumulator = types::PrimitiveType::Void;
    types::PrimitiveType output = types::PrimitiveType::Void;

    static ContractionTypes uniform(types::PrimitiveType type) {
        return {type, types::PrimitiveType::Void, type};
    }
 
    bool operator==(const ContractionTypes& other) const = default;
 
    void serialize_to_json(nlohmann::json& j) const;

    static ContractionTypes deserialize_from_json(const nlohmann::json& j, const std::string& legacy_field);
};
 
types::PrimitiveType deserialize_quantization(
    const nlohmann::json& j,
    const std::string& field_name,
    types::PrimitiveType default_value = QUANTIZATION_MATCH_INPUTS
);

class TensorNode : public math::MathNode {
public:
    TensorNode(
        size_t element_id,
        const DebugInfo& debug_info,
        const graph::Vertex vertex,
        data_flow::DataFlowGraph& parent,
        const data_flow::LibraryNodeCode& code,
        const std::vector<std::string>& outputs,
        const std::vector<std::string>& inputs,
        const data_flow::ImplementationType& impl_type
    );

    virtual void verify_data_types(const data_flow::DataFlowGraph& graph) const;

    void validate(const Function& function) const override;

    types::PrimitiveType primitive_type(const data_flow::DataFlowGraph& graph) const;

    virtual bool supports_integer_types() const = 0;
 
protected:
    static data_flow::TaskletCode get_integer_minmax_tasklet(types::PrimitiveType prim_type, bool is_max);
 
    void validate_shape_matches(
        const std::vector<symbolic::Expression>& required_shape, const TensorLayout& layout, const std::string& name
    ) const;
};
 
} // namespace tensor
} // namespace math
} // namespace sdfg