gaspersic/triton
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

[FRONTEND] check that tensors have power-of-two number of elements (#499)

Browse Source
This commit is contained in:
Philippe Tillet
2022-04-14 19:30:02 -07:00
committed by GitHub
parent 5c7122004c
commit 073be1d2ee
1 changed files with 3 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
python/triton/language/core.py
View File

@@ -380,6 +380,9 @@ class tensor:
self.numel = 1 self.numel = 1
for s in self.shape: for s in self.shape:
self.numel *= s self.numel *= s
is_pow2 = (self.numel and (not(self.numel & (self.numel - 1))))
if not is_pow2:
raise ValueError("Triton tensors must have a power-of-two number of elements")
self.numel = constexpr(self.numel) self.numel = constexpr(self.numel)
self.type = type # Tensor type (can be block_type) self.type = type # Tensor type (can be block_type)
# Following the practice in pytorch, dtype is scalar type # Following the practice in pytorch, dtype is scalar type