()
| 253 | @pytest.mark.gpu |
| 254 | @pytest.mark.skipif(not env.has_cuda(), reason="need cuda") |
| 255 | def test_cuda_make_int8(): |
| 256 | def check_cuda(n, value, lanes): |
| 257 | dtype = "int8" |
| 258 | dev = tvm.cuda(0) |
| 259 | const_value = tvm.tirx.const(value, dtype=dtype) |
| 260 | |
| 261 | @I.ir_module(s_tir=True) |
| 262 | class Module: |
| 263 | @T.prim_func(s_tir=True) |
| 264 | def main(A: T.Buffer((n, lanes), dtype)): |
| 265 | T.func_attr({"tirx.noalias": True}) |
| 266 | for i in T.thread_binding(n, thread="blockIdx.x"): |
| 267 | for j in T.vectorized(lanes): |
| 268 | with T.sblock("A"): |
| 269 | v_i, v_j = T.axis.remap("SS", [i, j]) |
| 270 | T.reads() |
| 271 | T.writes(A[v_i, v_j]) |
| 272 | A[v_i, v_j] = const_value |
| 273 | |
| 274 | fun = tvm.compile(Module, target="cuda") |
| 275 | |
| 276 | np_a = np.full((n, lanes), value, dtype=dtype) |
| 277 | a = tvm.runtime.empty(np_a.shape, dtype, dev) |
| 278 | fun(a) |
| 279 | np.testing.assert_equal(a.numpy(), np_a) |
| 280 | |
| 281 | check_cuda(64, np.uint8(0xAB).view(np.int8), 4) |
| 282 | check_cuda(64, 0, 4) |
| 283 | check_cuda(64, -3, 4) |
| 284 | check_cuda(64, np.uint8(0xAB).view(np.int8), 3) |
| 285 | check_cuda(64, 0, 3) |
| 286 | check_cuda(64, -3, 3) |
| 287 | check_cuda(64, np.uint8(0xAB).view(np.int8), 2) |
| 288 | check_cuda(64, 0, 2) |
| 289 | check_cuda(64, -3, 2) |
| 290 | |
| 291 | |
| 292 | @pytest.mark.gpu |
nothing calls this directly
no test coverage detected
searching dependent graphs…