From 073be1d2ee38353a6e71c535b3460bc160fd3822 Mon Sep 17 00:00:00 2001 From: Philippe Tillet Date: Thu, 14 Apr 2022 19:30:02 -0700 Subject: [PATCH] [FRONTEND] check that tensors have power-of-two number of elements (#499) --- python/triton/language/core.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/python/triton/language/core.py b/python/triton/language/core.py index 75d75e8ea..046e60e0e 100644 --- a/python/triton/language/core.py +++ b/python/triton/language/core.py @@ -380,6 +380,9 @@ class tensor: self.numel = 1 for s in self.shape: 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.type = type # Tensor type (can be block_type) # Following the practice in pytorch, dtype is scalar type