Tile kernel examples
These examples are the actual C++ test definitions, included directly from the test sources so the documented syntax cannot drift into a second DSL. All operator POCs use subtiles and Tile SSA. Scalar memory-access tests are separate, explicitly low-level tests.
The native TVMx bridge provides a correctness/reference schedule plus guarded
optimizations. An mma remains an MMA operation in TileIR; its reference TIRx
realization is a contraction loop. With an explicit device-capability opt-in,
eligible FP32 Metal group operations select native SIMD-group matrix atoms.
This is not a performance-parity claim. Safe pipeline cuts can use two-window
software prefetching; hardware-asynchronous transfers and warp specialization
remain future work.
Migrated examples and low precision
examples/compute/tile/kernels.h contains current execution-first counterparts
of the restored legacy benchmark: copy, add, SAXPY, clamp, exp, unweighted
RMSNorm, sum/max/min/absolute reductions, inclusive sum/max scan, transpose and
GEMM. They use current TensorView parameters, explicit stores and ordinary
staging values, not the old Tile frontend. The legacy SAXPY expression is
a / 0.4f + b; RMSNorm retains the old 1e-12f epsilon and no scale weight.
Scans use a library-composed blocked inclusive scan, not a new primitive or
a claim of hardware prefix-scan lowering.
The shared GEMM capture accepts gemm<half>(m, n, k, block) or
gemm<bf16>(...), with FP32 accumulation and explicit output conversion;
gemm<int8_t, int32_t>(...) chooses an integer accumulator. The same
execution structure serves each dtype. See precision semantics.
test_tile_migrated validates captures/oracles without a backend; the simd
and metal variants check 45 complete output/guard cases. test_tile_types
adds low-precision storage, conversion, mixed MMA and explicit FP8 rejection
tests. The opt-in benchmark_tile_migrated accepts the 15 operation names plus
gemm_fp16, gemm_bf16, and a separately labelled full-K gemm_direct, with
separate shape, block and timing arguments.
The migration report records comparisons
with the original lowering, including regressions and unsupported cases.
1. One memory-access convention
Expression |
Result / effect |
|---|---|
|
|
|
The same reference |
|
|
|
The same Tile load |
|
Explicit write |
|
A new value definition, never a memory write |
origin is coord(...). shape(...) accepts extents or kernel-local
axis("diagnostic_name", extent) values. Shared axis identities express
broadcasting and contraction; no axis name is predefined by the language.
The default bounds::zero fills out-of-bounds reads with zero and masks
out-of-bounds stores. bounds::assume is an explicit in-bounds caller contract.
For another fill value, use A(origin, shape).load(fallback).
C++23 uses native multi-argument operator[]. In C++20 the comma operator,
overloaded only for DSL coordinates plus a shape and optional bounds policy,
forms one typed selection argument. Both paths call the identical load builder:
auto x = A[coord(m0, k0), shape(m, k)];
auto y = B[coord(k0, n0), shape(k, n), bounds::zero];
C(coord(m0, n0), shape(m, n)).store(mma(x, y, acc));
Some compilers diagnose the pre-C++23 comma-subscript grammar as deprecated even when the comma is overloaded. With warning-as-error builds, suppress only that diagnostic in the C++20 client. No global comma overload or header-wide diagnostic suppression is installed.
The project minimum remains C++20. CMake adds the C++23 companion test when
the toolchain supports that language level; XMake exposes it as the opt-in
test_tile_values_cpp23 target.
2. Execution-first GEMM
The outer range enumerates logical output programs; the inner pipeline
enumerates K tiles. The shape of a, b, or acc does not create an
execution or memory hierarchy. Assignment to acc captures loop-carried SSA
state without loop.result().
An explicit outer.index(...) resolves within that Nest and its ancestors,
even when written inside a child. Reusing positional axes in a nested shape
does not make the child silently replace the parent’s coordinate.
Use nest.index(axis) for a named logical coordinate, or nest.index() for a
rank-one domain. Nest has no subscript overload; A[origin, shape] is only
the Tile-load syntax for tensor views.
Value selection uses ite(condition, true_value, false_value), matching the
SIMT DSL’s ite. Both Scalar and Tile operands use this order; the Tile DSL
does not expose the differently ordered select name.
struct GemmConfig {
uint32_t block_m{8u};
uint32_t block_n{8u};
uint32_t block_k{8u};
uint32_t window{2u};
};
// Host configuration is captured normally. Each capture/JIT with a different
// configuration records a new TileIR candidate, with no symbolic meta-DSL.
auto make_gemm(GemmConfig cfg) {
return tile_kernel("gemm", [=](TensorView<const float, 2> A,
TensorView<const float, 2> B,
TensorView<float, 2> C) {
auto gm = axis("block_m", ceil_div(A.extent<0>(), cfg.block_m));
auto gn = axis("block_n", ceil_div(B.extent<1>(), cfg.block_n));
auto kt = axis("k_tiles", ceil_div(A.extent<1>(), cfg.block_k));
auto m = axis("m", cfg.block_m);
auto n = axis("n", cfg.block_n);
auto k = axis("k", cfg.block_k);
for (auto &nest : parallel(shape(gm, gn))) {
auto m0 = nest.index(gm) * cfg.block_m;
auto n0 = nest.index(gn) * cfg.block_n;
auto acc = zeros<float>(shape(m, n));
for (auto &step : nest.pipeline(shape(kt), {.window = cfg.window,
.interval = 1u})) {
auto k0 = step.index() * cfg.block_k;
step.stage("load");
auto a = A[coord(m0, k0), shape(m, k)];
auto b = B[coord(k0, n0), shape(k, n)];
step.stage("compute");
acc = mma(a, b, acc);
}
C(coord(m0, n0), shape(m, n)).store(acc);
}
});
}
parallel(output blocks)
├─ acc: Tile<M,N>
└─ pipeline(K tiles)
├─ load: A[M,K], B[K,N] → Tile values
├─ compute: mma(A, B, acc) → next acc
└─ inferred carry → next K iteration
└─ C(origin, shape).store(acc)
2.1 Optional explicit Memory
The execution structure is unchanged. Two sibling resources are declared in
the output program and reused by the pipeline. store consumes a Tile and
load produces a snapshot; assigning a Tile never means storing to Memory.
The frontend infers each resource’s loop-carried MemoryState alongside acc.
The example uses padded rows for As and column-major local storage for Bs;
their logical MMA operand shapes remain unchanged.
// Explicit Memory is optional: use it only when stable address identity or
// a resource constraint is required. The ordinary make_gemm remains unchanged.
auto make_manual_gemm(GemmConfig cfg, exec::Scope scope = exec::Scope::AUTOMATIC,
mem::Resource resource = mem::auto_) {
return tile_kernel("manual_gemm", [=](TensorView<const float, 2> A,
TensorView<const float, 2> B,
TensorView<float, 2> C) {
auto gm = axis("block_m", ceil_div(A.extent<0>(), cfg.block_m));
auto gn = axis("block_n", ceil_div(B.extent<1>(), cfg.block_n));
auto kt = axis("k_tiles", ceil_div(A.extent<1>(), cfg.block_k));
auto m = axis("m", cfg.block_m);
auto n = axis("n", cfg.block_n);
auto k = axis("k", cfg.block_k);
for (auto &nest : parallel(shape(gm, gn), scope)) {
auto m0 = nest.index(gm) * cfg.block_m;
auto n0 = nest.index(gn) * cfg.block_n;
// Distinct local address maps; the logical operand shapes and
// their common execution owner do not change.
auto As = memory<float>(layout(shape(m, k), stride(cfg.block_k + 1, 1)), resource);
auto Bs = memory<float>(layout(shape(k, n), stride(1, cfg.block_k)), resource);
auto acc = zeros<float>(shape(m, n));
for (auto &step : nest.pipeline(shape(kt), {.window = cfg.window,
.interval = 1u})) {
auto k0 = step.index() * cfg.block_k;
step.stage("load");
As.store(A[coord(m0, k0), shape(m, k)]);
Bs.store(B[coord(k0, n0), shape(k, n)]);
step.stage("compute");
acc = mma(As.load(), Bs.load(), acc);
}
C(coord(m0, n0), shape(m, n)).store(acc);
}
});
}
For the current native reference mapper, choose exec::Scope::GROUP with
mem::shared on Metal, or exec::Scope::WORKER with mem::private_ on CPU.
Leaving the resource unspecified lets the mapper choose it. A resource
constraint does not change ownership; unsupported placements fail explicitly.
Use memory<T>(layout(shape, stride(...)), resource) for ordinary strided
storage, or pass a composed IndexMap for a custom address map. A whole-Tile
Memory layout must be total, in bounds, and injective. IndexMap::prove()
proves ordinary affine and structurally bit-linear layouts without enumerating
their elements. Gray-code/XOR swizzles, bit transposes, constant masks, and
power-of-two digit packing use exact GF(2) rank and image bounds. Remaining
maps use an exhaustive fallback of up to 1,048,576 logical points; an unknown
proof fails native realization, not structural representation. Range-aware
parallel writes and asynchronous pipeline versions are not implemented yet.
The native bridge does implement synchronous software prefetching across safe
stage cuts, with two versions for cross-phase iteration-local temporaries.
An outer mutable Memory resource is not silently renamed into an independent
per-iteration resource: cross-phase dependencies keep its execution ordered.
test_tile_tirx_pipeline exercises real versioned CPU/Metal execution, short
and ragged loops, multiple carried values, possible input/output aliasing,
late stores, and shared-capacity fallback. Hardware-asynchronous transfers and
warp specialization remain separate work; see the
implemented software-prefetch path.
test_tile_tirx_memory runs ragged manual GEMMs, multiple independently updated
resources, old-value snapshots, nested temporal regions, ancestor reads,
worker-private state, and resource/capacity rejection on CPU and Metal. It also
checks padded/transposed/composed-XOR layouts and logical shifts through the
sign bit, including physical allocation shapes, padding capacity, and empty
domains. LLVM also tests vector-private state and 513-by-2047 Memory snapshots
with identity, padded, and transposed layouts, above the enumeration budget.
The same large layouts export on Metal but correctly exceed its threadgroup
storage capacity; this is a resource constraint, not a failed layout proof.
Gray-code and digit-transpose Memory POCs exercise both CPU and Metal for
sizes 1, 31, 32, 257, and 1024, plus CPU execution at 2,097,152 elements. The
large case again exports on Metal and rejects its infeasible shared allocation.
Native index-expression tests compare compiled address calculations with the
TileIR evaluator. The canonical example above is independently captured
under both C++20 and C++23.
2.2 Optional matrix atom selection
The kernel syntax is unchanged. The native compiler’s cooperative_matrix
option asserts device support; Metal additionally requires
thread_warp_size=32 (Apple GPU family 7+). This is not inferred from a generic
Metal target. Eligible rank-two FP32 group MMA uses 8×8 SIMD-group matrix
instructions; other types, scopes, layouts, and shapes retain the reference
loops. Global ragged edges and transposed loaded Tiles are supported.
mma(a, b, acc, {.allow_reassociation = false}) also keeps the ordered
contraction; the default permits reassociation without reducing input precision.
test_tile_tirx_matrix verifies both emitted Metal instructions and numerical
results against a double-precision oracle. The native/PyTorch benchmark records
the capability request separately from generated matrix call counts, so a
fallback is not reported as a matrix-hardware measurement.
3. Elementwise bias + GELU + residual
Two-dimensional blocks load once. The one-dimensional bias broadcasts by its shared column axis; neither loads nor stores are hidden in assignment.
constexpr int64_t rows = 3;
constexpr int64_t columns = 8;
constexpr float gelu_scale = 0.7978845608028654f;
constexpr float gelu_cubic = 0.044715f;
auto definition = tile_kernel(
"tile_poc_bias_gelu_residual",
[=](TensorView<const float, 2> x,
TensorView<const float, 1> bias,
TensorView<const float, 2> residual,
TensorView<float, 2> y) {
constexpr auto bm = 2;
constexpr auto bn = 4;
auto gm = axis("block_m", ceil_div(x.extent<0>(), bm));
auto gn = axis("block_n", ceil_div(x.extent<1>(), bn));
auto m = axis("m", bm);
auto n = axis("n", bn);
for (auto &nest : parallel(shape(gm, gn))) {
auto origin = coord(nest.index(gm) * bm, nest.index(gn) * bn);
auto value = x[origin, shape(m, n)] + bias[coord(nest.index(gn) * bn), shape(n)];
auto cubic = value * value * value;
auto gelu = 0.5f * value * (1.0f + tanh(gelu_scale * (value + gelu_cubic * cubic)));
y(origin, shape(m, n)).store(gelu + residual[origin, shape(m, n)]);
}
});
4. Statistics and losses
Sum, mean, maximum, stable argmax, MSE, MAE, Huber, and binary cross entropy
compose the same Tile operations. reduce(value, axis, add) is a library
expression; it expands into an element region and the existing reduction nest.
The region form still yields a Nest, not a different kind of iterator element.
constexpr int64_t rows = 3;
constexpr int64_t columns = 5;
constexpr float huber_delta = 0.75f;
constexpr float epsilon = 1e-4f;
auto definition = tile_kernel(
"tile_poc_statistics_losses",
[=](TensorView<const float, 2> x,
TensorView<const float, 2> target,
TensorView<const float, 2> probability,
TensorView<float, 1> sum_out,
TensorView<float, 1> mean_out,
TensorView<float, 1> peak_out,
TensorView<int64_t, 1> argmax_out,
TensorView<float, 1> mse_out,
TensorView<float, 1> mae_out,
TensorView<float, 1> huber_out,
TensorView<float, 1> bce_out) {
auto row = axis("row", x.extent<0>());
auto r = axis("local_row", 1);
auto c = axis("column", x.extent<1>());
for (auto &nest : parallel(shape(row))) {
auto origin = coord(nest.index(), 0);
auto value = x[origin, shape(r, c)];
auto expected = target[origin, shape(r, c)];
auto probability_tile = probability[origin, shape(r, c)];
auto difference = value - expected;
auto magnitude = abs(difference);
auto p = min(max(probability_tile, epsilon), 1.0f - epsilon);
auto huber = ite(magnitude <= huber_delta,
0.5f * difference * difference,
huber_delta * (magnitude - 0.5f * huber_delta));
auto bce = -(expected * luisa::compute::tile::log(p) + (1.0f - expected) * luisa::compute::tile::log(1.0f - p));
auto sum = reduce(value, c, add);
auto denominator = static_cast<float>(x.extent<1>());
auto destination = coord(nest.index());
sum_out(destination, shape(r)).store(sum);
mean_out(destination, shape(r)).store(sum / denominator);
peak_out(destination, shape(r)).store(reduce(value, c, maximum));
argmax_out(destination, shape(r)).store(argmax(value, c));
mse_out(destination, shape(r)).store(reduce(difference * difference, c, add) / denominator);
mae_out(destination, shape(r)).store(reduce(magnitude, c, add) / denominator);
huber_out(destination, shape(r)).store(reduce(huber, c, add) / denominator);
bce_out(destination, shape(r)).store(reduce(bce, c, add) / denominator);
}
});
A whole-tensor reduction can use one logical program; larger reductions can write partial Tiles and invoke a second kernel. No implicit grid barrier is invented. The current small whole-tensor test uses the former.
5. Softmax, LayerNorm, and RMSNorm
constexpr int64_t rows = 4;
constexpr int64_t columns = 7;
constexpr float epsilon = 1e-5f;
auto definition = tile_kernel(
"tile_poc_softmax_norm",
[=](TensorView<const float, 2> x,
TensorView<float, 2> softmax_out,
TensorView<float, 2> layernorm_out,
TensorView<float, 2> rmsnorm_out) {
auto row = axis("row", x.extent<0>());
auto r = axis("local_row", 1);
auto c = axis("column", x.extent<1>());
for (auto &nest : parallel(shape(row))) {
auto origin = coord(nest.index(), 0);
auto value = x[origin, shape(r, c)];
auto peak = reduce(value, c, maximum);
auto exponentials = exp(value - peak);
auto denominator = static_cast<float>(x.extent<1>());
auto mean = reduce(value, c, add) / denominator;
auto mean_square = reduce(value * value, c, add) / denominator;
auto variance = max(mean_square - mean * mean, 0.0f);
softmax_out(origin, shape(r, c)).store(exponentials / reduce(exponentials, c, add));
layernorm_out(origin, shape(r, c)).store((value - mean) / sqrt(variance + epsilon));
rmsnorm_out(origin, shape(r, c)).store(value / sqrt(mean_square + epsilon));
}
});
The LayerNorm POC uses the simple mean-square formulation to exercise Tile broadcasting. Numerically robust Welford state is a library extension, not a claim that this reference formula is suitable for every input distribution. The same concise row-sum, softmax, RMSNorm and LayerNorm shapes now have an opt-in proof-driven Metal SIMD-group realization without changing their DSL source; its formal mapping, storage proof and measured limits are in TIRx Metal reductions.
6. Sparse cross entropy and gradient
gather semantically indexes an already-loaded Tile. Its portable reference
expansion uses pure Tile extraction. A target realization may replace an
immutable input snapshot with a guarded direct Tensor read only after proving
source immutability, path bounds and distributed worker ownership; this is an
optimization, not a different source operation.
constexpr int64_t rows = 4;
constexpr int64_t classes = 6;
auto definition = tile_kernel(
"tile_poc_sparse_softmax_cross_entropy",
[=](TensorView<const float, 2> logits,
TensorView<const int64_t, 1> labels,
TensorView<float, 1> losses,
TensorView<float, 2> gradient) {
auto row = axis("row", logits.extent<0>());
auto r = axis("local_row", 1);
auto c = axis("class", logits.extent<1>());
for (auto &nest : parallel(shape(row))) {
auto origin = coord(nest.index(), 0);
auto value = logits[origin, shape(r, c)];
auto label = labels[coord(nest.index()), shape(r)];
auto peak = reduce(value, c, maximum);
auto exponentials = exp(value - peak);
auto total = reduce(exponentials, c, add);
auto selected = gather(value, label, c);
losses(coord(nest.index()), shape(r)).store(luisa::compute::tile::log(total) + peak - selected);
auto one_hot = ite(iota(c) == label, 1.0f, 0.0f);
gradient(origin, shape(r, c)).store(exponentials / total - one_hot);
}
});
7. Causal FlashAttention with online softmax
Q is loaded once per query block. Each pipeline iteration loads a K/V block, forms scores with one MMA, rescales the carried row state, and applies the second MMA. Query/key and feature dimensions are explicit local identities. The causal mask and key-tail mask are separate from memory bounds.
constexpr int64_t batches = 1;
constexpr int64_t heads = 2;
constexpr int64_t queries = 4;
constexpr int64_t keys = 5;
constexpr int64_t channels = 4;
constexpr int64_t value_channels = 3;
constexpr float scale = 0.5f;
auto definition = tile_kernel(
"tile_poc_flash_attention",
[=](TensorView<const float, 4> q,
TensorView<const float, 4> k,
TensorView<const float, 4> v,
TensorView<float, 4> result) {
constexpr auto bq = 2;
constexpr auto bk = 3;
auto batch = axis("batch", result.extent<0>());
auto head = axis("head", result.extent<1>());
auto query_blocks = axis("query_blocks", ceil_div(result.extent<2>(), bq));
auto key_blocks = axis("key_blocks", ceil_div(k.extent<2>(), bk));
auto b = axis("local_batch", 1);
auto h = axis("local_head", 1);
auto m = axis("query", bq);
auto s = axis("key", bk);
auto d = axis("channel", q.extent<3>());
auto dv = axis("value_channel", result.extent<3>());
for (auto &nest : parallel(shape(batch, head, query_blocks))) {
auto b0 = nest.index(batch);
auto h0 = nest.index(head);
auto q0 = nest.index(query_blocks) * bq;
auto query = q[coord(b0, h0, q0, 0), shape(b, h, m, d)];
auto row_max = full<float>(shape(b, h, m), -1e30f);
auto row_sum = zeros<float>(shape(b, h, m));
auto acc = zeros<float>(shape(b, h, m, dv));
for (auto &step : nest.pipeline(shape(key_blocks), {.window = 2, .interval = 1})) {
auto k0 = step.index() * bk;
step.stage("load");
auto key = k[coord(b0, h0, k0, 0), shape(b, h, s, d)];
auto value = v[coord(b0, h0, k0, 0), shape(b, h, s, dv)];
step.stage("score");
auto score = mma(query, key, zeros<float>(shape(b, h, m, s))) * scale;
auto valid = (iota(s) + k0 < k.extent<2>()) && (iota(s) + k0 <= iota(m) + q0);
auto masked = ite(valid, score, -1e30f);
auto next_max = max(row_max, reduce(masked, s, maximum));
auto alpha = exp(row_max - next_max);
auto probability = ite(valid, exp(masked - next_max), 0.0f);
step.stage("update");
row_sum = row_sum * alpha + reduce(probability, s, add);
acc = mma(probability, value, acc * alpha);
row_max = next_max;
}
result(coord(b0, h0, q0, 0), shape(b, h, m, dv)).store(acc / row_sum);
}
});
8. CNN: padded, strided convolution
The input access names a complete receptive-field window. reindex maps
logical filter taps to that loaded window; multiplying by a weight Tile and
reducing the tap/input-channel axes is the direct-convolution library form.
A later im2col/MMA schedule need not change the memory-access syntax.
constexpr int64_t batch_size = 1;
constexpr int64_t input_height = 5;
constexpr int64_t input_width = 6;
constexpr int64_t input_channels = 2;
constexpr int64_t filter_height = 3;
constexpr int64_t filter_width = 3;
constexpr int64_t output_channels = 3;
constexpr int64_t stride = 2;
constexpr int64_t dilation = 1;
constexpr int64_t padding = 1;
constexpr int64_t output_height = 3;
constexpr int64_t output_width = 3;
auto definition = tile_kernel(
"tile_poc_conv2d",
[=](TensorView<const float, 4> x,
TensorView<const float, 4> weights,
TensorView<const float, 1> bias,
TensorView<float, 4> y) {
auto batch = axis("batch", y.extent<0>());
auto output_y = axis("output_y", y.extent<1>());
auto output_x = axis("output_x", y.extent<2>());
auto b = axis("local_batch", 1);
auto oy = axis("local_y", 1);
auto ox = axis("local_x", 1);
auto fy = axis("filter_y", weights.extent<0>());
auto fx = axis("filter_x", weights.extent<1>());
auto ic = axis("input_channel", x.extent<3>());
auto oc = axis("output_channel", y.extent<3>());
auto wy = axis("window_y", (weights.extent<0>() - 1) * dilation + 1);
auto wx = axis("window_x", (weights.extent<1>() - 1) * dilation + 1);
for (auto &nest : parallel(shape(batch, output_y, output_x))) {
auto n0 = nest.index(batch);
auto y0 = nest.index(output_y);
auto x0 = nest.index(output_x);
auto window = x[coord(n0, y0 * stride - padding, x0 * stride - padding, 0), shape(b, wy, wx, ic)];
auto taps = reindex(window, shape(b, fy, fx, ic), [&](const Nest &element) {
return coord(element.index(b), element.index(fy) * dilation, element.index(fx) * dilation, element.index(ic));
});
auto filter = weights[coord(0, 0, 0, 0), shape(fy, fx, ic, oc)];
auto bias_tile = bias[coord(0), shape(oc)];
auto value = reduce(taps * filter, shape(fy, fx, ic), add) + bias_tile;
auto result = reshape(max(value, 0.0f), shape(b, oy, ox, oc));
y(coord(n0, y0, x0, 0), shape(b, oy, ox, oc)).store(result);
}
});
9. Depthwise convolution and pooling
constexpr int64_t height = 5;
constexpr int64_t width = 6;
constexpr int64_t channels = 3;
constexpr int64_t filter_size = 3;
constexpr int64_t padding = 1;
auto definition = tile_kernel(
"tile_poc_depthwise_pool",
[](TensorView<const float, 3> source,
TensorView<const float, 3> weights,
TensorView<float, 3> convolution,
TensorView<float, 3> pool) {
auto y = axis("y", source.extent<0>());
auto x = axis("x", source.extent<1>());
auto c = axis("channel", source.extent<2>());
auto fy = axis("filter_y", weights.extent<0>());
auto fx = axis("filter_x", weights.extent<1>());
auto oy = axis("local_y", 1);
auto ox = axis("local_x", 1);
for (auto &nest : parallel(shape(y, x))) {
auto window = source[coord(nest.index(y) - padding, nest.index(x) - padding, 0), shape(fy, fx, c)];
auto filter = weights[coord(0, 0, 0), shape(fy, fx, c)];
auto conv = reduce(window * filter, shape(fy, fx), add);
auto pooled = reduce(window, shape(fy, fx), maximum);
auto destination = coord(nest.index(y), nest.index(x), 0);
convolution(destination, shape(oy, ox, c)).store(reshape(conv, shape(oy, ox, c)));
pool(destination, shape(oy, ox, c)).store(reshape(pooled, shape(oy, ox, c)));
}
});
This POC deliberately uses zero padding for both outputs. A conventional negative-infinity-padded max pool would use an explicit fallback on its load.
10. Traditional filters: Sobel and median
constexpr int64_t height = 6;
constexpr int64_t width = 7;
constexpr int64_t tap_count = 9;
auto definition = tile_kernel(
"tile_poc_sobel_median",
[](TensorView<const float, 2> source,
TensorView<float, 2> gradient_x,
TensorView<float, 2> gradient_y,
TensorView<float, 2> median) {
auto y = axis("y", source.extent<0>());
auto x = axis("x", source.extent<1>());
auto fy = axis("filter_y", 3);
auto fx = axis("filter_x", 3);
auto tap = axis("tap", tap_count);
auto oy = axis("local_y", 1);
auto ox = axis("local_x", 1);
for (auto &nest : parallel(shape(y, x))) {
auto window = source[coord(nest.index(y) - 1, nest.index(x) - 1), shape(fy, fx)];
auto dy = iota(fy) - 1;
auto dx = iota(fx) - 1;
auto weight_x = cast<float>(dx * ite(dy == 0, int64_t{2}, int64_t{1}));
auto weight_y = cast<float>(dy * ite(dx == 0, int64_t{2}, int64_t{1}));
auto gx = reduce(window * weight_x, shape(fy, fx), add);
auto gy = reduce(window * weight_y, shape(fy, fx), add);
auto ordered = luisa::compute::tile::sort(reshape(window, shape(tap)), tap);
auto middle = gather(ordered.values, full<int64_t>(IndexSpace{}, 4), tap);
auto destination = coord(nest.index(y), nest.index(x));
gradient_x(destination, shape(oy, ox)).store(gx);
gradient_y(destination, shape(oy, ox)).store(gy);
median(destination, shape(oy, ox)).store(middle);
}
});
The small-window median uses the same sorting library as Top-K. There is no median-specific primitive or execution scope.
11. Stable sort and Top-K
The topk library returns Tile values and source indices. Its current
quadratic composition orders finite values and breaks ties by source index.
A tuned network, radix sort, and a documented NaN policy remain future library
work; these are not hidden behind a claim of hardware-optimized sorting.
constexpr int64_t rows = 3;
constexpr int64_t columns = 8;
constexpr int64_t top_k = 3;
auto definition = tile_kernel(
"tile_poc_stable_sort",
[](TensorView<const float, 2> source,
TensorView<float, 2> sorted_value,
TensorView<int64_t, 2> sorted_index) {
auto row = axis("row", source.extent<0>());
auto r = axis("local_row", 1);
auto c = axis("column", source.extent<1>());
auto rank = axis("output_rank", sorted_value.extent<1>());
for (auto &nest : parallel(shape(row))) {
auto value = source[coord(nest.index(), 0), shape(r, c)];
auto ranked = topk(value, c, sorted_value.extent<1>());
auto destination = coord(nest.index(), 0);
sorted_value(destination, shape(r, rank)).store(ranked.values);
sorted_index(destination, shape(r, rank)).store(ranked.indices);
}
});
12. Segmented accumulation
This reference uses one program per bucket and a masked Tile reduction, including empty buckets. It is deterministic and needs no atomics. An atomic scatter implementation is a different schedule/library choice.
constexpr int64_t item_count = 13;
constexpr int64_t bucket_count = 5;
auto definition = tile_kernel(
"tile_poc_segmented_accumulation",
[](TensorView<const int64_t, 1> ids,
TensorView<const float, 1> values,
TensorView<float, 1> sums,
TensorView<int64_t, 1> counts) {
auto item = axis("item", ids.extent<0>());
auto bucket = axis("bucket", sums.extent<0>());
auto output = axis("output", 1);
for (auto &nest : parallel(shape(bucket))) {
auto id = ids[coord(0), shape(item)];
auto value = values[coord(0), shape(item)];
auto matches = id == nest.index();
auto sum = reduce(ite(matches, value, 0.0f), item, add);
auto count = reduce(cast<int64_t>(matches), item, add);
sums(coord(nest.index()), shape(output)).store(sum);
counts(coord(nest.index()), shape(output)).store(count);
}
});
13. Nested temporal composition
This test combines parallel → serial → pipeline → reduction. A single outer Tile variable is assigned inside the nested regions, exercising inferred state at each lexical boundary.
constexpr int64_t rows = 3;
constexpr int64_t phases = 2;
constexpr int64_t steps = 3;
constexpr int64_t lanes = 4;
auto definition = tile_kernel(
"tile_poc_all_structured_regions",
[](TensorView<const float, 4> source, TensorView<float, 1> result) {
auto row = axis("row", source.extent<0>());
auto phase = axis("phase", source.extent<1>());
auto step = axis("step", source.extent<2>());
auto lane = axis("lane", source.extent<3>());
auto r = axis("local_row", 1);
auto p = axis("local_phase", 1);
auto s = axis("local_step", 1);
for (auto &nest : parallel(shape(row))) {
auto sum = zeros<float>(shape(r));
for (auto &phase_nest : nest.serial(shape(phase))) {
for (auto &step_nest : phase_nest.pipeline(shape(step), {.window = 2, .interval = 1})) {
step_nest.stage("load");
auto value = source[coord(nest.index(row), phase_nest.index(phase), step_nest.index(step), 0), shape(r, p, s, lane)];
step_nest.stage("consume");
sum += reshape(reduce(value, lane, add), shape(r));
}
}
result(coord(nest.index(row)), shape(r)).store(sum);
}
});
14. Validation boundary
The same subscript syntax has compile/capture tests in C++20 and C++23.
Every operator POC is registered for native TVMx CPU and, when available, actual Metal execution. Metal failures do not silently fall back to CPU.
POC validation rejects scalar View loads/stores. Fine-grained scalar access remains available only through explicit low-level reference tests.
Numerical checks use independent host reference formulas, including tail tiles, boundary padding, masks, and stable ties.
Performance comparisons must use multiple shapes and report compile, transfer, cold call, and synchronized warm execution separately. Correctness tests do not establish tensor-core lowering or competitive performance.