refactor partition view

Signed-off-by: Boyan Li <boyanl@nvidia.com>

Author: blinxt

src/cuda/tile/_compile.py

@@ -38,6 +38,7 @@
 from cuda.tile._ir.typing_support import typeof_pyval, get_constant_value
 from cuda.tile._passes.ast2hir import get_function_hir
 from cuda.tile._passes.code_motion import hoist_loop_invariants
+from cuda.tile._passes.unhoist_partition_views import unhoist_partition_views
 from cuda.tile._passes.eliminate_assign_ops import eliminate_assign_ops
 from cuda.tile._passes.hir2ir import hir2ir
 from cuda.tile._passes.loop_split import split_loops
@@ -103,6 +104,11 @@ def _get_final_ir(pyfunc,
     # Otherwise, it may incorrectly hoist load operations out of the loop.
     hoist_loop_invariants(func_body)
 
+    # For version < V_13_3, MakePartitionView must be emitted inline before its consumer.
+    # Code motion may hoist it to an outer block; copy it back where needed.
+    if tileiras_version < BytecodeVersion.V_13_3:
+        unhoist_partition_views(func_body)
+
     split_loops(func_body)
     dead_code_elimination_pass(func_body)
     return ir.Function(func_body, func_hir.desc.name, func_hir.body.loc)

src/cuda/tile/_ir/ops.py

@@ -53,7 +53,7 @@
     BYTE_BITWIDTH, typeof_pyval, dtype_registry, loose_type_of_pyval, get_constant_value
 )
 from .type import (
-    TupleTy, TileTy, NoneType, BoundMethodTy, ArrayTy,
+    PartitionViewTy, TupleTy, TileTy, NoneType, BoundMethodTy, ArrayTy,
     ListTy, make_tile_ty, SliceType, DTypeConstructor, RangeIterType, Type,
     NONE, ModuleTy, TypeTy, LooselyTypedScalar, DTypeSpec, StringTy, InvalidType,
     array_size_type, ClosureTy, LiveCapturedScope, TokenTy,
@@ -2012,26 +2012,43 @@ def _check_load_store_hints(latency_value: int | None, allow_tma_value: bool | N
             raise TileTypeError(f"Allow TMA must be a boolean, got {allow_tma_value}")
 
 
+@dataclass(eq=False)
+class MakePartitionView(Operation, opcode="make_partition_view"):
+    array: Var = operand()
+
+    @override
+    def generate_bytecode(self, ctx: BytecodeContext) -> bc.Value:
+        partition_view_ty = self.result_var.get_type()
+        return bc.encode_MakePartitionViewOp(ctx.builder,
+                                             typeid(ctx.type_table, partition_view_ty),
+                                             ctx.get_value(self.array))
+
+
+def make_partition_view(array: Var, tile_shape: Sequence[int],
+                        order: Sequence[int],
+                        padding_mode: PaddingMode) -> Var:
+    array_ty = array.get_type()
+    assert isinstance(array_ty, ArrayTy)
+    view_ty = PartitionViewTy(array_ty, tuple(tile_shape), tuple(order), padding_mode)
+    return add_operation(MakePartitionView, view_ty, array=array)
+
+
 @dataclass(eq=False)
 class TileLoad(Operation, opcode="tile_load", memory_effect=MemoryEffect.LOAD):
-    order: tuple[int, ...] = attribute()
-    padding_mode: PaddingMode = attribute()
     latency: Optional[int] = attribute()
     allow_tma: Optional[bool] = attribute()
-    array: Var = operand()
+    view: Var = operand()
     index: tuple[Var, ...] = operand()
     token: Optional[Var] = operand(default=None)
 
     @override
     def generate_bytecode(self, ctx: BytecodeContext) -> tuple[bc.Value, bc.Value]:
         tile_type: TileTy = self.result_vars[0].get_type()
-        shape = tile_type.shape
-        partition = ctx.make_partition_view(self.array, self.order, shape, self.padding_mode)
         res, res_token = bc.encode_LoadViewTkoOp(
             ctx.builder,
             tile_type=typeid(ctx.type_table, tile_type),
             result_token_type=ctx.type_table.Token,
-            view=partition,
+            view=ctx.get_value(self.view),
             index=ctx.index_tuple(self.index),
             token=None if self.token is None else ctx.get_value(self.token),
             memory_ordering_semantics=bc.MemoryOrderingSemantics.WEAK,
@@ -2041,16 +2058,6 @@ def generate_bytecode(self, ctx: BytecodeContext) -> tuple[bc.Value, bc.Value]:
         return res, res_token
 
 
-def tile_load(array: Var, index: tuple[Var, ...], shape: Sequence[int], order: Sequence[int],
-              padding_mode: PaddingMode, latency: Optional[int],
-              allow_tma: Optional[bool]) -> tuple[Var, Var]:
-    res_ty = make_tile_ty(array.get_type().dtype, shape)
-    return add_operation(TileLoad, (res_ty, TokenTy()),
-                         array=array, index=index, order=tuple(order),
-                         padding_mode=padding_mode, latency=latency,
-                         allow_tma=allow_tma)
-
-
 @impl(ct.load)
 def tile_load_impl(array: Var, index: Var, shape: Var, order: Var,
                    padding_mode: Var, latency: Var, allow_tma: Var) -> Var:
@@ -2070,32 +2077,31 @@ def tile_load_impl(array: Var, index: Var, shape: Var, order: Var,
     latency = require_optional_constant_int(latency)
     allow_tma = require_optional_constant_bool(allow_tma)
     _check_load_store_hints(latency, allow_tma)
-    result, _token = tile_load(array, index_items, broadcasted_shape, order, padding_mode, latency,
-                               allow_tma)
+
+    view = make_partition_view(array, broadcasted_shape, order, padding_mode)
+    res_ty = make_tile_ty(array_ty.dtype, broadcasted_shape)
+    result, _token = add_operation(TileLoad, (res_ty, TokenTy()),
+                                   view=view, index=index_items, latency=latency,
+                                   allow_tma=allow_tma)
     return reshape(result, shape)
 
 
 @dataclass(eq=False)
 class TileStore(Operation, opcode="tile_store", memory_effect=MemoryEffect.STORE):
-    order: tuple[int, ...] = attribute()
     latency: Optional[int] = attribute()
     allow_tma: Optional[bool] = attribute()
-    array: Var = operand()
+    view: Var = operand()
     index: tuple[Var, ...] = operand()
     tile: Var = operand()
     token: Optional[Var] = operand(default=None)
 
     @override
     def generate_bytecode(self, ctx: BytecodeContext) -> bc.Value:
-        tile_ty = self.tile.get_type()
-        tile_shape = tile_ty.shape
-        partition = ctx.make_partition_view(self.array, self.order, tile_shape,
-                                            padding_mode=PaddingMode.UNDETERMINED)
         return bc.encode_StoreViewTkoOp(
             ctx.builder,
             result_token_type=ctx.type_table.Token,
             tile=ctx.get_value(self.tile),
-            view=partition,
+            view=ctx.get_value(self.view),
             index=ctx.index_tuple(self.index),
             token=None if self.token is None else ctx.get_value(self.token),
             memory_ordering_semantics=bc.MemoryOrderingSemantics.WEAK,
@@ -2104,17 +2110,6 @@ def generate_bytecode(self, ctx: BytecodeContext) -> bc.Value:
         )
 
 
-def tile_store(array: Var, index: tuple[Var, ...], tile: Var, order: Sequence[int],
-               latency: Optional[int], allow_tma: Optional[bool]) -> Var:
-    array_ty = array.get_type()
-    ty = require_tile_type(tile)
-    tile = astype(tile, array_ty.dtype)
-    if ty.ndim == 0:
-        tile = reshape(tile, (1,) * array_ty.ndim)
-    return add_operation(TileStore, (TokenTy(),), array=array, index=index, tile=tile,
-                         order=tuple(order), latency=latency, allow_tma=allow_tma)
-
-
 def _implicit_cast(src: Var, target_dtype: DType, error_context: str) -> Var:
     ty = require_tile_maybe_loose_type(src)
     try:
@@ -2130,20 +2125,25 @@ def _implicit_cast(src: Var, target_dtype: DType, error_context: str) -> Var:
 def tile_store_impl(array: Var, index: Var, tile: Var, order: Var,
                     latency: Var, allow_tma: Var):
     array_ty = require_array_type(array)
+    tile_ty = require_tile_type(tile)
     index_ty = require_index_or_index_tuple_type(index)
     index_items = index.get_aggregate().items if isinstance(index_ty, TupleTy) else (index,)
     if array_ty.ndim != len(index_items):
         raise TileTypeError(f"Index size {len(index_items)}"
                             f" does not match the array rank {array_ty.ndim}")
 
-    tile = _implicit_cast(tile, array_ty.dtype, "Stored tile is incompatible with array's dtype")
-
+    shape = tile_ty.shape
+    broadcasted_shape = (1,) * array_ty.ndim if len(shape) == 0 else shape
     order = require_constant_axis_order(order, array_ty.ndim)
     latency = require_optional_constant_int(latency)
     allow_tma = require_optional_constant_bool(allow_tma)
     _check_load_store_hints(latency, allow_tma)
 
-    [_token] = tile_store(array, index_items, tile, order, latency, allow_tma)
+    tile = _implicit_cast(tile, array_ty.dtype, "Stored tile is incompatible with array's dtype")
+    tile = reshape(tile, broadcasted_shape)
+    view = make_partition_view(array, broadcasted_shape, order, PaddingMode.UNDETERMINED)
+    [_token] = add_operation(TileStore, (TokenTy(),), view=view, index=index_items, tile=tile,
+                             latency=latency, allow_tma=allow_tma)
 
 
 @dataclass(eq=False)
@@ -2618,35 +2618,30 @@ def join_tokens(tokens: Tuple[Var, ...], *, block: Block, res: Var, loc: Loc) ->
 @dataclass(eq=False)
 class NumTiles(Operation, opcode="num_tiles"):
     axis: int = attribute()
-    shape: tuple[int, ...] = attribute()
-    order: tuple[int, ...] = attribute()
-    array: Var = operand()
+    view: Var = operand()
 
     @override
     def generate_bytecode(self, ctx: BytecodeContext):
-        pv = ctx.make_partition_view(self.array, self.order, self.shape,
-                                     padding_mode=PaddingMode.UNDETERMINED)
-        result_types = [ctx.type_table.tile(ctx.type_table.I32, ())] * len(self.shape)
-        values = bc.encode_GetIndexSpaceShapeOp(ctx.builder, result_types, pv)
+        view_ty: PartitionViewTy = self.view.get_type()
+        result_types = [ctx.type_table.tile(ctx.type_table.I32, ())] * len(view_ty.tile_shape)
+        values = bc.encode_GetIndexSpaceShapeOp(ctx.builder, result_types,
+                                                src=ctx.get_value(self.view))
         return values[self.axis]
 
 
-def num_tiles(array: Var, axis: int, shape: Sequence[int], order: Sequence[int]) -> Var:
-    return add_operation(NumTiles, make_tile_ty(datatype.default_int_type, ()), array=array,
-                         axis=axis, shape=tuple(shape), order=tuple(order))
-
-
 @impl(ct.num_tiles)
 def num_tiles_impl(array: Var, axis: Var, shape: Var, order: Var) -> Var:
     array_ty = require_array_type(array)
     axis = require_constant_int(axis)
     axis = normalize_axis(axis, array_ty.ndim)
     shape = require_constant_shape(shape, allow_single_int=True, expected_rank=array_ty.ndim,
                                    allow_0d_shape=True)
-    if len(shape) == 0:
-        shape = (1,) * array_ty.ndim
+    broadcasted_shape = (1,) * array_ty.ndim if len(shape) == 0 else shape
     order = require_constant_axis_order(order, array_ty.ndim)
-    return num_tiles(array, axis, shape, order)
+
+    view = make_partition_view(array, broadcasted_shape, order, PaddingMode.UNDETERMINED)
+    return add_operation(NumTiles, make_tile_ty(datatype.default_int_type, ()), view=view,
+                         axis=axis)
 
 
 def full_const(shape: Sequence[int], fill_value: int | float, dtype: DType) -> Var:

src/cuda/tile/_ir/type.py

@@ -13,6 +13,7 @@
 from typing import TYPE_CHECKING
 
 from cuda.tile._exception import Loc
+from cuda.tile._numeric_semantics import PaddingMode
 
 if TYPE_CHECKING:
     from cuda.tile._datatype import DType
@@ -358,6 +359,25 @@ def __str__(self):
         return f"Array[{self.dtype},{shape_str}:{strides_str}]"
 
 
+# ============== PartitionView Type ===============
+
+
+@dataclass(frozen=True)
+class PartitionViewTy(Type):
+    array_ty: ArrayTy
+    tile_shape: tuple[int, ...]
+    order: tuple[int, ...]
+    padding_mode: PaddingMode
+
+    @property
+    def dtype(self):
+        return self.array_ty.dtype
+
+    def __str__(self):
+        return (f"PartitionView[{self.array_ty},tile_shape={self.tile_shape},order={self.order},"
+                f"padding_mode={self.padding_mode}]")
+
+
 # ============== List Type ===============
 
 

src/cuda/tile/_ir2bytecode.py

@@ -5,12 +5,12 @@
 import functools
 import os
 from contextlib import contextmanager
-from typing import Dict, Tuple, Sequence, Any, Optional
+from typing import Dict, Tuple, Any, Optional
 
 from cuda.tile import _datatype as datatype
 from cuda.tile._bytecode.attribute import make_load_store_hints
 from cuda.tile._datatype import get_signedness
-from cuda.tile import DType, PaddingMode
+from cuda.tile import DType
 import cuda.tile._bytecode as bc
 from cuda.tile._compiler_options import CompilerOptions
 from cuda.tile._exception import TileInternalError, TileError, FunctionDesc
@@ -19,7 +19,10 @@
     padding_mode_to_bytecode, rounding_mode_to_bytecode,
     get_default_rounding_mode,
 )
-from cuda.tile._ir.type import Type, TileTy, PointerTy, TokenTy, TupleTy, ArrayTy, size_to_bytecode
+from cuda.tile._ir.type import (
+    PartitionViewTy, Type, TileTy, PointerTy, TokenTy, TupleTy, ArrayTy,
+    size_to_bytecode,
+)
 
 
 def dtype_typeid(tt: bc.TypeTable, dtype: datatype.DType | PointerTy) -> bc.TypeId:
@@ -49,6 +52,11 @@ def typeid(tt: bc.TypeTable, ty: Type) -> bc.TypeId:
         return tt.tile(dtype, shape)
     elif isinstance(ty, TokenTy):
         return tt.Token
+    elif isinstance(ty, PartitionViewTy):
+        padding_value = padding_mode_to_bytecode[ty.padding_mode]
+        assert isinstance(ty.array_ty, ArrayTy)
+        tv_id = tensor_view_typeid(tt, ty.array_ty)
+        return tt.partition_view(ty.tile_shape, tv_id, ty.order, padding_value)
     else:
         raise NotImplementedError(f"Lowering type '{ty}' is not supported")
 
@@ -387,20 +395,6 @@ def load_store_hints(self,
         load_store_hints = bc.LoadStoreHints(latency=latency, allow_tma=allow_tma)
         return make_load_store_hints({self.sm_arch: load_store_hints})
 
-    def make_partition_view(self,
-                            array: Var,
-                            order: Sequence[int],
-                            tile_shape: Sequence[int],
-                            padding_mode: PaddingMode) -> bc.Value:
-        padding_value = padding_mode_to_bytecode[padding_mode]
-        array_ty = self.typeof(array)
-        assert isinstance(array_ty, ArrayTy)
-        view_ty_id = tensor_view_typeid(self.type_table, array_ty)
-        partition_ty_id = self.type_table.partition_view(
-                tile_shape, view_ty_id, order, padding_value)
-        view = self.get_value(array)
-        return bc.encode_MakePartitionViewOp(self.builder, partition_ty_id, view)
-
 
 def generate_bytecode_for_block(ctx: BytecodeContext, block: Block):
     for op in block.operations:

src/cuda/tile/_passes/alias_analysis.py

@@ -7,7 +7,7 @@
 
 from cuda.tile._ir.ir import Var, Block
 from cuda.tile._ir.ops import Assign, GetArrayListItem, \
-    Loop, IfElse, Continue, Break, EndBranch, PointerOffset, \
+    Loop, IfElse, Continue, Break, EndBranch, MakePartitionView, PointerOffset, \
     TileBroadcast, TileReshape, MakeTensorView, MakeListView, AssumeDivBy, TileReduce, TileScan
 
 
@@ -108,6 +108,8 @@ def _analyze_aliases_in_block(block: Block,
                 alias_tracker[v.name] = ALIAS_UNIVERSE
         elif isinstance(op, MakeTensorView):
             _propagate(alias_tracker, op.base_ptr, op.result_var)
+        elif isinstance(op, MakePartitionView):
+            _propagate(alias_tracker, op.array, op.result_var)
         elif isinstance(op, MakeListView):
             _propagate(alias_tracker, op.base_ptr, op.result_var)
         elif isinstance(op, PointerOffset):

src/cuda/tile/_passes/token_order.py

@@ -112,7 +112,13 @@ def token_order_pass(root_block: Block, alias_result: AliasResult):
 
 
 def _get_input_var(op: Operation):
-    return op.array if "array" in op.operands else op.pointer
+    if "view" in op.operands:
+        return op.view
+    elif "pointer" in op.operands:
+        return op.pointer
+    else:
+        raise TileInternalError(f"Cannot determine input var for op {op}: "
+                                f"expected 'view' or 'pointer' operand")
 
 
 def _get_block_memory_effects(block: Block,
@@ -484,7 +490,7 @@ def is_idx_injective(idx_var: Var) -> bool:
         return loop_op.is_for_loop and idx_var.name == loop_op.induction_var.name
 
     return set(store_op for store_op in tile_store_candidates
-               if _get_input_var(store_op).get_type().elements_disjoint
+               if _get_input_var(store_op).get_type().array_ty.elements_disjoint
                and any(is_idx_injective(idx_var) for idx_var in store_op.index))
 
 

src/cuda/tile/_passes/unhoist_partition_views.py

@@ -0,0 +1,32 @@
+# SPDX-FileCopyrightText: Copyright (c) <2026> NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+import dataclasses
+from cuda.tile._ir.ir import Block, Mapper, Operation
+from cuda.tile._ir.ops import MakePartitionView, TileLoad, TileStore, NumTiles
+
+
+def unhoist_partition_views(root_block: Block):
+    def_info: dict[str, tuple[Operation, Block]] = {}
+    _unhoist(root_block, def_info)
+
+
+def _unhoist(block: Block, def_info: dict[str, tuple[Operation, Block]]):
+    new_block = block.empty_like_self()
+    for op in block:
+        if isinstance(op, (TileLoad, TileStore, NumTiles)):
+            view_def, def_block = def_info[op.view.name]
+            if isinstance(view_def, MakePartitionView) and def_block is not block:
+                mapper = Mapper(block.ctx)
+                new_block.append(view_def.clone(mapper))
+                op = dataclasses.replace(op, view=mapper.get_var(op.view))
+
+        for nested in op.nested_blocks:
+            _unhoist(nested, def_info)
+
+        new_block.append(op)
+        for v in op.result_vars:
+            def_info[v.name] = (op, block)
+
+    block[:] = new_block.detach_all()