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

74.81

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

#include "sdfg/data_flow/library_nodes/math/tensor/pooling_node.h"

#include "sdfg/analysis/analysis.h"
#include "sdfg/builder/structured_sdfg_builder.h"
#include "sdfg/data_flow/library_nodes/math/cmath/cmath_node.h"
#include "sdfg/data_flow/library_nodes/math/tensor/spatial_tensor_node.h"
#include "sdfg/data_flow/library_nodes/math/tensor/tensor_node.h"
#include "sdfg/symbolic/symbolic.h"
#include "sdfg/types/type.h"

namespace sdfg {
namespace math {
namespace tensor {

PoolingNode::PoolingNode(
    size_t element_id,
    const DebugInfo& debug_info,
    const graph::Vertex vertex,
    data_flow::DataFlowGraph& parent,
    PoolingMode mode,
    const std::vector<symbolic::Expression>& shape,
    const std::vector<symbolic::Expression>& kernel_shape,
    const std::vector<symbolic::Expression>& strides,
    const std::vector<symbolic::Expression>& pads,
    const std::vector<symbolic::Expression>& dilations,
    QuantizationType quantization,
    const data_flow::ImplementationType& impl_type
)
    : SpatialTensorNode(
          element_id,
          debug_info,
          vertex,
          parent,
          LibraryNodeType_Pooling,
          {},
          {"Y", "X"},
          impl_type,
          quantization,
          shape,
          kernel_shape,
          strides,
          pads,
          dilations
      ),
      mode_(mode) {}

void PoolingNode::validate(const Function& function) const {
    TensorNode::validate(function);

    if (kernel_shape_.empty()) {
        throw InvalidSDFGException("PoolingNode kernel_shape cannot be empty");
    }

    size_t spatial_dims = kernel_shape_.size();

    if (!strides_.empty() && strides_.size() != spatial_dims) {
        throw InvalidSDFGException("PoolingNode strides must match kernel spatial dimensions");
    }

    if (!pads_.empty() && pads_.size() != 2 * spatial_dims) {
        throw InvalidSDFGException("PoolingNode pads must have 2 * spatial dimensions");
    }

    if (!dilations_.empty() && dilations_.size() != spatial_dims) {
        throw InvalidSDFGException("PoolingNode dilations must match kernel spatial dimensions");
    }
}

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

    auto primitive_type = this->primitive_type(dataflow);
    types::Scalar scalar_type(primitive_type);

    auto x_edge = dataflow.in_edge_for_connector(*this, "X");
    if (!x_edge) {
        return context.unable();
    }

    auto y_edge = dataflow.in_edge_for_connector(*this, "Y");
    if (!y_edge) {
        return context.unable();
    }

    size_t spatial_dims = kernel_shape_.size();
    if (spatial_dims == 0) {
        return context.unable();
    }

    // Get strides (default to 1)
    std::vector<symbolic::Expression> strides_vec;
    for (size_t i = 0; i < spatial_dims; ++i) {
        if (i < strides_.size()) {
            strides_vec.push_back(strides_[i]);
        } else {
            strides_vec.push_back(symbolic::one());
        }
    }

    // Get padding (default to 0)
    std::vector<symbolic::Expression> pads_begin_vec, pads_end_vec;
    for (size_t i = 0; i < spatial_dims; ++i) {
        if (i < pads_.size()) {
            pads_begin_vec.push_back(pads_[i]);
        } else {
            pads_begin_vec.push_back(symbolic::zero());
        }
        if (spatial_dims + i < pads_.size()) {
            pads_end_vec.push_back(pads_[spatial_dims + i]);
        } else {
            pads_end_vec.push_back(symbolic::zero());
        }
    }

    // Get dilations (default to 1)
    std::vector<symbolic::Expression> dilations_vec;
    for (size_t i = 0; i < spatial_dims; ++i) {
        if (i < dilations_.size()) {
            dilations_vec.push_back(dilations_[i]);
        } else {
            dilations_vec.push_back(symbolic::one());
        }
    }

    // Input shape: [N, C, D0, D1, ..., Dn]
    symbolic::Expression N = shape_[0];
    symbolic::Expression C = shape_[1];
    std::vector<symbolic::Expression> input_spatial_dims;
    for (size_t i = 0; i < spatial_dims; ++i) {
        input_spatial_dims.push_back(shape_[2 + i]);
    }

    // Output spatial dimensions
    std::vector<symbolic::Expression> output_spatial_dims;
    for (size_t i = 0; i < spatial_dims; ++i) {
        auto d_in = input_spatial_dims[i];
        auto pad = symbolic::add(pads_begin_vec[i], pads_end_vec[i]);
        auto dk = symbolic::mul(dilations_vec[i], symbolic::sub(kernel_shape_[i], symbolic::one()));
        auto num = symbolic::sub(symbolic::add(d_in, pad), symbolic::add(dk, symbolic::one()));
        auto d_out = symbolic::add(symbolic::div(num, strides_vec[i]), symbolic::one());
        output_spatial_dims.push_back(d_out);
    }

    using Use = passes::LibNodeExpander::InputUse;
    auto standalone = context.replacement_requires_access_nodes({Use::IndirectWrite, Use::IndirectRead});

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

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

    structured_control_flow::Sequence* current_scope = &new_sequence;
    std::vector<symbolic::Expression> output_indices;
    std::vector<symbolic::Expression> output_spatial_vars;

    // Map over batch
    std::string n_str = builder.find_new_name("n");
    builder.add_container(n_str, types::Scalar(types::PrimitiveType::UInt64));
    auto n_var = symbolic::symbol(n_str);
    auto& map_n = builder.add_map(
        *current_scope,
        n_var,
        symbolic::Lt(n_var, N),
        symbolic::zero(),
        symbolic::add(n_var, symbolic::one()),
        structured_control_flow::ScheduleType_Sequential::create(),
        {},
        block.debug_info()
    );
    current_scope = &map_n.root();
    output_indices.push_back(n_var);

    // Map over channel
    std::string c_str = builder.find_new_name("c");
    builder.add_container(c_str, types::Scalar(types::PrimitiveType::UInt64));
    auto c_var = symbolic::symbol(c_str);
    auto& map_c = builder.add_map(
        *current_scope,
        c_var,
        symbolic::Lt(c_var, C),
        symbolic::zero(),
        symbolic::add(c_var, symbolic::one()),
        structured_control_flow::ScheduleType_Sequential::create(),
        {},
        block.debug_info()
    );
    current_scope = &map_c.root();
    output_indices.push_back(c_var);

    // Map over each output spatial dimension
    for (size_t i = 0; i < spatial_dims; ++i) {
        std::string od_str = builder.find_new_name("od" + std::to_string(i));
        builder.add_container(od_str, types::Scalar(types::PrimitiveType::UInt64));
        auto od_var = symbolic::symbol(od_str);
        auto& map_od = builder.add_map(
            *current_scope,
            od_var,
            symbolic::Lt(od_var, output_spatial_dims[i]),
            symbolic::zero(),
            symbolic::add(od_var, symbolic::one()),
            structured_control_flow::ScheduleType_Sequential::create(),
            {},
            block.debug_info()
        );
        current_scope = &map_od.root();
        output_indices.push_back(od_var);
        output_spatial_vars.push_back(od_var);
    }

    // Create accumulator
    std::string accum_var = builder.find_new_name("_pool_accum");
    builder.add_container(accum_var, scalar_type);

    // Initialize accumulator
    std::string init_value;
    if (mode_ == PoolingMode::Max) {
        // Use -INFINITY for float, type-min for integers
        if (types::is_integer(primitive_type)) {
            switch (primitive_type) {
                case types::PrimitiveType::Int8:
                    init_value = "INT8_MIN";
                    break;
                case types::PrimitiveType::Int16:
                    init_value = "INT16_MIN";
                    break;
                case types::PrimitiveType::Int32:
                    init_value = "INT32_MIN";
                    break;
                case types::PrimitiveType::Int64:
                    init_value = "INT64_MIN";
                    break;
                default:
                    init_value = "0";
                    break;
            }
        } else {
            init_value = "-INFINITY";
        }
    } else {
        // Sum / Avg: init to 0
        init_value = types::is_integer(primitive_type) ? "0" : "0.0";
    }

    auto& init_block = builder.add_block(*current_scope, {}, block.debug_info());
    auto& accum_init = builder.add_access(init_block, accum_var, block.debug_info());
    auto& zero_const = builder.add_constant(init_block, init_value, scalar_type, block.debug_info());
    auto& init_tasklet = builder.add_tasklet(init_block, data_flow::assign, "_out", {"_in"}, block.debug_info());
    builder.add_computational_memlet(init_block, zero_const, init_tasklet, "_in", {}, scalar_type, block.debug_info());
    builder.add_computational_memlet(init_block, init_tasklet, "_out", accum_init, {}, scalar_type, block.debug_info());

    // For loops over kernel spatial dimensions
    auto* loop_scope = current_scope;
    std::vector<symbolic::Expression> kernel_vars;
    for (size_t i = 0; i < spatial_dims; ++i) {
        std::string k_str = builder.find_new_name("k" + std::to_string(i));
        builder.add_container(k_str, types::Scalar(types::PrimitiveType::UInt64));
        auto k_var = symbolic::symbol(k_str);
        auto& for_k = builder.add_for(
            *loop_scope,
            k_var,
            symbolic::Lt(k_var, kernel_shape_[i]),
            symbolic::zero(),
            symbolic::add(k_var, symbolic::one()),
            {},
            block.debug_info()
        );
        loop_scope = &for_k.root();
        kernel_vars.push_back(k_var);
    }

    // Compute input spatial indices
    std::vector<symbolic::Expression> input_spatial_indices;
    for (size_t i = 0; i < spatial_dims; ++i) {
        auto k_dilated = symbolic::mul(kernel_vars[i], dilations_vec[i]);
        auto input_idx = symbolic::
            add(symbolic::sub(symbolic::mul(output_spatial_vars[i], strides_vec[i]), pads_begin_vec[i]), k_dilated);
        input_spatial_indices.push_back(input_idx);
    }

    // Add branching if padding is non-zero
    bool has_padding = false;
    for (auto padding : this->pads_) {
        if (!symbolic::eq(padding, symbolic::zero())) {
            has_padding = true;
            break;
        }
    }
    if (has_padding) {
        symbolic::Condition comp_condition = symbolic::__true__();
        for (size_t i = 0; i < spatial_dims; ++i) {
            comp_condition = symbolic::
                And(comp_condition,
                    symbolic::
                        And(symbolic::Lt(input_spatial_indices[i], input_spatial_dims[i]),
                            symbolic::Ge(input_spatial_indices[i], symbolic::zero())));
        }
        auto& branch = builder.add_if_else(*loop_scope, {}, block.debug_info());
        auto& comp_case = builder.add_case(branch, comp_condition, block.debug_info());
        loop_scope = &comp_case;
    }

    // Build X indices: [n, c, input_spatial...]
    std::vector<symbolic::Expression> x_indices_vec = {n_var, c_var};
    x_indices_vec.insert(x_indices_vec.end(), input_spatial_indices.begin(), input_spatial_indices.end());
    data_flow::Subset x_subset(x_indices_vec);

    // Computation block: accumulate
    auto& comp_block = builder.add_block(*loop_scope, {}, block.debug_info());
    auto& x_access = standalone->add_indirect_read_access(comp_block, 1);
    auto& accum_read = builder.add_access(comp_block, accum_var, block.debug_info());
    auto& accum_write = builder.add_access(comp_block, accum_var, block.debug_info());

    if (mode_ == PoolingMode::Max) {
        bool is_int = types::is_integer(primitive_type);
        if (is_int) {
            auto tasklet_code = TensorNode::get_integer_minmax_tasklet(primitive_type, true);
            auto& tasklet = builder.add_tasklet(comp_block, tasklet_code, "_out", {"_in1", "_in2"}, block.debug_info());
            builder.add_computational_memlet(
                comp_block, x_access, tasklet, "_in1", x_subset, x_edge->base_type(), block.debug_info()
            );
            builder
                .add_computational_memlet(comp_block, accum_read, tasklet, "_in2", {}, scalar_type, block.debug_info());
            builder
                .add_computational_memlet(comp_block, tasklet, "_out", accum_write, {}, scalar_type, block.debug_info());
        } else {
            auto& libnode = builder.add_library_node<
                math::cmath::CMathNode>(comp_block, block.debug_info(), cmath::CMathFunction::fmax, primitive_type);
            builder.add_computational_memlet(
                comp_block, x_access, libnode, "_in1", x_subset, x_edge->base_type(), block.debug_info()
            );
            builder
                .add_computational_memlet(comp_block, accum_read, libnode, "_in2", {}, scalar_type, block.debug_info());
            builder
                .add_computational_memlet(comp_block, libnode, "_out", accum_write, {}, scalar_type, block.debug_info());
        }
    } else {
        // Sum or Avg: accumulate with addition
        bool is_int = types::is_integer(primitive_type);
        data_flow::TaskletCode opcode = is_int ? data_flow::TaskletCode::int_add : data_flow::TaskletCode::fp_add;
        auto& tasklet = builder.add_tasklet(comp_block, opcode, "_out", {"_in1", "_in2"}, block.debug_info());
        builder.add_computational_memlet(
            comp_block, x_access, tasklet, "_in1", x_subset, x_edge->base_type(), block.debug_info()
        );
        builder.add_computational_memlet(comp_block, accum_read, tasklet, "_in2", {}, scalar_type, block.debug_info());
        builder.add_computational_memlet(comp_block, tasklet, "_out", accum_write, {}, scalar_type, block.debug_info());
    }

    // After kernel loops: write result to output
    data_flow::Subset y_subset(output_indices);

    auto& output_block = builder.add_block(*current_scope, {}, block.debug_info());
    auto& accum_final = builder.add_access(output_block, accum_var, block.debug_info());
    auto& y_access = standalone->add_indirect_write_access(output_block, 0);

    if (mode_ == PoolingMode::Avg) {
        // Divide by window size: product of kernel_shape dimensions
        // Create a temporary for the divisor
        std::string divisor_var = builder.find_new_name("_pool_divisor");
        builder.add_container(divisor_var, scalar_type);

        // Compute window size as product of kernel dimensions
        symbolic::Expression window_size = kernel_shape_[0];
        for (size_t i = 1; i < spatial_dims; ++i) {
            window_size = symbolic::mul(window_size, kernel_shape_[i]);
        }

        auto& divisor_const =
            builder.add_constant(output_block, window_size->__str__(), scalar_type, block.debug_info());
        auto& divisor_access = builder.add_access(output_block, divisor_var, block.debug_info());
        auto& divisor_assign =
            builder.add_tasklet(output_block, data_flow::assign, "_out", {"_in"}, block.debug_info());
        builder.add_computational_memlet(
            output_block, divisor_const, divisor_assign, "_in", {}, scalar_type, block.debug_info()
        );
        builder.add_computational_memlet(
            output_block, divisor_assign, "_out", divisor_access, {}, scalar_type, block.debug_info()
        );

        bool is_int = types::is_integer(primitive_type);
        data_flow::TaskletCode div_opcode = is_int ? data_flow::TaskletCode::int_sdiv : data_flow::TaskletCode::fp_div;
        auto& div_tasklet = builder.add_tasklet(output_block, div_opcode, "_out", {"_in1", "_in2"}, block.debug_info());
        builder
            .add_computational_memlet(output_block, accum_final, div_tasklet, "_in1", {}, scalar_type, block.debug_info());
        builder.add_computational_memlet(
            output_block, divisor_access, div_tasklet, "_in2", {}, scalar_type, block.debug_info()
        );
        builder.add_computational_memlet(
            output_block, div_tasklet, "_out", y_access, y_subset, y_edge->base_type(), y_edge->debug_info()
        );
    } else {
        // Max or Sum: just assign
        auto& assign_tasklet =
            builder.add_tasklet(output_block, data_flow::assign, "_out", {"_in"}, block.debug_info());
        builder.add_computational_memlet(
            output_block, accum_final, assign_tasklet, "_in", {}, scalar_type, block.debug_info()
        );
        builder.add_computational_memlet(
            output_block, assign_tasklet, "_out", y_access, y_subset, y_edge->base_type(), y_edge->debug_info()
        );
    }

    return standalone->successfully_expanded();
}

std::unique_ptr<data_flow::DataFlowNode> PoolingNode::
    clone(size_t element_id, const graph::Vertex vertex, data_flow::DataFlowGraph& parent) const {
    return std::unique_ptr<data_flow::DataFlowNode>(new PoolingNode(
        element_id,
        this->debug_info(),
        vertex,
        parent,
        mode_,
        shape_,
        kernel_shape_,
        strides_,
        pads_,
        dilations_,
        fixed_quantization_,
        implementation_type_
    ));
}

std::string PoolingNode::mode_to_string(PoolingMode mode) {
    switch (mode) {
        case PoolingMode::Max:
            return "max";
        case PoolingMode::Sum:
            return "sum";
        case PoolingMode::Avg:
            return "avg";
    }
    return "unknown";
}

PoolingMode PoolingNode::string_to_mode(const std::string& str) {
    if (str == "max") return PoolingMode::Max;
    if (str == "sum") return PoolingMode::Sum;
    if (str == "avg") return PoolingMode::Avg;
    throw InvalidSDFGException("Unknown pooling mode: " + str);
}

symbolic::Expression PoolingNode::flop() const {
    // Total output elements: N * C * prod(output_spatial_dim(i))
    auto output_elems = symbolic::mul(symbolic::mul(shape_[0], shape_[1]), output_spatial_volume());

    // Each output element reduces a full kernel window.
    auto kv = kernel_volume();

    switch (mode_) {
        case PoolingMode::Max:
            // max pooling: (kv - 1) comparisons per output element
            return symbolic::mul(output_elems, symbolic::sub(kv, symbolic::one()));
        case PoolingMode::Sum:
            // sum pooling: (kv - 1) additions per output element
            return symbolic::mul(output_elems, symbolic::sub(kv, symbolic::one()));
        case PoolingMode::Avg:
            // avg pooling: (kv - 1) additions + 1 division per output element
            return symbolic::mul(output_elems, kv);
        default:
            return symbolic::symbol("UnknownFlops_Pool_n" + std::to_string(element_id_));
    }
}

data_flow::PointerAccessType PoolingNode::pointer_access_type(int input_idx) const {
    if (input_idx == 0) {
        return data_flow::PointerAccessMeta::create_full_write_only(symbolic::__nullptr__(), true);
    } else if (input_idx == 1) {
        return data_flow::PointerAccessMeta::create_read_only(symbolic::__nullptr__(), true);
    } else {
        return TensorNode::pointer_access_type(input_idx);
    }
}

std::string PoolingNode::toStr() const {
    std::stringstream ss;
    ss << "Pooling(mode=" << mode_to_string(mode_) << ", ";
    SpatialTensorNode::operator<<(ss);
    ss << ")";
    return ss.str();
}

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

    j["mode"] = PoolingNode::mode_to_string(node.mode());

    fill_base_values(node, j);

    return j;
}

data_flow::LibraryNode& PoolingNodeSerializer::deserialize(
    const nlohmann::json& j, builder::StructuredSDFGBuilder& builder, structured_control_flow::Block& parent
) {
    assert(j.contains("mode"));

    auto base = deserialize_base_values(j);

    auto mode = PoolingNode::string_to_mode(j["mode"].get<std::string>());

    return builder.add_library_node<PoolingNode>(
        parent,
        base.debug_info,
        mode,
        base.shape,
        base.kernel_shape,
        base.strides,
        base.pads,
        base.dilations,
        base.quantization
    );
}

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

1	#include "sdfg/data_flow/library_nodes/math/tensor/pooling_node.h"
2
3	#include "sdfg/analysis/analysis.h"
4	#include "sdfg/builder/structured_sdfg_builder.h"
5	#include "sdfg/data_flow/library_nodes/math/cmath/cmath_node.h"
6	#include "sdfg/data_flow/library_nodes/math/tensor/spatial_tensor_node.h"
7	#include "sdfg/data_flow/library_nodes/math/tensor/tensor_node.h"
8	#include "sdfg/symbolic/symbolic.h"
9	#include "sdfg/types/type.h"
10
11	namespace sdfg {
12	namespace math {
13	namespace tensor {
14
15	PoolingNode::PoolingNode(
16	size_t element_id,
17	const DebugInfo& debug_info,
18	const graph::Vertex vertex,
19	data_flow::DataFlowGraph& parent,
20	PoolingMode mode,
21	const std::vector<symbolic::Expression>& shape,
22	const std::vector<symbolic::Expression>& kernel_shape,
23	const std::vector<symbolic::Expression>& strides,
24	const std::vector<symbolic::Expression>& pads,
25	const std::vector<symbolic::Expression>& dilations,
26	QuantizationType quantization,
27	const data_flow::ImplementationType& impl_type
28	)
29	: SpatialTensorNode(	15✔
30	element_id,	15✔
31	debug_info,	15✔
32	vertex,	15✔
33	parent,	15✔
34	LibraryNodeType_Pooling,	15✔
35	{},	15✔
36	{"Y", "X"},	15✔
37	impl_type,	15✔
38	quantization,	15✔
39	shape,	15✔
40	kernel_shape,	15✔
41	strides,	15✔
42	pads,	15✔
43	dilations	15✔
44	),	15✔
45	mode_(mode) {}	15✔
46
47	void PoolingNode::validate(const Function& function) const {	12✔
48	TensorNode::validate(function);	12✔
49
50	if (kernel_shape_.empty()) {	12✔
51	throw InvalidSDFGException("PoolingNode kernel_shape cannot be empty");	×
52	}	×
53
54	size_t spatial_dims = kernel_shape_.size();	12✔
55
56	if (!strides_.empty() && strides_.size() != spatial_dims) {	12✔
57	throw InvalidSDFGException("PoolingNode strides must match kernel spatial dimensions");	×
58	}	×
59
60	if (!pads_.empty() && pads_.size() != 2 * spatial_dims) {	12✔
61	throw InvalidSDFGException("PoolingNode pads must have 2 * spatial dimensions");	×
62	}	×
63
64	if (!dilations_.empty() && dilations_.size() != spatial_dims) {	12✔
65	throw InvalidSDFGException("PoolingNode dilations must match kernel spatial dimensions");	×
66	}	×
67	}	12✔
68
69	passes::LibNodeExpander::ExpandOutcome PoolingNode::
70	expand(passes::LibNodeExpander::ExpandContext& context, structured_control_flow::Block& block) {	9✔
71	auto& dataflow = this->get_parent();	9✔
72
73	auto primitive_type = this->primitive_type(dataflow);	9✔
74	types::Scalar scalar_type(primitive_type);	9✔
75
76	auto x_edge = dataflow.in_edge_for_connector(*this, "X");	9✔
77	if (!x_edge) {	9✔
NEW 78	return context.unable();	×
79	}	×
80
81	auto y_edge = dataflow.in_edge_for_connector(*this, "Y");	9✔
82	if (!y_edge) {	9✔
NEW 83	return context.unable();	×
84	}	×
85
86	size_t spatial_dims = kernel_shape_.size();	9✔
87	if (spatial_dims == 0) {	9✔
NEW 88	return context.unable();	×
89	}	×
90
91	// Get strides (default to 1)
92	std::vector<symbolic::Expression> strides_vec;	9✔
93	for (size_t i = 0; i < spatial_dims; ++i) {	26✔
94	if (i < strides_.size()) {	17✔
95	strides_vec.push_back(strides_[i]);	17✔
96	} else {	17✔
97	strides_vec.push_back(symbolic::one());	×
98	}	×
99	}	17✔
100
101	// Get padding (default to 0)
102	std::vector<symbolic::Expression> pads_begin_vec, pads_end_vec;	9✔
103	for (size_t i = 0; i < spatial_dims; ++i) {	26✔
104	if (i < pads_.size()) {	17✔
105	pads_begin_vec.push_back(pads_[i]);	2✔
106	} else {	15✔
107	pads_begin_vec.push_back(symbolic::zero());	15✔
108	}	15✔
109	if (spatial_dims + i < pads_.size()) {	17✔
110	pads_end_vec.push_back(pads_[spatial_dims + i]);	2✔
111	} else {	15✔
112	pads_end_vec.push_back(symbolic::zero());	15✔
113	}	15✔
114	}	17✔
115
116	// Get dilations (default to 1)
117	std::vector<symbolic::Expression> dilations_vec;	9✔
118	for (size_t i = 0; i < spatial_dims; ++i) {	26✔
119	if (i < dilations_.size()) {	17✔
120	dilations_vec.push_back(dilations_[i]);	2✔
121	} else {	15✔
122	dilations_vec.push_back(symbolic::one());	15✔
123	}	15✔
124	}	17✔
125
126	// Input shape: [N, C, D0, D1, ..., Dn]
127	symbolic::Expression N = shape_[0];	9✔
128	symbolic::Expression C = shape_[1];	9✔
129	std::vector<symbolic::Expression> input_spatial_dims;	9✔
130	for (size_t i = 0; i < spatial_dims; ++i) {	26✔
131	input_spatial_dims.push_back(shape_[2 + i]);	17✔
132	}	17✔
133
134	// Output spatial dimensions
135	std::vector<symbolic::Expression> output_spatial_dims;	9✔
136	for (size_t i = 0; i < spatial_dims; ++i) {	26✔
137	auto d_in = input_spatial_dims[i];	17✔
138	auto pad = symbolic::add(pads_begin_vec[i], pads_end_vec[i]);	17✔
139	auto dk = symbolic::mul(dilations_vec[i], symbolic::sub(kernel_shape_[i], symbolic::one()));	17✔
140	auto num = symbolic::sub(symbolic::add(d_in, pad), symbolic::add(dk, symbolic::one()));	17✔
141	auto d_out = symbolic::add(symbolic::div(num, strides_vec[i]), symbolic::one());	17✔
142	output_spatial_dims.push_back(d_out);	17✔
143	}	17✔
144
145	using Use = passes::LibNodeExpander::InputUse;	9✔
146	auto standalone = context.replacement_requires_access_nodes({Use::IndirectWrite, Use::IndirectRead});	9✔
147
148	if (!standalone) {	9✔
NEW 149	return context.unable();	×
NEW 150	}	×
151
152	auto& new_sequence = standalone->replace_with_sequence();	9✔
153	auto& builder = standalone->builder();	9✔
154
155	structured_control_flow::Sequence* current_scope = &new_sequence;	9✔
156	std::vector<symbolic::Expression> output_indices;	9✔
157	std::vector<symbolic::Expression> output_spatial_vars;	9✔
158
159	// Map over batch
160	std::string n_str = builder.find_new_name("n");	9✔
161	builder.add_container(n_str, types::Scalar(types::PrimitiveType::UInt64));	9✔
162	auto n_var = symbolic::symbol(n_str);	9✔
163	auto& map_n = builder.add_map(	9✔
164	*current_scope,	9✔
165	n_var,	9✔
166	symbolic::Lt(n_var, N),	9✔
167	symbolic::zero(),	9✔
168	symbolic::add(n_var, symbolic::one()),	9✔
169	structured_control_flow::ScheduleType_Sequential::create(),	9✔
170	{},	9✔
171	block.debug_info()	9✔
172	);	9✔
173	current_scope = &map_n.root();	9✔
174	output_indices.push_back(n_var);	9✔
175
176	// Map over channel
177	std::string c_str = builder.find_new_name("c");	9✔
178	builder.add_container(c_str, types::Scalar(types::PrimitiveType::UInt64));	9✔
179	auto c_var = symbolic::symbol(c_str);	9✔
180	auto& map_c = builder.add_map(	9✔
181	*current_scope,	9✔
182	c_var,	9✔
183	symbolic::Lt(c_var, C),	9✔
184	symbolic::zero(),	9✔
185	symbolic::add(c_var, symbolic::one()),	9✔
186	structured_control_flow::ScheduleType_Sequential::create(),	9✔
187	{},	9✔
188	block.debug_info()	9✔
189	);	9✔
190	current_scope = &map_c.root();	9✔
191	output_indices.push_back(c_var);	9✔
192
193	// Map over each output spatial dimension
194	for (size_t i = 0; i < spatial_dims; ++i) {	26✔
195	std::string od_str = builder.find_new_name("od" + std::to_string(i));	17✔
196	builder.add_container(od_str, types::Scalar(types::PrimitiveType::UInt64));	17✔
197	auto od_var = symbolic::symbol(od_str);	17✔
198	auto& map_od = builder.add_map(	17✔
199	*current_scope,	17✔
200	od_var,	17✔
201	symbolic::Lt(od_var, output_spatial_dims[i]),	17✔
202	symbolic::zero(),	17✔
203	symbolic::add(od_var, symbolic::one()),	17✔
204	structured_control_flow::ScheduleType_Sequential::create(),	17✔
205	{},	17✔
206	block.debug_info()	17✔
207	);	17✔
208	current_scope = &map_od.root();	17✔
209	output_indices.push_back(od_var);	17✔
210	output_spatial_vars.push_back(od_var);	17✔
211	}	17✔
212
213	// Create accumulator
214	std::string accum_var = builder.find_new_name("_pool_accum");	9✔
215	builder.add_container(accum_var, scalar_type);	9✔
216
217	// Initialize accumulator
218	std::string init_value;	9✔
219	if (mode_ == PoolingMode::Max) {	9✔
220	// Use -INFINITY for float, type-min for integers
221	if (types::is_integer(primitive_type)) {	7✔
222	switch (primitive_type) {	1✔
223	case types::PrimitiveType::Int8:	×
224	init_value = "INT8_MIN";	×
225	break;	×
226	case types::PrimitiveType::Int16:	×
227	init_value = "INT16_MIN";	×
228	break;	×
229	case types::PrimitiveType::Int32:	1✔
230	init_value = "INT32_MIN";	1✔
231	break;	1✔
232	case types::PrimitiveType::Int64:	×
233	init_value = "INT64_MIN";	×
234	break;	×
235	default:	×
236	init_value = "0";	×
237	break;	×
238	}	1✔
239	} else {	6✔
240	init_value = "-INFINITY";	6✔
241	}	6✔
242	} else {	7✔
243	// Sum / Avg: init to 0
244	init_value = types::is_integer(primitive_type) ? "0" : "0.0";	2✔
245	}	2✔
246
247	auto& init_block = builder.add_block(*current_scope, {}, block.debug_info());	9✔
248	auto& accum_init = builder.add_access(init_block, accum_var, block.debug_info());	9✔
249	auto& zero_const = builder.add_constant(init_block, init_value, scalar_type, block.debug_info());	9✔
250	auto& init_tasklet = builder.add_tasklet(init_block, data_flow::assign, "_out", {"_in"}, block.debug_info());	9✔
251	builder.add_computational_memlet(init_block, zero_const, init_tasklet, "_in", {}, scalar_type, block.debug_info());	9✔
252	builder.add_computational_memlet(init_block, init_tasklet, "_out", accum_init, {}, scalar_type, block.debug_info());	9✔
253
254	// For loops over kernel spatial dimensions
255	auto* loop_scope = current_scope;	9✔
256	std::vector<symbolic::Expression> kernel_vars;	9✔
257	for (size_t i = 0; i < spatial_dims; ++i) {	26✔
258	std::string k_str = builder.find_new_name("k" + std::to_string(i));	17✔
259	builder.add_container(k_str, types::Scalar(types::PrimitiveType::UInt64));	17✔
260	auto k_var = symbolic::symbol(k_str);	17✔
261	auto& for_k = builder.add_for(	17✔
262	*loop_scope,	17✔
263	k_var,	17✔
264	symbolic::Lt(k_var, kernel_shape_[i]),	17✔
265	symbolic::zero(),	17✔
266	symbolic::add(k_var, symbolic::one()),	17✔
267	{},	17✔
268	block.debug_info()	17✔
269	);	17✔
270	loop_scope = &for_k.root();	17✔
271	kernel_vars.push_back(k_var);	17✔
272	}	17✔
273
274	// Compute input spatial indices
275	std::vector<symbolic::Expression> input_spatial_indices;	9✔
276	for (size_t i = 0; i < spatial_dims; ++i) {	26✔
277	auto k_dilated = symbolic::mul(kernel_vars[i], dilations_vec[i]);	17✔
278	auto input_idx = symbolic::	17✔
279	add(symbolic::sub(symbolic::mul(output_spatial_vars[i], strides_vec[i]), pads_begin_vec[i]), k_dilated);	17✔
280	input_spatial_indices.push_back(input_idx);	17✔
281	}	17✔
282
283	// Add branching if padding is non-zero
284	bool has_padding = false;	9✔
285	for (auto padding : this->pads_) {	9✔
286	if (!symbolic::eq(padding, symbolic::zero())) {	1✔
287	has_padding = true;	1✔
288	break;	1✔
289	}	1✔
290	}	1✔
291	if (has_padding) {	9✔
292	symbolic::Condition comp_condition = symbolic::__true__();	1✔
293	for (size_t i = 0; i < spatial_dims; ++i) {	3✔
294	comp_condition = symbolic::	2✔
295	And(comp_condition,	2✔
296	symbolic::	2✔
297	And(symbolic::Lt(input_spatial_indices[i], input_spatial_dims[i]),	2✔
298	symbolic::Ge(input_spatial_indices[i], symbolic::zero())));	2✔
299	}	2✔
300	auto& branch = builder.add_if_else(*loop_scope, {}, block.debug_info());	1✔
301	auto& comp_case = builder.add_case(branch, comp_condition, block.debug_info());	1✔
302	loop_scope = &comp_case;	1✔
303	}	1✔
304
305	// Build X indices: [n, c, input_spatial...]
306	std::vector<symbolic::Expression> x_indices_vec = {n_var, c_var};	9✔
307	x_indices_vec.insert(x_indices_vec.end(), input_spatial_indices.begin(), input_spatial_indices.end());	9✔
308	data_flow::Subset x_subset(x_indices_vec);	9✔
309
310	// Computation block: accumulate
311	auto& comp_block = builder.add_block(*loop_scope, {}, block.debug_info());	9✔
312	auto& x_access = standalone->add_indirect_read_access(comp_block, 1);	9✔
313	auto& accum_read = builder.add_access(comp_block, accum_var, block.debug_info());	9✔
314	auto& accum_write = builder.add_access(comp_block, accum_var, block.debug_info());	9✔
315
316	if (mode_ == PoolingMode::Max) {	9✔
317	bool is_int = types::is_integer(primitive_type);	7✔
318	if (is_int) {	7✔
319	auto tasklet_code = TensorNode::get_integer_minmax_tasklet(primitive_type, true);	1✔
320	auto& tasklet = builder.add_tasklet(comp_block, tasklet_code, "_out", {"_in1", "_in2"}, block.debug_info());	1✔
321	builder.add_computational_memlet(	1✔
322	comp_block, x_access, tasklet, "_in1", x_subset, x_edge->base_type(), block.debug_info()	1✔
323	);	1✔
324	builder	1✔
325	.add_computational_memlet(comp_block, accum_read, tasklet, "_in2", {}, scalar_type, block.debug_info());	1✔
326	builder	1✔
327	.add_computational_memlet(comp_block, tasklet, "_out", accum_write, {}, scalar_type, block.debug_info());	1✔
328	} else {	6✔
329	auto& libnode = builder.add_library_node<	6✔
330	math::cmath::CMathNode>(comp_block, block.debug_info(), cmath::CMathFunction::fmax, primitive_type);	6✔
331	builder.add_computational_memlet(	6✔
332	comp_block, x_access, libnode, "_in1", x_subset, x_edge->base_type(), block.debug_info()	6✔
333	);	6✔
334	builder	6✔
335	.add_computational_memlet(comp_block, accum_read, libnode, "_in2", {}, scalar_type, block.debug_info());	6✔
336	builder	6✔
337	.add_computational_memlet(comp_block, libnode, "_out", accum_write, {}, scalar_type, block.debug_info());	6✔
338	}	6✔
339	} else {	7✔
340	// Sum or Avg: accumulate with addition
341	bool is_int = types::is_integer(primitive_type);	2✔
342	data_flow::TaskletCode opcode = is_int ? data_flow::TaskletCode::int_add : data_flow::TaskletCode::fp_add;	2✔
343	auto& tasklet = builder.add_tasklet(comp_block, opcode, "_out", {"_in1", "_in2"}, block.debug_info());	2✔
344	builder.add_computational_memlet(	2✔
345	comp_block, x_access, tasklet, "_in1", x_subset, x_edge->base_type(), block.debug_info()	2✔
346	);	2✔
347	builder.add_computational_memlet(comp_block, accum_read, tasklet, "_in2", {}, scalar_type, block.debug_info());	2✔
348	builder.add_computational_memlet(comp_block, tasklet, "_out", accum_write, {}, scalar_type, block.debug_info());	2✔
349	}	2✔
350
351	// After kernel loops: write result to output
352	data_flow::Subset y_subset(output_indices);	9✔
353
354	auto& output_block = builder.add_block(*current_scope, {}, block.debug_info());	9✔
355	auto& accum_final = builder.add_access(output_block, accum_var, block.debug_info());	9✔
356	auto& y_access = standalone->add_indirect_write_access(output_block, 0);	9✔
357
358	if (mode_ == PoolingMode::Avg) {	9✔
359	// Divide by window size: product of kernel_shape dimensions
360	// Create a temporary for the divisor
361	std::string divisor_var = builder.find_new_name("_pool_divisor");	1✔
362	builder.add_container(divisor_var, scalar_type);	1✔
363
364	// Compute window size as product of kernel dimensions
365	symbolic::Expression window_size = kernel_shape_[0];	1✔
366	for (size_t i = 1; i < spatial_dims; ++i) {	2✔
367	window_size = symbolic::mul(window_size, kernel_shape_[i]);	1✔
368	}	1✔
369
370	auto& divisor_const =	1✔
371	builder.add_constant(output_block, window_size->__str__(), scalar_type, block.debug_info());	1✔
372	auto& divisor_access = builder.add_access(output_block, divisor_var, block.debug_info());	1✔
373	auto& divisor_assign =	1✔
374	builder.add_tasklet(output_block, data_flow::assign, "_out", {"_in"}, block.debug_info());	1✔
375	builder.add_computational_memlet(	1✔
376	output_block, divisor_const, divisor_assign, "_in", {}, scalar_type, block.debug_info()	1✔
377	);	1✔
378	builder.add_computational_memlet(	1✔
379	output_block, divisor_assign, "_out", divisor_access, {}, scalar_type, block.debug_info()	1✔
380	);	1✔
381
382	bool is_int = types::is_integer(primitive_type);	1✔
383	data_flow::TaskletCode div_opcode = is_int ? data_flow::TaskletCode::int_sdiv : data_flow::TaskletCode::fp_div;	1✔
384	auto& div_tasklet = builder.add_tasklet(output_block, div_opcode, "_out", {"_in1", "_in2"}, block.debug_info());	1✔
385	builder	1✔
386	.add_computational_memlet(output_block, accum_final, div_tasklet, "_in1", {}, scalar_type, block.debug_info());	1✔
387	builder.add_computational_memlet(	1✔
388	output_block, divisor_access, div_tasklet, "_in2", {}, scalar_type, block.debug_info()	1✔
389	);	1✔
390	builder.add_computational_memlet(	1✔
391	output_block, div_tasklet, "_out", y_access, y_subset, y_edge->base_type(), y_edge->debug_info()	1✔
392	);	1✔
393	} else {	8✔
394	// Max or Sum: just assign
395	auto& assign_tasklet =	8✔
396	builder.add_tasklet(output_block, data_flow::assign, "_out", {"_in"}, block.debug_info());	8✔
397	builder.add_computational_memlet(	8✔
398	output_block, accum_final, assign_tasklet, "_in", {}, scalar_type, block.debug_info()	8✔
399	);	8✔
400	builder.add_computational_memlet(	8✔
401	output_block, assign_tasklet, "_out", y_access, y_subset, y_edge->base_type(), y_edge->debug_info()	8✔
402	);	8✔
403	}	8✔
404
405	return standalone->successfully_expanded();	9✔
406	}	9✔
407
408	std::unique_ptr<data_flow::DataFlowNode> PoolingNode::
409	clone(size_t element_id, const graph::Vertex vertex, data_flow::DataFlowGraph& parent) const {	×
410	return std::unique_ptr<data_flow::DataFlowNode>(new PoolingNode(	×
411	element_id,	×
412	this->debug_info(),	×
413	vertex,	×
414	parent,	×
415	mode_,	×
416	shape_,	×
417	kernel_shape_,	×
418	strides_,	×
419	pads_,	×
420	dilations_,	×
421	fixed_quantization_,	×
422	implementation_type_	×
423	));	×
424	}	×
425
426	std::string PoolingNode::mode_to_string(PoolingMode mode) {	3✔
427	switch (mode) {	3✔
428	case PoolingMode::Max:	1✔
429	return "max";	1✔
430	case PoolingMode::Sum:	1✔
431	return "sum";	1✔
432	case PoolingMode::Avg:	1✔
433	return "avg";	1✔
434	}	3✔
435	return "unknown";	×
436	}	3✔
437
438	PoolingMode PoolingNode::string_to_mode(const std::string& str) {	3✔
439	if (str == "max") return PoolingMode::Max;	3✔
440	if (str == "sum") return PoolingMode::Sum;	2✔
441	if (str == "avg") return PoolingMode::Avg;	1✔
442	throw InvalidSDFGException("Unknown pooling mode: " + str);	×
443	}	1✔
444
445	symbolic::Expression PoolingNode::flop() const {	×
446	// Total output elements: N * C * prod(output_spatial_dim(i))
447	auto output_elems = symbolic::mul(symbolic::mul(shape_[0], shape_[1]), output_spatial_volume());	×
448
449	// Each output element reduces a full kernel window.
450	auto kv = kernel_volume();	×
451
452	switch (mode_) {	×
453	case PoolingMode::Max:	×
454	// max pooling: (kv - 1) comparisons per output element
455	return symbolic::mul(output_elems, symbolic::sub(kv, symbolic::one()));	×
456	case PoolingMode::Sum:	×
457	// sum pooling: (kv - 1) additions per output element
458	return symbolic::mul(output_elems, symbolic::sub(kv, symbolic::one()));	×
459	case PoolingMode::Avg:	×
460	// avg pooling: (kv - 1) additions + 1 division per output element
461	return symbolic::mul(output_elems, kv);	×
462	default:	×
463	return symbolic::symbol("UnknownFlops_Pool_n" + std::to_string(element_id_));	×
464	}	×
465	}	×
466
467	data_flow::PointerAccessType PoolingNode::pointer_access_type(int input_idx) const {	×
468	if (input_idx == 0) {	×
469	return data_flow::PointerAccessMeta::create_full_write_only(symbolic::__nullptr__(), true);	×
470	} else if (input_idx == 1) {	×
471	return data_flow::PointerAccessMeta::create_read_only(symbolic::__nullptr__(), true);	×
472	} else {	×
473	return TensorNode::pointer_access_type(input_idx);	×
474	}	×
475	}	×
476
477	std::string PoolingNode::toStr() const {	×
478	std::stringstream ss;	×
479	ss << "Pooling(mode=" << mode_to_string(mode_) << ", ";	×
480	SpatialTensorNode::operator<<(ss);	×
481	ss << ")";	×
482	return ss.str();	×
483	}	×
484
485	nlohmann::json PoolingNodeSerializer::serialize(const data_flow::LibraryNode& library_node) {	×
486	const PoolingNode& node = static_cast<const PoolingNode&>(library_node);	×
487	nlohmann::json j;	×
488
489	j["mode"] = PoolingNode::mode_to_string(node.mode());	×
490
491	fill_base_values(node, j);	×
492
493	return j;	×
494	}	×
495
496	data_flow::LibraryNode& PoolingNodeSerializer::deserialize(
497	const nlohmann::json& j, builder::StructuredSDFGBuilder& builder, structured_control_flow::Block& parent
498	) {	×
499	assert(j.contains("mode"));	×
500
501	auto base = deserialize_base_values(j);	×
502
503	auto mode = PoolingNode::string_to_mode(j["mode"].get<std::string>());	×
504
505	return builder.add_library_node<PoolingNode>(	×
506	parent,	×
507	base.debug_info,	×
508	mode,	×
509	base.shape,	×
510	base.kernel_shape,	×
511	base.strides,	×
512	base.pads,	×
513	base.dilations,	×
514	base.quantization	×
515	);	×
516	}	×
517
518	} // namespace tensor
519	} // namespace math
520	} // namespace sdfg

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