spatial_tensor_node.h

Go to the documentation of this file.

#pragma once
 
#include <cassert>
#include <cstddef>
#include <vector>
#include "sdfg/data_flow/library_nodes/math/tensor/tensor_node.h"
#include "sdfg/serializer/json_serializer.h"
 
#include "sdfg/symbolic/symbolic.h"
 
namespace sdfg {
namespace math {
namespace tensor {

class SpatialTensorNode : public TensorNode {
protected:
    QuantizationType fixed_quantization_;
    std::vector<symbolic::Expression> shape_; 
    std::vector<symbolic::Expression> kernel_shape_; 
    std::vector<symbolic::Expression> strides_; 
    std::vector<symbolic::Expression> pads_; 
    std::vector<symbolic::Expression> dilations_; 
 
public:
    SpatialTensorNode(
        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,
        QuantizationType quantization,
        const std::vector<symbolic::Expression>& shape,
        const std::vector<symbolic::Expression>& kernel_shape,
        const std::vector<symbolic::Expression>& strides,
        const std::vector<symbolic::Expression>& pads,
        const std::vector<symbolic::Expression>& dilations
    );
 
    const std::vector<symbolic::Expression>& shape() const {
        return shape_;
    }
    const std::vector<symbolic::Expression>& kernel_shape() const {
        return kernel_shape_;
    }
    const std::vector<symbolic::Expression>& strides() const {
        return strides_;
    }
    const std::vector<symbolic::Expression>& pads() const {
        return pads_;
    }
    const std::vector<symbolic::Expression>& dilations() const {
        return dilations_;
    }
 
    QuantizationType quantization() const {
        return quantization(get_parent());
    }

    QuantizationType quantization(const data_flow::DataFlowGraph& dataflow) const;

    std::optional<QuantizationType> uniform_quantization(const data_flow::DataFlowGraph& dataflow) const;

    QuantizationType fixed_quantization() const;
 
    void set_fixed_quantization(const QuantizationType quant);
 
    symbolic::SymbolSet symbols() const override;
 
    void replace(const symbolic::Expression old_expression, const symbolic::Expression new_expression) override;
 
    void replace(const symbolic::ExpressionMapping& replacements) override;
    size_t num_spatial_dims() const;

    symbolic::Expression output_spatial_dim(size_t i) const;

    symbolic::Expression output_spatial_volume() const;

    symbolic::Expression kernel_volume() const;
 
    std::basic_ostream<char>& operator<<(std::basic_ostream<char>& os) const;
};
 
class SpatialTensorNodeBaseSerializer : public serializer::LibraryNodeSerializer {
public:
    struct BaseDeser {
        std::vector<symbolic::Expression> shape;
        std::vector<symbolic::Expression> kernel_shape;
        std::vector<symbolic::Expression> strides;
        std::vector<symbolic::Expression> pads;
        std::vector<symbolic::Expression> dilations;
        QuantizationType quantization;
        DebugInfo debug_info;
    };
 
    void fill_base_values(const SpatialTensorNode& node, nlohmann::json& j);
 
    BaseDeser deserialize_base_values(const nlohmann::json& j);
};
 
} // namespace tensor
} // namespace math
} // namespace sdfg