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

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