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()
