def f(arr): x, y, z = jit.grid(3) gx, gy, gz = jit.gridsize(3) if x == 0 and y == 0 and z == 0: arr[0] = gx arr[1] = gy arr[2] = gz
def f(arr): x = jit.grid(1) if x == 0: shapes = arr.shape strides = arr.strides arr[0, 0, 0] = shapes[0] arr[0, 0, 1] = shapes[1] arr[0, 0, 2] = shapes[2] arr[0, 0, 3] = strides[0] arr[0, 0, 4] = strides[1] arr[0, 0, 5] = strides[2]
def f(): x = jit.grid(n) # n can only be 1, 2, 3 (as int) # NOQA
def f(): x = jit.grid(2) y = cupy.int64(x) # <- x is a tuple # NOQA
def f(): x, = jit.grid(1) # cannot unpack an int
def f(arr1, arr2, k, m, n): x, y, z = jit.grid(3) if x < k and y < m and z < n: arr2[x, y, z] = arr1[x, y, z]
def f(arr1, arr2, n, m): x, y = jit.grid(2) # TODO(leofang): make it possible to write this: # if x < arr1.shape[0] and y < arr1.shape[1]: if x < n and y < m: arr2[x, y] = arr1[x, y]
def f(arr1, arr2): x = jit.grid(1) if x < arr1.size: arr2[x] = arr1[x]
def f(arr): x, y = jit.grid(2) gx, gy = jit.gridsize(2) if x == 0 and y == 0: arr[0] = gx arr[1] = gy
def f(arr): x = jit.grid(1) gx = jit.gridsize(1) if x == 0: arr[x] = gx
def f(arr): x = jit.grid(1) if x == 0: arr[0] = jit.warpsize
def f(arr): x = jit.grid(1) if x < arr.size: arr[x] = jit.laneid()