26812260585

Committed 02 Jun 2026 08:49AM UTC coverage: 60.823% (-0.05%) from 60.869%

Build # 26812260585

Build Type

Pull #725

github

Committed by

web-flow

Commit Message

Merge eceff48b9 into cd25c9278

Pull Request Pull Request #725: Tensor node backport

Coverage Stats

599 of 1255 new or added lines in 51 files covered. (47.73%)

541 existing lines in 46 files now uncovered.

35094 of 57699 relevant lines covered (60.82%)

11078.43 hits per line

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

45.1

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

#include "sdfg/data_flow/library_nodes/math/tensor/batchnorm_node.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/library_nodes/math/cmath/cmath_node.h"
#include "sdfg/data_flow/library_nodes/math/tensor/tensor_expansion_utils.h"
#include "sdfg/structured_control_flow/block.h"
#include "sdfg/structured_control_flow/structured_loop.h"

namespace sdfg::math::tensor {


BatchNormNode::BatchNormNode(
    size_t element_id,
    const DebugInfo& debug_info,
    graph::Vertex vertex,
    data_flow::DataFlowGraph& parent,
    TensorLayout layout,
    QuantizationType quantization,
    data_flow::ImplementationType impl_type
)
    : TensorNode(
          element_id,
          debug_info,
          vertex,
          parent,
          LibraryNodeType_BatchNorm,
          {},
          {"Batch", "Var", "E", "Gamma", "Beta", "epsilon", "B_out"},
          std::move(impl_type)
      ),
      layout_(std::move(layout)), quantization_(quantization) {}

symbolic::SymbolSet BatchNormNode::symbols() const {
    symbolic::SymbolSet syms;
    layout_.collect_symbols(syms);
    return syms;
}

types::PrimitiveType BatchNormNode::quantization() const { return quantization_; }

void BatchNormNode::set_quantization(const types::PrimitiveType quant) { quantization_ = quant; }

void BatchNormNode::replace(const symbolic::Expression old_expression, const symbolic::Expression new_expression) {
    layout_.replace_symbols(old_expression, new_expression);
}

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

std::string BatchNormNode::toStr() const { return "BatchNorm(" + layout_.toStr() + ")"; }

bool BatchNormNode::expand(builder::StructuredSDFGBuilder& builder, analysis::AnalysisManager& analysis_manager) {
    // CPU implementation of batchnorm:
    if (false) {
        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;

        auto batch_in = find_usable_input_access_node(dataflow, *this, "Batch");
        auto& data_type = batch_in.memlet->base_type();
        types::Scalar scalar_type(data_type.primitive_type());
        types::Tensor tensor_1d(scalar_type, {num_features()}, {symbolic::one()}); // TODO verify / get from inputs
        std::string temp_var_prefix = "_batchn_tmp";
        int tmp_idx = 0;
        auto var_in = find_usable_input_access_node(dataflow, *this, "Var");
        auto e_in = find_usable_input_access_node(dataflow, *this, "E");
        auto gamma_in = find_usable_input_access_node(dataflow, *this, "Gamma");
        auto beta_in = find_usable_input_access_node(dataflow, *this, "Beta");
        auto result_ptr_in = find_usable_input_access_node(dataflow, *this, "B_out");
        auto eps_in = find_usable_input_access_node(dataflow, *this, "epsilon");

        auto& new_sequence = builder.add_sequence_before(parent, block, transition.assignments(), debug_info());

        auto loop_dims = create_maps(builder, layout_.shape(), new_sequence);

        auto& c_dim = loop_dims.at(1);
        std::vector<symbolic::Expression> c_subset{c_dim.indvar};
        auto interm_name = builder.find_new_name("_b_sqrt_div");
        builder.add_container(interm_name, scalar_type);
        auto& inter_block = builder.add_block_before(
            c_dim.seq, static_cast<structured_control_flow::ControlFlowNode&>(loop_dims.at(2).loop), {}, DebugInfo()
        );

        auto& var_elem_in = builder.add_access(inter_block, var_in.name);
        data_flow::AccessNode& epsilon_const = eps_in.is_const
                                                   ? builder.add_constant(inter_block, eps_in.name, scalar_type)
                                                   : builder.add_access(inter_block, eps_in.name);

        auto& add_eps_op = builder.add_tasklet(inter_block, data_flow::fp_add, "_out", {"var", "eps"}, debug_info());

        builder.add_computational_memlet(inter_block, var_elem_in, add_eps_op, "var", c_subset, tensor_1d);
        builder.add_computational_memlet(inter_block, epsilon_const, add_eps_op, "eps", {}, scalar_type);

        auto tmp_eps_name = create_temp_var(builder, temp_var_prefix, tmp_idx++, scalar_type);
        auto& tmp_eps = builder.add_access(inter_block, tmp_eps_name);

        builder.add_computational_memlet(inter_block, add_eps_op, "_out", tmp_eps, {}, scalar_type);

        auto tmp_sqrt_name = create_temp_var(builder, temp_var_prefix, tmp_idx++, scalar_type);
        auto& tmp_sqrt = builder.add_access(inter_block, tmp_sqrt_name);

        auto& sqrt_op = builder.add_library_node<
            cmath::CMathNode>(inter_block, debug_info(), cmath::CMathFunction::sqrt, data_type.primitive_type());

        builder.add_computational_memlet(inter_block, tmp_eps, sqrt_op, "_in1", {}, scalar_type);

        builder.add_computational_memlet(inter_block, sqrt_op, "_out", tmp_sqrt, {}, scalar_type);

        auto& one_const = builder.add_constant(inter_block, "1.0", scalar_type);
        auto& div_op = builder.add_tasklet(inter_block, data_flow::fp_div, "_out", {"one", "sqrt"});
        builder.add_computational_memlet(inter_block, one_const, div_op, "one", {}, scalar_type);
        builder.add_computational_memlet(inter_block, tmp_sqrt, div_op, "sqrt", {}, scalar_type);

        auto& interm_store = builder.add_access(inter_block, interm_name);
        builder.add_computational_memlet(inter_block, div_op, "_out", interm_store, {}, scalar_type);

        auto& innermost_dim = loop_dims.at(layout_.dims() - 1);

        std::vector<symbolic::Expression> innermost_subset;
        for (auto& builder_map_dim : loop_dims) {
            innermost_subset.push_back(builder_map_dim.indvar);
        }

        auto& innermost_block = builder.add_block(innermost_dim.seq);
        auto& x_in = builder.add_access(innermost_block, batch_in.name);
        auto& interm_in = builder.add_access(innermost_block, interm_name);
        auto& e_elem_in = builder.add_access(innermost_block, e_in.name);
        auto& gamma_elem_in = builder.add_access(innermost_block, gamma_in.name);
        auto& beta_elem_in = builder.add_access(innermost_block, beta_in.name);

        auto& result_ptr_out_elem = builder.add_access(innermost_block, result_ptr_in.name);

        auto& sub_op = builder.add_tasklet(innermost_block, data_flow::fp_sub, "_out", {"x", "e"}, debug_info());

        builder.add_computational_memlet(innermost_block, x_in, sub_op, "x", innermost_subset, data_type);
        builder.add_computational_memlet(innermost_block, e_elem_in, sub_op, "e", c_subset, tensor_1d);
        auto tmp_sub_name = create_temp_var(builder, temp_var_prefix, tmp_idx++, scalar_type);
        auto& tmp_sub = builder.add_access(innermost_block, tmp_sub_name);
        builder.add_computational_memlet(innermost_block, sub_op, "_out", tmp_sub, {}, scalar_type);

        auto& mul_interm_op =
            builder.add_tasklet(innermost_block, data_flow::fp_mul, "_out", {"num", "den"}, debug_info());

        builder.add_computational_memlet(innermost_block, tmp_sub, mul_interm_op, "num", {}, scalar_type);
        builder.add_computational_memlet(innermost_block, interm_in, mul_interm_op, "den", {}, scalar_type);
        auto tmp_interm = create_temp_var(builder, temp_var_prefix, tmp_idx++, scalar_type);
        auto& tmp_mul_interm = builder.add_access(innermost_block, tmp_interm);
        builder.add_computational_memlet(innermost_block, mul_interm_op, "_out", tmp_mul_interm, {}, scalar_type);

        auto& mul_gamma_op =
            builder.add_tasklet(innermost_block, data_flow::fp_mul, "_out", {"frac", "g"}, debug_info());

        builder.add_computational_memlet(innermost_block, tmp_mul_interm, mul_gamma_op, "frac", {}, scalar_type);
        builder.add_computational_memlet(innermost_block, gamma_elem_in, mul_gamma_op, "g", c_subset, tensor_1d);

        auto tmp_gamma = create_temp_var(builder, temp_var_prefix, tmp_idx++, scalar_type);
        auto& tmp_mul_gamma = builder.add_access(innermost_block, tmp_gamma);
        builder.add_computational_memlet(innermost_block, mul_gamma_op, "_out", tmp_mul_gamma, {}, scalar_type);

        auto& add_beta_op = builder.add_tasklet(innermost_block, data_flow::fp_add, "_out", {"_in", "b"}, debug_info());

        builder.add_computational_memlet(innermost_block, tmp_mul_gamma, add_beta_op, "_in", {}, scalar_type);
        builder.add_computational_memlet(innermost_block, beta_elem_in, add_beta_op, "b", c_subset, tensor_1d);
        builder.add_computational_memlet(
            innermost_block, add_beta_op, "_out", result_ptr_out_elem, innermost_subset, data_type
        );

        batch_in.remove_old(builder, block);
        var_in.remove_old(builder, block);
        e_in.remove_old(builder, block);
        eps_in.remove_old(builder, block);
        gamma_in.remove_old(builder, block);
        beta_in.remove_old(builder, block);
        result_ptr_in.remove_old(builder, block);

        builder.remove_node(block, *this);
        assert(dataflow.nodes().size() == 0 && "At expand time, no other nodes may be in the same graph");
        builder.remove_child(parent, index + 1);

        return true;
    } else {
        // GPU implementation of batchnorm:
        // Move sqrt and division into the innermost loop to enable more parallelism.
        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;

        auto batch_in = find_usable_input_access_node(dataflow, *this, "Batch");
        auto& data_type = batch_in.memlet->base_type();
        types::Scalar scalar_type(data_type.primitive_type());
        types::Tensor tensor_1d(scalar_type, {num_features()}, {symbolic::one()});
        std::string temp_var_prefix = "_batchn_tmp";
        int tmp_idx = 0;
        auto var_in = find_usable_input_access_node(dataflow, *this, "Var");
        auto e_in = find_usable_input_access_node(dataflow, *this, "E");
        auto gamma_in = find_usable_input_access_node(dataflow, *this, "Gamma");
        auto beta_in = find_usable_input_access_node(dataflow, *this, "Beta");
        auto result_ptr_in = find_usable_input_access_node(dataflow, *this, "B_out");
        auto eps_in = find_usable_input_access_node(dataflow, *this, "epsilon");

        auto& new_sequence = builder.add_sequence_before(parent, block, transition.assignments(), debug_info());

        auto loop_dims = create_maps(builder, layout_.shape(), new_sequence);

        auto& c_dim = loop_dims.at(1);
        std::vector<symbolic::Expression> c_subset{c_dim.indvar};

        auto& innermost_dim = loop_dims.at(layout_.dims() - 1);

        std::vector<symbolic::Expression> innermost_subset;
        for (auto& builder_map_dim : loop_dims) {
            innermost_subset.push_back(builder_map_dim.indvar);
        }

        auto& innermost_block = builder.add_block(innermost_dim.seq);

        // Access nodes
        auto& x_in = builder.add_access(innermost_block, batch_in.name);
        auto& var_elem_in = builder.add_access(innermost_block, var_in.name);
        data_flow::AccessNode& epsilon_const = eps_in.is_const
                                                   ? builder.add_constant(innermost_block, eps_in.name, scalar_type)
                                                   : builder.add_access(innermost_block, eps_in.name);
        auto& e_elem_in = builder.add_access(innermost_block, e_in.name);
        auto& gamma_elem_in = builder.add_access(innermost_block, gamma_in.name);
        auto& beta_elem_in = builder.add_access(innermost_block, beta_in.name);
        auto& result_ptr_out_elem = builder.add_access(innermost_block, result_ptr_in.name);

        // var[c] + eps
        auto& add_eps_op =
            builder.add_tasklet(innermost_block, data_flow::fp_add, "_out", {"var", "eps"}, debug_info());
        builder.add_computational_memlet(innermost_block, var_elem_in, add_eps_op, "var", c_subset, tensor_1d);
        builder.add_computational_memlet(innermost_block, epsilon_const, add_eps_op, "eps", {}, scalar_type);
        auto tmp_eps_name = create_temp_var(builder, temp_var_prefix, tmp_idx++, scalar_type);
        auto& tmp_eps = builder.add_access(innermost_block, tmp_eps_name);
        builder.add_computational_memlet(innermost_block, add_eps_op, "_out", tmp_eps, {}, scalar_type);

        // sqrt(var[c] + eps)
        auto tmp_sqrt_name = create_temp_var(builder, temp_var_prefix, tmp_idx++, scalar_type);
        auto& tmp_sqrt = builder.add_access(innermost_block, tmp_sqrt_name);
        auto& sqrt_op = builder.add_library_node<
            cmath::CMathNode>(innermost_block, debug_info(), cmath::CMathFunction::sqrt, data_type.primitive_type());
        builder.add_computational_memlet(innermost_block, tmp_eps, sqrt_op, "_in1", {}, scalar_type);
        builder.add_computational_memlet(innermost_block, sqrt_op, "_out", tmp_sqrt, {}, scalar_type);

        // 1.0 / sqrt(var[c] + eps)
        auto& one_const = builder.add_constant(innermost_block, "1.0", scalar_type);
        auto& div_op = builder.add_tasklet(innermost_block, data_flow::fp_div, "_out", {"one", "sqrt"});
        builder.add_computational_memlet(innermost_block, one_const, div_op, "one", {}, scalar_type);
        builder.add_computational_memlet(innermost_block, tmp_sqrt, div_op, "sqrt", {}, scalar_type);
        auto interm_name = create_temp_var(builder, temp_var_prefix, tmp_idx++, scalar_type);
        auto& interm_store = builder.add_access(innermost_block, interm_name);
        builder.add_computational_memlet(innermost_block, div_op, "_out", interm_store, {}, scalar_type);

        // x - e[c]
        auto& sub_op = builder.add_tasklet(innermost_block, data_flow::fp_sub, "_out", {"x", "e"}, debug_info());
        builder.add_computational_memlet(innermost_block, x_in, sub_op, "x", innermost_subset, data_type);
        builder.add_computational_memlet(innermost_block, e_elem_in, sub_op, "e", c_subset, tensor_1d);
        auto tmp_sub_name = create_temp_var(builder, temp_var_prefix, tmp_idx++, scalar_type);
        auto& tmp_sub = builder.add_access(innermost_block, tmp_sub_name);
        builder.add_computational_memlet(innermost_block, sub_op, "_out", tmp_sub, {}, scalar_type);

        // (x - e[c]) * (1/sqrt(var[c]+eps))
        auto& mul_interm_op =
            builder.add_tasklet(innermost_block, data_flow::fp_mul, "_out", {"num", "den"}, debug_info());
        builder.add_computational_memlet(innermost_block, tmp_sub, mul_interm_op, "num", {}, scalar_type);
        builder.add_computational_memlet(innermost_block, interm_store, mul_interm_op, "den", {}, scalar_type);
        auto tmp_interm = create_temp_var(builder, temp_var_prefix, tmp_idx++, scalar_type);
        auto& tmp_mul_interm = builder.add_access(innermost_block, tmp_interm);
        builder.add_computational_memlet(innermost_block, mul_interm_op, "_out", tmp_mul_interm, {}, scalar_type);

        // * gamma[c]
        auto& mul_gamma_op =
            builder.add_tasklet(innermost_block, data_flow::fp_mul, "_out", {"frac", "g"}, debug_info());
        builder.add_computational_memlet(innermost_block, tmp_mul_interm, mul_gamma_op, "frac", {}, scalar_type);
        builder.add_computational_memlet(innermost_block, gamma_elem_in, mul_gamma_op, "g", c_subset, tensor_1d);
        auto tmp_gamma = create_temp_var(builder, temp_var_prefix, tmp_idx++, scalar_type);
        auto& tmp_mul_gamma = builder.add_access(innermost_block, tmp_gamma);
        builder.add_computational_memlet(innermost_block, mul_gamma_op, "_out", tmp_mul_gamma, {}, scalar_type);

        // + beta[c]
        auto& add_beta_op = builder.add_tasklet(innermost_block, data_flow::fp_add, "_out", {"_in", "b"}, debug_info());
        builder.add_computational_memlet(innermost_block, tmp_mul_gamma, add_beta_op, "_in", {}, scalar_type);
        builder.add_computational_memlet(innermost_block, beta_elem_in, add_beta_op, "b", c_subset, tensor_1d);
        builder.add_computational_memlet(
            innermost_block, add_beta_op, "_out", result_ptr_out_elem, innermost_subset, data_type
        );

        batch_in.remove_old(builder, block);
        var_in.remove_old(builder, block);
        e_in.remove_old(builder, block);
        eps_in.remove_old(builder, block);
        gamma_in.remove_old(builder, block);
        beta_in.remove_old(builder, block);
        result_ptr_in.remove_old(builder, block);

        builder.remove_node(block, *this);
        assert(dataflow.nodes().size() == 0 && "At expand time, no other nodes may be in the same graph");
        builder.remove_child(parent, index + 1);

        return true;
    }
}

symbolic::Expression BatchNormNode::flop() const {
    auto inner_elems = symbolic::mul(layout_.get_dim_innermost(0), layout_.get_dim_innermost(1));
    auto outer_elems = symbolic::mul(layout_.shape().at(0), layout_.shape().at(1));

    // (x-e) * sqrt_pre_calc * g + b = 4 flops
    auto inner_flops = symbolic::mul(symbolic::integer(4), inner_elems);
    // sqrt_pre_calc = 1/sqrt(var + eps) // 3 flops
    auto outer_flops = symbolic::mul(symbolic::add(inner_flops, symbolic::integer(3)), outer_elems);
    return outer_flops;
}

data_flow::PointerAccessType BatchNormNode::pointer_access_type(int input_idx) const {
    if (input_idx >= 0 && input_idx <= 4) {
        return data_flow::PointerAccessMeta::create_read_only(symbolic::__nullptr__(), true);
    } else if (input_idx == 6) {
        return data_flow::PointerAccessMeta::create_full_write_only(symbolic::__nullptr__(), true);
    } else {
        return TensorNode::pointer_access_type(input_idx);
    }
}

nlohmann::json BatchNormNodeSerializer::serialize(const data_flow::LibraryNode& library_node) {
    auto& node = static_cast<const BatchNormNode&>(library_node);
    nlohmann::json j;

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

    node.batch_layout().serialize_to_json(j["batch_layout"]);

    j["batch_quant"] = node.quantization();

    return j;
}

data_flow::LibraryNode& BatchNormNodeSerializer::deserialize(
    const nlohmann::json& j, builder::StructuredSDFGBuilder& builder, structured_control_flow::Block& parent
) {
    auto layout = TensorLayout::deserialize_from_json(j.at("batch_layout"));
    auto quant = j.at("batch_quant").get<types::PrimitiveType>();

    serializer::JSONSerializer serializer;
    auto deb_info = serializer.json_to_debug_info(j.at("debug_info"));

    return builder.add_library_node<BatchNormNode>(parent, deb_info, layout, quant);
}

} // namespace sdfg::math::tensor

1	#include "sdfg/data_flow/library_nodes/math/tensor/batchnorm_node.h"
2
3	#include "sdfg/analysis/scope_analysis.h"
4	#include "sdfg/builder/structured_sdfg_builder.h"
5	#include "sdfg/data_flow/access_node.h"
6	#include "sdfg/data_flow/library_nodes/math/cmath/cmath_node.h"
7	#include "sdfg/data_flow/library_nodes/math/tensor/tensor_expansion_utils.h"
8	#include "sdfg/structured_control_flow/block.h"
9	#include "sdfg/structured_control_flow/structured_loop.h"
10
11	namespace sdfg::math::tensor {
12
13
14	BatchNormNode::BatchNormNode(
15	size_t element_id,
16	const DebugInfo& debug_info,
17	graph::Vertex vertex,
18	data_flow::DataFlowGraph& parent,
19	TensorLayout layout,
20	QuantizationType quantization,
21	data_flow::ImplementationType impl_type
22	)
23	: TensorNode(	22✔
24	element_id,	22✔
25	debug_info,	22✔
26	vertex,	22✔
27	parent,	22✔
28	LibraryNodeType_BatchNorm,	22✔
29	{},	22✔
30	{"Batch", "Var", "E", "Gamma", "Beta", "epsilon", "B_out"},	22✔
31	std::move(impl_type)	22✔
32	),	22✔
33	layout_(std::move(layout)), quantization_(quantization) {}	22✔
34
35	symbolic::SymbolSet BatchNormNode::symbols() const {	24✔
36	symbolic::SymbolSet syms;	24✔
37	layout_.collect_symbols(syms);	24✔
38	return syms;	24✔
39	}	24✔
40
41	types::PrimitiveType BatchNormNode::quantization() const { return quantization_; }	×
42
43	void BatchNormNode::set_quantization(const types::PrimitiveType quant) { quantization_ = quant; }	×
44
45	void BatchNormNode::replace(const symbolic::Expression old_expression, const symbolic::Expression new_expression) {	×
46	layout_.replace_symbols(old_expression, new_expression);	×
47	}	×
48
49	std::unique_ptr<data_flow::DataFlowNode> BatchNormNode::
50	clone(size_t element_id, const graph::Vertex vertex, data_flow::DataFlowGraph& parent) const {	×
51	return std::unique_ptr<data_flow::DataFlowNode>(new BatchNormNode(	×
52	element_id, debug_info(), vertex, parent, this->layout_, this->quantization_, this->implementation_type_	×
53	));	×
54	}	×
55
56	std::string BatchNormNode::toStr() const { return "BatchNorm(" + layout_.toStr() + ")"; }	×
57
58	bool BatchNormNode::expand(builder::StructuredSDFGBuilder& builder, analysis::AnalysisManager& analysis_manager) {	1✔
59	// CPU implementation of batchnorm:
60	if (false) {	1✔
61	auto& dataflow = this->get_parent();	×
62	auto& block = static_cast<structured_control_flow::Block&>(*dataflow.get_parent());	×
63
64	auto& scope_analysis = analysis_manager.get<analysis::ScopeAnalysis>();	×
65	auto& parent = static_cast<structured_control_flow::Sequence&>(*scope_analysis.parent_scope(&block));	×
66	int index = parent.index(block);	×
67	auto& transition = parent.at(index).second;	×
68
69	auto batch_in = find_usable_input_access_node(dataflow, *this, "Batch");	×
70	auto& data_type = batch_in.memlet->base_type();	×
71	types::Scalar scalar_type(data_type.primitive_type());	×
72	types::Tensor tensor_1d(scalar_type, {num_features()}, {symbolic::one()}); // TODO verify / get from inputs	×
73	std::string temp_var_prefix = "_batchn_tmp";	×
74	int tmp_idx = 0;	×
75	auto var_in = find_usable_input_access_node(dataflow, *this, "Var");	×
76	auto e_in = find_usable_input_access_node(dataflow, *this, "E");	×
77	auto gamma_in = find_usable_input_access_node(dataflow, *this, "Gamma");	×
78	auto beta_in = find_usable_input_access_node(dataflow, *this, "Beta");	×
79	auto result_ptr_in = find_usable_input_access_node(dataflow, *this, "B_out");	×
80	auto eps_in = find_usable_input_access_node(dataflow, *this, "epsilon");	×
81
82	auto& new_sequence = builder.add_sequence_before(parent, block, transition.assignments(), debug_info());	×
83
84	auto loop_dims = create_maps(builder, layout_.shape(), new_sequence);	×
85
86	auto& c_dim = loop_dims.at(1);	×
87	std::vector<symbolic::Expression> c_subset{c_dim.indvar};	×
88	auto interm_name = builder.find_new_name("_b_sqrt_div");	×
89	builder.add_container(interm_name, scalar_type);	×
90	auto& inter_block = builder.add_block_before(	×
91	c_dim.seq, static_cast<structured_control_flow::ControlFlowNode&>(loop_dims.at(2).loop), {}, DebugInfo()	×
92	);	×
93
94	auto& var_elem_in = builder.add_access(inter_block, var_in.name);	×
95	data_flow::AccessNode& epsilon_const = eps_in.is_const	×
96	? builder.add_constant(inter_block, eps_in.name, scalar_type)	×
97	: builder.add_access(inter_block, eps_in.name);	×
98
99	auto& add_eps_op = builder.add_tasklet(inter_block, data_flow::fp_add, "_out", {"var", "eps"}, debug_info());	×
100
101	builder.add_computational_memlet(inter_block, var_elem_in, add_eps_op, "var", c_subset, tensor_1d);	×
102	builder.add_computational_memlet(inter_block, epsilon_const, add_eps_op, "eps", {}, scalar_type);	×
103
104	auto tmp_eps_name = create_temp_var(builder, temp_var_prefix, tmp_idx++, scalar_type);	×
105	auto& tmp_eps = builder.add_access(inter_block, tmp_eps_name);	×
106
107	builder.add_computational_memlet(inter_block, add_eps_op, "_out", tmp_eps, {}, scalar_type);	×
108
109	auto tmp_sqrt_name = create_temp_var(builder, temp_var_prefix, tmp_idx++, scalar_type);	×
110	auto& tmp_sqrt = builder.add_access(inter_block, tmp_sqrt_name);	×
111
112	auto& sqrt_op = builder.add_library_node<	×
113	cmath::CMathNode>(inter_block, debug_info(), cmath::CMathFunction::sqrt, data_type.primitive_type());	×
114
115	builder.add_computational_memlet(inter_block, tmp_eps, sqrt_op, "_in1", {}, scalar_type);	×
116
117	builder.add_computational_memlet(inter_block, sqrt_op, "_out", tmp_sqrt, {}, scalar_type);	×
118
119	auto& one_const = builder.add_constant(inter_block, "1.0", scalar_type);	×
120	auto& div_op = builder.add_tasklet(inter_block, data_flow::fp_div, "_out", {"one", "sqrt"});	×
121	builder.add_computational_memlet(inter_block, one_const, div_op, "one", {}, scalar_type);	×
122	builder.add_computational_memlet(inter_block, tmp_sqrt, div_op, "sqrt", {}, scalar_type);	×
123
124	auto& interm_store = builder.add_access(inter_block, interm_name);	×
125	builder.add_computational_memlet(inter_block, div_op, "_out", interm_store, {}, scalar_type);	×
126
127	auto& innermost_dim = loop_dims.at(layout_.dims() - 1);	×
128
129	std::vector<symbolic::Expression> innermost_subset;	×
130	for (auto& builder_map_dim : loop_dims) {	×
131	innermost_subset.push_back(builder_map_dim.indvar);	×
132	}	×
133
134	auto& innermost_block = builder.add_block(innermost_dim.seq);	×
135	auto& x_in = builder.add_access(innermost_block, batch_in.name);	×
136	auto& interm_in = builder.add_access(innermost_block, interm_name);	×
137	auto& e_elem_in = builder.add_access(innermost_block, e_in.name);	×
138	auto& gamma_elem_in = builder.add_access(innermost_block, gamma_in.name);	×
139	auto& beta_elem_in = builder.add_access(innermost_block, beta_in.name);	×
140
141	auto& result_ptr_out_elem = builder.add_access(innermost_block, result_ptr_in.name);	×
142
143	auto& sub_op = builder.add_tasklet(innermost_block, data_flow::fp_sub, "_out", {"x", "e"}, debug_info());	×
144
145	builder.add_computational_memlet(innermost_block, x_in, sub_op, "x", innermost_subset, data_type);	×
146	builder.add_computational_memlet(innermost_block, e_elem_in, sub_op, "e", c_subset, tensor_1d);	×
147	auto tmp_sub_name = create_temp_var(builder, temp_var_prefix, tmp_idx++, scalar_type);	×
148	auto& tmp_sub = builder.add_access(innermost_block, tmp_sub_name);	×
149	builder.add_computational_memlet(innermost_block, sub_op, "_out", tmp_sub, {}, scalar_type);	×
150
151	auto& mul_interm_op =	×
152	builder.add_tasklet(innermost_block, data_flow::fp_mul, "_out", {"num", "den"}, debug_info());	×
153
154	builder.add_computational_memlet(innermost_block, tmp_sub, mul_interm_op, "num", {}, scalar_type);	×
155	builder.add_computational_memlet(innermost_block, interm_in, mul_interm_op, "den", {}, scalar_type);	×
156	auto tmp_interm = create_temp_var(builder, temp_var_prefix, tmp_idx++, scalar_type);	×
157	auto& tmp_mul_interm = builder.add_access(innermost_block, tmp_interm);	×
158	builder.add_computational_memlet(innermost_block, mul_interm_op, "_out", tmp_mul_interm, {}, scalar_type);	×
159
160	auto& mul_gamma_op =	×
161	builder.add_tasklet(innermost_block, data_flow::fp_mul, "_out", {"frac", "g"}, debug_info());	×
162
163	builder.add_computational_memlet(innermost_block, tmp_mul_interm, mul_gamma_op, "frac", {}, scalar_type);	×
164	builder.add_computational_memlet(innermost_block, gamma_elem_in, mul_gamma_op, "g", c_subset, tensor_1d);	×
165
166	auto tmp_gamma = create_temp_var(builder, temp_var_prefix, tmp_idx++, scalar_type);	×
167	auto& tmp_mul_gamma = builder.add_access(innermost_block, tmp_gamma);	×
168	builder.add_computational_memlet(innermost_block, mul_gamma_op, "_out", tmp_mul_gamma, {}, scalar_type);	×
169
170	auto& add_beta_op = builder.add_tasklet(innermost_block, data_flow::fp_add, "_out", {"_in", "b"}, debug_info());	×
171
172	builder.add_computational_memlet(innermost_block, tmp_mul_gamma, add_beta_op, "_in", {}, scalar_type);	×
173	builder.add_computational_memlet(innermost_block, beta_elem_in, add_beta_op, "b", c_subset, tensor_1d);	×
174	builder.add_computational_memlet(	×
175	innermost_block, add_beta_op, "_out", result_ptr_out_elem, innermost_subset, data_type	×
176	);	×
177
178	batch_in.remove_old(builder, block);	×
179	var_in.remove_old(builder, block);	×
180	e_in.remove_old(builder, block);	×
181	eps_in.remove_old(builder, block);	×
182	gamma_in.remove_old(builder, block);	×
183	beta_in.remove_old(builder, block);	×
184	result_ptr_in.remove_old(builder, block);	×
185
186	builder.remove_node(block, *this);	×
187	assert(dataflow.nodes().size() == 0 && "At expand time, no other nodes may be in the same graph");	×
188	builder.remove_child(parent, index + 1);	×
189
190	return true;	×
191	} else {	1✔
192	// GPU implementation of batchnorm:
193	// Move sqrt and division into the innermost loop to enable more parallelism.
194	auto& dataflow = this->get_parent();	1✔
195	auto& block = static_cast<structured_control_flow::Block&>(*dataflow.get_parent());	1✔
196
197	auto& scope_analysis = analysis_manager.get<analysis::ScopeAnalysis>();	1✔
198	auto& parent = static_cast<structured_control_flow::Sequence&>(*scope_analysis.parent_scope(&block));	1✔
199	int index = parent.index(block);	1✔
200	auto& transition = parent.at(index).second;	1✔
201
202	auto batch_in = find_usable_input_access_node(dataflow, *this, "Batch");	1✔
203	auto& data_type = batch_in.memlet->base_type();	1✔
204	types::Scalar scalar_type(data_type.primitive_type());	1✔
205	types::Tensor tensor_1d(scalar_type, {num_features()}, {symbolic::one()});	1✔
206	std::string temp_var_prefix = "_batchn_tmp";	1✔
207	int tmp_idx = 0;	1✔
208	auto var_in = find_usable_input_access_node(dataflow, *this, "Var");	1✔
209	auto e_in = find_usable_input_access_node(dataflow, *this, "E");	1✔
210	auto gamma_in = find_usable_input_access_node(dataflow, *this, "Gamma");	1✔
211	auto beta_in = find_usable_input_access_node(dataflow, *this, "Beta");	1✔
212	auto result_ptr_in = find_usable_input_access_node(dataflow, *this, "B_out");	1✔
213	auto eps_in = find_usable_input_access_node(dataflow, *this, "epsilon");	1✔
214
215	auto& new_sequence = builder.add_sequence_before(parent, block, transition.assignments(), debug_info());	1✔
216
217	auto loop_dims = create_maps(builder, layout_.shape(), new_sequence);	1✔
218
219	auto& c_dim = loop_dims.at(1);	1✔
220	std::vector<symbolic::Expression> c_subset{c_dim.indvar};	1✔
221
222	auto& innermost_dim = loop_dims.at(layout_.dims() - 1);	1✔
223
224	std::vector<symbolic::Expression> innermost_subset;	1✔
225	for (auto& builder_map_dim : loop_dims) {	4✔
226	innermost_subset.push_back(builder_map_dim.indvar);	4✔
227	}	4✔
228
229	auto& innermost_block = builder.add_block(innermost_dim.seq);	1✔
230
231	// Access nodes
232	auto& x_in = builder.add_access(innermost_block, batch_in.name);	1✔
233	auto& var_elem_in = builder.add_access(innermost_block, var_in.name);	1✔
234	data_flow::AccessNode& epsilon_const = eps_in.is_const	1✔
235	? builder.add_constant(innermost_block, eps_in.name, scalar_type)	1✔
236	: builder.add_access(innermost_block, eps_in.name);	1✔
237	auto& e_elem_in = builder.add_access(innermost_block, e_in.name);	1✔
238	auto& gamma_elem_in = builder.add_access(innermost_block, gamma_in.name);	1✔
239	auto& beta_elem_in = builder.add_access(innermost_block, beta_in.name);	1✔
240	auto& result_ptr_out_elem = builder.add_access(innermost_block, result_ptr_in.name);	1✔
241
242	// var[c] + eps
243	auto& add_eps_op =	1✔
244	builder.add_tasklet(innermost_block, data_flow::fp_add, "_out", {"var", "eps"}, debug_info());	1✔
245	builder.add_computational_memlet(innermost_block, var_elem_in, add_eps_op, "var", c_subset, tensor_1d);	1✔
246	builder.add_computational_memlet(innermost_block, epsilon_const, add_eps_op, "eps", {}, scalar_type);	1✔
247	auto tmp_eps_name = create_temp_var(builder, temp_var_prefix, tmp_idx++, scalar_type);	1✔
248	auto& tmp_eps = builder.add_access(innermost_block, tmp_eps_name);	1✔
249	builder.add_computational_memlet(innermost_block, add_eps_op, "_out", tmp_eps, {}, scalar_type);	1✔
250
251	// sqrt(var[c] + eps)
252	auto tmp_sqrt_name = create_temp_var(builder, temp_var_prefix, tmp_idx++, scalar_type);	1✔
253	auto& tmp_sqrt = builder.add_access(innermost_block, tmp_sqrt_name);	1✔
254	auto& sqrt_op = builder.add_library_node<	1✔
255	cmath::CMathNode>(innermost_block, debug_info(), cmath::CMathFunction::sqrt, data_type.primitive_type());	1✔
256	builder.add_computational_memlet(innermost_block, tmp_eps, sqrt_op, "_in1", {}, scalar_type);	1✔
257	builder.add_computational_memlet(innermost_block, sqrt_op, "_out", tmp_sqrt, {}, scalar_type);	1✔
258
259	// 1.0 / sqrt(var[c] + eps)
260	auto& one_const = builder.add_constant(innermost_block, "1.0", scalar_type);	1✔
261	auto& div_op = builder.add_tasklet(innermost_block, data_flow::fp_div, "_out", {"one", "sqrt"});	1✔
262	builder.add_computational_memlet(innermost_block, one_const, div_op, "one", {}, scalar_type);	1✔
263	builder.add_computational_memlet(innermost_block, tmp_sqrt, div_op, "sqrt", {}, scalar_type);	1✔
264	auto interm_name = create_temp_var(builder, temp_var_prefix, tmp_idx++, scalar_type);	1✔
265	auto& interm_store = builder.add_access(innermost_block, interm_name);	1✔
266	builder.add_computational_memlet(innermost_block, div_op, "_out", interm_store, {}, scalar_type);	1✔
267
268	// x - e[c]
269	auto& sub_op = builder.add_tasklet(innermost_block, data_flow::fp_sub, "_out", {"x", "e"}, debug_info());	1✔
270	builder.add_computational_memlet(innermost_block, x_in, sub_op, "x", innermost_subset, data_type);	1✔
271	builder.add_computational_memlet(innermost_block, e_elem_in, sub_op, "e", c_subset, tensor_1d);	1✔
272	auto tmp_sub_name = create_temp_var(builder, temp_var_prefix, tmp_idx++, scalar_type);	1✔
273	auto& tmp_sub = builder.add_access(innermost_block, tmp_sub_name);	1✔
274	builder.add_computational_memlet(innermost_block, sub_op, "_out", tmp_sub, {}, scalar_type);	1✔
275
276	// (x - e[c]) * (1/sqrt(var[c]+eps))
277	auto& mul_interm_op =	1✔
278	builder.add_tasklet(innermost_block, data_flow::fp_mul, "_out", {"num", "den"}, debug_info());	1✔
279	builder.add_computational_memlet(innermost_block, tmp_sub, mul_interm_op, "num", {}, scalar_type);	1✔
280	builder.add_computational_memlet(innermost_block, interm_store, mul_interm_op, "den", {}, scalar_type);	1✔
281	auto tmp_interm = create_temp_var(builder, temp_var_prefix, tmp_idx++, scalar_type);	1✔
282	auto& tmp_mul_interm = builder.add_access(innermost_block, tmp_interm);	1✔
283	builder.add_computational_memlet(innermost_block, mul_interm_op, "_out", tmp_mul_interm, {}, scalar_type);	1✔
284
285	// * gamma[c]
286	auto& mul_gamma_op =	1✔
287	builder.add_tasklet(innermost_block, data_flow::fp_mul, "_out", {"frac", "g"}, debug_info());	1✔
288	builder.add_computational_memlet(innermost_block, tmp_mul_interm, mul_gamma_op, "frac", {}, scalar_type);	1✔
289	builder.add_computational_memlet(innermost_block, gamma_elem_in, mul_gamma_op, "g", c_subset, tensor_1d);	1✔
290	auto tmp_gamma = create_temp_var(builder, temp_var_prefix, tmp_idx++, scalar_type);	1✔
291	auto& tmp_mul_gamma = builder.add_access(innermost_block, tmp_gamma);	1✔
292	builder.add_computational_memlet(innermost_block, mul_gamma_op, "_out", tmp_mul_gamma, {}, scalar_type);	1✔
293
294	// + beta[c]
295	auto& add_beta_op = builder.add_tasklet(innermost_block, data_flow::fp_add, "_out", {"_in", "b"}, debug_info());	1✔
296	builder.add_computational_memlet(innermost_block, tmp_mul_gamma, add_beta_op, "_in", {}, scalar_type);	1✔
297	builder.add_computational_memlet(innermost_block, beta_elem_in, add_beta_op, "b", c_subset, tensor_1d);	1✔
298	builder.add_computational_memlet(	1✔
299	innermost_block, add_beta_op, "_out", result_ptr_out_elem, innermost_subset, data_type	1✔
300	);	1✔
301
302	batch_in.remove_old(builder, block);	1✔
303	var_in.remove_old(builder, block);	1✔
304	e_in.remove_old(builder, block);	1✔
305	eps_in.remove_old(builder, block);	1✔
306	gamma_in.remove_old(builder, block);	1✔
307	beta_in.remove_old(builder, block);	1✔
308	result_ptr_in.remove_old(builder, block);	1✔
309
310	builder.remove_node(block, *this);	1✔
311	assert(dataflow.nodes().size() == 0 && "At expand time, no other nodes may be in the same graph");	1✔
312	builder.remove_child(parent, index + 1);	1✔
313
314	return true;	1✔
315	}	1✔
316	}	1✔
317
318	symbolic::Expression BatchNormNode::flop() const {	×
319	auto inner_elems = symbolic::mul(layout_.get_dim_innermost(0), layout_.get_dim_innermost(1));	×
320	auto outer_elems = symbolic::mul(layout_.shape().at(0), layout_.shape().at(1));	×
321
322	// (x-e) * sqrt_pre_calc * g + b = 4 flops
323	auto inner_flops = symbolic::mul(symbolic::integer(4), inner_elems);	×
324	// sqrt_pre_calc = 1/sqrt(var + eps) // 3 flops
325	auto outer_flops = symbolic::mul(symbolic::add(inner_flops, symbolic::integer(3)), outer_elems);	×
326	return outer_flops;	×
327	}	×
328
NEW 329	data_flow::PointerAccessType BatchNormNode::pointer_access_type(int input_idx) const {	×
NEW 330	if (input_idx >= 0 && input_idx <= 4) {	×
NEW 331	return data_flow::PointerAccessMeta::create_read_only(symbolic::__nullptr__(), true);	×
NEW 332	} else if (input_idx == 6) {	×
NEW 333	return data_flow::PointerAccessMeta::create_full_write_only(symbolic::__nullptr__(), true);	×
NEW 334	} else {	×
NEW 335	return TensorNode::pointer_access_type(input_idx);	×
NEW 336	}	×
NEW 337	}	×
338
339	nlohmann::json BatchNormNodeSerializer::serialize(const data_flow::LibraryNode& library_node) {	×
340	auto& node = static_cast<const BatchNormNode&>(library_node);	×
341	nlohmann::json j;	×
342
343	j["code"] = node.code().value();	×
344
345	node.batch_layout().serialize_to_json(j["batch_layout"]);	×
346
347	j["batch_quant"] = node.quantization();	×
348
349	return j;	×
350	}	×
351
352	data_flow::LibraryNode& BatchNormNodeSerializer::deserialize(
353	const nlohmann::json& j, builder::StructuredSDFGBuilder& builder, structured_control_flow::Block& parent
354	) {	×
355	auto layout = TensorLayout::deserialize_from_json(j.at("batch_layout"));	×
356	auto quant = j.at("batch_quant").get<types::PrimitiveType>();	×
357
358	serializer::JSONSerializer serializer;	×
359	auto deb_info = serializer.json_to_debug_info(j.at("debug_info"));	×
360
361	return builder.add_library_node<BatchNormNode>(parent, deb_info, layout, quant);	×
362	}	×
363
364	} // namespace sdfg::math::tensor

daisytuner / docc / 26812260585

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