23891477606

Committed 02 Apr 2026 08:28AM UTC coverage: 64.48% (-0.006%) from 64.486%

Build # 23891477606

Build Type

push

github

Committed by

web-flow

Commit Message

Custom linalg.generic specialization to get more layers from MLIR (#628)

- Created a linalg_custom dialect
- Replaced `LinalgSpecializeGenericOpsPass` with `LinalgCustomSpecializeGenericOpsPass`: More `linalg.generic` operations are specialized to their explicit version
- Created ReLU operation (`linalg_custom.relu`) and sigmoid operation (`linalg_custom.sigmoid`) that are translated to their SDFG counterparts
- Also dump MLIR after conversion when DOCC_DEBUG is set
- A TaskletTensorNode, now, supports a scalar as an input

Coverage Stats

11 of 21 new or added lines in 1 file covered. (52.38%)

1 existing line in 1 file now uncovered.

28826 of 44705 relevant lines covered (64.48%)

468.77 hits per line

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

48.36

/sdfg/src/data_flow/library_nodes/math/tensor/elementwise_ops/tasklet_node.cpp

#include "sdfg/data_flow/library_nodes/math/tensor/elementwise_ops/tasklet_node.h"

#include <cstddef>
#include <memory>
#include <nlohmann/json_fwd.hpp>
#include <sstream>
#include <string>
#include <unordered_map>
#include <unordered_set>
#include <vector>

#include "sdfg/analysis/analysis.h"
#include "sdfg/analysis/scope_analysis.h"
#include "sdfg/builder/structured_sdfg_builder.h"
#include "sdfg/data_flow/access_node.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_node.h"
#include "sdfg/data_flow/tasklet.h"
#include "sdfg/element.h"
#include "sdfg/exceptions.h"
#include "sdfg/graph/graph.h"
#include "sdfg/serializer/json_serializer.h"
#include "sdfg/structured_control_flow/block.h"
#include "sdfg/structured_control_flow/sequence.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 {

TaskletTensorNode::TaskletTensorNode(
    size_t element_id,
    const DebugInfo& debug_info,
    const graph::Vertex vertex,
    data_flow::DataFlowGraph& parent,
    const data_flow::TaskletCode tasklet_code,
    const std::vector<std::string>& outputs,
    const std::vector<std::string>& inputs,
    const std::vector<symbolic::Expression>& shape
)
    : TensorNode(
          element_id,
          debug_info,
          vertex,
          parent,
          LibraryNodeType_TensorTasklet,
          outputs,
          inputs,
          data_flow::ImplementationType_NONE
      ),
      tasklet_code_(tasklet_code), shape_(shape) {}

void TaskletTensorNode::validate(const Function& function) const {
    auto& graph = this->get_parent();

    // Check that all input memlets are scalar or tensor of scalar
    for (auto& iedge : graph.in_edges(*this)) {
        if (iedge.base_type().type_id() != types::TypeID::Tensor &&
            iedge.base_type().type_id() != types::TypeID::Scalar) {
            throw InvalidSDFGException(
                "TensorNode: Input memlet must be of tensor type. Found type: " + iedge.base_type().print()
            );
        }
    }

    // Check that all output memlets are scalar or tensor of scalar
    for (auto& oedge : graph.out_edges(*this)) {
        if (oedge.base_type().type_id() != types::TypeID::Tensor &&
            oedge.base_type().type_id() != types::TypeID::Scalar) {
            throw InvalidSDFGException(
                "TensorNode: Output memlet must be of tensor type. Found type: " + oedge.base_type().print()
            );
        }
    }

    // Validate: inputs match arity
    if (data_flow::arity(this->tasklet_code()) != this->inputs_.size()) {
        throw InvalidSDFGException(
            "TaskletTensorNode (Code: " + std::to_string(this->tasklet_code()) +
            "): Invalid number of inputs. Expected " + std::to_string(data_flow::arity(this->tasklet_code())) +
            ", got " + std::to_string(this->inputs_.size())
        );
    }

    // Validate: inputs match type of operation
    for (auto& iedge : graph.in_edges(*this)) {
        auto input_type = iedge.result_type(function);
        if (is_integer(this->tasklet_code()) && !types::is_integer(input_type->primitive_type())) {
            throw InvalidSDFGException(
                "TaskletTensorNode (Code: " + std::to_string(this->tasklet_code()) +
                "): Integer operation with non-integer input type"
            );
        }
        if (is_floating_point(this->tasklet_code()) && !types::is_floating_point(input_type->primitive_type())) {
            throw InvalidSDFGException(
                "TaskletTensorNode (Code: " + std::to_string(this->tasklet_code()) +
                "): Floating point operation with integer input type"
            );
        }
    }

    auto& oedge = *graph.out_edges(*this).begin();
    auto& tensor_output = static_cast<const types::Tensor&>(oedge.base_type());
    if (tensor_output.shape().size() != this->shape_.size()) {
        throw InvalidSDFGException(
            "Library Node: Output tensor shape must match node shape. Output shape: " +
            std::to_string(tensor_output.shape().size()) + " Node shape: " + std::to_string(this->shape_.size())
        );
    }
    for (size_t i = 0; i < this->shape_.size(); ++i) {
        if (!symbolic::eq(tensor_output.shape().at(i), this->shape_.at(i))) {
            throw InvalidSDFGException(
                "Library Node: Output tensor shape does not match expected shape. Output shape: " +
                tensor_output.shape().at(i)->__str__() + " Expected shape: " + this->shape_.at(i)->__str__()
            );
        }
    }

    for (auto& iedge : graph.in_edges(*this)) {
        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;
        }

        // Case 2: Tensor input
        if (tensor_input.shape().size() != this->shape_.size()) {
            throw InvalidSDFGException(
                "Library Node: Input tensor shape must match node shape. Input shape: " +
                std::to_string(tensor_input.shape().size()) + " Node shape: " + std::to_string(this->shape_.size())
            );
        }
        for (size_t i = 0; i < this->shape_.size(); ++i) {
            if (!symbolic::eq(tensor_input.shape().at(i), this->shape_.at(i))) {
                throw InvalidSDFGException(
                    "Library Node: Input tensor shape does not match expected shape. Input shape: " +
                    tensor_input.shape().at(i)->__str__() + " Expected shape: " + this->shape_.at(i)->__str__()
                );
            }
        }
    }
}

data_flow::TaskletCode TaskletTensorNode::tasklet_code() const { return this->tasklet_code_; }

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

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

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

bool TaskletTensorNode::expand(builder::StructuredSDFGBuilder& builder, analysis::AnalysisManager& analysis_manager) {
    auto& dataflow = this->get_parent();
    auto& block = static_cast<structured_control_flow::Block&>(*dataflow.get_parent());
    if (dataflow.in_degree(*this) != data_flow::arity(this->tasklet_code()) || dataflow.out_degree(*this) != 1) {
        return false;
    }
    auto& scope_analysis = analysis_manager.get<analysis::ScopeAnalysis>();
    auto& parent = static_cast<structured_control_flow::Sequence&>(*scope_analysis.parent_scope(&block));
    int index = parent.index(block);
    auto& transition = parent.at(index).second;

    auto iedges = dataflow.in_edges_by_connector(*this);
    auto& oedge = *dataflow.out_edges(*this).begin();

    // Checks if legal
    std::unordered_set<data_flow::AccessNode*> input_nodes;
    for (auto* iedge : iedges) {
        input_nodes.insert(static_cast<data_flow::AccessNode*>(&iedge->src()));
    }
    auto& output_node = static_cast<data_flow::AccessNode&>(oedge.dst());
    for (auto* input_node : input_nodes) {
        if (dataflow.in_degree(*input_node) != 0) {
            return false;
        }
    }
    if (dataflow.out_degree(output_node) != 0) {
        return false;
    }

    // Add new graph after the current block
    auto& new_sequence = builder.add_sequence_before(parent, block, transition.assignments(), block.debug_info());

    // Add maps
    structured_control_flow::Sequence* last_scope = &new_sequence;
    structured_control_flow::Map* last_map = nullptr;
    std::vector<symbolic::Expression> loop_vars;

    for (size_t i = 0; i < this->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, this->shape_[i]);
        last_map = &builder.add_map(
            *last_scope,
            indvar,
            condition,
            init,
            update,
            structured_control_flow::ScheduleType_Sequential::create(),
            {},
            block.debug_info()
        );
        last_scope = &last_map->root();

        loop_vars.push_back(indvar);
    }

    // Add tasklet block
    auto& code_block = builder.add_block(*last_scope);
    auto& tasklet = builder.add_tasklet(code_block, this->tasklet_code(), this->output(0), this->inputs());
    std::unordered_map<std::string, data_flow::AccessNode*> container_map;
    for (size_t i = 0; i < iedges.size(); i++) {
        auto* access_node = static_cast<data_flow::AccessNode*>(&iedges[i]->src());
        auto& data = access_node->data();
        data_flow::AccessNode* new_access_node = nullptr;
        if (container_map.contains(data)) {
            new_access_node = container_map.at(data);
        } else {
            if (auto* constant_node = dynamic_cast<data_flow::ConstantNode*>(access_node)) {
                new_access_node =
                    &builder.add_constant(code_block, data, constant_node->type(), constant_node->debug_info());
            } else {
                new_access_node = &builder.add_access(code_block, data, access_node->debug_info());
            }
            container_map.insert({data, new_access_node});
        }
        const auto& iedge_base_type = static_cast<const types::Tensor&>(iedges[i]->base_type());
        if (builder.subject().exists(data) &&
            (iedge_base_type.is_scalar() || iedge_base_type.type_id() == types::TypeID::Scalar)) {
            builder.add_computational_memlet(
                code_block, *new_access_node, tasklet, this->input(i), {}, iedge_base_type, iedges[i]->debug_info()
            );
        } else {
            builder.add_computational_memlet(
                code_block,
                *new_access_node,
                tasklet,
                this->input(i),
                loop_vars,
                iedge_base_type,
                iedges[i]->debug_info()
            );
        }
    }
    auto& out_node = builder.add_access(code_block, output_node.data(), output_node.debug_info());
    builder.add_computational_memlet(
        code_block, tasklet, this->output(0), out_node, loop_vars, oedge.base_type(), oedge.debug_info()
    );

    // Clean up block
    for (auto* iedge : iedges) {
        builder.remove_memlet(block, *iedge);
    }
    builder.remove_memlet(block, oedge);
    for (auto* input_node : input_nodes) {
        builder.remove_node(block, *input_node);
    }
    builder.remove_node(block, output_node);
    builder.remove_node(block, *this);
    builder.remove_child(parent, index + 1);

    return true;
}

bool TaskletTensorNode::supports_integer_types() const { return data_flow::is_integer(this->tasklet_code()); }

std::unique_ptr<data_flow::DataFlowNode> TaskletTensorNode::
    clone(size_t element_id, const graph::Vertex vertex, data_flow::DataFlowGraph& parent) const {
    return std::unique_ptr<data_flow::DataFlowNode>(new TaskletTensorNode(
        element_id,
        this->debug_info(),
        vertex,
        parent,
        this->tasklet_code(),
        this->outputs(),
        this->inputs(),
        this->shape()
    ));
}

std::string TaskletTensorNode::toStr() const {
    std::stringstream stream;

    stream << this->code().value() << ": " << std::to_string(this->tasklet_code()) << ", [";
    for (size_t i = 0; i < this->shape().size(); i++) {
        if (i > 0) {
            stream << ", ";
        }
        stream << this->shape().at(i)->__str__();
    }
    stream << "]";

    return stream.str();
}

nlohmann::json TaskletTensorNodeSerializer::serialize(const data_flow::LibraryNode& library_node) {
    const auto& elem_node = static_cast<const TaskletTensorNode&>(library_node);
    nlohmann::json j;

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

    j["output"] = elem_node.output(0);

    j["inputs"] = nlohmann::json::array();
    for (auto& input : elem_node.inputs()) {
        j["inputs"].push_back(input);
    }

    j["tasklet_code"] = elem_node.tasklet_code();

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

    return j;
}

data_flow::LibraryNode& TaskletTensorNodeSerializer::deserialize(
    const nlohmann::json& j, builder::StructuredSDFGBuilder& builder, structured_control_flow::Block& parent
) {
    // Assertions for required fields
    assert(j.contains("element_id"));
    assert(j.contains("code"));
    assert(j.contains("debug_info"));
    assert(j.contains("output"));
    assert(j.contains("inputs"));
    assert(j.contains("tasklet_code"));
    assert(j.contains("shape"));

    auto code = j["code"].get<std::string>();

    std::vector<std::string> outputs({j["output"].get<std::string>()});

    std::vector<std::string> inputs;
    for (const auto& input : j["inputs"]) {
        inputs.push_back(input.get<std::string>());
    }

    auto tasklet_code = static_cast<data_flow::TaskletCode>(j["tasklet_code"].get<int>());

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

    // Extract debug info using JSONSerializer
    sdfg::serializer::JSONSerializer serializer;
    DebugInfo debug_info = serializer.json_to_debug_info(j["debug_info"]);

    return static_cast<TaskletTensorNode&>(builder.add_library_node<
                                           TaskletTensorNode>(parent, debug_info, tasklet_code, outputs, inputs, shape)
    );
}

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

1	#include "sdfg/data_flow/library_nodes/math/tensor/elementwise_ops/tasklet_node.h"
2
3	#include <cstddef>
4	#include <memory>
5	#include <nlohmann/json_fwd.hpp>
6	#include <sstream>
7	#include <string>
8	#include <unordered_map>
9	#include <unordered_set>
10	#include <vector>
11
12	#include "sdfg/analysis/analysis.h"
13	#include "sdfg/analysis/scope_analysis.h"
14	#include "sdfg/builder/structured_sdfg_builder.h"
15	#include "sdfg/data_flow/access_node.h"
16	#include "sdfg/data_flow/data_flow_graph.h"
17	#include "sdfg/data_flow/data_flow_node.h"
18	#include "sdfg/data_flow/library_node.h"
19	#include "sdfg/data_flow/library_nodes/math/tensor/tensor_node.h"
20	#include "sdfg/data_flow/tasklet.h"
21	#include "sdfg/element.h"
22	#include "sdfg/exceptions.h"
23	#include "sdfg/graph/graph.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/symbolic/symbolic.h"
28	#include "sdfg/types/scalar.h"
29	#include "sdfg/types/tensor.h"
30	#include "sdfg/types/type.h"
31
32	namespace sdfg {
33	namespace math {
34	namespace tensor {
35
36	TaskletTensorNode::TaskletTensorNode(
37	size_t element_id,
38	const DebugInfo& debug_info,
39	const graph::Vertex vertex,
40	data_flow::DataFlowGraph& parent,
41	const data_flow::TaskletCode tasklet_code,
42	const std::vector<std::string>& outputs,
43	const std::vector<std::string>& inputs,
44	const std::vector<symbolic::Expression>& shape
45	)
46	: TensorNode(	208✔
47	element_id,	208✔
48	debug_info,	208✔
49	vertex,	208✔
50	parent,	208✔
51	LibraryNodeType_TensorTasklet,	208✔
52	outputs,	208✔
53	inputs,	208✔
54	data_flow::ImplementationType_NONE	208✔
55	),	208✔
56	tasklet_code_(tasklet_code), shape_(shape) {}	208✔
57
58	void TaskletTensorNode::validate(const Function& function) const {	208✔
59	auto& graph = this->get_parent();	208✔
60
61	// Check that all input memlets are scalar or tensor of scalar
62	for (auto& iedge : graph.in_edges(*this)) {	408✔
63	if (iedge.base_type().type_id() != types::TypeID::Tensor &&	408✔
64	iedge.base_type().type_id() != types::TypeID::Scalar) {	408✔
NEW 65	throw InvalidSDFGException(	×
NEW 66	"TensorNode: Input memlet must be of tensor type. Found type: " + iedge.base_type().print()	×
NEW 67	);	×
NEW 68	}	×
69	}	408✔
70
71	// Check that all output memlets are scalar or tensor of scalar
72	for (auto& oedge : graph.out_edges(*this)) {	208✔
73	if (oedge.base_type().type_id() != types::TypeID::Tensor &&	208✔
74	oedge.base_type().type_id() != types::TypeID::Scalar) {	208✔
NEW 75	throw InvalidSDFGException(	×
NEW 76	"TensorNode: Output memlet must be of tensor type. Found type: " + oedge.base_type().print()	×
NEW 77	);	×
NEW 78	}	×
79	}	208✔
80
81	// Validate: inputs match arity
82	if (data_flow::arity(this->tasklet_code()) != this->inputs_.size()) {	208✔
83	throw InvalidSDFGException(	×
84	"TaskletTensorNode (Code: " + std::to_string(this->tasklet_code()) +	×
85	"): Invalid number of inputs. Expected " + std::to_string(data_flow::arity(this->tasklet_code())) +	×
86	", got " + std::to_string(this->inputs_.size())	×
87	);	×
88	}	×
89
90	// Validate: inputs match type of operation
91	for (auto& iedge : graph.in_edges(*this)) {	408✔
92	auto input_type = iedge.result_type(function);	408✔
93	if (is_integer(this->tasklet_code()) && !types::is_integer(input_type->primitive_type())) {	408✔
94	throw InvalidSDFGException(	×
95	"TaskletTensorNode (Code: " + std::to_string(this->tasklet_code()) +	×
96	"): Integer operation with non-integer input type"	×
97	);	×
98	}	×
99	if (is_floating_point(this->tasklet_code()) && !types::is_floating_point(input_type->primitive_type())) {	408✔
100	throw InvalidSDFGException(	×
101	"TaskletTensorNode (Code: " + std::to_string(this->tasklet_code()) +	×
102	"): Floating point operation with integer input type"	×
103	);	×
104	}	×
105	}	408✔
106
107	auto& oedge = graph.out_edges(this).begin();	208✔
108	auto& tensor_output = static_cast<const types::Tensor&>(oedge.base_type());	208✔
109	if (tensor_output.shape().size() != this->shape_.size()) {	208✔
110	throw InvalidSDFGException(	×
111	"Library Node: Output tensor shape must match node shape. Output shape: " +	×
112	std::to_string(tensor_output.shape().size()) + " Node shape: " + std::to_string(this->shape_.size())	×
113	);	×
114	}	×
115	for (size_t i = 0; i < this->shape_.size(); ++i) {	728✔
116	if (!symbolic::eq(tensor_output.shape().at(i), this->shape_.at(i))) {	520✔
117	throw InvalidSDFGException(	×
118	"Library Node: Output tensor shape does not match expected shape. Output shape: " +	×
119	tensor_output.shape().at(i)->__str__() + " Expected shape: " + this->shape_.at(i)->__str__()	×
120	);	×
121	}	×
122	}	520✔
123
124	for (auto& iedge : graph.in_edges(*this)) {	408✔
125	if (iedge.base_type().type_id() == types::TypeID::Scalar) {	408✔
NEW 126	continue;	×
NEW 127	}	×
128	auto& tensor_input = static_cast<const types::Tensor&>(iedge.base_type());	408✔
129	// Case 1: Scalar input is allowed as secondary input
130	if (tensor_input.is_scalar()) {	408✔
131	continue;	×
132	}	×
133
134	// Case 2: Tensor input
135	if (tensor_input.shape().size() != this->shape_.size()) {	408✔
136	throw InvalidSDFGException(	×
137	"Library Node: Input tensor shape must match node shape. Input shape: " +	×
138	std::to_string(tensor_input.shape().size()) + " Node shape: " + std::to_string(this->shape_.size())	×
139	);	×
140	}	×
141	for (size_t i = 0; i < this->shape_.size(); ++i) {	1,428✔
142	if (!symbolic::eq(tensor_input.shape().at(i), this->shape_.at(i))) {	1,020✔
143	throw InvalidSDFGException(	×
144	"Library Node: Input tensor shape does not match expected shape. Input shape: " +	×
145	tensor_input.shape().at(i)->__str__() + " Expected shape: " + this->shape_.at(i)->__str__()	×
146	);	×
147	}	×
148	}	1,020✔
149	}	408✔
150	}	208✔
151
152	data_flow::TaskletCode TaskletTensorNode::tasklet_code() const { return this->tasklet_code_; }	1,440✔
153
154	const std::vector<symbolic::Expression>& TaskletTensorNode::shape() const { return this->shape_; }	×
155
156	symbolic::SymbolSet TaskletTensorNode::symbols() const {	×
157	symbolic::SymbolSet syms;	×
158	for (const auto& dim : this->shape_) {	×
159	for (auto& atom : symbolic::atoms(dim)) {	×
160	syms.insert(atom);	×
161	}	×
162	}	×
163	return syms;	×
164	}	×
165
166	void TaskletTensorNode::replace(const symbolic::Expression old_expression, const symbolic::Expression new_expression) {	×
167	for (auto& dim : shape_) {	×
168	dim = symbolic::subs(dim, old_expression, new_expression);	×
169	}	×
170	}	×
171
172	bool TaskletTensorNode::expand(builder::StructuredSDFGBuilder& builder, analysis::AnalysisManager& analysis_manager) {	208✔
173	auto& dataflow = this->get_parent();	208✔
174	auto& block = static_cast<structured_control_flow::Block&>(*dataflow.get_parent());	208✔
175	if (dataflow.in_degree(this) != data_flow::arity(this->tasklet_code()) \|\| dataflow.out_degree(this) != 1) {	208✔
176	return false;	×
177	}	×
178	auto& scope_analysis = analysis_manager.get<analysis::ScopeAnalysis>();	208✔
179	auto& parent = static_cast<structured_control_flow::Sequence&>(*scope_analysis.parent_scope(&block));	208✔
180	int index = parent.index(block);	208✔
181	auto& transition = parent.at(index).second;	208✔
182
183	auto iedges = dataflow.in_edges_by_connector(*this);	208✔
184	auto& oedge = dataflow.out_edges(this).begin();	208✔
185
186	// Checks if legal
187	std::unordered_set<data_flow::AccessNode*> input_nodes;	208✔
188	for (auto* iedge : iedges) {	408✔
189	input_nodes.insert(static_cast<data_flow::AccessNode*>(&iedge->src()));	408✔
190	}	408✔
191	auto& output_node = static_cast<data_flow::AccessNode&>(oedge.dst());	208✔
192	for (auto* input_node : input_nodes) {	408✔
193	if (dataflow.in_degree(*input_node) != 0) {	408✔
194	return false;	×
195	}	×
196	}	408✔
197	if (dataflow.out_degree(output_node) != 0) {	208✔
198	return false;	×
199	}	×
200
201	// Add new graph after the current block
202	auto& new_sequence = builder.add_sequence_before(parent, block, transition.assignments(), block.debug_info());	208✔
203
204	// Add maps
205	structured_control_flow::Sequence* last_scope = &new_sequence;	208✔
206	structured_control_flow::Map* last_map = nullptr;	208✔
207	std::vector<symbolic::Expression> loop_vars;	208✔
208
209	for (size_t i = 0; i < this->shape_.size(); i++) {	728✔
210	std::string indvar_str = builder.find_new_name("_i");	520✔
211	builder.add_container(indvar_str, types::Scalar(types::PrimitiveType::UInt64));	520✔
212
213	auto indvar = symbolic::symbol(indvar_str);	520✔
214	auto init = symbolic::zero();	520✔
215	auto update = symbolic::add(indvar, symbolic::one());	520✔
216	auto condition = symbolic::Lt(indvar, this->shape_[i]);	520✔
217	last_map = &builder.add_map(	520✔
218	*last_scope,	520✔
219	indvar,	520✔
220	condition,	520✔
221	init,	520✔
222	update,	520✔
223	structured_control_flow::ScheduleType_Sequential::create(),	520✔
224	{},	520✔
225	block.debug_info()	520✔
226	);	520✔
227	last_scope = &last_map->root();	520✔
228
229	loop_vars.push_back(indvar);	520✔
230	}	520✔
231
232	// Add tasklet block
233	auto& code_block = builder.add_block(*last_scope);	208✔
234	auto& tasklet = builder.add_tasklet(code_block, this->tasklet_code(), this->output(0), this->inputs());	208✔
235	std::unordered_map<std::string, data_flow::AccessNode*> container_map;	208✔
236	for (size_t i = 0; i < iedges.size(); i++) {	616✔
237	auto* access_node = static_cast<data_flow::AccessNode*>(&iedges[i]->src());	408✔
238	auto& data = access_node->data();	408✔
239	data_flow::AccessNode* new_access_node = nullptr;	408✔
240	if (container_map.contains(data)) {	408✔
241	new_access_node = container_map.at(data);	×
242	} else {	408✔
243	if (auto* constant_node = dynamic_cast<data_flow::ConstantNode*>(access_node)) {	408✔
244	new_access_node =	×
245	&builder.add_constant(code_block, data, constant_node->type(), constant_node->debug_info());	×
246	} else {	408✔
247	new_access_node = &builder.add_access(code_block, data, access_node->debug_info());	408✔
248	}	408✔
249	container_map.insert({data, new_access_node});	408✔
250	}	408✔
251	const auto& iedge_base_type = static_cast<const types::Tensor&>(iedges[i]->base_type());	408✔
252	if (builder.subject().exists(data) &&	408✔
253	(iedge_base_type.is_scalar() \|\| iedge_base_type.type_id() == types::TypeID::Scalar)) {	408✔
254	builder.add_computational_memlet(	×
255	code_block, *new_access_node, tasklet, this->input(i), {}, iedge_base_type, iedges[i]->debug_info()	×
256	);	×
257	} else {	408✔
258	builder.add_computational_memlet(	408✔
259	code_block,	408✔
260	*new_access_node,	408✔
261	tasklet,	408✔
262	this->input(i),	408✔
263	loop_vars,	408✔
264	iedge_base_type,	408✔
265	iedges[i]->debug_info()	408✔
266	);	408✔
267	}	408✔
268	}	408✔
269	auto& out_node = builder.add_access(code_block, output_node.data(), output_node.debug_info());	208✔
270	builder.add_computational_memlet(	208✔
271	code_block, tasklet, this->output(0), out_node, loop_vars, oedge.base_type(), oedge.debug_info()	208✔
272	);	208✔
273
274	// Clean up block
275	for (auto* iedge : iedges) {	408✔
276	builder.remove_memlet(block, *iedge);	408✔
277	}	408✔
278	builder.remove_memlet(block, oedge);	208✔
279	for (auto* input_node : input_nodes) {	408✔
280	builder.remove_node(block, *input_node);	408✔
281	}	408✔
282	builder.remove_node(block, output_node);	208✔
283	builder.remove_node(block, *this);	208✔
284	builder.remove_child(parent, index + 1);	208✔
285
286	return true;	208✔
287	}	208✔
288
UNCOV 289	bool TaskletTensorNode::supports_integer_types() const { return data_flow::is_integer(this->tasklet_code()); }	×
290
291	std::unique_ptr<data_flow::DataFlowNode> TaskletTensorNode::
292	clone(size_t element_id, const graph::Vertex vertex, data_flow::DataFlowGraph& parent) const {	×
293	return std::unique_ptr<data_flow::DataFlowNode>(new TaskletTensorNode(	×
294	element_id,	×
295	this->debug_info(),	×
296	vertex,	×
297	parent,	×
298	this->tasklet_code(),	×
299	this->outputs(),	×
300	this->inputs(),	×
301	this->shape()	×
302	));	×
303	}	×
304
305	std::string TaskletTensorNode::toStr() const {	×
306	std::stringstream stream;	×
307
308	stream << this->code().value() << ": " << std::to_string(this->tasklet_code()) << ", [";	×
309	for (size_t i = 0; i < this->shape().size(); i++) {	×
310	if (i > 0) {	×
311	stream << ", ";	×
312	}	×
313	stream << this->shape().at(i)->__str__();	×
314	}	×
315	stream << "]";	×
316
317	return stream.str();	×
318	}	×
319
320	nlohmann::json TaskletTensorNodeSerializer::serialize(const data_flow::LibraryNode& library_node) {	×
321	const auto& elem_node = static_cast<const TaskletTensorNode&>(library_node);	×
322	nlohmann::json j;	×
323
324	j["code"] = elem_node.code().value();	×
325
326	j["output"] = elem_node.output(0);	×
327
328	j["inputs"] = nlohmann::json::array();	×
329	for (auto& input : elem_node.inputs()) {	×
330	j["inputs"].push_back(input);	×
331	}	×
332
333	j["tasklet_code"] = elem_node.tasklet_code();	×
334
335	serializer::JSONSerializer serializer;	×
336	j["shape"] = nlohmann::json::array();	×
337	for (auto& dim : elem_node.shape()) {	×
338	j["shape"].push_back(serializer.expression(dim));	×
339	}	×
340
341	return j;	×
342	}	×
343
344	data_flow::LibraryNode& TaskletTensorNodeSerializer::deserialize(
345	const nlohmann::json& j, builder::StructuredSDFGBuilder& builder, structured_control_flow::Block& parent
346	) {	×
347	// Assertions for required fields
348	assert(j.contains("element_id"));	×
349	assert(j.contains("code"));	×
350	assert(j.contains("debug_info"));	×
351	assert(j.contains("output"));	×
352	assert(j.contains("inputs"));	×
353	assert(j.contains("tasklet_code"));	×
354	assert(j.contains("shape"));	×
355
356	auto code = j["code"].get<std::string>();	×
357
358	std::vector<std::string> outputs({j["output"].get<std::string>()});	×
359
360	std::vector<std::string> inputs;	×
361	for (const auto& input : j["inputs"]) {	×
362	inputs.push_back(input.get<std::string>());	×
363	}	×
364
365	auto tasklet_code = static_cast<data_flow::TaskletCode>(j["tasklet_code"].get<int>());	×
366
367	std::vector<symbolic::Expression> shape;	×
368	for (const auto& dim : j["shape"]) {	×
369	shape.push_back(symbolic::parse(dim.get<std::string>()));	×
370	}	×
371
372	// Extract debug info using JSONSerializer
373	sdfg::serializer::JSONSerializer serializer;	×
374	DebugInfo debug_info = serializer.json_to_debug_info(j["debug_info"]);	×
375
376	return static_cast<TaskletTensorNode&>(builder.add_library_node<	×
377	TaskletTensorNode>(parent, debug_info, tasklet_code, outputs, inputs, shape)	×
378	);	×
379	}	×
380
381	} // namespace tensor
382	} // namespace math
383	} // namespace sdfg

daisytuner / docc / 23891477606

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