ozaki_contraction.h

Go to the documentation of this file.

#pragma once
 
#include <nlohmann/json_fwd.hpp>
#include <string>
#include <vector>
 
#include "sdfg/analysis/analysis.h"
#include "sdfg/builder/structured_sdfg_builder.h"
#include "sdfg/data_flow/library_node.h"
#include "sdfg/fp_math/transformations/quantize_contraction.h"
#include "sdfg/transformations/transformation.h"
 
namespace sdfg {
namespace fp_math {

enum class OzakiVariant {
    Slicing,
    Modular,
};

class OzakiContraction : public transformations::Transformation {
private:
    data_flow::LibraryNode& node_;
    OzakiVariant variant_;
    int splits_;
    Rescaling rescaling_;
 
public:
    static constexpr int MAX_SLICES = 10;

    static const std::vector<int64_t>& moduli();
    static constexpr int MAX_MODULI = 20;

    explicit OzakiContraction(
        data_flow::LibraryNode& node,
        OzakiVariant variant = OzakiVariant::Modular,
        int splits = 0,
        Rescaling rescaling = Rescaling::Exponents
    );
 
    virtual std::string name() const override;
 
    virtual bool
    can_be_applied(builder::StructuredSDFGBuilder& builder, analysis::AnalysisManager& analysis_manager) override;
 
    virtual void apply(builder::StructuredSDFGBuilder& builder, analysis::AnalysisManager& analysis_manager) override;
 
    virtual void to_json(nlohmann::json& j) const override;
 
    static OzakiContraction from_json(builder::StructuredSDFGBuilder& builder, const nlohmann::json& j);
 
    OzakiVariant variant() const {
        return variant_;
    }
 
    Rescaling rescaling() const {
        return rescaling_;
    }

    int effective_splits() const;

    static int modular_bits(int moduli, int64_t k);

    int matmuls() const;
};
 
} // namespace fp_math
} // namespace sdfg