MCPcopy Create free account
hub / github.com/NVIDIA/DALI / setup_cupy

Function setup_cupy

dali/test/python/test_dltensor_operator.py:40–70  ·  view source on GitHub ↗
()

Source from the content-addressed store, hash-verified

38
39
40def setup_cupy():
41 global cupy
42 global cupy_stream
43 global square_diff_kernel
44 global mix_channels_kernel
45 global gray_scale_kernel
46 import cupy as cupy
47
48 cupy_stream = cupy.cuda.Stream()
49 square_diff_kernel = cupy.ElementwiseKernel("T x, T y", "T z", "z = x*x - y*y", "square_diff")
50
51 mix_channels_kernel = cupy.ElementwiseKernel(
52 "uint8 x, uint8 y", "uint8 z", "z = (i % 3) ? x : y", "mix_channels"
53 )
54
55 gray_scale_kernel = cupy.RawKernel(
56 r"""
57 extern "C" __global__
58 void gray_scale(float *output, const unsigned char *input, long long height, long long width) {
59 int tidx = blockIdx.x * blockDim.x + threadIdx.x;
60 int tidy = blockIdx.y * blockDim.y + threadIdx.y;
61 if (tidx < width && tidy < height) {
62 float r = input[tidy * width + tidx] / 255.;
63 float g = input[tidy * width + tidx + 1] / 255.;
64 float b = input[tidy * width + tidx + 2] / 255.;
65 output[tidy * width + tidx] = 0.299 * r + 0.59 * g + 0.11 * b;
66 }
67 }
68 """,
69 "gray_scale",
70 )
71
72
73def random_seed():

Callers 2

test_cupyFunction · 0.85

Calls

no outgoing calls

Tested by

no test coverage detected