23086520240

Committed 14 Mar 2026 10:54AM UTC coverage: 63.927% (+0.3%) from 63.617%

Build # 23086520240

Build Type

Pull #568

github

Committed by

web-flow

Commit Message

Merge 21f66b0ec into 541bceab5

Pull Request Pull Request #568: Working on memory ownership & escape analysis

Coverage Stats

475 of 637 new or added lines in 28 files covered. (74.57%)

6 existing lines in 3 files now uncovered.

26010 of 40687 relevant lines covered (63.93%)

402.05 hits per line

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

0.0

/sdfg/src/data_flow/library_nodes/math/tensor/reduce_ops/softmax_node.cpp

#include "sdfg/data_flow/library_nodes/math/tensor/reduce_ops/softmax_node.h"
#include "sdfg/analysis/scope_analysis.h"
#include "sdfg/builder/structured_sdfg_builder.h"
#include "sdfg/data_flow/library_nodes/math/tensor/broadcast_node.h"
#include "sdfg/data_flow/library_nodes/math/tensor/elementwise_ops/div_node.h"
#include "sdfg/data_flow/library_nodes/math/tensor/elementwise_ops/exp_node.h"
#include "sdfg/data_flow/library_nodes/math/tensor/elementwise_ops/sub_node.h"
#include "sdfg/data_flow/library_nodes/math/tensor/reduce_ops/max_node.h"
#include "sdfg/data_flow/library_nodes/math/tensor/reduce_ops/sum_node.h"
#include "sdfg/data_flow/library_nodes/stdlib/malloc.h"
#include "sdfg/structured_control_flow/block.h"
#include "sdfg/structured_control_flow/for.h"
#include "sdfg/types/pointer.h"
#include "sdfg/types/scalar.h"
#include "sdfg/types/utils.h"

namespace sdfg {
namespace math {
namespace tensor {

SoftmaxNode::SoftmaxNode(
    size_t element_id,
    const DebugInfo& debug_info,
    const graph::Vertex vertex,
    data_flow::DataFlowGraph& parent,
    const std::vector<symbolic::Expression>& shape,
    const std::vector<int64_t>& axes,
    bool keepdims
)
    : ReduceNode(element_id, debug_info, vertex, parent, LibraryNodeType_Softmax, shape, axes, keepdims) {
    if (keepdims) {
        throw InvalidSDFGException("Unsupported attribute on library node: softmax");
    }
}

void SoftmaxNode::validate(const Function& function) const {}

bool SoftmaxNode::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) != 1 || 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& in_edge = *dataflow.in_edges(*this).begin();
    auto& out_edge = *dataflow.out_edges(*this).begin();
    auto& in_node = static_cast<data_flow::AccessNode&>(in_edge.src());
    auto& out_node = static_cast<data_flow::AccessNode&>(out_edge.dst());

    // Calculate reduced shape (for Max and Sum)
    std::vector<symbolic::Expression> reduced_shape;
    std::vector<int64_t> sorted_axes = axes_;
    // Normalize negative axes
    for (auto& axis : sorted_axes) {
        if (axis < 0) {
            axis = static_cast<int64_t>(shape_.size()) + axis;
        }
        // Validate axis is in bounds
        if (axis < 0 || axis >= static_cast<int64_t>(shape_.size())) {
            throw InvalidSDFGException(
                "Library Node: Axis value out of bounds. Axis: " + std::to_string(axis) +
                " Shape size: " + std::to_string(shape_.size())
            );
        }
    }
    std::sort(sorted_axes.begin(), sorted_axes.end());

    for (size_t i = 0; i < shape_.size(); ++i) {
        bool is_axis = false;
        for (auto axis : sorted_axes) {
            if (axis == (int64_t) i) {
                is_axis = true;
                break;
            }
        }

        if (is_axis) {
            reduced_shape.push_back(symbolic::one());
        } else {
            reduced_shape.push_back(shape_[i]);
        }
    }

    types::Scalar element_type(this->primitive_type(dataflow));
    types::Pointer pointer_type(element_type);

    // Type to store reduced results (e.g., max and sum)
    types::Tensor reduced_tensor_type(element_type, reduced_shape);
    // Type for broadcasted tensors (e.g., max and sum after broadcasting)
    // Compute broadcast strides: use strides from reduced_tensor_type, but set to zero for reduced dimensions
    symbolic::MultiExpression reduced_strides = reduced_tensor_type.strides();
    symbolic::MultiExpression broadcast_strides;
    for (size_t i = 0; i < shape_.size(); ++i) {
        bool is_reduced = std::find(sorted_axes.begin(), sorted_axes.end(), static_cast<int64_t>(i)) !=
                          sorted_axes.end();
        if (is_reduced) {
            broadcast_strides.push_back(symbolic::zero());
        } else {
            broadcast_strides.push_back(reduced_strides[i]);
        }
    }
    types::Tensor broadcast_tensor_type(element_type, shape_, broadcast_strides);

    // Temporary buffers
    std::string tmp_max_name = builder.find_new_name("_softmax_max");
    builder.add_container(tmp_max_name, pointer_type);

    std::string tmp_sub_name = builder.find_new_name("_softmax_sub");
    builder.add_container(tmp_sub_name, pointer_type);

    std::string tmp_exp_name = builder.find_new_name("_softmax_exp");
    builder.add_container(tmp_exp_name, pointer_type);

    std::string tmp_sum_name = builder.find_new_name("_softmax_sum");
    builder.add_container(tmp_sum_name, pointer_type);

    // Mallocs
    {
        symbolic::Expression bytes_elem = types::get_type_size(element_type, false);

        symbolic::Expression bytes_full = bytes_elem;
        for (auto& dim : this->shape_) {
            bytes_full = symbolic::mul(dim, bytes_full);
        }

        symbolic::Expression bytes_reduced = bytes_elem;
        for (auto& dim : reduced_shape) {
            bytes_reduced = symbolic::mul(dim, bytes_reduced);
        }

        auto& alloc_block = builder.add_block_before(parent, block, {}, this->debug_info());

        auto malloc_helper = [&](const std::string& name, const symbolic::Expression& size) {
            auto& access = builder.add_access(alloc_block, name);
            auto& malloc_node = builder.add_library_node<stdlib::MallocNode>(alloc_block, this->debug_info(), size);
            builder
                .add_computational_memlet(alloc_block, malloc_node, "_ret", access, {}, pointer_type, this->debug_info());
        };

        malloc_helper(tmp_max_name, bytes_reduced);
        malloc_helper(tmp_sub_name, bytes_full);
        malloc_helper(tmp_exp_name, bytes_full);
        malloc_helper(tmp_sum_name, bytes_reduced);
    }

    // 1. Max(X) -> TmpMax
    {
        auto& max_block = builder.add_block_before(parent, block, {}, this->debug_info());
        auto& max_node =
            builder.add_library_node<MaxNode>(max_block, this->debug_info(), this->shape_, this->axes_, true);

        auto& in_access = builder.add_access(max_block, in_node.data());
        auto& out_access = builder.add_access(max_block, tmp_max_name);

        builder
            .add_computational_memlet(max_block, in_access, max_node, "X", {}, in_edge.base_type(), this->debug_info());
        builder
            .add_computational_memlet(max_block, max_node, "Y", out_access, {}, reduced_tensor_type, this->debug_info());
    }

    // 2. Sub(X, TmpMaxBcast) -> TmpSub
    {
        auto& sub_block = builder.add_block_before(parent, block, {}, this->debug_info());
        auto& sub_node = builder.add_library_node<SubNode>(sub_block, this->debug_info(), this->shape_);

        auto& in1_access = builder.add_access(sub_block, in_node.data());
        auto& in2_access = builder.add_access(sub_block, tmp_max_name);
        auto& out_access = builder.add_access(sub_block, tmp_sub_name);

        builder
            .add_computational_memlet(sub_block, in1_access, sub_node, "A", {}, in_edge.base_type(), this->debug_info());
        builder
            .add_computational_memlet(sub_block, in2_access, sub_node, "B", {}, broadcast_tensor_type, this->debug_info());
        builder
            .add_computational_memlet(sub_block, sub_node, "C", out_access, {}, in_edge.base_type(), this->debug_info());
    }

    // 3. Exp(TmpSub) -> TmpExp
    {
        auto& exp_block = builder.add_block_before(parent, block, {}, this->debug_info());
        auto& exp_node = builder.add_library_node<ExpNode>(exp_block, this->debug_info(), this->shape_);

        auto& in_access = builder.add_access(exp_block, tmp_sub_name);
        auto& out_access = builder.add_access(exp_block, tmp_exp_name);

        builder
            .add_computational_memlet(exp_block, in_access, exp_node, "X", {}, in_edge.base_type(), this->debug_info());
        builder
            .add_computational_memlet(exp_block, exp_node, "Y", out_access, {}, in_edge.base_type(), this->debug_info());
    }

    // 4. Sum(TmpExp) -> TmpSum
    {
        auto& sum_block = builder.add_block_before(parent, block, {}, this->debug_info());
        auto& sum_node =
            builder.add_library_node<SumNode>(sum_block, this->debug_info(), this->shape_, this->axes_, true);

        auto& in_access = builder.add_access(sum_block, tmp_exp_name);
        auto& out_access = builder.add_access(sum_block, tmp_sum_name);

        builder
            .add_computational_memlet(sum_block, in_access, sum_node, "X", {}, in_edge.base_type(), this->debug_info());
        builder
            .add_computational_memlet(sum_block, sum_node, "Y", out_access, {}, reduced_tensor_type, this->debug_info());
    }

    // 5. Div(TmpExp, TmpSum) -> Output
    {
        auto& div_block = builder.add_block_before(parent, block, {}, this->debug_info());
        auto& div_node = builder.add_library_node<DivNode>(div_block, this->debug_info(), this->shape_);

        auto& in1_access = builder.add_access(div_block, tmp_exp_name);
        auto& in2_access = builder.add_access(div_block, tmp_sum_name);
        auto& out_access = builder.add_access(div_block, out_node.data());

        builder
            .add_computational_memlet(div_block, in1_access, div_node, "A", {}, in_edge.base_type(), this->debug_info());
        builder
            .add_computational_memlet(div_block, in2_access, div_node, "B", {}, broadcast_tensor_type, this->debug_info());
        builder
            .add_computational_memlet(div_block, div_node, "C", out_access, {}, out_edge.base_type(), this->debug_info());
    }

    // Cleanup
    builder.remove_memlet(block, in_edge);
    builder.remove_memlet(block, out_edge);
    builder.remove_node(block, in_node);
    builder.remove_node(block, out_node);
    builder.remove_node(block, *this);

    int last_index = parent.index(block);
    builder.remove_child(parent, last_index);

    return true;
}

std::unique_ptr<data_flow::DataFlowNode> SoftmaxNode::
    clone(size_t element_id, const graph::Vertex vertex, data_flow::DataFlowGraph& parent) const {
    return std::unique_ptr<
        data_flow::DataFlowNode>(new SoftmaxNode(element_id, this->debug_info(), vertex, parent, this->shape_, this->axes_)
    );
}

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

1	#include "sdfg/data_flow/library_nodes/math/tensor/reduce_ops/softmax_node.h"
2	#include "sdfg/analysis/scope_analysis.h"
3	#include "sdfg/builder/structured_sdfg_builder.h"
4	#include "sdfg/data_flow/library_nodes/math/tensor/broadcast_node.h"
5	#include "sdfg/data_flow/library_nodes/math/tensor/elementwise_ops/div_node.h"
6	#include "sdfg/data_flow/library_nodes/math/tensor/elementwise_ops/exp_node.h"
7	#include "sdfg/data_flow/library_nodes/math/tensor/elementwise_ops/sub_node.h"
8	#include "sdfg/data_flow/library_nodes/math/tensor/reduce_ops/max_node.h"
9	#include "sdfg/data_flow/library_nodes/math/tensor/reduce_ops/sum_node.h"
10	#include "sdfg/data_flow/library_nodes/stdlib/malloc.h"
11	#include "sdfg/structured_control_flow/block.h"
12	#include "sdfg/structured_control_flow/for.h"
13	#include "sdfg/types/pointer.h"
14	#include "sdfg/types/scalar.h"
15	#include "sdfg/types/utils.h"
16
17	namespace sdfg {
18	namespace math {
19	namespace tensor {
20
21	SoftmaxNode::SoftmaxNode(
22	size_t element_id,
23	const DebugInfo& debug_info,
24	const graph::Vertex vertex,
25	data_flow::DataFlowGraph& parent,
26	const std::vector<symbolic::Expression>& shape,
27	const std::vector<int64_t>& axes,
28	bool keepdims
29	)
30	: ReduceNode(element_id, debug_info, vertex, parent, LibraryNodeType_Softmax, shape, axes, keepdims) {	×
31	if (keepdims) {	×
32	throw InvalidSDFGException("Unsupported attribute on library node: softmax");	×
33	}	×
34	}	×
35
36	void SoftmaxNode::validate(const Function& function) const {}	×
37
38	bool SoftmaxNode::expand(builder::StructuredSDFGBuilder& builder, analysis::AnalysisManager& analysis_manager) {	×
39	auto& dataflow = this->get_parent();	×
40	auto& block = static_cast<structured_control_flow::Block&>(*dataflow.get_parent());	×
41
42	if (dataflow.in_degree(this) != 1 \|\| dataflow.out_degree(this) != 1) {	×
43	return false;	×
44	}	×
45
46	auto& scope_analysis = analysis_manager.get<analysis::ScopeAnalysis>();	×
47	auto& parent = static_cast<structured_control_flow::Sequence&>(*scope_analysis.parent_scope(&block));	×
48	int index = parent.index(block);	×
49
50	auto& in_edge = dataflow.in_edges(this).begin();	×
51	auto& out_edge = dataflow.out_edges(this).begin();	×
52	auto& in_node = static_cast<data_flow::AccessNode&>(in_edge.src());	×
53	auto& out_node = static_cast<data_flow::AccessNode&>(out_edge.dst());	×
54
55	// Calculate reduced shape (for Max and Sum)
56	std::vector<symbolic::Expression> reduced_shape;	×
57	std::vector<int64_t> sorted_axes = axes_;	×
58	// Normalize negative axes
59	for (auto& axis : sorted_axes) {	×
60	if (axis < 0) {	×
61	axis = static_cast<int64_t>(shape_.size()) + axis;	×
62	}	×
63	// Validate axis is in bounds
64	if (axis < 0 \|\| axis >= static_cast<int64_t>(shape_.size())) {	×
65	throw InvalidSDFGException(	×
66	"Library Node: Axis value out of bounds. Axis: " + std::to_string(axis) +	×
67	" Shape size: " + std::to_string(shape_.size())	×
68	);	×
69	}	×
70	}	×
71	std::sort(sorted_axes.begin(), sorted_axes.end());	×
72
73	for (size_t i = 0; i < shape_.size(); ++i) {	×
74	bool is_axis = false;	×
75	for (auto axis : sorted_axes) {	×
76	if (axis == (int64_t) i) {	×
77	is_axis = true;	×
78	break;	×
79	}	×
80	}	×
81
82	if (is_axis) {	×
83	reduced_shape.push_back(symbolic::one());	×
84	} else {	×
85	reduced_shape.push_back(shape_[i]);	×
86	}	×
87	}	×
88
89	types::Scalar element_type(this->primitive_type(dataflow));	×
90	types::Pointer pointer_type(element_type);	×
91
92	// Type to store reduced results (e.g., max and sum)
NEW 93	types::Tensor reduced_tensor_type(element_type, reduced_shape);	×
94	// Type for broadcasted tensors (e.g., max and sum after broadcasting)
95	// Compute broadcast strides: use strides from reduced_tensor_type, but set to zero for reduced dimensions
NEW 96	symbolic::MultiExpression reduced_strides = reduced_tensor_type.strides();	×
NEW 97	symbolic::MultiExpression broadcast_strides;	×
NEW 98	for (size_t i = 0; i < shape_.size(); ++i) {	×
NEW 99	bool is_reduced = std::find(sorted_axes.begin(), sorted_axes.end(), static_cast<int64_t>(i)) !=	×
NEW 100	sorted_axes.end();	×
NEW 101	if (is_reduced) {	×
NEW 102	broadcast_strides.push_back(symbolic::zero());	×
NEW 103	} else {	×
NEW 104	broadcast_strides.push_back(reduced_strides[i]);	×
NEW 105	}	×
NEW 106	}	×
NEW 107	types::Tensor broadcast_tensor_type(element_type, shape_, broadcast_strides);	×
108
109	// Temporary buffers
110	std::string tmp_max_name = builder.find_new_name("_softmax_max");	×
111	builder.add_container(tmp_max_name, pointer_type);	×
112
113	std::string tmp_sub_name = builder.find_new_name("_softmax_sub");	×
114	builder.add_container(tmp_sub_name, pointer_type);	×
115
116	std::string tmp_exp_name = builder.find_new_name("_softmax_exp");	×
117	builder.add_container(tmp_exp_name, pointer_type);	×
118
119	std::string tmp_sum_name = builder.find_new_name("_softmax_sum");	×
120	builder.add_container(tmp_sum_name, pointer_type);	×
121
122	// Mallocs
123	{	×
124	symbolic::Expression bytes_elem = types::get_type_size(element_type, false);	×
125
126	symbolic::Expression bytes_full = bytes_elem;	×
127	for (auto& dim : this->shape_) {	×
128	bytes_full = symbolic::mul(dim, bytes_full);	×
129	}	×
130
131	symbolic::Expression bytes_reduced = bytes_elem;	×
132	for (auto& dim : reduced_shape) {	×
133	bytes_reduced = symbolic::mul(dim, bytes_reduced);	×
134	}	×
135
136	auto& alloc_block = builder.add_block_before(parent, block, {}, this->debug_info());	×
137
138	auto malloc_helper = [&](const std::string& name, const symbolic::Expression& size) {	×
139	auto& access = builder.add_access(alloc_block, name);	×
140	auto& malloc_node = builder.add_library_node<stdlib::MallocNode>(alloc_block, this->debug_info(), size);	×
141	builder	×
142	.add_computational_memlet(alloc_block, malloc_node, "_ret", access, {}, pointer_type, this->debug_info());	×
143	};	×
144
145	malloc_helper(tmp_max_name, bytes_reduced);	×
146	malloc_helper(tmp_sub_name, bytes_full);	×
147	malloc_helper(tmp_exp_name, bytes_full);	×
148	malloc_helper(tmp_sum_name, bytes_reduced);	×
149	}	×
150
151	// 1. Max(X) -> TmpMax
152	{	×
153	auto& max_block = builder.add_block_before(parent, block, {}, this->debug_info());	×
154	auto& max_node =	×
155	builder.add_library_node<MaxNode>(max_block, this->debug_info(), this->shape_, this->axes_, true);	×
156
157	auto& in_access = builder.add_access(max_block, in_node.data());	×
158	auto& out_access = builder.add_access(max_block, tmp_max_name);	×
159
160	builder	×
161	.add_computational_memlet(max_block, in_access, max_node, "X", {}, in_edge.base_type(), this->debug_info());	×
162	builder	×
NEW 163	.add_computational_memlet(max_block, max_node, "Y", out_access, {}, reduced_tensor_type, this->debug_info());	×
164	}	×
165
166	// 2. Sub(X, TmpMaxBcast) -> TmpSub
167	{	×
168	auto& sub_block = builder.add_block_before(parent, block, {}, this->debug_info());	×
169	auto& sub_node = builder.add_library_node<SubNode>(sub_block, this->debug_info(), this->shape_);	×
170
171	auto& in1_access = builder.add_access(sub_block, in_node.data());	×
NEW 172	auto& in2_access = builder.add_access(sub_block, tmp_max_name);	×
173	auto& out_access = builder.add_access(sub_block, tmp_sub_name);	×
174
175	builder	×
176	.add_computational_memlet(sub_block, in1_access, sub_node, "A", {}, in_edge.base_type(), this->debug_info());	×
177	builder	×
NEW 178	.add_computational_memlet(sub_block, in2_access, sub_node, "B", {}, broadcast_tensor_type, this->debug_info());	×
179	builder	×
NEW 180	.add_computational_memlet(sub_block, sub_node, "C", out_access, {}, in_edge.base_type(), this->debug_info());	×
181	}	×
182
183	// 3. Exp(TmpSub) -> TmpExp
184	{	×
185	auto& exp_block = builder.add_block_before(parent, block, {}, this->debug_info());	×
186	auto& exp_node = builder.add_library_node<ExpNode>(exp_block, this->debug_info(), this->shape_);	×
187
188	auto& in_access = builder.add_access(exp_block, tmp_sub_name);	×
189	auto& out_access = builder.add_access(exp_block, tmp_exp_name);	×
190
191	builder	×
NEW 192	.add_computational_memlet(exp_block, in_access, exp_node, "X", {}, in_edge.base_type(), this->debug_info());	×
193	builder	×
NEW 194	.add_computational_memlet(exp_block, exp_node, "Y", out_access, {}, in_edge.base_type(), this->debug_info());	×
195	}	×
196
197	// 4. Sum(TmpExp) -> TmpSum
198	{	×
199	auto& sum_block = builder.add_block_before(parent, block, {}, this->debug_info());	×
200	auto& sum_node =	×
201	builder.add_library_node<SumNode>(sum_block, this->debug_info(), this->shape_, this->axes_, true);	×
202
203	auto& in_access = builder.add_access(sum_block, tmp_exp_name);	×
204	auto& out_access = builder.add_access(sum_block, tmp_sum_name);	×
205
206	builder	×
NEW 207	.add_computational_memlet(sum_block, in_access, sum_node, "X", {}, in_edge.base_type(), this->debug_info());	×
208	builder	×
NEW 209	.add_computational_memlet(sum_block, sum_node, "Y", out_access, {}, reduced_tensor_type, this->debug_info());	×
210	}	×
211
212	// 5. Div(TmpExp, TmpSum) -> Output
213	{	×
214	auto& div_block = builder.add_block_before(parent, block, {}, this->debug_info());	×
215	auto& div_node = builder.add_library_node<DivNode>(div_block, this->debug_info(), this->shape_);	×
216
217	auto& in1_access = builder.add_access(div_block, tmp_exp_name);	×
NEW 218	auto& in2_access = builder.add_access(div_block, tmp_sum_name);	×
219	auto& out_access = builder.add_access(div_block, out_node.data());	×
220
221	builder	×
NEW 222	.add_computational_memlet(div_block, in1_access, div_node, "A", {}, in_edge.base_type(), this->debug_info());	×
223	builder	×
NEW 224	.add_computational_memlet(div_block, in2_access, div_node, "B", {}, broadcast_tensor_type, this->debug_info());	×
225	builder	×
226	.add_computational_memlet(div_block, div_node, "C", out_access, {}, out_edge.base_type(), this->debug_info());	×
227	}	×
228
229	// Cleanup
230	builder.remove_memlet(block, in_edge);	×
231	builder.remove_memlet(block, out_edge);	×
232	builder.remove_node(block, in_node);	×
233	builder.remove_node(block, out_node);	×
234	builder.remove_node(block, *this);	×
235
236	int last_index = parent.index(block);	×
237	builder.remove_child(parent, last_index);	×
238
239	return true;	×
240	}	×
241
242	std::unique_ptr<data_flow::DataFlowNode> SoftmaxNode::
243	clone(size_t element_id, const graph::Vertex vertex, data_flow::DataFlowGraph& parent) const {	×
244	return std::unique_ptr<	×
245	data_flow::DataFlowNode>(new SoftmaxNode(element_id, this->debug_info(), vertex, parent, this->shape_, this->axes_)	×
246	);	×
247	}	×
248
249	} // namespace tensor
250	} // namespace math
251	} // namespace sdfg

daisytuner / docc / 23086520240

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