30634500673

Committed 31 Jul 2026 01:27PM UTC coverage: 64.806% (+0.1%) from 64.71%

Build # 30634500673

Build Type

push

github

Committed by

web-flow

Commit Message

Merge pull request #920 from daisytuner/pytorch-any-where

Big & small stuff to get segformer working

Coverage Stats

486 of 604 new or added lines in 13 files covered. (80.46%)

24 existing lines in 2 files now uncovered.

45593 of 70353 relevant lines covered (64.81%)

723.15 hits per line

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

85.71

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

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

#include <cstddef>
#include <memory>
#include <nlohmann/json_fwd.hpp>
#include <sstream>
#include <string>

#include "sdfg/builder/structured_sdfg_builder.h"
#include "sdfg/data_flow/data_flow_graph.h"
#include "sdfg/data_flow/data_flow_node.h"
#include "sdfg/data_flow/library_node.h"
#include "sdfg/data_flow/library_nodes/math/tensor/tensor_layout.h"
#include "sdfg/data_flow/library_nodes/math/tensor/tensor_node.h"
#include "sdfg/data_flow/memlet.h"
#include "sdfg/data_flow/tasklet.h"
#include "sdfg/element.h"
#include "sdfg/exceptions.h"
#include "sdfg/function.h"
#include "sdfg/graph/graph.h"
#include "sdfg/passes/expansion/lib_node_expander.h"
#include "sdfg/serializer/json_serializer.h"
#include "sdfg/structured_control_flow/block.h"
#include "sdfg/structured_control_flow/sequence.h"
#include "sdfg/structured_control_flow/structured_loop.h"
#include "sdfg/symbolic/symbolic.h"
#include "sdfg/types/scalar.h"
#include "sdfg/types/tensor.h"
#include "sdfg/types/type.h"

namespace sdfg {
namespace math {
namespace tensor {

void ConditionalTensorCopyNode::validate_equal_shapes(const TensorLayout& layout1, const TensorLayout& layout2) const {
    int dims = layout1.dims();
    if (dims == 0) {
        return;
    }
    if (dims != layout2.dims()) {
        throw InvalidSDFGException(
            "ConditionalTensorCopyNode: Shapes mismatch: " + layout1.toStr() + " != " + layout2.toStr()
        );
    }
    for (int i = 0; i < dims; i++) {
        if (!symbolic::eq(layout1.get_dim(i), layout2.get_dim(i)) &&
            !symbolic::eq(layout1.get_dim(i), symbolic::one())) {
            throw InvalidSDFGException(
                "ConditionalTensorCopyNode: Shapes mismatch: " + layout1.toStr() + " != " + layout2.toStr()
            );
        }
    }
}

ConditionalTensorCopyNode::ConditionalTensorCopyNode(
    size_t element_id,
    const DebugInfo& debug_info,
    const graph::Vertex vertex,
    data_flow::DataFlowGraph& parent,
    const TensorLayout& layout_mask,
    const TensorLayout& layout_x1,
    const TensorLayout& layout_x2,
    const TensorLayout& layout_y,
    const data_flow::ImplementationType& impl_type
)
    : TensorNode(
          element_id,
          debug_info,
          vertex,
          parent,
          LibraryNodeType_ConditionalTensorCopy,
          {},
          {"Mask", "X1", "X2", "Y"},
          impl_type
      ),
      layout_mask_(layout_mask), layout_x1_(layout_x1), layout_x2_(layout_x2), layout_y_(layout_y) {}


const TensorLayout& ConditionalTensorCopyNode::layout_mask() const { return this->layout_mask_; }

const TensorLayout& ConditionalTensorCopyNode::layout_x1() const { return this->layout_x1_; }

const TensorLayout& ConditionalTensorCopyNode::layout_x2() const { return this->layout_x2_; }

const TensorLayout& ConditionalTensorCopyNode::layout_y() const { return this->layout_y_; }

void ConditionalTensorCopyNode::validate(const Function& function) const {
    auto& graph = this->get_parent();

    // Check presence of in and out edges
    if (graph.out_degree(*this) != 0) {
        throw InvalidSDFGException(
            "ConditionalTensorCopyNode: Expected no outputs but got: " + std::to_string(graph.out_degree(*this))
        );
    }
    const data_flow::Memlet* mask_iedge = graph.in_edge_for_connector(*this, "Mask");
    if (!mask_iedge) {
        throw InvalidSDFGException("ConditionalTensorCopyNode: No memlet connected at connector: Mask");
    }
    const data_flow::Memlet* x1_iedge = graph.in_edge_for_connector(*this, "X1");
    if (!x1_iedge) {
        throw InvalidSDFGException("ConditionalTensorCopyNode: No memlet connected at connector: X1");
    }
    const data_flow::Memlet* x2_iedge = graph.in_edge_for_connector(*this, "X2");
    if (!x2_iedge) {
        throw InvalidSDFGException("ConditionalTensorCopyNode: No memlet connected at connector: X2");
    }
    const data_flow::Memlet* y_iedge = graph.in_edge_for_connector(*this, "Y");
    if (!y_iedge) {
        throw InvalidSDFGException("ConditionalTensorCopyNode: No memlet connected at connector: Y");
    }

    // Check that the in edges have tensor types as base types
    if (mask_iedge->base_type().type_id() != types::TypeID::Tensor) {
        throw InvalidSDFGException(
            "ConditionalTensorCopyNode: Expected tensor type at connector 'Mask' but got: " +
            mask_iedge->base_type().print()
        );
    }
    if (x1_iedge->base_type().type_id() != types::TypeID::Tensor) {
        throw InvalidSDFGException(
            "ConditionalTensorCopyNode: Expected tensor type at connector 'X1' but got: " +
            x1_iedge->base_type().print()
        );
    }
    if (x2_iedge->base_type().type_id() != types::TypeID::Tensor) {
        throw InvalidSDFGException(
            "ConditionalTensorCopyNode: Expected tensor type at connector 'X2' but got: " +
            x2_iedge->base_type().print()
        );
    }
    if (y_iedge->base_type().type_id() != types::TypeID::Tensor) {
        throw InvalidSDFGException(
            "ConditionalTensorCopyNode: Expected tensor type at connector 'Y' but got: " + y_iedge->base_type().print()
        );
    }

    const types::Tensor& mask_tensor = static_cast<const types::Tensor&>(mask_iedge->base_type());
    const types::Tensor& x1_tensor = static_cast<const types::Tensor&>(x1_iedge->base_type());
    const types::Tensor& x2_tensor = static_cast<const types::Tensor&>(x2_iedge->base_type());
    const types::Tensor& y_tensor = static_cast<const types::Tensor&>(y_iedge->base_type());

    // Check that the tensor layouts match with the tensor types on the edges
    if (mask_tensor.layout() != this->layout_mask_) {
        throw InvalidSDFGException(
            "ConditionalTensorCopyNode: Provided tensor layout does not match the memlet tensor type for connector "
            "'Mask': " +
            mask_tensor.layout().toStr() + " != " + this->layout_mask_.toStr()
        );
    }
    if (x1_tensor.layout() != this->layout_x1_) {
        throw InvalidSDFGException(
            "ConditionalTensorCopyNode: Provided tensor layout does not match the memlet tensor type for connector "
            "'X1': " +
            x1_tensor.layout().toStr() + " != " + this->layout_x1_.toStr()
        );
    }
    if (x2_tensor.layout() != this->layout_x2_) {
        throw InvalidSDFGException(
            "ConditionalTensorCopyNode: Provided tensor layout does not match the memlet tensor type for connector "
            "'X2': " +
            x2_tensor.layout().toStr() + " != " + this->layout_x2_.toStr()
        );
    }
    if (y_tensor.layout() != this->layout_y_) {
        throw InvalidSDFGException(
            "ConditionalTensorCopyNode: Provided tensor layout does not match the memlet tensor type for connector "
            "'Y': " +
            y_tensor.layout().toStr() + " != " + this->layout_y_.toStr()
        );
    }

    // Check that all tensor layouts have the same shape
    this->validate_equal_shapes(this->layout_mask_, this->layout_y_);
    this->validate_equal_shapes(this->layout_x1_, this->layout_y_);
    this->validate_equal_shapes(this->layout_x2_, this->layout_y_);

    // Check that the tensor element type of the memlet for connector Mask is bool
    if (mask_tensor.primitive_type() != types::PrimitiveType::Bool) {
        throw InvalidSDFGException(
            "ConditionalTensorCopyNode: Expected a boolean element type for tensor type at connector 'Mask' but got: " +
            mask_tensor.element_type().print()
        );
    }

    // Check that the other tensor have the same element types
    types::PrimitiveType prim = y_tensor.primitive_type();
    if (x1_tensor.primitive_type() != prim || x2_tensor.primitive_type() != prim) {
        throw InvalidSDFGException(
            "ConditionalTensorCopyNode: Expected the same primitive types but got: " +
            x1_tensor.element_type().print() + ", " + x2_tensor.element_type().print() + ", and " +
            y_tensor.element_type().print()
        );
    }
}

bool ConditionalTensorCopyNode::supports_integer_types() const { return true; }

using Dir = passes::LibNodeExpander::InputUse;

passes::LibNodeExpander::ExpandOutcome ConditionalTensorCopyNode::
    expand(passes::LibNodeExpander::ExpandContext& context, structured_control_flow::Block& block) {
    auto standalone =
        context
            .replacement_requires_access_nodes({Dir::IndirectRead, Dir::IndirectRead, Dir::IndirectRead, Dir::IndirectWrite}
            );

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

    auto& builder = standalone->builder();
    auto& dfg = this->get_parent();
    const auto* iedge_mask = dfg.in_edge_for_connector(*this, "Mask");
    if (!iedge_mask) {
        throw InvalidSDFGException("ConditionalTensorCopyNode: Cannot get in edge for connector 'Mask'");
    }
    const auto* iedge_x1 = dfg.in_edge_for_connector(*this, "X1");
    if (!iedge_x1) {
        throw InvalidSDFGException("ConditionalTensorCopyNode: Cannot get in edge for connector 'X1'");
    }
    const auto* iedge_x2 = dfg.in_edge_for_connector(*this, "X2");
    if (!iedge_x2) {
        throw InvalidSDFGException("ConditionalTensorCopyNode: Cannot get in edge for connector 'X2'");
    }
    const auto* iedge_y = dfg.in_edge_for_connector(*this, "Y");
    if (!iedge_y) {
        throw InvalidSDFGException("ConditionalTensorCopyNode: Cannot get in edge for connector 'Y'");
    }

    // Add a graph after the current block
    auto& new_sequence = standalone->replace_with_sequence();

    // Add map nest over shape
    types::Scalar indvar_type(types::PrimitiveType::UInt64);
    structured_control_flow::Sequence* current_seq = &new_sequence;
    data_flow::Subset subset;
    subset.reserve(this->layout_y_.dims());
    for (auto dim : this->layout_y_.shape()) {
        auto indvar_container = builder.find_new_name("_i");
        builder.add_container(indvar_container, indvar_type);
        auto indvar = symbolic::symbol(indvar_container);
        subset.push_back(indvar);
        auto& map = builder.add_map(
            *current_seq,
            indvar,
            symbolic::Lt(indvar, dim),
            symbolic::zero(),
            symbolic::add(indvar, symbolic::one()),
            structured_control_flow::ScheduleType_Sequential::create(),
            this->debug_info_
        );
        current_seq = &map.root();
    }

    // Load local variable from mask
    int mask_dims = this->layout_mask_.dims();
    data_flow::Subset mask_subset;
    mask_subset.reserve(mask_dims);
    for (int i = 0; i < mask_dims; i++) {
        if (symbolic::eq(this->layout_mask_.get_dim(i), symbolic::one())) {
            mask_subset.push_back(symbolic::zero());
        } else {
            mask_subset.push_back(subset[i]);
        }
    }
    auto cond_copy_container = builder.find_new_name("tmp_cond_copy_");
    types::Scalar bool_type(types::PrimitiveType::Bool);
    builder.add_container(cond_copy_container, bool_type);
    auto& mask_block = builder.add_block(*current_seq, this->debug_info_);
    {
        auto& mask_access = standalone->add_scalar_input_access(mask_block, MASK_INPUT_IDX);
        auto& cond_copy_access = builder.add_access(mask_block, cond_copy_container, this->debug_info_);
        auto& tasklet =
            builder.add_tasklet(mask_block, data_flow::TaskletCode::assign, "_out", {"_in"}, this->debug_info_);
        builder.add_computational_memlet(
            mask_block, mask_access, tasklet, "_in", mask_subset, iedge_mask->base_type(), iedge_mask->debug_info()
        );
        builder.add_computational_memlet(mask_block, tasklet, "_out", cond_copy_access, {}, this->debug_info_);
    }

    // Add branch over local mask variable
    auto& if_else = builder.add_if_else(*current_seq, this->debug_info_);
    auto cond_copy = symbolic::symbol(cond_copy_container);
    auto& case_true = builder.add_case(if_else, symbolic::Eq(cond_copy, symbolic::__true__()), this->debug_info_);
    auto& case_false = builder.add_case(if_else, symbolic::Eq(cond_copy, symbolic::__false__()), this->debug_info_);

    // Fill true case: Copy value from X1 to Y
    int x1_dims = this->layout_x1_.dims();
    data_flow::Subset x1_subset;
    x1_subset.reserve(x1_dims);
    for (int i = 0; i < x1_dims; i++) {
        if (symbolic::eq(this->layout_x1_.get_dim(i), symbolic::one())) {
            x1_subset.push_back(symbolic::zero());
        } else {
            x1_subset.push_back(subset[i]);
        }
    }
    auto& block_true = builder.add_block(case_true, this->debug_info_);
    {
        auto& x1_access = standalone->add_scalar_input_access(block_true, X1_INPUT_IDX);
        auto& y_access = standalone->add_scalar_input_access(block_true, Y_INPUT_IDX);
        auto& tasklet =
            builder.add_tasklet(block_true, data_flow::TaskletCode::assign, "_out", {"_in"}, this->debug_info_);
        builder.add_computational_memlet(
            block_true, x1_access, tasklet, "_in", x1_subset, iedge_x1->base_type(), iedge_x1->debug_info()
        );
        builder.add_computational_memlet(
            block_true, tasklet, "_out", y_access, subset, iedge_y->base_type(), iedge_y->debug_info()
        );
    }

    // Fill false case: Copy value from X2 to Y
    int x2_dims = this->layout_x2_.dims();
    data_flow::Subset x2_subset;
    x2_subset.reserve(x2_dims);
    for (int i = 0; i < x2_dims; i++) {
        if (symbolic::eq(this->layout_x2_.get_dim(i), symbolic::one())) {
            x2_subset.push_back(symbolic::zero());
        } else {
            x2_subset.push_back(subset[i]);
        }
    }
    auto& block_false = builder.add_block(case_false, this->debug_info_);
    {
        auto& x2_access = standalone->add_scalar_input_access(block_false, X2_INPUT_IDX);
        auto& y_access = standalone->add_scalar_input_access(block_false, Y_INPUT_IDX);
        auto& tasklet =
            builder.add_tasklet(block_false, data_flow::TaskletCode::assign, "_out", {"_in"}, this->debug_info_);
        builder.add_computational_memlet(
            block_false, x2_access, tasklet, "_in", x2_subset, iedge_x2->base_type(), iedge_x2->debug_info()
        );
        builder.add_computational_memlet(
            block_false, tasklet, "_out", y_access, subset, iedge_y->base_type(), iedge_y->debug_info()
        );
    }

    return standalone->successfully_expanded();
}

std::string ConditionalTensorCopyNode::toStr() const {
    std::stringstream stream;

    stream << "ConditionalTensorCopyNode(shape: ";
    TensorLayout::emit_symbolic_list(stream, this->layout_y_.shape());
    stream << ")";

    return stream.str();
}

symbolic::SymbolSet ConditionalTensorCopyNode::symbols() const {
    symbolic::SymbolSet syms;
    this->layout_mask_.collect_symbols(syms);
    this->layout_x1_.collect_symbols(syms);
    this->layout_x2_.collect_symbols(syms);
    this->layout_y_.collect_symbols(syms);
    return syms;
}

symbolic::Expression ConditionalTensorCopyNode::flop() const { return symbolic::zero(); }

std::unique_ptr<data_flow::DataFlowNode> ConditionalTensorCopyNode::
    clone(size_t element_id, const graph::Vertex vertex, data_flow::DataFlowGraph& parent) const {
    return std::make_unique<ConditionalTensorCopyNode>(
        element_id,
        this->debug_info_,
        vertex,
        parent,
        this->layout_mask_,
        this->layout_x1_,
        this->layout_x2_,
        this->layout_y_
    );
}

void ConditionalTensorCopyNode::
    replace(const symbolic::Expression old_expression, const symbolic::Expression new_expression) {
    this->layout_mask_.replace_symbols(old_expression, new_expression);
    this->layout_x1_.replace_symbols(old_expression, new_expression);
    this->layout_x2_.replace_symbols(old_expression, new_expression);
    this->layout_y_.replace_symbols(old_expression, new_expression);
}

void ConditionalTensorCopyNode::replace(const symbolic::ExpressionMapping& replacements) {
    this->layout_mask_.replace_symbols(replacements);
    this->layout_x1_.replace_symbols(replacements);
    this->layout_x2_.replace_symbols(replacements);
    this->layout_y_.replace_symbols(replacements);
}

nlohmann::json ConditionalTensorCopyNodeSerializer::serialize(const data_flow::LibraryNode& library_node) {
    const ConditionalTensorCopyNode& conditional_copy_node = static_cast<const ConditionalTensorCopyNode&>(library_node
    );
    nlohmann::json j;

    j["code"] = conditional_copy_node.code().value();
    conditional_copy_node.layout_mask().serialize_to_json(j["layout_mask"]);
    conditional_copy_node.layout_x1().serialize_to_json(j["layout_x1"]);
    conditional_copy_node.layout_x2().serialize_to_json(j["layout_x2"]);
    conditional_copy_node.layout_y().serialize_to_json(j["layout_y"]);

    return j;
}

data_flow::LibraryNode& ConditionalTensorCopyNodeSerializer::deserialize(
    const nlohmann::json& j, builder::StructuredSDFGBuilder& builder, structured_control_flow::Block& parent
) {
    assert(j.contains("element_id"));
    assert(j.contains("code"));
    assert(j.contains("layout_mask"));
    assert(j.contains("layout_x1"));
    assert(j.contains("layout_x2"));
    assert(j.contains("layout_y"));
    assert(j.contains("debug_info"));

    TensorLayout layout_mask = TensorLayout::deserialize_from_json(j.at("layout_mask"));
    TensorLayout layout_x1 = TensorLayout::deserialize_from_json(j.at("layout_x1"));
    TensorLayout layout_x2 = TensorLayout::deserialize_from_json(j.at("layout_x2"));
    TensorLayout layout_y = TensorLayout::deserialize_from_json(j.at("layout_y"));

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

    return builder
        .add_library_node<ConditionalTensorCopyNode>(parent, debug_info, layout_mask, layout_x1, layout_x2, layout_y);
}

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

1	#include "sdfg/data_flow/library_nodes/math/tensor/conditional_copy_node.h"
2
3	#include <cstddef>
4	#include <memory>
5	#include <nlohmann/json_fwd.hpp>
6	#include <sstream>
7	#include <string>
8
9	#include "sdfg/builder/structured_sdfg_builder.h"
10	#include "sdfg/data_flow/data_flow_graph.h"
11	#include "sdfg/data_flow/data_flow_node.h"
12	#include "sdfg/data_flow/library_node.h"
13	#include "sdfg/data_flow/library_nodes/math/tensor/tensor_layout.h"
14	#include "sdfg/data_flow/library_nodes/math/tensor/tensor_node.h"
15	#include "sdfg/data_flow/memlet.h"
16	#include "sdfg/data_flow/tasklet.h"
17	#include "sdfg/element.h"
18	#include "sdfg/exceptions.h"
19	#include "sdfg/function.h"
20	#include "sdfg/graph/graph.h"
21	#include "sdfg/passes/expansion/lib_node_expander.h"
22	#include "sdfg/serializer/json_serializer.h"
23	#include "sdfg/structured_control_flow/block.h"
24	#include "sdfg/structured_control_flow/sequence.h"
25	#include "sdfg/structured_control_flow/structured_loop.h"
26	#include "sdfg/symbolic/symbolic.h"
27	#include "sdfg/types/scalar.h"
28	#include "sdfg/types/tensor.h"
29	#include "sdfg/types/type.h"
30
31	namespace sdfg {
32	namespace math {
33	namespace tensor {
34
35	void ConditionalTensorCopyNode::validate_equal_shapes(const TensorLayout& layout1, const TensorLayout& layout2) const {	28✔
36	int dims = layout1.dims();	28✔
37	if (dims == 0) {	28✔
NEW 38	return;	×
NEW 39	}	×
40	if (dims != layout2.dims()) {	28✔
41	throw InvalidSDFGException(	1✔
42	"ConditionalTensorCopyNode: Shapes mismatch: " + layout1.toStr() + " != " + layout2.toStr()	1✔
43	);	1✔
44	}	1✔
45	for (int i = 0; i < dims; i++) {	86✔
46	if (!symbolic::eq(layout1.get_dim(i), layout2.get_dim(i)) &&	60✔
47	!symbolic::eq(layout1.get_dim(i), symbolic::one())) {	60✔
48	throw InvalidSDFGException(	1✔
49	"ConditionalTensorCopyNode: Shapes mismatch: " + layout1.toStr() + " != " + layout2.toStr()	1✔
50	);	1✔
51	}	1✔
52	}	60✔
53	}	27✔
54
55	ConditionalTensorCopyNode::ConditionalTensorCopyNode(
56	size_t element_id,
57	const DebugInfo& debug_info,
58	const graph::Vertex vertex,
59	data_flow::DataFlowGraph& parent,
60	const TensorLayout& layout_mask,
61	const TensorLayout& layout_x1,
62	const TensorLayout& layout_x2,
63	const TensorLayout& layout_y,
64	const data_flow::ImplementationType& impl_type
65	)
66	: TensorNode(	16✔
67	element_id,	16✔
68	debug_info,	16✔
69	vertex,	16✔
70	parent,	16✔
71	LibraryNodeType_ConditionalTensorCopy,	16✔
72	{},	16✔
73	{"Mask", "X1", "X2", "Y"},	16✔
74	impl_type	16✔
75	),	16✔
76	layout_mask_(layout_mask), layout_x1_(layout_x1), layout_x2_(layout_x2), layout_y_(layout_y) {}	16✔
77
78
79	const TensorLayout& ConditionalTensorCopyNode::layout_mask() const { return this->layout_mask_; }	1✔
80
81	const TensorLayout& ConditionalTensorCopyNode::layout_x1() const { return this->layout_x1_; }	1✔
82
83	const TensorLayout& ConditionalTensorCopyNode::layout_x2() const { return this->layout_x2_; }	1✔
84
85	const TensorLayout& ConditionalTensorCopyNode::layout_y() const { return this->layout_y_; }	1✔
86
87	void ConditionalTensorCopyNode::validate(const Function& function) const {	18✔
88	auto& graph = this->get_parent();	18✔
89
90	// Check presence of in and out edges
91	if (graph.out_degree(*this) != 0) {	18✔
NEW 92	throw InvalidSDFGException(	×
NEW 93	"ConditionalTensorCopyNode: Expected no outputs but got: " + std::to_string(graph.out_degree(*this))	×
NEW 94	);	×
NEW 95	}	×
96	const data_flow::Memlet* mask_iedge = graph.in_edge_for_connector(*this, "Mask");	18✔
97	if (!mask_iedge) {	18✔
NEW 98	throw InvalidSDFGException("ConditionalTensorCopyNode: No memlet connected at connector: Mask");	×
NEW 99	}	×
100	const data_flow::Memlet* x1_iedge = graph.in_edge_for_connector(*this, "X1");	18✔
101	if (!x1_iedge) {	18✔
NEW 102	throw InvalidSDFGException("ConditionalTensorCopyNode: No memlet connected at connector: X1");	×
NEW 103	}	×
104	const data_flow::Memlet* x2_iedge = graph.in_edge_for_connector(*this, "X2");	18✔
105	if (!x2_iedge) {	18✔
NEW 106	throw InvalidSDFGException("ConditionalTensorCopyNode: No memlet connected at connector: X2");	×
NEW 107	}	×
108	const data_flow::Memlet* y_iedge = graph.in_edge_for_connector(*this, "Y");	18✔
109	if (!y_iedge) {	18✔
NEW 110	throw InvalidSDFGException("ConditionalTensorCopyNode: No memlet connected at connector: Y");	×
NEW 111	}	×
112
113	// Check that the in edges have tensor types as base types
114	if (mask_iedge->base_type().type_id() != types::TypeID::Tensor) {	18✔
115	throw InvalidSDFGException(	1✔
116	"ConditionalTensorCopyNode: Expected tensor type at connector 'Mask' but got: " +	1✔
117	mask_iedge->base_type().print()	1✔
118	);	1✔
119	}	1✔
120	if (x1_iedge->base_type().type_id() != types::TypeID::Tensor) {	17✔
121	throw InvalidSDFGException(	1✔
122	"ConditionalTensorCopyNode: Expected tensor type at connector 'X1' but got: " +	1✔
123	x1_iedge->base_type().print()	1✔
124	);	1✔
125	}	1✔
126	if (x2_iedge->base_type().type_id() != types::TypeID::Tensor) {	16✔
127	throw InvalidSDFGException(	1✔
128	"ConditionalTensorCopyNode: Expected tensor type at connector 'X2' but got: " +	1✔
129	x2_iedge->base_type().print()	1✔
130	);	1✔
131	}	1✔
132	if (y_iedge->base_type().type_id() != types::TypeID::Tensor) {	15✔
133	throw InvalidSDFGException(	1✔
134	"ConditionalTensorCopyNode: Expected tensor type at connector 'Y' but got: " + y_iedge->base_type().print()	1✔
135	);	1✔
136	}	1✔
137
138	const types::Tensor& mask_tensor = static_cast<const types::Tensor&>(mask_iedge->base_type());	14✔
139	const types::Tensor& x1_tensor = static_cast<const types::Tensor&>(x1_iedge->base_type());	14✔
140	const types::Tensor& x2_tensor = static_cast<const types::Tensor&>(x2_iedge->base_type());	14✔
141	const types::Tensor& y_tensor = static_cast<const types::Tensor&>(y_iedge->base_type());	14✔
142
143	// Check that the tensor layouts match with the tensor types on the edges
144	if (mask_tensor.layout() != this->layout_mask_) {	14✔
145	throw InvalidSDFGException(	1✔
146	"ConditionalTensorCopyNode: Provided tensor layout does not match the memlet tensor type for connector "	1✔
147	"'Mask': " +	1✔
148	mask_tensor.layout().toStr() + " != " + this->layout_mask_.toStr()	1✔
149	);	1✔
150	}	1✔
151	if (x1_tensor.layout() != this->layout_x1_) {	13✔
152	throw InvalidSDFGException(	1✔
153	"ConditionalTensorCopyNode: Provided tensor layout does not match the memlet tensor type for connector "	1✔
154	"'X1': " +	1✔
155	x1_tensor.layout().toStr() + " != " + this->layout_x1_.toStr()	1✔
156	);	1✔
157	}	1✔
158	if (x2_tensor.layout() != this->layout_x2_) {	12✔
159	throw InvalidSDFGException(	1✔
160	"ConditionalTensorCopyNode: Provided tensor layout does not match the memlet tensor type for connector "	1✔
161	"'X2': " +	1✔
162	x2_tensor.layout().toStr() + " != " + this->layout_x2_.toStr()	1✔
163	);	1✔
164	}	1✔
165	if (y_tensor.layout() != this->layout_y_) {	11✔
166	throw InvalidSDFGException(	1✔
167	"ConditionalTensorCopyNode: Provided tensor layout does not match the memlet tensor type for connector "	1✔
168	"'Y': " +	1✔
169	y_tensor.layout().toStr() + " != " + this->layout_y_.toStr()	1✔
170	);	1✔
171	}	1✔
172
173	// Check that all tensor layouts have the same shape
174	this->validate_equal_shapes(this->layout_mask_, this->layout_y_);	10✔
175	this->validate_equal_shapes(this->layout_x1_, this->layout_y_);	10✔
176	this->validate_equal_shapes(this->layout_x2_, this->layout_y_);	10✔
177
178	// Check that the tensor element type of the memlet for connector Mask is bool
179	if (mask_tensor.primitive_type() != types::PrimitiveType::Bool) {	10✔
180	throw InvalidSDFGException(	1✔
181	"ConditionalTensorCopyNode: Expected a boolean element type for tensor type at connector 'Mask' but got: " +	1✔
182	mask_tensor.element_type().print()	1✔
183	);	1✔
184	}	1✔
185
186	// Check that the other tensor have the same element types
187	types::PrimitiveType prim = y_tensor.primitive_type();	9✔
188	if (x1_tensor.primitive_type() != prim \|\| x2_tensor.primitive_type() != prim) {	9✔
189	throw InvalidSDFGException(	1✔
190	"ConditionalTensorCopyNode: Expected the same primitive types but got: " +	1✔
191	x1_tensor.element_type().print() + ", " + x2_tensor.element_type().print() + ", and " +	1✔
192	y_tensor.element_type().print()	1✔
193	);	1✔
194	}	1✔
195	}	9✔
196
197	bool ConditionalTensorCopyNode::supports_integer_types() const { return true; }	1✔
198
199	using Dir = passes::LibNodeExpander::InputUse;
200
201	passes::LibNodeExpander::ExpandOutcome ConditionalTensorCopyNode::
202	expand(passes::LibNodeExpander::ExpandContext& context, structured_control_flow::Block& block) {	1✔
203	auto standalone =	1✔
204	context	1✔
205	.replacement_requires_access_nodes({Dir::IndirectRead, Dir::IndirectRead, Dir::IndirectRead, Dir::IndirectWrite}	1✔
206	);	1✔
207
208	if (!standalone) {	1✔
NEW 209	return context.unable();	×
NEW 210	}	×
211
212	auto& builder = standalone->builder();	1✔
213	auto& dfg = this->get_parent();	1✔
214	const auto* iedge_mask = dfg.in_edge_for_connector(*this, "Mask");	1✔
215	if (!iedge_mask) {	1✔
NEW 216	throw InvalidSDFGException("ConditionalTensorCopyNode: Cannot get in edge for connector 'Mask'");	×
NEW 217	}	×
218	const auto* iedge_x1 = dfg.in_edge_for_connector(*this, "X1");	1✔
219	if (!iedge_x1) {	1✔
NEW 220	throw InvalidSDFGException("ConditionalTensorCopyNode: Cannot get in edge for connector 'X1'");	×
NEW 221	}	×
222	const auto* iedge_x2 = dfg.in_edge_for_connector(*this, "X2");	1✔
223	if (!iedge_x2) {	1✔
NEW 224	throw InvalidSDFGException("ConditionalTensorCopyNode: Cannot get in edge for connector 'X2'");	×
NEW 225	}	×
226	const auto* iedge_y = dfg.in_edge_for_connector(*this, "Y");	1✔
227	if (!iedge_y) {	1✔
NEW 228	throw InvalidSDFGException("ConditionalTensorCopyNode: Cannot get in edge for connector 'Y'");	×
NEW 229	}	×
230
231	// Add a graph after the current block
232	auto& new_sequence = standalone->replace_with_sequence();	1✔
233
234	// Add map nest over shape
235	types::Scalar indvar_type(types::PrimitiveType::UInt64);	1✔
236	structured_control_flow::Sequence* current_seq = &new_sequence;	1✔
237	data_flow::Subset subset;	1✔
238	subset.reserve(this->layout_y_.dims());	1✔
239	for (auto dim : this->layout_y_.shape()) {	2✔
240	auto indvar_container = builder.find_new_name("_i");	2✔
241	builder.add_container(indvar_container, indvar_type);	2✔
242	auto indvar = symbolic::symbol(indvar_container);	2✔
243	subset.push_back(indvar);	2✔
244	auto& map = builder.add_map(	2✔
245	*current_seq,	2✔
246	indvar,	2✔
247	symbolic::Lt(indvar, dim),	2✔
248	symbolic::zero(),	2✔
249	symbolic::add(indvar, symbolic::one()),	2✔
250	structured_control_flow::ScheduleType_Sequential::create(),	2✔
251	this->debug_info_	2✔
252	);	2✔
253	current_seq = &map.root();	2✔
254	}	2✔
255
256	// Load local variable from mask
257	int mask_dims = this->layout_mask_.dims();	1✔
258	data_flow::Subset mask_subset;	1✔
259	mask_subset.reserve(mask_dims);	1✔
260	for (int i = 0; i < mask_dims; i++) {	3✔
261	if (symbolic::eq(this->layout_mask_.get_dim(i), symbolic::one())) {	2✔
NEW 262	mask_subset.push_back(symbolic::zero());	×
263	} else {	2✔
264	mask_subset.push_back(subset[i]);	2✔
265	}	2✔
266	}	2✔
267	auto cond_copy_container = builder.find_new_name("tmp_cond_copy_");	1✔
268	types::Scalar bool_type(types::PrimitiveType::Bool);	1✔
269	builder.add_container(cond_copy_container, bool_type);	1✔
270	auto& mask_block = builder.add_block(*current_seq, this->debug_info_);	1✔
271	{	1✔
272	auto& mask_access = standalone->add_scalar_input_access(mask_block, MASK_INPUT_IDX);	1✔
273	auto& cond_copy_access = builder.add_access(mask_block, cond_copy_container, this->debug_info_);	1✔
274	auto& tasklet =	1✔
275	builder.add_tasklet(mask_block, data_flow::TaskletCode::assign, "_out", {"_in"}, this->debug_info_);	1✔
276	builder.add_computational_memlet(	1✔
277	mask_block, mask_access, tasklet, "_in", mask_subset, iedge_mask->base_type(), iedge_mask->debug_info()	1✔
278	);	1✔
279	builder.add_computational_memlet(mask_block, tasklet, "_out", cond_copy_access, {}, this->debug_info_);	1✔
280	}	1✔
281
282	// Add branch over local mask variable
283	auto& if_else = builder.add_if_else(*current_seq, this->debug_info_);	1✔
284	auto cond_copy = symbolic::symbol(cond_copy_container);	1✔
285	auto& case_true = builder.add_case(if_else, symbolic::Eq(cond_copy, symbolic::__true__()), this->debug_info_);	1✔
286	auto& case_false = builder.add_case(if_else, symbolic::Eq(cond_copy, symbolic::__false__()), this->debug_info_);	1✔
287
288	// Fill true case: Copy value from X1 to Y
289	int x1_dims = this->layout_x1_.dims();	1✔
290	data_flow::Subset x1_subset;	1✔
291	x1_subset.reserve(x1_dims);	1✔
292	for (int i = 0; i < x1_dims; i++) {	3✔
293	if (symbolic::eq(this->layout_x1_.get_dim(i), symbolic::one())) {	2✔
NEW 294	x1_subset.push_back(symbolic::zero());	×
295	} else {	2✔
296	x1_subset.push_back(subset[i]);	2✔
297	}	2✔
298	}	2✔
299	auto& block_true = builder.add_block(case_true, this->debug_info_);	1✔
300	{	1✔
301	auto& x1_access = standalone->add_scalar_input_access(block_true, X1_INPUT_IDX);	1✔
302	auto& y_access = standalone->add_scalar_input_access(block_true, Y_INPUT_IDX);	1✔
303	auto& tasklet =	1✔
304	builder.add_tasklet(block_true, data_flow::TaskletCode::assign, "_out", {"_in"}, this->debug_info_);	1✔
305	builder.add_computational_memlet(	1✔
306	block_true, x1_access, tasklet, "_in", x1_subset, iedge_x1->base_type(), iedge_x1->debug_info()	1✔
307	);	1✔
308	builder.add_computational_memlet(	1✔
309	block_true, tasklet, "_out", y_access, subset, iedge_y->base_type(), iedge_y->debug_info()	1✔
310	);	1✔
311	}	1✔
312
313	// Fill false case: Copy value from X2 to Y
314	int x2_dims = this->layout_x2_.dims();	1✔
315	data_flow::Subset x2_subset;	1✔
316	x2_subset.reserve(x2_dims);	1✔
317	for (int i = 0; i < x2_dims; i++) {	3✔
318	if (symbolic::eq(this->layout_x2_.get_dim(i), symbolic::one())) {	2✔
NEW 319	x2_subset.push_back(symbolic::zero());	×
320	} else {	2✔
321	x2_subset.push_back(subset[i]);	2✔
322	}	2✔
323	}	2✔
324	auto& block_false = builder.add_block(case_false, this->debug_info_);	1✔
325	{	1✔
326	auto& x2_access = standalone->add_scalar_input_access(block_false, X2_INPUT_IDX);	1✔
327	auto& y_access = standalone->add_scalar_input_access(block_false, Y_INPUT_IDX);	1✔
328	auto& tasklet =	1✔
329	builder.add_tasklet(block_false, data_flow::TaskletCode::assign, "_out", {"_in"}, this->debug_info_);	1✔
330	builder.add_computational_memlet(	1✔
331	block_false, x2_access, tasklet, "_in", x2_subset, iedge_x2->base_type(), iedge_x2->debug_info()	1✔
332	);	1✔
333	builder.add_computational_memlet(	1✔
334	block_false, tasklet, "_out", y_access, subset, iedge_y->base_type(), iedge_y->debug_info()	1✔
335	);	1✔
336	}	1✔
337
338	return standalone->successfully_expanded();	1✔
339	}	1✔
340
NEW 341	std::string ConditionalTensorCopyNode::toStr() const {	×
NEW 342	std::stringstream stream;	×
343
NEW 344	stream << "ConditionalTensorCopyNode(shape: ";	×
NEW 345	TensorLayout::emit_symbolic_list(stream, this->layout_y_.shape());	×
NEW 346	stream << ")";	×
347
NEW 348	return stream.str();	×
NEW 349	}	×
350
351	symbolic::SymbolSet ConditionalTensorCopyNode::symbols() const {	2✔
352	symbolic::SymbolSet syms;	2✔
353	this->layout_mask_.collect_symbols(syms);	2✔
354	this->layout_x1_.collect_symbols(syms);	2✔
355	this->layout_x2_.collect_symbols(syms);	2✔
356	this->layout_y_.collect_symbols(syms);	2✔
357	return syms;	2✔
358	}	2✔
359
360	symbolic::Expression ConditionalTensorCopyNode::flop() const { return symbolic::zero(); }	1✔
361
362	std::unique_ptr<data_flow::DataFlowNode> ConditionalTensorCopyNode::
NEW 363	clone(size_t element_id, const graph::Vertex vertex, data_flow::DataFlowGraph& parent) const {	×
NEW 364	return std::make_unique<ConditionalTensorCopyNode>(	×
NEW 365	element_id,	×
NEW 366	this->debug_info_,	×
NEW 367	vertex,	×
NEW 368	parent,	×
NEW 369	this->layout_mask_,	×
NEW 370	this->layout_x1_,	×
NEW 371	this->layout_x2_,	×
NEW 372	this->layout_y_	×
NEW 373	);	×
NEW 374	}	×
375
376	void ConditionalTensorCopyNode::
377	replace(const symbolic::Expression old_expression, const symbolic::Expression new_expression) {	1✔
378	this->layout_mask_.replace_symbols(old_expression, new_expression);	1✔
379	this->layout_x1_.replace_symbols(old_expression, new_expression);	1✔
380	this->layout_x2_.replace_symbols(old_expression, new_expression);	1✔
381	this->layout_y_.replace_symbols(old_expression, new_expression);	1✔
382	}	1✔
383
384	void ConditionalTensorCopyNode::replace(const symbolic::ExpressionMapping& replacements) {	1✔
385	this->layout_mask_.replace_symbols(replacements);	1✔
386	this->layout_x1_.replace_symbols(replacements);	1✔
387	this->layout_x2_.replace_symbols(replacements);	1✔
388	this->layout_y_.replace_symbols(replacements);	1✔
389	}	1✔
390
391	nlohmann::json ConditionalTensorCopyNodeSerializer::serialize(const data_flow::LibraryNode& library_node) {	1✔
392	const ConditionalTensorCopyNode& conditional_copy_node = static_cast<const ConditionalTensorCopyNode&>(library_node	1✔
393	);	1✔
394	nlohmann::json j;	1✔
395
396	j["code"] = conditional_copy_node.code().value();	1✔
397	conditional_copy_node.layout_mask().serialize_to_json(j["layout_mask"]);	1✔
398	conditional_copy_node.layout_x1().serialize_to_json(j["layout_x1"]);	1✔
399	conditional_copy_node.layout_x2().serialize_to_json(j["layout_x2"]);	1✔
400	conditional_copy_node.layout_y().serialize_to_json(j["layout_y"]);	1✔
401
402	return j;	1✔
403	}	1✔
404
405	data_flow::LibraryNode& ConditionalTensorCopyNodeSerializer::deserialize(
406	const nlohmann::json& j, builder::StructuredSDFGBuilder& builder, structured_control_flow::Block& parent
407	) {	1✔
408	assert(j.contains("element_id"));	1✔
409	assert(j.contains("code"));	1✔
410	assert(j.contains("layout_mask"));	1✔
411	assert(j.contains("layout_x1"));	1✔
412	assert(j.contains("layout_x2"));	1✔
413	assert(j.contains("layout_y"));	1✔
414	assert(j.contains("debug_info"));	1✔
415
416	TensorLayout layout_mask = TensorLayout::deserialize_from_json(j.at("layout_mask"));	1✔
417	TensorLayout layout_x1 = TensorLayout::deserialize_from_json(j.at("layout_x1"));	1✔
418	TensorLayout layout_x2 = TensorLayout::deserialize_from_json(j.at("layout_x2"));	1✔
419	TensorLayout layout_y = TensorLayout::deserialize_from_json(j.at("layout_y"));	1✔
420
421	serializer::JSONSerializer serializer;	1✔
422	DebugInfo debug_info = serializer.json_to_debug_info(j.at("debug_info"));	1✔
423
424	return builder	1✔
425	.add_library_node<ConditionalTensorCopyNode>(parent, debug_info, layout_mask, layout_x1, layout_x2, layout_y);	1✔
426	}	1✔
427
428	} // namespace tensor
429	} // namespace math
430	} // namespace sdfg

daisytuner / docc / 30634500673

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