mirror of
https://github.com/tile-ai/tilelang.git
synced 2026-10-22 06:50:25 +08:00
57 lines
2.0 KiB
Python
57 lines
2.0 KiB
Python
import numpy as np
|
|
|
|
import tilelang
|
|
import tilelang.language as T
|
|
import tilelang.testing
|
|
from tilelang import tvm
|
|
from tilelang.quantize.quantization import (
|
|
_tir_u8_to_f8_e4m3_to_f16,
|
|
_tir_u8_to_f8_e4m3_to_f16_naive,
|
|
)
|
|
|
|
|
|
def _e4m3_to_float16(code: int) -> np.float16:
|
|
sign = -1.0 if code & 0x80 else 1.0
|
|
exponent = (code >> 3) & 0xF
|
|
mantissa = code & 0x7
|
|
|
|
if exponent == 0:
|
|
return np.float16(sign * np.ldexp(mantissa / 8.0, -6))
|
|
if exponent == 0xF and mantissa == 0x7:
|
|
return np.float16(np.nan)
|
|
return np.float16(sign * np.ldexp(1.0 + mantissa / 8.0, exponent - 7))
|
|
|
|
|
|
@tilelang.testing.requires_llvm
|
|
def test_u8_to_f8_e4m3_to_f16_all_codes():
|
|
@T.prim_func
|
|
def decode(
|
|
values: T.Tensor((256,), "uint8"),
|
|
naive: T.Tensor((256,), "uint16"),
|
|
optimized: T.Tensor((256,), "uint16"),
|
|
):
|
|
for i in T.serial(256):
|
|
naive[i] = T.reinterpret(_tir_u8_to_f8_e4m3_to_f16_naive(8, values[i], T.float16), T.uint16)
|
|
optimized[i] = T.reinterpret(_tir_u8_to_f8_e4m3_to_f16(8, values[i], T.float16), T.uint16)
|
|
|
|
built = tvm.compile(decode, target="llvm")
|
|
values = tvm.runtime.tensor(np.arange(256, dtype=np.uint8))
|
|
naive = tvm.runtime.empty((256,), "uint16", tvm.cpu())
|
|
optimized = tvm.runtime.empty((256,), "uint16", tvm.cpu())
|
|
built(values, naive, optimized)
|
|
|
|
naive_bits = naive.numpy()
|
|
optimized_bits = optimized.numpy()
|
|
expected = np.array([_e4m3_to_float16(code) for code in range(256)], dtype=np.float16)
|
|
finite = ~np.isnan(expected)
|
|
|
|
np.testing.assert_array_equal(naive_bits[finite], expected[finite].view(np.uint16))
|
|
np.testing.assert_array_equal(optimized_bits[finite], expected[finite].view(np.uint16))
|
|
np.testing.assert_array_equal(np.isnan(naive_bits.view(np.float16)), np.isnan(expected))
|
|
np.testing.assert_array_equal(np.isnan(optimized_bits.view(np.float16)), np.isnan(expected))
|
|
np.testing.assert_array_equal(naive_bits, optimized_bits)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
tilelang.testing.main()
|