17637380013

Committed 11 Sep 2025 07:29AM UTC coverage: 59.755% (+0.6%) from 59.145%

Build # 17637380013

Build Type

push

github

Committed by

web-flow

Commit Message

New debug info (#210)

* initial draft

* update data structure and construction logic

* finalize DebugInfo draft

* fix tests

* Update serializer and fix tests

* fix append bug

* update data structure

* sdfg builder update

* const ref vectors

* update implementation and partial tests

* compiling state

* update serializer interface

* update dot test

* reset interface to debug_info in json to maintain compatibility with tools

* first review batch

* second batch of changes

* merge fixes

Run Details

777 of 1111 new or added lines in 46 files covered. (69.94%)

11 existing lines in 11 files now uncovered.

9755 of 16325 relevant lines covered (59.75%)

115.06 hits per line

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

0.0

/src/data_flow/library_nodes/math/ml/layer_normalization.cpp

#include "sdfg/data_flow/library_nodes/math/ml/layer_normalization.h"

#include "sdfg/analysis/analysis.h"
#include "sdfg/analysis/scope_analysis.h"
#include "sdfg/builder/structured_sdfg_builder.h"

namespace sdfg {
namespace math {
namespace ml {

LayerNormalizationNode::LayerNormalizationNode(
    size_t element_id,
    const DebugInfoRegion &debug_info,
    const graph::Vertex vertex,
    data_flow::DataFlowGraph &parent,
    int axis,
    const std::string &epsilon
)
    : MathNode(
          element_id,
          debug_info,
          vertex,
          parent,
          LibraryNodeType_LayerNormalization,
          {"Y"},
          {"X", "Scale", "B"},
          data_flow::ImplementationType_NONE
      ),
      axis_(axis), epsilon_(epsilon) {}

void LayerNormalizationNode::validate(const Function &) const { /* TODO */ }

bool LayerNormalizationNode::expand(builder::StructuredSDFGBuilder &builder, analysis::AnalysisManager &analysis_manager) {
    auto &dataflow = this->get_parent();
    auto &block = static_cast<structured_control_flow::Block &>(*dataflow.get_parent());

    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;

    // Locate edges
    const data_flow::Memlet *iedge_input = nullptr;
    const data_flow::Memlet *iedge_scale = nullptr;
    const data_flow::Memlet *iedge_bias = nullptr;
    const data_flow::Memlet *oedge_output = nullptr;
    for (const auto &edge : dataflow.in_edges(*this)) {
        if (edge.dst_conn() == "X") {
            iedge_input = &edge;
        } else if (edge.dst_conn() == "Scale") {
            iedge_scale = &edge;
        } else if (edge.dst_conn() == "B") {
            iedge_bias = &edge;
        }
    }
    for (const auto &edge : dataflow.out_edges(*this)) {
        if (edge.src_conn() == "Y") {
            oedge_output = &edge;
        }
    }
    if (!iedge_input || !iedge_scale || !oedge_output) return false;

    std::string input_name = static_cast<const data_flow::AccessNode &>(iedge_input->src()).data();
    std::string scale_name = static_cast<const data_flow::AccessNode &>(iedge_scale->src()).data();
    std::string bias_name;
    if (iedge_bias) {
        bias_name = static_cast<const data_flow::AccessNode &>(iedge_bias->src()).data();
    }
    std::string output_name = static_cast<const data_flow::AccessNode &>(oedge_output->dst()).data();

    // Create new sequence before
    auto &new_sequence = builder.add_sequence_before(
        parent, block, transition.assignments(), builder.debug_info().get_region(block.debug_info().indices())
    );
    structured_control_flow::Sequence *last_scope = &new_sequence;

    // Create maps over output subset dims (parallel dims)
    data_flow::Subset domain_begin = oedge_output->begin_subset();
    data_flow::Subset domain_end = oedge_output->end_subset();

    std::vector<symbolic::Expression> loop_syms;
    structured_control_flow::Map *last_map = nullptr;
    for (size_t d = 0; d < domain_begin.size(); ++d) {
        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 = domain_begin[d];
        auto update = symbolic::add(indvar, symbolic::one());
        auto cond = symbolic::Lt(indvar, symbolic::add(domain_end[d], symbolic::one()));
        last_map = &builder.add_map(
            *last_scope,
            indvar,
            cond,
            init,
            update,
            structured_control_flow::ScheduleType_Sequential::create(),
            {},
            builder.subject().debug_info().get_region(block.debug_info().indices())
        );
        last_scope = &last_map->root();
        loop_syms.push_back(indvar);
    }

    // Initialize temp variables for mean and variance
    std::string mean_temp_name = builder.find_new_name("_mean_tmp");
    std::string var_temp_name = builder.find_new_name("_var_tmp");
    std::string count_temp_name = builder.find_new_name("_count_tmp");
    types::Scalar temp_type(types::PrimitiveType::Float);
    builder.add_container(mean_temp_name, temp_type);
    builder.add_container(var_temp_name, temp_type);
    builder.add_container(count_temp_name, temp_type);

    auto &init_block = builder.add_block(*last_scope);
    auto &init_mean_tasklet = builder.add_tasklet(init_block, data_flow::TaskletCode::assign, "_out", {"0.0f"});
    auto &init_var_tasklet = builder.add_tasklet(init_block, data_flow::TaskletCode::assign, "_out", {"0.0f"});
    auto &init_count_tasklet = builder.add_tasklet(init_block, data_flow::TaskletCode::assign, "_out", {"0.0f"});
    auto &mean_access_init = builder.add_access(init_block, mean_temp_name);
    auto &var_access_init = builder.add_access(init_block, var_temp_name);
    auto &count_access_init = builder.add_access(init_block, count_temp_name);
    builder.add_computational_memlet(init_block, init_mean_tasklet, "_out", mean_access_init, {}, temp_type);
    builder.add_computational_memlet(init_block, init_var_tasklet, "_out", var_access_init, {}, temp_type);
    builder.add_computational_memlet(init_block, init_count_tasklet, "_out", count_access_init, {}, temp_type);

    // Add reduction for loop to compute mean and variance
    symbolic::Expression red_begin;
    symbolic::Expression red_end;
    if (axis_ >= 0) {
        red_begin = iedge_input->begin_subset()[axis_];
        red_end = iedge_input->end_subset()[axis_];
    } else {
        red_begin = iedge_input->begin_subset().back();
        red_end = iedge_input->end_subset().back();
    }
    std::string red_name = builder.find_new_name("_i");
    builder.add_container(red_name, types::Scalar(types::PrimitiveType::UInt64));
    auto red_indvar = symbolic::symbol(red_name);
    auto red_init = red_begin;
    auto red_update = symbolic::add(red_indvar, symbolic::one());
    auto red_cond = symbolic::Lt(red_indvar, symbolic::add(red_end, symbolic::one()));
    auto red_map = &builder.add_for(
        *last_scope,
        red_indvar,
        red_cond,
        red_init,
        red_update,
        {},
        builder.debug_info().get_region(block.debug_info().indices())
    );

    // Create innermost block for mean and variance computation
    auto &compute_block = builder.add_block(red_map->root());

    // Create access nodes
    auto &input_access = builder.add_access(compute_block, input_name);
    auto &mean_access_in = builder.add_access(compute_block, mean_temp_name);
    auto &mean_access_out = builder.add_access(compute_block, mean_temp_name);
    auto &var_access_in = builder.add_access(compute_block, var_temp_name);
    auto &var_access_out = builder.add_access(compute_block, var_temp_name);
    auto &count_access_in = builder.add_access(compute_block, count_temp_name);
    auto &count_access_out = builder.add_access(compute_block, count_temp_name);

    // Create index expressions for input
    std::vector<symbolic::Expression> input_subset = loop_syms;

    // Replace the reduction axis index with the reduction variable for input
    if (axis_ >= 0 && axis_ < static_cast<int>(input_subset.size())) {
        input_subset.insert(input_subset.begin() + axis_, red_indvar);
    } else if (axis_ < 0) {
        input_subset.push_back(red_indvar);
    }

    // Add to mean (reduction)
    auto &add_mean_tasklet = builder.add_tasklet(compute_block, data_flow::TaskletCode::add, "_out", {"_in1", "_in2"});
    builder.add_computational_memlet(
        compute_block, input_access, add_mean_tasklet, "_in1", input_subset, iedge_input->base_type()
    );
    builder.add_computational_memlet(compute_block, mean_access_in, add_mean_tasklet, "_in2", {}, temp_type);
    builder.add_computational_memlet(compute_block, add_mean_tasklet, "_out", mean_access_out, {}, temp_type);

    // Add to variance (reduction of squared values)
    auto &square_tasklet = builder.add_tasklet(compute_block, data_flow::TaskletCode::mul, "_out", {"_in1", "_in2"});
    auto &square_temp_access = builder.add_access(compute_block, builder.find_new_name("_square_tmp"));
    builder.add_computational_memlet(
        compute_block, input_access, square_tasklet, "_in1", input_subset, iedge_input->base_type()
    );
    builder.add_computational_memlet(
        compute_block, input_access, square_tasklet, "_in2", input_subset, iedge_input->base_type()
    );
    builder.add_computational_memlet(compute_block, square_tasklet, "_out", square_temp_access, {}, temp_type);

    auto &add_var_tasklet = builder.add_tasklet(compute_block, data_flow::TaskletCode::add, "_out", {"_in1", "_in2"});
    builder.add_computational_memlet(compute_block, square_temp_access, add_var_tasklet, "_in1", {}, temp_type);
    builder.add_computational_memlet(compute_block, var_access_in, add_var_tasklet, "_in2", {}, temp_type);
    builder.add_computational_memlet(compute_block, add_var_tasklet, "_out", var_access_out, {}, temp_type);

    // Increment count
    auto &add_count_tasklet = builder.add_tasklet(compute_block, data_flow::TaskletCode::add, "_out", {"_in1", "1.0f"});
    builder.add_computational_memlet(compute_block, count_access_in, add_count_tasklet, "_in1", {}, temp_type);
    builder.add_computational_memlet(compute_block, add_count_tasklet, "_out", count_access_out, {}, temp_type);

    // Create normalization block
    auto &norm_block = builder.add_block(*last_scope);

    // Create access nodes for normalization
    auto &mean_access_norm = builder.add_access(norm_block, mean_temp_name);
    auto &var_access_norm = builder.add_access(norm_block, var_temp_name);
    auto &count_access_norm = builder.add_access(norm_block, count_temp_name);
    auto &input_access_norm = builder.add_access(norm_block, input_name);
    auto &scale_access_norm = builder.add_access(norm_block, scale_name);
    auto &output_access_norm = builder.add_access(norm_block, output_name);

    // Compute mean by dividing sum by count
    auto &div_mean_tasklet = builder.add_tasklet(norm_block, data_flow::TaskletCode::div, "_out", {"_in1", "_in2"});
    auto &mean_result_access = builder.add_access(norm_block, builder.find_new_name("_mean_result"));
    builder.add_computational_memlet(norm_block, mean_access_norm, div_mean_tasklet, "_in1", {}, temp_type);
    builder.add_computational_memlet(norm_block, count_access_norm, div_mean_tasklet, "_in2", {}, temp_type);
    builder.add_computational_memlet(norm_block, div_mean_tasklet, "_out", mean_result_access, {}, temp_type);

    // Compute variance by dividing sum of squares by count and subtracting mean squared
    auto &div_var_tasklet = builder.add_tasklet(norm_block, data_flow::TaskletCode::div, "_out", {"_in1", "_in2"});
    auto &var_div_access = builder.add_access(norm_block, builder.find_new_name("_var_div"));
    builder.add_computational_memlet(norm_block, var_access_norm, div_var_tasklet, "_in1", {}, temp_type);
    builder.add_computational_memlet(norm_block, count_access_norm, div_var_tasklet, "_in2", {}, temp_type);
    builder.add_computational_memlet(norm_block, div_var_tasklet, "_out", var_div_access, {}, temp_type);

    auto &square_mean_tasklet = builder.add_tasklet(norm_block, data_flow::TaskletCode::mul, "_out", {"_in", "_in"});
    auto &mean_squared_access = builder.add_access(norm_block, builder.find_new_name("_mean_squared"));
    builder.add_computational_memlet(norm_block, mean_result_access, square_mean_tasklet, "_in", {}, temp_type);
    builder.add_computational_memlet(norm_block, square_mean_tasklet, "_out", mean_squared_access, {}, temp_type);

    auto &sub_var_tasklet = builder.add_tasklet(norm_block, data_flow::TaskletCode::sub, "_out", {"_in1", "_in2"});
    auto &var_result_access = builder.add_access(norm_block, builder.find_new_name("_var_result"));
    builder.add_computational_memlet(norm_block, var_div_access, sub_var_tasklet, "_in1", {}, temp_type);
    builder.add_computational_memlet(norm_block, mean_squared_access, sub_var_tasklet, "_in2", {}, temp_type);
    builder.add_computational_memlet(norm_block, sub_var_tasklet, "_out", var_result_access, {}, temp_type);

    // Add epsilon to variance and compute standard deviation
    auto &add_epsilon_tasklet =
        builder.add_tasklet(norm_block, data_flow::TaskletCode::add, "_out", {"_in1", epsilon_});
    auto &var_eps_access = builder.add_access(norm_block, builder.find_new_name("_var_eps"));
    builder.add_computational_memlet(norm_block, var_result_access, add_epsilon_tasklet, "_in1", {}, temp_type);
    builder.add_computational_memlet(norm_block, add_epsilon_tasklet, "_out", var_eps_access, {}, temp_type);

    auto &sqrt_tasklet = builder.add_tasklet(norm_block, data_flow::TaskletCode::sqrt, "_out", {"_in"});
    auto &std_dev_access = builder.add_access(norm_block, builder.find_new_name("_std_dev"));
    builder.add_computational_memlet(norm_block, var_eps_access, sqrt_tasklet, "_in", {}, temp_type);
    builder.add_computational_memlet(norm_block, sqrt_tasklet, "_out", std_dev_access, {}, temp_type);

    // Normalize: (x - mean) / std_dev
    auto &sub_norm_tasklet = builder.add_tasklet(norm_block, data_flow::TaskletCode::sub, "_out", {"_in1", "_in2"});
    auto &centered_access = builder.add_access(norm_block, builder.find_new_name("_centered"));
    builder.add_computational_memlet(
        norm_block, input_access_norm, sub_norm_tasklet, "_in1", loop_syms, iedge_input->base_type()
    );
    builder.add_computational_memlet(norm_block, mean_result_access, sub_norm_tasklet, "_in2", {}, temp_type);
    builder.add_computational_memlet(norm_block, sub_norm_tasklet, "_out", centered_access, {}, temp_type);

    auto &div_norm_tasklet = builder.add_tasklet(norm_block, data_flow::TaskletCode::div, "_out", {"_in1", "_in2"});
    auto &normalized_access = builder.add_access(norm_block, builder.find_new_name("_normalized"));
    builder.add_computational_memlet(norm_block, centered_access, div_norm_tasklet, "_in1", {}, temp_type);
    builder.add_computational_memlet(norm_block, std_dev_access, div_norm_tasklet, "_in2", {}, temp_type);
    builder.add_computational_memlet(norm_block, div_norm_tasklet, "_out", normalized_access, {}, temp_type);

    // Apply scale and bias: scale * normalized + bias (if bias is provided)
    auto &mul_scale_tasklet = builder.add_tasklet(norm_block, data_flow::TaskletCode::mul, "_out", {"_in1", "_in2"});
    auto &scaled_access = builder.add_access(norm_block, builder.find_new_name("_scaled"));
    builder.add_computational_memlet(norm_block, normalized_access, mul_scale_tasklet, "_in1", {}, temp_type);
    builder.add_computational_memlet(
        norm_block, scale_access_norm, mul_scale_tasklet, "_in2", loop_syms, iedge_scale->base_type()
    );
    builder.add_computational_memlet(norm_block, mul_scale_tasklet, "_out", scaled_access, {}, temp_type);

    if (iedge_bias) {
        // Add bias if provided
        auto &bias_access_norm = builder.add_access(norm_block, bias_name);
        auto &add_bias_tasklet = builder.add_tasklet(norm_block, data_flow::TaskletCode::add, "_out", {"_in1", "_in2"});
        builder.add_computational_memlet(norm_block, scaled_access, add_bias_tasklet, "_in1", {}, temp_type);
        builder.add_computational_memlet(
            norm_block, bias_access_norm, add_bias_tasklet, "_in2", loop_syms, iedge_bias->base_type()
        );
        builder.add_computational_memlet(
            norm_block, add_bias_tasklet, "_out", output_access_norm, loop_syms, oedge_output->base_type()
        );
    } else {
        // No bias, just assign scaled result to output
        auto &assign_tasklet = builder.add_tasklet(norm_block, data_flow::TaskletCode::assign, "_out", {"_in"});
        builder.add_computational_memlet(norm_block, scaled_access, assign_tasklet, "_in", {}, temp_type);
        builder.add_computational_memlet(
            norm_block, assign_tasklet, "_out", output_access_norm, loop_syms, oedge_output->base_type()
        );
    }

    // Cleanup old block
    builder.remove_memlet(block, *iedge_input);
    builder.remove_memlet(block, *iedge_scale);
    if (iedge_bias) {
        builder.remove_memlet(block, *iedge_bias);
    }
    builder.remove_memlet(block, *oedge_output);
    builder.remove_node(block, *this);
    builder.remove_child(parent, index + 1);

    return true;
}

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

nlohmann::json LayerNormalizationNodeSerializer::serialize(const data_flow::LibraryNode &library_node) {
    const LayerNormalizationNode &node = static_cast<const LayerNormalizationNode &>(library_node);
    nlohmann::json j;

    j["code"] = node.code().value();
    j["axis"] = node.axis();
    j["epsilon"] = node.epsilon();

    return j;
}

data_flow::LibraryNode &LayerNormalizationNodeSerializer::deserialize(
    const nlohmann::json &j, builder::StructuredSDFGBuilder &builder, structured_control_flow::Block &parent
) {
    auto code = j["code"].get<std::string>();
    if (code != LibraryNodeType_LayerNormalization.value()) {
        throw std::runtime_error("Invalid library node code");
    }

    sdfg::serializer::JSONSerializer serializer;
    DebugInfoRegion debug_info = serializer.json_to_debug_info_region(j["debug_info"], builder.debug_info());

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

    return builder.add_library_node<LayerNormalizationNode>(parent, debug_info, axis, epsilon);
}

} // namespace ml
} // namespace math
} // namespace sdfg

1	#include "sdfg/data_flow/library_nodes/math/ml/layer_normalization.h"
2
3	#include "sdfg/analysis/analysis.h"
4	#include "sdfg/analysis/scope_analysis.h"
5	#include "sdfg/builder/structured_sdfg_builder.h"
6
7	namespace sdfg {
8	namespace math {
9	namespace ml {
10
11	LayerNormalizationNode::LayerNormalizationNode(	×
12	size_t element_id,
13	const DebugInfoRegion &debug_info,
14	const graph::Vertex vertex,
15	data_flow::DataFlowGraph &parent,
16	int axis,
17	const std::string &epsilon
18	)
19	: MathNode(	×
20	element_id,	×
21	debug_info,	×
22	vertex,	×
23	parent,	×
24	LibraryNodeType_LayerNormalization,
25	{"Y"},	×
26	{"X", "Scale", "B"},	×
27	data_flow::ImplementationType_NONE
28	),
29	axis_(axis), epsilon_(epsilon) {}	×
30
31	void LayerNormalizationNode::validate(const Function &) const { /* TODO */ }	×
32
33	bool LayerNormalizationNode::expand(builder::StructuredSDFGBuilder &builder, analysis::AnalysisManager &analysis_manager) {	×
34	auto &dataflow = this->get_parent();	×
35	auto &block = static_cast<structured_control_flow::Block &>(*dataflow.get_parent());	×
36
37	auto &scope_analysis = analysis_manager.get<analysis::ScopeAnalysis>();	×
38	auto &parent = static_cast<structured_control_flow::Sequence &>(*scope_analysis.parent_scope(&block));	×
39	int index = parent.index(block);	×
40	auto &transition = parent.at(index).second;	×
41
42	// Locate edges
43	const data_flow::Memlet *iedge_input = nullptr;	×
44	const data_flow::Memlet *iedge_scale = nullptr;	×
45	const data_flow::Memlet *iedge_bias = nullptr;	×
46	const data_flow::Memlet *oedge_output = nullptr;	×
47	for (const auto &edge : dataflow.in_edges(*this)) {	×
48	if (edge.dst_conn() == "X") {	×
49	iedge_input = &edge;	×
50	} else if (edge.dst_conn() == "Scale") {	×
51	iedge_scale = &edge;	×
52	} else if (edge.dst_conn() == "B") {	×
53	iedge_bias = &edge;	×
54	}	×
55	}
56	for (const auto &edge : dataflow.out_edges(*this)) {	×
57	if (edge.src_conn() == "Y") {	×
58	oedge_output = &edge;	×
59	}	×
60	}
61	if (!iedge_input \|\| !iedge_scale \|\| !oedge_output) return false;	×
62
63	std::string input_name = static_cast<const data_flow::AccessNode &>(iedge_input->src()).data();	×
64	std::string scale_name = static_cast<const data_flow::AccessNode &>(iedge_scale->src()).data();	×
65	std::string bias_name;	×
66	if (iedge_bias) {	×
67	bias_name = static_cast<const data_flow::AccessNode &>(iedge_bias->src()).data();	×
68	}	×
69	std::string output_name = static_cast<const data_flow::AccessNode &>(oedge_output->dst()).data();	×
70
71	// Create new sequence before
NEW 72	auto &new_sequence = builder.add_sequence_before(	×
NEW 73	parent, block, transition.assignments(), builder.debug_info().get_region(block.debug_info().indices())	×
74	);
UNCOV 75	structured_control_flow::Sequence *last_scope = &new_sequence;	×
76
77	// Create maps over output subset dims (parallel dims)
78	data_flow::Subset domain_begin = oedge_output->begin_subset();	×
79	data_flow::Subset domain_end = oedge_output->end_subset();	×
80
81	std::vector<symbolic::Expression> loop_syms;	×
82	structured_control_flow::Map *last_map = nullptr;	×
83	for (size_t d = 0; d < domain_begin.size(); ++d) {	×
84	std::string indvar_str = builder.find_new_name("_i");	×
85	builder.add_container(indvar_str, types::Scalar(types::PrimitiveType::UInt64));	×
86	auto indvar = symbolic::symbol(indvar_str);	×
87	auto init = domain_begin[d];	×
88	auto update = symbolic::add(indvar, symbolic::one());	×
89	auto cond = symbolic::Lt(indvar, symbolic::add(domain_end[d], symbolic::one()));	×
90	last_map = &builder.add_map(	×
91	*last_scope,	×
92	indvar,
93	cond,
94	init,
95	update,
96	structured_control_flow::ScheduleType_Sequential::create(),	×
97	{},	×
NEW 98	builder.subject().debug_info().get_region(block.debug_info().indices())	×
99	);
100	last_scope = &last_map->root();	×
101	loop_syms.push_back(indvar);	×
102	}	×
103
104	// Initialize temp variables for mean and variance
105	std::string mean_temp_name = builder.find_new_name("_mean_tmp");	×
106	std::string var_temp_name = builder.find_new_name("_var_tmp");	×
107	std::string count_temp_name = builder.find_new_name("_count_tmp");	×
108	types::Scalar temp_type(types::PrimitiveType::Float);	×
109	builder.add_container(mean_temp_name, temp_type);	×
110	builder.add_container(var_temp_name, temp_type);	×
111	builder.add_container(count_temp_name, temp_type);	×
112
113	auto &init_block = builder.add_block(*last_scope);	×
114	auto &init_mean_tasklet = builder.add_tasklet(init_block, data_flow::TaskletCode::assign, "_out", {"0.0f"});	×
115	auto &init_var_tasklet = builder.add_tasklet(init_block, data_flow::TaskletCode::assign, "_out", {"0.0f"});	×
116	auto &init_count_tasklet = builder.add_tasklet(init_block, data_flow::TaskletCode::assign, "_out", {"0.0f"});	×
117	auto &mean_access_init = builder.add_access(init_block, mean_temp_name);	×
118	auto &var_access_init = builder.add_access(init_block, var_temp_name);	×
119	auto &count_access_init = builder.add_access(init_block, count_temp_name);	×
120	builder.add_computational_memlet(init_block, init_mean_tasklet, "_out", mean_access_init, {}, temp_type);	×
121	builder.add_computational_memlet(init_block, init_var_tasklet, "_out", var_access_init, {}, temp_type);	×
122	builder.add_computational_memlet(init_block, init_count_tasklet, "_out", count_access_init, {}, temp_type);	×
123
124	// Add reduction for loop to compute mean and variance
125	symbolic::Expression red_begin;	×
126	symbolic::Expression red_end;	×
127	if (axis_ >= 0) {	×
128	red_begin = iedge_input->begin_subset()[axis_];	×
129	red_end = iedge_input->end_subset()[axis_];	×
130	} else {	×
131	red_begin = iedge_input->begin_subset().back();	×
132	red_end = iedge_input->end_subset().back();	×
133	}
134	std::string red_name = builder.find_new_name("_i");	×
135	builder.add_container(red_name, types::Scalar(types::PrimitiveType::UInt64));	×
136	auto red_indvar = symbolic::symbol(red_name);	×
137	auto red_init = red_begin;	×
138	auto red_update = symbolic::add(red_indvar, symbolic::one());	×
139	auto red_cond = symbolic::Lt(red_indvar, symbolic::add(red_end, symbolic::one()));	×
NEW 140	auto red_map = &builder.add_for(	×
NEW 141	*last_scope,	×
142	red_indvar,
143	red_cond,
144	red_init,
145	red_update,
NEW 146	{},	×
NEW 147	builder.debug_info().get_region(block.debug_info().indices())	×
148	);
149
150	// Create innermost block for mean and variance computation
151	auto &compute_block = builder.add_block(red_map->root());	×
152
153	// Create access nodes
154	auto &input_access = builder.add_access(compute_block, input_name);	×
155	auto &mean_access_in = builder.add_access(compute_block, mean_temp_name);	×
156	auto &mean_access_out = builder.add_access(compute_block, mean_temp_name);	×
157	auto &var_access_in = builder.add_access(compute_block, var_temp_name);	×
158	auto &var_access_out = builder.add_access(compute_block, var_temp_name);	×
159	auto &count_access_in = builder.add_access(compute_block, count_temp_name);	×
160	auto &count_access_out = builder.add_access(compute_block, count_temp_name);	×
161
162	// Create index expressions for input
163	std::vector<symbolic::Expression> input_subset = loop_syms;	×
164
165	// Replace the reduction axis index with the reduction variable for input
166	if (axis_ >= 0 && axis_ < static_cast<int>(input_subset.size())) {	×
167	input_subset.insert(input_subset.begin() + axis_, red_indvar);	×
168	} else if (axis_ < 0) {	×
169	input_subset.push_back(red_indvar);	×
170	}	×
171
172	// Add to mean (reduction)
173	auto &add_mean_tasklet = builder.add_tasklet(compute_block, data_flow::TaskletCode::add, "_out", {"_in1", "_in2"});	×
174	builder.add_computational_memlet(	×
175	compute_block, input_access, add_mean_tasklet, "_in1", input_subset, iedge_input->base_type()	×
176	);
177	builder.add_computational_memlet(compute_block, mean_access_in, add_mean_tasklet, "_in2", {}, temp_type);	×
178	builder.add_computational_memlet(compute_block, add_mean_tasklet, "_out", mean_access_out, {}, temp_type);	×
179
180	// Add to variance (reduction of squared values)
181	auto &square_tasklet = builder.add_tasklet(compute_block, data_flow::TaskletCode::mul, "_out", {"_in1", "_in2"});	×
182	auto &square_temp_access = builder.add_access(compute_block, builder.find_new_name("_square_tmp"));	×
183	builder.add_computational_memlet(	×
184	compute_block, input_access, square_tasklet, "_in1", input_subset, iedge_input->base_type()	×
185	);
186	builder.add_computational_memlet(	×
187	compute_block, input_access, square_tasklet, "_in2", input_subset, iedge_input->base_type()	×
188	);
189	builder.add_computational_memlet(compute_block, square_tasklet, "_out", square_temp_access, {}, temp_type);	×
190
191	auto &add_var_tasklet = builder.add_tasklet(compute_block, data_flow::TaskletCode::add, "_out", {"_in1", "_in2"});	×
192	builder.add_computational_memlet(compute_block, square_temp_access, add_var_tasklet, "_in1", {}, temp_type);	×
193	builder.add_computational_memlet(compute_block, var_access_in, add_var_tasklet, "_in2", {}, temp_type);	×
194	builder.add_computational_memlet(compute_block, add_var_tasklet, "_out", var_access_out, {}, temp_type);	×
195
196	// Increment count
197	auto &add_count_tasklet = builder.add_tasklet(compute_block, data_flow::TaskletCode::add, "_out", {"_in1", "1.0f"});	×
198	builder.add_computational_memlet(compute_block, count_access_in, add_count_tasklet, "_in1", {}, temp_type);	×
199	builder.add_computational_memlet(compute_block, add_count_tasklet, "_out", count_access_out, {}, temp_type);	×
200
201	// Create normalization block
202	auto &norm_block = builder.add_block(*last_scope);	×
203
204	// Create access nodes for normalization
205	auto &mean_access_norm = builder.add_access(norm_block, mean_temp_name);	×
206	auto &var_access_norm = builder.add_access(norm_block, var_temp_name);	×
207	auto &count_access_norm = builder.add_access(norm_block, count_temp_name);	×
208	auto &input_access_norm = builder.add_access(norm_block, input_name);	×
209	auto &scale_access_norm = builder.add_access(norm_block, scale_name);	×
210	auto &output_access_norm = builder.add_access(norm_block, output_name);	×
211
212	// Compute mean by dividing sum by count
213	auto &div_mean_tasklet = builder.add_tasklet(norm_block, data_flow::TaskletCode::div, "_out", {"_in1", "_in2"});	×
214	auto &mean_result_access = builder.add_access(norm_block, builder.find_new_name("_mean_result"));	×
215	builder.add_computational_memlet(norm_block, mean_access_norm, div_mean_tasklet, "_in1", {}, temp_type);	×
216	builder.add_computational_memlet(norm_block, count_access_norm, div_mean_tasklet, "_in2", {}, temp_type);	×
217	builder.add_computational_memlet(norm_block, div_mean_tasklet, "_out", mean_result_access, {}, temp_type);	×
218
219	// Compute variance by dividing sum of squares by count and subtracting mean squared
220	auto &div_var_tasklet = builder.add_tasklet(norm_block, data_flow::TaskletCode::div, "_out", {"_in1", "_in2"});	×
221	auto &var_div_access = builder.add_access(norm_block, builder.find_new_name("_var_div"));	×
222	builder.add_computational_memlet(norm_block, var_access_norm, div_var_tasklet, "_in1", {}, temp_type);	×
223	builder.add_computational_memlet(norm_block, count_access_norm, div_var_tasklet, "_in2", {}, temp_type);	×
224	builder.add_computational_memlet(norm_block, div_var_tasklet, "_out", var_div_access, {}, temp_type);	×
225
226	auto &square_mean_tasklet = builder.add_tasklet(norm_block, data_flow::TaskletCode::mul, "_out", {"_in", "_in"});	×
227	auto &mean_squared_access = builder.add_access(norm_block, builder.find_new_name("_mean_squared"));	×
228	builder.add_computational_memlet(norm_block, mean_result_access, square_mean_tasklet, "_in", {}, temp_type);	×
229	builder.add_computational_memlet(norm_block, square_mean_tasklet, "_out", mean_squared_access, {}, temp_type);	×
230
231	auto &sub_var_tasklet = builder.add_tasklet(norm_block, data_flow::TaskletCode::sub, "_out", {"_in1", "_in2"});	×
232	auto &var_result_access = builder.add_access(norm_block, builder.find_new_name("_var_result"));	×
233	builder.add_computational_memlet(norm_block, var_div_access, sub_var_tasklet, "_in1", {}, temp_type);	×
234	builder.add_computational_memlet(norm_block, mean_squared_access, sub_var_tasklet, "_in2", {}, temp_type);	×
235	builder.add_computational_memlet(norm_block, sub_var_tasklet, "_out", var_result_access, {}, temp_type);	×
236
237	// Add epsilon to variance and compute standard deviation
238	auto &add_epsilon_tasklet =	×
239	builder.add_tasklet(norm_block, data_flow::TaskletCode::add, "_out", {"_in1", epsilon_});	×
240	auto &var_eps_access = builder.add_access(norm_block, builder.find_new_name("_var_eps"));	×
241	builder.add_computational_memlet(norm_block, var_result_access, add_epsilon_tasklet, "_in1", {}, temp_type);	×
242	builder.add_computational_memlet(norm_block, add_epsilon_tasklet, "_out", var_eps_access, {}, temp_type);	×
243
244	auto &sqrt_tasklet = builder.add_tasklet(norm_block, data_flow::TaskletCode::sqrt, "_out", {"_in"});	×
245	auto &std_dev_access = builder.add_access(norm_block, builder.find_new_name("_std_dev"));	×
246	builder.add_computational_memlet(norm_block, var_eps_access, sqrt_tasklet, "_in", {}, temp_type);	×
247	builder.add_computational_memlet(norm_block, sqrt_tasklet, "_out", std_dev_access, {}, temp_type);	×
248
249	// Normalize: (x - mean) / std_dev
250	auto &sub_norm_tasklet = builder.add_tasklet(norm_block, data_flow::TaskletCode::sub, "_out", {"_in1", "_in2"});	×
251	auto &centered_access = builder.add_access(norm_block, builder.find_new_name("_centered"));	×
252	builder.add_computational_memlet(	×
253	norm_block, input_access_norm, sub_norm_tasklet, "_in1", loop_syms, iedge_input->base_type()	×
254	);
255	builder.add_computational_memlet(norm_block, mean_result_access, sub_norm_tasklet, "_in2", {}, temp_type);	×
256	builder.add_computational_memlet(norm_block, sub_norm_tasklet, "_out", centered_access, {}, temp_type);	×
257
258	auto &div_norm_tasklet = builder.add_tasklet(norm_block, data_flow::TaskletCode::div, "_out", {"_in1", "_in2"});	×
259	auto &normalized_access = builder.add_access(norm_block, builder.find_new_name("_normalized"));	×
260	builder.add_computational_memlet(norm_block, centered_access, div_norm_tasklet, "_in1", {}, temp_type);	×
261	builder.add_computational_memlet(norm_block, std_dev_access, div_norm_tasklet, "_in2", {}, temp_type);	×
262	builder.add_computational_memlet(norm_block, div_norm_tasklet, "_out", normalized_access, {}, temp_type);	×
263
264	// Apply scale and bias: scale * normalized + bias (if bias is provided)
265	auto &mul_scale_tasklet = builder.add_tasklet(norm_block, data_flow::TaskletCode::mul, "_out", {"_in1", "_in2"});	×
266	auto &scaled_access = builder.add_access(norm_block, builder.find_new_name("_scaled"));	×
267	builder.add_computational_memlet(norm_block, normalized_access, mul_scale_tasklet, "_in1", {}, temp_type);	×
268	builder.add_computational_memlet(	×
269	norm_block, scale_access_norm, mul_scale_tasklet, "_in2", loop_syms, iedge_scale->base_type()	×
270	);
271	builder.add_computational_memlet(norm_block, mul_scale_tasklet, "_out", scaled_access, {}, temp_type);	×
272
273	if (iedge_bias) {	×
274	// Add bias if provided
275	auto &bias_access_norm = builder.add_access(norm_block, bias_name);	×
276	auto &add_bias_tasklet = builder.add_tasklet(norm_block, data_flow::TaskletCode::add, "_out", {"_in1", "_in2"});	×
277	builder.add_computational_memlet(norm_block, scaled_access, add_bias_tasklet, "_in1", {}, temp_type);	×
278	builder.add_computational_memlet(	×
279	norm_block, bias_access_norm, add_bias_tasklet, "_in2", loop_syms, iedge_bias->base_type()	×
280	);
281	builder.add_computational_memlet(	×
282	norm_block, add_bias_tasklet, "_out", output_access_norm, loop_syms, oedge_output->base_type()	×
283	);
284	} else {	×
285	// No bias, just assign scaled result to output
286	auto &assign_tasklet = builder.add_tasklet(norm_block, data_flow::TaskletCode::assign, "_out", {"_in"});	×
287	builder.add_computational_memlet(norm_block, scaled_access, assign_tasklet, "_in", {}, temp_type);	×
288	builder.add_computational_memlet(	×
289	norm_block, assign_tasklet, "_out", output_access_norm, loop_syms, oedge_output->base_type()	×
290	);
291	}
292
293	// Cleanup old block
294	builder.remove_memlet(block, *iedge_input);	×
295	builder.remove_memlet(block, *iedge_scale);	×
296	if (iedge_bias) {	×
297	builder.remove_memlet(block, *iedge_bias);	×
298	}	×
299	builder.remove_memlet(block, *oedge_output);	×
300	builder.remove_node(block, *this);	×
301	builder.remove_child(parent, index + 1);	×
302
303	return true;	×
304	}	×
305
306	std::unique_ptr<data_flow::DataFlowNode> LayerNormalizationNode::
307	clone(size_t element_id, const graph::Vertex vertex, data_flow::DataFlowGraph &parent) const {	×
308	return std::unique_ptr<data_flow::DataFlowNode>(	×
309	new LayerNormalizationNode(element_id, this->debug_info(), vertex, parent, axis_, epsilon_)	×
310	);
311	}	×
312
313	nlohmann::json LayerNormalizationNodeSerializer::serialize(const data_flow::LibraryNode &library_node) {	×
314	const LayerNormalizationNode &node = static_cast<const LayerNormalizationNode &>(library_node);	×
315	nlohmann::json j;	×
316
317	j["code"] = node.code().value();	×
318	j["axis"] = node.axis();	×
319	j["epsilon"] = node.epsilon();	×
320
321	return j;	×
322	}	×
323
324	data_flow::LibraryNode &LayerNormalizationNodeSerializer::deserialize(	×
325	const nlohmann::json &j, builder::StructuredSDFGBuilder &builder, structured_control_flow::Block &parent
326	) {
327	auto code = j["code"].get<std::string>();	×
328	if (code != LibraryNodeType_LayerNormalization.value()) {	×
329	throw std::runtime_error("Invalid library node code");	×
330	}
331
332	sdfg::serializer::JSONSerializer serializer;	×
NEW 333	DebugInfoRegion debug_info = serializer.json_to_debug_info_region(j["debug_info"], builder.debug_info());	×
334
335	auto axis = j["axis"].get<int>();	×
336	auto epsilon = j["epsilon"].get<std::string>();	×
337
338	return builder.add_library_node<LayerNormalizationNode>(parent, debug_info, axis, epsilon);	×
339	}	×
340
341	} // namespace ml
342	} // namespace math
343	} // namespace sdfg

daisytuner / sdfglib / 17637380013

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