"""默认的 Megacore 模式下，原生 XLA 把 f32[64,128] 乘以 2：观察 XLA 怎样把工作分给两个 TensorCore。"""
import tpu_init
tpu_init.initialize_one_chip()

import jax
import jax.numpy as jnp
import numpy as np

import tpuasm_tools

def main() -> None:
    x = jnp.arange(64 * 128, dtype=jnp.float32).reshape(64, 128)
    compiled = tpuasm_tools.compile(lambda x: x * 2.0, x)
    np.testing.assert_array_equal(np.asarray(compiled(x)), np.asarray(x) * 2.0)
    print('数值检查通过')
    listing = tpuasm_tools.kernel_listing(compiled)
    print('\n# fusion 段的指令统计')
    tpuasm_tools.print_mnemonic_counts(listing)
    print('\n# fusion 段')
    print(listing)

if __name__ == '__main__':
    main()