28793626242

Committed 06 Jul 2026 01:04PM UTC coverage: 62.962% (+0.2%) from 62.801%

Build # 28793626242

Build Type

Pull #740

github

Committed by

web-flow

Commit Message

Merge b7e03d389 into 23e67a4ab

Pull Request Pull Request #740: Expand pass

Coverage Stats

594 of 962 new or added lines in 31 files covered. (61.75%)

36 existing lines in 10 files now uncovered.

40557 of 64415 relevant lines covered (62.96%)

972.09 hits per line

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

50.0

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

#include "sdfg/data_flow/library_nodes/math/tensor/batchnorm_node.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);
}

void BatchNormNode::replace(const symbolic::ExpressionMapping& replacements) { layout_.replace_symbols(replacements); }

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() + ")"; }

passes::LibNodeExpander::ExpandOutcome BatchNormNode::
    expand(passes::LibNodeExpander::ExpandContext& context, structured_control_flow::Block& block) {
    auto& dataflow = this->get_parent();

    auto* batch_iedge = dataflow.in_edge_for_connector(*this, "Batch");
    auto& data_type = batch_iedge->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;

    //{"Batch", "Var", "E", "Gamma", "Beta", "epsilon", "B_out"},

    constexpr auto BATCH_IDX = 0;
    constexpr auto VAR_IDX = 1;
    constexpr auto E_IDX = 2;
    constexpr auto GAMMA_IDX = 3;
    constexpr auto BETA_IDX = 4;
    constexpr auto EPS_IDX = 5;
    constexpr auto B_OUT_IDX = 6;
    using Use = passes::LibNodeExpander::InputUse;
    auto standalone = context.replacement_requires_access_nodes(
        {Use::IndirectRead,
         Use::IndirectRead,
         Use::IndirectRead,
         Use::IndirectRead,
         Use::IndirectRead,
         Use::Scalar,
         Use::IndirectWrite}
    );

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

    auto& new_sequence = standalone->replace_with_sequence();
    auto& builder = standalone->builder();

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


    // CPU implementation of batchnorm:
    if (false) {
        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 = standalone->add_indirect_read_access(inter_block, VAR_IDX);
        data_flow::AccessNode& epsilon_const = standalone->add_scalar_input_access(inter_block, EPS_IDX);

        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 = standalone->add_indirect_read_access(innermost_block, BATCH_IDX);
        auto& interm_in = builder.add_access(innermost_block, interm_name);
        auto& e_elem_in = standalone->add_indirect_read_access(innermost_block, E_IDX);
        auto& gamma_elem_in = standalone->add_indirect_read_access(innermost_block, GAMMA_IDX);
        auto& beta_elem_in = standalone->add_indirect_read_access(innermost_block, BETA_IDX);

        auto& result_ptr_out_elem = standalone->add_indirect_write_access(innermost_block, B_OUT_IDX);

        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
        );

        return standalone->successfully_expanded();
    } else {
        // GPU implementation of batchnorm:
        // Move sqrt and division into the innermost loop to enable more parallelism.

        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 = standalone->add_indirect_read_access(innermost_block, BATCH_IDX);
        auto& var_elem_in = standalone->add_indirect_read_access(innermost_block, VAR_IDX);
        data_flow::AccessNode& epsilon_const = standalone->add_scalar_input_access(innermost_block, EPS_IDX);
        auto& e_elem_in = standalone->add_indirect_read_access(innermost_block, E_IDX);
        auto& gamma_elem_in = standalone->add_indirect_read_access(innermost_block, GAMMA_IDX);
        auto& beta_elem_in = standalone->add_indirect_read_access(innermost_block, BETA_IDX);
        auto& result_ptr_out_elem = standalone->add_indirect_write_access(innermost_block, B_OUT_IDX);

        // 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
        );

        return standalone->successfully_expanded();
    }
}

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

daisytuner / docc / 28793626242

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