matmul_node.h

Go to the documentation of this file.

 
#pragma once
 
#include <optional>
 
#include "sdfg/data_flow/library_nodes/math/tensor/tensor_layout.h"
#include "sdfg/data_flow/library_nodes/math/tensor/tensor_node.h"
 
#include "sdfg/serializer/json_serializer.h"
#include "sdfg/symbolic/symbolic.h"
 
namespace sdfg {
namespace math {
namespace tensor {
 
inline data_flow::LibraryNodeCode LibraryNodeType_MatMul("ml::MatMul");

class MatMulNode : public TensorNode {
private:
    ContractionTypes types_;
    TensorLayout layout_a_;
    TensorLayout layout_b_;
    TensorLayout layout_y_;

    static bool has_basic_strides(symbolic::MultiExpression shape, symbolic::MultiExpression strides);
    static bool has_transposed_strides(symbolic::MultiExpression shape, symbolic::MultiExpression strides);
 
public:
    MatMulNode(
        size_t element_id,
        const DebugInfo& debug_info,
        const graph::Vertex vertex,
        data_flow::DataFlowGraph& parent,
        const TensorLayout& layout_a,
        const TensorLayout& layout_b,
        QuantizationType quantization = QUANTIZATION_MATCH_INPUTS,
        const TensorLayout* layout_y = nullptr,
        const data_flow::ImplementationType& impl_type = data_flow::ImplementationType_NONE
    );

    MatMulNode(
        size_t element_id,
        const DebugInfo& debug_info,
        const graph::Vertex vertex,
        data_flow::DataFlowGraph& parent,
        const TensorLayout& layout_a,
        const TensorLayout& layout_b,
        const ContractionTypes& types,
        const TensorLayout* layout_y = nullptr,
        const data_flow::ImplementationType& impl_type = data_flow::ImplementationType_NONE
    );
 
    static auto constexpr Y_INPUT_IDX = 0;
    static auto constexpr A_INPUT_IDX = 1;
    static auto constexpr B_INPUT_IDX = 2;

    static constexpr const char* OPERAND_METADATA_KEY = "ml.matmul.operand";
 

    symbolic::Expression m() const;

    symbolic::Expression n() const;

    const ContractionTypes& types() const;
 
    void set_types(const ContractionTypes& types);

    types::PrimitiveType input_type() const;

    types::PrimitiveType accumulator_type() const;

    types::PrimitiveType output_type() const;

    std::optional<types::PrimitiveType> uniform_quantization(const data_flow::DataFlowGraph& dataflow) const;
 
    [[deprecated("Use input_type(), accumulator_type(), output_type() or uniform_quantization() instead")]]
    QuantizationType quantization() const {
        return quantization(get_parent());
    }
 
    [[deprecated("Use input_type(), accumulator_type(), output_type() or uniform_quantization() instead")]]
    QuantizationType quantization(const data_flow::DataFlowGraph& dataflow) const;
 
    [[deprecated("Use input_type() instead")]]
    types::PrimitiveType input_quantization() const;
 
    [[deprecated("Use output_type() instead")]]
    types::PrimitiveType output_quantization() const;

    [[deprecated("Use types() instead")]]
    QuantizationType fixed_quantization() const;
 
    [[deprecated("Use set_types(ContractionTypes::uniform(...)) instead")]]
    void set_fixed_quantization(const QuantizationType quant);
 
    const TensorLayout& layout_a() const;
 
    const TensorLayout& layout_b() const;
 
    const TensorLayout& layout_y() const;
 
    TensorLayout& layout_a() {
        return layout_a_;
    }
 
    TensorLayout& layout_b() {
        return layout_b_;
    }
 
    TensorLayout& layout_y() {
        return layout_y_;
    }
 
    static TensorLayout get_linear_result_layout(const TensorLayout& layout_a, const TensorLayout& layout_b);

    symbolic::Expression k() const;
 
    void validate(const Function& function) const override;
 
    void verify_data_types(const data_flow::DataFlowGraph& graph) const override;

    virtual passes::LibNodeExpander::ExpandOutcome
    expand(passes::LibNodeExpander::ExpandContext&, structured_control_flow::Block& block) override;
 
    symbolic::SymbolSet symbols() const override;
 
    void replace(const symbolic::Expression old_expression, const symbolic::Expression new_expression) override;
 
    void replace(const symbolic::ExpressionMapping& replacements) override;
 
    std::unique_ptr<data_flow::DataFlowNode>
    clone(size_t element_id, const graph::Vertex vertex, data_flow::DataFlowGraph& parent) const override;
 
    std::string toStr() const override;
 
    symbolic::Expression flop() const override;
 
    data_flow::PointerAccessType pointer_access_type(int input_idx) const override;
 
    bool relocalize_operand(int input_idx, const TensorLayout& packed) override;
    bool can_relocalize_operand(int input_idx, const TensorLayout& packed) const override;
 
    bool supports_integer_types() const override {
        return true;
    }
};

class MatMulNodeSerializer : public serializer::LibraryNodeSerializer {
public:
    nlohmann::json serialize(const data_flow::LibraryNode& library_node) override;
 
    data_flow::LibraryNode& deserialize(
        const nlohmann::json& j, builder::StructuredSDFGBuilder& builder, structured_control_flow::Block& parent
    ) override;
};
 
class MatMulNaiveExpander : public passes::CodeLibNodeExpander<MatMulNode> {
public:
    MatMulNaiveExpander() : passes::CodeLibNodeExpander<MatMulNode>(LibraryNodeType_MatMul) {
    }
 
    passes::LibNodeExpander::ExpandOutcome handle_expand(
        passes::LibNodeExpander::ExpandContext& context, structured_control_flow::Block& block, MatMulNode& node
    ) const override;
};
 
class MatMulBlasExpander : public passes::CodeLibNodeExpander<MatMulNode> {
private:
    bool supports_half_; //< Not all BLAS implementations support half precision (hgemm)
 
public:
    MatMulBlasExpander(bool supports_half = false)
        : passes::CodeLibNodeExpander<MatMulNode>(LibraryNodeType_MatMul), supports_half_(supports_half) {
    }
 
 
    passes::LibNodeExpander::ExpandOutcome handle_expand(
        passes::LibNodeExpander::ExpandContext& context, structured_control_flow::Block& block, MatMulNode& node
    ) const override;
};
 
} // namespace tensor
} // namespace math
} // namespace sdfg