from llvmlite import ir

from numba.cuda.cudadrv import nvvm
from numba.cuda.testing import unittest, ContextResettingTestCase
from numba.cuda.testing import skip_on_cudasim


@skip_on_cudasim("Inline PTX cannot be used in the simulator")
class TestCudaInlineAsm(ContextResettingTestCase):
    def test_inline_rsqrt(self):
        mod = ir.Module(__name__)
        mod.triple = "nvptx64-nvidia-cuda"
        nvvm.add_ir_version(mod)
        fnty = ir.FunctionType(ir.VoidType(), [ir.PointerType(ir.FloatType())])
        fn = ir.Function(mod, fnty, "cu_rsqrt")
        bldr = ir.IRBuilder(fn.append_basic_block("entry"))

        rsqrt_approx_fnty = ir.FunctionType(ir.FloatType(), [ir.FloatType()])
        inlineasm = ir.InlineAsm(
            rsqrt_approx_fnty,
            "rsqrt.approx.f32 $0, $1;",
            "=f,f",
            side_effect=True,
        )
        val = bldr.load(fn.args[0])
        res = bldr.call(inlineasm, [val])

        bldr.store(res, fn.args[0])
        bldr.ret_void()

        # generate ptx
        mod.data_layout = nvvm.NVVM().data_layout
        nvvm.set_cuda_kernel(fn)
        nvvmir = str(mod)
        ptx = nvvm.compile_ir(nvvmir)
        self.assertTrue("rsqrt.approx.f32" in str(ptx))


if __name__ == "__main__":
    unittest.main()
