28806128926

Committed 06 Jul 2026 04:16PM UTC coverage: 62.96%. First build

Build # 28806128926

Build Type

push

github

Committed by

web-flow

Commit Message

New LibNodeExpansion pass (#740)

* Switched expansion "pipeline" to new LibNodeExpansionPass that can recursively expand using a LibNodeExpander impl.
- removed access to analysis_manager from expansion methods, as those would currently not be appropriately maintained between nodes
+ New expansion API handles more of the boilerplate code (checking for standalone, creating boundary access nodes, removing old elements)
* Also migrated sdfg-json-to-c.cpp to not using the expansion pipeline anymore
* toStr() for ReduceNodes and giving PyStructuredSDFG access to the output_dir for additional dumping
~ DotVisualizer : Fix on nested sequences.
+ DotVisualizer: visualize empty sequences
* DotVisualizer default-enabled show block/loop ids
 * while MathNodes still have an expand-method, the new infrastructure is based upon "Expander" classes. The MathNodeExpander just redirects to the method for now
 ~ Broadcast node still used old ptr-output semantics
 * updated tensor & blas node expand to new expand API, that handles more of the boilerplate code (checking, removing old nodes, creating standalone-replacement nodes)
 - removed Transpose Node. Was unused and on old ptr-semantics
 + StructuredSDFGBuilder.add_sequence_at, add_for_at, add_map_at
 * switched StructuredSDFGBuilder internally to use ptr of Assignments to express using default assignments when needed
 * updated tests to use the new expand_single_math_node helper function, instead of the method directly
 + pass and expand-single helper functions are ready to use other expanders
 + EinsumExpansionPass is basically the legacy ExpansionPass, only restricted to EinsumNodes. It still needs to run in a pipeline to ensure finding multiple EinsumNodes per block. This is temporary, because Einsum is the only node that already supported splitting a block, which the new expansion disallows, as it should handle it itself when needed (but that part is not yet implemented)

Coverage Stats

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

40554 of 64412 relevant lines covered (62.96%)

963.11 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 / 28806128926

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