From 2b0f2b4880e7084028115452802f804410a6f7b0 Mon Sep 17 00:00:00 2001 From: Tim Besard Date: Wed, 22 Jul 2026 19:23:14 +0200 Subject: [PATCH 1/5] Add Tile IR StridedView encoding and intrinsics --- src/bytecode/encodings.jl | 16 +++ src/bytecode/types.jl | 19 +++ src/compiler/analysis/alias.jl | 6 +- src/compiler/analysis/effects.jl | 2 + src/compiler/intrinsics.jl | 2 +- src/compiler/intrinsics/views.jl | 197 +++++++++++++++++++++++++++++++ src/language/types.jl | 45 +++++++ test/bytecode.jl | 26 ++++ 8 files changed, 310 insertions(+), 3 deletions(-) diff --git a/src/bytecode/encodings.jl b/src/bytecode/encodings.jl index 1473c8fc..3e9de02f 100644 --- a/src/bytecode/encodings.jl +++ b/src/bytecode/encodings.jl @@ -100,6 +100,7 @@ module Opcode const UnpackOp = 112 # since 13.3 # 113 (AllocaOp) not implemented const MmaFScaledOp = 114 # since 13.3 + const MakeStridedViewOp = 116 # since 13.3 const InsertOp = 118 # since 13.4 end @@ -462,6 +463,21 @@ function encode_MakePartitionViewOp!(cb::CodeBuilder, result_type::TypeId, tenso return new_op!(cb) end +""" + encode_MakeStridedViewOp!(cb, result_type, tensor_view) -> Value + +Create a strided view from a tensor view. +Opcode: 116 (Tile IR v13.3+) +""" +function encode_MakeStridedViewOp!(cb::CodeBuilder, result_type::TypeId, tensor_view::Value) + cb.version >= v"13.3" || + throw(IRError("MakeStridedViewOp requires Tile IR v13.3+, got v$(cb.version)")) + encode_varint!(cb.buf, Opcode.MakeStridedViewOp) + encode_typeid!(cb.buf, result_type) + encode_operand!(cb.buf, tensor_view) + return new_op!(cb) +end + """ encode_GetIndexSpaceShapeOp!(cb, result_types, partition_view) -> Tuple{Value...} diff --git a/src/bytecode/types.jl b/src/bytecode/types.jl index 14bef7cf..a19e107b 100644 --- a/src/bytecode/types.jl +++ b/src/bytecode/types.jl @@ -46,6 +46,7 @@ module CompositeType const TensorView = UInt8(0x0e) const PartitionView = UInt8(0x0f) const Func = UInt8(0x10) + const StridedView = UInt8(0x15) # since 13.3 end # Dynamic shape marker @@ -206,6 +207,24 @@ function partition_view_type!(table::TypeTable, _get_or_create!(table, buf) end +function strided_view_type!(table::TypeTable, + tile_shape::TileShape, + traversal_strides::TileShape, + tensor_view::TypeId, + dim_map::AbstractVector{<:Integer}, + padding_value::PaddingValue.T) + table.version >= v"13.3" || + throw(IRError("StridedView requires Tile IR bytecode v13.3+, got v$(table.version)")) + buf = UInt8[CompositeType.StridedView] + encode_optional_flags!(buf, padding_value) + encode_int_list!(buf, collect(tile_shape), 4) + encode_int_list!(buf, collect(traversal_strides), 4) + encode_varint!(buf, tensor_view.id) + encode_int_list!(buf, dim_map, 4) + encode_optional_padding_byte!(buf, padding_value) + _get_or_create!(table, buf) +end + function function_type!(table::TypeTable, param_types::AbstractVector{TypeId}, result_types::AbstractVector{TypeId}) diff --git a/src/compiler/analysis/alias.jl b/src/compiler/analysis/alias.jl index 889778f0..2e815f9e 100644 --- a/src/compiler/analysis/alias.jl +++ b/src/compiler/analysis/alias.jl @@ -78,7 +78,8 @@ function transfer(a::AliasAnalysis, r::DataflowResult, @nospecialize(func), # View constructors and pointer passthroughs: propagate from the source # operand. `make_tensor_view(::Type{T}, ptr, sizes, strides)` — alias source - # is the ptr (operand 2). `make_partition_view(tv, ...)` — operand 1. + # is the ptr (operand 2). Tile-view constructors take their TensorView as + # operand 1. if is_view_constructor(func) || is_pointer_passthrough(func) src_idx = func === Intrinsics.make_tensor_view ? 2 : 1 length(ops) >= src_idx && return operand_value(a, r, ops[src_idx]) @@ -101,7 +102,8 @@ These propagate alias identity from their first operand. """ function is_view_constructor(func) return func === Intrinsics.make_tensor_view || - func === Intrinsics.make_partition_view + func === Intrinsics.make_partition_view || + func === Intrinsics.make_strided_view end function is_pointer_passthrough(func) diff --git a/src/compiler/analysis/effects.jl b/src/compiler/analysis/effects.jl index 0d18c97b..d6c3ef78 100644 --- a/src/compiler/analysis/effects.jl +++ b/src/compiler/analysis/effects.jl @@ -21,9 +21,11 @@ externally observable side effect). """ function classify_memory_op(resolved_func) if resolved_func === Intrinsics.load_partition_view || + resolved_func === Intrinsics.load_strided_view || resolved_func === Intrinsics.load_ptr_tko return MEM_LOAD elseif resolved_func === Intrinsics.store_partition_view || + resolved_func === Intrinsics.store_strided_view || resolved_func === Intrinsics.store_ptr_tko return MEM_STORE elseif resolved_func === Intrinsics.print_tko diff --git a/src/compiler/intrinsics.jl b/src/compiler/intrinsics.jl index 73189514..a757791e 100644 --- a/src/compiler/intrinsics.jl +++ b/src/compiler/intrinsics.jl @@ -5,7 +5,7 @@ module Intrinsics using Base: compilerbarrier, inferencebarrier -using ..cuTile: Tile, TileArray, Constant, TensorView, PartitionView +using ..cuTile: Tile, TileArray, Constant, TensorView, PartitionView, StridedView using ..cuTile: Signedness, ComparisonPredicate, ComparisonOrdering using ..cuTile: IdentityVal, FloatIdentityVal, IntegerIdentityVal diff --git a/src/compiler/intrinsics/views.jl b/src/compiler/intrinsics/views.jl index d682c2f3..b274111e 100644 --- a/src/compiler/intrinsics/views.jl +++ b/src/compiler/intrinsics/views.jl @@ -135,6 +135,68 @@ function emit_intrinsic!(ctx::CGCtx, ::typeof(Intrinsics.load_partition_view), a return CGVal(tile_val, tile_type, Tile{elem_type, TupleType(julia_shape)}, tile_shape) end +""" + Intrinsics.load_strided_view(sv::StridedView{T,N,Shape,Steps}, ...) + +Token-ordered load from a `StridedView`. It uses the same Tile IR load op as +`load_partition_view`, but has a distinct intrinsic identity so analyses never +assume that distinct tile indices access disjoint memory. +""" +@intrinsic load_strided_view(sv, latency, allow_tma, indices, check_bounds) +function tfunc(𝕃, ::typeof(Intrinsics.load_strided_view), @nospecialize(sv), @nospecialize args...) + sv_type = CC.widenconst(sv) + sv_type <: StridedView || return nothing + sv_type isa DataType || return nothing + length(sv_type.parameters) >= 4 || return nothing + T = eltype(sv_type) + Shape = sv_type.parameters[3] + Shape isa Type || return nothing + return Tile{T, Shape} +end +function emit_intrinsic!(ctx::CGCtx, ::typeof(Intrinsics.load_strided_view), args) + cb = ctx.cb + tt = ctx.tt + input_token = extract_token_arg!(ctx, args) + + sv_arg = emit_value!(ctx, args[1]) + sv_arg === nothing && throw(IRError("load_strided_view() requires a StridedView argument")) + sv_arg.v === nothing && throw(IRError("load_strided_view() requires a materialized StridedView")) + sv_arg.constant === nothing && throw(IRError("load_strided_view(): StridedView missing ndim info")) + ndim = something(sv_arg.constant) + + sv_type = CC.widenconst(sv_arg.jltype) + elem_type = eltype(sv_type) + tile_shape = RowMajorShape(ColMajorShape(size(sv_type))) + dtype = lookup_dtype!(tt, elem_type) + tile_type = tile_type!(tt, dtype, tile_shape) + token_type = Token(tt) + + latency = @something get_constant(ctx, args[2]) throw(IRError("load_strided_view(): latency must be a compile-time constant")) + allow_tma = @something get_constant(ctx, args[3]) throw(IRError("load_strided_view(): allow_tma must be a compile-time constant")) + allow_tma_val = allow_tma isa Bool ? allow_tma : true + check_bounds = @something get_constant(ctx, args[5]) throw(IRError("load_strided_view(): check_bounds must be a compile-time constant")) + + index_tvs = resolve_tuple(ctx, args[4], "load_strided_view indices") + index_vals = Value[tv.v for tv in index_tvs] + index_jl_types = Type[tv.jltype for tv in index_tvs] + unique_types = unique(index_jl_types) + length(unique_types) <= 1 || throw(IRError("All index types must match, got: $unique_types")) + isempty(unique_types) && ndim > 0 && throw(IRError("load_strided_view(): indices required for $(ndim)D view")) + index_jl_type = isempty(unique_types) ? Int32 : unique_types[1] + index_type = tile_type_for_julia!(ctx, index_jl_type) + index_vals = pad_indices(ctx, index_vals, ndim, index_type, index_jl_type) + reverse!(index_vals) + + optimization_hints = create_optimization_hints(ctx, latency, allow_tma_val) + tile_val, result_token = encode_LoadViewTkoOp!( + cb, tile_type, token_type, sv_arg.v, index_vals; + token=input_token, optimization_hints, inbounds=fill(!check_bounds, ndim)) + ctx.result_tokens[ctx.current_ssa_idx] = result_token + + julia_shape = ColMajorShape(tile_shape) + return CGVal(tile_val, tile_type, Tile{elem_type, TupleType(julia_shape)}, tile_shape) +end + function pad_indices(ctx::CGCtx, index_vals::Vector{Value}, ndim::Int, idx_type::TypeId, idx_jl_type::Type) while length(index_vals) < ndim idx_bytes = reinterpret(UInt8, [eltype(idx_jl_type)(0)]) @@ -207,6 +269,71 @@ function emit_intrinsic!(ctx::CGCtx, ::typeof(Intrinsics.make_partition_view), a CGVal(partition, pv_type, PartitionView{elem_type, ndim, Tuple{shape...}}, RowMajorShape(()), nothing, Some(ndim), nothing) end +""" + Intrinsics.make_strided_view(tensor_view, tile_shape, traversal_strides, + padding_value, dim_map) + +Construct a Tile IR `StridedView` whose `steps` control the distance between +successive tile origins. The arguments follow the Tile IR type fields; the +Julia-facing `eachtile` layer translates its column-major `step` and default +dimension order at this boundary. +""" +@intrinsic make_strided_view(tensor_view, tile_shape, traversal_strides, padding_value, dim_map) +function tfunc(𝕃, ::typeof(Intrinsics.make_strided_view), @nospecialize(tensor_view), + @nospecialize(tile_shape_arg), @nospecialize(traversal_strides_arg), @nospecialize args...) + tv_type = CC.widenconst(tensor_view) + tv_type <: TensorView || return nothing + isa(tile_shape_arg, CC.Const) || return nothing + isa(traversal_strides_arg, CC.Const) || return nothing + shape = tile_shape_arg.val + strides = traversal_strides_arg.val + shape isa Tuple && strides isa Tuple || return nothing + length(shape) == length(strides) || return nothing + T = eltype(tv_type) + N = ndims(tv_type) + return StridedView{T, N, Tuple{shape...}, Tuple{strides...}} +end +function emit_intrinsic!(ctx::CGCtx, ::typeof(Intrinsics.make_strided_view), args) + tensor_view = emit_value!(ctx, args[1]) + tensor_view === nothing && throw(IRError("make_strided_view() requires a TensorView argument")) + ctx.tt.version >= v"13.3" || + throw(IRError("make_strided_view requires Tile IR bytecode v13.3+, got v$(ctx.tt.version)")) + + shape = @something get_constant(ctx, args[2]) throw(IRError("make_strided_view() tile_shape must be a compile-time constant")) + shape isa Tuple || throw(IRError("make_strided_view() shape must be a tuple, got $(typeof(shape))")) + validate_tile_shape(collect(Int, shape), "eachtile") + tile_shape = RowMajorShape(ColMajorShape(shape)) + + strides = @something get_constant(ctx, args[3]) throw(IRError("make_strided_view() traversal_strides must be a compile-time constant")) + strides isa Tuple || throw(IRError("make_strided_view() traversal_strides must be a tuple, got $(typeof(strides))")) + length(strides) == length(tile_shape) || + throw(IRError("make_strided_view(): expected $(length(tile_shape)) traversal strides, got $(length(strides))")) + all(stride -> stride isa Integer && stride > 0, strides) || + throw(IRError("make_strided_view(): traversal_strides must be strictly positive integers, got $strides")) + traversal_strides = RowMajorShape(ColMajorShape(strides)) + + padding_value = convert_enum(PaddingValue, + @something get_constant(ctx, args[4]) throw(IRError("padding_mode must be a compile-time constant"))) + + ndim = length(tile_shape) + dim_map_val = @something get_constant(ctx, args[5]) throw(IRError("make_strided_view() dim_map must be a compile-time constant")) + if dim_map_val === nothing + dim_map = collect(0:ndim-1) + else + validate_axis_order(dim_map_val, ndim, 1, "eachtile") + julia_dim_map = collect(Int, map(p -> p - 1, dim_map_val)) + dim_map = [ndim - 1 - julia_dim_map[ndim - i] for i in 0:ndim-1] + end + + sv_type = strided_view_type!(ctx.tt, tile_shape, traversal_strides, tensor_view.type_id, dim_map, + padding_value) + strided = encode_MakeStridedViewOp!(ctx.cb, sv_type, tensor_view.v) + elem_type = eltype(tensor_view.jltype) + CGVal(strided, sv_type, + StridedView{elem_type, ndim, Tuple{shape...}, Tuple{strides...}}, + RowMajorShape(()), nothing, Some(ndim), nothing) +end + """ compute_tensor_view_strides(array_spec, ndim) -> Vector{Int64} @@ -497,3 +624,73 @@ function emit_intrinsic!(ctx::CGCtx, ::typeof(Intrinsics.store_partition_view), return nothing end + +""" + Intrinsics.store_strided_view(sv::StridedView{T,N,Shape,Steps}, ...) + +Token-ordered store to a `StridedView`. This deliberately has a different +intrinsic identity from `store_partition_view`: overlapping windows must keep +their loop-carried token dependency even when their tile indices differ. +""" +@intrinsic store_strided_view(sv::StridedView{T, N, Shape, Steps}, + tile::Tile{T}, + latency::Union{Int, Nothing}, + allow_tma::Bool, + indices::NTuple{M, <:Integer}, + check_bounds::Bool) where {T, N, Shape, Steps, M} +tfunc(𝕃, ::typeof(Intrinsics.store_strided_view), @nospecialize args...) = Nothing +efunc(::typeof(Intrinsics.store_strided_view), effects::CC.Effects) = + CC.Effects(effects; effect_free=CC.ALWAYS_FALSE) +function emit_intrinsic!(ctx::CGCtx, ::typeof(Intrinsics.store_strided_view), args) + cb = ctx.cb + tt = ctx.tt + input_token = extract_token_arg!(ctx, args) + + sv_arg = emit_value!(ctx, args[1]) + sv_arg === nothing && throw(IRError("store_strided_view() requires a StridedView argument")) + sv_arg.v === nothing && throw(IRError("store_strided_view() requires a materialized StridedView")) + sv_arg.constant === nothing && throw(IRError("store_strided_view(): StridedView missing ndim info")) + ndim = something(sv_arg.constant) + + tile_tv = emit_value!(ctx, args[2]) + tile_tv === nothing && throw(IRError("store_strided_view() requires a tile argument")) + tile_shape = tile_tv.shape + tile_shape === nothing && throw(IRError("Cannot determine tile shape for store_strided_view()")) + elem_type = eltype(CC.widenconst(tile_tv.jltype)) + dtype = lookup_dtype!(tt, elem_type) + + tile_val = tile_tv.v + actual_ndim = ndim + actual_tile_shape = tile_shape + if length(tile_shape) == 0 + actual_ndim = 1 + actual_tile_shape = RowMajorShape([1]) + tile_1d_type = tile_type!(tt, dtype, actual_tile_shape) + tile_val = encode_ReshapeOp!(cb, tile_1d_type, tile_val) + end + + latency = @something get_constant(ctx, args[3]) throw(IRError("store_strided_view(): latency must be a compile-time constant")) + allow_tma = @something get_constant(ctx, args[4]) throw(IRError("store_strided_view(): allow_tma must be a compile-time constant")) + allow_tma_val = allow_tma isa Bool ? allow_tma : true + check_bounds = @something get_constant(ctx, args[6]) throw(IRError("store_strided_view(): check_bounds must be a compile-time constant")) + + index_tvs = resolve_tuple(ctx, args[5], "store_strided_view indices") + index_vals = Value[tv.v for tv in index_tvs] + index_jl_types = Type[tv.jltype for tv in index_tvs] + unique_types = unique(index_jl_types) + length(unique_types) <= 1 || throw(IRError("All index types must match, got: $unique_types")) + isempty(unique_types) && actual_ndim > 0 && + throw(IRError("store_strided_view(): indices required for $(actual_ndim)D view")) + index_jl_type = isempty(unique_types) ? Int32 : unique_types[1] + index_type = tile_type_for_julia!(ctx, index_jl_type) + index_vals = pad_indices(ctx, index_vals, actual_ndim, index_type, index_jl_type) + reverse!(index_vals) + + optimization_hints = create_optimization_hints(ctx, latency, allow_tma_val) + token_type = Token(tt) + result_token = encode_StoreViewTkoOp!( + cb, token_type, tile_val, sv_arg.v, index_vals; + token=input_token, optimization_hints, inbounds=fill(!check_bounds, actual_ndim)) + ctx.result_tokens[ctx.current_ssa_idx] = result_token + return nothing +end diff --git a/src/language/types.jl b/src/language/types.jl index df584d7e..dae36823 100644 --- a/src/language/types.jl +++ b/src/language/types.jl @@ -391,6 +391,51 @@ Base.size(::Type{<:PartitionView{<:Any, <:Any, Shape}}, d::Integer) where {Shape Base.size(pv::PartitionView) = size(typeof(pv)) Base.size(pv::PartitionView, d::Integer) = size(typeof(pv), d) +""" + StridedView{T, N, Shape, Steps} + +Opaque Tile IR view with a fixed tile shape and a distinct traversal stride +between adjacent tile origins. It is an internal compiler carrier; public +array-of-tiles operations use `eachtile` instead. +""" +mutable struct StridedView{T, N, Shape, Steps} end + +Base.eltype(::Type{<:StridedView{T}}) where {T} = T +Base.eltype(::StridedView{T}) where {T} = T +Base.ndims(::Type{<:StridedView{<:Any, N}}) where {N} = N +Base.ndims(::StridedView{<:Any, N}) where {N} = N +Base.size(::Type{<:StridedView{<:Any, <:Any, Shape}}) where {Shape} = Tuple(Shape.parameters) +Base.size(::Type{<:StridedView{<:Any, <:Any, Shape}}, d::Integer) where {Shape} = Shape.parameters[d] +Base.size(sv::StridedView) = size(typeof(sv)) +Base.size(sv::StridedView, d::Integer) = size(typeof(sv), d) + +""" + TiledView{A, RequestedShape, ViewShape, Step, Padding} + +Internal immutable array-of-tiles wrapper returned by `eachtile`. Shape and +step live in the type because Tile IR view types require compile-time values; +the parent is the sole runtime field. +""" +struct TiledView{A, RequestedShape, ViewShape, Step, Padding} + parent::A +end + +Base.parent(tiles::TiledView) = tiles.parent +Base.ndims(::Type{<:TiledView{A}}) where {A} = ndims(A) +Base.ndims(tiles::TiledView) = ndims(typeof(tiles)) + +_tiled_view_steps(::Type{<:TiledView{A, RequestedShape, ViewShape, Step}}) where + {A, RequestedShape, ViewShape, Step} = Tuple(Step.parameters) + +function Base.size(tiles::TiledView) + steps = _tiled_view_steps(typeof(tiles)) + ntuple(i -> cld(size(tiles.parent, i), Int32(steps[i])), Val(ndims(tiles))) +end +function Base.size(tiles::TiledView, d::Integer) + d < 1 && error("arraysize: dimension out of range") + d > ndims(tiles) ? Int32(1) : size(tiles)[d] +end + """ Constant{T, V} diff --git a/test/bytecode.jl b/test/bytecode.jl index b2583a74..2b6f07e2 100644 --- a/test/bytecode.jl +++ b/test/bytecode.jl @@ -44,3 +44,29 @@ end cuTile.Value(1), cuTile.Value[]) end end + +@testset "Tile IR v13.3 StridedView encodings" begin + tt = cuTile.TypeTable(; version=v"13.3") + tensor_view = cuTile.TypeId(7) + strided = cuTile.strided_view_type!( + tt, cuTile.RowMajorShape([4, 8]), cuTile.RowMajorShape([2, 3]), + tensor_view, [0, 1], cuTile.PaddingValue.Zero) + @test strided == cuTile.TypeId(2) + @test last(cuTile.items(tt)).first == UInt8[ + 0x15, 0x01, + 0x02, 0x04, 0x00, 0x00, 0x00, 0x08, 0x00, 0x00, 0x00, + 0x02, 0x02, 0x00, 0x00, 0x00, 0x03, 0x00, 0x00, 0x00, + 0x07, + 0x02, 0x00, 0x00, 0x00, 0x00, 0x01, 0x00, 0x00, 0x00, + 0x00, + ] + @test_throws "v13.3+" cuTile.strided_view_type!( + cuTile.TypeTable(; version=v"13.2"), cuTile.RowMajorShape([4]), + cuTile.RowMajorShape([2]), tensor_view, [0], cuTile.PaddingValue.Missing) + + cb = make_builder(v"13.3") + @test cuTile.encode_MakeStridedViewOp!(cb, cuTile.TypeId(5), cuTile.Value(2)) == cuTile.Value(0) + @test cb.buf == UInt8[116, 5, 2] + @test_throws "v13.3+" cuTile.encode_MakeStridedViewOp!( + make_builder(v"13.2"), cuTile.TypeId(5), cuTile.Value(2)) +end From 42c84e95a8f62fdc56f7e77bfb89bc27651cd977 Mon Sep 17 00:00:00 2001 From: Tim Besard Date: Wed, 22 Jul 2026 19:26:43 +0200 Subject: [PATCH 2/5] Export eachtile for stepped tile windows --- README.md | 35 ++++++++-- src/language/operations.jl | 124 ++++++++++++++++++++++++++++++++++++ test/codegen/token_order.jl | 20 ++++++ test/codegen/views.jl | 41 ++++++++++++ test/device/views.jl | 75 ++++++++++++++++++++++ test/types.jl | 22 +++++++ 6 files changed, 311 insertions(+), 6 deletions(-) diff --git a/README.md b/README.md index da1d3b48..64043244 100644 --- a/README.md +++ b/README.md @@ -133,6 +133,7 @@ uses standard Julia syntax and is overlaid on `Base`. |-----------|-------------| | `ct.load(arr; index, shape, ...)` | Load a tile from array | | `ct.store(arr; index, tile, ...)` | Store a tile to array | +| `eachtile(arr, shape; step=...)` | Array-of-tiles view with controllable tile origins | | `ct.gather(arr, indices; ...)` | Gather elements by index tile | | `ct.scatter(arr, indices, tile; ...)` | Scatter elements by index tile | @@ -324,12 +325,12 @@ end | `@view arr[r1:r2, :, ...]` / `view(arr, ...)` | Sub-range view of a `TileArray` | `@view` and `view` derive a sub-range `TileArray` from an existing one. Each -index must be `:` or a `UnitRange` (e.g. `i:j`); other forms (`StepRange`, -scalar `Int`, `CartesianIndex`, ...) are rejected at compile time. The result -is itself a `TileArray` and can be passed to `ct.load`/`ct.store` (or sliced -again). A runtime assert verifies that each range starts at ≥ 1; the upper -bound is not checked because Julia's range construction already clamps -`last(r) >= first(r) - 1`. +index must be `:`, a `UnitRange` (e.g. `i:j`), or a positive `StepRange` (e.g. +`i:s:j`); scalar `Int` and `CartesianIndex` forms are rejected at compile +time. A StepRange changes the element stride inside the resulting TileArray. +The result can be passed to `ct.load`/`ct.store` (or sliced again). Runtime +asserts verify that ranges start at ≥ 1 and have a positive step; negative +steps cannot be represented by Tile IR TensorViews. ```julia function rowsum(a, b, r1::Int32, r2::Int32) @@ -348,6 +349,28 @@ end | `transpose(arr)` | 2D transpose (`permutedims(arr, (2, 1))`) | | `reshape(arr, dims)` | Column-major reshape, requires contiguous source | +### Tile Windows + +`eachtile` creates a small, indexable device-side collection of fixed-shape +tiles. Its indices are 1-based and `size(eachtile(a, shape; step), d)` is +`cld(size(a, d), step[d])`. `step` controls tile origins, not the element +stride inside a tile: + +```julia +adjacent = eachtile(a, (8, 8)) # default: step == (8, 8) +overlap = eachtile(a, (8, 8); step=(4, 8)) # neighboring windows overlap +gapped = eachtile(a, (8, 8); step=(16, 8)) # gaps between row windows + +tile = overlap[2, 1] +overlap[2, 1] = tile +``` + +Use `ct.load` and `ct.store` for `check_bounds`, `latency`, and `allow_tma` +controls. Equal shape and step use the ordinary Tile IR partition view; +unequal values require Tile IR bytecode v13.3 or newer. This is distinct from +`@view a[1:2:end, :]`, which steps individual elements rather than tile +origins. + ### Atomics | Operation | Description | |-----------|-------------| diff --git a/src/language/operations.jl b/src/language/operations.jl index a1b1e2cd..393720b2 100644 --- a/src/language/operations.jl +++ b/src/language/operations.jl @@ -261,6 +261,7 @@ Base.reshape(arr::TileArray, dims::Int...) = reshape(arr, Val(dims)) =============================================================================# public bid, num_blocks, num_tiles, load, store, gather, scatter, Rounding +export eachtile """ Padding mode for load operations. @@ -340,6 +341,129 @@ end return :($kept) end +function _normalize_eachtile_shape(shape::Tuple, rank::Int, name::String) + length(shape) == rank && return shape + if length(shape) < rank + return (shape..., ntuple(_ -> 1, rank - length(shape))...) + end + trailing = length(shape) - something(findlast(!=(1), shape), 0) + trailing >= length(shape) - rank || + throw(ArgumentError("eachtile: cannot squeeze $name $shape to rank $rank; " * + "only trailing singleton dimensions may be removed")) + return shape[1:rank] +end + +@generated function _eachtile(a::A, ::Val{Requested}, ::Val{Step}, ::Val{Padding}) where + {A<:TileArray, Requested, Step, Padding} + Requested isa Tuple || throw(ArgumentError( + "eachtile: tile_shape must be a compile-time tuple, got $(typeof(Requested))")) + rank = ndims(A) + view_shape = _normalize_eachtile_shape(Requested, rank, "tile_shape") + normalized_step = if Step === nothing + view_shape + elseif Step isa Tuple + _normalize_eachtile_shape(Step, rank, "step") + else + throw(ArgumentError("eachtile: step must be a compile-time tuple or nothing, got $(typeof(Step))")) + end + all(step -> step isa Integer && step > 0, normalized_step) || + throw(ArgumentError("eachtile: step must contain strictly positive integers, got $normalized_step")) + result_type = TiledView{A, Tuple{Requested...}, Tuple{view_shape...}, + Tuple{normalized_step...}, Padding} + return :($result_type(a)) +end + +_tiled_view_requested_shape(::TiledView{A, RequestedShape}) where {A, RequestedShape} = + Tuple(RequestedShape.parameters) +_tiled_view_shape(::TiledView{A, RequestedShape, ViewShape}) where + {A, RequestedShape, ViewShape} = Tuple(ViewShape.parameters) +_tiled_view_step(::TiledView{A, RequestedShape, ViewShape, Step}) where + {A, RequestedShape, ViewShape, Step} = Tuple(Step.parameters) + +""" + eachtile(a, tile_shape; step=nothing, + padding_mode=PaddingMode.Undetermined) + +Create an immutable device-side collection of fixed-shape tiles over `a`. +`step` controls the distance between adjacent tile origins: the default and +`step == tile_shape` create adjacent partitions, smaller values overlap, and +larger values leave gaps. Unlike `@view a[1:2:end]`, which changes element +strides inside an array, `eachtile` changes tile origins. + +Unequal tile shape and step require Tile IR bytecode v13.3 or newer. Both are +compile-time tuples; tile indices are 1-based and partial edge tiles use the +given padding mode on loads and clipped stores. +""" +@inline function eachtile(a::TileArray, tile_shape::Tuple; + step::Union{Tuple, Nothing}=nothing, + padding_mode::PaddingMode.T=PaddingMode.Undetermined) + _eachtile(a, Val(tile_shape), Val(step), Val(padding_mode)) +end + +@inline function _make_tile_view(tiles::TiledView{A, RequestedShape, Shape, Shape, Padding}) where + {A, RequestedShape, Shape, Padding} + parent = tiles.parent + tv = Intrinsics.make_tensor_view(typeof(parent), parent.ptr, parent.sizes, parent.strides) + Intrinsics.make_partition_view(tv, _tiled_view_shape(tiles), Padding, nothing) +end +@inline function _make_tile_view(tiles::TiledView{A, RequestedShape, Shape, Step, Padding}) where + {A, RequestedShape, Shape, Step, Padding} + parent = tiles.parent + tv = Intrinsics.make_tensor_view(typeof(parent), parent.ptr, parent.sizes, parent.strides) + Intrinsics.make_strided_view(tv, _tiled_view_shape(tiles), _tiled_view_step(tiles), Padding, nothing) +end + +@inline _load_tile_view(view::PartitionView, latency, allow_tma, indices, check_bounds) = + Intrinsics.load_partition_view(view, latency, allow_tma, indices, check_bounds) +@inline _load_tile_view(view::StridedView, latency, allow_tma, indices, check_bounds) = + Intrinsics.load_strided_view(view, latency, allow_tma, indices, check_bounds) +@inline _store_tile_view(view::PartitionView, tile, latency, allow_tma, indices, check_bounds) = + Intrinsics.store_partition_view(view, tile, latency, allow_tma, indices, check_bounds) +@inline _store_tile_view(view::StridedView, tile, latency, allow_tma, indices, check_bounds) = + Intrinsics.store_strided_view(view, tile, latency, allow_tma, indices, check_bounds) + +@inline function load(tiles::TiledView{A}, index::NTuple{N, <:Integer}; + check_bounds::Bool=true, + latency::Union{Int, Nothing}=nothing, + allow_tma::Union{Bool, Nothing}=nothing) where {A, N} + N == ndims(tiles) || throw(ArgumentError("eachtile: expected $(ndims(tiles)) tile indices, got $N")) + view = _make_tile_view(tiles) + tile = _load_tile_view(view, latency, allow_tma, promote(index...) .- One(), check_bounds) + reshape(tile, _tiled_view_requested_shape(tiles)) +end +@inline function load(tiles::TiledView, index::Integer; kwargs...) + load(tiles, (index,); kwargs...) +end +@inline function load(tiles::TiledView; index, kwargs...) + load(tiles, index; kwargs...) +end + +@inline function store(tiles::TiledView{A}, index::NTuple{N, <:Integer}, tile::Tile{T}; + check_bounds::Bool=true, + latency::Union{Int, Nothing}=nothing, + allow_tma::Union{Bool, Nothing}=nothing) where {A, N, T} + N == ndims(tiles) || throw(ArgumentError("eachtile: expected $(ndims(tiles)) tile indices, got $N")) + reshaped = _reshape_to_rank(tile, Val(ndims(tiles))) + view = _make_tile_view(tiles) + _store_tile_view(view, reshaped, latency, allow_tma, promote(index...) .- One(), check_bounds) + return tile +end +@inline function store(tiles::TiledView, index::Integer, tile::Tile; kwargs...) + store(tiles, (index,), tile; kwargs...) +end +@inline function store(tiles::TiledView; index, tile::Tile, kwargs...) + store(tiles, index, tile; kwargs...) +end + +@overlay function Base.getindex(tiles::TiledView{A}, indices::Vararg{Integer, N}) where {A, N} + load(tiles, indices) +end +Base.Experimental.@consistent_overlay cuTileMethodTable function Base.setindex!( + tiles::TiledView{A}, tile::Tile, indices::Vararg{Integer, N}) where {A, N} + store(tiles, indices, tile) + return +end + # Reshape a tile to match target rank N, preserving data layout. @inline function _reshape_to_rank(tile::Tile, ::Val{N}) where {N} new_shape = _match_shape(Val(size(tile)), Val(N)) diff --git a/test/codegen/token_order.jl b/test/codegen/token_order.jl index e15b1c8a..f7445e15 100644 --- a/test/codegen/token_order.jl +++ b/test/codegen/token_order.jl @@ -124,6 +124,26 @@ end end end +@testset "token_order — StridedView stores keep token carry" begin + # Tile indices are injective, but overlapping windows are not disjoint in + # memory. `store_strided_view` therefore must never take the partition + # store parallelization path. + @test @filecheck begin + @check_label "entry" + @check "make_strided_view" + @check "iter_values" + @check "store_view_tko" + code_tiled(Tuple{AT, AT, Int32}) do a, b, n + windows = eachtile(b, (16,); step=(8,)) + for i in 1:n + tile = ct.load(a, i, (16,)) + ct.store(windows, i, tile) + end + return + end + end +end + @testset "layout_may_alias_internally" begin lmai = ct.layout_may_alias_internally @test !lmai((Int32(4),), (Int32(1),)) # contiguous 1-D diff --git a/test/codegen/views.jl b/test/codegen/views.jl index 2b74ea65..de93425f 100644 --- a/test/codegen/views.jl +++ b/test/codegen/views.jl @@ -77,3 +77,44 @@ end @check "assert {{.*}}reshape: TileArray must be contiguous" end end + +@testset "eachtile — partition and strided views" begin + # Equal shape and step preserves the pre-v13.3 PartitionView path. + @test @filecheck begin + @check_label "entry" + @check "make_partition_view" + @check_not "make_strided_view" + code_tiled(Tuple{ct.TileArray{Float32,1,spec1d}, ct.TileArray{Float32,1,spec1d}}; + bytecode_version=v"13.2") do a, b + src = eachtile(a, (8,)) + dst = eachtile(b, (8,); step=(8,)) + dst[1] = src[1] + return + end + end + + # Tile IR is row-major, so Julia shape/step `(8,4)`/`(3,2)` are reversed + # in the emitted StridedView type. + @test @filecheck begin + @check_label "entry" + @check "make_strided_view" + @check "tile=(4x8), traversal_strides=[2,3]" + @check "load_view_tko" + @check "store_view_tko" + code_tiled(Tuple{ct.TileArray{Float32,2,spec2d}, ct.TileArray{Float32,2,spec2d}}; + bytecode_version=v"13.3") do a, b + src = eachtile(a, (8, 4); step=(3, 2), padding_mode=ct.PaddingMode.Zero) + dst = eachtile(b, (8, 4); step=(3, 2)) + tile = ct.load(src, (2, 3); check_bounds=true, latency=3, allow_tma=false) + ct.store(dst, (2, 3), tile; check_bounds=true, latency=3, allow_tma=false) + return + end + end + + @test_throws "v13.3+" code_tiled( + Tuple{ct.TileArray{Float32,1,spec1d}}; bytecode_version=v"13.2") do a + tiles = eachtile(a, (8,); step=(4,)) + ct.store(tiles, 1, ct.load(tiles, 1)) + return + end +end diff --git a/test/device/views.jl b/test/device/views.jl index edc62771..154216b3 100644 --- a/test/device/views.jl +++ b/test/device/views.jl @@ -73,3 +73,78 @@ end expected = permutedims(Array(a), (3, 1, 2)) @test Array(b)[1:2, 1:2, 1:2] == expected[1:2, 1:2, 1:2] end + +@testset "eachtile — adjacent, overlapping, and gapped windows" begin + function copy_windows(a::ct.TileArray{Float32,1}, b::ct.TileArray{Float32,1}, n::Int32) + src = eachtile(a, (4,); step=(2,)) + dst = eachtile(b, (4,); step=(2,)) + for i in 1:n + dst[i] = src[i] + end + return + end + + a = CUDA.rand(Float32, 16) + b = CUDA.zeros(Float32, 16) + @cuda backend=cuTile copy_windows(a, b, Int32(8)) + @test Array(b) == Array(a) + + function copy_gaps(a::ct.TileArray{Float32,1}, b::ct.TileArray{Float32,1}) + src = eachtile(a, (4,); step=(8,)) + dst = eachtile(b, (4,); step=(8,)) + for i in 1:size(src, 1) + dst[i] = src[i] + end + return + end + + fill!(b, 0) + @cuda backend=cuTile copy_gaps(a, b) + expected = zeros(Float32, 16) + expected[1:4] .= Array(a)[1:4] + expected[9:12] .= Array(a)[9:12] + @test Array(b) == expected +end + +@testset "eachtile — asymmetric 2D windows" begin + function copy_window(a::ct.TileArray{Float32,2}, b::ct.TileArray{Float32,2}) + src = eachtile(a, (4, 8); step=(3, 2), padding_mode=ct.PaddingMode.Zero) + dst = eachtile(b, (4, 8); step=(3, 2)) + ct.store(dst, (2, 3), ct.load(src, (2, 3))) + return + end + + a = CUDA.rand(Float32, 12, 12) + b = CUDA.zeros(Float32, 12, 12) + @cuda backend=cuTile copy_window(a, b) + expected = zeros(Float32, 12, 12) + expected[4:7, 5:12] .= Array(a)[4:7, 5:12] + @test Array(b) == expected +end + +@testset "eachtile — partial edges and requested-rank normalization" begin + function copy_partial(a::ct.TileArray{Float32,1}, b::ct.TileArray{Float32,1}) + src = eachtile(a, (4,); step=(4,), padding_mode=ct.PaddingMode.Zero) + ct.store(b, 1, ct.load(src, 2)) + return + end + + a = CUDA.rand(Float32, 6) + b = CUDA.zeros(Float32, 4) + @cuda backend=cuTile copy_partial(a, b) + @test Array(b) == vcat(Array(a)[5:6], zeros(Float32, 2)) + + function copy_normalized(a::ct.TileArray{Float32,2}, b::ct.TileArray{Float32,2}) + src = eachtile(a, (4,); step=(2,)) + dst = eachtile(b, (4,); step=(2,)) + dst[2, 1] = src[2, 1] + return + end + + a2 = CUDA.rand(Float32, 8, 1) + b2 = CUDA.zeros(Float32, 8, 1) + @cuda backend=cuTile copy_normalized(a2, b2) + expected = zeros(Float32, 8, 1) + expected[3:6, 1] .= Array(a2)[3:6, 1] + @test Array(b2) == expected +end diff --git a/test/types.jl b/test/types.jl index 43d8670d..fafcbddb 100644 --- a/test/types.jl +++ b/test/types.jl @@ -62,6 +62,28 @@ end @test size(PV, 2) == 32 end +@testset "StridedView" begin + SV = cuTile.StridedView{Float32, 2, Tuple{16, 32}, Tuple{8, 4}} + @test eltype(SV) == Float32 + @test ndims(SV) == 2 + @test size(SV) == (16, 32) + @test size(SV, 1) == 16 + @test size(SV, 2) == 32 +end + +@testset "eachtile" begin + a = ct.TileArray(Ptr{Float32}(0), (Int32(16), Int32(10)), (Int32(1), Int32(16))) + tiles = eachtile(a, (8, 4); step=(3, 2)) + @test parent(tiles) === a + @test ndims(tiles) == 2 + @test size(tiles) == (Int32(6), Int32(5)) + @test size(tiles, 1) == Int32(6) + @test size(tiles, 2) == Int32(5) + @test :eachtile in names(cuTile) + @test_throws "strictly positive" eachtile(a, (8, 4); step=(0, 2)) + @test_throws "cannot squeeze step" eachtile(a, (8, 4); step=(2, 2, 2)) +end + @testset "TensorView" begin @test eltype(cuTile.TensorView{Float32, 2}) == Float32 @test eltype(cuTile.TensorView{Float64, 3}) == Float64 From e59546de09b5dcbc6b1835a9093ecdc68086a747 Mon Sep 17 00:00:00 2001 From: Tim Besard Date: Thu, 23 Jul 2026 11:06:51 +0200 Subject: [PATCH 3/5] Deduplicate view intrinsic codegen and generalize get_index_space_shape MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The load/store emit code for partition and strided views was near-identical: both lower to the same load_view_tko/store_view_tko ops. Factor it into shared emit_view_load!/emit_view_store! helpers (and a shared tfunc helper), keeping the four distinct intrinsic identities that token_order_pass! relies on to treat overlapping strided windows conservatively. Also let get_index_space_shape accept a StridedView operand — the Tile IR op takes any tile view since v13.1 — and drop the make_strided_view version check that merely duplicated the bytecode-layer v13.3 gates. Co-Authored-By: Claude Fable 5 --- src/compiler/intrinsics/views.jl | 366 ++++++++++++------------------- 1 file changed, 139 insertions(+), 227 deletions(-) diff --git a/src/compiler/intrinsics/views.jl b/src/compiler/intrinsics/views.jl index b274111e..59ed046c 100644 --- a/src/compiler/intrinsics/views.jl +++ b/src/compiler/intrinsics/views.jl @@ -3,33 +3,35 @@ """ - Intrinsics.get_index_space_shape(pv::PartitionView, axis::Integer) -> Int32 + Intrinsics.get_index_space_shape(view, axis::Integer) -> Int32 -Returns the size of `pv`'s index space along `axis` (i.e. how many tiles -fit along that dimension); lowers to `cuda_tile.get_index_space_shape`. +Returns the size of `view`'s index space along `axis` (i.e. how many tiles +fit along that dimension); lowers to `cuda_tile.get_index_space_shape`. `view` +may be a `PartitionView` or a `StridedView` — the Tile IR op accepts any tile +view since v13.1. `axis` is 0-indexed in Julia order and must be a compile-time constant. The Tile IR op returns the full shape; the codegen picks the requested axis (in row-major order). """ -@intrinsic get_index_space_shape(pv, axis) -tfunc(𝕃, ::typeof(Intrinsics.get_index_space_shape), @nospecialize(pv), @nospecialize(axis)) = Int32 +@intrinsic get_index_space_shape(view, axis) +tfunc(𝕃, ::typeof(Intrinsics.get_index_space_shape), @nospecialize(view), @nospecialize(axis)) = Int32 function emit_intrinsic!(ctx::CGCtx, ::typeof(Intrinsics.get_index_space_shape), args) cb = ctx.cb tt = ctx.tt - # args: (partition_view, axis) - pv_arg = emit_value!(ctx, args[1]) - pv_arg === nothing && throw(IRError("get_index_space_shape() requires a PartitionView argument")) - pv_arg.v === nothing && throw(IRError("get_index_space_shape() requires a materialized PartitionView")) + # args: (view, axis) — view is a PartitionView or StridedView + view_arg = emit_value!(ctx, args[1]) + view_arg === nothing && throw(IRError("get_index_space_shape() requires a view argument")) + view_arg.v === nothing && throw(IRError("get_index_space_shape() requires a materialized view")) # Get axis (0-indexed Julia) and flip to Tile IR order axis = @something get_constant(ctx, args[2]) throw(IRError("get_index_space_shape() axis must be a compile-time constant")) axis = Int(axis) - # Get ndim from the PartitionView constant field - pv_arg.constant === nothing && throw(IRError("get_index_space_shape(): PartitionView missing ndim info")) - ndim = something(pv_arg.constant) + # Get ndim from the view's constant field + view_arg.constant === nothing && throw(IRError("get_index_space_shape(): view missing ndim info")) + ndim = something(view_arg.constant) # Flip axis for row-major Tile IR: Julia dim 0 → Tile IR dim ndim-1 tileir_axis = ndim - 1 - axis @@ -39,7 +41,7 @@ function emit_intrinsic!(ctx::CGCtx, ::typeof(Intrinsics.get_index_space_shape), result_types = fill(scalar_i32, ndim) # Emit GetIndexSpaceShapeOp - shape_vals = encode_GetIndexSpaceShapeOp!(cb, result_types, pv_arg.v) + shape_vals = encode_GetIndexSpaceShapeOp!(cb, result_types, view_arg.v) # Return the value for the requested axis (in Tile IR order) # shape_vals is a single Value when ndim == 1, otherwise a Tuple @@ -49,70 +51,68 @@ end # TODO: cuda_tile.get_tensor_shape -""" - Intrinsics.load_partition_view(pv::PartitionView{T,N,Shape}, - latency::Union{Int,Nothing}, - allow_tma::Bool, - indices::NTuple{M,<:Integer}) -> Tile{T,Shape} - -Token-ordered load of a tile from a `PartitionView`; lowers to -`cuda_tile.load_view_tko`. - -`latency` and `allow_tma` are compile-time hints. `indices` is in Julia -order and is reversed/zero-padded to match `pv`'s index space rank -before emission. The token argument is appended by `token_order_pass!` -and is not part of the user-visible signature. -""" -@intrinsic load_partition_view(pv, latency, allow_tma, indices, check_bounds) -function tfunc(𝕃, ::typeof(Intrinsics.load_partition_view), @nospecialize(pv), @nospecialize args...) - pv_type = CC.widenconst(pv) - pv_type <: PartitionView || return nothing - pv_type isa DataType || return nothing - length(pv_type.parameters) >= 3 || return nothing - T = eltype(pv_type) - Shape = pv_type.parameters[3] +#----------------------------------------------------------------------------- +# Shared codegen for the view load/store intrinsics +# +# `load_partition_view`/`load_strided_view` (and the two store variants) keep +# distinct intrinsic identities so token-ordering analyses can treat +# overlapping strided windows conservatively (see `get_parallel_stores` in +# transform/token_order.jl). Their bodies are otherwise identical — both lower +# to the same `load_view_tko`/`store_view_tko` ops — so the codegen lives in +# these shared helpers, keyed only by the intrinsic name for error messages. +#----------------------------------------------------------------------------- + +# Return type for the view-load intrinsics: `Tile{eltype, Shape}` where `Shape` +# is the view's tile-shape parameter (index 3 of both `PartitionView{T,N,Shape}` +# and `StridedView{T,N,Shape,Steps}`). `min_params` guards a fully-parameterized +# view type (3 for partition, 4 for strided). +function view_load_return_type(@nospecialize(view), min_params::Int) + view_type = CC.widenconst(view) + view_type isa DataType || return nothing + (view_type <: PartitionView || view_type <: StridedView) || return nothing + length(view_type.parameters) >= min_params || return nothing + Shape = view_type.parameters[3] Shape isa Type || return nothing - return Tile{T, Shape} + return Tile{eltype(view_type), Shape} end -function emit_intrinsic!(ctx::CGCtx, ::typeof(Intrinsics.load_partition_view), args) + +function emit_view_load!(ctx::CGCtx, args, name::String) cb = ctx.cb tt = ctx.tt - # Extract input token from last arg (added by token_order_pass!) + # Input token appended by token_order_pass!. input_token = extract_token_arg!(ctx, args) - # args: (partition_view, latency, allow_tma, indices, check_bounds) - pv_arg = emit_value!(ctx, args[1]) - pv_arg === nothing && throw(IRError("load_partition_view() requires a PartitionView argument")) - pv_arg.v === nothing && throw(IRError("load_partition_view() requires a materialized PartitionView")) + # args: (view, latency, allow_tma, indices, check_bounds) + view_arg = emit_value!(ctx, args[1]) + view_arg === nothing && throw(IRError("$name() requires a view argument")) + view_arg.v === nothing && throw(IRError("$name() requires a materialized view")) - # Get ndim from PartitionView constant field - pv_arg.constant === nothing && throw(IRError("load_partition_view(): PartitionView missing ndim info")) - ndim = something(pv_arg.constant) + # ndim from the view's constant field + view_arg.constant === nothing && throw(IRError("$name(): view missing ndim info")) + ndim = something(view_arg.constant) - # Extract tile shape from PartitionView type (PartitionView{T, N, Shape}) - # Reverse to Tile IR row-major order - pv_type = CC.widenconst(pv_arg.jltype) - elem_type = eltype(pv_type) - tile_shape = RowMajorShape(ColMajorShape(size(pv_type))) + # Tile shape from the view type, reversed to Tile IR row-major order. + view_type = CC.widenconst(view_arg.jltype) + elem_type = eltype(view_type) + tile_shape = RowMajorShape(ColMajorShape(size(view_type))) dtype = lookup_dtype!(tt, elem_type) tile_type = tile_type!(tt, dtype, tile_shape) token_type = Token(tt) - latency = @something get_constant(ctx, args[2]) throw(IRError("load_partition_view(): latency must be a compile-time constant")) - allow_tma = @something get_constant(ctx, args[3]) throw(IRError("load_partition_view(): allow_tma must be a compile-time constant")) + latency = @something get_constant(ctx, args[2]) throw(IRError("$name(): latency must be a compile-time constant")) + allow_tma = @something get_constant(ctx, args[3]) throw(IRError("$name(): allow_tma must be a compile-time constant")) allow_tma_val = allow_tma isa Bool ? allow_tma : true - check_bounds = @something get_constant(ctx, args[5]) throw(IRError("load_partition_view(): check_bounds must be a compile-time constant")) + check_bounds = @something get_constant(ctx, args[5]) throw(IRError("$name(): check_bounds must be a compile-time constant")) - # Extract indices - index_tvs = resolve_tuple(ctx, args[4], "load_partition_view indices") + index_tvs = resolve_tuple(ctx, args[4], "$name indices") index_vals = Value[tv.v for tv in index_tvs] index_jl_types = Type[tv.jltype for tv in index_tvs] unique_types = unique(index_jl_types) length(unique_types) <= 1 || throw(IRError("All index types must match, got: $unique_types")) - isempty(unique_types) && ndim > 0 && throw(IRError("load_partition_view(): indices required for $(ndim)D view")) + isempty(unique_types) && ndim > 0 && throw(IRError("$name(): indices required for $(ndim)D view")) index_jl_type = isempty(unique_types) ? Int32 : unique_types[1] # Int32 only for 0D case index_type = tile_type_for_julia!(ctx, index_jl_type) @@ -120,11 +120,10 @@ function emit_intrinsic!(ctx::CGCtx, ::typeof(Intrinsics.load_partition_view), a index_vals = pad_indices(ctx, index_vals, ndim, index_type, index_jl_type) reverse!(index_vals) - # Create optimization hints if provided optimization_hints = create_optimization_hints(ctx, latency, allow_tma_val) tile_val, result_token = encode_LoadViewTkoOp!( - cb, tile_type, token_type, pv_arg.v, index_vals; + cb, tile_type, token_type, view_arg.v, index_vals; token = input_token, optimization_hints, inbounds=fill(!check_bounds, ndim) ) @@ -135,68 +134,102 @@ function emit_intrinsic!(ctx::CGCtx, ::typeof(Intrinsics.load_partition_view), a return CGVal(tile_val, tile_type, Tile{elem_type, TupleType(julia_shape)}, tile_shape) end -""" - Intrinsics.load_strided_view(sv::StridedView{T,N,Shape,Steps}, ...) - -Token-ordered load from a `StridedView`. It uses the same Tile IR load op as -`load_partition_view`, but has a distinct intrinsic identity so analyses never -assume that distinct tile indices access disjoint memory. -""" -@intrinsic load_strided_view(sv, latency, allow_tma, indices, check_bounds) -function tfunc(𝕃, ::typeof(Intrinsics.load_strided_view), @nospecialize(sv), @nospecialize args...) - sv_type = CC.widenconst(sv) - sv_type <: StridedView || return nothing - sv_type isa DataType || return nothing - length(sv_type.parameters) >= 4 || return nothing - T = eltype(sv_type) - Shape = sv_type.parameters[3] - Shape isa Type || return nothing - return Tile{T, Shape} -end -function emit_intrinsic!(ctx::CGCtx, ::typeof(Intrinsics.load_strided_view), args) +function emit_view_store!(ctx::CGCtx, args, name::String) cb = ctx.cb tt = ctx.tt + + # Input token appended by token_order_pass!. input_token = extract_token_arg!(ctx, args) - sv_arg = emit_value!(ctx, args[1]) - sv_arg === nothing && throw(IRError("load_strided_view() requires a StridedView argument")) - sv_arg.v === nothing && throw(IRError("load_strided_view() requires a materialized StridedView")) - sv_arg.constant === nothing && throw(IRError("load_strided_view(): StridedView missing ndim info")) - ndim = something(sv_arg.constant) + # args: (view, tile, latency, allow_tma, indices, check_bounds) + view_arg = emit_value!(ctx, args[1]) + view_arg === nothing && throw(IRError("$name() requires a view argument")) + view_arg.v === nothing && throw(IRError("$name() requires a materialized view")) + + view_arg.constant === nothing && throw(IRError("$name(): view missing ndim info")) + ndim = something(view_arg.constant) + + tile_tv = emit_value!(ctx, args[2]) + tile_tv === nothing && throw(IRError("$name() requires a tile argument")) + tile_shape = tile_tv.shape + tile_shape === nothing && throw(IRError("Cannot determine tile shape for $name()")) - sv_type = CC.widenconst(sv_arg.jltype) - elem_type = eltype(sv_type) - tile_shape = RowMajorShape(ColMajorShape(size(sv_type))) + elem_type = eltype(CC.widenconst(tile_tv.jltype)) dtype = lookup_dtype!(tt, elem_type) - tile_type = tile_type!(tt, dtype, tile_shape) - token_type = Token(tt) - latency = @something get_constant(ctx, args[2]) throw(IRError("load_strided_view(): latency must be a compile-time constant")) - allow_tma = @something get_constant(ctx, args[3]) throw(IRError("load_strided_view(): allow_tma must be a compile-time constant")) + # 0-D scalar stores reshape to 1-D (views require at least 1-D). + tile_val = tile_tv.v + actual_ndim = ndim + if length(tile_shape) == 0 + actual_ndim = 1 + tile_1d_type = tile_type!(tt, dtype, RowMajorShape([1])) + tile_val = encode_ReshapeOp!(cb, tile_1d_type, tile_val) + end + + latency = @something get_constant(ctx, args[3]) throw(IRError("$name(): latency must be a compile-time constant")) + allow_tma = @something get_constant(ctx, args[4]) throw(IRError("$name(): allow_tma must be a compile-time constant")) allow_tma_val = allow_tma isa Bool ? allow_tma : true - check_bounds = @something get_constant(ctx, args[5]) throw(IRError("load_strided_view(): check_bounds must be a compile-time constant")) + check_bounds = @something get_constant(ctx, args[6]) throw(IRError("$name(): check_bounds must be a compile-time constant")) - index_tvs = resolve_tuple(ctx, args[4], "load_strided_view indices") + index_tvs = resolve_tuple(ctx, args[5], "$name indices") index_vals = Value[tv.v for tv in index_tvs] index_jl_types = Type[tv.jltype for tv in index_tvs] + unique_types = unique(index_jl_types) length(unique_types) <= 1 || throw(IRError("All index types must match, got: $unique_types")) - isempty(unique_types) && ndim > 0 && throw(IRError("load_strided_view(): indices required for $(ndim)D view")) - index_jl_type = isempty(unique_types) ? Int32 : unique_types[1] + isempty(unique_types) && actual_ndim > 0 && throw(IRError("$name(): indices required for $(actual_ndim)D view")) + index_jl_type = isempty(unique_types) ? Int32 : unique_types[1] # Int32 only for 0D case index_type = tile_type_for_julia!(ctx, index_jl_type) - index_vals = pad_indices(ctx, index_vals, ndim, index_type, index_jl_type) + + index_vals = pad_indices(ctx, index_vals, actual_ndim, index_type, index_jl_type) reverse!(index_vals) optimization_hints = create_optimization_hints(ctx, latency, allow_tma_val) - tile_val, result_token = encode_LoadViewTkoOp!( - cb, tile_type, token_type, sv_arg.v, index_vals; - token=input_token, optimization_hints, inbounds=fill(!check_bounds, ndim)) - ctx.result_tokens[ctx.current_ssa_idx] = result_token + token_type = Token(tt) - julia_shape = ColMajorShape(tile_shape) - return CGVal(tile_val, tile_type, Tile{elem_type, TupleType(julia_shape)}, tile_shape) + result_token = encode_StoreViewTkoOp!( + cb, token_type, tile_val, view_arg.v, index_vals; + token = input_token, optimization_hints, inbounds=fill(!check_bounds, actual_ndim) + ) + + # Store result token for TokenResultNode + ctx.result_tokens[ctx.current_ssa_idx] = result_token + return nothing end +""" + Intrinsics.load_partition_view(pv::PartitionView{T,N,Shape}, + latency::Union{Int,Nothing}, + allow_tma::Bool, + indices::NTuple{M,<:Integer}) -> Tile{T,Shape} + +Token-ordered load of a tile from a `PartitionView`; lowers to +`cuda_tile.load_view_tko`. + +`latency` and `allow_tma` are compile-time hints. `indices` is in Julia +order and is reversed/zero-padded to match `pv`'s index space rank +before emission. The token argument is appended by `token_order_pass!` +and is not part of the user-visible signature. +""" +@intrinsic load_partition_view(pv, latency, allow_tma, indices, check_bounds) +tfunc(𝕃, ::typeof(Intrinsics.load_partition_view), @nospecialize(pv), @nospecialize args...) = + view_load_return_type(pv, 3) +emit_intrinsic!(ctx::CGCtx, ::typeof(Intrinsics.load_partition_view), args) = + emit_view_load!(ctx, args, "load_partition_view") + +""" + Intrinsics.load_strided_view(sv::StridedView{T,N,Shape,Steps}, ...) + +Token-ordered load from a `StridedView`. It uses the same Tile IR load op as +`load_partition_view`, but has a distinct intrinsic identity so analyses never +assume that distinct tile indices access disjoint memory. +""" +@intrinsic load_strided_view(sv, latency, allow_tma, indices, check_bounds) +tfunc(𝕃, ::typeof(Intrinsics.load_strided_view), @nospecialize(sv), @nospecialize args...) = + view_load_return_type(sv, 4) +emit_intrinsic!(ctx::CGCtx, ::typeof(Intrinsics.load_strided_view), args) = + emit_view_load!(ctx, args, "load_strided_view") + function pad_indices(ctx::CGCtx, index_vals::Vector{Value}, ndim::Int, idx_type::TypeId, idx_jl_type::Type) while length(index_vals) < ndim idx_bytes = reinterpret(UInt8, [eltype(idx_jl_type)(0)]) @@ -296,8 +329,8 @@ end function emit_intrinsic!(ctx::CGCtx, ::typeof(Intrinsics.make_strided_view), args) tensor_view = emit_value!(ctx, args[1]) tensor_view === nothing && throw(IRError("make_strided_view() requires a TensorView argument")) - ctx.tt.version >= v"13.3" || - throw(IRError("make_strided_view requires Tile IR bytecode v13.3+, got v$(ctx.tt.version)")) + # The v13.3 gate is enforced at the bytecode layer by `strided_view_type!` + # and `encode_MakeStridedViewOp!` below; no need to re-check it here. shape = @something get_constant(ctx, args[2]) throw(IRError("make_strided_view() tile_shape must be a compile-time constant")) shape isa Tuple || throw(IRError("make_strided_view() shape must be a tuple, got $(typeof(shape))")) @@ -552,78 +585,8 @@ views require at least 1-D. The token argument is appended by tfunc(𝕃, ::typeof(Intrinsics.store_partition_view), @nospecialize args...) = Nothing efunc(::typeof(Intrinsics.store_partition_view), effects::CC.Effects) = CC.Effects(effects; effect_free=CC.ALWAYS_FALSE) -function emit_intrinsic!(ctx::CGCtx, ::typeof(Intrinsics.store_partition_view), args) - cb = ctx.cb - tt = ctx.tt - - # Extract input token from last arg (added by token_order_pass!) - input_token = extract_token_arg!(ctx, args) - - # args: (partition_view, tile, latency, allow_tma, indices) - pv_arg = emit_value!(ctx, args[1]) - pv_arg === nothing && throw(IRError("store_partition_view() requires a PartitionView argument")) - pv_arg.v === nothing && throw(IRError("store_partition_view() requires a materialized PartitionView")) - - # Get ndim from PartitionView constant field - pv_arg.constant === nothing && throw(IRError("store_partition_view(): PartitionView missing ndim info")) - ndim = something(pv_arg.constant) - - # Get tile value - tile_tv = emit_value!(ctx, args[2]) - tile_tv === nothing && throw(IRError("store_partition_view() requires a tile argument")) - tile_shape = tile_tv.shape - tile_shape === nothing && throw(IRError("Cannot determine tile shape for store_partition_view()")) - - elem_type = eltype(CC.widenconst(tile_tv.jltype)) - dtype = lookup_dtype!(tt, elem_type) - - # Handle 0D scalar stores by reshaping to 1D (partition views require at least 1D) - tile_val = tile_tv.v - actual_ndim = ndim - actual_tile_shape = tile_shape - if length(tile_shape) == 0 - actual_ndim = 1 - actual_tile_shape = RowMajorShape([1]) - tile_1d_type = tile_type!(tt, dtype, actual_tile_shape) - tile_val = encode_ReshapeOp!(cb, tile_1d_type, tile_val) - end - - # Extract optimization hints (args[3] = latency, args[4] = allow_tma) - latency = @something get_constant(ctx, args[3]) throw(IRError("store_partition_view(): latency must be a compile-time constant")) - allow_tma = @something get_constant(ctx, args[4]) throw(IRError("store_partition_view(): allow_tma must be a compile-time constant")) - allow_tma_val = allow_tma isa Bool ? allow_tma : true - check_bounds = @something get_constant(ctx, args[6]) throw(IRError("store_partition_view(): check_bounds must be a compile-time constant")) - - # Extract indices - index_tvs = resolve_tuple(ctx, args[5], "store_partition_view indices") - index_vals = Value[tv.v for tv in index_tvs] - index_jl_types = Type[tv.jltype for tv in index_tvs] - - unique_types = unique(index_jl_types) - length(unique_types) <= 1 || throw(IRError("All index types must match, got: $unique_types")) - isempty(unique_types) && actual_ndim > 0 && throw(IRError("store_partition_view(): indices required for $(actual_ndim)D view")) - index_jl_type = isempty(unique_types) ? Int32 : unique_types[1] # Int32 only for 0D case - index_type = tile_type_for_julia!(ctx, index_jl_type) - - # Pad indices if needed, then reverse for Tile IR row-major order - index_vals = pad_indices(ctx, index_vals, actual_ndim, index_type, index_jl_type) - reverse!(index_vals) - - # Create optimization hints if provided - optimization_hints = create_optimization_hints(ctx, latency, allow_tma_val) - - token_type = Token(tt) - - result_token = encode_StoreViewTkoOp!( - cb, token_type, tile_val, pv_arg.v, index_vals; - token = input_token, optimization_hints, inbounds=fill(!check_bounds, actual_ndim) - ) - - # Store result token for TokenResultNode - ctx.result_tokens[ctx.current_ssa_idx] = result_token - - return nothing -end +emit_intrinsic!(ctx::CGCtx, ::typeof(Intrinsics.store_partition_view), args) = + emit_view_store!(ctx, args, "store_partition_view") """ Intrinsics.store_strided_view(sv::StridedView{T,N,Shape,Steps}, ...) @@ -641,56 +604,5 @@ their loop-carried token dependency even when their tile indices differ. tfunc(𝕃, ::typeof(Intrinsics.store_strided_view), @nospecialize args...) = Nothing efunc(::typeof(Intrinsics.store_strided_view), effects::CC.Effects) = CC.Effects(effects; effect_free=CC.ALWAYS_FALSE) -function emit_intrinsic!(ctx::CGCtx, ::typeof(Intrinsics.store_strided_view), args) - cb = ctx.cb - tt = ctx.tt - input_token = extract_token_arg!(ctx, args) - - sv_arg = emit_value!(ctx, args[1]) - sv_arg === nothing && throw(IRError("store_strided_view() requires a StridedView argument")) - sv_arg.v === nothing && throw(IRError("store_strided_view() requires a materialized StridedView")) - sv_arg.constant === nothing && throw(IRError("store_strided_view(): StridedView missing ndim info")) - ndim = something(sv_arg.constant) - - tile_tv = emit_value!(ctx, args[2]) - tile_tv === nothing && throw(IRError("store_strided_view() requires a tile argument")) - tile_shape = tile_tv.shape - tile_shape === nothing && throw(IRError("Cannot determine tile shape for store_strided_view()")) - elem_type = eltype(CC.widenconst(tile_tv.jltype)) - dtype = lookup_dtype!(tt, elem_type) - - tile_val = tile_tv.v - actual_ndim = ndim - actual_tile_shape = tile_shape - if length(tile_shape) == 0 - actual_ndim = 1 - actual_tile_shape = RowMajorShape([1]) - tile_1d_type = tile_type!(tt, dtype, actual_tile_shape) - tile_val = encode_ReshapeOp!(cb, tile_1d_type, tile_val) - end - - latency = @something get_constant(ctx, args[3]) throw(IRError("store_strided_view(): latency must be a compile-time constant")) - allow_tma = @something get_constant(ctx, args[4]) throw(IRError("store_strided_view(): allow_tma must be a compile-time constant")) - allow_tma_val = allow_tma isa Bool ? allow_tma : true - check_bounds = @something get_constant(ctx, args[6]) throw(IRError("store_strided_view(): check_bounds must be a compile-time constant")) - - index_tvs = resolve_tuple(ctx, args[5], "store_strided_view indices") - index_vals = Value[tv.v for tv in index_tvs] - index_jl_types = Type[tv.jltype for tv in index_tvs] - unique_types = unique(index_jl_types) - length(unique_types) <= 1 || throw(IRError("All index types must match, got: $unique_types")) - isempty(unique_types) && actual_ndim > 0 && - throw(IRError("store_strided_view(): indices required for $(actual_ndim)D view")) - index_jl_type = isempty(unique_types) ? Int32 : unique_types[1] - index_type = tile_type_for_julia!(ctx, index_jl_type) - index_vals = pad_indices(ctx, index_vals, actual_ndim, index_type, index_jl_type) - reverse!(index_vals) - - optimization_hints = create_optimization_hints(ctx, latency, allow_tma_val) - token_type = Token(tt) - result_token = encode_StoreViewTkoOp!( - cb, token_type, tile_val, sv_arg.v, index_vals; - token=input_token, optimization_hints, inbounds=fill(!check_bounds, actual_ndim)) - ctx.result_tokens[ctx.current_ssa_idx] = result_token - return nothing -end +emit_intrinsic!(ctx::CGCtx, ::typeof(Intrinsics.store_strided_view), args) = + emit_view_store!(ctx, args, "store_strided_view") From 3f3e6f50cdd31bb4cacd10bba55c7fb5960e714f Mon Sep 17 00:00:00 2001 From: Tim Besard Date: Thu, 23 Jul 2026 11:08:25 +0200 Subject: [PATCH 4/5] Polish eachtile after review - Drop underscore prefixes from the eachtile helpers and collapse the duplicated step accessor; the type-parameter accessors now live on the type with instance forwarders. - Define Base.eltype on TiledView: the element is a whole tile of the user-requested shape. - Require `step` to match `tile_shape` in length instead of padding short steps with 1s, which silently meant "overlap with step 1" on padded dimensions; a short shape/step pair now pads the trailing dimensions together. - Overlay Base.size on TiledView in kernels to defer the tile count to the backend via get_index_space_shape (matching cuTile Python's num_tiles lowering) instead of baking in the host cld formula; the host method keeps cld for launch-grid sizing. - Return the collection from setindex!, per convention. Co-Authored-By: Claude Fable 5 --- README.md | 8 ++-- src/language/operations.jl | 80 +++++++++++++++++++++++++------------- src/language/types.jl | 30 +++++++++++++- test/codegen/views.jl | 20 ++++++++++ test/device/views.jl | 2 + test/types.jl | 16 +++++++- 6 files changed, 124 insertions(+), 32 deletions(-) diff --git a/README.md b/README.md index 64043244..be890a1b 100644 --- a/README.md +++ b/README.md @@ -352,9 +352,11 @@ end ### Tile Windows `eachtile` creates a small, indexable device-side collection of fixed-shape -tiles. Its indices are 1-based and `size(eachtile(a, shape; step), d)` is -`cld(size(a, d), step[d])`. `step` controls tile origins, not the element -stride inside a tile: +tiles. Its indices are 1-based and `step` (one entry per tile dimension) +controls tile origins, not the element stride inside a tile. `size(tiles, d)` +is the number of tiles along `d`: on the host it computes +`cld(size(a, d), step[d])` for launch-grid sizing, while inside a kernel it +queries the Tile IR backend for the authoritative index-space count: ```julia adjacent = eachtile(a, (8, 8)) # default: step == (8, 8) diff --git a/src/language/operations.jl b/src/language/operations.jl index 393720b2..edfc6a13 100644 --- a/src/language/operations.jl +++ b/src/language/operations.jl @@ -341,7 +341,9 @@ end return :($kept) end -function _normalize_eachtile_shape(shape::Tuple, rank::Int, name::String) +# Pad a short tile shape with trailing 1s, or squeeze trailing singletons, to +# reach the array rank. Rejects a shape whose non-singleton tail can't be dropped. +function normalize_eachtile_shape(shape::Tuple, rank::Int, name::String) length(shape) == rank && return shape if length(shape) < rank return (shape..., ntuple(_ -> 1, rank - length(shape))...) @@ -353,16 +355,36 @@ function _normalize_eachtile_shape(shape::Tuple, rank::Int, name::String) return shape[1:rank] end -@generated function _eachtile(a::A, ::Val{Requested}, ::Val{Step}, ::Val{Padding}) where +# Normalize a user `step` (already length-matched to the user `tile_shape`) to +# the array rank, following the same pad/squeeze as the tile shape. Padded +# (trailing) dimensions default to the corresponding `view_shape` entry — the +# "step defaults to tile shape" rule applied per dimension — so a short +# `tile_shape`/`step` pair stays consistent (padded dims are unit-step +# singletons). Squeezed trailing dimensions (already validated as singletons by +# the tile-shape normalization) simply drop their step entries. +function normalize_eachtile_step(step::Tuple, view_shape::Tuple, rank::Int) + M = length(step) + M == rank && return step + M < rank && return (step..., ntuple(i -> view_shape[M + i], rank - M)...) + return step[1:rank] +end + +@generated function build_tiled_view(a::A, ::Val{Requested}, ::Val{Step}, ::Val{Padding}) where {A<:TileArray, Requested, Step, Padding} Requested isa Tuple || throw(ArgumentError( "eachtile: tile_shape must be a compile-time tuple, got $(typeof(Requested))")) rank = ndims(A) - view_shape = _normalize_eachtile_shape(Requested, rank, "tile_shape") + view_shape = normalize_eachtile_shape(Requested, rank, "tile_shape") normalized_step = if Step === nothing view_shape elseif Step isa Tuple - _normalize_eachtile_shape(Step, rank, "step") + # A short `step` must not silently overlap the padded tile dimensions + # (Python requires `traversal_steps` at full rank). Require it to carry + # one entry per user tile dimension, then normalize both together. + length(Step) == length(Requested) || throw(ArgumentError( + "eachtile: step $Step must match tile_shape $Requested in length " * + "($(length(Step)) vs $(length(Requested))); pass one step per tile dimension")) + normalize_eachtile_step(Step, view_shape, rank) else throw(ArgumentError("eachtile: step must be a compile-time tuple or nothing, got $(typeof(Step))")) end @@ -373,13 +395,6 @@ end return :($result_type(a)) end -_tiled_view_requested_shape(::TiledView{A, RequestedShape}) where {A, RequestedShape} = - Tuple(RequestedShape.parameters) -_tiled_view_shape(::TiledView{A, RequestedShape, ViewShape}) where - {A, RequestedShape, ViewShape} = Tuple(ViewShape.parameters) -_tiled_view_step(::TiledView{A, RequestedShape, ViewShape, Step}) where - {A, RequestedShape, ViewShape, Step} = Tuple(Step.parameters) - """ eachtile(a, tile_shape; step=nothing, padding_mode=PaddingMode.Undetermined) @@ -397,29 +412,29 @@ given padding mode on loads and clipped stores. @inline function eachtile(a::TileArray, tile_shape::Tuple; step::Union{Tuple, Nothing}=nothing, padding_mode::PaddingMode.T=PaddingMode.Undetermined) - _eachtile(a, Val(tile_shape), Val(step), Val(padding_mode)) + build_tiled_view(a, Val(tile_shape), Val(step), Val(padding_mode)) end -@inline function _make_tile_view(tiles::TiledView{A, RequestedShape, Shape, Shape, Padding}) where +@inline function make_tile_view(tiles::TiledView{A, RequestedShape, Shape, Shape, Padding}) where {A, RequestedShape, Shape, Padding} parent = tiles.parent tv = Intrinsics.make_tensor_view(typeof(parent), parent.ptr, parent.sizes, parent.strides) - Intrinsics.make_partition_view(tv, _tiled_view_shape(tiles), Padding, nothing) + Intrinsics.make_partition_view(tv, tiled_view_shape(tiles), Padding, nothing) end -@inline function _make_tile_view(tiles::TiledView{A, RequestedShape, Shape, Step, Padding}) where +@inline function make_tile_view(tiles::TiledView{A, RequestedShape, Shape, Step, Padding}) where {A, RequestedShape, Shape, Step, Padding} parent = tiles.parent tv = Intrinsics.make_tensor_view(typeof(parent), parent.ptr, parent.sizes, parent.strides) - Intrinsics.make_strided_view(tv, _tiled_view_shape(tiles), _tiled_view_step(tiles), Padding, nothing) + Intrinsics.make_strided_view(tv, tiled_view_shape(tiles), tiled_view_step(tiles), Padding, nothing) end -@inline _load_tile_view(view::PartitionView, latency, allow_tma, indices, check_bounds) = +@inline load_tile_view(view::PartitionView, latency, allow_tma, indices, check_bounds) = Intrinsics.load_partition_view(view, latency, allow_tma, indices, check_bounds) -@inline _load_tile_view(view::StridedView, latency, allow_tma, indices, check_bounds) = +@inline load_tile_view(view::StridedView, latency, allow_tma, indices, check_bounds) = Intrinsics.load_strided_view(view, latency, allow_tma, indices, check_bounds) -@inline _store_tile_view(view::PartitionView, tile, latency, allow_tma, indices, check_bounds) = +@inline store_tile_view(view::PartitionView, tile, latency, allow_tma, indices, check_bounds) = Intrinsics.store_partition_view(view, tile, latency, allow_tma, indices, check_bounds) -@inline _store_tile_view(view::StridedView, tile, latency, allow_tma, indices, check_bounds) = +@inline store_tile_view(view::StridedView, tile, latency, allow_tma, indices, check_bounds) = Intrinsics.store_strided_view(view, tile, latency, allow_tma, indices, check_bounds) @inline function load(tiles::TiledView{A}, index::NTuple{N, <:Integer}; @@ -427,9 +442,9 @@ end latency::Union{Int, Nothing}=nothing, allow_tma::Union{Bool, Nothing}=nothing) where {A, N} N == ndims(tiles) || throw(ArgumentError("eachtile: expected $(ndims(tiles)) tile indices, got $N")) - view = _make_tile_view(tiles) - tile = _load_tile_view(view, latency, allow_tma, promote(index...) .- One(), check_bounds) - reshape(tile, _tiled_view_requested_shape(tiles)) + view = make_tile_view(tiles) + tile = load_tile_view(view, latency, allow_tma, promote(index...) .- One(), check_bounds) + reshape(tile, tiled_view_requested_shape(tiles)) end @inline function load(tiles::TiledView, index::Integer; kwargs...) load(tiles, (index,); kwargs...) @@ -444,8 +459,8 @@ end allow_tma::Union{Bool, Nothing}=nothing) where {A, N, T} N == ndims(tiles) || throw(ArgumentError("eachtile: expected $(ndims(tiles)) tile indices, got $N")) reshaped = _reshape_to_rank(tile, Val(ndims(tiles))) - view = _make_tile_view(tiles) - _store_tile_view(view, reshaped, latency, allow_tma, promote(index...) .- One(), check_bounds) + view = make_tile_view(tiles) + store_tile_view(view, reshaped, latency, allow_tma, promote(index...) .- One(), check_bounds) return tile end @inline function store(tiles::TiledView, index::Integer, tile::Tile; kwargs...) @@ -455,13 +470,26 @@ end store(tiles, index, tile; kwargs...) end +# In kernels, `size` on a TiledView defers the tile count to the Tile IR +# backend (`get_index_space_shape`) instead of baking in the host `cld` +# formula; this matches cuTile Python's `TiledView.num_tiles` lowering. +@overlay function Base.size(tiles::TiledView) + view = make_tile_view(tiles) + ntuple(i -> Intrinsics.get_index_space_shape(view, i - One()), Val(ndims(tiles))) +end +@overlay function Base.size(tiles::TiledView, d::Integer) + d > ndims(tiles) && return Int32(1) + view = make_tile_view(tiles) + Intrinsics.get_index_space_shape(view, d - One()) +end + @overlay function Base.getindex(tiles::TiledView{A}, indices::Vararg{Integer, N}) where {A, N} load(tiles, indices) end Base.Experimental.@consistent_overlay cuTileMethodTable function Base.setindex!( tiles::TiledView{A}, tile::Tile, indices::Vararg{Integer, N}) where {A, N} store(tiles, indices, tile) - return + return tiles end # Reshape a tile to match target rank N, preserving data layout. diff --git a/src/language/types.jl b/src/language/types.jl index dae36823..09ac45e9 100644 --- a/src/language/types.jl +++ b/src/language/types.jl @@ -424,11 +424,37 @@ Base.parent(tiles::TiledView) = tiles.parent Base.ndims(::Type{<:TiledView{A}}) where {A} = ndims(A) Base.ndims(tiles::TiledView) = ndims(typeof(tiles)) -_tiled_view_steps(::Type{<:TiledView{A, RequestedShape, ViewShape, Step}}) where +# The element of a `TiledView` is a whole tile of the user-requested shape: +# `load` reshapes the view tile back to `RequestedShape`, which may be a lower +# rank than the rank-normalized `ViewShape`. +Base.eltype(::Type{<:TiledView{A, RequestedShape}}) where + {T, A<:TileArray{T}, RequestedShape} = Tile{T, RequestedShape} +Base.eltype(tiles::TiledView) = eltype(typeof(tiles)) + +# Type-parameter accessors. Defined on the type (single source of truth) with +# an instance forwarder for convenience. +tiled_view_requested_shape(::Type{<:TiledView{A, RequestedShape}}) where {A, RequestedShape} = + Tuple(RequestedShape.parameters) +tiled_view_shape(::Type{<:TiledView{A, RequestedShape, ViewShape}}) where + {A, RequestedShape, ViewShape} = Tuple(ViewShape.parameters) +tiled_view_step(::Type{<:TiledView{A, RequestedShape, ViewShape, Step}}) where {A, RequestedShape, ViewShape, Step} = Tuple(Step.parameters) +tiled_view_requested_shape(tiles::TiledView) = tiled_view_requested_shape(typeof(tiles)) +tiled_view_shape(tiles::TiledView) = tiled_view_shape(typeof(tiles)) +tiled_view_step(tiles::TiledView) = tiled_view_step(typeof(tiles)) +""" + size(tiles::TiledView[, d]) + +Number of tiles along each axis. On the host this is computed as +`cld(size(parent, d), step[d])` for launch-grid sizing. Inside kernels the +same call is overlaid to query the Tile IR backend via +`get_index_space_shape` (matching cuTile Python's `TiledView.num_tiles` +lowering), since the strided index-space formula is the backend's contract, +not ours. +""" function Base.size(tiles::TiledView) - steps = _tiled_view_steps(typeof(tiles)) + steps = tiled_view_step(typeof(tiles)) ntuple(i -> cld(size(tiles.parent, i), Int32(steps[i])), Val(ndims(tiles))) end function Base.size(tiles::TiledView, d::Integer) diff --git a/test/codegen/views.jl b/test/codegen/views.jl index de93425f..fa99d2d8 100644 --- a/test/codegen/views.jl +++ b/test/codegen/views.jl @@ -118,3 +118,23 @@ end return end end + +@testset "eachtile — device size queries the backend index space" begin + # `size(tiles, d)` computes `cld` on the host, but in kernels it is + # overlaid to defer to `get_index_space_shape` (matching cuTile Python's + # `num_tiles` lowering). Feed the result into a load index so it isn't + # DCE'd. + @test @filecheck begin + @check_label "entry" + @check "make_strided_view" + @check "get_index_space_shape" + @check "load_view_tko" + @check "store_view_tko" + code_tiled(Tuple{ct.TileArray{Float32,1,spec1d}}; bytecode_version=v"13.3") do a + tiles = eachtile(a, (8,); step=(4,)) + n = size(tiles, 1) + ct.store(a, 1, ct.load(tiles, n)) + return + end + end +end diff --git a/test/device/views.jl b/test/device/views.jl index 154216b3..02f59d0d 100644 --- a/test/device/views.jl +++ b/test/device/views.jl @@ -92,6 +92,8 @@ end function copy_gaps(a::ct.TileArray{Float32,1}, b::ct.TileArray{Float32,1}) src = eachtile(a, (4,); step=(8,)) dst = eachtile(b, (4,); step=(8,)) + # On device, `size` is overlaid to query the backend index space + # (get_index_space_shape) rather than baking in `cld`. for i in 1:size(src, 1) dst[i] = src[i] end diff --git a/test/types.jl b/test/types.jl index fafcbddb..97933335 100644 --- a/test/types.jl +++ b/test/types.jl @@ -80,8 +80,22 @@ end @test size(tiles, 1) == Int32(6) @test size(tiles, 2) == Int32(5) @test :eachtile in names(cuTile) + + # The element of the collection is a whole tile of the requested shape. + @test eltype(tiles) == ct.Tile{Float32, Tuple{8, 4}} + @test eltype(typeof(tiles)) == ct.Tile{Float32, Tuple{8, 4}} + @test_throws "strictly positive" eachtile(a, (8, 4); step=(0, 2)) - @test_throws "cannot squeeze step" eachtile(a, (8, 4); step=(2, 2, 2)) + # `step` must carry one entry per user tile dimension: a mismatched length + # is rejected rather than silently overlapping the padded dimensions. + @test_throws "must match tile_shape" eachtile(a, (8, 4); step=(2, 2, 2)) + @test_throws "must match tile_shape" eachtile(a, (4, 4); step=(2,)) + + # A short tile_shape/step pair on a higher-rank array pads consistently: + # the trailing dimension is a unit-step singleton, not a step-1 overlap. + short = eachtile(a, (4,); step=(2,)) + @test cuTile.tiled_view_shape(short) == (4, 1) + @test cuTile.tiled_view_step(short) == (2, 1) end @testset "TensorView" begin From 2277e7bda6d5ca3d9d30a2eeecd2b859b62c884a Mon Sep 17 00:00:00 2001 From: Tim Besard Date: Thu, 23 Jul 2026 11:08:44 +0200 Subject: [PATCH 5/5] Support dimension order in eachtile Expose the view dim_map through the same 1-indexed `order` kwarg that ct.load/ct.store already provide: tile_shape[i], step[i], and tile index i describe tile dimension i, which maps to array dimension order[i]. The permutation is carried as a sixth TiledView type parameter and threaded into both the partition and strided view constructors, making the previously unreachable make_strided_view dim_map arm exercisable from the public API. The host-side size(tiles) pairs step[i] with array dim order[i]; the device-side overlay needs no change since get_index_space_shape respects the view's dim_map in the backend. Co-Authored-By: Claude Fable 5 --- README.md | 5 +++-- src/language/operations.jl | 30 ++++++++++++++++++++++-------- src/language/types.jl | 19 +++++++++++++------ test/codegen/views.jl | 30 ++++++++++++++++++++++++++++++ test/device/views.jl | 20 ++++++++++++++++++++ test/types.jl | 7 +++++++ 6 files changed, 95 insertions(+), 16 deletions(-) diff --git a/README.md b/README.md index be890a1b..5eb21263 100644 --- a/README.md +++ b/README.md @@ -367,8 +367,9 @@ tile = overlap[2, 1] overlap[2, 1] = tile ``` -Use `ct.load` and `ct.store` for `check_bounds`, `latency`, and `allow_tma` -controls. Equal shape and step use the ordinary Tile IR partition view; +`eachtile` also accepts the same `order` kwarg as `ct.load`/`ct.store` to +permute which array dimension each tile dimension walks. Use `ct.load` and +`ct.store` for `check_bounds`, `latency`, and `allow_tma` controls. Equal shape and step use the ordinary Tile IR partition view; unequal values require Tile IR bytecode v13.3 or newer. This is distinct from `@view a[1:2:end, :]`, which steps individual elements rather than tile origins. diff --git a/src/language/operations.jl b/src/language/operations.jl index edfc6a13..a9b3ed70 100644 --- a/src/language/operations.jl +++ b/src/language/operations.jl @@ -369,8 +369,9 @@ function normalize_eachtile_step(step::Tuple, view_shape::Tuple, rank::Int) return step[1:rank] end -@generated function build_tiled_view(a::A, ::Val{Requested}, ::Val{Step}, ::Val{Padding}) where - {A<:TileArray, Requested, Step, Padding} +@generated function build_tiled_view(a::A, ::Val{Requested}, ::Val{Step}, ::Val{Padding}, + ::Val{Order}) where + {A<:TileArray, Requested, Step, Padding, Order} Requested isa Tuple || throw(ArgumentError( "eachtile: tile_shape must be a compile-time tuple, got $(typeof(Requested))")) rank = ndims(A) @@ -390,13 +391,19 @@ end end all(step -> step isa Integer && step > 0, normalized_step) || throw(ArgumentError("eachtile: step must contain strictly positive integers, got $normalized_step")) + if Order !== nothing + Order isa Tuple || + throw(ArgumentError("eachtile: order must be a compile-time tuple or nothing, got $(typeof(Order))")) + length(Order) == rank && all(o -> o isa Integer, Order) && isperm(Order) || + throw(ArgumentError("eachtile: order must be a permutation of 1:$rank, got $Order")) + end result_type = TiledView{A, Tuple{Requested...}, Tuple{view_shape...}, - Tuple{normalized_step...}, Padding} + Tuple{normalized_step...}, Padding, Order} return :($result_type(a)) end """ - eachtile(a, tile_shape; step=nothing, + eachtile(a, tile_shape; step=nothing, order=nothing, padding_mode=PaddingMode.Undetermined) Create an immutable device-side collection of fixed-shape tiles over `a`. @@ -405,27 +412,34 @@ Create an immutable device-side collection of fixed-shape tiles over `a`. larger values leave gaps. Unlike `@view a[1:2:end]`, which changes element strides inside an array, `eachtile` changes tile origins. -Unequal tile shape and step require Tile IR bytecode v13.3 or newer. Both are +`order` is the same 1-indexed logical-to-physical dimension mapping as the +`order` kwarg of [`load`](@ref)/[`store`](@ref): `tile_shape[i]`, `step[i]`, +and tile index `i` describe tile dimension `i`, which maps to array dimension +`order[i]`. + +Unequal tile shape and step require Tile IR bytecode v13.3 or newer. All are compile-time tuples; tile indices are 1-based and partial edge tiles use the given padding mode on loads and clipped stores. """ @inline function eachtile(a::TileArray, tile_shape::Tuple; step::Union{Tuple, Nothing}=nothing, + order::Union{Tuple, Nothing}=nothing, padding_mode::PaddingMode.T=PaddingMode.Undetermined) - build_tiled_view(a, Val(tile_shape), Val(step), Val(padding_mode)) + build_tiled_view(a, Val(tile_shape), Val(step), Val(padding_mode), Val(order)) end @inline function make_tile_view(tiles::TiledView{A, RequestedShape, Shape, Shape, Padding}) where {A, RequestedShape, Shape, Padding} parent = tiles.parent tv = Intrinsics.make_tensor_view(typeof(parent), parent.ptr, parent.sizes, parent.strides) - Intrinsics.make_partition_view(tv, tiled_view_shape(tiles), Padding, nothing) + Intrinsics.make_partition_view(tv, tiled_view_shape(tiles), Padding, tiled_view_order(tiles)) end @inline function make_tile_view(tiles::TiledView{A, RequestedShape, Shape, Step, Padding}) where {A, RequestedShape, Shape, Step, Padding} parent = tiles.parent tv = Intrinsics.make_tensor_view(typeof(parent), parent.ptr, parent.sizes, parent.strides) - Intrinsics.make_strided_view(tv, tiled_view_shape(tiles), tiled_view_step(tiles), Padding, nothing) + Intrinsics.make_strided_view(tv, tiled_view_shape(tiles), tiled_view_step(tiles), Padding, + tiled_view_order(tiles)) end @inline load_tile_view(view::PartitionView, latency, allow_tma, indices, check_bounds) = diff --git a/src/language/types.jl b/src/language/types.jl index 09ac45e9..477d7152 100644 --- a/src/language/types.jl +++ b/src/language/types.jl @@ -410,13 +410,15 @@ Base.size(sv::StridedView) = size(typeof(sv)) Base.size(sv::StridedView, d::Integer) = size(typeof(sv), d) """ - TiledView{A, RequestedShape, ViewShape, Step, Padding} + TiledView{A, RequestedShape, ViewShape, Step, Padding, Order} -Internal immutable array-of-tiles wrapper returned by `eachtile`. Shape and -step live in the type because Tile IR view types require compile-time values; -the parent is the sole runtime field. +Internal immutable array-of-tiles wrapper returned by `eachtile`. Shape, step, +and order live in the type because Tile IR view types require compile-time +values; the parent is the sole runtime field. `Order` is the 1-indexed +tile-dim-to-array-dim permutation (or `nothing` for identity), matching the +`order` kwarg of `ct.load`/`ct.store`. """ -struct TiledView{A, RequestedShape, ViewShape, Step, Padding} +struct TiledView{A, RequestedShape, ViewShape, Step, Padding, Order} parent::A end @@ -439,9 +441,12 @@ tiled_view_shape(::Type{<:TiledView{A, RequestedShape, ViewShape}}) where {A, RequestedShape, ViewShape} = Tuple(ViewShape.parameters) tiled_view_step(::Type{<:TiledView{A, RequestedShape, ViewShape, Step}}) where {A, RequestedShape, ViewShape, Step} = Tuple(Step.parameters) +tiled_view_order(::Type{<:TiledView{A, RequestedShape, ViewShape, Step, Padding, Order}}) where + {A, RequestedShape, ViewShape, Step, Padding, Order} = Order tiled_view_requested_shape(tiles::TiledView) = tiled_view_requested_shape(typeof(tiles)) tiled_view_shape(tiles::TiledView) = tiled_view_shape(typeof(tiles)) tiled_view_step(tiles::TiledView) = tiled_view_step(typeof(tiles)) +tiled_view_order(tiles::TiledView) = tiled_view_order(typeof(tiles)) """ size(tiles::TiledView[, d]) @@ -455,7 +460,9 @@ not ours. """ function Base.size(tiles::TiledView) steps = tiled_view_step(typeof(tiles)) - ntuple(i -> cld(size(tiles.parent, i), Int32(steps[i])), Val(ndims(tiles))) + order = tiled_view_order(typeof(tiles)) + dims = something(order, ntuple(identity, Val(ndims(tiles)))) + ntuple(i -> cld(size(tiles.parent, dims[i]), Int32(steps[i])), Val(ndims(tiles))) end function Base.size(tiles::TiledView, d::Integer) d < 1 && error("arraysize: dimension out of range") diff --git a/test/codegen/views.jl b/test/codegen/views.jl index fa99d2d8..8b3a4397 100644 --- a/test/codegen/views.jl +++ b/test/codegen/views.jl @@ -138,3 +138,33 @@ end end end end + +@testset "eachtile — order permutes the view dim_map" begin + # `order` follows ct.load/ct.store: tile dim i maps to array dim order[i]. + # Julia 1-indexed order (2, 1) reverses to the row-major dim_map [1, 0]. + # Equal shape/step keeps the PartitionView arm. + @test @filecheck begin + @check_label "entry" + @check "make_partition_view" + @check "dim_map=[1, 0]" + @check "load_view_tko" + code_tiled(Tuple{ct.TileArray{Float32,2,spec2d}}; bytecode_version=v"13.3") do a + tiles = eachtile(a, (4, 8); order=(2, 1)) + ct.store(a, (1, 1), ct.load(tiles, (1, 1))) + return + end + end + + # Unequal step exercises the StridedView arm. + @test @filecheck begin + @check_label "entry" + @check "make_strided_view" + @check "dim_map=[1, 0]" + @check "load_view_tko" + code_tiled(Tuple{ct.TileArray{Float32,2,spec2d}}; bytecode_version=v"13.3") do a + tiles = eachtile(a, (4, 8); step=(2, 8), order=(2, 1)) + ct.store(a, (1, 1), ct.load(tiles, (1, 1))) + return + end + end +end diff --git a/test/device/views.jl b/test/device/views.jl index 02f59d0d..cb350d7d 100644 --- a/test/device/views.jl +++ b/test/device/views.jl @@ -124,6 +124,26 @@ end @test Array(b) == expected end +@testset "eachtile — order permutes tile dimensions" begin + # order=(2, 1): tile dim 1 (extent 4, step 3) walks array dim 2, tile dim 2 + # (extent 8, step 2) walks array dim 1. Index (2, 3) selects the window + # cols 4:7 (origin (2-1)*3) and rows 5:12 (origin (3-1)*2). A load/store + # roundtrip through identically-parameterized views copies that window. + function copy_permuted(a::ct.TileArray{Float32,2}, b::ct.TileArray{Float32,2}) + src = eachtile(a, (4, 8); step=(3, 2), order=(2, 1), padding_mode=ct.PaddingMode.Zero) + dst = eachtile(b, (4, 8); step=(3, 2), order=(2, 1)) + dst[2, 3] = src[2, 3] + return + end + + a = CUDA.rand(Float32, 12, 12) + b = CUDA.zeros(Float32, 12, 12) + @cuda backend=cuTile copy_permuted(a, b) + expected = zeros(Float32, 12, 12) + expected[5:12, 4:7] .= Array(a)[5:12, 4:7] + @test Array(b) == expected +end + @testset "eachtile — partial edges and requested-rank normalization" begin function copy_partial(a::ct.TileArray{Float32,1}, b::ct.TileArray{Float32,1}) src = eachtile(a, (4,); step=(4,), padding_mode=ct.PaddingMode.Zero) diff --git a/test/types.jl b/test/types.jl index 97933335..9f94ac60 100644 --- a/test/types.jl +++ b/test/types.jl @@ -96,6 +96,13 @@ end short = eachtile(a, (4,); step=(2,)) @test cuTile.tiled_view_shape(short) == (4, 1) @test cuTile.tiled_view_step(short) == (2, 1) + + # `order` permutes which array dimension each tile dimension walks, so the + # host tile count pairs step[i] with array dim order[i]. a is 16×10. + permuted = eachtile(a, (8, 4); step=(3, 2), order=(2, 1)) + @test size(permuted) == (Int32(4), Int32(8)) # cld(10, 3), cld(16, 2) + @test_throws "must be a permutation" eachtile(a, (8, 4); order=(1, 1)) + @test_throws "must be a permutation" eachtile(a, (8, 4); order=(1,)) end @testset "TensorView" begin