tvm_tirx.tirx.transform.FP8ComputeLegalize

tvm_tirx.tirx.transform.FP8ComputeLegalize(promote_dtype='float32')

Promote unsupported FP8 computations and internal temporaries.

Inserts conversions at FP8 boundaries. Leaves functions unchanged when their CUDA target reports native FP8 support.

Parameters:

promote_dtype (str) – The data type we promote fp8 to, options: float16/float32.

Returns:

fpass (tvm_tirx.ir.transform.Pass) – The result pass

Examples

import tvm_tirx as tvm
from tvm_tirx.script import ir as I, tirx as T
from tvm_tirx.tirx import transform as tr

@I.ir_module
class Before:
    @T.function(private=True)
    def main():
        A = T.alloc_tensor(ty_args=[T.Tensor((1,), "float8_e5m2")])
        A[0] = T.float8_e5m2(1)
        T.evaluate(A[0] + T.float8_e5m2(2))

@I.ir_module
class Expected:
    @T.function(private=True)
    def main():
        A = T.alloc_tensor(ty_args=[T.Tensor((1,), T.float16)])
        A[0] = T.float16(1.0)
        T.evaluate(A[0] + T.float16(2.0))

def run_pass():
    return tr.FP8ComputeLegalize("float16")(Before)

tvm.ir.assert_structural_equal(run_pass(), Expected)