from __future__ import annotations


__copyright__ = """
Copyright 2008-2021 Andreas Kloeckner
Copyright 2021 NVIDIA Corporation
"""

import numpy as np
import numpy.linalg as la
import pytest

import pycuda.driver as drv
import pycuda.gpuarray as gpuarray
from pycuda.compiler import SourceModule
from pycuda.tools import dtype_to_ctype, mark_cuda_test


class TestDriver:
    @mark_cuda_test
    def test_memory(self):
        z = np.random.randn(400).astype(np.float32)
        new_z = drv.from_device_like(drv.to_device(z), z)
        assert la.norm(new_z - z) == 0

    @mark_cuda_test
    def test_simple_kernel(self):
        mod = SourceModule(
            """
        __global__ void multiply_them(float *dest, float *a, float *b)
        {
          const int i = threadIdx.x;
          dest[i] = a[i] * b[i];
        }
        """
        )

        multiply_them = mod.get_function("multiply_them")

        a = np.random.randn(400).astype(np.float32)
        b = np.random.randn(400).astype(np.float32)

        dest = np.zeros_like(a)
        multiply_them(drv.Out(dest), drv.In(a), drv.In(b), block=(400, 1, 1))
        assert la.norm(dest - a * b) == 0

    @mark_cuda_test
    def test_simple_kernel_2(self):
        mod = SourceModule(
            """
        __global__ void multiply_them(float *dest, float *a, float *b)
        {
          const int i = threadIdx.x;
          dest[i] = a[i] * b[i];
        }
        """
        )

        multiply_them = mod.get_function("multiply_them")

        a = np.random.randn(400).astype(np.float32)
        b = np.random.randn(400).astype(np.float32)
        a_gpu = drv.to_device(a)
        b_gpu = drv.to_device(b)

        dest = np.zeros_like(a)
        multiply_them(drv.Out(dest), a_gpu, b_gpu, block=(400, 1, 1))
        assert la.norm(dest - a * b) == 0

        drv.Context.synchronize()
        # now try with offsets
        dest = np.zeros_like(a)
        multiply_them(
            drv.Out(dest), np.intp(a_gpu) + a.itemsize, b_gpu, block=(399, 1, 1)
        )

        assert la.norm(dest[:-1] - a[1:] * b[:-1]) == 0

    @mark_cuda_test
    def test_vector_types(self):
        mod = SourceModule(
            """
        __global__ void set_them(float3 *dest, float3 x)
        {
          const int i = threadIdx.x;
          dest[i] = x;
        }
        """
        )

        set_them = mod.get_function("set_them")
        a = gpuarray.vec.make_float3(1, 2, 3)
        dest = np.empty((400), gpuarray.vec.float3)

        set_them(drv.Out(dest), a, block=(400, 1, 1))
        assert (dest == a).all()

    @mark_cuda_test
    def test_streamed_kernel(self):
        # this differs from the "simple_kernel" case in that *all* computation
        # and data copying is asynchronous. Observe how this necessitates the
        # use of page-locked memory.

        mod = SourceModule(
            """
        __global__ void multiply_them(float *dest, float *a, float *b)
        {
          const int i = threadIdx.x*blockDim.y + threadIdx.y;
          dest[i] = a[i] * b[i];
        }
        """
        )

        multiply_them = mod.get_function("multiply_them")

        shape = (32, 8)
        a = drv.pagelocked_zeros(shape, dtype=np.float32)
        b = drv.pagelocked_zeros(shape, dtype=np.float32)
        a[:] = np.random.randn(*shape)
        b[:] = np.random.randn(*shape)

        a_gpu = drv.mem_alloc(a.nbytes)
        b_gpu = drv.mem_alloc(b.nbytes)

        strm = drv.Stream()
        drv.memcpy_htod_async(a_gpu, a, strm)
        drv.memcpy_htod_async(b_gpu, b, strm)
        strm.synchronize()

        dest = drv.pagelocked_empty_like(a)
        multiply_them(drv.Out(dest), a_gpu, b_gpu, block=(*shape, 1), stream=strm)
        strm.synchronize()

        drv.memcpy_dtoh_async(a, a_gpu, strm)
        drv.memcpy_dtoh_async(b, b_gpu, strm)
        strm.synchronize()

        assert la.norm(dest - a * b) == 0

    @mark_cuda_test
    def test_gpuarray(self):
        a = np.arange(200000, dtype=np.float32)
        b = a + 17
        import pycuda.gpuarray as gpuarray

        a_g = gpuarray.to_gpu(a)
        b_g = gpuarray.to_gpu(b)
        diff = (a_g - 3 * b_g + (-a_g)).get() - (a - 3 * b + (-a))
        assert la.norm(diff) == 0

        diff = (a_g * b_g).get() - a * b
        assert la.norm(diff) == 0

    @mark_cuda_test
    def test_gpuarray_cai(self):
        a = np.zeros(10, dtype=np.float32)
        a_g = gpuarray.to_gpu(a)
        cai = a_g.__cuda_array_interface__
        ptr = cai["data"][0]
        masked = cai["data"][1]

        assert cai["shape"] == a.shape
        assert cai["strides"] == a.strides
        assert cai["typestr"] == a.dtype.str
        assert isinstance(ptr, int)
        assert ptr != 0
        assert not masked
        assert cai["stream"] is None
        assert cai["version"] == 3

    @mark_cuda_test
    def donottest_cublas_mixing(self):
        self.test_streamed_kernel()

        import pycuda.blas as blas

        shape = (10,)
        a = blas.ones(shape, dtype=np.float32)
        b = 33 * blas.ones(shape, dtype=np.float32)
        assert ((-a + b).from_gpu() == 32).all()

        self.test_streamed_kernel()

    @mark_cuda_test
    def test_2d_texture(self):
        if drv.get_driver_version() // 1000 >= 12:
            pytest.skip("texture references were removed in CUDA 12")

        mod = SourceModule(
            """
        texture<float, 2, cudaReadModeElementType> mtx_tex;

        __global__ void copy_texture(float *dest)
        {
          int row = threadIdx.x;
          int col = threadIdx.y;
          int w = blockDim.y;
          dest[row*w+col] = tex2D(mtx_tex, row, col);
        }
        """
        )

        copy_texture = mod.get_function("copy_texture")
        mtx_tex = mod.get_texref("mtx_tex")

        shape = (3, 4)
        a = np.random.randn(*shape).astype(np.float32)
        drv.matrix_to_texref(a, mtx_tex, order="F")

        dest = np.zeros(shape, dtype=np.float32)
        copy_texture(drv.Out(dest), block=(*shape, 1), texrefs=[mtx_tex])
        assert la.norm(dest - a) == 0

    @mark_cuda_test
    def test_multiple_2d_textures(self):
        if drv.get_driver_version() // 1000 >= 12:
            pytest.skip("texture references were removed in CUDA 12")

        mod = SourceModule(
            """
        texture<float, 2, cudaReadModeElementType> mtx_tex;
        texture<float, 2, cudaReadModeElementType> mtx2_tex;

        __global__ void copy_texture(float *dest)
        {
          int row = threadIdx.x;
          int col = threadIdx.y;
          int w = blockDim.y;
          dest[row*w+col] =
              tex2D(mtx_tex, row, col)
              +
              tex2D(mtx2_tex, row, col);
        }
        """
        )

        copy_texture = mod.get_function("copy_texture")
        mtx_tex = mod.get_texref("mtx_tex")
        mtx2_tex = mod.get_texref("mtx2_tex")

        shape = (3, 4)
        a = np.random.randn(*shape).astype(np.float32)
        b = np.random.randn(*shape).astype(np.float32)
        drv.matrix_to_texref(a, mtx_tex, order="F")
        drv.matrix_to_texref(b, mtx2_tex, order="F")

        dest = np.zeros(shape, dtype=np.float32)
        copy_texture(drv.Out(dest), block=(*shape, 1), texrefs=[mtx_tex, mtx2_tex])
        assert la.norm(dest - a - b) < 1e-6

    @mark_cuda_test
    def test_multichannel_2d_texture(self):
        if drv.get_driver_version() // 1000 >= 12:
            pytest.skip("texture references were removed in CUDA 12")

        mod = SourceModule(
            """
        #define CHANNELS 4
        texture<float4, 2, cudaReadModeElementType> mtx_tex;

        __global__ void copy_texture(float *dest)
        {
          int row = threadIdx.x;
          int col = threadIdx.y;
          int w = blockDim.y;
          float4 texval = tex2D(mtx_tex, row, col);
          dest[(row*w+col)*CHANNELS + 0] = texval.x;
          dest[(row*w+col)*CHANNELS + 1] = texval.y;
          dest[(row*w+col)*CHANNELS + 2] = texval.z;
          dest[(row*w+col)*CHANNELS + 3] = texval.w;
        }
        """
        )

        copy_texture = mod.get_function("copy_texture")
        mtx_tex = mod.get_texref("mtx_tex")

        shape = (5, 6)
        channels = 4
        a = np.asarray(
            np.random.randn(*((channels, *shape))), dtype=np.float32, order="F"
        )
        drv.bind_array_to_texref(drv.make_multichannel_2d_array(a, order="F"), mtx_tex)

        dest = np.zeros((*shape, channels), dtype=np.float32)
        copy_texture(drv.Out(dest), block=(*shape, 1), texrefs=[mtx_tex])
        reshaped_a = a.transpose(1, 2, 0)
        # print reshaped_a
        # print dest
        assert la.norm(dest - reshaped_a) == 0

    @mark_cuda_test
    def test_multichannel_linear_texture(self):
        if drv.get_driver_version() // 1000 >= 12:
            pytest.skip("texture references were removed in CUDA 12")

        mod = SourceModule(
            """
        #define CHANNELS 4
        texture<float4, 1, cudaReadModeElementType> mtx_tex;

        __global__ void copy_texture(float *dest)
        {
          int i = threadIdx.x+blockDim.x*threadIdx.y;
          float4 texval = tex1Dfetch(mtx_tex, i);
          dest[i*CHANNELS + 0] = texval.x;
          dest[i*CHANNELS + 1] = texval.y;
          dest[i*CHANNELS + 2] = texval.z;
          dest[i*CHANNELS + 3] = texval.w;
        }
        """
        )

        copy_texture = mod.get_function("copy_texture")
        mtx_tex = mod.get_texref("mtx_tex")

        shape = (16, 16)
        channels = 4
        a = np.random.randn(*((*shape, channels))).astype(np.float32)
        a_gpu = drv.to_device(a)
        mtx_tex.set_address(a_gpu, a.nbytes)
        mtx_tex.set_format(drv.array_format.FLOAT, 4)

        dest = np.zeros((*shape, channels), dtype=np.float32)
        copy_texture(drv.Out(dest), block=(*shape, 1), texrefs=[mtx_tex])
        # print a
        # print dest
        assert la.norm(dest - a) == 0

    @mark_cuda_test
    def test_2d_fp_textures(self):
        if drv.get_driver_version() // 1000 >= 12:
            pytest.skip("texture references were removed in CUDA 12")

        orden = "F"
        npoints = 32

        for prec in [np.int16, np.float32, np.float64, np.complex64, np.complex128]:
            prec_str = dtype_to_ctype(prec)
            if prec == np.complex64:
                fpName_str = "fp_tex_cfloat"
            elif prec == np.complex128:
                fpName_str = "fp_tex_cdouble"
            elif prec == np.float64:
                fpName_str = "fp_tex_double"
            else:
                fpName_str = prec_str
            A_cpu = np.zeros([npoints, npoints], order=orden, dtype=prec)
            A_cpu[:] = np.random.rand(npoints, npoints)[:]
            A_gpu = gpuarray.zeros(A_cpu.shape, dtype=prec, order=orden)

            myKern = """
            #include <pycuda-helpers.hpp>
            texture<fpName, 2, cudaReadModeElementType> mtx_tex;

            __global__ void copy_texture(cuPres *dest)
            {
              int row = blockIdx.x*blockDim.x + threadIdx.x;
              int col = blockIdx.y*blockDim.y + threadIdx.y;

              dest[row + col*blockDim.x*gridDim.x] = fp_tex2D(mtx_tex, col, row);
            }
            """
            myKern = myKern.replace("fpName", fpName_str)
            myKern = myKern.replace("cuPres", prec_str)
            mod = SourceModule(myKern)

            copy_texture = mod.get_function("copy_texture")
            mtx_tex = mod.get_texref("mtx_tex")
            cuBlock = (16, 16, 1)
            if cuBlock[0] > npoints:
                cuBlock = (npoints, npoints, 1)
            cuGrid = (
                npoints // cuBlock[0] + 1 * (npoints % cuBlock[0] != 0),
                npoints // cuBlock[1] + 1 * (npoints % cuBlock[1] != 0),
                1,
            )
            copy_texture.prepare("P", texrefs=[mtx_tex])
            cudaArray = drv.np_to_array(A_cpu, orden, allowSurfaceBind=False)
            mtx_tex.set_array(cudaArray)
            copy_texture.prepared_call(cuGrid, cuBlock, A_gpu.gpudata)
            assert np.sum(np.abs(A_gpu.get() - np.transpose(A_cpu))) == np.array(
                0, dtype=prec
            )
            A_gpu.gpudata.free()

    @mark_cuda_test
    def test_2d_fp_textures_layered(self):
        if drv.get_driver_version() // 1000 >= 12:
            pytest.skip("texture references were removed in CUDA 12")

        orden = "F"
        npoints = 32

        for prec in [np.int16, np.float32, np.float64, np.complex64, np.complex128]:
            prec_str = dtype_to_ctype(prec)
            if prec == np.complex64:
                fpName_str = "fp_tex_cfloat"
            elif prec == np.complex128:
                fpName_str = "fp_tex_cdouble"
            elif prec == np.float64:
                fpName_str = "fp_tex_double"
            else:
                fpName_str = prec_str
            A_cpu = np.zeros([npoints, npoints], order=orden, dtype=prec)
            A_cpu[:] = np.random.rand(npoints, npoints)[:]
            A_gpu = gpuarray.zeros(A_cpu.shape, dtype=prec, order=orden)

            myKern = """
            #include <pycuda-helpers.hpp>
            texture<fpName, cudaTextureType2DLayered, cudaReadModeElementType> mtx_tex;

            __global__ void copy_texture(cuPres *dest)
            {
              int row = blockIdx.x*blockDim.x + threadIdx.x;
              int col = blockIdx.y*blockDim.y + threadIdx.y;

              dest[row + col*blockDim.x*gridDim.x] = fp_tex2DLayered(mtx_tex, col, row, 1);
            }
            """
            myKern = myKern.replace("fpName", fpName_str)
            myKern = myKern.replace("cuPres", prec_str)
            mod = SourceModule(myKern)

            copy_texture = mod.get_function("copy_texture")
            mtx_tex = mod.get_texref("mtx_tex")
            cuBlock = (16, 16, 1)
            if cuBlock[0] > npoints:
                cuBlock = (npoints, npoints, 1)
            cuGrid = (
                npoints // cuBlock[0] + 1 * (npoints % cuBlock[0] != 0),
                npoints // cuBlock[1] + 1 * (npoints % cuBlock[1] != 0),
                1,
            )
            copy_texture.prepare("P", texrefs=[mtx_tex])
            cudaArray = drv.np_to_array(A_cpu, orden, allowSurfaceBind=True)
            mtx_tex.set_array(cudaArray)
            copy_texture.prepared_call(cuGrid, cuBlock, A_gpu.gpudata)
            assert np.sum(np.abs(A_gpu.get() - np.transpose(A_cpu))) == np.array(
                0, dtype=prec
            )
            A_gpu.gpudata.free()

    @mark_cuda_test
    def test_3d_fp_textures(self):
        if drv.get_driver_version() // 1000 >= 12:
            pytest.skip("texture references were removed in CUDA 12")

        orden = "C"
        npoints = 32

        for prec in [np.int16, np.float32, np.float64, np.complex64, np.complex128]:
            prec_str = dtype_to_ctype(prec)
            if prec == np.complex64:
                fpName_str = "fp_tex_cfloat"
            elif prec == np.complex128:
                fpName_str = "fp_tex_cdouble"
            elif prec == np.float64:
                fpName_str = "fp_tex_double"
            else:
                fpName_str = prec_str
            A_cpu = np.zeros([npoints, npoints, npoints], order=orden, dtype=prec)
            A_cpu[:] = np.random.rand(npoints, npoints, npoints)[:]
            A_gpu = gpuarray.zeros(A_cpu.shape, dtype=prec, order=orden)

            myKern = """
            #include <pycuda-helpers.hpp>
            texture<fpName, 3, cudaReadModeElementType> mtx_tex;

            __global__ void copy_texture(cuPres *dest)
            {
              int row   = blockIdx.x*blockDim.x + threadIdx.x;
              int col   = blockIdx.y*blockDim.y + threadIdx.y;
              int slice = blockIdx.z*blockDim.z + threadIdx.z;
              dest[row + col*blockDim.x*gridDim.x
                + slice*blockDim.x*gridDim.x*blockDim.y*gridDim.y] = fp_tex3D(
                    mtx_tex, slice, col, row);
            }
            """
            myKern = myKern.replace("fpName", fpName_str)
            myKern = myKern.replace("cuPres", prec_str)
            mod = SourceModule(myKern)

            copy_texture = mod.get_function("copy_texture")
            mtx_tex = mod.get_texref("mtx_tex")
            cuBlock = (8, 8, 8)
            if cuBlock[0] > npoints:
                cuBlock = (npoints, npoints, npoints)
            cuGrid = (
                npoints // cuBlock[0] + 1 * (npoints % cuBlock[0] != 0),
                npoints // cuBlock[1] + 1 * (npoints % cuBlock[1] != 0),
                npoints // cuBlock[2] + 1 * (npoints % cuBlock[1] != 0),
            )
            copy_texture.prepare("P", texrefs=[mtx_tex])
            cudaArray = drv.np_to_array(A_cpu, orden, allowSurfaceBind=False)
            mtx_tex.set_array(cudaArray)
            copy_texture.prepared_call(cuGrid, cuBlock, A_gpu.gpudata)
            assert np.sum(np.abs(A_gpu.get() - np.transpose(A_cpu))) == np.array(
                0, dtype=prec
            )
            A_gpu.gpudata.free()

    @mark_cuda_test
    def test_3d_fp_surfaces(self):
        if drv.get_driver_version() // 1000 >= 12:
            pytest.skip("surface references were removed in CUDA 12")

        orden = "C"
        npoints = 32

        for prec in [np.int16, np.float32, np.float64, np.complex64, np.complex128]:
            prec_str = dtype_to_ctype(prec)
            if prec == np.complex64:
                fpName_str = "fp_tex_cfloat"
                A_cpu = np.zeros([npoints, npoints, npoints], order=orden, dtype=prec)
                A_cpu[:].real = np.random.rand(npoints, npoints, npoints)[:]
                A_cpu[:].imag = np.random.rand(npoints, npoints, npoints)[:]
            elif prec == np.complex128:
                fpName_str = "fp_tex_cdouble"
                A_cpu = np.zeros([npoints, npoints, npoints], order=orden, dtype=prec)
                A_cpu[:].real = np.random.rand(npoints, npoints, npoints)[:]
                A_cpu[:].imag = np.random.rand(npoints, npoints, npoints)[:]
            elif prec == np.float64:
                fpName_str = "fp_tex_double"
                A_cpu = np.zeros([npoints, npoints, npoints], order=orden, dtype=prec)
                A_cpu[:] = np.random.rand(npoints, npoints, npoints)[:]
            else:
                fpName_str = prec_str
                A_cpu = np.zeros([npoints, npoints, npoints], order=orden, dtype=prec)
                A_cpu[:] = np.random.rand(npoints, npoints, npoints)[:] * 100.0

            A_gpu = gpuarray.to_gpu(A_cpu)  # Array randomized

            myKernRW = """
            #include <pycuda-helpers.hpp>

            surface<void, cudaSurfaceType3D> mtx_tex;

            __global__ void copy_texture(cuPres *dest, int rw)
            {
              int row   = blockIdx.x*blockDim.x + threadIdx.x;
              int col   = blockIdx.y*blockDim.y + threadIdx.y;
              int slice = blockIdx.z*blockDim.z + threadIdx.z;
              int tid = row + col*blockDim.x*gridDim.x
                  + slice*blockDim.x*gridDim.x*blockDim.y*gridDim.y;
              if (rw==0){
              cuPres aux = dest[tid];
              fp_surf3Dwrite(aux, mtx_tex, row, col, slice,cudaBoundaryModeClamp);}
              else {
              cuPres aux = 0;
              fp_surf3Dread(&aux, mtx_tex, slice, col, row, cudaBoundaryModeClamp);
              dest[tid] = aux;
              }
            }
            """
            myKernRW = myKernRW.replace("fpName", fpName_str)
            myKernRW = myKernRW.replace("cuPres", prec_str)
            modW = SourceModule(myKernRW)

            copy_texture = modW.get_function("copy_texture")
            mtx_tex = modW.get_surfref("mtx_tex")
            cuBlock = (8, 8, 8)
            if cuBlock[0] > npoints:
                cuBlock = (npoints, npoints, npoints)
            cuGrid = (
                npoints // cuBlock[0] + 1 * (npoints % cuBlock[0] != 0),
                npoints // cuBlock[1] + 1 * (npoints % cuBlock[1] != 0),
                npoints // cuBlock[2] + 1 * (npoints % cuBlock[1] != 0),
            )
            copy_texture.prepare("Pi")  # ,texrefs=[mtx_tex])
            A_gpu2 = gpuarray.zeros_like(A_gpu)  # To initialize surface with zeros
            cudaArray = drv.gpuarray_to_array(A_gpu2, orden, allowSurfaceBind=True)
            A_cpu = A_gpu.get()  # To remember original array
            mtx_tex.set_array(cudaArray)
            copy_texture.prepared_call(
                cuGrid, cuBlock, A_gpu.gpudata, np.int32(0)
            )  # Write random array
            copy_texture.prepared_call(
                cuGrid, cuBlock, A_gpu.gpudata, np.int32(1)
            )  # Read, but transposed
            assert np.sum(np.abs(A_gpu.get() - np.transpose(A_cpu))) == np.array(
                0, dtype=prec
            )
            A_gpu.gpudata.free()

    @mark_cuda_test
    def test_2d_fp_surfaces(self):
        if drv.get_driver_version() // 1000 >= 12:
            pytest.skip("surface references were removed in CUDA 12")

        orden = "C"
        npoints = 32

        for prec in [np.int16, np.float32, np.float64, np.complex64, np.complex128]:
            prec_str = dtype_to_ctype(prec)
            if prec == np.complex64:
                fpName_str = "fp_tex_cfloat"
            elif prec == np.complex128:
                fpName_str = "fp_tex_cdouble"
            elif prec == np.float64:
                fpName_str = "fp_tex_double"
            else:
                fpName_str = prec_str
            A_cpu = np.zeros([npoints, npoints], order=orden, dtype=prec)
            A_cpu[:] = np.random.rand(npoints, npoints)[:]
            A_gpu = gpuarray.to_gpu(A_cpu)  # Array randomized

            myKernRW = """
            #include <pycuda-helpers.hpp>

            surface<void, cudaSurfaceType2DLayered> mtx_tex;

            __global__ void copy_texture(cuPres *dest, int rw)
            {
              int row   = blockIdx.x*blockDim.x + threadIdx.x;
              int col   = blockIdx.y*blockDim.y + threadIdx.y;
              int layer = 1;
              int tid = row + col*blockDim.x*gridDim.x ;
              if (rw==0){
              cuPres aux = dest[tid];
              fp_surf2DLayeredwrite(aux, mtx_tex, row, col, layer,cudaBoundaryModeClamp);}
              else {
              cuPres aux = 0;
              fp_surf2DLayeredread(&aux, mtx_tex, col, row, layer, cudaBoundaryModeClamp);
              dest[tid] = aux;
              }
            }
            """
            myKernRW = myKernRW.replace("fpName", fpName_str)
            myKernRW = myKernRW.replace("cuPres", prec_str)
            modW = SourceModule(myKernRW)

            copy_texture = modW.get_function("copy_texture")
            mtx_tex = modW.get_surfref("mtx_tex")
            cuBlock = (8, 8, 1)
            if cuBlock[0] > npoints:
                cuBlock = (npoints, npoints, 1)
            cuGrid = (
                npoints // cuBlock[0] + 1 * (npoints % cuBlock[0] != 0),
                npoints // cuBlock[1] + 1 * (npoints % cuBlock[1] != 0),
                1,
            )
            copy_texture.prepare("Pi")  # ,texrefs=[mtx_tex])
            A_gpu2 = gpuarray.zeros_like(A_gpu)  # To initialize surface with zeros
            cudaArray = drv.gpuarray_to_array(A_gpu2, orden, allowSurfaceBind=True)
            A_cpu = A_gpu.get()  # To remember original array
            mtx_tex.set_array(cudaArray)
            copy_texture.prepared_call(
                cuGrid, cuBlock, A_gpu.gpudata, np.int32(0)
            )  # Write random array
            copy_texture.prepared_call(
                cuGrid, cuBlock, A_gpu.gpudata, np.int32(1)
            )  # Read, but transposed
            assert np.sum(np.abs(A_gpu.get() - np.transpose(A_cpu))) == np.array(
                0, dtype=prec
            )
            A_gpu.gpudata.free()

    @mark_cuda_test
    def test_large_smem(self):
        n = 4000
        mod = SourceModule(
            """
        #include <stdio.h>

        __global__ void kernel(int *d_data)
        {
        __shared__ int sdata[%d];
        sdata[threadIdx.x] = threadIdx.x;
        d_data[threadIdx.x] = sdata[threadIdx.x];
        }
        """
            % n
        )

        kernel = mod.get_function("kernel")

        import pycuda.gpuarray as gpuarray

        arg = gpuarray.zeros((n,), dtype=np.float32)

        kernel(
            arg,
            block=(
                1,
                1,
                1,
            ),
        )

    @mark_cuda_test
    def test_bitlog(self):
        from pycuda.tools import bitlog2

        assert bitlog2(17) == 4
        assert bitlog2(0xAFFE) == 15
        assert bitlog2(0x3AFFE) == 17
        assert bitlog2(0xCC3AFFE) == 27

    @mark_cuda_test
    def test_mempool_2(self):
        from random import randrange

        from pycuda.tools import DeviceMemoryPool
        pool = DeviceMemoryPool()

        for _i in range(2000):
            s = randrange(1 << 31) >> randrange(32)
            bin_nr = pool.bin_number(s)
            asize = pool.alloc_size(bin_nr)

            assert asize >= s, s
            assert pool.bin_number(asize) == bin_nr, s
            assert asize < asize * (1 + 1 / 8)

    @mark_cuda_test
    def test_mempool(self):
        from pycuda.tools import DeviceMemoryPool, bitlog2

        pool = DeviceMemoryPool()
        queue = []
        free, _total = drv.mem_get_info()

        e0 = bitlog2(free)

        for e in range(e0 - 6, e0 - 4):
            for _i in range(100):
                queue.append(pool.allocate(1 << e))
                if len(queue) > 10:
                    queue.pop(0)
        del queue
        pool.stop_holding()

    @mark_cuda_test
    def test_multi_context(self):
        if drv.get_version() < (2, 0, 0):
            return
        if drv.get_version() >= (2, 2, 0) and drv.get_version() < (8,):
            if drv.Context.get_device().compute_mode == drv.compute_mode.EXCLUSIVE:
                return

        mem_a = drv.mem_alloc(50)
        ctx2 = drv.Context.get_device().make_context()
        mem_b = drv.mem_alloc(60)

        del mem_a
        del mem_b
        ctx2.detach()

    @mark_cuda_test
    def test_additional_primary_context(self):
        if drv.get_version() < (2, 0, 0):
            return
        if drv.get_version() >= (2, 2, 0) and drv.get_version() < (8,):
            if drv.Context.get_device().compute_mode == drv.compute_mode.EXCLUSIVE:
                return

        mem_a = drv.mem_alloc(50)
        pctx = drv.Context.get_device().retain_primary_context()
        pctx.push()
        mem_b = drv.mem_alloc(60)

        del mem_a
        del mem_b
        pctx.detach()

    @mark_cuda_test
    def test_3d_texture(self):
        if drv.get_driver_version() // 1000 >= 12:
            pytest.skip("texture references were removed in CUDA 12")

        # adapted from code by Nicolas Pinto
        w = 2
        h = 4
        d = 8
        shape = (w, h, d)

        a = np.asarray(np.random.randn(*shape), dtype=np.float32, order="F")

        descr = drv.ArrayDescriptor3D()
        descr.width = w
        descr.height = h
        descr.depth = d
        descr.format = drv.dtype_to_array_format(a.dtype)
        descr.num_channels = 1
        descr.flags = 0

        ary = drv.Array(descr)

        copy = drv.Memcpy3D()
        copy.set_src_host(a)
        copy.set_dst_array(ary)
        copy.width_in_bytes = copy.src_pitch = a.strides[1]
        copy.src_height = copy.height = h
        copy.depth = d

        copy()

        mod = SourceModule(
            """
        texture<float, 3, cudaReadModeElementType> mtx_tex;

        __global__ void copy_texture(float *dest)
        {
          int x = threadIdx.x;
          int y = threadIdx.y;
          int z = threadIdx.z;
          int dx = blockDim.x;
          int dy = blockDim.y;
          int i = (z*dy + y)*dx + x;
          dest[i] = tex3D(mtx_tex, x, y, z);
          //dest[i] = x;
        }
        """
        )

        copy_texture = mod.get_function("copy_texture")
        mtx_tex = mod.get_texref("mtx_tex")

        mtx_tex.set_array(ary)

        dest = np.zeros(shape, dtype=np.float32, order="F")
        copy_texture(drv.Out(dest), block=shape, texrefs=[mtx_tex])
        assert la.norm(dest - a) == 0

    @mark_cuda_test
    def test_prepared_invocation(self):
        a = np.random.randn(4, 4).astype(np.float32)
        a_gpu = drv.mem_alloc(a.size * a.dtype.itemsize)

        drv.memcpy_htod(a_gpu, a)

        mod = SourceModule(
            """
            __global__ void doublify(float *a)
            {
              int idx = threadIdx.x + threadIdx.y*blockDim.x;
              a[idx] *= 2;
            }
            """
        )

        func = mod.get_function("doublify")
        func.prepare("P")
        func.prepared_call((1, 1), (4, 4, 1), a_gpu, shared_size=20)
        a_doubled = np.empty_like(a)
        drv.memcpy_dtoh(a_doubled, a_gpu)
        print(a)
        print(a_doubled)
        assert la.norm(a_doubled - 2 * a) == 0

        # now with offsets
        func.prepare("P")
        a_quadrupled = np.empty_like(a)
        func.prepared_call((1, 1), (15, 1, 1), int(a_gpu) + a.dtype.itemsize)
        drv.memcpy_dtoh(a_quadrupled, a_gpu)
        assert la.norm(a_quadrupled[1:] - 4 * a[1:]) == 0

    @mark_cuda_test
    def test_prepared_with_vector(self):
        cuda_source = r"""
        __global__ void cuda_function(float3 input)
        {
        float3 result = make_float3(input.x, input.y, input.z);
        }
        """

        mod = SourceModule(cuda_source, cache_dir=False, keep=False)

        kernel = mod.get_function("cuda_function")
        arg_types = [gpuarray.vec.float3]

        kernel.prepare(arg_types)
        kernel.prepared_call(
            (1, 1, 1), (1, 1, 1), gpuarray.vec.make_float3(0.0, 1.0, 2.0)
        )

    @mark_cuda_test
    def test_fp_textures(self):
        if drv.get_driver_version() // 1000 >= 12:
            pytest.skip("texture references were removed in CUDA 12")

        if drv.Context.get_device().compute_capability() < (1, 3):
            return

        for tp in [np.float32, np.float64]:
            tp_cstr = dtype_to_ctype(tp)
            mod = SourceModule(
                """
            #include <pycuda-helpers.hpp>

            texture<fp_tex_%(tp)s, 1, cudaReadModeElementType> my_tex;

            __global__ void copy_texture(%(tp)s *dest)
            {
              int i = threadIdx.x;
              dest[i] = fp_tex1Dfetch(my_tex, i);
            }
            """
                % {"tp": tp_cstr}
            )

            copy_texture = mod.get_function("copy_texture")
            my_tex = mod.get_texref("my_tex")

            shape = (384,)
            a = np.random.randn(*shape).astype(tp)
            a_gpu = gpuarray.to_gpu(a)
            a_gpu.bind_to_texref_ext(my_tex, allow_double_hack=True)

            dest = np.zeros(shape, dtype=tp)
            copy_texture(
                drv.Out(dest),
                block=(*shape, 1, 1),
                texrefs=[my_tex],
            )

            assert la.norm(dest - a) == 0

    @mark_cuda_test
    def test_constant_memory(self):
        # contributed by Andrew Wagner

        module = SourceModule(
            """
        __constant__ float const_array[32];

        __global__ void copy_constant_into_global(float* global_result_array)
        {
            global_result_array[threadIdx.x] = const_array[threadIdx.x];
        }
        """
        )

        copy_constant_into_global = module.get_function("copy_constant_into_global")
        const_array, _ = module.get_global("const_array")

        host_array = np.random.randint(0, 255, (32,)).astype(np.float32)

        global_result_array = drv.mem_alloc_like(host_array)
        drv.memcpy_htod(const_array, host_array)

        copy_constant_into_global(global_result_array, grid=(1, 1), block=(32, 1, 1))

        host_result_array = np.zeros_like(host_array)
        drv.memcpy_dtoh(host_result_array, global_result_array)

        assert (host_result_array == host_array).all

    @mark_cuda_test
    def test_register_host_memory(self):
        pytest.xfail("known issue: must fix array creation")

        if drv.get_version() < (4,):
            from py.test import skip

            skip("register_host_memory only exists on CUDA 4.0 and later")

        import sys

        if sys.platform == "darwin":
            from py.test import skip

            skip("register_host_memory is not supported on OS X")

        a = drv.aligned_empty((2 ** 20,), np.float64)
        a_pin = drv.register_host_memory(a)

        gpu_ary = drv.mem_alloc_like(a)
        stream = drv.Stream()
        drv.memcpy_htod_async(gpu_ary, a_pin, stream)
        drv.Context.synchronize()

    @mark_cuda_test
    # https://github.com/inducer/pycuda/issues/45
    def test_recursive_launch(self):
        # Test contributed by Aditya Avinash Atluri

        if drv.Context.get_device().compute_capability() < (3, 5):
            from pytest import skip

            skip("need compute capability 3.5 or higher for dynamic parallelism")

        cuda_string = """
        __device__ void  saxpy(double* s, float a, long* p, int b, long* q)
        {
            int tx = threadIdx.x;
            s[tx] = a*p[tx]+b*q[tx];
        }

        __global__ void  sub(long* p, long* q, long* d)
        {
            int tx = threadIdx.x;
            p[tx] = 2*p[tx];
            d[tx] = p[tx]-q[tx];
        }

        __device__ long add(long p, long q)
        {
            p = p+1;
            return p+q;
        }

        __global__ void math(long* a, long* b, long* c, long* d, long* e, double* f)
        {
            int tx = threadIdx.x;
            __shared__ long x[100];
            x[tx] = a[tx + 0];
            __shared__ long y[100];
            y[tx] = b[tx + 0];
            c[tx]=add(x[tx],y[tx]);
            dim3 dimGrid_sub(1,1,1);
            dim3 dimBlock_sub(100,1,1);
            sub<<<dimGrid_sub,dimBlock_sub>>>(a,b,d);
            saxpy(f,1.0345,x,-2,y);
        }
        """

        def math(a, b, c, d, e, f):
            a_gpu = drv.mem_alloc(a.nbytes)
            b_gpu = drv.mem_alloc(b.nbytes)
            c_gpu = drv.mem_alloc(c.nbytes)
            d_gpu = drv.mem_alloc(d.nbytes)
            e_gpu = drv.mem_alloc(e.nbytes)
            f_gpu = drv.mem_alloc(f.nbytes)

            drv.memcpy_htod(a_gpu, a)
            drv.memcpy_htod(b_gpu, b)

            from pycuda.compiler import DynamicSourceModule

            mod = DynamicSourceModule(cuda_string, keep=True)

            func = mod.get_function("math")
            func(
                a_gpu,
                b_gpu,
                c_gpu,
                d_gpu,
                e_gpu,
                f_gpu,
                block=(100, 1, 1),
                grid=(1, 1, 1),
            )

            drv.memcpy_dtoh(c, c_gpu)
            drv.memcpy_dtoh(d, d_gpu)
            drv.memcpy_dtoh(e, e_gpu)
            drv.memcpy_dtoh(f, f_gpu)

            # print(c,d,e,f)

        a = np.random.randint(10, size=100)
        b = np.random.randint(10, size=100)
        c = np.empty_like(a)
        d = np.empty_like(a)
        e = np.empty_like(a)
        f = np.array(a, dtype="d")

        math(a, b, c, d, e, f)

    @mark_cuda_test
    def test_jit_link_module(self):
        from pycuda.compiler import DEFAULT_NVCC_FLAGS

        if drv.Context.get_device().compute_capability() < (3, 5):
            from pytest import skip

            skip("need compute capability 3.5 or higher for dynamic parallelism")

        test_outer_cu = """#include <cstdio>
        __global__ void test_kernel() {
            extern __global__ void test_kernel_inner();
            printf("Hello outer world!\\n");
            test_kernel_inner<<<2, 1>>>();
        }"""

        test_inner_cu = """#include <cstdio>
        __global__ void test_kernel_inner() {
            printf("  Hello inner world!\\n");
        }"""

        from pycuda.compiler import DynamicModule

        mod = DynamicModule()
        mod.add_source(
            test_outer_cu,
            nvcc_options=(["-rdc=true", "-lcudadevrt", *DEFAULT_NVCC_FLAGS]),
        )
        mod.add_source(
            test_inner_cu,
            nvcc_options=(["-rdc=true", "-lcudadevrt", *DEFAULT_NVCC_FLAGS]),
        )
        mod.add_stdlib("cudadevrt")
        mod.link()

        test_kernel = mod.get_function("test_kernel")
        test_kernel(grid=(2, 1), block=(1, 1, 1))


class WrappedAllocation(drv.PointerHolderBase):
    def __init__(self, wrapped):
        self.wrapped = wrapped
        super().__init__()

    def get_pointer(self):
        return int(self.wrapped)


@mark_cuda_test
def test_pointer_holder_base():
    alloc = WrappedAllocation(drv.mem_alloc(1024))
    ary = gpuarray.GPUArray((1024,), np.uint8, gpudata=alloc)
    print(ary.get())


# A class to emulate an object from outside PyCUDA that implements the CUDA
# Array Interface
class CudaArrayInterfaceImpl:
    def __init__(self, size, itemsize, dtype):
        self._shape = (size,)
        self._strides = (itemsize,)
        self._typestr = dtype.str
        self._ptr = drv.mem_alloc(size * itemsize)

    @property
    def __cuda_array_interface__(self):
        return {
            "shape": self._shape,
            "strides": self._strides,
            "typestr": self._typestr,
            "data": (int(self._ptr), False),
            "stream": None,
            "version": 3
        }

    @property
    def ptr(self):
        return self._ptr


@mark_cuda_test
def test_pass_cai_array():
    dtype = np.int32
    size = 1024
    np_array = np.arange(size, dtype=dtype)
    cai_array = CudaArrayInterfaceImpl(size, np_array.itemsize, np_array.dtype)

    mod = SourceModule(
        """
    __global__ void gpu_arange(int *x)
    {
      const int i = threadIdx.x;
      x[i] = i;
    }
    """
    )

    gpu_arange = mod.get_function("gpu_arange")
    gpu_arange(cai_array, grid=(1,), block=(size, 1, 1))

    host_array = np.empty_like(np_array)
    drv.memcpy_dtoh(host_array, cai_array.ptr)
    assert (host_array == np_array).all()


def test_import_pyopencl_before_pycuda():
    try:
        import pyopencl  # noqa
    except ImportError:
        return
    import pycuda.driver  # noqa


if __name__ == "__main__":
    # make sure that import failures get reported, instead of skipping the tests.
    import pycuda.autoinit  # noqa

    import sys

    if len(sys.argv) > 1:
        exec(sys.argv[1])
    else:
        from pytest import main

        main([__file__])
