29936384397

Committed 22 Jul 2026 04:04PM UTC coverage: 64.072% (+0.3%) from 63.782%

Build # 29936384397

Build Type

Pull #868

github

Committed by

web-flow

Commit Message

Merge 67eee0fe8 into d833d292d

Pull Request Pull Request #868: Add tensor concatenation to PyTorch frontend (and MLIR frontend)

Coverage Stats

352 of 454 new or added lines in 19 files covered. (77.53%)

4 existing lines in 2 files now uncovered.

42634 of 66541 relevant lines covered (64.07%)

739.99 hits per line

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

79.13

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

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

#include <cstddef>
#include <memory>
#include <sstream>
#include <string>
#include <vector>

#include <nlohmann/json_fwd.hpp>

#include "sdfg/builder/structured_sdfg_builder.h"
#include "sdfg/data_flow/data_flow_graph.h"
#include "sdfg/data_flow/data_flow_node.h"
#include "sdfg/data_flow/library_node.h"
#include "sdfg/data_flow/library_nodes/math/tensor/tensor_layout.h"
#include "sdfg/data_flow/library_nodes/math/tensor/tensor_node.h"
#include "sdfg/data_flow/memlet.h"
#include "sdfg/data_flow/tasklet.h"
#include "sdfg/element.h"
#include "sdfg/exceptions.h"
#include "sdfg/function.h"
#include "sdfg/graph/graph.h"
#include "sdfg/passes/expansion/lib_node_expander.h"
#include "sdfg/serializer/json_serializer.h"
#include "sdfg/structured_control_flow/block.h"
#include "sdfg/structured_control_flow/sequence.h"
#include "sdfg/structured_control_flow/structured_loop.h"
#include "sdfg/symbolic/symbolic.h"
#include "sdfg/types/scalar.h"
#include "sdfg/types/tensor.h"
#include "sdfg/types/type.h"

namespace sdfg {
namespace math {
namespace tensor {

ConcatNode::ConcatNode(
    size_t element_id,
    const DebugInfo& debug_info,
    const graph::Vertex vertex,
    data_flow::DataFlowGraph& parent,
    const std::string& result,
    const TensorLayout& result_layout,
    const std::vector<std::string>& tensors,
    const std::vector<TensorLayout>& tensor_layouts,
    long long dim,
    const data_flow::ImplementationType& impl_type
)
    : TensorNode(element_id, debug_info, vertex, parent, LibraryNodeType_TensorConcat, {}, tensors, impl_type),
      result_layout_(result_layout), tensor_layouts_(tensor_layouts), dim_(dim) {
    this->inputs_.push_back(result);
}

const std::string& ConcatNode::result() const { return this->inputs_.back(); }

const TensorLayout& ConcatNode::result_layout() const { return this->result_layout_; }

std::vector<std::string> ConcatNode::tensors() const {
    return std::vector<std::string>(this->inputs_.begin(), this->inputs_.end() - 1);
}

const std::vector<TensorLayout>& ConcatNode::tensor_layouts() const { return this->tensor_layouts_; }

long long ConcatNode::dim() const { return this->dim_; }

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

    auto& graph = this->get_parent();

    if (graph.out_degree(*this) != 0) {
        throw InvalidSDFGException("ConcatNode: Expected no outputs but got: " + std::to_string(graph.out_degree(*this)));
    }

    // Check that the tensor layouts match with the tensor types on the edges
    const data_flow::Memlet* result_iedge = graph.in_edge_for_connector(*this, this->result());
    if (!result_iedge) {
        throw InvalidSDFGException("ConcatNode: No memlet connected at connector: " + this->result());
    }
    if (result_iedge->base_type().type_id() != types::TypeID::Tensor) {
        throw InvalidSDFGException(
            "ConcatNode: Expected tensor type at connector '" + this->result() +
            "' but got: " + result_iedge->base_type().print()
        );
    }
    const types::Tensor& result_tensor = static_cast<const types::Tensor&>(result_iedge->base_type());
    if (result_tensor.layout() != this->result_layout_) {
        throw InvalidSDFGException(
            "ConcatNode: Provided tensor layout does not match the tensor type on the edge for connector '" +
            this->result() + "': " + result_tensor.layout().toStr() + " != " + this->result_layout_.toStr()
        );
    }
    for (long long i = 0; i < this->tensor_layouts_.size(); i++) {
        const std::string& tensor = this->inputs_[i];
        const data_flow::Memlet* tensor_iedge = graph.in_edge_for_connector(*this, tensor);
        if (!tensor_iedge) {
            throw InvalidSDFGException("ConcatNode: No memlet connected at connector: " + tensor);
        }
        if (tensor_iedge->base_type().type_id() != types::TypeID::Tensor) {
            throw InvalidSDFGException(
                "ConcatNode: Expected tensor type at connector '" + tensor +
                "' but got: " + tensor_iedge->base_type().print()
            );
        }
        const types::Tensor& tensor_tensor = static_cast<const types::Tensor&>(tensor_iedge->base_type());
        if (tensor_tensor.layout() != this->tensor_layouts_[i]) {
            throw InvalidSDFGException(
                "ConcatNode: Provided tensor layout does not match the tensor type on the edge for connector '" +
                tensor + "': " + tensor_tensor.layout().toStr() + " != " + this->tensor_layouts_[i].toStr()
            );
        }
    }

    // Check that the cat dimension is < the tensor layout dimensions
    if (this->dim_ < 0 || this->dim_ >= this->result_layout_.dims()) {
        throw InvalidSDFGException(
            "ConcatNode: Cat dimension must be in [0, " + std::to_string(this->result_layout_.dims() - 1) +
            "] but got: " + std::to_string(this->dim_)
        );
    }

    // Check that all tensor layout have the same shape except in the cat dimension
    symbolic::Expression cat_dim = symbolic::zero();
    for (const TensorLayout& tensor_layout : this->tensor_layouts_) {
        int dims = tensor_layout.dims();
        if (dims != this->result_layout_.dims()) {
            throw InvalidSDFGException(
                "ConcatNode: Tensor layouts have different dimensions: " + tensor_layout.toStr() +
                " != " + this->result_layout_.toStr()
            );
        }
        for (long long i = 0; i < dims; i++) {
            if (i == this->dim_) {
                cat_dim = symbolic::add(cat_dim, tensor_layout.get_dim(i));
                continue; // Skip the cat dimension
            }
            if (!symbolic::eq(tensor_layout.get_dim(i), this->result_layout_.get_dim(i))) {
                throw InvalidSDFGException(
                    "ConcatNode: Shapes do not match at position " + std::to_string(i) + ": " + tensor_layout.toStr() +
                    " != " + this->result_layout_.toStr()
                );
            }
        }
    }
    if (!symbolic::eq(cat_dim, this->result_layout_.get_dim(this->dim_))) {
        throw InvalidSDFGException(
            "ConcatNode: Cat dimension does not match with the sum of tensor dimensions: " + cat_dim->__str__() +
            " != " + this->result_layout_.get_dim(this->dim_)->__str__()
        );
    }
}

bool ConcatNode::supports_integer_types() const { return true; }

using Dir = passes::LibNodeExpander::InputUse;

passes::LibNodeExpander::ExpandOutcome ConcatNode::
    expand(passes::LibNodeExpander::ExpandContext& context, structured_control_flow::Block& block) {
    std::vector<Dir> access_dirs(this->inputs_.size(), Dir::IndirectRead);
    access_dirs.back() = Dir::IndirectWrite;
    auto standalone = context.replacement_requires_access_nodes(access_dirs);

    if (!standalone) {
        return context.unable();
    }

    auto& builder = standalone->builder();

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

    auto& dfg = this->get_parent();

    types::Scalar indvar_type(types::PrimitiveType::UInt64);
    structured_control_flow::Sequence* current_seq = &new_sequence;
    data_flow::Subset subset;
    subset.reserve(this->result_layout_.dims());
    for (auto dim : this->result_layout_.shape()) {
        auto indvar_container = builder.find_new_name("_i");
        builder.add_container(indvar_container, indvar_type);
        auto indvar = symbolic::symbol(indvar_container);
        subset.push_back(indvar);
        auto& map = builder.add_map(
            *current_seq,
            indvar,
            symbolic::Lt(indvar, dim),
            symbolic::zero(),
            symbolic::add(indvar, symbolic::one()),
            structured_control_flow::ScheduleType_Sequential::create(),
            {},
            this->debug_info_
        );
        current_seq = &map.root();
    }

    auto& if_else = builder.add_if_else(*current_seq, {}, this->debug_info_);
    auto dim_indvar = subset.at(this->dim_);
    const auto* iedge_result = dfg.in_edge_for_connector(*this, this->result());
    if (!iedge_result) {
        throw InvalidSDFGException("ConcatNode: Cannot get in edge for connector: " + this->result());
    }

    symbolic::Expression offset = symbolic::zero();
    for (size_t i = 0; i < this->tensor_layouts_.size(); i++) {
        auto new_offset = symbolic::add(offset, this->tensor_layouts_[i].get_dim(this->dim_));
        auto condition = symbolic::And(symbolic::Ge(dim_indvar, offset), symbolic::Lt(dim_indvar, new_offset));
        auto& case_seq = builder.add_case(if_else, condition, this->debug_info_);

        data_flow::Subset offset_subset(subset);
        offset_subset[this->dim_] = symbolic::sub(dim_indvar, offset);

        auto& block = builder.add_block(case_seq, {}, this->debug_info_);
        auto& tensor_access = standalone->add_scalar_input_access(block, i);
        auto& result_access = standalone->add_scalar_input_access(block, this->tensor_layouts_.size());
        auto& tasklet = builder.add_tasklet(block, data_flow::TaskletCode::assign, "_out", {"_in"}, this->debug_info_);

        const auto* iedge_tensor = dfg.in_edge_for_connector(*this, this->inputs_[i]);
        if (!iedge_tensor) {
            throw InvalidSDFGException("ConcatNode: Cannot get in edge for connector: " + this->inputs_[i]);
        }
        builder.add_computational_memlet(
            block, tensor_access, tasklet, "_in", offset_subset, iedge_tensor->base_type(), iedge_tensor->debug_info()
        );
        builder.add_computational_memlet(
            block, tasklet, "_out", result_access, subset, iedge_result->base_type(), iedge_result->debug_info()
        );

        offset = new_offset;
    }

    return standalone->successfully_expanded();
}

std::string ConcatNode::toStr() const {
    std::stringstream stream;
    stream << "ConcatNode(" << this->result() << ": [";
    for (long long i = 0; i < this->result_layout_.dims(); i++) {
        if (i > 0) {
            stream << ",";
        }
        stream << this->result_layout_.get_dim(i)->__str__();
    }
    stream << "], {";
    for (long long i = 0; i < this->tensor_layouts_.size(); i++) {
        if (i > 0) {
            stream << ", ";
        }
        stream << this->inputs_[i] << ": [";
        for (long long j = 0; j < this->tensor_layouts_[i].dims(); j++) {
            if (j > 0) {
                stream << ",";
            }
            stream << this->tensor_layouts_[i].get_dim(j)->__str__();
        }
        stream << "]";
    }
    stream << "}, dim: " << this->dim_ << ")";
    return stream.str();
}

symbolic::SymbolSet ConcatNode::symbols() const {
    symbolic::SymbolSet syms;
    this->result_layout_.collect_symbols(syms);
    for (const TensorLayout& tensor_layout : this->tensor_layouts_) {
        tensor_layout.collect_symbols(syms);
    }
    return syms;
}

symbolic::Expression ConcatNode::flop() const { return symbolic::zero(); }

std::unique_ptr<data_flow::DataFlowNode> ConcatNode::
    clone(size_t element_id, const graph::Vertex vertex, data_flow::DataFlowGraph& parent) const {
    return std::make_unique<ConcatNode>(
        element_id,
        this->debug_info_,
        vertex,
        parent,
        this->result(),
        this->result_layout_,
        this->tensors(),
        this->tensor_layouts_,
        this->dim_,
        this->implementation_type_
    );
}

void ConcatNode::replace(const symbolic::Expression old_expression, const symbolic::Expression new_expression) {
    this->result_layout_.replace_symbols(old_expression, new_expression);
    for (TensorLayout& tensor_layout : this->tensor_layouts_) {
        tensor_layout.replace_symbols(old_expression, new_expression);
    }
}

void ConcatNode::replace(const symbolic::ExpressionMapping& replacements) {
    this->result_layout_.replace_symbols(replacements);
    for (TensorLayout& tensor_layout : this->tensor_layouts_) {
        tensor_layout.replace_symbols(replacements);
    }
}

nlohmann::json ConcatNodeSerializer::serialize(const data_flow::LibraryNode& library_node) {
    const ConcatNode& concat_node = static_cast<const ConcatNode&>(library_node);
    nlohmann::json j;

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

    j["result"] = concat_node.result();
    concat_node.result_layout().serialize_to_json(j["result_layout"]);

    j["tensors"] = nlohmann::json::array();
    for (std::string& tensor : concat_node.tensors()) {
        j["tensors"].push_back(tensor);
    }
    j["tensor_layouts"] = nlohmann::json::array();
    for (const TensorLayout& tensor_layout : concat_node.tensor_layouts()) {
        nlohmann::json tensor_layout_j;
        tensor_layout.serialize_to_json(tensor_layout_j);
        j["tensor_layouts"].push_back(tensor_layout_j);
    }

    j["dim"] = concat_node.dim();

    return j;
}

data_flow::LibraryNode& ConcatNodeSerializer::deserialize(
    const nlohmann::json& j, builder::StructuredSDFGBuilder& builder, structured_control_flow::Block& parent
) {
    assert(j.contains("element_id"));
    assert(j.contains("code"));
    assert(j.contains("result"));
    assert(j.contains("result_layout"));
    assert(j.contains("tensors"));
    assert(j.contains("tensor_layouts"));
    assert(j.contains("dim"));
    assert(j.contains("debug_info"));

    std::string result = j.at("result").get<std::string>();
    TensorLayout result_layout = TensorLayout::deserialize_from_json(j.at("result_layout"));

    std::vector<std::string> tensors = j.at("tensors").get<std::vector<std::string>>();
    std::vector<TensorLayout> tensor_layouts;
    tensor_layouts.reserve(j.at("tensor_layouts").size());
    for (long long i = 0; i < j.at("tensor_layouts").size(); i++) {
        tensor_layouts.push_back(TensorLayout::deserialize_from_json(j.at("tensor_layouts").at(i)));
    }

    long long dim = j.at("dim").get<long long>();

    sdfg::serializer::JSONSerializer serializer;
    DebugInfo debug_info = serializer.json_to_debug_info(j.at("debug_info"));

    return builder.add_library_node<ConcatNode>(parent, debug_info, result, result_layout, tensors, tensor_layouts, dim);
}

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

1	#include "sdfg/data_flow/library_nodes/math/tensor/concat_node.h"
2
3	#include <cstddef>
4	#include <memory>
5	#include <sstream>
6	#include <string>
7	#include <vector>
8
9	#include <nlohmann/json_fwd.hpp>
10
11	#include "sdfg/builder/structured_sdfg_builder.h"
12	#include "sdfg/data_flow/data_flow_graph.h"
13	#include "sdfg/data_flow/data_flow_node.h"
14	#include "sdfg/data_flow/library_node.h"
15	#include "sdfg/data_flow/library_nodes/math/tensor/tensor_layout.h"
16	#include "sdfg/data_flow/library_nodes/math/tensor/tensor_node.h"
17	#include "sdfg/data_flow/memlet.h"
18	#include "sdfg/data_flow/tasklet.h"
19	#include "sdfg/element.h"
20	#include "sdfg/exceptions.h"
21	#include "sdfg/function.h"
22	#include "sdfg/graph/graph.h"
23	#include "sdfg/passes/expansion/lib_node_expander.h"
24	#include "sdfg/serializer/json_serializer.h"
25	#include "sdfg/structured_control_flow/block.h"
26	#include "sdfg/structured_control_flow/sequence.h"
27	#include "sdfg/structured_control_flow/structured_loop.h"
28	#include "sdfg/symbolic/symbolic.h"
29	#include "sdfg/types/scalar.h"
30	#include "sdfg/types/tensor.h"
31	#include "sdfg/types/type.h"
32
33	namespace sdfg {
34	namespace math {
35	namespace tensor {
36
37	ConcatNode::ConcatNode(
38	size_t element_id,
39	const DebugInfo& debug_info,
40	const graph::Vertex vertex,
41	data_flow::DataFlowGraph& parent,
42	const std::string& result,
43	const TensorLayout& result_layout,
44	const std::vector<std::string>& tensors,
45	const std::vector<TensorLayout>& tensor_layouts,
46	long long dim,
47	const data_flow::ImplementationType& impl_type
48	)
49	: TensorNode(element_id, debug_info, vertex, parent, LibraryNodeType_TensorConcat, {}, tensors, impl_type),	14✔
50	result_layout_(result_layout), tensor_layouts_(tensor_layouts), dim_(dim) {	14✔
51	this->inputs_.push_back(result);	14✔
52	}	14✔
53
54	const std::string& ConcatNode::result() const { return this->inputs_.back(); }	21✔
55
56	const TensorLayout& ConcatNode::result_layout() const { return this->result_layout_; }	1✔
57
58	std::vector<std::string> ConcatNode::tensors() const {	1✔
59	return std::vector<std::string>(this->inputs_.begin(), this->inputs_.end() - 1);	1✔
60	}	1✔
61
62	const std::vector<TensorLayout>& ConcatNode::tensor_layouts() const { return this->tensor_layouts_; }	13✔
63
64	long long ConcatNode::dim() const { return this->dim_; }	13✔
65
66	void ConcatNode::validate(const Function& function) const {	15✔
67	TensorNode::validate(function);	15✔
68
69	auto& graph = this->get_parent();	15✔
70
71	if (graph.out_degree(*this) != 0) {	15✔
NEW 72	throw InvalidSDFGException("ConcatNode: Expected no outputs but got: " + std::to_string(graph.out_degree(*this)));	×
NEW 73	}	×
74
75	// Check that the tensor layouts match with the tensor types on the edges
76	const data_flow::Memlet* result_iedge = graph.in_edge_for_connector(*this, this->result());	15✔
77	if (!result_iedge) {	15✔
NEW 78	throw InvalidSDFGException("ConcatNode: No memlet connected at connector: " + this->result());	×
NEW 79	}	×
80	if (result_iedge->base_type().type_id() != types::TypeID::Tensor) {	15✔
81	throw InvalidSDFGException(	1✔
82	"ConcatNode: Expected tensor type at connector '" + this->result() +	1✔
83	"' but got: " + result_iedge->base_type().print()	1✔
84	);	1✔
85	}	1✔
86	const types::Tensor& result_tensor = static_cast<const types::Tensor&>(result_iedge->base_type());	14✔
87	if (result_tensor.layout() != this->result_layout_) {	14✔
88	throw InvalidSDFGException(	1✔
89	"ConcatNode: Provided tensor layout does not match the tensor type on the edge for connector '" +	1✔
90	this->result() + "': " + result_tensor.layout().toStr() + " != " + this->result_layout_.toStr()	1✔
91	);	1✔
92	}	1✔
93	for (long long i = 0; i < this->tensor_layouts_.size(); i++) {	35✔
94	const std::string& tensor = this->inputs_[i];	24✔
95	const data_flow::Memlet* tensor_iedge = graph.in_edge_for_connector(*this, tensor);	24✔
96	if (!tensor_iedge) {	24✔
NEW 97	throw InvalidSDFGException("ConcatNode: No memlet connected at connector: " + tensor);	×
NEW 98	}	×
99	if (tensor_iedge->base_type().type_id() != types::TypeID::Tensor) {	24✔
100	throw InvalidSDFGException(	1✔
101	"ConcatNode: Expected tensor type at connector '" + tensor +	1✔
102	"' but got: " + tensor_iedge->base_type().print()	1✔
103	);	1✔
104	}	1✔
105	const types::Tensor& tensor_tensor = static_cast<const types::Tensor&>(tensor_iedge->base_type());	23✔
106	if (tensor_tensor.layout() != this->tensor_layouts_[i]) {	23✔
107	throw InvalidSDFGException(	1✔
108	"ConcatNode: Provided tensor layout does not match the tensor type on the edge for connector '" +	1✔
109	tensor + "': " + tensor_tensor.layout().toStr() + " != " + this->tensor_layouts_[i].toStr()	1✔
110	);	1✔
111	}	1✔
112	}	23✔
113
114	// Check that the cat dimension is < the tensor layout dimensions
115	if (this->dim_ < 0 \|\| this->dim_ >= this->result_layout_.dims()) {	11✔
116	throw InvalidSDFGException(	1✔
117	"ConcatNode: Cat dimension must be in [0, " + std::to_string(this->result_layout_.dims() - 1) +	1✔
118	"] but got: " + std::to_string(this->dim_)	1✔
119	);	1✔
120	}	1✔
121
122	// Check that all tensor layout have the same shape except in the cat dimension
123	symbolic::Expression cat_dim = symbolic::zero();	10✔
124	for (const TensorLayout& tensor_layout : this->tensor_layouts_) {	19✔
125	int dims = tensor_layout.dims();	19✔
126	if (dims != this->result_layout_.dims()) {	19✔
127	throw InvalidSDFGException(	1✔
128	"ConcatNode: Tensor layouts have different dimensions: " + tensor_layout.toStr() +	1✔
129	" != " + this->result_layout_.toStr()	1✔
130	);	1✔
131	}	1✔
132	for (long long i = 0; i < dims; i++) {	69✔
133	if (i == this->dim_) {	52✔
134	cat_dim = symbolic::add(cat_dim, tensor_layout.get_dim(i));	17✔
135	continue; // Skip the cat dimension	17✔
136	}	17✔
137	if (!symbolic::eq(tensor_layout.get_dim(i), this->result_layout_.get_dim(i))) {	35✔
138	throw InvalidSDFGException(	1✔
139	"ConcatNode: Shapes do not match at position " + std::to_string(i) + ": " + tensor_layout.toStr() +	1✔
140	" != " + this->result_layout_.toStr()	1✔
141	);	1✔
142	}	1✔
143	}	35✔
144	}	18✔
145	if (!symbolic::eq(cat_dim, this->result_layout_.get_dim(this->dim_))) {	8✔
146	throw InvalidSDFGException(	1✔
147	"ConcatNode: Cat dimension does not match with the sum of tensor dimensions: " + cat_dim->__str__() +	1✔
148	" != " + this->result_layout_.get_dim(this->dim_)->__str__()	1✔
149	);	1✔
150	}	1✔
151	}	8✔
152
153	bool ConcatNode::supports_integer_types() const { return true; }	15✔
154
155	using Dir = passes::LibNodeExpander::InputUse;
156
157	passes::LibNodeExpander::ExpandOutcome ConcatNode::
158	expand(passes::LibNodeExpander::ExpandContext& context, structured_control_flow::Block& block) {	1✔
159	std::vector<Dir> access_dirs(this->inputs_.size(), Dir::IndirectRead);	1✔
160	access_dirs.back() = Dir::IndirectWrite;	1✔
161	auto standalone = context.replacement_requires_access_nodes(access_dirs);	1✔
162
163	if (!standalone) {	1✔
NEW 164	return context.unable();	×
NEW 165	}	×
166
167	auto& builder = standalone->builder();	1✔
168
169	// Add a graph after the current block
170	auto& new_sequence = standalone->replace_with_sequence();	1✔
171
172	auto& dfg = this->get_parent();	1✔
173
174	types::Scalar indvar_type(types::PrimitiveType::UInt64);	1✔
175	structured_control_flow::Sequence* current_seq = &new_sequence;	1✔
176	data_flow::Subset subset;	1✔
177	subset.reserve(this->result_layout_.dims());	1✔
178	for (auto dim : this->result_layout_.shape()) {	3✔
179	auto indvar_container = builder.find_new_name("_i");	3✔
180	builder.add_container(indvar_container, indvar_type);	3✔
181	auto indvar = symbolic::symbol(indvar_container);	3✔
182	subset.push_back(indvar);	3✔
183	auto& map = builder.add_map(	3✔
184	*current_seq,	3✔
185	indvar,	3✔
186	symbolic::Lt(indvar, dim),	3✔
187	symbolic::zero(),	3✔
188	symbolic::add(indvar, symbolic::one()),	3✔
189	structured_control_flow::ScheduleType_Sequential::create(),	3✔
190	{},	3✔
191	this->debug_info_	3✔
192	);	3✔
193	current_seq = &map.root();	3✔
194	}	3✔
195
196	auto& if_else = builder.add_if_else(*current_seq, {}, this->debug_info_);	1✔
197	auto dim_indvar = subset.at(this->dim_);	1✔
198	const auto* iedge_result = dfg.in_edge_for_connector(*this, this->result());	1✔
199	if (!iedge_result) {	1✔
NEW 200	throw InvalidSDFGException("ConcatNode: Cannot get in edge for connector: " + this->result());	×
NEW 201	}	×
202
203	symbolic::Expression offset = symbolic::zero();	1✔
204	for (size_t i = 0; i < this->tensor_layouts_.size(); i++) {	3✔
205	auto new_offset = symbolic::add(offset, this->tensor_layouts_[i].get_dim(this->dim_));	2✔
206	auto condition = symbolic::And(symbolic::Ge(dim_indvar, offset), symbolic::Lt(dim_indvar, new_offset));	2✔
207	auto& case_seq = builder.add_case(if_else, condition, this->debug_info_);	2✔
208
209	data_flow::Subset offset_subset(subset);	2✔
210	offset_subset[this->dim_] = symbolic::sub(dim_indvar, offset);	2✔
211
212	auto& block = builder.add_block(case_seq, {}, this->debug_info_);	2✔
213	auto& tensor_access = standalone->add_scalar_input_access(block, i);	2✔
214	auto& result_access = standalone->add_scalar_input_access(block, this->tensor_layouts_.size());	2✔
215	auto& tasklet = builder.add_tasklet(block, data_flow::TaskletCode::assign, "_out", {"_in"}, this->debug_info_);	2✔
216
217	const auto* iedge_tensor = dfg.in_edge_for_connector(*this, this->inputs_[i]);	2✔
218	if (!iedge_tensor) {	2✔
NEW 219	throw InvalidSDFGException("ConcatNode: Cannot get in edge for connector: " + this->inputs_[i]);	×
NEW 220	}	×
221	builder.add_computational_memlet(	2✔
222	block, tensor_access, tasklet, "_in", offset_subset, iedge_tensor->base_type(), iedge_tensor->debug_info()	2✔
223	);	2✔
224	builder.add_computational_memlet(	2✔
225	block, tasklet, "_out", result_access, subset, iedge_result->base_type(), iedge_result->debug_info()	2✔
226	);	2✔
227
228	offset = new_offset;	2✔
229	}	2✔
230
231	return standalone->successfully_expanded();	1✔
232	}	1✔
233
NEW 234	std::string ConcatNode::toStr() const {	×
NEW 235	std::stringstream stream;	×
NEW 236	stream << "ConcatNode(" << this->result() << ": [";	×
NEW 237	for (long long i = 0; i < this->result_layout_.dims(); i++) {	×
NEW 238	if (i > 0) {	×
NEW 239	stream << ",";	×
NEW 240	}	×
NEW 241	stream << this->result_layout_.get_dim(i)->__str__();	×
NEW 242	}	×
NEW 243	stream << "], {";	×
NEW 244	for (long long i = 0; i < this->tensor_layouts_.size(); i++) {	×
NEW 245	if (i > 0) {	×
NEW 246	stream << ", ";	×
NEW 247	}	×
NEW 248	stream << this->inputs_[i] << ": [";	×
NEW 249	for (long long j = 0; j < this->tensor_layouts_[i].dims(); j++) {	×
NEW 250	if (j > 0) {	×
NEW 251	stream << ",";	×
NEW 252	}	×
NEW 253	stream << this->tensor_layouts_[i].get_dim(j)->__str__();	×
NEW 254	}	×
NEW 255	stream << "]";	×
NEW 256	}	×
NEW 257	stream << "}, dim: " << this->dim_ << ")";	×
NEW 258	return stream.str();	×
NEW 259	}	×
260
261	symbolic::SymbolSet ConcatNode::symbols() const {	2✔
262	symbolic::SymbolSet syms;	2✔
263	this->result_layout_.collect_symbols(syms);	2✔
264	for (const TensorLayout& tensor_layout : this->tensor_layouts_) {	4✔
265	tensor_layout.collect_symbols(syms);	4✔
266	}	4✔
267	return syms;	2✔
268	}	2✔
269
NEW 270	symbolic::Expression ConcatNode::flop() const { return symbolic::zero(); }	×
271
272	std::unique_ptr<data_flow::DataFlowNode> ConcatNode::
NEW 273	clone(size_t element_id, const graph::Vertex vertex, data_flow::DataFlowGraph& parent) const {	×
NEW 274	return std::make_unique<ConcatNode>(	×
NEW 275	element_id,	×
NEW 276	this->debug_info_,	×
NEW 277	vertex,	×
NEW 278	parent,	×
NEW 279	this->result(),	×
NEW 280	this->result_layout_,	×
NEW 281	this->tensors(),	×
NEW 282	this->tensor_layouts_,	×
NEW 283	this->dim_,	×
NEW 284	this->implementation_type_	×
NEW 285	);	×
NEW 286	}	×
287
288	void ConcatNode::replace(const symbolic::Expression old_expression, const symbolic::Expression new_expression) {	1✔
289	this->result_layout_.replace_symbols(old_expression, new_expression);	1✔
290	for (TensorLayout& tensor_layout : this->tensor_layouts_) {	2✔
291	tensor_layout.replace_symbols(old_expression, new_expression);	2✔
292	}	2✔
293	}	1✔
294
295	void ConcatNode::replace(const symbolic::ExpressionMapping& replacements) {	1✔
296	this->result_layout_.replace_symbols(replacements);	1✔
297	for (TensorLayout& tensor_layout : this->tensor_layouts_) {	2✔
298	tensor_layout.replace_symbols(replacements);	2✔
299	}	2✔
300	}	1✔
301
302	nlohmann::json ConcatNodeSerializer::serialize(const data_flow::LibraryNode& library_node) {	1✔
303	const ConcatNode& concat_node = static_cast<const ConcatNode&>(library_node);	1✔
304	nlohmann::json j;	1✔
305
306	j["code"] = concat_node.code().value();	1✔
307
308	j["result"] = concat_node.result();	1✔
309	concat_node.result_layout().serialize_to_json(j["result_layout"]);	1✔
310
311	j["tensors"] = nlohmann::json::array();	1✔
312	for (std::string& tensor : concat_node.tensors()) {	2✔
313	j["tensors"].push_back(tensor);	2✔
314	}	2✔
315	j["tensor_layouts"] = nlohmann::json::array();	1✔
316	for (const TensorLayout& tensor_layout : concat_node.tensor_layouts()) {	2✔
317	nlohmann::json tensor_layout_j;	2✔
318	tensor_layout.serialize_to_json(tensor_layout_j);	2✔
319	j["tensor_layouts"].push_back(tensor_layout_j);	2✔
320	}	2✔
321
322	j["dim"] = concat_node.dim();	1✔
323
324	return j;	1✔
325	}	1✔
326
327	data_flow::LibraryNode& ConcatNodeSerializer::deserialize(
328	const nlohmann::json& j, builder::StructuredSDFGBuilder& builder, structured_control_flow::Block& parent
329	) {	1✔
330	assert(j.contains("element_id"));	1✔
331	assert(j.contains("code"));	1✔
332	assert(j.contains("result"));	1✔
333	assert(j.contains("result_layout"));	1✔
334	assert(j.contains("tensors"));	1✔
335	assert(j.contains("tensor_layouts"));	1✔
336	assert(j.contains("dim"));	1✔
337	assert(j.contains("debug_info"));	1✔
338
339	std::string result = j.at("result").get<std::string>();	1✔
340	TensorLayout result_layout = TensorLayout::deserialize_from_json(j.at("result_layout"));	1✔
341
342	std::vector<std::string> tensors = j.at("tensors").get<std::vector<std::string>>();	1✔
343	std::vector<TensorLayout> tensor_layouts;	1✔
344	tensor_layouts.reserve(j.at("tensor_layouts").size());	1✔
345	for (long long i = 0; i < j.at("tensor_layouts").size(); i++) {	3✔
346	tensor_layouts.push_back(TensorLayout::deserialize_from_json(j.at("tensor_layouts").at(i)));	2✔
347	}	2✔
348
349	long long dim = j.at("dim").get<long long>();	1✔
350
351	sdfg::serializer::JSONSerializer serializer;	1✔
352	DebugInfo debug_info = serializer.json_to_debug_info(j.at("debug_info"));	1✔
353
354	return builder.add_library_node<ConcatNode>(parent, debug_info, result, result_layout, tensors, tensor_layouts, dim);	1✔
355	}	1✔
356
357	} // namespace tensor
358	} // namespace math
359	} // namespace sdfg

daisytuner / docc / 29936384397

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