17655994367

Committed 11 Sep 2025 08:06PM UTC coverage: 60.451% (+1.1%) from 59.335%

Build # 17655994367

Build Type

Pull #219

github

Committed by

web-flow

Commit Message

Merge 769d2cc0b into 6c1992b40

Pull Request Pull Request #219: stdlib Library Nodes and ConstantNodes

Run Details

457 of 1629 new or added lines in 81 files covered. (28.05%)

93 existing lines in 35 files now uncovered.

9382 of 15520 relevant lines covered (60.45%)

107.26 hits per line

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

57.05

/src/data_flow/library_nodes/math/blas/gemm.cpp

#include "sdfg/data_flow/library_nodes/math/blas/gemm.h"

#include "sdfg/analysis/analysis.h"
#include "sdfg/builder/structured_sdfg_builder.h"

#include "sdfg/analysis/scope_analysis.h"

namespace sdfg {
namespace math {
namespace blas {

GEMMNode::GEMMNode(
    size_t element_id,
    const DebugInfo& debug_info,
    const graph::Vertex vertex,
    data_flow::DataFlowGraph& parent,
    const data_flow::ImplementationType& implementation_type,
    const BLAS_Precision& precision,
    const BLAS_Layout& layout,
    const BLAS_Transpose& trans_a,
    const BLAS_Transpose& trans_b,
    symbolic::Expression m,
    symbolic::Expression n,
    symbolic::Expression k,
    symbolic::Expression lda,
    symbolic::Expression ldb,
    symbolic::Expression ldc
)
    : MathNode(
          element_id,
          debug_info,
          vertex,
          parent,
          LibraryNodeType_GEMM,
          {"C"},
          {"A", "B", "C", "alpha", "beta"},
          implementation_type
      ),
      precision_(precision), layout_(layout), trans_a_(trans_a), trans_b_(trans_b), m_(m), n_(n), k_(k), lda_(lda),
      ldb_(ldb), ldc_(ldc) {}

BLAS_Precision GEMMNode::precision() const { return this->precision_; };

BLAS_Layout GEMMNode::layout() const { return this->layout_; };

BLAS_Transpose GEMMNode::trans_a() const { return this->trans_a_; };

BLAS_Transpose GEMMNode::trans_b() const { return this->trans_b_; };

symbolic::Expression GEMMNode::m() const { return this->m_; };

symbolic::Expression GEMMNode::n() const { return this->n_; };

symbolic::Expression GEMMNode::k() const { return this->k_; };

symbolic::Expression GEMMNode::lda() const { return this->lda_; };

symbolic::Expression GEMMNode::ldb() const { return this->ldb_; };

symbolic::Expression GEMMNode::ldc() const { return this->ldc_; };

void GEMMNode::validate(const Function& function) const {}

types::PrimitiveType GEMMNode::scalar_primitive() const {
    switch (this->precision_) {
        case BLAS_Precision::s:
            return types::PrimitiveType::Float;
        case BLAS_Precision::d:
            return types::PrimitiveType::Double;
        case BLAS_Precision::h:
            return types::PrimitiveType::Half;
        default:
            return types::PrimitiveType::Void;
    }
}

bool GEMMNode::expand(builder::StructuredSDFGBuilder& builder, analysis::AnalysisManager& analysis_manager) {
    auto& scope_analysis = analysis_manager.get<analysis::ScopeAnalysis>();

    auto& dataflow = this->get_parent();
    auto& block = static_cast<structured_control_flow::Block&>(*dataflow.get_parent());
    auto& parent = static_cast<structured_control_flow::Sequence&>(*scope_analysis.parent_scope(&block));
    int index = parent.index(block);
    auto& transition = parent.at(index).second;

    if (trans_a_ != BLAS_Transpose::No || trans_b_ != BLAS_Transpose::No) {
        return false;
    }

    auto primitive_type = scalar_primitive();
    if (primitive_type == types::PrimitiveType::Void) {
        return false;
    }

    types::Scalar scalar_type(primitive_type);

    auto in_edges = dataflow.in_edges(*this);
    auto in_edges_it = in_edges.begin();

    data_flow::Memlet* iedge_a = nullptr;
    data_flow::Memlet* iedge_b = nullptr;
    data_flow::Memlet* iedge_c = nullptr;
    data_flow::Memlet* alpha_edge = nullptr;
    data_flow::Memlet* beta_edge = nullptr;
    while (in_edges_it != in_edges.end()) {
        auto& edge = *in_edges_it;
        auto dst_conn = edge.dst_conn();
        if (dst_conn == "A") {
            iedge_a = &edge;
        } else if (dst_conn == "B") {
            iedge_b = &edge;
        } else if (dst_conn == "C") {
            iedge_c = &edge;
        } else if (dst_conn == "alpha") {
            alpha_edge = &edge;
        } else if (dst_conn == "beta") {
            beta_edge = &edge;
        } else {
            throw InvalidSDFGException("GEMMNode has unexpected input: " + dst_conn);
        }
        ++in_edges_it;
    }

    auto& oedge = *dataflow.out_edges(*this).begin();

    // Checks if legal
    auto* input_node_a = static_cast<data_flow::AccessNode*>(&iedge_a->src());
    auto* input_node_b = static_cast<data_flow::AccessNode*>(&iedge_b->src());
    auto* input_node_c = static_cast<data_flow::AccessNode*>(&iedge_c->src());
    auto* output_node = static_cast<data_flow::AccessNode*>(&oedge.dst());
    auto* alpha_node = static_cast<data_flow::AccessNode*>(&alpha_edge->src());
    auto* beta_node = static_cast<data_flow::AccessNode*>(&beta_edge->src());

    // we must be the only thing in this block, as we do not support splitting a block into pre, expanded lib-node, post
    if (!input_node_a || dataflow.in_degree(*input_node_a) != 0 || !input_node_b ||
        dataflow.in_degree(*input_node_b) != 0 || !input_node_c || dataflow.in_degree(*input_node_c) != 0 ||
        !output_node || dataflow.out_degree(*output_node) != 0) {
        return false; // data nodes are not standalone
    }
    if (dataflow.in_degree(*alpha_node) != 0 || dataflow.in_degree(*beta_node) != 0) {
        return false; // alpha and beta are not standalone
    }
    for (auto* nd : dataflow.data_nodes()) {
        if (nd != input_node_a && nd != input_node_b && nd != input_node_c && nd != output_node &&
            (!alpha_node || nd != alpha_node) && (!beta_node || nd != beta_node)) {
            return false; // there are other nodes in here that we could not preserve correctly
        }
    }

    auto& A_var = input_node_a->data();
    auto& B_var = input_node_b->data();
    auto& C_in_var = input_node_c->data();
    auto& C_out_var = output_node->data();


    // Add new graph after the current block
    auto& new_sequence = builder.add_sequence_before(parent, block, transition.assignments(), block.debug_info());

    // Add maps
    std::vector<symbolic::Expression> indvar_ends{this->m(), this->n(), this->k()};
    data_flow::Subset new_subset;
    structured_control_flow::Sequence* last_scope = &new_sequence;
    structured_control_flow::Map* last_map = nullptr;
    structured_control_flow::Map* output_loop = nullptr;
    std::vector<std::string> indvar_names{"_i", "_j", "_k"};

    std::string sum_var = builder.find_new_name("_sum");
    builder.add_container(sum_var, scalar_type);

    for (size_t i = 0; i < 3; i++) {
        auto dim_begin = symbolic::zero();
        auto& dim_end = indvar_ends[i];

        std::string indvar_str = builder.find_new_name(indvar_names[i]);
        builder.add_container(indvar_str, types::Scalar(types::PrimitiveType::UInt64));

        auto indvar = symbolic::symbol(indvar_str);
        auto init = dim_begin;
        auto update = symbolic::add(indvar, symbolic::one());
        auto condition = symbolic::Lt(indvar, dim_end);
        last_map = &builder.add_map(
            *last_scope,
            indvar,
            condition,
            init,
            update,
            structured_control_flow::ScheduleType_Sequential::create(),
            {},
            block.debug_info()
        );
        last_scope = &last_map->root();

        if (i == 1) {
            output_loop = last_map;
        }

        new_subset.push_back(indvar);
    }


    // Add code
    auto& init_block = builder.add_block_before(output_loop->root(), *last_map, {}, block.debug_info());
    auto& sum_init = builder.add_access(init_block, sum_var, block.debug_info());

    auto& zero_node = builder.add_constant(init_block, "0.0", alpha_edge->base_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_node, init_tasklet, "_in", {}, block.debug_info());
    builder.add_computational_memlet(init_block, init_tasklet, "_out", sum_init, {}, block.debug_info());

    auto& code_block = builder.add_block(*last_scope, {}, block.debug_info());
    auto& input_node_a_new = builder.add_access(code_block, A_var, input_node_a->debug_info());
    auto& input_node_b_new = builder.add_access(code_block, B_var, input_node_b->debug_info());

    auto& core_fma =
        builder.add_tasklet(code_block, data_flow::fma, "_out", {"_in1", "_in2", "_in3"}, block.debug_info());
    auto& sum_in = builder.add_access(code_block, sum_var, block.debug_info());
    auto& sum_out = builder.add_access(code_block, sum_var, block.debug_info());
    builder.add_computational_memlet(code_block, sum_in, core_fma, "_in3", {}, block.debug_info());

    symbolic::Expression a_idx = symbolic::add(symbolic::mul(lda(), new_subset[0]), new_subset[2]);
    builder.add_computational_memlet(
        code_block, input_node_a_new, core_fma, "_in1", {a_idx}, iedge_a->base_type(), iedge_a->debug_info()
    );
    symbolic::Expression b_idx = symbolic::add(symbolic::mul(ldb(), new_subset[2]), new_subset[1]);
    builder.add_computational_memlet(
        code_block, input_node_b_new, core_fma, "_in2", {b_idx}, iedge_b->base_type(), iedge_b->debug_info()
    );
    builder.add_computational_memlet(code_block, core_fma, "_out", sum_out, {}, oedge.debug_info());

    auto& flush_block = builder.add_block_after(output_loop->root(), *last_map, {}, block.debug_info());
    auto& sum_final = builder.add_access(flush_block, sum_var, block.debug_info());
    auto& input_node_c_new = builder.add_access(flush_block, C_in_var, input_node_c->debug_info());
    symbolic::Expression c_idx = symbolic::add(symbolic::mul(ldc(), new_subset[0]), new_subset[1]);

    auto& scale_sum_tasklet =
        builder.add_tasklet(flush_block, data_flow::mul, "_out", {"_in1", "_in2"}, block.debug_info());
    builder.add_computational_memlet(flush_block, sum_final, scale_sum_tasklet, "_in1", {}, block.debug_info());
    if (auto const_node = dynamic_cast<data_flow::ConstantNode*>(alpha_node)) {
        auto& alpha_node_new =
            builder.add_constant(flush_block, const_node->data(), const_node->type(), block.debug_info());
        builder.add_computational_memlet(flush_block, alpha_node_new, scale_sum_tasklet, "_in2", {}, block.debug_info());
    } else {
        auto& alpha_node_new = builder.add_access(flush_block, alpha_node->data(), block.debug_info());
        builder.add_computational_memlet(flush_block, alpha_node_new, scale_sum_tasklet, "_in2", {}, block.debug_info());
    }

    std::string scaled_sum_temp = builder.find_new_name("scaled_sum_temp");
    builder.add_container(scaled_sum_temp, scalar_type);
    auto& scaled_sum_final = builder.add_access(flush_block, scaled_sum_temp, block.debug_info());
    builder.add_computational_memlet(
        flush_block, scale_sum_tasklet, "_out", scaled_sum_final, {}, scalar_type, block.debug_info()
    );

    auto& scale_input_tasklet =
        builder.add_tasklet(flush_block, data_flow::mul, "_out", {"_in1", "_in2"}, block.debug_info());
    builder.add_computational_memlet(
        flush_block, input_node_c_new, scale_input_tasklet, "_in1", {c_idx}, iedge_c->base_type(), iedge_c->debug_info()
    );
    if (auto const_node = dynamic_cast<data_flow::ConstantNode*>(beta_node)) {
        auto& beta_node_new =
            builder.add_constant(flush_block, const_node->data(), const_node->type(), block.debug_info());
        builder
            .add_computational_memlet(flush_block, beta_node_new, scale_input_tasklet, "_in2", {}, block.debug_info());
    } else {
        auto& beta_node_new = builder.add_access(flush_block, beta_node->data(), block.debug_info());
        builder
            .add_computational_memlet(flush_block, beta_node_new, scale_input_tasklet, "_in2", {}, block.debug_info());
    }

    std::string scaled_input_temp = builder.find_new_name("scaled_input_temp");
    builder.add_container(scaled_input_temp, scalar_type);
    auto& scaled_input_c = builder.add_access(flush_block, scaled_input_temp, block.debug_info());
    builder.add_computational_memlet(
        flush_block, scale_input_tasklet, "_out", scaled_input_c, {}, scalar_type, block.debug_info()
    );

    auto& flush_add_tasklet =
        builder.add_tasklet(flush_block, data_flow::add, "_out", {"_in1", "_in2"}, block.debug_info());
    auto& output_node_new = builder.add_access(flush_block, C_out_var, output_node->debug_info());
    builder.add_computational_memlet(
        flush_block, scaled_sum_final, flush_add_tasklet, "_in1", {}, scalar_type, block.debug_info()
    );
    builder.add_computational_memlet(
        flush_block, scaled_input_c, flush_add_tasklet, "_in2", {}, scalar_type, block.debug_info()
    );
    builder.add_computational_memlet(
        flush_block, flush_add_tasklet, "_out", output_node_new, {c_idx}, iedge_c->base_type(), iedge_c->debug_info()
    );


    // Clean up block
    builder.remove_memlet(block, *iedge_a);
    builder.remove_memlet(block, *iedge_b);
    builder.remove_memlet(block, *iedge_c);
    builder.remove_memlet(block, *alpha_edge);
    builder.remove_node(block, *alpha_node);
    builder.remove_memlet(block, *beta_edge);
    builder.remove_node(block, *beta_node);
    builder.remove_memlet(block, oedge);
    builder.remove_node(block, *input_node_a);
    builder.remove_node(block, *input_node_b);
    builder.remove_node(block, *input_node_c);
    builder.remove_node(block, *output_node);
    builder.remove_node(block, *this);
    builder.remove_child(parent, index + 1);

    return true;
}

std::unique_ptr<data_flow::DataFlowNode> GEMMNode::
    clone(size_t element_id, const graph::Vertex vertex, data_flow::DataFlowGraph& parent) const {
    auto node_clone = std::unique_ptr<GEMMNode>(new GEMMNode(
        element_id,
        this->debug_info(),
        vertex,
        parent,
        this->implementation_type_,
        this->precision_,
        this->layout_,
        this->trans_a_,
        this->trans_b_,
        this->m_,
        this->n_,
        this->k_,
        this->lda_,
        this->ldb_,
        this->ldc_
    ));
    return std::move(node_clone);
}

std::string GEMMNode::toStr() const {
    return LibraryNode::toStr() + "(" + static_cast<char>(precision_) + ", " +
           std::string(BLAS_Layout_to_short_string(layout_)) + ", " + BLAS_Transpose_to_char(trans_a_) +
           BLAS_Transpose_to_char(trans_b_) + ", " + m_->__str__() + ", " + n_->__str__() + ", " + k_->__str__() +
           ", " + lda_->__str__() + ", " + ldb_->__str__() + ", " + ldc_->__str__() + ")";
}

nlohmann::json GEMMNodeSerializer::serialize(const data_flow::LibraryNode& library_node) {
    const GEMMNode& gemm_node = static_cast<const GEMMNode&>(library_node);
    nlohmann::json j;

    serializer::JSONSerializer serializer;
    j["code"] = gemm_node.code().value();
    j["precision"] = gemm_node.precision();
    j["layout"] = gemm_node.layout();
    j["trans_a"] = gemm_node.trans_a();
    j["trans_b"] = gemm_node.trans_b();
    j["m"] = serializer.expression(gemm_node.m());
    j["n"] = serializer.expression(gemm_node.n());
    j["k"] = serializer.expression(gemm_node.k());
    j["lda"] = serializer.expression(gemm_node.lda());
    j["ldb"] = serializer.expression(gemm_node.ldb());
    j["ldc"] = serializer.expression(gemm_node.ldc());

    return j;
}

data_flow::LibraryNode& GEMMNodeSerializer::deserialize(
    const nlohmann::json& j, builder::StructuredSDFGBuilder& builder, structured_control_flow::Block& parent
) {
    // Assertions for required fields
    assert(j.contains("element_id"));
    assert(j.contains("code"));
    assert(j.contains("debug_info"));

    auto code = j["code"].get<std::string>();
    if (code != LibraryNodeType_GEMM.value()) {
        throw std::runtime_error("Invalid library node code");
    }

    // Extract debug info using JSONSerializer
    sdfg::serializer::JSONSerializer serializer;
    DebugInfo debug_info = serializer.json_to_debug_info(j["debug_info"]);

    auto precision = j.at("precision").get<BLAS_Precision>();
    auto layout = j.at("layout").get<BLAS_Layout>();
    auto trans_a = j.at("trans_a").get<BLAS_Transpose>();
    auto trans_b = j.at("trans_b").get<BLAS_Transpose>();
    auto m = SymEngine::Expression(j.at("m"));
    auto n = SymEngine::Expression(j.at("n"));
    auto k = SymEngine::Expression(j.at("k"));
    auto lda = SymEngine::Expression(j.at("lda"));
    auto ldb = SymEngine::Expression(j.at("ldb"));
    auto ldc = SymEngine::Expression(j.at("ldc"));

    auto implementation_type = j.at("implementation_type").get<std::string>();

    return builder.add_library_node<
        GEMMNode>(parent, debug_info, implementation_type, precision, layout, trans_a, trans_b, m, n, k, lda, ldb, ldc);
}

GEMMNodeDispatcher_BLAS::GEMMNodeDispatcher_BLAS(
    codegen::LanguageExtension& language_extension,
    const Function& function,
    const data_flow::DataFlowGraph& data_flow_graph,
    const GEMMNode& node
)
    : codegen::LibraryNodeDispatcher(language_extension, function, data_flow_graph, node) {}

void GEMMNodeDispatcher_BLAS::dispatch_code(
    codegen::PrettyPrinter& stream,
    codegen::PrettyPrinter& globals_stream,
    codegen::CodeSnippetFactory& library_snippet_factory
) {
    stream << "{" << std::endl;
    stream.setIndent(stream.indent() + 4);

    auto& gemm_node = static_cast<const GEMMNode&>(this->node_);

    sdfg::types::Scalar base_type(types::PrimitiveType::Void);
    switch (gemm_node.precision()) {
        case BLAS_Precision::h:
            base_type = types::Scalar(types::PrimitiveType::Half);
            break;
        case BLAS_Precision::s:
            base_type = types::Scalar(types::PrimitiveType::Float);
            break;
        case BLAS_Precision::d:
            base_type = types::Scalar(types::PrimitiveType::Double);
            break;
        default:
            throw std::runtime_error("Invalid BLAS_Precision value");
    }

    stream << "cblas_" << BLAS_Precision_to_string(gemm_node.precision()) << "gemm(";
    stream.setIndent(stream.indent() + 4);
    stream << BLAS_Layout_to_string(gemm_node.layout());
    stream << ", ";
    stream << BLAS_Transpose_to_string(gemm_node.trans_a());
    stream << ", ";
    stream << BLAS_Transpose_to_string(gemm_node.trans_b());
    stream << ", ";
    stream << this->language_extension_.expression(gemm_node.m());
    stream << ", ";
    stream << this->language_extension_.expression(gemm_node.n());
    stream << ", ";
    stream << this->language_extension_.expression(gemm_node.k());
    stream << ", ";
    stream << "alpha";
    stream << ", ";
    stream << "A";
    stream << ", ";
    stream << this->language_extension_.expression(gemm_node.lda());
    stream << ", ";
    stream << "B";
    stream << ", ";
    stream << this->language_extension_.expression(gemm_node.ldb());
    stream << ", ";
    stream << "beta";
    stream << ", ";
    stream << "C";
    stream << ", ";
    stream << this->language_extension_.expression(gemm_node.ldc());

    stream.setIndent(stream.indent() - 4);
    stream << ");" << std::endl;

    stream.setIndent(stream.indent() - 4);
    stream << "}" << std::endl;
}

GEMMNodeDispatcher_CUBLAS::GEMMNodeDispatcher_CUBLAS(
    codegen::LanguageExtension& language_extension,
    const Function& function,
    const data_flow::DataFlowGraph& data_flow_graph,
    const GEMMNode& node
)
    : codegen::LibraryNodeDispatcher(language_extension, function, data_flow_graph, node) {}

void GEMMNodeDispatcher_CUBLAS::dispatch_code(
    codegen::PrettyPrinter& stream,
    codegen::PrettyPrinter& globals_stream,
    codegen::CodeSnippetFactory& library_snippet_factory
) {
    throw std::runtime_error("GEMMNodeDispatcher_CUBLAS not implemented");
}

} // namespace blas
} // namespace math
} // namespace sdfg

1	#include "sdfg/data_flow/library_nodes/math/blas/gemm.h"
2
3	#include "sdfg/analysis/analysis.h"
4	#include "sdfg/builder/structured_sdfg_builder.h"
5
6	#include "sdfg/analysis/scope_analysis.h"
7
8	namespace sdfg {
9	namespace math {
10	namespace blas {
11
12	GEMMNode::GEMMNode(	1✔
13	size_t element_id,
14	const DebugInfo& debug_info,
15	const graph::Vertex vertex,
16	data_flow::DataFlowGraph& parent,
17	const data_flow::ImplementationType& implementation_type,
18	const BLAS_Precision& precision,
19	const BLAS_Layout& layout,
20	const BLAS_Transpose& trans_a,
21	const BLAS_Transpose& trans_b,
22	symbolic::Expression m,
23	symbolic::Expression n,
24	symbolic::Expression k,
25	symbolic::Expression lda,
26	symbolic::Expression ldb,
27	symbolic::Expression ldc
28	)
29	: MathNode(	1✔
30	element_id,	1✔
31	debug_info,	1✔
32	vertex,	1✔
33	parent,	1✔
34	LibraryNodeType_GEMM,
35	{"C"},	1✔
36	{"A", "B", "C", "alpha", "beta"},	1✔
37	implementation_type	1✔
38	),
39	precision_(precision), layout_(layout), trans_a_(trans_a), trans_b_(trans_b), m_(m), n_(n), k_(k), lda_(lda),	1✔
40	ldb_(ldb), ldc_(ldc) {}	2✔
41
42	BLAS_Precision GEMMNode::precision() const { return this->precision_; };	×
43
44	BLAS_Layout GEMMNode::layout() const { return this->layout_; };	×
45
46	BLAS_Transpose GEMMNode::trans_a() const { return this->trans_a_; };	×
47
48	BLAS_Transpose GEMMNode::trans_b() const { return this->trans_b_; };	×
49
50	symbolic::Expression GEMMNode::m() const { return this->m_; };	1✔
51
52	symbolic::Expression GEMMNode::n() const { return this->n_; };	1✔
53
54	symbolic::Expression GEMMNode::k() const { return this->k_; };	1✔
55
56	symbolic::Expression GEMMNode::lda() const { return this->lda_; };	1✔
57
58	symbolic::Expression GEMMNode::ldb() const { return this->ldb_; };	1✔
59
60	symbolic::Expression GEMMNode::ldc() const { return this->ldc_; };	1✔
61
62	void GEMMNode::validate(const Function& function) const {}	1✔
63
64	types::PrimitiveType GEMMNode::scalar_primitive() const {	1✔
65	switch (this->precision_) {	1✔
66	case BLAS_Precision::s:
67	return types::PrimitiveType::Float;	1✔
68	case BLAS_Precision::d:
69	return types::PrimitiveType::Double;	×
70	case BLAS_Precision::h:
71	return types::PrimitiveType::Half;	×
72	default:
73	return types::PrimitiveType::Void;	×
74	}
75	}	1✔
76
77	bool GEMMNode::expand(builder::StructuredSDFGBuilder& builder, analysis::AnalysisManager& analysis_manager) {	1✔
78	auto& scope_analysis = analysis_manager.get<analysis::ScopeAnalysis>();	1✔
79
80	auto& dataflow = this->get_parent();	1✔
81	auto& block = static_cast<structured_control_flow::Block&>(*dataflow.get_parent());	1✔
82	auto& parent = static_cast<structured_control_flow::Sequence&>(*scope_analysis.parent_scope(&block));	1✔
83	int index = parent.index(block);	1✔
84	auto& transition = parent.at(index).second;	1✔
85
86	if (trans_a_ != BLAS_Transpose::No \|\| trans_b_ != BLAS_Transpose::No) {	1✔
87	return false;	×
88	}
89
90	auto primitive_type = scalar_primitive();	1✔
91	if (primitive_type == types::PrimitiveType::Void) {	1✔
92	return false;	×
93	}
94
95	types::Scalar scalar_type(primitive_type);	1✔
96
97	auto in_edges = dataflow.in_edges(*this);	1✔
98	auto in_edges_it = in_edges.begin();	1✔
99
100	data_flow::Memlet* iedge_a = nullptr;	1✔
101	data_flow::Memlet* iedge_b = nullptr;	1✔
102	data_flow::Memlet* iedge_c = nullptr;	1✔
103	data_flow::Memlet* alpha_edge = nullptr;	1✔
104	data_flow::Memlet* beta_edge = nullptr;	1✔
105	while (in_edges_it != in_edges.end()) {	6✔
106	auto& edge = *in_edges_it;	5✔
107	auto dst_conn = edge.dst_conn();	5✔
108	if (dst_conn == "A") {	5✔
109	iedge_a = &edge;	1✔
110	} else if (dst_conn == "B") {	5✔
111	iedge_b = &edge;	1✔
112	} else if (dst_conn == "C") {	4✔
113	iedge_c = &edge;	1✔
114	} else if (dst_conn == "alpha") {	3✔
115	alpha_edge = &edge;	1✔
116	} else if (dst_conn == "beta") {	2✔
117	beta_edge = &edge;	1✔
118	} else {	1✔
119	throw InvalidSDFGException("GEMMNode has unexpected input: " + dst_conn);	×
120	}
121	++in_edges_it;	5✔
122	}	5✔
123
124	auto& oedge = dataflow.out_edges(this).begin();	1✔
125
126	// Checks if legal
127	auto* input_node_a = static_cast<data_flow::AccessNode*>(&iedge_a->src());	1✔
128	auto* input_node_b = static_cast<data_flow::AccessNode*>(&iedge_b->src());	1✔
129	auto* input_node_c = static_cast<data_flow::AccessNode*>(&iedge_c->src());	1✔
130	auto* output_node = static_cast<data_flow::AccessNode*>(&oedge.dst());	1✔
131	auto* alpha_node = static_cast<data_flow::AccessNode*>(&alpha_edge->src());	1✔
132	auto* beta_node = static_cast<data_flow::AccessNode*>(&beta_edge->src());	1✔
133
134	// we must be the only thing in this block, as we do not support splitting a block into pre, expanded lib-node, post
135	if (!input_node_a \|\| dataflow.in_degree(*input_node_a) != 0 \|\| !input_node_b \|\|	2✔
136	dataflow.in_degree(input_node_b) != 0 \|\| !input_node_c \|\| dataflow.in_degree(input_node_c) != 0 \|\|	1✔
137	!output_node \|\| dataflow.out_degree(*output_node) != 0) {	1✔
138	return false; // data nodes are not standalone	×
139	}
140	if (dataflow.in_degree(alpha_node) != 0 \|\| dataflow.in_degree(beta_node) != 0) {	1✔
141	return false; // alpha and beta are not standalone	×
142	}
143	for (auto* nd : dataflow.data_nodes()) {	7✔
144	if (nd != input_node_a && nd != input_node_b && nd != input_node_c && nd != output_node &&	7✔
145	(!alpha_node \|\| nd != alpha_node) && (!beta_node \|\| nd != beta_node)) {	2✔
146	return false; // there are other nodes in here that we could not preserve correctly	×
147	}
148	}
149
150	auto& A_var = input_node_a->data();	1✔
151	auto& B_var = input_node_b->data();	1✔
152	auto& C_in_var = input_node_c->data();	1✔
153	auto& C_out_var = output_node->data();	1✔
154
155
156	// Add new graph after the current block
157	auto& new_sequence = builder.add_sequence_before(parent, block, transition.assignments(), block.debug_info());	1✔
158
159	// Add maps
160	std::vector<symbolic::Expression> indvar_ends{this->m(), this->n(), this->k()};	1✔
161	data_flow::Subset new_subset;	1✔
162	structured_control_flow::Sequence* last_scope = &new_sequence;	1✔
163	structured_control_flow::Map* last_map = nullptr;	1✔
164	structured_control_flow::Map* output_loop = nullptr;	1✔
165	std::vector<std::string> indvar_names{"_i", "_j", "_k"};	1✔
166
167	std::string sum_var = builder.find_new_name("_sum");	1✔
168	builder.add_container(sum_var, scalar_type);	1✔
169
170	for (size_t i = 0; i < 3; i++) {	4✔
171	auto dim_begin = symbolic::zero();	3✔
172	auto& dim_end = indvar_ends[i];	3✔
173
174	std::string indvar_str = builder.find_new_name(indvar_names[i]);	3✔
175	builder.add_container(indvar_str, types::Scalar(types::PrimitiveType::UInt64));	3✔
176
177	auto indvar = symbolic::symbol(indvar_str);	3✔
178	auto init = dim_begin;	3✔
179	auto update = symbolic::add(indvar, symbolic::one());	3✔
180	auto condition = symbolic::Lt(indvar, dim_end);	3✔
181	last_map = &builder.add_map(	6✔
182	*last_scope,	3✔
183	indvar,
184	condition,
185	init,	3✔
186	update,
187	structured_control_flow::ScheduleType_Sequential::create(),	3✔
188	{},	3✔
189	block.debug_info()	3✔
190	);
191	last_scope = &last_map->root();	3✔
192
193	if (i == 1) {	3✔
194	output_loop = last_map;	1✔
195	}	1✔
196
197	new_subset.push_back(indvar);	3✔
198	}	3✔
199
200
201	// Add code
202	auto& init_block = builder.add_block_before(output_loop->root(), *last_map, {}, block.debug_info());	1✔
203	auto& sum_init = builder.add_access(init_block, sum_var, block.debug_info());	1✔
204
205	auto& zero_node = builder.add_constant(init_block, "0.0", alpha_edge->base_type(), block.debug_info());	1✔
206	auto& init_tasklet = builder.add_tasklet(init_block, data_flow::assign, "_out", {"_in"}, block.debug_info());	1✔
207	builder.add_computational_memlet(init_block, zero_node, init_tasklet, "_in", {}, block.debug_info());	1✔
208	builder.add_computational_memlet(init_block, init_tasklet, "_out", sum_init, {}, block.debug_info());	1✔
209
210	auto& code_block = builder.add_block(*last_scope, {}, block.debug_info());	1✔
211	auto& input_node_a_new = builder.add_access(code_block, A_var, input_node_a->debug_info());	1✔
212	auto& input_node_b_new = builder.add_access(code_block, B_var, input_node_b->debug_info());	1✔
213
214	auto& core_fma =	1✔
215	builder.add_tasklet(code_block, data_flow::fma, "_out", {"_in1", "_in2", "_in3"}, block.debug_info());	1✔
216	auto& sum_in = builder.add_access(code_block, sum_var, block.debug_info());	1✔
217	auto& sum_out = builder.add_access(code_block, sum_var, block.debug_info());	1✔
218	builder.add_computational_memlet(code_block, sum_in, core_fma, "_in3", {}, block.debug_info());	1✔
219
220	symbolic::Expression a_idx = symbolic::add(symbolic::mul(lda(), new_subset[0]), new_subset[2]);	1✔
221	builder.add_computational_memlet(	2✔
222	code_block, input_node_a_new, core_fma, "_in1", {a_idx}, iedge_a->base_type(), iedge_a->debug_info()	1✔
223	);
224	symbolic::Expression b_idx = symbolic::add(symbolic::mul(ldb(), new_subset[2]), new_subset[1]);	1✔
225	builder.add_computational_memlet(	2✔
226	code_block, input_node_b_new, core_fma, "_in2", {b_idx}, iedge_b->base_type(), iedge_b->debug_info()	1✔
227	);
228	builder.add_computational_memlet(code_block, core_fma, "_out", sum_out, {}, oedge.debug_info());	1✔
229
230	auto& flush_block = builder.add_block_after(output_loop->root(), *last_map, {}, block.debug_info());	1✔
231	auto& sum_final = builder.add_access(flush_block, sum_var, block.debug_info());	1✔
232	auto& input_node_c_new = builder.add_access(flush_block, C_in_var, input_node_c->debug_info());	1✔
233	symbolic::Expression c_idx = symbolic::add(symbolic::mul(ldc(), new_subset[0]), new_subset[1]);	1✔
234
235	auto& scale_sum_tasklet =	1✔
236	builder.add_tasklet(flush_block, data_flow::mul, "_out", {"_in1", "_in2"}, block.debug_info());	1✔
237	builder.add_computational_memlet(flush_block, sum_final, scale_sum_tasklet, "_in1", {}, block.debug_info());	1✔
238	if (auto const_node = dynamic_cast<data_flow::ConstantNode*>(alpha_node)) {	1✔
239	auto& alpha_node_new =	1✔
240	builder.add_constant(flush_block, const_node->data(), const_node->type(), block.debug_info());	1✔
241	builder.add_computational_memlet(flush_block, alpha_node_new, scale_sum_tasklet, "_in2", {}, block.debug_info());	1✔
242	} else {	1✔
243	auto& alpha_node_new = builder.add_access(flush_block, alpha_node->data(), block.debug_info());	×
NEW 244	builder.add_computational_memlet(flush_block, alpha_node_new, scale_sum_tasklet, "_in2", {}, block.debug_info());	×
245	}
246
247	std::string scaled_sum_temp = builder.find_new_name("scaled_sum_temp");	1✔
248	builder.add_container(scaled_sum_temp, scalar_type);	1✔
249	auto& scaled_sum_final = builder.add_access(flush_block, scaled_sum_temp, block.debug_info());	1✔
250	builder.add_computational_memlet(	2✔
251	flush_block, scale_sum_tasklet, "_out", scaled_sum_final, {}, scalar_type, block.debug_info()	1✔
252	);
253
254	auto& scale_input_tasklet =	1✔
255	builder.add_tasklet(flush_block, data_flow::mul, "_out", {"_in1", "_in2"}, block.debug_info());	1✔
256	builder.add_computational_memlet(	2✔
257	flush_block, input_node_c_new, scale_input_tasklet, "_in1", {c_idx}, iedge_c->base_type(), iedge_c->debug_info()	1✔
258	);
259	if (auto const_node = dynamic_cast<data_flow::ConstantNode*>(beta_node)) {	1✔
260	auto& beta_node_new =	1✔
261	builder.add_constant(flush_block, const_node->data(), const_node->type(), block.debug_info());	1✔
262	builder	2✔
263	.add_computational_memlet(flush_block, beta_node_new, scale_input_tasklet, "_in2", {}, block.debug_info());	1✔
264	} else {	1✔
265	auto& beta_node_new = builder.add_access(flush_block, beta_node->data(), block.debug_info());	×
NEW 266	builder	×
NEW 267	.add_computational_memlet(flush_block, beta_node_new, scale_input_tasklet, "_in2", {}, block.debug_info());	×
268	}
269
270	std::string scaled_input_temp = builder.find_new_name("scaled_input_temp");	1✔
271	builder.add_container(scaled_input_temp, scalar_type);	1✔
272	auto& scaled_input_c = builder.add_access(flush_block, scaled_input_temp, block.debug_info());	1✔
273	builder.add_computational_memlet(	2✔
274	flush_block, scale_input_tasklet, "_out", scaled_input_c, {}, scalar_type, block.debug_info()	1✔
275	);
276
277	auto& flush_add_tasklet =	1✔
278	builder.add_tasklet(flush_block, data_flow::add, "_out", {"_in1", "_in2"}, block.debug_info());	1✔
279	auto& output_node_new = builder.add_access(flush_block, C_out_var, output_node->debug_info());	1✔
280	builder.add_computational_memlet(	2✔
281	flush_block, scaled_sum_final, flush_add_tasklet, "_in1", {}, scalar_type, block.debug_info()	1✔
282	);
283	builder.add_computational_memlet(	2✔
284	flush_block, scaled_input_c, flush_add_tasklet, "_in2", {}, scalar_type, block.debug_info()	1✔
285	);
286	builder.add_computational_memlet(	2✔
287	flush_block, flush_add_tasklet, "_out", output_node_new, {c_idx}, iedge_c->base_type(), iedge_c->debug_info()	1✔
288	);
289
290
291	// Clean up block
292	builder.remove_memlet(block, *iedge_a);	1✔
293	builder.remove_memlet(block, *iedge_b);	1✔
294	builder.remove_memlet(block, *iedge_c);	1✔
295	builder.remove_memlet(block, *alpha_edge);	1✔
296	builder.remove_node(block, *alpha_node);	1✔
297	builder.remove_memlet(block, *beta_edge);	1✔
298	builder.remove_node(block, *beta_node);	1✔
299	builder.remove_memlet(block, oedge);	1✔
300	builder.remove_node(block, *input_node_a);	1✔
301	builder.remove_node(block, *input_node_b);	1✔
302	builder.remove_node(block, *input_node_c);	1✔
303	builder.remove_node(block, *output_node);	1✔
304	builder.remove_node(block, *this);	1✔
305	builder.remove_child(parent, index + 1);	1✔
306
307	return true;	1✔
308	}	1✔
309
310	std::unique_ptr<data_flow::DataFlowNode> GEMMNode::
311	clone(size_t element_id, const graph::Vertex vertex, data_flow::DataFlowGraph& parent) const {	×
312	auto node_clone = std::unique_ptr<GEMMNode>(new GEMMNode(	×
313	element_id,	×
314	this->debug_info(),	×
315	vertex,	×
316	parent,	×
317	this->implementation_type_,	×
318	this->precision_,	×
319	this->layout_,	×
320	this->trans_a_,	×
321	this->trans_b_,	×
322	this->m_,	×
323	this->n_,	×
324	this->k_,	×
325	this->lda_,	×
326	this->ldb_,	×
NEW 327	this->ldc_	×
328	));
329	return std::move(node_clone);	×
330	}	×
331
UNCOV 332	std::string GEMMNode::toStr() const {	×
UNCOV 333	return LibraryNode::toStr() + "(" + static_cast<char>(precision_) + ", " +	×
UNCOV 334	std::string(BLAS_Layout_to_short_string(layout_)) + ", " + BLAS_Transpose_to_char(trans_a_) +	×
UNCOV 335	BLAS_Transpose_to_char(trans_b_) + ", " + m_->__str__() + ", " + n_->__str__() + ", " + k_->__str__() +	×
UNCOV 336	", " + lda_->__str__() + ", " + ldb_->__str__() + ", " + ldc_->__str__() + ")";	×
337	}	×
338
339	nlohmann::json GEMMNodeSerializer::serialize(const data_flow::LibraryNode& library_node) {	×
340	const GEMMNode& gemm_node = static_cast<const GEMMNode&>(library_node);	×
341	nlohmann::json j;	×
342
343	serializer::JSONSerializer serializer;	×
344	j["code"] = gemm_node.code().value();	×
345	j["precision"] = gemm_node.precision();	×
346	j["layout"] = gemm_node.layout();	×
347	j["trans_a"] = gemm_node.trans_a();	×
348	j["trans_b"] = gemm_node.trans_b();	×
349	j["m"] = serializer.expression(gemm_node.m());	×
350	j["n"] = serializer.expression(gemm_node.n());	×
351	j["k"] = serializer.expression(gemm_node.k());	×
352	j["lda"] = serializer.expression(gemm_node.lda());	×
353	j["ldb"] = serializer.expression(gemm_node.ldb());	×
354	j["ldc"] = serializer.expression(gemm_node.ldc());	×
355
356	return j;	×
357	}	×
358
359	data_flow::LibraryNode& GEMMNodeSerializer::deserialize(	×
360	const nlohmann::json& j, builder::StructuredSDFGBuilder& builder, structured_control_flow::Block& parent
361	) {
362	// Assertions for required fields
363	assert(j.contains("element_id"));	×
364	assert(j.contains("code"));	×
365	assert(j.contains("debug_info"));	×
366
367	auto code = j["code"].get<std::string>();	×
368	if (code != LibraryNodeType_GEMM.value()) {	×
369	throw std::runtime_error("Invalid library node code");	×
370	}
371
372	// Extract debug info using JSONSerializer
373	sdfg::serializer::JSONSerializer serializer;	×
374	DebugInfo debug_info = serializer.json_to_debug_info(j["debug_info"]);	×
375
376	auto precision = j.at("precision").get<BLAS_Precision>();	×
377	auto layout = j.at("layout").get<BLAS_Layout>();	×
378	auto trans_a = j.at("trans_a").get<BLAS_Transpose>();	×
379	auto trans_b = j.at("trans_b").get<BLAS_Transpose>();	×
380	auto m = SymEngine::Expression(j.at("m"));	×
381	auto n = SymEngine::Expression(j.at("n"));	×
382	auto k = SymEngine::Expression(j.at("k"));	×
383	auto lda = SymEngine::Expression(j.at("lda"));	×
384	auto ldb = SymEngine::Expression(j.at("ldb"));	×
385	auto ldc = SymEngine::Expression(j.at("ldc"));	×
386
387	auto implementation_type = j.at("implementation_type").get<std::string>();	×
388
NEW 389	return builder.add_library_node<	×
NEW 390	GEMMNode>(parent, debug_info, implementation_type, precision, layout, trans_a, trans_b, m, n, k, lda, ldb, ldc);	×
UNCOV 391	}	×
392
393	GEMMNodeDispatcher_BLAS::GEMMNodeDispatcher_BLAS(	×
394	codegen::LanguageExtension& language_extension,
395	const Function& function,
396	const data_flow::DataFlowGraph& data_flow_graph,
397	const GEMMNode& node
398	)
399	: codegen::LibraryNodeDispatcher(language_extension, function, data_flow_graph, node) {}	×
400
NEW 401	void GEMMNodeDispatcher_BLAS::dispatch_code(	×
402	codegen::PrettyPrinter& stream,
403	codegen::PrettyPrinter& globals_stream,
404	codegen::CodeSnippetFactory& library_snippet_factory
405	) {
406	stream << "{" << std::endl;	×
407	stream.setIndent(stream.indent() + 4);	×
408
409	auto& gemm_node = static_cast<const GEMMNode&>(this->node_);	×
410
411	sdfg::types::Scalar base_type(types::PrimitiveType::Void);	×
412	switch (gemm_node.precision()) {	×
413	case BLAS_Precision::h:
414	base_type = types::Scalar(types::PrimitiveType::Half);	×
415	break;	×
416	case BLAS_Precision::s:
417	base_type = types::Scalar(types::PrimitiveType::Float);	×
418	break;	×
419	case BLAS_Precision::d:
420	base_type = types::Scalar(types::PrimitiveType::Double);	×
421	break;	×
422	default:
423	throw std::runtime_error("Invalid BLAS_Precision value");	×
424	}
425
426	stream << "cblas_" << BLAS_Precision_to_string(gemm_node.precision()) << "gemm(";	×
427	stream.setIndent(stream.indent() + 4);	×
428	stream << BLAS_Layout_to_string(gemm_node.layout());	×
429	stream << ", ";	×
430	stream << BLAS_Transpose_to_string(gemm_node.trans_a());	×
431	stream << ", ";	×
432	stream << BLAS_Transpose_to_string(gemm_node.trans_b());	×
433	stream << ", ";	×
434	stream << this->language_extension_.expression(gemm_node.m());	×
435	stream << ", ";	×
436	stream << this->language_extension_.expression(gemm_node.n());	×
437	stream << ", ";	×
438	stream << this->language_extension_.expression(gemm_node.k());	×
439	stream << ", ";	×
440	stream << "alpha";	×
441	stream << ", ";	×
442	stream << "A";	×
443	stream << ", ";	×
444	stream << this->language_extension_.expression(gemm_node.lda());	×
445	stream << ", ";	×
446	stream << "B";	×
447	stream << ", ";	×
448	stream << this->language_extension_.expression(gemm_node.ldb());	×
449	stream << ", ";	×
450	stream << "beta";	×
451	stream << ", ";	×
452	stream << "C";	×
453	stream << ", ";	×
454	stream << this->language_extension_.expression(gemm_node.ldc());	×
455
456	stream.setIndent(stream.indent() - 4);	×
457	stream << ");" << std::endl;	×
458
459	stream.setIndent(stream.indent() - 4);	×
460	stream << "}" << std::endl;	×
461	}	×
462
463	GEMMNodeDispatcher_CUBLAS::GEMMNodeDispatcher_CUBLAS(	×
464	codegen::LanguageExtension& language_extension,
465	const Function& function,
466	const data_flow::DataFlowGraph& data_flow_graph,
467	const GEMMNode& node
468	)
469	: codegen::LibraryNodeDispatcher(language_extension, function, data_flow_graph, node) {}	×
470
NEW 471	void GEMMNodeDispatcher_CUBLAS::dispatch_code(	×
472	codegen::PrettyPrinter& stream,
473	codegen::PrettyPrinter& globals_stream,
474	codegen::CodeSnippetFactory& library_snippet_factory
475	) {
476	throw std::runtime_error("GEMMNodeDispatcher_CUBLAS not implemented");	×
477	}	×
478
479	} // namespace blas
480	} // namespace math
481	} // namespace sdfg

daisytuner / sdfglib / 17655994367

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