[lang] Remove power-of-two constraint on vector types

Some NVVM intrinsics take vectors with non power of 2 size, so I guess
we ought to remove this constraint.
I don't see a big downside--the user might want to parameterize a
kernel over many vector types and just get as many
wide loads/stores as they can, even if it's not perfectly divisible by
the widest possible vector.

Signed-off-by: Asher Mancinelli <amancinelli@nvidia.com>

Author: ashermancinelli

experimental/cuda-lang/src/cuda/lang/_ir/type.py

@@ -31,16 +31,10 @@
 from cuda.lang._exception import TileTypeError
 
 
-def _is_power_of_2(value: int) -> bool:
-    assert isinstance(value, int)
-    return value > 0 and value & (value - 1) == 0
-
-
 def is_vector_ty(ty: Type) -> bool:
     return (
         isinstance(ty, TileTy)
         and len(ty.shape) == 1
-        and _is_power_of_2(ty.shape[0])
     )
 
 
@@ -49,10 +43,6 @@ def make_vector_ty(dtype: DType, length: int) -> TileTy:
         raise TileTypeError(
             f"Expected vector length to be an int, got {type(length).__name__}"
         )
-    if not _is_power_of_2(length):
-        raise TileTypeError(
-            f"Expected vector length to be a positive power of two, got {length}"
-        )
     return TileTy(dtype, (length,))
 
 

experimental/cuda-lang/test/test_vectors.py

@@ -9,7 +9,6 @@
 
 import cuda.lang as cl
 from cuda.lang._datatype import to_torch_dtype
-from cuda.lang._exception import TileTypeError
 from cuda.tile import static_iter
 
 
@@ -295,11 +294,10 @@ def kernel(out):
     torch.testing.assert_close(out.cpu(), expected)
 
 
-def test_pointer_vector_count_must_be_power_of_two():
+def test_pointer_vector_count_can_be_non_power_of_two():
     @cl.kernel
     def kernel(out):
         out.get_base_pointer().load(count=3, alignment=4)
 
     out = torch.zeros(3, dtype=torch.int32).cuda()
-    with pytest.raises(TileTypeError, match="positive power of two"):
-        cl.launch(torch.cuda.current_stream(), (1,), (1,), kernel, (out,))
+    cl.launch(torch.cuda.current_stream(), (1,), (1,), kernel, (out,))