reduce_node.h

Go to the documentation of this file.

 
#pragma once
 
#include <optional>
#include "sdfg/data_flow/library_nodes/math/tensor/tensor_node.h"
 
#include "sdfg/codegen/dispatchers/block_dispatcher.h"
#include "sdfg/serializer/json_serializer.h"
 
namespace sdfg {
namespace math {
namespace tensor {

class ReduceNode : public TensorNode {
protected:
    std::vector<symbolic::Expression> shape_; 
    std::vector<int64_t> axes_; 
    bool keepdims_; 
 
public:
    ReduceNode(
        size_t element_id,
        const DebugInfo& debug_info,
        const graph::Vertex vertex,
        data_flow::DataFlowGraph& parent,
        const data_flow::LibraryNodeCode& code,
        const std::vector<symbolic::Expression>& shape,
        const std::vector<int64_t>& axes,
        bool keepdims
    );
 
    static auto constexpr RESULT_PTR_IDX = 0;
    static auto constexpr X_INPUT_IDX = 1;
 
    void validate(const Function& function) const override;

    const std::vector<symbolic::Expression>& shape() const {
        return shape_;
    }

    const std::vector<int64_t>& axes() const {
        return axes_;
    }

    bool keepdims() const {
        return keepdims_;
    }
 
    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;

    passes::LibNodeExpander::ExpandOutcome expand(passes::LibNodeExpander::ExpandContext& context, Block& block) override;

    virtual bool expand_reduction(
        passes::LibNodeExpander::AccessNodeExpand& expansion,
        builder::StructuredSDFGBuilder& builder,
        structured_control_flow::Sequence& body,
        const types::Tensor& input_type,
        const types::Tensor& output_type,
        const data_flow::Subset& input_subset,
        const data_flow::Subset& output_subset
    ) = 0;

    virtual std::string identity(types::PrimitiveType primitive_type) const = 0;

    virtual std::optional<structured_control_flow::ReductionOperation> reduction_operation() const {
        return std::nullopt;
    }
 
    data_flow::PointerAccessType pointer_access_type(int input_idx) const override;
 
    std::string toStr() const override;
 
protected:
    virtual passes::LibNodeExpander::ExpandOutcome expand_inner(
        passes::LibNodeExpander::AccessNodeExpand& expansion,
        structured_control_flow::Block& block,
        const data_flow::Memlet* iedge_input,
        const data_flow::Memlet* iedge_result,
        const std::vector<symbolic::Expression>& output_shape,
        const std::vector<int64_t>& sorted_axes
    );
};
 
template<typename T>
class ReduceNodeSerializer : public serializer::LibraryNodeSerializer {
public:
    nlohmann::json serialize(const data_flow::LibraryNode& library_node) override {
        const ReduceNode& reduce_node = static_cast<const ReduceNode&>(library_node);
        nlohmann::json j;
 
        j["code"] = reduce_node.code().value();
 
        serializer::JSONSerializer serializer;
        j["shape"] = nlohmann::json::array();
        for (auto& dim : reduce_node.shape()) {
            j["shape"].push_back(serializer.expression(dim));
        }
 
        j["axes"] = reduce_node.axes();
        j["keepdims"] = reduce_node.keepdims();
 
        return j;
    }
 
    data_flow::LibraryNode& deserialize(
        const nlohmann::json& j, builder::StructuredSDFGBuilder& builder, structured_control_flow::Block& parent
    ) override {
        // Assertions for required fields
        assert(j.contains("element_id"));
        assert(j.contains("code"));
        assert(j.contains("debug_info"));
        assert(j.contains("shape"));
        assert(j.contains("axes"));
        assert(j.contains("keepdims"));
 
        auto code = j["code"].get<std::string>();
 
        std::vector<symbolic::Expression> shape;
        for (const auto& dim : j["shape"]) {
            shape.push_back(symbolic::parse(dim.get<std::string>()));
        }
 
        std::vector<int64_t> axes = j["axes"].get<std::vector<int64_t>>();
        bool keepdims = j["keepdims"].get<bool>();
 
        // Extract debug info using JSONSerializer
        sdfg::serializer::JSONSerializer serializer;
        DebugInfo debug_info = serializer.json_to_debug_info(j["debug_info"]);
 
        return static_cast<ReduceNode&>(builder.add_library_node<T>(parent, debug_info, shape, axes, keepdims));
    }
};
 
} // namespace tensor
} // namespace math
} // namespace sdfg