Expose assume_divby

Signed-off-by: Qiqi Xiao <qiqix@nvidia.com>

Author: xiaoqiqi177

changelog.d/assume-divby.md

@@ -0,0 +1,4 @@
+<!--- SPDX-FileCopyrightText: Copyright (c) <2026> NVIDIA CORPORATION & AFFILIATES. All rights reserved. -->
+<!--- SPDX-License-Identifier: Apache-2.0 -->
+
+- Added `ct.assume_divisible_by(x, divisor)`, a compiler hint that declares an integer scalar to be divisible by a constant. The divisibility fact is propagated through arithmetic, e.g., allowing the compiler to prove alignment for derived indices and pointer offsets to emit wider memory operations than it could with unknown divisibility.

docs/source/operations.rst

@@ -214,6 +214,7 @@ Utility
    printf
    print
    assert_
+   assume_divisible_by
 
 
 .. _operations-metaprogramming:

docs/source/performance.rst

@@ -10,7 +10,8 @@ Performance Tuning
 Several performance tuning techniques are available in cuTile:
 
 * architecture-specific configuration values, using :class:`ByTarget`;
-* load/store hints such as ``latency`` and ``allow_tma``.
+* load/store hints such as ``latency`` and ``allow_tma``;
+* divisibility hints via :func:`assume_divisible_by`.
 
 
 Architecture-specific configuration
@@ -53,6 +54,43 @@ Example
     :end-before: example-end
 
 
+.. _divisibility-hints:
+
+Divisibility hints
+------------------
+
+:func:`assume_divisible_by` is a compiler hint that declares an integer
+scalar to be divisible by a compile-time constant. No check is performed at
+runtime:
+
+.. code-block:: python
+
+    n = ct.assume_divisible_by(n, 16)
+
+The compiler propagates the divisibility metadata through arithmetic operations — so
+that derived indices and pointer offsets inherit the same fact. This matters
+most when a runtime scalar is used to compute a dynamic array slice:
+
+.. code-block:: python
+
+    @ct.kernel
+    def kernel(x, dim_offset: int, dim_size: int):
+        dim_offset = ct.assume_divisible_by(dim_offset, 16)
+        dim_size   = ct.assume_divisible_by(dim_size,   16)
+        start = ct.bid(0) * dim_offset
+        sub_x = x.slice(axis=0, start=start, stop=start + dim_size)
+        tile  = ct.load(sub_x, index=(0,), shape=(128,))
+        ct.store(sub_x, index=(0,), tile=tile)
+
+Without the hints, the compiler treats ``dim_offset`` and ``dim_size`` as
+fully unknown and cannot prove alignment for the derived view. With the
+hints, it can infer alignment all the way into the view's base address and
+shape, enabling wider memory operations.
+
+The hint is a programmer declaration, not an enforcement. Behavior is undefined
+if ``x`` is not actually divisible by ``divisor`` at runtime.
+
+
 .. _autotuning:
 
 Autotuning

src/cuda/tile/__init__.py

@@ -74,6 +74,7 @@
     argmax,
     argmin,
     assert_,
+    assume_divisible_by,
     astile,
     astype,
     atan2,
@@ -236,6 +237,7 @@
     "argmax",
     "argmin",
     "assert_",
+    "assume_divisible_by",
     "astile",
     "astype",
     "atan2",

src/cuda/tile/_ir/ops.py

@@ -2132,6 +2132,19 @@ def assume_div_by(x: Var, divisor: int | None) -> Var:
     return add_operation(AssumeDivBy, x.get_type(), x=x, divisor=divisor)
 
 
+@impl(ct.assume_divisible_by)
+def assume_divisible_by_impl(x: Var, divisor: Var) -> Var:
+    ty = x.get_type()
+    if not is_0d_tile(ty, is_integral):
+        raise TileTypeError(
+            f"`assume_divisible_by` requires an integer scalar, got {ty}")
+    divisor_val = require_constant_int(divisor)
+    if divisor_val < 1:
+        raise TileTypeError(
+            f"`assume_divisible_by` requires a positive divisor, got {divisor_val}")
+    return assume_div_by(x, divisor_val)
+
+
 @dataclass(eq=False)
 class AssumeBounded(Operation, opcode="assume_bounded"):
     lower_bound: int | None = attribute()

src/cuda/tile/_stub.py

@@ -1183,6 +1183,36 @@ def kernel(x):
     """
 
 
+@stub
+def assume_divisible_by(x: Scalar, divisor: int) -> Scalar:
+    """Declares that ``x`` is divisible by ``divisor``.
+
+    The caller is responsible for the correctness of the claim;
+    behavior is undefined if ``x`` is not actually divisible by
+    ``divisor`` at runtime.
+
+    Args:
+        x: An integer scalar.
+        divisor (const int): The assumed divisor. Must be a positive integer constant.
+
+    Returns:
+        An integer scalar. ``x`` value unchanged.
+
+    Examples:
+
+        .. testcode::
+            :template: kernel_wrapper.py
+
+            n = ct.bid(0) + 128
+            n = ct.assume_divisible_by(n, 128)
+            print(n)
+
+        .. testoutput::
+
+            128
+    """
+
+
 @stub
 def load(array: Array, /,
          index: Shape,

test/test_propagate_divby.py

@@ -8,9 +8,11 @@
 from cuda.tile._ir.ir import Block
 from cuda.tile._compile import compile_tile
 from cuda.tile.compilation import (
-        ParameterConstraint, ArrayConstraint, ListConstraint, ConstantConstraint, KernelSignature
+        ParameterConstraint, ArrayConstraint, ListConstraint,
+        ConstantConstraint, ScalarConstraint, KernelSignature
 )
 from cuda.tile._cext import CallingConvention
+from cuda.tile._exception import TileTypeError
 from typing import Sequence
 
 
@@ -104,6 +106,109 @@ def kernel(x):
     assert get_op_divby(body, MakeTensorView) == [{}]
 
 
+# --- ct.assume_divisible_by ---
+
+def test_assume_divisible_by_emits_op():
+    def kernel(x, n: int):
+        n = ct.assume_divisible_by(n, 16)
+        ct.store(x, (n,), 0)
+
+    body = get_ir(kernel, (
+        array_arg(ndim=1, base_div=1, stride_const=(1,)),
+        ScalarConstraint(ct.int32),
+    ))
+    ops = [op.divisor for op in body.traverse() if isinstance(op, AssumeDivBy)]
+    assert ops == [16]
+
+
+def test_assume_divisible_by_propagates_to_dynamic_slice():
+    def kernel(x, start_factor: int, extent: int):
+        start_factor = ct.assume_divisible_by(start_factor, 32)
+        extent = ct.assume_divisible_by(extent, 32)
+        start = ct.bid(0) * start_factor
+        stop = start + extent
+        sub_x = x.slice(axis=0, start=start, stop=stop)
+        tile = ct.load(sub_x, index=(0,), shape=(1,))
+        ct.store(sub_x, index=(0,), tile=tile)
+
+    body = get_ir(kernel, (
+        array_arg(dtype=ct.bfloat16, base_div=16, stride_const=(1,)),
+        ScalarConstraint(ct.int32),
+        ScalarConstraint(ct.int32),
+    ))
+
+    divby = get_op_divby(body, MakeTensorView)[0]
+    assert divby.get('base_ptr') == 16 and divby.get('shape[0]') == 32
+
+
+def test_assume_divisible_by_divisor_one_is_noop():
+    def kernel(x, n: int):
+        n = ct.assume_divisible_by(n, 1)
+        ct.store(x, (n,), 0)
+
+    body = get_ir(kernel, (
+        array_arg(ndim=1, base_div=1, stride_const=(1,)),
+        ScalarConstraint(ct.int32),
+    ))
+    ops = [op for op in body.traverse() if isinstance(op, AssumeDivBy)]
+    assert ops == []
+
+
+def test_assume_divisible_by_non_power_of_two_divisor():
+    # divisor=12 has largest power-of-2 factor 4.
+    # The propagate_divby pass extracts that power-of-2 when inserting
+    # AssumeDivBy before MakeTensorView, so shape[0] ends up with divisor=4.
+    def kernel(x, extent: int):
+        extent = ct.assume_divisible_by(extent, 12)
+        sub_x = x.slice(axis=0, start=0, stop=extent)
+        tile = ct.load(sub_x, index=(0,), shape=(1,))
+        ct.store(sub_x, index=(0,), tile=tile)
+
+    body = get_ir(kernel, (
+        array_arg(ndim=1, base_div=1, stride_const=(1,)),
+        ScalarConstraint(ct.int32),
+    ))
+    divby = get_op_divby(body, MakeTensorView)[0]
+    assert divby.get('shape[0]') == 4
+
+
+def test_assume_divisible_by_type_error_on_float():
+    def kernel(x, f: float):
+        f = ct.assume_divisible_by(f, 16)
+        ct.store(x, (0,), 0)
+
+    with pytest.raises(TileTypeError, match="integer scalar"):
+        get_ir(kernel, (
+            array_arg(ndim=1, stride_const=(1,)),
+            ScalarConstraint(ct.float32),
+        ))
+
+
+def test_assume_divisible_by_error_on_nonconstant_divisor():
+    def kernel(x, n: int, d: int):
+        n = ct.assume_divisible_by(n, d)
+        ct.store(x, (n,), 0)
+
+    with pytest.raises(TileTypeError, match="integer constant"):
+        get_ir(kernel, (
+            array_arg(ndim=1, stride_const=(1,)),
+            ScalarConstraint(ct.int32),
+            ScalarConstraint(ct.int32),
+        ))
+
+
+def test_assume_divisible_by_error_on_nonpositive_divisor():
+    def kernel(x, n: int):
+        n = ct.assume_divisible_by(n, 0)
+        ct.store(x, (n,), 0)
+
+    with pytest.raises(TileTypeError, match="positive divisor"):
+        get_ir(kernel, (
+            array_arg(ndim=1, stride_const=(1,)),
+            ScalarConstraint(ct.int32),
+        ))
+
+
 # --- Control flow propagation ---
 
 def test_if_else():