mirror of
https://github.com/ROCm/ROCm.git
synced 2026-04-05 03:01:17 -04:00
a3c39d8fbe
Triton has supported different codegen backends for different devices, so enabling the unit test cases to support different devices also makes sense. Otherwise, the third-party backend might have to intrusively change the Triton test cases.
22 lines
399 B
Python
22 lines
399 B
Python
|
|
from __future__ import annotations
|
|
|
|
import torch
|
|
|
|
import triton
|
|
import triton.language as tl
|
|
|
|
|
|
def test_annotations(device):
|
|
|
|
@triton.jit
|
|
def _kernel(X: torch.Tensor, N: int, BLOCK_SIZE: tl.constexpr):
|
|
pass
|
|
|
|
x = torch.empty(1, device=device)
|
|
_kernel[(1,)](x, x.shape[0], 32)
|
|
try:
|
|
_kernel[(1,)](x.shape[0], x.shape[0], 32)
|
|
except AttributeError:
|
|
pass
|