warp-lang 1.5.1__py3-none-manylinux2014_aarch64.whl → 1.6.0__py3-none-manylinux2014_aarch64.whl

warp/builtins.py

@@ -1707,64 +1707,98 @@ add_builtin(
 
 # ------------------
 # Tile-based primitives
-shared_memory_id = 0
+
+
+def tile_unpack_shape(arg_values):
+    shape = arg_values["shape"]
+
+    if not isinstance(shape, tuple):
+        # promote to tuple
+        shape = (shape,)
+
+    # check that components are constants
+    for d in shape:
+        if d is None:
+            raise ValueError("Tile functions require shape to be a compile time constant.")
+
+    return shape
+
+
+def tile_unpack_offset(arg_values, ndim=0):
+    if "offset" in arg_values:
+        offset = arg_values["offset"]
+    else:
+        offset = (0,) * ndim
+
+    if isinstance(offset, tuple):
+        return offset
+    else:
+        # promote to tuple
+        return (offset,)
 
 
 def tile_zeros_value_func(arg_types: Mapping[str, type], arg_values: Mapping[str, Any]):
     # return generic type (for doc builds)
     if arg_types is None:
-        return Tile(dtype=Any, M=Any, N=Any)
-
-    if "m" not in arg_values:
-        raise RuntimeError("'m' keyword argument must be specified when calling tile_zeros() function")
+        return Tile(dtype=Any, shape=Any)
 
-    if "n" not in arg_values:
-        raise RuntimeError("'n' keyword argument must be specified when calling tile_zeros() function")
+    shape = tile_unpack_shape(arg_values)
 
     if "dtype" not in arg_values:
-        raise RuntimeError("'dtype' keyword argument must be specified when calling tile_zeros() function")
+        raise TypeError("tile_zeros() missing required keyword argument 'dtype'")
 
     if "storage" not in arg_values:
-        raise ValueError("'storage' keyword not provided for tile_zeros")
+        raise TypeError("tile_zeros() missing required keyword argument 'storage'")
 
     if arg_values["storage"] not in {"shared", "register"}:
-        raise ValueError(
-            f"'storage' keyword argument must be either 'shared' or 'register', got {arg_values['storage']}"
-        )
+        raise ValueError(f"Invalid value for 'storage': {arg_values['storage']!r}. Expected 'shared' or 'register'.")
 
-    m, n = arg_values["m"], arg_values["n"]
     dtype = arg_values["dtype"]
 
-    return TileZeros(dtype=dtype, M=m, N=n, storage=arg_values["storage"])
+    return TileZeros(dtype=dtype, shape=shape, storage=arg_values["storage"])
 
 
 def tile_zeros_dispatch_func(arg_types: Mapping[str, type], return_type: Any, arg_values: Mapping[str, Var]):
-    m, n, dtype = arg_values["m"], arg_values["n"], arg_values["dtype"]
+    shape = tile_unpack_shape(arg_values)
+    dtype = arg_values["dtype"]
 
     template_args = []
     template_args.append(dtype)
-    template_args.append(m.constant)
-    template_args.append(n.constant)
+    for d in shape:
+        template_args.append(d.constant)
 
     return ([], template_args)
 
 
 add_builtin(
     "tile_zeros",
-    input_types={"m": int, "n": int, "dtype": Any, "storage": str},
-    defaults={"storage": "register"},
+    input_types={"shape": Tuple[int, ...], "dtype": Any, "storage": str},
+    defaults={"storage": "register", "dtype": float},
     value_func=tile_zeros_value_func,
     dispatch_func=tile_zeros_dispatch_func,
     variadic=False,
     missing_grad=True,
-    doc="""Allocates a tile of zero-initialized items.
+    doc="""Allocate a tile of zero-initialized items.
 
-    :param m: Size of the first dimension of the output tile
-    :param n: Size of the second dimension of the output tile
-    :param dtype: Datatype of output tile's elements
+    :param shape: Shape of the output tile
+    :param dtype: Data type of output tile's elements (default float)
     :param storage: The storage location for the tile: ``"register"`` for registers
       (default) or ``"shared"`` for shared memory.
-    :returns: A zero-initialized tile with ``shape=(m,n)`` and the specified datatype""",
+    :returns: A zero-initialized tile with shape and data type as specified""",
+    group="Tile Primitives",
+    export=False,
+)
+
+# overload for scalar shape
+add_builtin(
+    "tile_zeros",
+    input_types={"shape": int, "dtype": Any, "storage": str},
+    defaults={"storage": "register", "dtype": float},
+    value_func=tile_zeros_value_func,
+    dispatch_func=tile_zeros_dispatch_func,
+    variadic=False,
+    missing_grad=True,
+    hidden=True,
     group="Tile Primitives",
     export=False,
 )
@@ -1773,54 +1807,63 @@ add_builtin(
 def tile_ones_value_func(arg_types: Mapping[str, type], arg_values: Mapping[str, Any]):
     # return generic type (for doc builds)
     if arg_types is None:
-        return Tile(dtype=Any, M=Any, N=Any)
+        return Tile(dtype=Any, shape=Any)
 
-    if "m" not in arg_values:
-        raise RuntimeError("'m' keyword argument must be specified when calling tile_zeros() function")
-
-    if "n" not in arg_values:
-        raise RuntimeError("'n' keyword argument must be specified when calling tile_zeros() function")
+    shape = tile_unpack_shape(arg_values)
 
     if "dtype" not in arg_values:
-        raise RuntimeError("'dtype' keyword argument must be specified when calling tile_zeros() function")
+        raise TypeError("tile_ones() missing required keyword argument 'dtype'")
+
+    if "storage" not in arg_values:
+        raise TypeError("tile_ones() missing required keyword argument 'storage'")
 
     if arg_values["storage"] not in {"shared", "register"}:
-        raise ValueError(
-            f"'storage' keyword argument must be either 'shared' or 'register', got {arg_values['storage']}"
-        )
+        raise ValueError(f"Invalid value for 'storage': {arg_values['storage']!r}. Expected 'shared' or 'register'.")
 
-    m, n = arg_values["m"], arg_values["n"]
     dtype = arg_values["dtype"]
 
-    return TileZeros(dtype=dtype, M=m, N=n, storage=arg_values["storage"])
+    return TileOnes(dtype=dtype, shape=shape, storage=arg_values["storage"])
 
 
 def tile_ones_dispatch_func(arg_types: Mapping[str, type], return_type: Any, arg_values: Mapping[str, Var]):
-    m, n, dtype = arg_values["m"], arg_values["n"], arg_values["dtype"]
+    shape = tile_unpack_shape(arg_values)
+    dtype = arg_values["dtype"]
 
     template_args = []
     template_args.append(dtype)
-    template_args.append(m.constant)
-    template_args.append(n.constant)
+    for d in shape:
+        template_args.append(d.constant)
 
     return ([], template_args)
 
 
 add_builtin(
     "tile_ones",
-    input_types={"m": int, "n": int, "dtype": Any, "storage": str},
+    input_types={"shape": Tuple[int, ...], "dtype": Any, "storage": str},
     defaults={"storage": "register"},
     value_func=tile_ones_value_func,
     dispatch_func=tile_ones_dispatch_func,
     missing_grad=True,
-    doc="""Allocates a tile of one-initialized items.
+    doc="""Allocate a tile of one-initialized items.
 
-    :param m: Size of the first dimension of the output tile
-    :param n: Size of the second dimension of the output tile
-    :param dtype: Datatype of output tile's elements
+    :param shape: Shape of the output tile
+    :param dtype: Data type of output tile's elements
     :param storage: The storage location for the tile: ``"register"`` for registers
       (default) or ``"shared"`` for shared memory.
-    :returns: A one-initialized tile with ``shape=(m,n)`` and the specified dtype""",
+    :returns: A one-initialized tile with shape and data type as specified""",
+    group="Tile Primitives",
+    export=False,
+)
+
+# overload for scalar shape
+add_builtin(
+    "tile_ones",
+    input_types={"shape": int, "dtype": Any, "storage": str},
+    defaults={"storage": "register"},
+    value_func=tile_ones_value_func,
+    dispatch_func=tile_ones_dispatch_func,
+    missing_grad=True,
+    hidden=True,
     group="Tile Primitives",
     export=False,
 )
@@ -1829,14 +1872,16 @@ add_builtin(
 def tile_arange_value_func(arg_types: Mapping[str, type], arg_values: Mapping[str, Any]):
     # return generic type (for doc builds)
     if arg_types is None:
-        return Tile(dtype=Any, M=Any, N=Any)
+        return Tile(dtype=Any, shape=Any)
+
+    if "args" not in arg_values:
+        raise TypeError("tile_arange() requires at least one positional argument specifying the range")
+
+    args = arg_values["args"]
 
     start = 0
     stop = 0
     step = 1
-    dtype = int
-
-    args = arg_values["args"]
 
     if len(args) == 1:
         start = 0
@@ -1852,8 +1897,7 @@ def tile_arange_value_func(arg_types: Mapping[str, type], arg_values: Mapping[st
         step = args[2]
 
     if start is None or stop is None or step is None:
-        print(args)
-        raise RuntimeError("wp.tile_arange() arguments must be compile time constants")
+        raise RuntimeError("tile_arange() arguments must be compile time constants")
 
     if "dtype" in arg_values:
         dtype = arg_values["dtype"]
@@ -1861,26 +1905,37 @@ def tile_arange_value_func(arg_types: Mapping[str, type], arg_values: Mapping[st
         dtype = float
 
     if arg_values["storage"] not in {"shared", "register"}:
-        raise ValueError(
-            f"'storage' keyword argument must be either 'shared' or 'register', got {arg_values['storage']}"
-        )
+        raise ValueError(f"Invalid value for 'storage': {arg_values['storage']!r}. Expected 'shared' or 'register'.")
 
     return TileRange(dtype=dtype, start=start, stop=stop, step=step, storage=arg_values["storage"])
 
 
 def tile_arange_dispatch_func(arg_types: Mapping[str, type], return_type: Any, arg_values: Mapping[str, Var]):
-    m, n, dtype = return_type.M, return_type.N, return_type.dtype
+    size, dtype = return_type.size, return_type.dtype
 
     template_args = []
     template_args.append(dtype)
-    template_args.append(m)
-    template_args.append(n)
+    template_args.append(size)
+
+    if "args" not in arg_values:
+        raise TypeError("tile_arange() requires at least one positional argument specifying the range")
+
+    args = arg_values["args"]
 
-    # todo: it is somewhat redundant to create new vars here since some of start,stop,step
-    # already exist depending on which form the function was called by the user
-    start = warp.codegen.Var(label=None, type=return_type.dtype, constant=return_type.start)
-    stop = warp.codegen.Var(label=None, type=return_type.dtype, constant=return_type.stop)
-    step = warp.codegen.Var(label=None, type=return_type.dtype, constant=return_type.step)
+    if len(args) == 1:
+        start = warp.codegen.Var(label=None, type=return_type.dtype, constant=return_type.start)
+        stop = args[0]
+        step = warp.codegen.Var(label=None, type=return_type.dtype, constant=return_type.step)
+    elif len(args) == 2:
+        start = args[0]
+        stop = args[1]
+        step = warp.codegen.Var(label=None, type=return_type.dtype, constant=return_type.step)
+    elif len(args) == 3:
+        start = args[0]
+        stop = args[1]
+        step = args[2]
+    else:
+        raise TypeError(f"tile_arange() accepts at most 3 positional arguments, got {len(args)}")
 
     function_args = []
     function_args.append(start)
@@ -1898,7 +1953,7 @@ add_builtin(
     dispatch_func=tile_arange_dispatch_func,
     variadic=True,
     missing_grad=True,
-    doc="""Generates a tile of linearly spaced elements.
+    doc="""Generate a tile of linearly spaced elements.
 
     :param args: Variable-length positional arguments, interpreted as:
 
@@ -1906,246 +1961,157 @@ add_builtin(
         - ``(start, stop)``: Generates values from ``start`` to ``stop - 1``
         - ``(start, stop, step)``: Generates values from ``start`` to ``stop - 1`` with a step size
 
-    :param dtype: Datatype of output tile's elements (optional, default: int)
+    :param dtype: Data type of output tile's elements (optional, default: ``float``)
     :param storage: The storage location for the tile: ``"register"`` for registers
       (default) or ``"shared"`` for shared memory.
-    :returns: A tile with ``shape=(1,n)`` with linearly spaced elements of specified dtype""",
+    :returns: A tile with ``shape=(n)`` with linearly spaced elements of specified data type""",
     group="Tile Primitives",
     export=False,
 )
 
 
-def tile_load_1d_value_func(arg_types, arg_values):
-    # return generic type (for doc builds)
+def tile_load_tuple_value_func(arg_types: Mapping[str, type], arg_values: Mapping[str, Any]):
     if arg_types is None:
-        return Tile(dtype=Any, M=Any, N=Any)
-
-    if not is_array(arg_types["a"]):
-        raise RuntimeError("tile_load() argument 0 must be an array")
+        return array(dtype=Scalar)
 
-    if arg_types["a"].ndim != 1:
-        raise RuntimeError(
-            "tile_load() argument 0 must be 1-dimensional if using the ``wp.tile_load(array, i, n)`` syntax."
-        )
+    a = arg_types["a"]
 
-    if not type_is_int(arg_types["i"]):
-        raise RuntimeError("tile_load() argument 1 must be an integer")
+    shape = tile_unpack_shape(arg_values)
+    offset = tile_unpack_offset(arg_values, a.ndim)
 
-    if "n" not in arg_values:
-        raise RuntimeError("'n' keyword argument must be specified when calling tile_load() function")
+    if a.ndim != len(shape):
+        raise ValueError(
+            f"tile_load() array argument must have same number of dimensions as the tile shape, trying to perform an {len(shape)} dimensional load from an array with {a.ndim} dimensions."
+        )
 
-    if arg_values["storage"] not in {"shared", "register"}:
+    if a.ndim != len(offset):
         raise ValueError(
-            f"'storage' keyword argument must be either 'shared' or 'register', got {arg_values['storage']}"
+            f"tile_load() offset argument must have the same number of dimensions as the array to load from, got {len(offset)} indices for an array with {a.ndim} dimensions"
         )
 
-    a = arg_types["a"]
-    _m, n = 1, arg_values["n"]
+    if arg_values["storage"] not in {"shared", "register"}:
+        raise ValueError(f"Invalid value for 'storage': {arg_values['storage']!r}. Expected 'shared' or 'register'.")
 
-    return TileLoad(a, 1, n, arg_values["storage"])
+    return Tile(dtype=a.dtype, shape=shape, storage=arg_values["storage"])
 
 
-def tile_load_1d_dispatch_func(arg_types: Mapping[str, type], return_type: Any, arg_values: Mapping[str, Var]):
-    array = arg_values["a"]
-    i = arg_values["i"]
-    n = arg_values["n"].constant
-    dtype = arg_values["a"].type.dtype
+def tile_load_tuple_dispatch_func(input_types: Mapping[str, type], return_type: Any, args: Mapping[str, Var]):
+    a = args["a"]
+    shape = tile_unpack_shape(args)
+    offset = tile_unpack_offset(args, a.type.ndim)
 
-    template_args = []
-    template_args.append(dtype)
-    template_args.append(n)
+    func_args = (a, *offset)
+    template_args = (d.constant for d in shape)
 
-    return ((array, i), template_args)
+    return (func_args, template_args)
 
 
 add_builtin(
     "tile_load",
-    input_types={"a": array(dtype=Any), "i": int, "n": int, "storage": str},
-    defaults={"storage": "register"},
-    value_func=tile_load_1d_value_func,
-    dispatch_func=tile_load_1d_dispatch_func,
+    input_types={"a": array(dtype=Any), "shape": Tuple[int, ...], "offset": Tuple[int, ...], "storage": str},
+    value_func=tile_load_tuple_value_func,
+    dispatch_func=tile_load_tuple_dispatch_func,
+    defaults={"offset": None, "storage": "register"},
     variadic=False,
-    doc="""Loads a 1D tile from a global memory array.
+    doc="""Loads a tile from a global memory array.
 
     This method will cooperatively load a tile from global memory using all threads in the block.
 
     :param a: The source array in global memory
-    :param i: Offset in the source array measured in multiples of ``n``, i.e.: ``offset=i*n``
-    :param n: The number of elements in the tile
+    :param shape: Shape of the tile to load, must have the same number of dimensions as ``a``
+    :param offset: Offset in the source array to begin reading from (optional)
     :param storage: The storage location for the tile: ``"register"`` for registers
       (default) or ``"shared"`` for shared memory.
-    :returns: A tile with ``shape=(1,n)`` and dtype the same as the source array""",
+    :returns: A tile with shape as specified and data type the same as the source array""",
     group="Tile Primitives",
     export=False,
 )
 
-
-def tile_load_2d_value_func(arg_types, arg_values):
-    # return generic type (for doc builds)
-    if arg_types is None:
-        return Tile(dtype=Any, M=Any, N=Any)
-
-    if not is_array(arg_types["a"]):
-        raise RuntimeError("tile_load() argument 0 must be an array")
-
-    if arg_types["a"].ndim != 2:
-        raise RuntimeError(
-            "tile_load() argument 0 must be 2-dimensional if using the ``wp.tile_load(array, i, j, m, n)`` syntax."
-        )
-
-    if not type_is_int(arg_types["i"]):
-        raise RuntimeError("tile_load() argument 1 must be an integer")
-
-    if not type_is_int(arg_types["j"]):
-        raise RuntimeError("tile_load() argument 1 must be an integer")
-
-    if "m" not in arg_values:
-        raise RuntimeError("'m' keyword argument must be specified when calling tile_load() function")
-
-    if "n" not in arg_values:
-        raise RuntimeError("'n' keyword argument must be specified when calling tile_load() function")
-
-    if arg_values["storage"] not in {"shared", "register"}:
-        raise ValueError(
-            f"'storage' keyword argument must be either 'shared' or 'register', got {arg_values['storage']}"
-        )
-
-    a = arg_types["a"]
-    m, n = arg_values["m"], arg_values["n"]
-
-    return TileLoad(a, m, n, arg_values["storage"])
-
-
-def tile_load_2d_dispatch_func(arg_types: Mapping[str, type], return_type: Any, arg_values: Mapping[str, Var]):
-    array = arg_values["a"]
-    i, j = arg_values["i"], arg_values["j"]
-    m, n = arg_values["m"].constant, arg_values["n"].constant
-    dtype = arg_values["a"].type.dtype
-
-    template_args = []
-    template_args.append(dtype)
-    template_args.append(m)
-    template_args.append(n)
-
-    return ((array, i, j), template_args)
-
-
+# overload for scalar shape
 add_builtin(
     "tile_load",
-    input_types={"a": array(dtype=Any), "i": int, "j": int, "m": int, "n": int, "storage": str},
-    defaults={"storage": "register"},
-    value_func=tile_load_2d_value_func,
-    dispatch_func=tile_load_2d_dispatch_func,
-    variadic=False,
-    doc="""Loads a 2D tile from a global memory array.
-
-    This method will cooperatively load a tile from global memory using all threads in the block.
-
-    :param a: The source array in global memory
-    :param i: Offset in the source array measured in multiples of ``m``, i.e.: ``row=i*m``
-    :param j: Offset in the source array measured in multiples of ``n``, i.e.; ``col=j*n``
-    :param m: The size of the tile's first dimension
-    :param n: The size of the tile's second dimension
-    :param storage: The storage location for the tile: ``"register"`` for registers
-      (default) or ``"shared"`` for shared memory.
-    :returns: A tile with ``shape=(m,n)`` and dtype the same as the source array""",
+    input_types={"a": array(dtype=Any), "shape": int, "offset": int, "storage": str},
+    value_func=tile_load_tuple_value_func,
+    dispatch_func=tile_load_tuple_dispatch_func,
+    defaults={"offset": None, "storage": "register"},
     group="Tile Primitives",
+    hidden=True,
     export=False,
 )
 
 
-def tile_store_1d_value_func(arg_types, arg_values):
+def tile_store_value_func(arg_types, arg_values):
     # return generic type (for doc builds)
     if arg_types is None:
         return None
 
-    if len(arg_types) != 3:
-        raise RuntimeError("tile_store() requires 3 positional args")
+    a = arg_types["a"]
+    t = arg_types["t"]
 
-    if not is_array(arg_types["a"]):
-        raise RuntimeError("tile_store() argument 0 must be an array")
+    c = tile_unpack_offset(arg_types, a.ndim)
 
-    if arg_types["a"].ndim != 1:
-        raise RuntimeError(
-            "tile_load() argument 0 must be a 1-dimensional array if using the ``wp.tile_store(array, i, t)`` syntax."
+    if len(c) != a.ndim:
+        raise ValueError(
+            f"tile_store() 'a' argument must have {len(c)} dimensions, "
+            f"calculated based on the provided offset arguments, but got {a.ndim} dimensions."
         )
 
-    if not type_is_int(arg_types["i"]):
-        raise RuntimeError("tile_store() argument 1 must be an integer")
-
-    if not is_tile(arg_types["t"]):
-        raise RuntimeError("tile_store() argument 2 must be a tile")
+    if len(t.shape) != a.ndim:
+        raise ValueError(
+            f"tile_store() 'a' argument must have the same number of dimensions as the 't' argument, "
+            f"but got {a.ndim} dimensions for 'a' and {len(t.shape)} dimensions for 't'"
+        )
 
     if not types_equal(arg_types["a"].dtype, arg_types["t"].dtype):
-        raise RuntimeError("tile_store() destination array must have same type as source tile")
+        raise TypeError(
+            f"tile_store() 'a' and 't' arguments must have the same dtype, got {arg_types['a'].dtype} and {arg_types['t'].dtype}"
+        )
 
     return None
 
 
+def tile_store_dispatch_func(input_types: Mapping[str, type], return_type: Any, args: Mapping[str, Var]):
+    a = args["a"]
+    t = args["t"]
+
+    offset = tile_unpack_offset(args, a.type.ndim)
+
+    func_args = (a, *offset, t)
+    template_args = []
+
+    return (func_args, template_args)
+
+
 add_builtin(
     "tile_store",
-    input_types={"a": array(dtype=Any), "i": int, "t": Tile(dtype=Any, M=Any, N=Any)},
-    value_func=tile_store_1d_value_func,
+    input_types={"a": array(dtype=Any), "t": Tile(dtype=Any, shape=Any), "offset": Tuple[int, ...]},
+    value_func=tile_store_value_func,
+    dispatch_func=tile_store_dispatch_func,
+    defaults={"offset": None},
     variadic=False,
     skip_replay=True,
-    doc="""Stores a 1D tile to a global memory array.
+    doc="""Store a tile to a global memory array.
 
     This method will cooperatively store a tile to global memory using all threads in the block.
 
     :param a: The destination array in global memory
-    :param i: Offset in the destination array measured in multiples of ``n``, i.e.: ``offset=i*n``
-    :param t: The source tile to store data from, must have the same dtype as the destination array""",
+    :param t: The source tile to store data from, must have the same data type and number of dimensions as the destination array
+    :param offset: Offset in the destination array (optional)""",
     group="Tile Primitives",
     export=False,
 )
 
-
-def tile_store_2d_value_func(arg_types, arg_values):
-    # return generic type (for doc builds)
-    if arg_types is None:
-        return None
-
-    if len(arg_types) != 4:
-        raise RuntimeError("tile_store() requires 4 positional args")
-
-    if not is_array(arg_types["a"]):
-        raise RuntimeError("tile_store() argument 0 must be an array")
-
-    if arg_types["a"].ndim != 2:
-        raise RuntimeError(
-            "tile_load() argument 0 must be a 2-dimensional array if using the ``wp.tile_store(array, i, j, t)`` syntax."
-        )
-
-    if not type_is_int(arg_types["i"]):
-        raise RuntimeError("tile_store() argument 1 must be an integer")
-
-    if not type_is_int(arg_types["j"]):
-        raise RuntimeError("tile_store() argument 2 must be an integer")
-
-    if not is_tile(arg_types["t"]):
-        raise RuntimeError("tile_store() argument 3 must be a tile")
-
-    if not types_equal(arg_types["a"].dtype, arg_types["t"].dtype):
-        raise RuntimeError("tile_store() destination array must have same type as source tile")
-
-    return None
-
-
+# overload for scalar offset
 add_builtin(
     "tile_store",
-    input_types={"a": array(dtype=Any), "i": int, "j": int, "t": Tile(dtype=Any, M=Any, N=Any)},
-    value_func=tile_store_2d_value_func,
+    input_types={"a": array(dtype=Any), "t": Tile(dtype=Any, shape=Any), "offset": int},
+    value_func=tile_store_value_func,
+    dispatch_func=tile_store_dispatch_func,
+    defaults={"offset": None},
     variadic=False,
     skip_replay=True,
-    doc="""Stores a tile to a global memory array.
-
-    This method will cooperatively store a tile to global memory using all threads in the block.
-
-    :param a: The destination array in global memory
-    :param i: Offset in the destination array measured in multiples of ``m``, i.e.: ``row=i*m``
-    :param j: Offset in the destination array measured in multiples of ``n``, i.e.; ``col=j*n``
-    :param t: The source tile to store data from, must have the same dtype as the destination array""",
     group="Tile Primitives",
+    hidden=True,
     export=False,
 )
 
@@ -2153,130 +2119,219 @@ add_builtin(
 def tile_atomic_add_value_func(arg_types, arg_values):
     # return generic type (for doc builds)
     if arg_types is None:
-        return Tile(dtype=Any, M=Any, N=Any)
+        return Tile(dtype=Any, shape=Any)
+
+    a = arg_types["a"]
+    t = arg_types["t"]
+
+    c = tile_unpack_offset(arg_types, a.ndim)
+    if len(c) != a.ndim:
+        raise ValueError(
+            f"tile_atomic_add() 'a' argument must have {len(c)} dimensions, "
+            f"calculated based on the provided offset arguments, but got {a.ndim} dimensions."
+        )
+
+    if a.ndim != len(t.shape):
+        raise ValueError(
+            f"tile_atomic_add() 'a' argument must have the same number of dimensions as the 't' argument, "
+            f"but got {a.ndim} dimensions for 'a' and {len(t.shape)} dimensions for 't'"
+        )
 
-    if len(arg_types) != 4:
-        raise RuntimeError("tile_atomic_add() requires 4 positional args")
+    if not types_equal(arg_types["a"].dtype, arg_types["t"].dtype):
+        raise TypeError(
+            f"tile_atomic_add() 'a' and 't' arguments must have the same dtype, got {arg_types['a'].dtype} and {arg_types['t'].dtype}"
+        )
 
-    if not is_array(arg_types["a"]):
-        raise RuntimeError("tile_atomic_add() argument 0 must be an array")
+    return Tile(dtype=arg_types["t"].dtype, shape=arg_types["t"].shape)
 
-    if not type_is_int(arg_types["x"]):
-        raise RuntimeError("tile_atomic_add() argument 1 must be an integer")
 
-    if not type_is_int(arg_types["y"]):
-        raise RuntimeError("tile_atomic_add() argument 2 must be an integer")
+def tile_atomic_add_dispatch_func(input_types: Mapping[str, type], return_type: Any, args: Mapping[str, Var]):
+    a = args["a"]
+    t = args["t"]
 
-    if not is_tile(arg_types["t"]):
-        raise RuntimeError("tile_atomic_add() argument 3 must be a tile")
+    offset = tile_unpack_offset(args, a.type.ndim)
 
-    if arg_types["a"].dtype != arg_types["t"].dtype:
-        raise RuntimeError("tile_atomic_add() tile dtype and array dtype must match")
+    func_args = (a, *offset, t)
+    template_args = []
 
-    return Tile(dtype=arg_types["t"].dtype, M=arg_types["t"].M, N=arg_types["t"].N)
+    return (func_args, template_args)
 
 
 add_builtin(
     "tile_atomic_add",
-    input_types={"a": array(dtype=Any), "x": int, "y": int, "t": Tile(dtype=Any, M=Any, N=Any)},
+    input_types={"a": array(dtype=Any), "t": Tile(dtype=Any, shape=Any), "offset": Tuple[int, ...]},
     value_func=tile_atomic_add_value_func,
-    variadic=True,
+    dispatch_func=tile_atomic_add_dispatch_func,
+    defaults={"offset": None},
+    variadic=False,
     skip_replay=True,
-    doc="""Atomically add a tile to the array `a`, each element will be updated atomically.
+    doc="""Atomically add a 1D tile to the array `a`, each element will be updated atomically.
 
     :param a: Array in global memory, should have the same ``dtype`` as the input tile
-    :param x: Offset in the destination array measured in multiples of ``m``, i.e.: ``i=x*M`` where ``M`` is the first tile dimension
-    :param y: Offset in the destination array measured in multiples of ``n``, i.e.: ``j=y*N`` where ``N`` is the second tile dimension
     :param t: Source tile to add to the destination array
-    :returns: A tile with the same dimensions and type as the source tile, holding the original value of the destination elements""",
+    :param offset: Offset in the destination array (optional)
+    :returns: A tile with the same dimensions and data type as the source tile, holding the original value of the destination elements""",
     group="Tile Primitives",
     export=False,
 )
 
+# overload for scalar offset
+add_builtin(
+    "tile_atomic_add",
+    input_types={"a": array(dtype=Any), "t": Tile(dtype=Any, shape=Any), "offset": int},
+    value_func=tile_atomic_add_value_func,
+    dispatch_func=tile_atomic_add_dispatch_func,
+    defaults={"offset": None},
+    variadic=False,
+    skip_replay=True,
+    group="Tile Primitives",
+    hidden=True,
+    export=False,
+)
+
 
 def tile_view_value_func(arg_types, arg_values):
     # return generic type (for doc builds)
     if arg_types is None:
-        return Tile(dtype=Any, M=Any, N=Any)
+        return Tile(dtype=Any, shape=Any)
 
     tile = arg_types["t"]
+    offset = arg_types["offset"]
 
-    if "m" not in arg_values:
-        m = 1
-    else:
-        m = arg_values["m"]
+    if len(offset) > len(tile.shape):
+        raise ValueError(f"tile_view() specified too many offset coordinates {len(offset)} > {len(tile.shape)}")
 
-    if "n" not in arg_values:
-        n = tile.N
+    if "shape" in arg_values:
+        # if shape is specified take it directly, e.g.:
+        # tile_view(t, offset=(i,j), shape=(m,n))
+        shape = arg_values["shape"]
+        strides = tile.strides
+
+        if len(shape) != len(tile.shape):
+            raise ValueError(
+                f"tile_view() if shape is specified it must have same number of dimensions as source tile, expected {len(tile.shape)}, got {len(shape)}"
+            )
     else:
-        n = arg_values["n"]
+        # if not specified, then take output shape from unspecified src dimensions
+        # e.g.: tile[i] will return a whole row of a 2D tile
+        shape = tile.shape[len(offset) :]
+        strides = tile.strides[len(offset) :]
 
-    if m > tile.M or n > tile.N:
-        raise RuntimeError(
-            f"Trying to view a tile subrange with dimensions ({m}, {n}) which is larger than source tile with dimensions ({tile.M}, {tile.N})"
-        )
+    assert len(shape) == len(strides)
 
     # force source tile to shared memory
     tile.storage = "shared"
 
-    output = Tile(dtype=tile.dtype, M=m, N=n, strides=tile.strides, layout=tile.layout, storage="shared", owner=False)
+    output = Tile(dtype=tile.dtype, shape=shape, strides=strides, layout=tile.layout, storage="shared", owner=False)
     return output
 
 
 def tile_view_dispatch_func(arg_types: Mapping[str, type], return_type: Any, arg_values: Mapping[str, Var]):
     tile = arg_values["t"]
-    i = arg_values["i"]
-
-    if "j" not in arg_values:
-        j = warp.codegen.Var(label=None, type=int, constant=0)
-    else:
-        j = arg_values["j"]
+    coord = arg_values["offset"]
 
-    template_args = []
-    template_args.append(return_type.M)
-    template_args.append(return_type.N)
+    # zero-pad coord to match source array
+    view_coord = [0] * len(tile.type.shape)
+    for i in range(len(coord)):
+        view_coord[i] = coord[i]
 
-    return ((tile, i, j), template_args)
+    return ((tile, *view_coord), (return_type,))
 
 
 add_builtin(
     "tile_view",
-    input_types={"t": Tile(dtype=Any, M=Any, N=Any), "i": int, "j": int, "m": int, "n": int},
+    input_types={"t": Tile(dtype=Any, shape=Any), "offset": Tuple[int, ...], "shape": Tuple[int, ...]},
     value_func=tile_view_value_func,
     dispatch_func=tile_view_dispatch_func,
-    defaults={"j": None, "m": None, "n": None},
-    variadic=True,
-    doc="""Return a subrange of a given tile from coordinates (i,j) to (i+m, j+n).
+    defaults={"shape": None},
+    variadic=False,
+    doc="""Return a slice of a given tile [offset, offset+shape], if shape is not specified it will be inferred from the unspecified offset dimensions.
 
     :param t: Input tile to extract a subrange from
-    :param i: Offset in the source tile along the first dimension
-    :param j: Offset in the source tile along the second dimensions
-    :param m: Size of the subrange to return along the first dimension
-    :param n: Size of the subrange to return along the second dimension
-    :returns: A tile with dimensions (m,n) and the same datatype as the input tile""",
+    :param offset: Offset in the source tile
+    :param shape: Shape of the returned slice
+    :returns: A tile with dimensions given by the specified shape or the remaining source tile dimensions""",
     group="Tile Primitives",
+    missing_grad=True,
     export=False,
 )
 
 
 def tile_assign_value_func(arg_types, arg_values):
-    # return generic type (for doc builds)
+    if arg_types is None:
+        return None
+
+    # force the destination tile to shared memory
+    arg_types["dst"].storage = "shared"
     return None
 
 
+def tile_assign_dispatch_func(input_types: Mapping[str, type], return_type: Any, args: Mapping[str, Var]):
+    dst = args["dst"]
+    src = args["src"]
+
+    offset = tile_unpack_offset(args, len(dst.type.shape))
+
+    func_args = (dst, src, *offset)
+    template_args = []
+
+    return (func_args, template_args)
+
+
 add_builtin(
     "tile_assign",
-    input_types={"dst": Tile(dtype=Any, M=Any, N=Any), "i": int, "j": int, "src": Tile(dtype=Any, M=Any, N=Any)},
+    input_types={"dst": Tile(dtype=Any, shape=Any), "src": Tile(dtype=Any, shape=Any), "offset": Tuple[int, ...]},
+    value_func=tile_assign_value_func,
+    dispatch_func=tile_assign_dispatch_func,
+    defaults={"offset": None},
+    doc="""Assign a tile to a subrange of a destination tile.
+
+    :param dst: The destination tile to assign to
+    :param src: The source tile to read values from
+    :param offset: Offset in the destination tile to write to""",
+    group="Tile Primitives",
+    export=False,
+)
+
+# handles expressions like tile[i,j] = 1.0
+add_builtin(
+    "assign",
+    input_types={"dst": Tile(dtype=Any, shape=Any), "i": int, "src": Scalar},
+    value_func=tile_assign_value_func,
+    group="Tile Primitives",
+    export=False,
+    hidden=True,
+    missing_grad=True,
+)
+
+add_builtin(
+    "assign",
+    input_types={"dst": Tile(dtype=Any, shape=Any), "i": int, "j": int, "src": Scalar},
     value_func=tile_assign_value_func,
-    #    dispatch_func=tile_assign_dispatch_func,
-    doc="""Assign a tile to a subrange of a destination tile at coordinates (i,j).
+    group="Tile Primitives",
+    export=False,
+    hidden=True,
+    missing_grad=True,
+)
+
+add_builtin(
+    "assign",
+    input_types={"dst": Tile(dtype=Any, shape=Any), "i": int, "j": int, "k": int, "src": Scalar},
+    value_func=tile_assign_value_func,
+    group="Tile Primitives",
+    export=False,
+    hidden=True,
+    missing_grad=True,
+)
 
-    :param t: The destination tile to assign to
-    :param i: Offset in the source tile along the first dimension
-    :param j: Offset in the source tile along the second dimensions
-    :param src: The source tile to read values from""",
+add_builtin(
+    "assign",
+    input_types={"dst": Tile(dtype=Any, shape=Any), "i": int, "j": int, "k": int, "l": int, "src": Scalar},
+    value_func=tile_assign_value_func,
     group="Tile Primitives",
     export=False,
+    hidden=True,
+    missing_grad=True,
 )
 
 
@@ -2286,7 +2341,7 @@ def tile_value_func(arg_types, arg_values):
         return Tile
 
     if len(arg_types) != 1:
-        raise RuntimeError("tile() requires 1 positional arg")
+        raise TypeError(f"tile() takes exactly 1 positional argument but {len(arg_types)} were given")
 
     dtype = None
     length = None
@@ -2294,11 +2349,12 @@ def tile_value_func(arg_types, arg_values):
     if type_is_vector(arg_types["x"]):
         dtype = arg_types["x"]._wp_scalar_type_
         length = arg_types["x"]._shape_[0]
+        shape = (length, warp.codegen.options["block_dim"])
     else:
         dtype = arg_types["x"]
-        length = 1
+        shape = (warp.codegen.options["block_dim"],)
 
-    return Tile(dtype=dtype, M=length, N=warp.codegen.options["block_dim"], op="tile")
+    return Tile(dtype=dtype, shape=shape, op="tile")
 
 
 add_builtin(
@@ -2306,14 +2362,14 @@ add_builtin(
     input_types={"x": Any},
     value_func=tile_value_func,
     variadic=True,
-    doc="""Constructs a new Tile from per-thread kernel values.
+    doc="""Construct a new tile from per-thread kernel values.
 
     This function converts values computed using scalar kernel code to a tile representation for input into collective operations.
 
     * If the input value is a scalar, then the resulting tile has ``shape=(1, block_dim)``
     * If the input value is a vector, then the resulting tile has ``shape=(length(vector), block_dim)``
 
-    :param x: A per-thread local value, e.g.: scalar, vector, or matrix.
+    :param x: A per-thread local value, e.g. scalar, vector, or matrix.
     :returns: A tile with first dimension according to the value type length and a second dimension equal to ``block_dim``
 
     This example shows how to create a linear sequence from thread variables:
@@ -2332,7 +2388,7 @@ add_builtin(
 
     .. code-block:: text
 
-        tile(m=1, n=16, storage=register) = [[0 2 4 6 8 ...]]
+        [0 2 4 6 8 10 12 14 16 18 20 22 24 26 28 30] = tile(shape=(16), storage=register)
 
     """,
     group="Tile Primitives",
@@ -2346,38 +2402,40 @@ def untile_value_func(arg_types, arg_values):
         return Scalar
 
     if len(arg_types) != 1:
-        raise RuntimeError("untile() requires 1 positional arg")
+        raise TypeError(f"untile() takes exactly 1 positional argument but {len(arg_types)} were given")
 
     t = arg_types["a"]
 
     if not is_tile(t):
-        raise RuntimeError(f"untile() accepts arguments of type tile only, got {arg_types[0]}")
+        raise TypeError(f"untile() argument must be a tile, got {t!r}")
 
-    if t.N != warp.codegen.options["block_dim"]:
-        raise RuntimeError(
-            f"untile() argument must have the same length as the block width, got {t.N}, expected {warp.codegen.options['block_dim']}"
+    if t.shape[-1] != warp.codegen.options["block_dim"]:
+        raise ValueError(
+            f"untile() argument last dimension {t.shape[-1]} does not match the expected block width {warp.codegen.options['block_dim']}"
         )
 
-    if t.M == 1:
+    if len(t.shape) == 1:
         return t.dtype
-    elif t.M > 1:
-        return warp.types.vector(t.M, t.dtype)
+    elif len(t.shape) == 2:
+        return warp.types.vector(t.shape[0], t.dtype)
+    else:
+        raise ValueError(f"untile() argument must have a positive size in dimension 0, but got {t.shape[0]}")
 
 
 add_builtin(
     "untile",
-    input_types={"a": Tile(dtype=Any, M=Any, N=Any)},
+    input_types={"a": Tile(dtype=Any, shape=Any)},
     value_func=untile_value_func,
     variadic=True,
-    doc="""Convert a Tile back to per-thread values.
+    doc="""Convert a tile back to per-thread values.
 
     This function converts a block-wide tile back to per-thread values.
 
-    * If the input tile is 1-dimensional then the resulting value will be a per-thread scalar
-    * If the input tile is 2-dimensional then the resulting value will be a per-thread vector of length M
+    * If the input tile is 1D, then the resulting value will be a per-thread scalar
+    * If the input tile is 2D, then the resulting value will be a per-thread vector of length M
 
     :param a: A tile with dimensions ``shape=(M, block_dim)``
-    :returns: A single value per-thread with the same dtype as the tile
+    :returns: A single value per-thread with the same data type as the tile
 
     This example shows how to create a linear sequence from thread variables:
 
@@ -2418,21 +2476,58 @@ def tile_extract_value_func(arg_types, arg_values):
     if arg_types is None:
         return Scalar
 
-    if len(arg_types) != 3:
-        raise RuntimeError("tile_extract() requires 3 positional args")
-
-    if not is_tile(arg_types["a"]):
-        raise RuntimeError("tile_extract() argument 0 must be a tile")
+    # force the input tile to shared memory
+    arg_types["a"].storage = "shared"
 
     return arg_types["a"].dtype
 
 
 add_builtin(
     "tile_extract",
-    input_types={"a": Tile(dtype=Any, M=Any, N=Any), "i": int, "j": int},
+    input_types={"a": Tile(dtype=Any, shape=Any), "i": int},
     value_func=tile_extract_value_func,
-    variadic=True,
-    doc="""Extracts a single element from the tile and returns it as a scalar type.
+    variadic=False,
+    doc="""Extract a single element from the tile and return it as a scalar type.
+
+    This function will extract an element from the tile and broadcast its value to all threads in the block.
+
+    Note that this may incur additional synchronization if the source tile is a register tile.
+
+    :param a: Tile to extract the element from
+    :param i: Coordinate of element on first dimension
+    :returns: The value of the element at the specified tile location with the same data type as the input tile""",
+    group="Tile Primitives",
+    hidden=True,
+    export=False,
+)
+
+
+add_builtin(
+    "tile_extract",
+    input_types={"a": Tile(dtype=Any, shape=Any), "i": int, "j": int},
+    value_func=tile_extract_value_func,
+    variadic=False,
+    doc="""Extract a single element from the tile and return it as a scalar type.
+
+    This function will extract an element from the tile and broadcast its value to all threads in the block.
+
+    Note that this may incur additional synchronization if the source tile is a register tile.
+
+    :param a: Tile to extract the element from
+    :param i: Coordinate of element on first dimension
+    :param j: Coordinate of element on the second dimension
+    :returns: The value of the element at the specified tile location with the same data type as the input tile""",
+    group="Tile Primitives",
+    hidden=True,
+    export=False,
+)
+
+add_builtin(
+    "tile_extract",
+    input_types={"a": Tile(dtype=Any, shape=Any), "i": int, "j": int, "k": int},
+    value_func=tile_extract_value_func,
+    variadic=False,
+    doc="""Extract a single element from the tile and return it as a scalar type.
 
     This function will extract an element from the tile and broadcast its value to all threads in the block.
 
@@ -2441,8 +2536,32 @@ add_builtin(
     :param a: Tile to extract the element from
     :param i: Coordinate of element on first dimension
     :param j: Coordinate of element on the second dimension
-    :returns: The value of the element at the specified tile location, with the same type as the input tile's per-element dtype""",
+    :param k: Coordinate of element on the third dimension
+    :returns: The value of the element at the specified tile location with the same data type as the input tile""",
     group="Tile Primitives",
+    hidden=True,
+    export=False,
+)
+
+add_builtin(
+    "tile_extract",
+    input_types={"a": Tile(dtype=Any, shape=Any), "i": int, "j": int, "k": int, "l": int},
+    value_func=tile_extract_value_func,
+    variadic=False,
+    doc="""Extract a single element from the tile and return it as a scalar type.
+
+    This function will extract an element from the tile and broadcast its value to all threads in the block.
+
+    Note that this may incur additional synchronization if the source tile is a register tile.
+
+    :param a: Tile to extract the element from
+    :param i: Coordinate of element on first dimension
+    :param j: Coordinate of element on the second dimension
+    :param k: Coordinate of element on the third dimension
+    :param l: Coordinate of element on the fourth dimension
+    :returns: The value of the element at the specified tile location, with the same data type as the input tile""",
+    group="Tile Primitives",
+    hidden=True,
     export=False,
 )
 
@@ -2453,12 +2572,12 @@ def tile_transpose_value_func(arg_types, arg_values):
         return Tile
 
     if len(arg_types) != 1:
-        raise RuntimeError("tile_transpose() requires 1 positional args")
+        raise TypeError(f"tile_transpose() takes exactly 1 positional argument but {len(arg_types)} were given")
 
     t = arg_types["a"]
 
     if not is_tile(t):
-        raise RuntimeError("tile_transpose() argument 0 must be a tile")
+        raise TypeError(f"tile_transpose() argument must be a tile, got {t!r}")
 
     layout = None
 
@@ -2473,8 +2592,7 @@ def tile_transpose_value_func(arg_types, arg_values):
 
     return Tile(
         dtype=t.dtype,
-        M=t.N,
-        N=t.M,
+        shape=t.shape[::-1],
         op="transpose",
         storage=t.storage,
         strides=t.strides[::-1],
@@ -2485,12 +2603,13 @@ def tile_transpose_value_func(arg_types, arg_values):
 
 add_builtin(
     "tile_transpose",
-    input_types={"a": Tile(dtype=Any, M=Any, N=Any)},
+    input_types={"a": Tile(dtype=Any, shape=Any)},
     value_func=tile_transpose_value_func,
     variadic=True,
     doc="""Transpose a tile.
 
-    For shared memory tiles this operation will alias the input tile, register tiles will first be transferred to shared memory before transposition.
+    For shared memory tiles, this operation will alias the input tile.
+    Register tiles will first be transferred to shared memory before transposition.
 
     :param a: Tile to transpose with ``shape=(M,N)``
     :returns: Tile with ``shape=(N,M)``""",
@@ -2504,41 +2623,36 @@ def tile_broadcast_value_func(arg_types, arg_values):
     if arg_types is None:
         return Tile
 
-    if len(arg_types) != 3:
-        raise RuntimeError("tile_broadcast() requires 1 positional args")
-
     t = arg_types["a"]
-    m = arg_values["m"]
-    n = arg_values["n"]
 
-    if not is_tile(t):
-        raise RuntimeError("tile_broadcast() argument 0 must be a tile")
+    # target shape and strides
+    target_shape = tile_unpack_shape(arg_values)
+    target_strides = [0] * len(target_shape)
 
-    # try to broadcast last dimension
-    if t.N == 1:
-        stride_n = 0
-    elif t.N == n:
-        stride_n = t.strides[1]
-    else:
-        raise RuntimeError(
-            f"Broadcast dimension must be 1 or match destination, shape(src) = {t.m, t.n}, shape(dest) = {m, n}"
-        )
+    offset = len(target_shape) - len(t.shape)
 
-    # try to broadcast first dimension
-    if t.M == 1:
-        stride_m = 0
-    elif t.M == m:
-        stride_m = t.strides[0]
-    else:
-        raise RuntimeError(
-            f"Broadcast dimension must be 1 or match destination, shape(src) = {t.m, t.n}, shape(dest) = {m, n}"
-        )
+    # compute target strides
+    for i in reversed(range(len(target_shape))):
+        j = i - offset
+
+        if j < 0:
+            target_strides[i] = 0
+        else:
+            # try to broadcast each dimension
+            if t.shape[j] == 1:
+                target_strides[i] = 0
+            elif t.shape[j] == target_shape[i]:
+                target_strides[i] = t.strides[j]
+            else:
+                raise ValueError(
+                    f"tile_broadcast() cannot broadcast dimension {t.shape[j]} into {target_shape[i]} at index {i}"
+                )
 
     # force the input tile to shared memory
     t.storage = "shared"
 
     tile_type = Tile(
-        dtype=t.dtype, M=m, N=n, op="broadcast", storage=t.storage, strides=(stride_m, stride_n), owner=False
+        dtype=t.dtype, shape=target_shape, op="broadcast", storage=t.storage, strides=target_strides, owner=False
     )
     return tile_type
 
@@ -2547,8 +2661,8 @@ def tile_broadcast_dispatch_func(arg_types: Mapping[str, type], return_type: Any
     tile = arg_values["a"]
 
     template_args = []
-    template_args.append(return_type.M)
-    template_args.append(return_type.N)
+    template_args.append(return_type.shape[0])
+    template_args.append(return_type.shape[1])
     template_args.append(return_type.strides[0])
     template_args.append(return_type.strides[1])
 
@@ -2557,15 +2671,18 @@ def tile_broadcast_dispatch_func(arg_types: Mapping[str, type], return_type: Any
 
 add_builtin(
     "tile_broadcast",
-    input_types={"a": Tile(dtype=Any, M=Any, N=Any), "m": int, "n": int},
+    input_types={"a": Tile(dtype=Any, shape=Any), "shape": Tuple[int, ...]},
     value_func=tile_broadcast_value_func,
     dispatch_func=tile_broadcast_dispatch_func,
-    variadic=True,
+    variadic=False,
     doc="""Broadcast a tile.
 
-    This function will attempt to broadcast the input tile ``a`` to the destination shape (m, n), broadcasting follows NumPy broadcast rules.
+    This function will attempt to broadcast the input tile ``a`` to the destination shape (m, n).
+
+    Broadcasting follows NumPy broadcast rules.
 
     :param a: Tile to broadcast
+    :param shape: The shape to broadcast to
     :returns: Tile with broadcast ``shape=(m, n)``""",
     group="Tile Primitives",
     export=False,
@@ -2575,19 +2692,10 @@ add_builtin(
 def tile_matmul_value_func(arg_types, arg_values):
     # return generic type (for doc builds)
     if arg_types is None:
-        return Tile(dtype=Any, M=Any, N=Any)
+        return Tile(dtype=Any, shape=Any)
 
     if len(arg_types) != 3:
-        raise RuntimeError("tile_matmul() requires 4 positional args")
-
-    if not is_tile(arg_types["a"]):
-        raise RuntimeError("tile_matmul() argument 0 must be a tile")
-
-    if not is_tile(arg_types["b"]):
-        raise RuntimeError("tile_matmul() argument 1 must be an tile")
-
-    if not isinstance(arg_types["out"], Tile):
-        raise RuntimeError("tile_matmul() output argument must be a tile")
+        raise TypeError(f"tile_matmul() takes exactly 3 positional arguments but {len(arg_types)} were given")
 
     return None
 
@@ -2622,17 +2730,17 @@ add_builtin(
 def tile_sum_value_func(arg_types, arg_values):
     # return generic type (for doc builds)
     if arg_types is None:
-        return Tile(dtype=Any, M=1, N=1)
+        return Tile(dtype=Any, shape=(1,))
 
     if len(arg_types) != 1:
-        raise RuntimeError("tile_sum() requires 1 positional args")
+        raise TypeError(f"tile_sum() takes exactly 1 positional argument but {len(arg_types)} were given")
 
     a = arg_types["a"]
 
     if not is_tile(a):
-        raise RuntimeError("tile_sum() argument 0 must be a tile")
+        raise TypeError(f"tile_sum() argument must be a tile, got {a!r}")
 
-    return Tile(dtype=a.dtype, M=1, N=1, op="sum")
+    return Tile(dtype=a.dtype, shape=(1,), op="sum")
 
 
 add_builtin(
@@ -2643,7 +2751,7 @@ add_builtin(
     doc="""Cooperatively compute the sum of the tile elements using all threads in the block.
 
     :param a: The tile to compute the sum of
-    :returns: A single-element tile with dimensions of (1,1) holding the sum
+    :returns: A single-element tile holding the sum
 
     Example:
 
@@ -2652,7 +2760,7 @@ add_builtin(
         @wp.kernel
         def compute():
 
-            t = wp.tile_ones(dtype=float, m=16, n=16)
+            t = wp.tile_ones(dtype=float, shape=(16, 16))
             s = wp.tile_sum(t)
 
             print(s)
@@ -2663,7 +2771,7 @@ add_builtin(
 
     .. code-block:: text
 
-        tile(m=1, n=1, storage=register) = [[256]]
+        [256] = tile(shape=(1), storage=register)
 
     """,
     group="Tile Primitives",
@@ -2674,17 +2782,17 @@ add_builtin(
 def tile_min_value_func(arg_types, arg_values):
     # return generic type (for doc builds)
     if arg_types is None:
-        return Tile(dtype=Any, M=1, N=1)
+        return Tile(dtype=Any, shape=(1,))
 
     if len(arg_types) != 1:
-        raise RuntimeError("tile_min() requires 1 positional args")
+        raise TypeError(f"tile_min() takes exactly 1 positional argument but {len(arg_types)} were given")
 
     a = arg_types["a"]
 
     if not is_tile(a):
-        raise RuntimeError("tile_min() argument 0 must be a tile")
+        raise TypeError(f"tile_min() argument must be a tile, got {a!r}")
 
-    return Tile(dtype=a.dtype, M=1, N=1, op="min")
+    return Tile(dtype=a.dtype, shape=(1,), op="min")
 
 
 add_builtin(
@@ -2695,7 +2803,7 @@ add_builtin(
     doc="""Cooperatively compute the minimum of the tile elements using all threads in the block.
 
     :param a: The tile to compute the minimum of
-    :returns: A single-element tile with dimensions of (1,1) holding the minimum value
+    :returns: A single-element tile holding the minimum value
 
     Example:
 
@@ -2716,7 +2824,7 @@ add_builtin(
 
     .. code-block:: text
 
-        tile(m=1, n=1, storage=register) = [[64 ]]
+        [64] = tile(shape=(1), storage=register)
 
     """,
     group="Tile Primitives",
@@ -2727,28 +2835,28 @@ add_builtin(
 def tile_max_value_func(arg_types, arg_values):
     # return generic type (for doc builds)
     if arg_types is None:
-        return Tile(dtype=Any, M=1, N=1)
+        return Tile(dtype=Any, shape=(1,))
 
     if len(arg_types) != 1:
-        raise RuntimeError("tile_max() requires 1 positional args")
+        raise TypeError(f"tile_max() takes exactly 1 positional argument but {len(arg_types)} were given")
 
     a = arg_types["a"]
 
     if not is_tile(a):
-        raise RuntimeError("tile_max() argument 0 must be a tile")
+        raise TypeError(f"tile_max() argument must be a tile, got {a!r}")
 
-    return Tile(dtype=a.dtype, M=1, N=1, op="min")
+    return Tile(dtype=a.dtype, shape=(1,), op="min")
 
 
 add_builtin(
     "tile_max",
-    input_types={"a": Tile},
+    input_types={"a": Tile(dtype=Any, shape=Any)},
     value_func=tile_max_value_func,
-    variadic=True,
+    variadic=False,
     doc="""Cooperatively compute the maximum of the tile elements using all threads in the block.
 
     :param a: The tile to compute the maximum from
-    :returns: A single-element tile with dimensions of (1,1) holding the maximum value
+    :returns: A single-element tile holding the maximum value
 
     Example:
 
@@ -2768,7 +2876,7 @@ add_builtin(
 
     .. code-block:: text
 
-        tile(m=1, n=1, storage=register) = [[127 ]]
+        [127] = tile(shape=(1), storage=register)
 
     """,
     group="Tile Primitives",
@@ -2779,15 +2887,14 @@ add_builtin(
 # does type propagation for load()
 def tile_reduce_value_func(arg_types, arg_values):
     if arg_types is None:
-        return Tile(dtype=Any, M=Any, N=Any)
+        return Tile(dtype=Any, shape=(1,))
 
     a = arg_types["a"]
 
-    # check all args are tiles
     if not is_tile(a):
-        raise RuntimeError(f"tile_reduce() arguments must be tiles, got type {a}")
+        raise TypeError(f"tile_reduce() 'a' argument must be a tile, got {a!r}")
 
-    return Tile(dtype=a.dtype, M=1, N=1, op="reduce")
+    return Tile(dtype=a.dtype, shape=(1,), op="reduce")
 
 
 def tile_reduce_dispatch_func(input_types: Mapping[str, type], return_type: Any, args: Mapping[str, Var]):
@@ -2798,7 +2905,7 @@ def tile_reduce_dispatch_func(input_types: Mapping[str, type], return_type: Any,
 
 add_builtin(
     "tile_reduce",
-    input_types={"op": Callable, "a": Tile(dtype=Any, M=Any, N=Any)},
+    input_types={"op": Callable, "a": Tile(dtype=Any, shape=Any)},
     value_func=tile_reduce_value_func,
     native_func="tile_reduce",
     doc="""Apply a custom reduction operator across the tile.
@@ -2806,8 +2913,8 @@ add_builtin(
     This function cooperatively performs a reduction using the provided operator across the tile.
 
     :param op: A callable function that accepts two arguments and returns one argument, may be a user function or builtin
-    :param a: The input tile, the operator (or one of its overloads) must be able to accept the tile's dtype
-    :returns: A single-element tile with ``shape=(1,1)`` with the same datatype as the input tile.
+    :param a: The input tile, the operator (or one of its overloads) must be able to accept the tile's data type
+    :returns: A single-element tile with the same data type as the input tile.
 
     Example:
 
@@ -2827,7 +2934,7 @@ add_builtin(
 
     .. code-block:: text
 
-        tile(m=1, n=1, storage=register) = [[362880]]
+        [362880] = tile(shape=(1), storage=register)
     """,
     group="Tile Primitives",
     export=False,
@@ -2839,26 +2946,19 @@ add_builtin(
 # does type propagation for load()
 def tile_unary_map_value_func(arg_types, arg_values):
     if arg_types is None:
-        return Tile(dtype=Any, M=Any, N=Any)
+        return Tile(dtype=Any, shape=Any)
 
     a = arg_types["a"]
 
-    # check all args are tiles
     if not is_tile(a):
-        raise RuntimeError(f"tile_map() arguments must be tiles, got type {a}")
+        raise TypeError(f"tile_map() 'a' argument must be a tile, got {a!r}")
 
     return TileUnaryMap(a)
 
 
-# def tile_map_dispatch_func(input_types: Mapping[str, type], return_type: Any, args: Mapping[str, Var]):
-#    func_args = (args["op"], *args["args"])
-#    template_args = ()
-#    return (func_args, template_args)
-
-
 add_builtin(
     "tile_map",
-    input_types={"op": Callable, "a": Tile(dtype=Any, M=Any, N=Any)},
+    input_types={"op": Callable, "a": Tile(dtype=Any, shape=Any)},
     value_func=tile_unary_map_value_func,
     # dispatch_func=tile_map_dispatch_func,
     # variadic=True,
@@ -2868,8 +2968,8 @@ add_builtin(
     This function cooperatively applies a unary function to each element of the tile using all threads in the block.
 
     :param op: A callable function that accepts one argument and returns one argument, may be a user function or builtin
-    :param a: The input tile, the operator (or one of its overloads) must be able to accept the tile's dtype
-    :returns: A tile with the same dimensions and datatype as the input tile.
+    :param a: The input tile, the operator (or one of its overloads) must be able to accept the tile's data type
+    :returns: A tile with the same dimensions and data type as the input tile.
 
     Example:
 
@@ -2889,7 +2989,7 @@ add_builtin(
 
     .. code-block:: text
 
-        tile(m=1, n=10, storage=register) = [[0 0.0998334 0.198669 0.29552 ...]]
+        [0 0.0998334 0.198669 0.29552 0.389418 0.479426 0.564642 0.644218 0.717356 0.783327] = tile(shape=(10), storage=register)
     """,
     group="Tile Primitives",
     export=False,
@@ -2898,34 +2998,37 @@ add_builtin(
 
 def tile_binary_map_value_func(arg_types, arg_values):
     if arg_types is None:
-        return Tile(dtype=Any, M=Any, N=Any)
+        return Tile(dtype=Any, shape=Any)
 
     a = arg_types["a"]
     b = arg_types["b"]
 
     # check all args are tiles
     if not is_tile(a):
-        raise RuntimeError(f"tile_map() arguments must be tiles, got type {a}")
+        raise TypeError(f"tile_map() 'a' argument must be a tile, got {a!r}")
 
     if not is_tile(b):
-        raise RuntimeError(f"tile_map() arguments must be tiles, got type {b}")
+        raise TypeError(f"tile_map() 'b' argument must be a tile, got {b!r}")
 
-    # use first argument to define output type
+    # ensure types equal
     if not types_equal(a.dtype, b.dtype):
-        raise RuntimeError(f"tile_map() arguments must all have the same type {a.dtype} != {b.dtype}")
+        raise TypeError(f"tile_map() arguments must have the same dtype, got {a.dtype} and {b.dtype}")
 
-    if a.M != b.M:
-        raise RuntimeError(f"tile_map() arguments must all have the same m dimension {a.M} != {b.M}")
+    if len(a.shape) != len(b.shape):
+        raise ValueError(
+            f"tile_map() shapes must have the same number of dimensions, got {len(a.shape)} and {len(b.shape)}"
+        )
 
-    if a.N != b.N:
-        raise RuntimeError(f"tile_map() arguments must all have the same n dimension {a.N} != {b.N}")
+    for i in range(len(a.shape)):
+        if a.shape[i] != b.shape[i]:
+            raise ValueError(f"tile_map() shapes do not match on dimension {i}, got {a.shape[i]} and {b.shape[i]}")
 
     return TileBinaryMap(a, b)
 
 
 add_builtin(
     "tile_map",
-    input_types={"op": Callable, "a": Tile(dtype=Any, M=Any, N=Any), "b": Tile(dtype=Any, M=Any, N=Any)},
+    input_types={"op": Callable, "a": Tile(dtype=Any, shape=Any), "b": Tile(dtype=Any, shape=Any)},
     value_func=tile_binary_map_value_func,
     # dispatch_func=tile_map_dispatch_func,
     # variadic=True,
@@ -2948,7 +3051,7 @@ add_builtin(
         def compute():
 
             a = wp.tile_arange(0.0, 1.0, 0.1, dtype=float)
-            b = wp.tile_ones(m=1, n=10, dtype=float)
+            b = wp.tile_ones(shape=10, dtype=float)
 
             s = wp.tile_map(wp.add, a, b)
 
@@ -2960,7 +3063,7 @@ add_builtin(
 
     .. code-block:: text
 
-        tile(m=1, n=10, storage=register) = [[1 1.1 1.2 1.3 1.4 1.5 1.6 1.7 1.8 1.9]]""",
+        [1 1.1 1.2 1.3 1.4 1.5 1.6 1.7 1.8 1.9] = tile(shape=(10), storage=register)""",
     group="Tile Primitives",
     export=False,
 )
@@ -3075,6 +3178,18 @@ add_builtin(
 )
 
 
+def mlp_dispatch_func(input_types: Mapping[str, type], return_type: Any, args: Mapping[str, Var]):
+    warp.utils.warn(
+        "wp.mlp() is deprecated and will be removed in a future\nversion. Use tile primitives instead.",
+        category=DeprecationWarning,
+    )
+
+    func_args = tuple(args.values())
+    template_args = ()
+
+    return (func_args, template_args)
+
+
 add_builtin(
     "mlp",
     input_types={
@@ -3086,9 +3201,13 @@ add_builtin(
         "out": array(dtype=float, ndim=2),
     },
     value_type=None,
+    dispatch_func=mlp_dispatch_func,
     skip_replay=True,
     doc="""Evaluate a multi-layer perceptron (MLP) layer in the form: ``out = act(weights*x + bias)``.
 
+    .. deprecated:: 1.6
+        Use :doc:`tile primitives </modules/tiles>` instead.
+
     :param weights: A layer's network weights with dimensions ``(m, n)``.
     :param bias: An array with dimensions ``(n)``.
     :param activation: A ``wp.func`` function that takes a single scalar float as input and returns a scalar float as output
@@ -5008,6 +5127,43 @@ add_builtin(
 )
 
 
+# implements vector[idx] += scalar
+add_builtin(
+    "augassign_add",
+    input_types={"a": vector(length=Any, dtype=Scalar), "i": int, "value": Scalar},
+    value_type=None,
+    hidden=True,
+    group="Utility",
+)
+
+# implements quaternion[idx] += scalar
+add_builtin(
+    "augassign_add",
+    input_types={"a": quaternion(dtype=Scalar), "i": int, "value": Scalar},
+    value_type=None,
+    hidden=True,
+    group="Utility",
+)
+
+# implements vector[idx] -= scalar
+add_builtin(
+    "augassign_sub",
+    input_types={"a": vector(length=Any, dtype=Scalar), "i": int, "value": Scalar},
+    value_type=None,
+    hidden=True,
+    group="Utility",
+)
+
+# implements quaternion[idx] -= scalar
+add_builtin(
+    "augassign_sub",
+    input_types={"a": quaternion(dtype=Scalar), "i": int, "value": Scalar},
+    value_type=None,
+    hidden=True,
+    group="Utility",
+)
+
+
 def matrix_index_row_value_func(arg_types: Mapping[str, type], arg_values: Mapping[str, Any]):
     mat_type = arg_types["a"]
     row_type = mat_type._wp_row_type_
@@ -5057,22 +5213,42 @@ def matrix_vector_sametype(arg_types: Mapping[str, Any]):
     return mat_size == vec_size and mat_type == vec_type
 
 
-# implements matrix[i,j] = scalar
+# implements matrix[i,j] = scalar
+add_builtin(
+    "assign",
+    input_types={"a": matrix(shape=(Any, Any), dtype=Scalar), "i": int, "j": int, "value": Scalar},
+    value_func=matrix_assign_value_func,
+    hidden=True,
+    group="Utility",
+)
+
+
+# implements matrix[i] = vector
+add_builtin(
+    "assign",
+    input_types={"a": matrix(shape=(Any, Any), dtype=Scalar), "i": int, "value": vector(length=Any, dtype=Scalar)},
+    constraint=matrix_vector_sametype,
+    value_func=matrix_assign_value_func,
+    hidden=True,
+    group="Utility",
+)
+
+
+# implements matrix[i,j] += scalar
 add_builtin(
-    "assign",
+    "augassign_add",
     input_types={"a": matrix(shape=(Any, Any), dtype=Scalar), "i": int, "j": int, "value": Scalar},
-    value_func=matrix_assign_value_func,
+    value_type=None,
     hidden=True,
     group="Utility",
 )
 
 
-# implements matrix[i] = vector
+# implements matrix[i,j] -= scalar
 add_builtin(
-    "assign",
-    input_types={"a": matrix(shape=(Any, Any), dtype=Scalar), "i": int, "value": vector(length=Any, dtype=Scalar)},
-    constraint=matrix_vector_sametype,
-    value_func=matrix_assign_value_func,
+    "augassign_sub",
+    input_types={"a": matrix(shape=(Any, Any), dtype=Scalar), "i": int, "j": int, "value": Scalar},
+    value_type=None,
     hidden=True,
     group="Utility",
 )
@@ -5644,19 +5820,19 @@ add_builtin("unot", input_types={"a": array(dtype=Any)}, value_type=builtins.boo
 # Tile operators
 def tile_unary_value_func(arg_types, arg_values):
     if arg_types is None:
-        return Tile(dtype=Any, M=Any, N=Any)
+        return Tile(dtype=Any, shape=Any)
 
     t = arg_types["x"]
 
     if not is_tile(t):
-        raise RuntimeError("Expected tile for unary expression")
+        raise TypeError(f"Expected tile for unary expression, got {t}")
 
     return TileUnaryMap(t)
 
 
 def tile_scalar_mul_value_func(arg_types, arg_values):
     if arg_types is None:
-        return Tile(dtype=Any, M=Any, N=Any)
+        return Tile(dtype=Any, shape=Any)
 
     x = arg_types["x"]
     y = arg_types["y"]
@@ -5664,25 +5840,21 @@ def tile_scalar_mul_value_func(arg_types, arg_values):
     # tile*scalar
     if is_tile(x):
         if x.dtype != y:
-            raise RuntimeError(
-                "Scalar factor should have the same type as tile for tile*scalar, tile type: {x} scalar type: {y}"
-            )
+            raise TypeError(f"Scalar factor type {y} does not match tile type {x.dtype} for tile*scalar")
 
-        return TileBinaryMap(x, TileConstant(y, x.M, x.N))
+        return TileBinaryMap(x, TileConstant(y, x.shape))
 
     # scalar*tile
     if is_tile(y):
         if y.dtype != x:
-            raise RuntimeError(
-                "Scalar factor should have the same type as tile for scalar*tile, tile type: {x} scalar type: {y}"
-            )
+            raise TypeError(f"Scalar factor type {x} does not match tile type {y.dtype} for scalar*tile")
 
-        return TileBinaryMap(TileConstant(x, y.M, y.N), y)
+        return TileBinaryMap(TileConstant(x, y.shape), y)
 
 
 add_builtin(
     "neg",
-    input_types={"x": Tile(dtype=Any, M=Any, N=Any)},
+    input_types={"x": Tile(dtype=Any, shape=Any)},
     value_func=tile_unary_value_func,
     doc="Negate each element of a tile",
     export=False,
@@ -5692,7 +5864,7 @@ add_builtin(
 
 add_builtin(
     "add",
-    input_types={"a": Tile(dtype=Any, M=Any, N=Any), "b": Tile(dtype=Any, M=Any, N=Any)},
+    input_types={"a": Tile(dtype=Any, shape=Any), "b": Tile(dtype=Any, shape=Any)},
     value_func=tile_binary_map_value_func,
     # dispatch_func=tile_map_dispatch_func,
     # variadic=True,
@@ -5702,9 +5874,22 @@ add_builtin(
     export=False,
 )
 
+add_builtin(
+    "sub",
+    input_types={"a": Tile(dtype=Any, shape=Any), "b": Tile(dtype=Any, shape=Any)},
+    value_func=tile_binary_map_value_func,
+    # dispatch_func=tile_map_dispatch_func,
+    # variadic=True,
+    native_func="tile_sub",
+    doc="Subtract each element b from a",
+    group="Tile Primitives",
+    export=False,
+)
+
+
 add_builtin(
     "mul",
-    input_types={"x": Tile(dtype=Any, M=Any, N=Any), "y": Scalar},
+    input_types={"x": Tile(dtype=Any, shape=Any), "y": Scalar},
     value_func=tile_scalar_mul_value_func,
     doc="Multiply each element of a tile by a scalar",
     export=False,
@@ -5714,7 +5899,7 @@ add_builtin(
 
 add_builtin(
     "mul",
-    input_types={"x": Scalar, "y": Tile(dtype=Any, M=Any, N=Any)},
+    input_types={"x": Scalar, "y": Tile(dtype=Any, shape=Any)},
     value_func=tile_scalar_mul_value_func,
     doc="Multiply each element of a tile by a scalar",
     export=False,
@@ -5723,6 +5908,70 @@ add_builtin(
 )
 
 
+def tile_diag_add_value_func(arg_types, arg_values):
+    if arg_types is None:
+        return Tile(dtype=Any, shape=Any)
+
+    a = arg_types["a"]
+    d = arg_types["d"]
+
+    if not is_tile(a):
+        raise TypeError(f"tile_diag_add() 'a' argument must be a tile, got {a!r}")
+
+    if not is_tile(d):
+        raise TypeError(f"tile_diag_add() 'd' argument must be a tile, got {d!r}")
+
+    if not types_equal(a.dtype, d.dtype):
+        raise TypeError(f"tile_diag_add() arguments must have the same dtype, got {a.dtype} and {d.dtype}")
+
+    if len(a.shape) != 2:
+        raise TypeError("tile_diag_add() argument 'a' must be a 2D tile")
+
+    if len(d.shape) != 1:
+        raise TypeError("tile_diag_add() argument 'd' must be a 1D tile")
+
+    if a.shape[0] != a.shape[1]:
+        raise ValueError("tile_diag_add() 'a' argument must be square")
+
+    if a.shape[0] != d.shape[0]:
+        raise ValueError(
+            f"tile_diag_add() 'd' argument must have the same number of elements as the number of rows in 'a', "
+            f"got {d.shape[0]} elements in 'd' and {a.shape[0]} rows in 'a'"
+        )
+
+    # use first argument to define output type
+    return Tile(dtype=a.dtype, shape=a.shape, storage="shared")
+
+
+def tile_diag_add_lto_dispatch_func(
+    arg_types: Mapping[str, type],
+    return_type: Any,
+    return_values: List[Var],
+    arg_values: Mapping[str, Var],
+    options: Mapping[str, Any],
+    builder: warp.context.ModuleBuilder,
+):
+    a = arg_values["a"]
+    d = arg_values["d"]
+    # force the storage type of the input variables to shared memory
+    a.type.storage = "shared"
+    d.type.storage = "shared"
+    out = return_values[0]
+    return ((a, d, out), [], [], 0)
+
+
+add_builtin(
+    "tile_diag_add",
+    input_types={"a": Tile(dtype=Any, shape=Any), "d": Tile(dtype=Any, shape=Any)},
+    value_func=tile_diag_add_value_func,
+    lto_dispatch_func=tile_diag_add_lto_dispatch_func,
+    native_func="tile_diag_add",
+    doc="Add a square matrix and a diagonal matrix 'd' represented as a 1D tile",
+    group="Tile Primitives",
+    export=False,
+)
+
+
 ##
 ## MathDx, LTOIR-based, Tile functions
 ##
@@ -5734,24 +5983,25 @@ add_builtin(
 def tile_matmul_generic_value_func(arg_types, arg_values):
     # return generic type (for doc builds)
     if arg_types is None:
-        return Tile(dtype=Any, M=Any, N=Any)
+        return Tile(dtype=Any, shape=Any)
 
     a = arg_types["a"]
     b = arg_types["b"]
 
     if not is_tile(a):
-        raise RuntimeError("tile_matmul() argument 0 must be a tile")
+        raise TypeError(f"tile_matmul() 'a' argument must be a tile, got {a!r}")
+
     if not is_tile(b):
-        raise RuntimeError("tile_matmul() argument 1 must be an tile")
+        raise TypeError(f"tile_matmul() 'b' argument must be a tile, got {b!r}")
 
     # out = wp.tile_matmul(a, b)
     if len(arg_types) == 2:
-        return Tile(dtype=a.dtype, M=a.M, N=b.N, storage="shared")
+        return Tile(dtype=a.dtype, shape=(a.shape[0], b.shape[1]), storage="shared")
 
     # wp.tile_matmul(a, b, out)
     elif len(arg_types) == 3:
         if not is_tile(arg_types["out"]):
-            raise RuntimeError("tile_matmul() output argument must be a tile")
+            raise TypeError(f"tile_matmul() 'out' argument must be a tile, got {arg_types['out']!r}")
 
     return None
 
@@ -5774,16 +6024,20 @@ def tile_matmul_generic_lto_dispatch_func(
         accumulate = 1  # for tile_matmul(a,b,c) case we want to add to c value
         out = arg_values["out"]
 
-    if any(not is_tile(arg.type) for arg in [a, b, out]):
-        raise RuntimeError("tile_matmul() requires three Tile arguments")
+    if not is_tile(out.type):
+        raise TypeError(f"tile_matmul() 'out' argument must be a tile, got {out!r}")
 
     if any(arg.type.dtype not in [float16, float32, float64, vec2h, vec2f, vec2d] for arg in [a, b, out]):
-        raise RuntimeError(
+        raise TypeError(
             "tile_matmul() arguments must be tiles of float16, float32 or float64, vec2h, vec2f, vec2d entries"
         )
 
-    if (a.type.N != b.type.M) or (a.type.M != out.type.M) or (b.type.N != out.type.N):
-        raise RuntimeError("tile_matmul(A, B, C) requires sizes of A, B and C to be consistent for a matmul")
+    if (
+        (a.type.shape[1] != b.type.shape[0])
+        or (a.type.shape[0] != out.type.shape[0])
+        or (b.type.shape[1] != out.type.shape[1])
+    ):
+        raise ValueError("tile_matmul(A, B, C) requires sizes of A, B and C to be consistent for a matmul")
 
     # set the storage type to the inputs to shared
     a.type.storage = "shared"
@@ -5805,18 +6059,18 @@ def tile_matmul_generic_lto_dispatch_func(
             return ("wp::vec2f", 5, 1)
         if dtype == vec2d:
             return ("wp::vec2d", 6, 1)
-        raise RuntimeError("Unsupported input type in tile_matmul")
+        raise TypeError("Unsupported input type in tile_matmul")
 
     def cublasdx_arrangement_map(layout):
         if layout == "colmajor":
             return 0  # CUBLASDX_ARRANGEMENT_COL_MAJOR
         if layout == "rowmajor":
             return 1  # CUBLASDX_ARRANGEMENT_ROW_MAJOR
-        raise RuntimeError("Unsupported layout in tile_matmul")
+        raise ValueError("Unsupported layout in tile_matmul")
 
     # generate the LTO
-    M, K = a.type.M, a.type.N
-    _, N = b.type.M, b.type.N
+    M, K = a.type.shape[0], a.type.shape[1]
+    _, N = b.type.shape[0], b.type.shape[1]
     num_threads = options["block_dim"]
     arch = options["output_arch"]
 
@@ -5829,7 +6083,8 @@ def tile_matmul_generic_lto_dispatch_func(
         c_arrangement = cublasdx_arrangement_map(clayout)
 
         if a_type != b_type or a_type != c_type:
-            raise RuntimeError("time_matmul(A, B, C) requires all inputs to be real or complex")
+            raise TypeError("time_matmul(A, B, C) requires all inputs to be real or complex")
+
         element_type = a_type
 
         lto_symbol = f"dot_{M}_{N}_{K}_{arch}_{num_threads}_{a_arrangement}_{b_arrangement}_{c_arrangement}_{a_prec}_{b_prec}_{c_prec}_{element_type}"
@@ -5924,15 +6179,16 @@ def tile_matmul_generic_lto_dispatch_func(
         ),
         template_args,
         [lto_forward, lto_backward_A, lto_backward_B],
+        0,
     )
 
 
 add_builtin(
     "tile_matmul",
     input_types={
-        "a": Tile(dtype=Any, M=Any, N=Any),
-        "b": Tile(dtype=Any, M=Any, N=Any),
-        "out": Tile(dtype=Any, M=Any, N=Any),
+        "a": Tile(dtype=Any, shape=Any),
+        "b": Tile(dtype=Any, shape=Any),
+        "out": Tile(dtype=Any, shape=Any),
     },
     value_func=tile_matmul_generic_value_func,
     lto_dispatch_func=tile_matmul_generic_lto_dispatch_func,
@@ -5956,7 +6212,7 @@ add_builtin(
 
 add_builtin(
     "tile_matmul",
-    input_types={"a": Tile(dtype=Any, M=Any, N=Any), "b": Tile(dtype=Any, M=Any, N=Any)},
+    input_types={"a": Tile(dtype=Any, shape=Any), "b": Tile(dtype=Any, shape=Any)},
     value_func=tile_matmul_generic_value_func,
     lto_dispatch_func=tile_matmul_generic_lto_dispatch_func,
     variadic=False,
@@ -5983,16 +6239,23 @@ add_builtin(
 ##
 def tile_fft_generic_value_func(arg_types, arg_values):
     if arg_types is None:
-        return Tile(dtype=Any, M=Any, N=Any)
+        return Tile(dtype=Any, shape=Any)
 
     if len(arg_types) != 1:
-        raise RuntimeError("tile_fft() requires 1 positional args")
+        raise TypeError(f"tile_fft() takes exactly 1 positional argument but {len(arg_types)} were given")
 
-    if not is_tile(arg_types["inout"]):
-        raise RuntimeError("tile_fft() argument 0 must be a tile")
+    inout = arg_types["inout"]
 
-    if arg_types["inout"].storage != "register":
-        raise RuntimeError("tile_fft() input/output argument must have register memory storage")
+    if not is_tile(inout):
+        raise TypeError(f"tile_fft() argument must be a tile, got {inout!r}")
+
+    if inout.storage != "register":
+        raise ValueError(f"tile_fft() argument must have 'register' storage, got {inout.storage}")
+
+    if inout.dtype not in [vec2f, vec2d]:
+        raise TypeError(
+            f"tile_fft() argument must be a tile of vec2f or vec2d (interpreted as complex) entries, got {inout.dtype!r}"
+        )
 
     return None
 
@@ -6009,19 +6272,13 @@ def tile_fft_generic_lto_dispatch_func(
     inout = arg_values["inout"]
     inout.type.storage = "register"
 
-    if not is_tile(inout.type):
-        raise RuntimeError("tile_fft() arguments must be a single tile with register storage")
-
-    if inout.type.dtype not in [vec2f, vec2d]:
-        raise RuntimeError("tile_fft() argument must be a tile of vec2f or vec2d (interpreted as complex) entries")
-
     # see libcufftdx.hpp
     if direction == "forward":
         dir = 0  # CUFFTDX_DIRECTION_FORWARD
     elif direction == "inverse":
         dir = 1  # CUFFTDX_DIRECTION_INVERSE
     else:
-        raise RuntimeError("Invalid direction")
+        raise ValueError(f"Invalid direction: {direction!r}.  Expected 'forward' or 'inverse'.")
 
     if inout.type.dtype == vec2f:
         dtype = "wp::vec2f"
@@ -6030,10 +6287,10 @@ def tile_fft_generic_lto_dispatch_func(
         dtype = "wp::vec2d"
         precision = 6  # COMMONDX_PRECISION_F64
     else:
-        raise RuntimeError("Unsupported datatype")
+        raise TypeError(f"Unsupported data type, got {dtype!r}")
 
     # M FFTs of size N each
-    batch, size = inout.type.M, inout.type.N
+    batch, size = inout.type.shape[0], inout.type.shape[1]
     num_threads = options["block_dim"]
     arch = options["output_arch"]
     ept = size // num_threads
@@ -6065,7 +6322,7 @@ def tile_fft_generic_lto_dispatch_func(
         lto_code.close()
         if lto_code_path.exists():
             lto_code_path.unlink()
-        raise RuntimeError("Failed to compile tile_matmul")
+        raise RuntimeError("Failed to compile tile_fft")
 
     with open(lto_code.name, "rb") as f:
         lto_code_data = f.read()
@@ -6075,17 +6332,20 @@ def tile_fft_generic_lto_dispatch_func(
 
     builder.ltoirs[lto_symbol] = lto_code_data
 
+    shared_memory_bytes = Tile.round_up(shared_memory_size.value)
+
     return (
         (
             Var(lto_symbol, str, False, True, False),
             Var(dtype, str, False, True, False),
-            Var(str(shared_memory_size.value), str, False, True, False),
+            Var(str(shared_memory_bytes), str, False, True, False),
             Var(str(batch), str, False, True, False),
             Var(str(ept), str, False, True, False),
             inout,
         ),
         [],
         [lto_code_data],
+        shared_memory_bytes,
     )
 
 
@@ -6099,6 +6359,8 @@ add_builtin(
 
     This function cooperatively computes the forward FFT on a tile of data inplace, treating each row individually.
 
+    Note that computing the adjoint is not yet supported.
+
     Supported datatypes are:
         * vec2f, vec2d
 
@@ -6118,6 +6380,8 @@ add_builtin(
 
     This function cooperatively computes the inverse FFT on a tile of data inplace, treating each row individually.
 
+    Note that computing the adjoint is not yet supported.
+
     Supported datatypes are:
         * vec2f, vec2d
 
@@ -6127,6 +6391,283 @@ add_builtin(
     namespace="",
 )
 
+
+##
+## Cholesky
+##
+def tile_cholesky_generic_value_func(arg_types, arg_values):
+    if arg_types is None:
+        return Tile(dtype=Any, shape=Any)
+
+    if len(arg_types) != 1:
+        raise TypeError("tile_cholesky() requires 1 positional args")
+
+    a = arg_types["A"]
+
+    if not is_tile(a):
+        raise TypeError(f"tile_cholesky() argument must be a tile, got {a!r}")
+
+    if len(a.shape) != 2:
+        raise ValueError("tile_cholesky() argumust must be a 2D tile")
+
+    if a.shape[0] != a.shape[1]:
+        raise ValueError("tile_cholesky() argument must be square")
+
+    return Tile(dtype=a.dtype, shape=a.shape, storage="shared")
+
+
+cusolver_function_map = {"getrf": 0, "getrf_no_pivot": 1, "potrf": 2, "potrs": 3}
+
+cusolver_type_map = {float32: ("wp::float32", 5), float64: ("wp::float64", 6)}
+
+cusolver_fill_mode_map = {"upper": 0, "lower": 1}
+
+
+def tile_cholesky_generic_lto_dispatch_func(
+    arg_types: Mapping[str, type],
+    return_type: Any,
+    return_values: List[Var],
+    arg_values: Mapping[str, Var],
+    options: Mapping[str, Any],
+    builder: warp.context.ModuleBuilder,
+):
+    a = arg_values["A"]
+    # force source tile to shared memory
+    a.type.storage = "shared"
+
+    if a.type.dtype not in cusolver_type_map.keys():
+        raise TypeError("tile_cholesky() argument must be a tile of float32 or float64 entries")
+
+    if len(return_values) != 1:
+        raise TypeError("tile_cholesky() returns one output")
+    out = return_values[0]
+
+    dtype, precision_enum = cusolver_type_map[a.type.dtype]
+
+    # We already ensured a is square in tile_cholesky_generic_value_func()
+    M, N = a.type.shape[0], a.type.shape[1]
+    if out.type.shape[0] != M or out.type.shape[1] != M:
+        raise ValueError("tile_cholesky() output tile must be square")
+
+    num_threads = options["block_dim"]
+    arch = options["output_arch"]
+    lto_symbol = f"potrf_{M}_{N}_{arch}_{precision_enum}"
+
+    # early out if LTO for this combination already exists for this module
+    if lto_symbol in builder.ltoirs:
+        return lto_symbol, builder.ltoirs[lto_symbol]
+
+    # otherwise compile LTO
+    lto_code = tempfile.NamedTemporaryFile(prefix="warp", delete=False)
+    universal_fatbin_code = tempfile.NamedTemporaryFile(prefix="warp", delete=False)
+
+    # cuSOLVERDx only support col-major input/outputs,
+    # so we use upper to mimic a row-major input
+    result = warp.context.runtime.core.cuda_compile_solver(
+        universal_fatbin_code.name.encode("utf-8"),
+        lto_code.name.encode("utf-8"),
+        lto_symbol.encode("utf-8"),
+        0,
+        None,
+        None,
+        arch,
+        M,
+        N,
+        cusolver_function_map["potrf"],
+        precision_enum,
+        cusolver_fill_mode_map["upper"],
+        num_threads,
+    )
+
+    if not result:
+        for f in [lto_code, universal_fatbin_code]:
+            f.close()
+            if Path(f.name).exists():
+                Path(f.name).unlink()
+        raise RuntimeError("Failed to compile tile_cholesky")
+
+    else:
+        with open(lto_code.name, "rb") as f:
+            lto_code_data = f.read()
+        with open(universal_fatbin_code.name, "rb") as f:
+            universal_fatbin_code_data = f.read()
+        for f in [lto_code, universal_fatbin_code]:
+            f.close()
+            Path(f.name).unlink()
+
+    builder.ltoirs[lto_symbol] = lto_code_data
+    builder.ltoirs_decl[lto_symbol] = f"void {lto_symbol}({dtype}*, unsigned);"
+    builder.fatbins["cholesky"] = universal_fatbin_code_data
+
+    return ((Var(lto_symbol, str, False, True, False), a, out), [], [lto_code_data], 0)
+
+
+add_builtin(
+    "tile_cholesky",
+    input_types={"A": Tile},
+    value_func=tile_cholesky_generic_value_func,
+    lto_dispatch_func=tile_cholesky_generic_lto_dispatch_func,
+    variadic=True,
+    doc="""Compute the Cholesky factorization L of a matrix A.
+    L is lower triangular and satisfies LL^T = A.
+
+    Note that computing the adjoint is not yet supported.
+
+    Supported datatypes are:
+        * float32
+        * float64
+
+    :param A: A square, symmetric positive-definite, matrix.
+    :returns L: A square, lower triangular, matrix, such that LL^T = A""",
+    group="Tile Primitives",
+    export=False,
+    namespace="",
+)
+
+
+def tile_cholesky_solve_generic_value_func(arg_types, arg_values):
+    if arg_types is None:
+        return None
+
+    if len(arg_types) != 2:
+        raise TypeError("tile_cholesky_solve() requires exactly 2 positional args")
+
+    l = arg_types["L"]
+    x = arg_types["x"]
+
+    if not is_tile(l):
+        raise TypeError(f"tile_cholesky_solve() 'L' argument must be a tile, got {l!r}")
+
+    if not is_tile(x):
+        raise TypeError(f"tile_cholesky_solve() 'x' argument must be a tile, got {l!r}")
+
+    if not types_equal(l.dtype, x.dtype):
+        raise TypeError(f"tile_cholesky_solve() arguments must have the same dtype, got {l.dtype} and {x.dtype}")
+
+    if l.shape[0] != l.shape[1]:
+        raise ValueError("tile_cholesky_solve() 'L' argument must be square")
+
+    if len(x.shape) != 1:
+        raise TypeError("tile_cholesky_solve() 'x' argument must be a 1D tile")
+
+    if x.shape[0] != l.shape[0]:
+        raise ValueError(
+            f"tile_cholesky_solve() 'x' argument must have the same number of elements as the number of rows in 'L', "
+            f"got {x.shape[0]} elements in 'x' and {l.shape[0]} rows in 'L'"
+        )
+
+    return Tile(dtype=l.dtype, shape=x.shape, storage="shared")
+
+
+def tile_cholesky_solve_generic_lto_dispatch_func(
+    arg_types: Mapping[str, type],
+    return_type: Any,
+    return_values: List[Var],
+    arg_values: Mapping[str, Var],
+    options: Mapping[str, Any],
+    builder: warp.context.ModuleBuilder,
+):
+    L = arg_values["L"]
+    x = arg_values["x"]
+    # force the storage type of the input variables to shared memory
+    L.type.storage = "shared"
+    x.type.storage = "shared"
+
+    if len(return_values) != 1:
+        raise TypeError(f"tile_cholesky_solve() must return exactly one value, got {len(return_values)}")
+
+    y = return_values[0]
+
+    if any(T not in cusolver_type_map.keys() for T in [x.type.dtype, L.type.dtype]):
+        raise TypeError("tile_cholesky_solve() arguments be tiles of float64 or float32")
+
+    dtype, precision_enum = cusolver_type_map[L.type.dtype]
+    M, N = L.type.shape[0], L.type.shape[1]
+
+    if len(y.type.shape) != 1:
+        raise TypeError("tile_cholesky_solve() output vector must be 1D")
+
+    if y.type.shape[0] != M:
+        raise ValueError(
+            "tile_cholesky_solve() output vector must have same number of elements as the number of rows in 'L' "
+            f"got {y.type.shape[0]} elements in output and {M} rows in 'L'"
+        )
+
+    num_threads = options["block_dim"]
+    arch = options["output_arch"]
+    lto_symbol = f"potrs_{M}_{N}_{arch}_{precision_enum}"
+
+    # early out if LTO for this combination already exists for this module
+    if lto_symbol in builder.ltoirs:
+        return lto_symbol, builder.ltoirs[lto_symbol]
+
+    # otherwise compile LTO
+    lto_code = tempfile.NamedTemporaryFile(prefix="warp", delete=False)
+    universal_fatbin_code = tempfile.NamedTemporaryFile(prefix="warp", delete=False)
+
+    # cuSOLVERDx only support col-major input/outputs,
+    # so we use upper to mimic a row-major input
+    result = warp.context.runtime.core.cuda_compile_solver(
+        universal_fatbin_code.name.encode("utf-8"),
+        lto_code.name.encode("utf-8"),
+        lto_symbol.encode("utf-8"),
+        0,
+        None,
+        None,
+        arch,
+        M,
+        N,
+        cusolver_function_map["potrs"],
+        precision_enum,
+        cusolver_fill_mode_map["upper"],
+        num_threads,
+    )
+
+    if not result:
+        for f in [lto_code, universal_fatbin_code]:
+            f.close()
+            if Path(f.name).exists():
+                Path(f.name).unlink()
+        raise RuntimeError("Failed to compile tile_cholesky_solve")
+
+    else:
+        with open(lto_code.name, "rb") as f:
+            lto_code_data = f.read()
+        with open(universal_fatbin_code.name, "rb") as f:
+            universal_fatbin_code_data = f.read()
+        for f in [lto_code, universal_fatbin_code]:
+            f.close()
+            Path(f.name).unlink()
+
+    builder.ltoirs[lto_symbol] = lto_code_data
+    builder.ltoirs_decl[lto_symbol] = f"void {lto_symbol}({dtype}*, {dtype}*);"
+    builder.fatbins["cholesky"] = universal_fatbin_code_data
+
+    return ((Var(lto_symbol, str, False, True, False), L, x, y), [], [lto_code_data], 0)
+
+
+add_builtin(
+    "tile_cholesky_solve",
+    input_types={"L": Tile, "x": Tile},
+    value_func=tile_cholesky_solve_generic_value_func,
+    lto_dispatch_func=tile_cholesky_solve_generic_lto_dispatch_func,
+    variadic=True,
+    doc="""With L such that LL^T = A, solve for x in Ax = y
+
+    Note that computing the adjoint is not yet supported.
+
+    Supported datatypes are:
+        * float32
+        * float64
+
+    :param L: A square, lower triangular, matrix, such that LL^T = A
+    :param x: An 1D tile of length M
+    :returns y: An 1D tile of length M such that LL^T y = x""",
+    group="Tile Primitives",
+    export=False,
+    namespace="",
+)
+
 # ---------------------------------
 # Code Generation
 
@@ -6134,7 +6675,7 @@ add_builtin(
     "static",
     input_types={"expr": Any},
     value_type=Any,
-    doc="""Evaluates a static Python expression and replaces it with its result.
+    doc="""Evaluate a static Python expression and replaces it with its result.
 
     See the :ref:`code generation guide <static_expressions>` for more details.
 
@@ -6158,3 +6699,58 @@ def static(expr):
         which includes constant variables and variables captured in the current closure in which the function or kernel is implemented.
     """
     return expr
+
+
+add_builtin(
+    "len",
+    input_types={"a": vector(length=Any, dtype=Scalar)},
+    value_type=int,
+    doc="Return the number of elements in a vector.",
+    group="Utility",
+    export=False,
+)
+
+add_builtin(
+    "len",
+    input_types={"a": quaternion(dtype=Scalar)},
+    value_type=int,
+    doc="Return the number of elements in a quaternion.",
+    group="Utility",
+    export=False,
+)
+
+add_builtin(
+    "len",
+    input_types={"a": matrix(shape=(Any, Any), dtype=Scalar)},
+    value_type=int,
+    doc="Return the number of rows in a matrix.",
+    group="Utility",
+    export=False,
+)
+
+add_builtin(
+    "len",
+    input_types={"a": transformation(dtype=Float)},
+    value_type=int,
+    doc="Return the number of elements in a transformation.",
+    group="Utility",
+    export=False,
+)
+
+add_builtin(
+    "len",
+    input_types={"a": array(dtype=Any)},
+    value_type=int,
+    doc="Return the size of the first dimension in an array.",
+    group="Utility",
+    export=False,
+)
+
+add_builtin(
+    "len",
+    input_types={"a": Tile(dtype=Any, shape=Any)},
+    value_type=int,
+    doc="Return the number of rows in a tile.",
+    group="Utility",
+    export=False,
+)