28806128926

Committed 06 Jul 2026 04:16PM UTC coverage: 62.96%. First build

Build # 28806128926

Build Type

push

github

Committed by

web-flow

Commit Message

New LibNodeExpansion pass (#740)

* Switched expansion "pipeline" to new LibNodeExpansionPass that can recursively expand using a LibNodeExpander impl.
- removed access to analysis_manager from expansion methods, as those would currently not be appropriately maintained between nodes
+ New expansion API handles more of the boilerplate code (checking for standalone, creating boundary access nodes, removing old elements)
* Also migrated sdfg-json-to-c.cpp to not using the expansion pipeline anymore
* toStr() for ReduceNodes and giving PyStructuredSDFG access to the output_dir for additional dumping
~ DotVisualizer : Fix on nested sequences.
+ DotVisualizer: visualize empty sequences
* DotVisualizer default-enabled show block/loop ids
 * while MathNodes still have an expand-method, the new infrastructure is based upon "Expander" classes. The MathNodeExpander just redirects to the method for now
 ~ Broadcast node still used old ptr-output semantics
 * updated tensor & blas node expand to new expand API, that handles more of the boilerplate code (checking, removing old nodes, creating standalone-replacement nodes)
 - removed Transpose Node. Was unused and on old ptr-semantics
 + StructuredSDFGBuilder.add_sequence_at, add_for_at, add_map_at
 * switched StructuredSDFGBuilder internally to use ptr of Assignments to express using default assignments when needed
 * updated tests to use the new expand_single_math_node helper function, instead of the method directly
 + pass and expand-single helper functions are ready to use other expanders
 + EinsumExpansionPass is basically the legacy ExpansionPass, only restricted to EinsumNodes. It still needs to run in a pipeline to ensure finding multiple EinsumNodes per block. This is temporary, because Einsum is the only node that already supported splitting a block, which the new expansion disallows, as it should handle it itself when needed (but that part is not yet implemented)

Coverage Stats

594 of 962 new or added lines in 31 files covered. (61.75%)

40554 of 64412 relevant lines covered (62.96%)

963.11 hits per line

Source File
Press 'n' to go to next uncovered line, 'b' for previous

57.86

/sdfg/src/data_flow/library_nodes/math/tensor/elementwise_node.cpp

#include "sdfg/data_flow/library_nodes/math/tensor/elementwise_node.h"

#include "sdfg/analysis/analysis.h"
#include "sdfg/builder/structured_sdfg_builder.h"
#include "sdfg/data_flow/tasklet.h"
#include "sdfg/types/type.h"

namespace sdfg {
namespace math {
namespace tensor {

ElementWiseDataflowTensorNode::ElementWiseDataflowTensorNode(
    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::string& modified_tensor_conn,
    const std::vector<std::string>& tensor_inputs,
    QuantizationType quantization,
    const data_flow::ImplementationType& impl_type
)
    : TensorNode(
          element_id,
          debug_info,
          vertex,
          parent,
          code,
          {},
          build_input_conns(modified_tensor_conn, tensor_inputs),
          impl_type
      ),
      fixed_quantization_(quantization), shape_(shape) {}

std::vector<std::string> ElementWiseDataflowTensorNode::
    build_input_conns(const std::string& modified_tensor_conn, const std::vector<std::string>& tensor_inputs) {
    std::vector<std::string> input_conns;
    input_conns.reserve(1 + input_conns.size());
    input_conns.push_back(modified_tensor_conn);
    input_conns.insert(input_conns.end(), tensor_inputs.begin(), tensor_inputs.end());
    return input_conns;
}

types::PrimitiveType ElementWiseDataflowTensorNode::fixed_quantization() const { return fixed_quantization_; }

void ElementWiseDataflowTensorNode::set_fixed_quantization(const QuantizationType quant) {
    fixed_quantization_ = quant;
}

types::PrimitiveType ElementWiseDataflowTensorNode::quantization(const data_flow::DataFlowGraph& data_flow_graph
) const {
    if (fixed_quantization_ != QUANTIZATION_MATCH_INPUTS) {
        return fixed_quantization_;
    } else {
        return this->primitive_type(data_flow_graph);
    }
}

std::optional<types::PrimitiveType> ElementWiseDataflowTensorNode::uniform_quantization(const data_flow::DataFlowGraph&
                                                                                            data_flow_graph) const {
    if (fixed_quantization_ != QUANTIZATION_MATCH_INPUTS) {
        auto inferred = this->primitive_type(data_flow_graph);
        if (inferred == fixed_quantization_) {
            return fixed_quantization_;
        } else {
            return std::nullopt;
        }
    } else {
        return this->primitive_type(data_flow_graph);
    }
}

void ElementWiseDataflowTensorNode::validate_target_tensor(const data_flow::DataFlowGraph& graph) const {
    auto* target_ptr_edge = graph.in_edge_for_connector(*this, inputs_.at(0));
    auto& tensor_output = static_cast<const types::Tensor&>(target_ptr_edge->base_type());

    validate_shape_matches(shape_, tensor_output.layout(), "output tensor");
}

void ElementWiseDataflowTensorNode::validate_all_input_tensors(const data_flow::DataFlowGraph& graph) const {
    for (int i = 1; i < tensor_input_count(); ++i) {
        auto* iedge = graph.in_edge_for_connector(*this, inputs_.at(i));
        if (!iedge) {
            throw InvalidSDFGException(
                "On libNode #" + std::to_string(element_id()) + ": input " + inputs_.at(i) + " is not connected"
            );
        }
        if (iedge->base_type().type_id() == types::TypeID::Scalar) {
            continue;
        }
        auto& tensor_input = static_cast<const types::Tensor&>(iedge->base_type());
        // Case 1: Scalar input is allowed as secondary input
        if (tensor_input.is_scalar()) {
            continue;
        }

        // currently no arbitrary broadcast support! but could be added
        validate_shape_matches(shape_, tensor_input.layout(), "input " + inputs_.at(i));
    }
}

void ElementWiseDataflowTensorNode::validate_non_tensor_inputs(const data_flow::DataFlowGraph& graph) const {
    for (int i = tensor_input_count(); i < inputs_.size(); ++i) {
        auto* iedge = graph.in_edge_for_connector(*this, inputs_.at(i));
        if (!iedge) {
            if (i < mandatory_input_count()) {
                throw InvalidSDFGException(
                    "On libNode #" + std::to_string(element_id()) + ": input " + inputs_.at(i) + " is not connected"
                );
            } else {
                continue;
            }
        }
        if (iedge->base_type().type_id() != types::TypeID::Scalar) {
            throw InvalidSDFGException(
                "On libNode #" + std::to_string(element_id()) + ": input " + inputs_.at(i) + " is not scalar"
            );
        }
    }
}

void ElementWiseDataflowTensorNode::validate(const Function& function) const {
    TensorNode::validate(function);

    auto& graph = this->get_parent();

    validate_target_tensor(graph);

    validate_all_input_tensors(graph);

    validate_non_tensor_inputs(graph);
}

symbolic::SymbolSet ElementWiseDataflowTensorNode::symbols() const {
    symbolic::SymbolSet syms;
    for (const auto& dim : shape_) {
        for (auto& atom : symbolic::atoms(dim)) {
            syms.insert(atom);
        }
    }
    return syms;
}

void ElementWiseDataflowTensorNode::
    replace(const symbolic::Expression old_expression, const symbolic::Expression new_expression) {
    for (auto& dim : shape_) {
        dim = symbolic::subs(dim, old_expression, new_expression);
    }
}

void ElementWiseDataflowTensorNode::replace(const symbolic::ExpressionMapping& replacements) {
    for (auto& dim : shape_) {
        dim = symbolic::subs(dim, replacements);
    }
}

std::pair<structured_control_flow::Sequence*, std::vector<symbolic::Expression>> ElementWiseDataflowTensorNode::
    add_eltwise_scope(
        builder::StructuredSDFGBuilder& builder,
        const DebugInfo& scope_deb_info,
        Sequence& parent,
        const std::vector<symbolic::Expression>& shape
    ) {
    // Add maps
    data_flow::Subset new_subset;
    std::vector<symbolic::Expression> loop_vars;
    structured_control_flow::Sequence* last_scope = &parent;
    structured_control_flow::Map* last_map = nullptr;

    for (size_t i = 0; i < shape.size(); i++) {
        std::string indvar_str = builder.find_new_name("_i");
        builder.add_container(indvar_str, types::Scalar(types::PrimitiveType::UInt64));

        auto indvar = symbolic::symbol(indvar_str);
        auto init = symbolic::zero();
        auto update = symbolic::add(indvar, symbolic::one());
        auto condition = symbolic::Lt(indvar, shape.at(i));
        last_map = &builder.add_map(
            *last_scope,
            indvar,
            condition,
            init,
            update,
            structured_control_flow::ScheduleType_Sequential::create(),
            {},
            scope_deb_info
        );
        last_scope = &last_map->root();

        loop_vars.push_back(indvar);
    }
    return {last_scope, loop_vars};
}

std::unique_ptr<types::IType> ElementWiseDataflowTensorNode::access_type(const std::pair<
                                                                         types::PrimitiveType,
                                                                         const TensorLayout*>& pair) {
    if (pair.second) {
        return std::make_unique<types::Tensor>(pair.first, *pair.second);
    } else {
        return std::make_unique<types::Scalar>(pair.first);
    }
}

bool ElementWiseDataflowTensorNode::create_input(
    builder::StructuredSDFGBuilder& builder,
    structured_control_flow::Block& block,
    const data_flow::AccessNode& org_src,
    const std::pair<types::PrimitiveType, const TensorLayout*>& src_type,
    const ElementInput& needed_input,
    const std::vector<symbolic::Expression>& eltwise_subset,
    std::unordered_map<const data_flow::AccessNode*, data_flow::AccessNode*>& new_node_mapping
) {
    auto* new_consumer = needed_input.consumer;
    if (new_consumer) {
        if (src_type.first != needed_input.required_type) {
            throw InvalidSDFGException(
                "Input " + std::to_string(needed_input.input_conn_index) + " on node #" +
                std::to_string(new_consumer->element_id()) + " is required as " +
                types::primitive_type_to_string(needed_input.required_type) + " but provided as " +
                types::primitive_type_to_string(src_type.first)
            );
        }
        auto existing_input_it = new_node_mapping.find(&org_src);
        data_flow::AccessNode* input_node;
        std::vector<symbolic::Expression> empty_subset;
        const std::vector<symbolic::Expression>* memlet_subset;
        if (src_type.second && !src_type.second->is_scalar()) {
            memlet_subset = &eltwise_subset;
        } else {
            memlet_subset = &empty_subset;
        }
        auto new_type = access_type(src_type);
        if (existing_input_it != new_node_mapping.end()) {
            input_node = existing_input_it->second;
        } else {
            if (org_src.is_constant()) {
                types::Scalar const_type(src_type.first);
                input_node = &builder.add_constant(block, org_src.data(), const_type);
            } else {
                input_node = &builder.add_access(block, org_src.data());
            }
            new_node_mapping.emplace(&org_src, input_node);
        }

        builder.add_computational_memlet(
            block,
            *input_node,
            *new_consumer,
            new_consumer->input(needed_input.input_conn_index),
            *memlet_subset,
            *new_type
        );
        return true;
    } else {
        return false;
    }
}

void ElementWiseDataflowTensorNode::create_output(
    builder::StructuredSDFGBuilder& builder,
    structured_control_flow::Block& block,
    const data_flow::AccessNode& org_dst,
    const types::Tensor& dst_type,
    const ElementOutput& provided_output,
    const std::vector<symbolic::Expression>& eltwise_subset
) {
    auto* producer = provided_output.producer;
    if (dst_type.primitive_type() != provided_output.type) {
        throw InvalidSDFGException(
            "Output " + std::to_string(provided_output.output_conn_index) + " on node #" +
            std::to_string(producer->element_id()) + " is provided as " +
            types::primitive_type_to_string(provided_output.type) + " but required as " +
            types::primitive_type_to_string(dst_type.primitive_type())
        );
    }
    auto& output_node = builder.add_access(block, org_dst.data());
    builder.add_computational_memlet(
        block, *producer, producer->output(provided_output.output_conn_index), output_node, eltwise_subset, dst_type
    );
}

passes::LibNodeExpander::ExpandOutcome ElementWiseDataflowTensorNode::
    expand(passes::LibNodeExpander::ExpandContext& context, Block& block) {
    auto& dataflow = this->get_parent();

    auto* output_tensor_iedge = dataflow.in_edge_for_connector(*this, inputs_.at(0));
    if (!output_tensor_iedge) {
        return context.unable();
    }
    auto& target_tensor = static_cast<const types::Tensor&>(output_tensor_iedge->base_type());
    std::vector<const data_flow::Memlet*> iedges;
    std::vector<const data_flow::AccessNode*> inputs_sa;
    std::vector<std::pair<types::PrimitiveType, const TensorLayout*>> input_types;
    iedges.reserve(inputs_.size() - 1);
    using Dir = passes::LibNodeExpander::InputUse;
    std::vector<Dir> access_dirs{Dir::IndirectWrite};
    for (int i = 1; i < this->inputs_.size(); ++i) {
        auto* iedge = dataflow.in_edge_for_connector(*this, inputs_.at(i));
        if (!iedge) {
            if (i < mandatory_input_count()) {
                return context.unable();
            } else {
                continue;
            }
        }
        iedges.push_back(iedge);
        auto* input_sa = dataflow.find_standalone_entry(iedge);
        if (!input_sa) {
            return context.unable();
        }
        inputs_sa.push_back(input_sa);
        auto& input_type = iedge->base_type();
        if (input_type.type_id() == types::TypeID::Scalar) {
            access_dirs.push_back(Dir::Scalar);
            input_types.emplace_back(input_type.primitive_type(), nullptr);
        } else {
            access_dirs.push_back(Dir::IndirectRead);
            auto& tensor_type = static_cast<const types::Tensor&>(iedge->base_type());
            input_types.emplace_back(input_type.primitive_type(), &tensor_type.layout());
        }
    }

    auto* output_tensor_sa = dataflow.find_standalone_entry(output_tensor_iedge);
    if (!output_tensor_sa) {
        return context.unable();
    }

    auto standalone = context.replacement_requires_access_nodes(access_dirs);

    if (standalone) {
        auto& builder = standalone->builder();

        // Add new graph after the current block
        auto& new_sequence = standalone->replace_with_sequence();

        auto [eltw_scope, loop_vars] = add_eltwise_scope(builder, block.debug_info(), new_sequence, shape_);

        std::vector<tensor::ElementWiseDataflowTensorNode::ElementInput> eltwise_inputs;
        eltwise_inputs.reserve(inputs_.size() - 1);
        for (int i = 0; i < input_types.size(); ++i) {
            eltwise_inputs.push_back({.required_type = input_types.at(i).first});
        }

        auto& new_block = builder.add_block(*eltw_scope);

        auto produced_output =
            expand_operation_dataflow(builder, new_block, eltwise_inputs, target_tensor.primitive_type());
        if (!produced_output.producer) {
            return context.unable();
        }

        std::unordered_map<const data_flow::AccessNode*, data_flow::AccessNode*> new_node_mapping;

        // for all old input edge, remove old, create new
        for (int i = 0; i < iedges.size(); ++i) {
            create_input(
                builder, new_block, *inputs_sa.at(i), input_types.at(i), eltwise_inputs.at(i), loop_vars, new_node_mapping
            );
        }
        create_output(builder, new_block, *output_tensor_sa, target_tensor, produced_output, loop_vars);

        return standalone->successfully_expanded();
    } else {
        return context.unable();
    }
}

data_flow::PointerAccessType ElementWiseDataflowTensorNode::pointer_access_type(int input_idx) const {
    if (input_idx == 0) {
        return data_flow::PointerAccessMeta::create_full_write_only(symbolic::__nullptr__(), true);
    } else if (input_idx < tensor_input_count()) {
        return data_flow::PointerAccessMeta::create_read_only(symbolic::__nullptr__(), true);
    } else {
        return TensorNode::pointer_access_type(input_idx);
    }
}

data_flow::AccessNode& ElementWiseDataflowTensorNode::create_tmp_access_node(
    builder::StructuredSDFGBuilder& builder,
    structured_control_flow::Block& block,
    const std::string& prefix,
    const types::IType& type
) const {
    auto cont = builder.find_new_name(prefix);
    builder.add_container(cont, type);
    auto& output_node_add = builder.add_access(block, cont);
    return output_node_add;
}

nlohmann::json BaseElementWiseDataflowTensorNodeSerializer::serialize(const data_flow::LibraryNode& library_node) {
    const ElementWiseDataflowTensorNode& elem_node = static_cast<const ElementWiseDataflowTensorNode&>(library_node);
    nlohmann::json j;

    j["code"] = elem_node.code().value();

    serializer::JSONSerializer serializer;
    j["shape"] = nlohmann::json::array();
    for (auto& dim : elem_node.shape()) {
        j["shape"].push_back(serializer.expression(dim));
    }

    j["result_quant"] = elem_node.fixed_quantization();

    return j;
}

BaseElementWiseDataflowTensorNodeSerializer::BaseDeser BaseElementWiseDataflowTensorNodeSerializer::
    deserialize_base_values(const nlohmann::json& j) {
    assert(j.contains("element_id"));
    assert(j.contains("code"));
    assert(j.contains("debug_info"));

    std::vector<symbolic::Expression> shape;
    if (j.contains("shape")) {
        for (const auto& dim : j["shape"]) {
            shape.push_back(symbolic::parse(dim.get<std::string>()));
        }
    }

    serializer::JSONSerializer serializer;
    auto debug_info = serializer.json_to_debug_info(j["debug_info"]);
    return {
        .shape = shape,
        .quantization = deserialize_quantization(j, "result_quant", QUANTIZATION_MATCH_INPUTS),
        .debug_info = debug_info
    };
}

} // namespace tensor
} // namespace math
} // namespace sdfg

1	#include "sdfg/data_flow/library_nodes/math/tensor/elementwise_node.h"
2
3	#include "sdfg/analysis/analysis.h"
4	#include "sdfg/builder/structured_sdfg_builder.h"
5	#include "sdfg/data_flow/tasklet.h"
6	#include "sdfg/types/type.h"
7
8	namespace sdfg {
9	namespace math {
10	namespace tensor {
11
12	ElementWiseDataflowTensorNode::ElementWiseDataflowTensorNode(
13	size_t element_id,
14	const DebugInfo& debug_info,
15	const graph::Vertex vertex,
16	data_flow::DataFlowGraph& parent,
17	const data_flow::LibraryNodeCode& code,
18	const std::vector<symbolic::Expression>& shape,
19	const std::string& modified_tensor_conn,
20	const std::vector<std::string>& tensor_inputs,
21	QuantizationType quantization,
22	const data_flow::ImplementationType& impl_type
23	)
24	: TensorNode(	558✔
25	element_id,	558✔
26	debug_info,	558✔
27	vertex,	558✔
28	parent,	558✔
29	code,	558✔
30	{},	558✔
31	build_input_conns(modified_tensor_conn, tensor_inputs),	558✔
32	impl_type	558✔
33	),	558✔
34	fixed_quantization_(quantization), shape_(shape) {}	558✔
35
36	std::vector<std::string> ElementWiseDataflowTensorNode::
37	build_input_conns(const std::string& modified_tensor_conn, const std::vector<std::string>& tensor_inputs) {	558✔
38	std::vector<std::string> input_conns;	558✔
39	input_conns.reserve(1 + input_conns.size());	558✔
40	input_conns.push_back(modified_tensor_conn);	558✔
41	input_conns.insert(input_conns.end(), tensor_inputs.begin(), tensor_inputs.end());	558✔
42	return input_conns;	558✔
43	}	558✔
44
45	types::PrimitiveType ElementWiseDataflowTensorNode::fixed_quantization() const { return fixed_quantization_; }	×
46
47	void ElementWiseDataflowTensorNode::set_fixed_quantization(const QuantizationType quant) {	×
48	fixed_quantization_ = quant;	×
49	}	×
50
51	types::PrimitiveType ElementWiseDataflowTensorNode::quantization(const data_flow::DataFlowGraph& data_flow_graph
52	) const {	×
53	if (fixed_quantization_ != QUANTIZATION_MATCH_INPUTS) {	×
54	return fixed_quantization_;	×
55	} else {	×
56	return this->primitive_type(data_flow_graph);	×
57	}	×
58	}	×
59
60	std::optional<types::PrimitiveType> ElementWiseDataflowTensorNode::uniform_quantization(const data_flow::DataFlowGraph&
61	data_flow_graph) const {	×
62	if (fixed_quantization_ != QUANTIZATION_MATCH_INPUTS) {	×
63	auto inferred = this->primitive_type(data_flow_graph);	×
64	if (inferred == fixed_quantization_) {	×
65	return fixed_quantization_;	×
66	} else {	×
67	return std::nullopt;	×
68	}	×
69	} else {	×
70	return this->primitive_type(data_flow_graph);	×
71	}	×
72	}	×
73
74	void ElementWiseDataflowTensorNode::validate_target_tensor(const data_flow::DataFlowGraph& graph) const {	551✔
75	auto* target_ptr_edge = graph.in_edge_for_connector(*this, inputs_.at(0));	551✔
76	auto& tensor_output = static_cast<const types::Tensor&>(target_ptr_edge->base_type());	551✔
77
78	validate_shape_matches(shape_, tensor_output.layout(), "output tensor");	551✔
79	}	551✔
80
81	void ElementWiseDataflowTensorNode::validate_all_input_tensors(const data_flow::DataFlowGraph& graph) const {	579✔
82	for (int i = 1; i < tensor_input_count(); ++i) {	1,443✔
83	auto* iedge = graph.in_edge_for_connector(*this, inputs_.at(i));	864✔
84	if (!iedge) {	864✔
85	throw InvalidSDFGException(	×
86	"On libNode #" + std::to_string(element_id()) + ": input " + inputs_.at(i) + " is not connected"	×
87	);	×
88	}	×
89	if (iedge->base_type().type_id() == types::TypeID::Scalar) {	864✔
90	continue;	×
91	}	×
92	auto& tensor_input = static_cast<const types::Tensor&>(iedge->base_type());	864✔
93	// Case 1: Scalar input is allowed as secondary input
94	if (tensor_input.is_scalar()) {	864✔
95	continue;	1✔
96	}	1✔
97
98	// currently no arbitrary broadcast support! but could be added
99	validate_shape_matches(shape_, tensor_input.layout(), "input " + inputs_.at(i));	863✔
100	}	863✔
101	}	579✔
102
103	void ElementWiseDataflowTensorNode::validate_non_tensor_inputs(const data_flow::DataFlowGraph& graph) const {	323✔
104	for (int i = tensor_input_count(); i < inputs_.size(); ++i) {	323✔
105	auto* iedge = graph.in_edge_for_connector(*this, inputs_.at(i));	×
106	if (!iedge) {	×
107	if (i < mandatory_input_count()) {	×
108	throw InvalidSDFGException(	×
109	"On libNode #" + std::to_string(element_id()) + ": input " + inputs_.at(i) + " is not connected"	×
110	);	×
111	} else {	×
112	continue;	×
113	}	×
114	}	×
115	if (iedge->base_type().type_id() != types::TypeID::Scalar) {	×
116	throw InvalidSDFGException(	×
117	"On libNode #" + std::to_string(element_id()) + ": input " + inputs_.at(i) + " is not scalar"	×
118	);	×
119	}	×
120	}	×
121	}	323✔
122
123	void ElementWiseDataflowTensorNode::validate(const Function& function) const {	325✔
124	TensorNode::validate(function);	325✔
125
126	auto& graph = this->get_parent();	325✔
127
128	validate_target_tensor(graph);	325✔
129
130	validate_all_input_tensors(graph);	325✔
131
132	validate_non_tensor_inputs(graph);	325✔
133	}	325✔
134
135	symbolic::SymbolSet ElementWiseDataflowTensorNode::symbols() const {	17✔
136	symbolic::SymbolSet syms;	17✔
137	for (const auto& dim : shape_) {	68✔
138	for (auto& atom : symbolic::atoms(dim)) {	68✔
139	syms.insert(atom);	×
140	}	×
141	}	68✔
142	return syms;	17✔
143	}	17✔
144
145	void ElementWiseDataflowTensorNode::
146	replace(const symbolic::Expression old_expression, const symbolic::Expression new_expression) {	×
147	for (auto& dim : shape_) {	×
148	dim = symbolic::subs(dim, old_expression, new_expression);	×
149	}	×
150	}	×
151
152	void ElementWiseDataflowTensorNode::replace(const symbolic::ExpressionMapping& replacements) {	×
153	for (auto& dim : shape_) {	×
154	dim = symbolic::subs(dim, replacements);	×
155	}	×
156	}	×
157
158	std::pair<structured_control_flow::Sequence*, std::vector<symbolic::Expression>> ElementWiseDataflowTensorNode::
159	add_eltwise_scope(
160	builder::StructuredSDFGBuilder& builder,
161	const DebugInfo& scope_deb_info,
162	Sequence& parent,
163	const std::vector<symbolic::Expression>& shape
164	) {	534✔
165	// Add maps
166	data_flow::Subset new_subset;	534✔
167	std::vector<symbolic::Expression> loop_vars;	534✔
168	structured_control_flow::Sequence* last_scope = &parent;	534✔
169	structured_control_flow::Map* last_map = nullptr;	534✔
170
171	for (size_t i = 0; i < shape.size(); i++) {	1,814✔
172	std::string indvar_str = builder.find_new_name("_i");	1,280✔
173	builder.add_container(indvar_str, types::Scalar(types::PrimitiveType::UInt64));	1,280✔
174
175	auto indvar = symbolic::symbol(indvar_str);	1,280✔
176	auto init = symbolic::zero();	1,280✔
177	auto update = symbolic::add(indvar, symbolic::one());	1,280✔
178	auto condition = symbolic::Lt(indvar, shape.at(i));	1,280✔
179	last_map = &builder.add_map(	1,280✔
180	*last_scope,	1,280✔
181	indvar,	1,280✔
182	condition,	1,280✔
183	init,	1,280✔
184	update,	1,280✔
185	structured_control_flow::ScheduleType_Sequential::create(),	1,280✔
186	{},	1,280✔
187	scope_deb_info	1,280✔
188	);	1,280✔
189	last_scope = &last_map->root();	1,280✔
190
191	loop_vars.push_back(indvar);	1,280✔
192	}	1,280✔
193	return {last_scope, loop_vars};	534✔
194	}	534✔
195
196	std::unique_ptr<types::IType> ElementWiseDataflowTensorNode::access_type(const std::pair<
197	types::PrimitiveType,
198	const TensorLayout*>& pair) {	836✔
199	if (pair.second) {	836✔
200	return std::make_unique<types::Tensor>(pair.first, *pair.second);	836✔
201	} else {	836✔
202	return std::make_unique<types::Scalar>(pair.first);	×
203	}	×
204	}	836✔
205
206	bool ElementWiseDataflowTensorNode::create_input(
207	builder::StructuredSDFGBuilder& builder,
208	structured_control_flow::Block& block,
209	const data_flow::AccessNode& org_src,
210	const std::pair<types::PrimitiveType, const TensorLayout*>& src_type,
211	const ElementInput& needed_input,
212	const std::vector<symbolic::Expression>& eltwise_subset,
213	std::unordered_map<const data_flow::AccessNode, data_flow::AccessNode>& new_node_mapping
214	) {	836✔
215	auto* new_consumer = needed_input.consumer;	836✔
216	if (new_consumer) {	836✔
217	if (src_type.first != needed_input.required_type) {	836✔
218	throw InvalidSDFGException(	×
219	"Input " + std::to_string(needed_input.input_conn_index) + " on node #" +	×
220	std::to_string(new_consumer->element_id()) + " is required as " +	×
221	types::primitive_type_to_string(needed_input.required_type) + " but provided as " +	×
222	types::primitive_type_to_string(src_type.first)	×
223	);	×
224	}	×
225	auto existing_input_it = new_node_mapping.find(&org_src);	836✔
226	data_flow::AccessNode* input_node;	836✔
227	std::vector<symbolic::Expression> empty_subset;	836✔
228	const std::vector<symbolic::Expression>* memlet_subset;	836✔
229	if (src_type.second && !src_type.second->is_scalar()) {	836✔
230	memlet_subset = &eltwise_subset;	817✔
231	} else {	817✔
232	memlet_subset = &empty_subset;	19✔
233	}	19✔
234	auto new_type = access_type(src_type);	836✔
235	if (existing_input_it != new_node_mapping.end()) {	836✔
236	input_node = existing_input_it->second;	4✔
237	} else {	832✔
238	if (org_src.is_constant()) {	832✔
239	types::Scalar const_type(src_type.first);	×
240	input_node = &builder.add_constant(block, org_src.data(), const_type);	×
241	} else {	832✔
242	input_node = &builder.add_access(block, org_src.data());	832✔
243	}	832✔
244	new_node_mapping.emplace(&org_src, input_node);	832✔
245	}	832✔
246
247	builder.add_computational_memlet(	836✔
248	block,	836✔
249	*input_node,	836✔
250	*new_consumer,	836✔
251	new_consumer->input(needed_input.input_conn_index),	836✔
252	*memlet_subset,	836✔
253	*new_type	836✔
254	);	836✔
255	return true;	836✔
256	} else {	836✔
257	return false;	×
258	}	×
259	}	836✔
260
261	void ElementWiseDataflowTensorNode::create_output(
262	builder::StructuredSDFGBuilder& builder,
263	structured_control_flow::Block& block,
264	const data_flow::AccessNode& org_dst,
265	const types::Tensor& dst_type,
266	const ElementOutput& provided_output,
267	const std::vector<symbolic::Expression>& eltwise_subset
268	) {	534✔
269	auto* producer = provided_output.producer;	534✔
270	if (dst_type.primitive_type() != provided_output.type) {	534✔
271	throw InvalidSDFGException(	×
272	"Output " + std::to_string(provided_output.output_conn_index) + " on node #" +	×
273	std::to_string(producer->element_id()) + " is provided as " +	×
274	types::primitive_type_to_string(provided_output.type) + " but required as " +	×
275	types::primitive_type_to_string(dst_type.primitive_type())	×
276	);	×
277	}	×
278	auto& output_node = builder.add_access(block, org_dst.data());	534✔
279	builder.add_computational_memlet(	534✔
280	block, *producer, producer->output(provided_output.output_conn_index), output_node, eltwise_subset, dst_type	534✔
281	);	534✔
282	}	534✔
283
284	passes::LibNodeExpander::ExpandOutcome ElementWiseDataflowTensorNode::
285	expand(passes::LibNodeExpander::ExpandContext& context, Block& block) {	534✔
286	auto& dataflow = this->get_parent();	534✔
287
288	auto* output_tensor_iedge = dataflow.in_edge_for_connector(*this, inputs_.at(0));	534✔
289	if (!output_tensor_iedge) {	534✔
NEW 290	return context.unable();	×
291	}	×
292	auto& target_tensor = static_cast<const types::Tensor&>(output_tensor_iedge->base_type());	534✔
293	std::vector<const data_flow::Memlet*> iedges;	534✔
294	std::vector<const data_flow::AccessNode*> inputs_sa;	534✔
295	std::vector<std::pair<types::PrimitiveType, const TensorLayout*>> input_types;	534✔
296	iedges.reserve(inputs_.size() - 1);	534✔
297	using Dir = passes::LibNodeExpander::InputUse;	534✔
298	std::vector<Dir> access_dirs{Dir::IndirectWrite};	534✔
299	for (int i = 1; i < this->inputs_.size(); ++i) {	1,370✔
300	auto* iedge = dataflow.in_edge_for_connector(*this, inputs_.at(i));	836✔
301	if (!iedge) {	836✔
302	if (i < mandatory_input_count()) {	×
NEW 303	return context.unable();	×
304	} else {	×
305	continue;	×
306	}	×
307	}	×
308	iedges.push_back(iedge);	836✔
309	auto* input_sa = dataflow.find_standalone_entry(iedge);	836✔
310	if (!input_sa) {	836✔
NEW 311	return context.unable();	×
312	}	×
313	inputs_sa.push_back(input_sa);	836✔
314	auto& input_type = iedge->base_type();	836✔
315	if (input_type.type_id() == types::TypeID::Scalar) {	836✔
NEW 316	access_dirs.push_back(Dir::Scalar);	×
317	input_types.emplace_back(input_type.primitive_type(), nullptr);	×
318	} else {	836✔
319	access_dirs.push_back(Dir::IndirectRead);	836✔
320	auto& tensor_type = static_cast<const types::Tensor&>(iedge->base_type());	836✔
321	input_types.emplace_back(input_type.primitive_type(), &tensor_type.layout());	836✔
322	}	836✔
323	}	836✔
324
325	auto* output_tensor_sa = dataflow.find_standalone_entry(output_tensor_iedge);	534✔
326	if (!output_tensor_sa) {	534✔
NEW 327	return context.unable();	×
328	}	×
329
330	auto standalone = context.replacement_requires_access_nodes(access_dirs);	534✔
331
332	if (standalone) {	534✔
333	auto& builder = standalone->builder();	534✔
334
335	// Add new graph after the current block
336	auto& new_sequence = standalone->replace_with_sequence();	534✔
337
338	auto [eltw_scope, loop_vars] = add_eltwise_scope(builder, block.debug_info(), new_sequence, shape_);	534✔
339
340	std::vector<tensor::ElementWiseDataflowTensorNode::ElementInput> eltwise_inputs;	534✔
341	eltwise_inputs.reserve(inputs_.size() - 1);	534✔
342	for (int i = 0; i < input_types.size(); ++i) {	1,370✔
343	eltwise_inputs.push_back({.required_type = input_types.at(i).first});	836✔
344	}	836✔
345
346	auto& new_block = builder.add_block(*eltw_scope);	534✔
347
348	auto produced_output =	534✔
349	expand_operation_dataflow(builder, new_block, eltwise_inputs, target_tensor.primitive_type());	534✔
350	if (!produced_output.producer) {	534✔
NEW 351	return context.unable();	×
NEW 352	}	×
353
354	std::unordered_map<const data_flow::AccessNode, data_flow::AccessNode> new_node_mapping;	534✔
355
356	// for all old input edge, remove old, create new
357	for (int i = 0; i < iedges.size(); ++i) {	1,370✔
358	create_input(	836✔
359	builder, new_block, *inputs_sa.at(i), input_types.at(i), eltwise_inputs.at(i), loop_vars, new_node_mapping	836✔
360	);	836✔
361	}	836✔
362	create_output(builder, new_block, *output_tensor_sa, target_tensor, produced_output, loop_vars);	534✔
363
364	return standalone->successfully_expanded();	534✔
365	} else {	534✔
NEW 366	return context.unable();	×
NEW 367	}	×
368	}	534✔
369
370	data_flow::PointerAccessType ElementWiseDataflowTensorNode::pointer_access_type(int input_idx) const {	×
371	if (input_idx == 0) {	×
372	return data_flow::PointerAccessMeta::create_full_write_only(symbolic::__nullptr__(), true);	×
373	} else if (input_idx < tensor_input_count()) {	×
374	return data_flow::PointerAccessMeta::create_read_only(symbolic::__nullptr__(), true);	×
375	} else {	×
376	return TensorNode::pointer_access_type(input_idx);	×
377	}	×
378	}	×
379
380	data_flow::AccessNode& ElementWiseDataflowTensorNode::create_tmp_access_node(
381	builder::StructuredSDFGBuilder& builder,
382	structured_control_flow::Block& block,
383	const std::string& prefix,
384	const types::IType& type
385	) const {	12✔
386	auto cont = builder.find_new_name(prefix);	12✔
387	builder.add_container(cont, type);	12✔
388	auto& output_node_add = builder.add_access(block, cont);	12✔
389	return output_node_add;	12✔
390	}	12✔
391
392	nlohmann::json BaseElementWiseDataflowTensorNodeSerializer::serialize(const data_flow::LibraryNode& library_node) {	×
393	const ElementWiseDataflowTensorNode& elem_node = static_cast<const ElementWiseDataflowTensorNode&>(library_node);	×
394	nlohmann::json j;	×
395
396	j["code"] = elem_node.code().value();	×
397
398	serializer::JSONSerializer serializer;	×
399	j["shape"] = nlohmann::json::array();	×
400	for (auto& dim : elem_node.shape()) {	×
401	j["shape"].push_back(serializer.expression(dim));	×
402	}	×
403
404	j["result_quant"] = elem_node.fixed_quantization();	×
405
406	return j;	×
407	}	×
408
409	BaseElementWiseDataflowTensorNodeSerializer::BaseDeser BaseElementWiseDataflowTensorNodeSerializer::
410	deserialize_base_values(const nlohmann::json& j) {	×
411	assert(j.contains("element_id"));	×
412	assert(j.contains("code"));	×
413	assert(j.contains("debug_info"));	×
414
415	std::vector<symbolic::Expression> shape;	×
416	if (j.contains("shape")) {	×
417	for (const auto& dim : j["shape"]) {	×
418	shape.push_back(symbolic::parse(dim.get<std::string>()));	×
419	}	×
420	}	×
421
422	serializer::JSONSerializer serializer;	×
423	auto debug_info = serializer.json_to_debug_info(j["debug_info"]);	×
424	return {	×
425	.shape = shape,	×
426	.quantization = deserialize_quantization(j, "result_quant", QUANTIZATION_MATCH_INPUTS),	×
427	.debug_info = debug_info	×
428	};	×
429	}	×
430
431	} // namespace tensor
432	} // namespace math
433	} // namespace sdfg

daisytuner / docc / 28806128926

Source File Press 'n' to go to next uncovered line, 'b' for previous

Source File
Press 'n' to go to next uncovered line, 'b' for previous