Emulates a floating-point matmul (MatMul or GEMM node) with exact Int8 matmuls (Ozaki scheme), at about the input's precision. More...

#include <sdfg/fp_math/transformations/ozaki_contraction.h>

Inheritance diagram for sdfg::fp_math::OzakiContraction:

[legend]

Public Member Functions
	OzakiContraction (data_flow::LibraryNode &node, OzakiVariant variant=OzakiVariant::Modular, int splits=0, Rescaling rescaling=Rescaling::Exponents)
virtual std::string	name () const override
	Get the name of this transformation.
virtual bool	can_be_applied (builder::StructuredSDFGBuilder &builder, analysis::AnalysisManager &analysis_manager) override
	Check if this transformation can be safely applied.
virtual void	apply (builder::StructuredSDFGBuilder &builder, analysis::AnalysisManager &analysis_manager) override
	Apply this transformation to the SDFG.
virtual void	to_json (nlohmann::json &j) const override
	Serialize this transformation to JSON.
OzakiVariant	variant () const
Rescaling	rescaling () const
int	effective_splits () const
	The slices or moduli the transformation uses for `node`, resolving 0 to the default.
int	matmuls () const
	Int8 matmuls the transformation computes.
Public Member Functions inherited from sdfg::transformations::Transformation
virtual	~Transformation ()=default
virtual void	set_report (PassReportConsumer *report)
	Set the report consumer for transformation feedback.
virtual bool	try_apply (builder::StructuredSDFGBuilder &builder, analysis::AnalysisManager &analysis_manager)
	Attempt to apply the transformation if possible.

Static Public Member Functions
static const std::vector< int64_t > &	moduli ()
	Odd primes below 256, largest first; the first MAX_MODULI serve as the moduli of the modular variant.
static OzakiContraction	from_json (builder::StructuredSDFGBuilder &builder, const nlohmann::json &j)
static int	modular_bits (int moduli, int64_t k)
	Bits each of A and B keeps per row or column in the modular variant with `moduli` and reduction `k`.

Static Public Attributes
static constexpr int	MAX_SLICES = 10
static constexpr int	MAX_MODULI = 20

Additional Inherited Members
Protected Attributes inherited from sdfg::transformations::Transformation
PassReportConsumer *	report_ = nullptr

Detailed Description

Emulates a floating-point matmul (MatMul or GEMM node) with exact Int8 matmuls (Ozaki scheme), at about the input's precision.

Every row of A and column of B is scaled by a power of two, which is exact, and represented in Int8 matrices whose products the target computes exactly with Int32 accumulation:

Slicing (Ozaki-I). A[i, :] / 2^e_i = sum_t 2^-7t A_t[i, :] with Int8 slices A_t, and likewise for B. Y = 2^(e_i + f_j) sum_{t+u <= s+1} 2^-7(t+u) (A_t @ B_u), dropping the products below the precision.
Modular (Ozaki-II). A' = trunc(2^q A[i, :] / 2^e_i) and B' are integers below 2^q; their product is exact modulo P = prod m_l as long as K 2^2q < P / 2. Each modulus takes one product of the Int8 residues; Garner's algorithm turns the residues of the result into mixed-radix digits, which are summed in Double.

Both run in Double, the scales applied through their exponents (Rescaling::Exponents). With Rescaling::Direct, slicing instead runs in the contraction's own type with nothing beyond the arithmetic QuantizeContraction's direct rescaling needs, for targets without Double, floor or exact powers of two: with q = 2^7 - 1, s the largest magnitude of a row of A (or column of B) and r_1 = q x / s, the slices are d_t = rint(r_t) and r_{t+1} = 2q (r_t - d_t), each in [-q, q], so x = (s / q) sum_t (2q)^(1-t) d_t up to the last remainder, and Y = (s_A[i] / q) (s_B[j] / q) sum_{t+u <= s+1} (2q)^(2-t-u) (A_t @ B_u). Each slice is computed by a nest of its own that reads the slices before it back from memory, as a target needs that rounds only when it stores Int8. Its accuracy is that of the contraction's type, and it applies to batched MatMuls too, with scales per matrix.

Accuracy is relative to the magnitudes of each row of A and column of B: elements much smaller than their row's largest lose bits, as in all Ozaki schemes. The Int8 matmuls stay library nodes for the target to map.

Applies to a 2-D Float or Double MatMul, or a row-major GEMM (whose alpha and beta it applies to the result), with a constant K small enough for exact Int32 accumulation; with direct rescaling, slicing only, also to a MatMul that repeats over batch dimensions.

Constructor & Destructor Documentation

◆ OzakiContraction()

sdfg::fp_math::OzakiContraction::OzakiContraction	(	data_flow::LibraryNode &	node,
		OzakiVariant	variant = OzakiVariant::Modular,
		int	splits = 0,
		Rescaling	rescaling = Rescaling::Exponents )

explicit

Parameters

splits	slices (Slicing) or moduli (Modular); 0 picks as many as the output precision needs
rescaling	Direct computes the slices and the result in the contraction's own type; Slicing only

Member Function Documentation

◆ apply()

virtual void sdfg::fp_math::OzakiContraction::apply	(	builder::StructuredSDFGBuilder &	builder,
		analysis::AnalysisManager &	analysis_manager )

overridevirtual

Apply this transformation to the SDFG.

Parameters

builder	The SDFG builder
analysis_manager	The analysis manager (may be invalidated)

Implements sdfg::transformations::Transformation.

◆ can_be_applied()

virtual bool sdfg::fp_math::OzakiContraction::can_be_applied	(	builder::StructuredSDFGBuilder &	builder,
		analysis::AnalysisManager &	analysis_manager )

overridevirtual

Check if this transformation can be safely applied.

Parameters

builder	The SDFG builder
analysis_manager	The analysis manager providing necessary analyses

Returns: true if all preconditions are met

Implements sdfg::transformations::Transformation.

◆ effective_splits()

int sdfg::fp_math::OzakiContraction::effective_splits ( ) const

The slices or moduli the transformation uses for node, resolving 0 to the default.

◆ from_json()

OzakiContraction sdfg::fp_math::OzakiContraction::from_json	(	builder::StructuredSDFGBuilder &	builder,
		const nlohmann::json &	j )

static

◆ matmuls()

int sdfg::fp_math::OzakiContraction::matmuls ( ) const

Int8 matmuls the transformation computes.

◆ modular_bits()

int sdfg::fp_math::OzakiContraction::modular_bits	(	int	moduli,
		int64_t	k )

static

Bits each of A and B keeps per row or column in the modular variant with moduli and reduction k.

◆ moduli()

const std::vector< int64_t > & sdfg::fp_math::OzakiContraction::moduli ( )

static

Odd primes below 256, largest first; the first MAX_MODULI serve as the moduli of the modular variant.

◆ name()

virtual std::string sdfg::fp_math::OzakiContraction::name ( ) const

overridevirtual

Get the name of this transformation.

Returns: The transformation name (e.g., "LoopInterchange")

Implements sdfg::transformations::Transformation.

◆ rescaling()

Rescaling sdfg::fp_math::OzakiContraction::rescaling ( ) const

inline

◆ to_json()

virtual void sdfg::fp_math::OzakiContraction::to_json ( nlohmann::json & j ) const

overridevirtual

Serialize this transformation to JSON.

Parameters

j	JSON object to populate with transformation description

Exceptions

std::logic_error if not implemented

Implements sdfg::transformations::Transformation.

◆ variant()

OzakiVariant sdfg::fp_math::OzakiContraction::variant ( ) const

inline

Member Data Documentation

◆ MAX_MODULI

int sdfg::fp_math::OzakiContraction::MAX_MODULI = 20

staticconstexpr

◆ MAX_SLICES

int sdfg::fp_math::OzakiContraction::MAX_SLICES = 10

staticconstexpr

The documentation for this class was generated from the following file:

opt/include/sdfg/fp_math/transformations/ozaki_contraction.h

Public Member Functions

Static Public Member Functions

Static Public Attributes

Additional Inherited Members

Detailed Description

Constructor & Destructor Documentation

◆ OzakiContraction()

Member Function Documentation

◆ apply()

◆ can_be_applied()

◆ effective_splits()

◆ from_json()

◆ matmuls()

◆ modular_bits()

◆ moduli()

◆ name()

◆ rescaling()

◆ to_json()

◆ variant()

Member Data Documentation

◆ MAX_MODULI

◆ MAX_SLICES