tvm_tirx.tirx.transform.FP8StorageLegalize

tvm_tirx.tirx.transform.FP8StorageLegalize()

Represent FP8 storage as uint8 while preserving its bit patterns.

Run after compute legalization. Functions whose CUDA target reports native FP8 support are unchanged.

Returns:

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

Examples

After packed-API lowering, pointer storage uses integer elements.

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.Ptr(T.float8_e5m2)):
        T.evaluate(0)

@I.ir_module
class Expected:
    @T.function(private=True)
    def main(A: T.Ptr(T.uint8)):
        pass

def run_pass():
    return tr.FP8StorageLegalize()(Before)

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