feat(loss): add cross_entropy_forward for Ascend 950-ops-tilelang-AtomGit

Adocs/cross_entropy.md+19-0

@@ -0,0 +1,19 @@
	1	+# Cross entropy 前向
	G G gong-zhihao15 hours ago md 文档改一下里面的描述吧，看着ai味道有点重
	2	+
	3	+接口为 `cann_ops_tilelang.cross_entropy_forward(logits, targets, *, ignore_index=-100, reduction="none")`。硬标签交叉熵，输入 logits 为 `[N,V]` 的 FP32、FP16 或 BF16，targets 为同设备连续 int64 `[N]`，输出 FP32。none 返回 `[N]`，sum 和 mean 返回标量。logits 每行连续，行间可有非重叠间隔。
	4	+
	5	+有效标签位于 `[0,V)`，ignore_index 标签整行忽略。调用者保证有效标签范围和有限 logits。空输入 none 为空，sum 为0，mean 为NaN。全忽略 mean 为NaN。只提供前向计算，返回值不参与自动求导。当前接口由算子分配独立输出。
	6	+
	7	+稳定计算基于逐行最大值、指数和与答案分数。大输入使用分块统计、联合归约或核内部分累计，mean 路径按数据范围选择预缩放来避免中间和溢出。具体路径按输入元数据、核心数和UB预算选取。
	8	+
	9	+精度测试使用FP64参考，输出FP32，逐元素检查误差和NaN/Inf掩码。测试覆盖三种类型、三种reduction、空输入、ignore和非连续行距。
	10	+
	11	+运行 `OPS_TILELANG_TEST_LEVEL=2 pytest tests/loss -m "not benchmark"`。运行性能测试使用 `pytest tests/loss -m benchmark --run-benchmark`。基线与显存画像通过仓内刷新命令生成。
	12	+
	13	+仓内benchmark沿用统一msprof口径。公共接口端到端性能须另行测量调用开始至当前流完成，包含调用内分配、提交和等待。重复输入和冲刷条件分别报告。
	14	+
	15	+最大性能档N=65536、V=129280、FP32在验收环境中峰值约32577 MiB。测试框架根据生成画像将大显存用例分到第二阶段，运行前应检查可用HBM。
	16	+
	17	+性能基线按仓库a5类别记录，来自验收机器的同机复核。不同配置的CI设备应重新建立对应基线。显存画像可供pytest分类，实际并发策略由所使用的测试入口决定。
	18	+
	19	+保留原公共接口 `cross_entropy_forward` 的 eager 调用。torch 自定义入口为 `torch.ops.cann_ops_tilelang.cross_entropy_forward.default`，可用 `torch.compile(..., backend='npugraph_ex', fullgraph=True, dynamic=False)` 编译，并在 `torch.no_grad()` 下显式执行前向。fake 实现仅根据元数据推导 FP32 输出的形状和设备。该入口仅提供无梯度前向，不支持 backward。

Msrc/cann_ops_tilelang/__init__.py+5-1

	@@ -1,6 +1,9 @@
1	-"""High-performance Ascend operators implemented with TileLang."""	1	+# Copyright (c) 2026 ops-tilelang contributors
		2	+# SPDX-License-Identifier: MIT
		3	+# See LICENSE in the repository root for the full license text.
2		4
3	__all__ = [	5	__all__ = [
		6	+ 'cross_entropy_forward',
4	'engram_fused_weight',	7	'engram_fused_weight',
5	'engram_gate_bwd',	8	'engram_gate_bwd',
6	'engram_gate_fwd',	9	'engram_gate_fwd',
	@@ -113,6 +116,7 @@ from .engram import (
113	engram_sinkhorn_momentum_update,	116	engram_sinkhorn_momentum_update,
114	engram_sinkhorn_step,	117	engram_sinkhorn_step,
115	)	118	)
		119	+from .loss import cross_entropy_forward
116	from .mhc import (	120	from .mhc import (
117	expand_to_mhc_backward,	121	expand_to_mhc_backward,
118	expand_to_mhc_forward,	122	expand_to_mhc_forward,

Asrc/cann_ops_tilelang/loss/__init__.py+7-0

@@ -0,0 +1,7 @@
	1	+# Copyright (c) 2026 ops-tilelang contributors
	2	+# SPDX-License-Identifier: MIT
	3	+# See LICENSE in the repository root for the full license text.
	4	+
	5	+__all__ = ['cross_entropy_forward']
	6	+
	7	+from .cross_entropy import cross_entropy_forward

Asrc/cann_ops_tilelang/loss/_cross_entropy_chunk.py+342-0

@@ -0,0 +1,342 @@
	1	+# Copyright (c) 2026 ops-tilelang contributors
	2	+# SPDX-License-Identifier: MIT
	3	+# See LICENSE in the repository root for the full license text.
	4	+
	5	+from functools import cache
	6	+
	7	+import tilelang.ascend.language as T
	8	+
	9	+from ._cross_entropy_common import _compile
	10	+from ._cross_entropy_simd import _block_exponential_lanes, _block_maximum, _plain_exponential_lanes
	11	+
	12	+
	13	+@cache
	14	+def _chunk_kernel(dtype, block, cores):
	15	+ (n, v, k) = (T.dynamic(name, dtype='int64') for name in ('n', 'v', 'k'))
	16	+ stride = T.dynamic('stride', dtype='int64')
	17	+ narrow = dtype != 'float32'
	18	+ stages = 2
	19	+
	20	+ block_maximum = _block_maximum(dtype)
	21	+ exponential_lanes = _plain_exponential_lanes(narrow)
	22	+
	23	+ @T.macro
	24	+ def _plain_chunk_task(buffer_x, buffer_max, buffer_sum, x_ub, row_max, row_sum, row, part):
	25	+ valid = T.min(T.int64(block), v - part * block)
	26	+ T.copy(buffer_x[row, slice(part * block, part * block + valid)], x_ub[:valid])
	27	+ with T.SimdVF():
	28	+ full = T.simd.pset(32)
	29	+ neg = T.simd.vdup(-T.infinity('float32'), T.float32, full)
	30	+ block_max = block_maximum(x_ub, valid, full, neg)
	31	+ m = block_max
	32	+ ls = exponential_lanes(x_ub, valid, m, full, neg)
	33	+ s = T.simd.vcadd(ls, full)
	34	+ T.simd.vsts(row_max[0], m, full, 'ONEPT_B32')
	35	+ T.simd.vsts(row_sum[0], s, full, 'ONEPT_B32')
	36	+ T.copy(row_max[0:1], buffer_max[row, slice(part, part + 1)])
	37	+ T.copy(row_sum[0:1], buffer_sum[row, slice(part, part + 1)])
	38	+
	39	+ @T.prim_func
	40	+ def chunk(
	41	+ buffer_x: T.StridedTensor((n, v), (stride, 1), dtype),
	42	+ buffer_max: T.Tensor((n, k), 'float32'),
	43	+ buffer_sum: T.Tensor((n, k), 'float32'),
	44	+ ):
	45	+ with T.Kernel(cores) as core:
	46	+ x_ub = T.alloc_shared((block + 128,), dtype)
	47	+ row_max = T.alloc_shared((1,), 'float32')
	48	+ row_sum = T.alloc_shared((1,), 'float32')
	49	+ T.annotate_buffer_versions({x_ub: stages, row_max: stages, row_sum: stages})
	50	+ for wave in T.Pipelined(T.ceildiv(n * k, cores), num_stages=stages):
	51	+ task = wave * cores + core
	52	+ if task < n * k:
	53	+ row = task // k
	54	+ part = task % k
	55	+ _plain_chunk_task(buffer_x, buffer_max, buffer_sum, x_ub, row_max, row_sum, row, part)
	56	+
	57	+ return _compile(f'chunk_{dtype}_{block}_{cores}', chunk)
	58	+
	59	+
	60	+@cache
	61	+def _fully_simd_fused_kernel(dtype, block, cores, slots):
	62	+ (n, v) = (T.dynamic(name, dtype='int32') for name in ('n', 'v'))
	63	+ k = T.ceildiv(v, block)
	64	+ stride = T.dynamic('stride', dtype='int64')
	65	+ narrow = dtype != 'float32'
	66	+ stages = 2
	67	+ out_size = T.dynamic('out_size', dtype='int32')
	68	+ row_capacity = (2 * cores + 63) // 64 * 64
	69	+
	70	+ block_maximum = _block_maximum(dtype)
	71	+ exponential_lanes = _block_exponential_lanes(dtype)
	72	+
	73	+ @T.macro
	74	+ def _label_words(labels, label_index, full):
	75	+ low = T.simd.vgather2(
	76	+ T.access_ptr(labels[0], 'r', extent=4 * cores + 64),
	77	+ T.reinterpret(label_index, 'uint32x64'),
	78	+ full,
	79	+ )
	80	+ high = T.simd.vgather2(
	81	+ T.access_ptr(labels[0], 'r', extent=4 * cores + 64),
	82	+ T.reinterpret(T.simd.vadd(label_index, T.simd.vdup(1, T.int32, full), full), 'uint32x64'),
	83	+ full,
	84	+ )
	85	+ return low, high
	86	+
	87	+ @T.macro
	88	+ def _chunk_statistics_vf(x_ub, labels, stats, row, part, valid, wave):
	89	+ with T.SimdVF():
	90	+ full = T.simd.pset(32)
	91	+ neg = T.simd.vdup(-T.infinity('float32'), T.float32, full)
	92	+ block_max = block_maximum(x_ub, valid, full, neg)
	93	+ m = block_max
	94	+ ls = exponential_lanes(x_ub, valid, m, full, neg)
	95	+ label_index = T.simd.vdup(row * 2, T.int32, full)
	96	+ low, high = _label_words(labels, label_index, full)
	97	+ in_block = T.simd.pand(
	98	+ T.simd.vcmps(low, part * block, full, 'ge'),
	99	+ T.simd.vcmps(low, part * block + valid, full, 'lt'),
	100	+ full,
	101	+ )
	102	+ in_block = T.simd.pand(in_block, T.simd.vcmps(high, 0, full, 'eq'), full)
	103	+ offset = T.simd.vsub(low, T.simd.vdup(part * block, T.int32, full), full)
	104	+ safe_offset = T.simd.vsel(offset, T.simd.vdup(0, T.int32, full), in_block)
	105	+ if narrow:
	106	+ gather_full = T.simd.pset(16)
	107	+ gather_index = T.simd.vdupv(T.simd.vcvt(safe_offset, T.uint16, full), gather_full)
	108	+ answer = T.simd.vcvt(
	109	+ T.simd.vgather2(T.access_ptr(x_ub[0], 'r', extent=block), gather_index, gather_full),
	110	+ T.float32,
	111	+ part=0,
	112	+ )
	113	+ else:
	114	+ answer = T.simd.vgather2(
	115	+ T.access_ptr(x_ub[0], 'r', extent=block), T.reinterpret(safe_offset, 'uint32x64'), full
	116	+ )
	117	+ answer = T.simd.vsel(answer, T.simd.vdup(0.0, T.float32, full), in_block)
	118	+ s = T.simd.vcadd(ls, full)
	119	+ T.simd.vsts(stats[wave * 8], m, full, 'ONEPT_B32')
	120	+ T.simd.vsts(stats[wave * 8 + 1], s, full, 'ONEPT_B32')
	121	+ T.simd.vsts(stats[wave * 8 + 2], answer, full, 'ONEPT_B32')
	122	+
	123	+ @T.macro
	124	+ def _row_lanes(n, group, full):
	125	+ active = T.simd.update_mask(T.min(64, n - group * 64), 32)
	126	+ rows = T.simd.vadd(T.simd.vci(0, 'int32'), T.simd.vdup(group * 64, T.int32, full), full)
	127	+ rows = T.simd.vsel(rows, T.simd.vdup(0, T.int32, full), active)
	128	+ return active, rows
	129	+
	130	+ @T.macro
	131	+ def _label_ignored(low, high, ignore, full):
	132	+ ignored = T.simd.pand(
	133	+ T.simd.vcmps(low, T.cast(ignore, 'int32'), full, 'eq'),
	134	+ T.simd.vcmps(high, T.cast(ignore >> 32, 'int32'), full, 'eq'),
	135	+ full,
	136	+ )
	137	+ return ignored
	138	+
	139	+ @T.macro
	140	+ def _valid_row_count(labels, n, ignore, full, zero, one):
	141	+ count_lanes = T.simd.alloc_var(T.float32)
	142	+ count_lanes = zero
	143	+ for group in T.serial(T.ceildiv(n, 64)):
	144	+ active, rows = _row_lanes(n, group, full)
	145	+ label_index = T.simd.vmul(rows, T.simd.vdup(2, T.int32, full), full)
	146	+ low, high = _label_words(labels, label_index, full)
	147	+ ignored = _label_ignored(low, high, ignore, full)
	148	+ valid_count = T.simd.vsel(zero, one, ignored)
	149	+ valid_count = T.simd.vsel(valid_count, zero, active)
	150	+ count_lanes = T.simd.vadd(count_lanes, valid_count, full)
	151	+ count = T.simd.vdupv(T.simd.vcadd(count_lanes, full), full)
	152	+ return count
	153	+
	154	+ @T.macro
	155	+ def _corrected_task_quotient(task0, full):
	156	+ q0 = T.simd.vcvt(
	157	+ T.simd.vmul(T.simd.vcvt(task0, T.float32, full), T.simd.vdup(1.0 / cores, T.float32, full), full),
	158	+ T.int32,
	159	+ full,
	160	+ round='ROUND_Z',
	161	+ )
	162	+ rem0 = T.simd.vsub(task0, T.simd.vmul(q0, T.simd.vdup(cores, T.int32, full), full), full)
	163	+ q0 = T.simd.vsel(
	164	+ T.simd.vadd(q0, T.simd.vdup(1, T.int32, full), full),
	165	+ q0,
	166	+ T.simd.vcmps(rem0, cores, full, 'ge'),
	167	+ )
	168	+ q0 = T.simd.vsel(T.simd.vsub(q0, T.simd.vdup(1, T.int32, full), full), q0, T.simd.vcmps(rem0, 0, full, 'lt'))
	169	+ return q0
	170	+
	171	+ @T.macro
	172	+ def _loss_from_merged_statistics(mode, maximum, answer, total, inverse, full):
	173	+ value = T.simd.alloc_var(T.float32)
	174	+ if mode == 2:
	175	+ half = T.simd.vdup(0.5, T.float32, full)
	176	+ gap = T.simd.vsub(T.simd.vmul(maximum, half, full), T.simd.vmul(answer, half, full), full)
	177	+ value = T.simd.vadd(
	178	+ T.simd.vmul(T.simd.vln(total, full), inverse, full),
	179	+ T.simd.vmul(gap, T.simd.vmul(T.simd.vdup(2.0, T.float32, full), inverse, full), full),
	180	+ full,
	181	+ )
	182	+ else:
	183	+ value = T.simd.vsub(T.simd.vln(total, full), T.simd.vsub(answer, maximum, full), full)
	184	+ return value
	185	+
	186	+ @T.macro
	187	+ def _group_origin(rows, full):
	188	+ task0 = T.simd.vmul(rows, T.simd.vdup(k, T.int32, full), full)
	189	+ q0 = _corrected_task_quotient(task0, full)
	190	+ owner0 = T.simd.vsub(task0, T.simd.vmul(q0, T.simd.vdup(cores, T.int32, full), full), full)
	191	+ index0 = T.simd.vadd(
	192	+ T.simd.vmul(owner0, T.simd.vdup(slots * 8, T.int32, full), full),
	193	+ T.simd.vmul(q0, T.simd.vdup(8, T.int32, full), full),
	194	+ full,
	195	+ )
	196	+ return owner0, index0
	197	+
	198	+ @T.macro
	199	+ def _merged_max_load(merged, index, full):
	200	+ chunk_max = T.simd.vgather2(
	201	+ T.access_ptr(merged[0, 0], 'r', extent=cores * slots * 8),
	202	+ T.reinterpret(index, 'uint32x64'),
	203	+ full,
	204	+ )
	205	+ return chunk_max
	206	+
	207	+ @T.macro
	208	+ def _merged_chunk_loads(merged, index, full):
	209	+ chunk_max = _merged_max_load(merged, index, full)
	210	+ chunk_sum = T.simd.vgather2(
	211	+ T.access_ptr(merged[0, 0], 'r', extent=cores * slots * 8),
	212	+ T.reinterpret(T.simd.vadd(index, T.simd.vdup(1, T.int32, full), full), 'uint32x64'),
	213	+ full,
	214	+ )
	215	+ chunk_answer = T.simd.vgather2(
	216	+ T.access_ptr(merged[0, 0], 'r', extent=cores * slots * 8),
	217	+ T.reinterpret(T.simd.vadd(index, T.simd.vdup(2, T.int32, full), full), 'uint32x64'),
	218	+ full,
	219	+ )
	220	+ return chunk_max, chunk_sum, chunk_answer
	221	+
	222	+ @T.macro
	223	+ def _masked_group_value(labels, rows, ignore, value, zero, active, full):
	224	+ label_index = T.simd.vmul(rows, T.simd.vdup(2, T.int32, full), full)
	225	+ low, high = _label_words(labels, label_index, full)
	226	+ ignored = _label_ignored(low, high, ignore, full)
	227	+ masked = T.simd.vsel(zero, value, ignored)
	228	+ return T.simd.vsel(masked, zero, active)
	229	+
	230	+ @T.macro
	231	+ def _merged_group(merged, labels, rows, active, ignore, mode, inverse, full, zero):
	232	+ maximum = T.simd.alloc_var(T.float32)
	233	+ maximum = T.simd.vdup(-T.infinity('float32'), T.float32, full)
	234	+ owner0, index0 = _group_origin(rows, full)
	235	+ owner = T.simd.alloc_var(T.int32)
	236	+ index = T.simd.alloc_var(T.int32)
	237	+ owner = owner0
	238	+ index = index0
	239	+ for _ in T.serial(k):
	240	+ chunk_max = _merged_max_load(merged, index, full)
	241	+ maximum = T.simd.vmax(maximum, chunk_max, full)
	242	+ owner = T.simd.vadd(owner, T.simd.vdup(1, T.int32, full), full)
	243	+ wrap = T.simd.vcmps(owner, cores, full, 'ge')
	244	+ owner = T.simd.vsel(T.simd.vsub(owner, T.simd.vdup(cores, T.int32, full), full), owner, wrap)
	245	+ index = T.simd.vadd(index, T.simd.vdup(slots * 8, T.int32, full), full)
	246	+ index = T.simd.vsel(
	247	+ T.simd.vsub(index, T.simd.vdup(cores * slots * 8 - 8, T.int32, full), full), index, wrap
	248	+ )
	249	+ total = T.simd.alloc_var(T.float32)
	250	+ answer = T.simd.alloc_var(T.float32)
	251	+ total = zero
	252	+ answer = zero
	253	+ owner = owner0
	254	+ index = index0
	255	+ for _ in T.serial(k):
	256	+ chunk_max, chunk_sum, chunk_answer = _merged_chunk_loads(merged, index, full)
	257	+ if narrow:
	258	+ total = T.simd.vadd(total, T.simd.vmul(chunk_sum, T.simd.vexpdif(chunk_max, maximum, full), full), full)
	259	+ else:
	260	+ total = T.simd.vadd(
	261	+ total,
	262	+ T.simd.vmul(chunk_sum, T.simd.vexp(T.simd.vsub(chunk_max, maximum, full), full), full),
	263	+ full,
	264	+ )
	265	+ answer = T.simd.vadd(answer, chunk_answer, full)
	266	+ owner = T.simd.vadd(owner, T.simd.vdup(1, T.int32, full), full)
	267	+ wrap = T.simd.vcmps(owner, cores, full, 'ge')
	268	+ owner = T.simd.vsel(T.simd.vsub(owner, T.simd.vdup(cores, T.int32, full), full), owner, wrap)
	269	+ index = T.simd.vadd(index, T.simd.vdup(slots * 8, T.int32, full), full)
	270	+ index = T.simd.vsel(
	271	+ T.simd.vsub(index, T.simd.vdup(cores * slots * 8 - 8, T.int32, full), full), index, wrap
	272	+ )
	273	+ value = _loss_from_merged_statistics(mode, maximum, answer, total, inverse, full)
	274	+ value = _masked_group_value(labels, rows, ignore, value, zero, active, full)
	275	+ return value
	276	+
	277	+ @T.macro
	278	+ def _merge_statistics_vf(merged, labels, output, n, ignore, mode):
	279	+ with T.SimdVF():
	280	+ full = T.simd.pset(32)
	281	+ zero = T.simd.vdup(0.0, T.float32, full)
	282	+ one = T.simd.vdup(1.0, T.float32, full)
	283	+ count = _valid_row_count(labels, n, ignore, full, zero, one)
	284	+ inverse = T.simd.vdiv(one, count, full)
	285	+ combined = T.simd.alloc_var(T.float32)
	286	+ combined = zero
	287	+ for group in T.serial(T.ceildiv(n, 64)):
	288	+ active, rows = _row_lanes(n, group, full)
	289	+ value = _merged_group(merged, labels, rows, active, ignore, mode, inverse, full, zero)
	290	+ if mode == 0:
	291	+ T.simd.vsts(output[group * 64], value, active)
	292	+ else:
	293	+ combined = T.simd.vadd(combined, value, full)
	294	+ if mode != 0:
	295	+ result = T.simd.alloc_var(T.float32)
	296	+ result = T.simd.vcadd(combined, full)
	297	+ if mode == 2:
	298	+ result = T.simd.vsel(
	299	+ T.simd.vdup(T.float32(float('nan')), T.float32, full),
	300	+ result,
	301	+ T.simd.vcmps(count, 0.0, full, 'eq'),
	302	+ )
	303	+ T.simd.vsts(output[0], result, full, 'ONEPT_B32')
	304	+
	305	+ @T.prim_func
	306	+ def chunk(
	307	+ buffer_x: T.StridedTensor((n, v), (stride, 1), dtype),
	308	+ buffer_y: T.Tensor((n * 2,), 'int32'),
	309	+ buffer_workspace: T.Tensor((cores, slots * 8), 'float32'),
	310	+ buffer_loss: T.Tensor((out_size,), 'float32'),
	311	+ ignore: T.int64,
	312	+ mode: T.int32,
	313	+ ):
	314	+ with T.Kernel(cores) as core:
	315	+ x_ub = T.alloc_shared((block + 128,), dtype)
	316	+ stats = T.alloc_shared((slots * 8,), 'float32')
	317	+ labels = T.alloc_shared((4 * cores + 64,), 'int32')
	318	+ merged = T.alloc_shared((cores, slots * 8), 'float32')
	319	+ output = T.alloc_shared((row_capacity,), 'float32')
	320	+ T.annotate_buffer_versions({x_ub: stages, stats: 1, labels: 1, merged: 1, output: 1})
	321	+ T.copy(buffer_y[: n * 2], labels[: n * 2])
	322	+ for wave in T.Pipelined(T.ceildiv(n * k, cores), num_stages=stages):
	323	+ task = wave * cores + core
	324	+ if task < n * k:
	325	+ row = T.alloc_var('int32', init=task // k)
	326	+ part = T.alloc_var('int32', init=task % k)
	327	+ valid = T.alloc_var('int32', init=T.min(T.int32(block), v - part * block))
	328	+ T.copy(buffer_x[row, part * block], x_ub[:valid])
	329	+ _chunk_statistics_vf(x_ub, labels, stats, row, part, valid, wave)
	330	+ with T.PerCoreTask():
	331	+ T.copy(stats[: slots * 8], buffer_workspace[core, : slots * 8], l2_cache_ctrl=0)
	332	+ T.ascend_sync_inter_arrive('PIPE_MTE3', 15)
	333	+ T.ascend_sync_inter_wait('PIPE_S', 15)
	334	+ if core == 0:
	335	+ T.copy(buffer_workspace[:cores, : slots * 8], merged)
	336	+ _merge_statistics_vf(merged, labels, output, n, ignore, mode)
	337	+ if mode == 0:
	338	+ T.copy(output[:n], buffer_loss[:n], l2_cache_ctrl=0)
	339	+ else:
	340	+ T.copy(output[:1], buffer_loss[:1], l2_cache_ctrl=0)
	341	+
	342	+ return _compile(f'row_fully_simd_fused_{dtype}_{block}_{cores}_{slots}', chunk)

Asrc/cann_ops_tilelang/loss/_cross_entropy_common.py+26-0

@@ -0,0 +1,26 @@
	1	+# Copyright (c) 2026 ops-tilelang contributors
	2	+# SPDX-License-Identifier: MIT
	3	+# See LICENSE in the repository root for the full license text.
	4	+
	5	+
	6	+import tilelang
	7	+
	8	+VOCAB_BLOCK = 24576
	9	+VOCAB_BYTES = 98304
	10	+ROW_BLOCK = 256
	11	+GRID_CORES = 64
	12	+SMALL_MAX_LANE_WORK = 16
	13	+SIMD_ROW_MAX_WAVES = 2
	14	+BALANCED_FUSED_BLOCK = True
	15	+FUSED_UB_BUDGET = 208 * 1024
	16	+COMPILED_KERNELS = {}
	17	+
	18	+
	19	+def _compile(name, program):
	20	+ compiled = tilelang.compile(
	21	+ program,
	22	+ target='ascend',
	23	+ pass_configs={tilelang.PassConfigKey.TL_ENABLE_AUTO_SCHEDULE: name.startswith(('chunk', 'row_'))},
	24	+ )
	25	+ COMPILED_KERNELS[name] = compiled
	26	+ return compiled

Asrc/cann_ops_tilelang/loss/_cross_entropy_counted.py+217-0

@@ -0,0 +1,217 @@
	1	+# Copyright (c) 2026 ops-tilelang contributors
	2	+# SPDX-License-Identifier: MIT
	3	+# See LICENSE in the repository root for the full license text.
	4	+
	5	+from functools import cache
	6	+
	7	+import tilelang.ascend.language as T
	8	+
	9	+from ._cross_entropy_common import _compile
	10	+from ._cross_entropy_simd import _block_exponential_sum, _block_maximum
	11	+
	12	+
	13	+def _row_statistics_update(dtype):
	14	+ narrow = dtype != 'float32'
	15	+ block_maximum = _block_maximum(dtype)
	16	+ block_exponential_sum = _block_exponential_sum(dtype)
	17	+
	18	+ @T.macro
	19	+ def update_row_statistics(x_ub, valid, row_max, row_sum):
	20	+ full = T.simd.pset(32)
	21	+ neg = T.simd.vdup(-T.infinity('float32'), T.float32, full)
	22	+ block_max = block_maximum(x_ub, valid, full, neg)
	23	+ run_max = T.simd.vld(row_max[0], dist='BRC_B32')
	24	+ run_sum = T.simd.vld(row_sum[0], dist='BRC_B32')
	25	+ m = T.simd.vmax(block_max, run_max, full)
	26	+ block_sum = block_exponential_sum(x_ub, valid, m, full, neg)
	27	+ s = T.simd.vadd(
	28	+ T.simd.vmul(
	29	+ run_sum,
	30	+ T.simd.vexpdif(run_max, m, full) if narrow else T.simd.vexp(T.simd.vsub(run_max, m, full), full),
	31	+ full,
	32	+ ),
	33	+ block_sum,
	34	+ full,
	35	+ )
	36	+ T.simd.vsts(row_max[0], m, full, 'ONEPT_B32')
	37	+ T.simd.vsts(row_sum[0], s, full, 'ONEPT_B32')
	38	+
	39	+ return update_row_statistics
	40	+
	41	+
	42	+@cache
	43	+def _counted_row_kernel(dtype, block, cores, scaled):
	44	+ (n, v, k) = (T.dynamic(name, dtype='int64') for name in ('n', 'v', 'k'))
	45	+ stride = T.dynamic('stride', dtype='int64')
	46	+ stages = 2
	47	+
	48	+ update_row_statistics = _row_statistics_update(dtype)
	49	+
	50	+ @T.macro
	51	+ def counted_statistics_block(buffer_x, row, part, x_ub, row_max, row_sum):
	52	+ valid = T.min(T.int64(block), v - part * block)
	53	+ T.copy(buffer_x[row, part * block], x_ub[:valid])
	54	+ with T.SimdVF():
	55	+ update_row_statistics(x_ub, valid, row_max, row_sum)
	56	+
	57	+ @T.macro
	58	+ def counted_store_row_loss(buffer_x, buffer_y, buffer_loss, buffer_counts, row, ignore, row_max, row_sum):
	59	+ for _ in T.Parallel(1):
	60	+ if buffer_y[row] == ignore:
	61	+ buffer_counts[row] = 0
	62	+ buffer_loss[row] = T.float32(0)
	63	+ else:
	64	+ buffer_counts[row] = 1
	65	+ if scaled:
	66	+ half_gap = row_max[0] * T.float32(0.5) - T.cast(
	67	+ buffer_x[row, buffer_y[row]], 'float32'
	68	+ ) * T.float32(0.5)
	69	+ inv_n = T.float32(1) / T.cast(n, 'float32')
	70	+ buffer_loss[row] = T.log(row_sum[0]) * inv_n + half_gap * (T.float32(2) * inv_n)
	71	+ else:
	72	+ buffer_loss[row] = T.log(row_sum[0]) - (
	73	+ T.cast(buffer_x[row, buffer_y[row]], 'float32') - row_max[0]
	74	+ )
	75	+
	76	+ @T.macro
	77	+ def _counted_row_loss_body(buffer_x, buffer_y, buffer_loss, buffer_counts, ignore, core):
	78	+ x_ub = T.alloc_shared((block + 128,), dtype)
	79	+ row_max = T.alloc_shared((1,), 'float32')
	80	+ row_sum = T.alloc_shared((1,), 'float32')
	81	+ T.annotate_buffer_versions({x_ub: stages, row_max: 1, row_sum: 1})
	82	+ for wave in T.serial(T.ceildiv(n, cores)):
	83	+ row = wave * cores + core
	84	+ if row < n:
	85	+ with T.SimdVF():
	86	+ full_init = T.simd.pset(32)
	87	+ T.simd.vsts(
	88	+ row_max[0],
	89	+ T.simd.vdup(-T.infinity('float32'), T.float32, full_init),
	90	+ full_init,
	91	+ 'ONEPT_B32',
	92	+ )
	93	+ T.simd.vsts(row_sum[0], T.simd.vdup(0.0, T.float32, full_init), full_init, 'ONEPT_B32')
	94	+ for part in T.Pipelined(T.ceildiv(v, block), num_stages=stages):
	95	+ counted_statistics_block(buffer_x, row, part, x_ub, row_max, row_sum)
	96	+ with T.SimtVF(threads=32):
	97	+ counted_store_row_loss(
	98	+ buffer_x, buffer_y, buffer_loss, buffer_counts, row, ignore, row_max, row_sum
	99	+ )
	100	+
	101	+ @T.prim_func
	102	+ def row_loss(
	103	+ buffer_x: T.StridedTensor((n, v), (stride, 1), dtype),
	104	+ buffer_y: T.Tensor((n,), 'int64'),
	105	+ buffer_loss: T.Tensor((n,), 'float32'),
	106	+ buffer_counts: T.Tensor((n,), 'int32'),
	107	+ ignore: T.int64,
	108	+ ):
	109	+ with T.Kernel(cores) as core:
	110	+ _counted_row_loss_body(buffer_x, buffer_y, buffer_loss, buffer_counts, ignore, core)
	111	+
	112	+ return _compile(f'row_count_{dtype}_{block}_{cores}_{scaled}', row_loss)
	113	+
	114	+
	115	+@cache
	116	+def _partial_launch_kernel(dtype, block, cores, scaled):
	117	+ (n, v, k) = (T.dynamic(name, dtype='int64') for name in ('n', 'v', 'k'))
	118	+ stride = T.dynamic('stride', dtype='int64')
	119	+ stages = 2
	120	+
	121	+ update_row_statistics = _row_statistics_update(dtype)
	122	+
	123	+ @T.macro
	124	+ def partial_statistics_block(buffer_x, row, part, x_ub, row_max, row_sum):
	125	+ valid = T.min(T.int64(block), v - part * block)
	126	+ T.copy(buffer_x[row, part * block], x_ub[:valid])
	127	+ with T.SimdVF():
	128	+ update_row_statistics(x_ub, valid, row_max, row_sum)
	129	+
	130	+ @T.macro
	131	+ def partial_accumulate_row_loss(
	132	+ buffer_x, buffer_y, row, ignore, row_max, row_sum, partial_count, partial_sum, compensation
	133	+ ):
	134	+ for _ in T.Parallel(1):
	135	+ if buffer_y[row] != ignore:
	136	+ partial_count[0] = partial_count[0] + 1
	137	+ half_gap = row_max[0] * T.float32(0.5) - T.cast(buffer_x[row, buffer_y[row]], 'float32') * T.float32(
	138	+ 0.5
	139	+ )
	140	+ inv_n = T.float32(1) / T.cast(n, 'float32')
	141	+ row_loss_scaled = T.log(row_sum[0]) * inv_n + half_gap * (T.float32(2) * inv_n)
	142	+ old_sum = partial_sum[0]
	143	+ corrected = row_loss_scaled - compensation[0]
	144	+ updated = T.call_pure_extern('float32', 'fmaf', old_sum, T.float32(1), corrected)
	145	+ new_compensation = updated - old_sum - corrected
	146	+ compensation[0] = new_compensation
	147	+ partial_sum[0] = updated
	148	+
	149	+ @T.macro
	150	+ def partial_reduce_and_store(buffer_loss, buffer_counts, buffer_out):
	151	+ sums = T.alloc_reducer((1,), 'float32', op='sum')
	152	+ counts = T.alloc_reducer((1,), 'int32', op='sum')
	153	+ T.reducer_init(sums)
	154	+ T.reducer_init(counts)
	155	+ for i in T.Parallel(cores):
	156	+ T.reducer_update(sums[0], buffer_loss[i])
	157	+ T.reducer_update(counts[0], buffer_counts[i])
	158	+ total = T.alloc_fragment((1,), 'float32')
	159	+ count = T.alloc_fragment((1,), 'int32')
	160	+ T.finalize_reducer(sums, total)
	161	+ T.finalize_reducer(counts, count)
	162	+ for _ in T.Parallel(1):
	163	+ if scaled:
	164	+ buffer_out[0] = total[0] * (T.cast(n, 'float32') / T.cast(count[0], 'float32'))
	165	+ else:
	166	+ buffer_out[0] = total[0] * T.cast(n, 'float32')
	167	+
	168	+ @T.macro
	169	+ def _partial_launch_loss_body(buffer_x, buffer_y, buffer_loss, buffer_counts, buffer_out, ignore, core):
	170	+ x_ub = T.alloc_shared((block + 128,), dtype)
	171	+ row_max = T.alloc_shared((1,), 'float32')
	172	+ row_sum = T.alloc_shared((1,), 'float32')
	173	+ T.annotate_buffer_versions({x_ub: stages, row_max: 1, row_sum: 1})
	174	+ partial_sum = T.alloc_shared((1,), 'float32')
	175	+ partial_count = T.alloc_shared((1,), 'int32')
	176	+ compensation = T.alloc_shared((1,), 'float32')
	177	+ T.annotate_buffer_versions({partial_sum: 1, partial_count: 1, compensation: 1})
	178	+ partial_sum[0] = T.float32(0)
	179	+ partial_count[0] = 0
	180	+ compensation[0] = T.float32(0)
	181	+ for wave in T.serial(T.ceildiv(n, cores)):
	182	+ row = wave * cores + core
	183	+ if row < n:
	184	+ with T.SimdVF():
	185	+ row_max[0] = -T.infinity('float32')
	186	+ row_sum[0] = T.float32(0)
	187	+ for part in T.Pipelined(T.ceildiv(v, block), num_stages=stages):
	188	+ partial_statistics_block(buffer_x, row, part, x_ub, row_max, row_sum)
	189	+ with T.SimtVF(threads=32):
	190	+ partial_accumulate_row_loss(
	191	+ buffer_x, buffer_y, row, ignore, row_max, row_sum, partial_count, partial_sum, compensation
	192	+ )
	193	+ with T.PerCoreTask():
	194	+ with T.SimtVF(threads=32):
	195	+ for _ in T.Parallel(1):
	196	+ buffer_loss[core] = partial_sum[0]
	197	+ buffer_counts[core] = partial_count[0]
	198	+ T.ascend_threadfence()
	199	+ T.ascend_cross_core_set_flag(0, 'PIPE_V', 14)
	200	+ T.ascend_cross_core_wait_flag(0, 'PIPE_V', 14)
	201	+ if core == 0:
	202	+ with T.SimtVF(threads=256):
	203	+ partial_reduce_and_store(buffer_loss, buffer_counts, buffer_out)
	204	+
	205	+ @T.prim_func
	206	+ def row_loss(
	207	+ buffer_x: T.StridedTensor((n, v), (stride, 1), dtype),
	208	+ buffer_y: T.Tensor((n,), 'int64'),
	209	+ buffer_loss: T.Tensor((cores,), 'float32'),
	210	+ buffer_counts: T.Tensor((cores,), 'int32'),
	211	+ buffer_out: T.Tensor((1,), 'float32'),
	212	+ ignore: T.int64,
	213	+ ):
	214	+ with T.Kernel(cores) as core:
	215	+ _partial_launch_loss_body(buffer_x, buffer_y, buffer_loss, buffer_counts, buffer_out, ignore, core)
	216	+
	217	+ return _compile(f'row_partial_compensated_{dtype}_{block}_{cores}_{scaled}', row_loss)

Asrc/cann_ops_tilelang/loss/_cross_entropy_reduce.py+382-0

@@ -0,0 +1,382 @@
	1	+# Copyright (c) 2026 ops-tilelang contributors
	2	+# SPDX-License-Identifier: MIT
	3	+# See LICENSE in the repository root for the full license text.
	4	+
	5	+from functools import cache
	6	+
	7	+import tilelang.ascend.language as T
	8	+
	9	+from ._cross_entropy_common import ROW_BLOCK, _compile
	10	+
	11	+
	12	+def _make_small_loss_batch(n, v, width, rows):
	13	+ @T.macro
	14	+ def loss_batch(buffer_x, buffer_y, buffer_loss, buffer_total, start, ignore, mean):
	15	+ values = T.alloc_fragment((rows, width), 'float32')
	16	+ exponential = T.alloc_fragment((rows, width), 'float32')
	17	+ maximum = T.alloc_fragment((rows,), 'float32')
	18	+ total = T.alloc_fragment((rows,), 'float32')
	19	+ T.annotate_layout(
	20	+ {
	21	+ values: T.Fragment(
	22	+ (rows, width),
	23	+ forward_thread_fn=lambda r, c: r * 32 + c % 32,
	24	+ forward_index_fn=lambda r, c: c // 32,
	25	+ ),
	26	+ exponential: T.Fragment(
	27	+ (rows, width),
	28	+ forward_thread_fn=lambda r, c: r * 32 + c % 32,
	29	+ forward_index_fn=lambda r, c: c // 32,
	30	+ ),
	31	+ }
	32	+ )
	33	+ for r, col in T.Parallel(rows, width):
	34	+ values[r, col] = -T.infinity('float32')
	35	+ if start + r < n and col < v:
	36	+ values[r, col] = T.cast(buffer_x[start + r, col], 'float32')
	37	+ T.reduce_max(values, maximum, dim=1)
	38	+ for r, col in T.Parallel(rows, width):
	39	+ exponential[r, col] = T.float32(0)
	40	+ if start + r < n and col < v:
	41	+ exponential[r, col] = T.exp(values[r, col] - maximum[r])
	42	+ T.reduce_sum(exponential, total, dim=1)
	43	+ for r in T.Parallel(rows):
	44	+ row = start + r
	45	+ if row < n:
	46	+ if buffer_y[row] == ignore:
	47	+ buffer_loss[row] = T.float32(0)
	48	+ elif mean != 0:
	49	+ half_gap = maximum[r] * T.float32(0.5) - T.cast(
	50	+ buffer_x[row, buffer_y[row]], 'float32'
	51	+ ) * T.float32(0.5)
	52	+ inv_count = T.float32(1) / T.cast(buffer_total[0], 'float32')
	53	+ buffer_loss[row] = T.log(total[r]) * inv_count + half_gap * (T.float32(2) * inv_count)
	54	+ else:
	55	+ buffer_loss[row] = T.log(total[r]) - (T.cast(buffer_x[row, buffer_y[row]], 'float32') - maximum[r])
	56	+
	57	+ return loss_batch
	58	+
	59	+
	60	+@cache
	61	+def _batched_small_kernel(dtype, width, rows, cores):
	62	+ (n, v) = (T.dynamic(name, dtype='int64') for name in ('n', 'v'))
	63	+ stride = T.dynamic('stride', dtype='int64')
	64	+ loss_batch = _make_small_loss_batch(n, v, width, rows)
	65	+
	66	+ @T.prim_func
	67	+ def small_loss(
	68	+ buffer_x: T.StridedTensor((n, v), (stride, 1), dtype),
	69	+ buffer_y: T.Tensor((n,), 'int64'),
	70	+ buffer_loss: T.Tensor((n,), 'float32'),
	71	+ buffer_total: T.Tensor((1,), 'int64'),
	72	+ ignore: T.int64,
	73	+ mean: T.int32,
	74	+ ):
	75	+ with T.Kernel(cores) as core, T.SimtVF(threads=32 * rows):
	76	+ for wave in T.serial(T.ceildiv(n, cores * rows)):
	77	+ start = (wave * cores + core) * rows
	78	+ loss_batch(buffer_x, buffer_y, buffer_loss, buffer_total, start, ignore, mean)
	79	+
	80	+ return _compile(f'small_{dtype}_{width}_{rows}_{cores}', small_loss)
	81	+
	82	+
	83	+def _make_valid_count_tile(n, cores, acc_dtype):
	84	+ @T.macro
	85	+ def count_nonignored(buffer_y, row, ignore, count):
	86	+ if buffer_y[row] != ignore:
	87	+ T.reducer_update(count[0], T.cast(1, acc_dtype))
	88	+
	89	+ @T.macro
	90	+ def count_tile(buffer_y, count, tile, core, ignore):
	91	+ for i in T.Parallel(ROW_BLOCK):
	92	+ row = (tile * cores + core) * ROW_BLOCK + i
	93	+ if row < n:
	94	+ count_nonignored(buffer_y, row, ignore, count)
	95	+
	96	+ return count_tile
	97	+
	98	+
	99	+@cache
	100	+def _valid_count_kernel(wide=False, cores=1):
	101	+ acc_dtype = 'int64' if wide else 'int32'
	102	+ threads = 32 if wide else 256
	103	+ n = T.dynamic('n', dtype='int64')
	104	+ count_tile = _make_valid_count_tile(n, cores, acc_dtype)
	105	+
	106	+ @T.prim_func
	107	+ def count_valid(buffer_y: T.Tensor((n,), 'int64'), buffer_total: T.Tensor((cores,), 'int64'), ignore: T.int64):
	108	+ with T.Kernel(cores) as core, T.SimtVF(threads=threads):
	109	+ count = T.alloc_reducer((1,), acc_dtype, op='sum')
	110	+ T.reducer_init(count)
	111	+ for tile in T.serial(T.ceildiv(n, ROW_BLOCK * cores)):
	112	+ count_tile(buffer_y, count, tile, core, ignore)
	113	+ result = T.alloc_fragment((1,), acc_dtype)
	114	+ T.finalize_reducer(count, result)
	115	+ for i in T.Parallel(1):
	116	+ buffer_total[core] = T.cast(result[i], 'int64')
	117	+
	118	+ return _compile(f'count_valid_{acc_dtype}_{cores}', count_valid)
	119	+
	120	+
	121	+def _make_count_sum_tile(n):
	122	+ @T.macro
	123	+ def sum_tile(partial_counts, count, tile):
	124	+ for i in T.Parallel(32):
	125	+ row = tile * 32 + i
	126	+ if row < n:
	127	+ T.reducer_update(count[0], partial_counts[row])
	128	+
	129	+ return sum_tile
	130	+
	131	+
	132	+@cache
	133	+def _valid_count_merge_kernel():
	134	+ n = T.dynamic('n', dtype='int64')
	135	+ sum_tile = _make_count_sum_tile(n)
	136	+
	137	+ @T.prim_func
	138	+ def count_sum(partial_counts: T.Tensor((n,), 'int64'), total_count: T.Tensor((1,), 'int64')):
	139	+ with T.Kernel(1) as _, T.SimtVF(threads=32):
	140	+ count = T.alloc_reducer((1,), 'int64', op='sum')
	141	+ T.reducer_init(count)
	142	+ for tile in T.serial(T.ceildiv(n, 32)):
	143	+ sum_tile(partial_counts, count, tile)
	144	+ result = T.alloc_fragment((1,), 'int64')
	145	+ T.finalize_reducer(count, result)
	146	+ for _ in T.Parallel(1):
	147	+ total_count[0] = result[0]
	148	+
	149	+ return _compile('count_sum', count_sum)
	150	+
	151	+
	152	+def _make_merge_row(n, k, width):
	153	+ @T.macro
	154	+ def store_loss(buffer_x, buffer_y, buffer_loss, buffer_total, row, mean, m, s):
	155	+ for _ in T.Parallel(1):
	156	+ if mean != 0:
	157	+ half_gap = m[0] * T.float32(0.5) - T.cast(buffer_x[row, buffer_y[row]], 'float32') * T.float32(0.5)
	158	+ inv_count = T.alloc_fragment((1,), 'float32')
	159	+ if n == 1:
	160	+ inv_count[0] = T.float32(1)
	161	+ else:
	162	+ inv_count[0] = T.float32(1) / T.cast(buffer_total[0], 'float32')
	163	+ buffer_loss[row] = T.log(s[0]) * inv_count[0] + half_gap * (T.float32(2) * inv_count[0])
	164	+ else:
	165	+ buffer_loss[row] = T.log(s[0]) - (T.cast(buffer_x[row, buffer_y[row]], 'float32') - m[0])
	166	+
	167	+ @T.macro
	168	+ def merge_row(buffer_x, buffer_y, buffer_max, buffer_sum, buffer_loss, ignore, buffer_total, mean, row):
	169	+ if buffer_y[row] == ignore:
	170	+ for _ in T.Parallel(1):
	171	+ if mean != 0 and n == 1:
	172	+ buffer_loss[row] = T.float32(float('nan'))
	173	+ else:
	174	+ buffer_loss[row] = T.float32(0)
	175	+ else:
	176	+ mx = T.alloc_reducer((1,), 'float32', op='max')
	177	+ T.reducer_init(mx)
	178	+ for i in T.Parallel(width):
	179	+ if i < k:
	180	+ T.reducer_update(mx[0], buffer_max[row, i])
	181	+ m = T.alloc_fragment((1,), 'float32')
	182	+ T.finalize_reducer(mx, m)
	183	+ sm = T.alloc_reducer((1,), 'float32', op='sum')
	184	+ T.reducer_init(sm)
	185	+ for i in T.Parallel(width):
	186	+ if i < k:
	187	+ T.reducer_update(sm[0], buffer_sum[row, i] * T.exp(buffer_max[row, i] - m[0]))
	188	+ s = T.alloc_fragment((1,), 'float32')
	189	+ T.finalize_reducer(sm, s)
	190	+ store_loss(buffer_x, buffer_y, buffer_loss, buffer_total, row, mean, m, s)
	191	+
	192	+ return merge_row
	193	+
	194	+
	195	+@cache
	196	+def _merge_kernel(dtype, width, cores):
	197	+ (n, v, k) = (T.dynamic(name, dtype='int64') for name in ('n', 'v', 'k'))
	198	+ stride = T.dynamic('stride', dtype='int64')
	199	+ merge_row = _make_merge_row(n, k, width)
	200	+
	201	+ @T.prim_func
	202	+ def merge(
	203	+ buffer_x: T.StridedTensor((n, v), (stride, 1), dtype),
	204	+ buffer_y: T.Tensor((n,), 'int64'),
	205	+ buffer_max: T.Tensor((n, k), 'float32'),
	206	+ buffer_sum: T.Tensor((n, k), 'float32'),
	207	+ buffer_loss: T.Tensor((n,), 'float32'),
	208	+ ignore: T.int64,
	209	+ buffer_total: T.Tensor((1,), 'int64'),
	210	+ mean: T.int32,
	211	+ ):
	212	+ with T.Kernel(cores) as core, T.SimtVF(threads=32):
	213	+ for wave in T.serial(T.ceildiv(n, cores)):
	214	+ row = wave * cores + core
	215	+ if row < n:
	216	+ merge_row(buffer_x, buffer_y, buffer_max, buffer_sum, buffer_loss, ignore, buffer_total, mean, row)
	217	+
	218	+ return _compile(f'merge_{dtype}_{width}_{cores}', merge)
	219	+
	220	+
	221	+def _make_reduce_rows_block(n, capacity):
	222	+ @T.macro
	223	+ def reduce_block(buffer_loss, buffer_out, buffer_total, mean, block):
	224	+ sums = T.alloc_reducer((1,), 'float32', op='sum')
	225	+ T.reducer_init(sums)
	226	+ for i in T.Parallel(capacity):
	227	+ row = block * capacity + i
	228	+ if row < n:
	229	+ T.reducer_update(sums[0], buffer_loss[row])
	230	+ total = T.alloc_fragment((1,), 'float32')
	231	+ T.finalize_reducer(sums, total)
	232	+ for _ in T.Parallel(1):
	233	+ if mean != 0:
	234	+ buffer_out[block] = total[0] / T.cast(T.min(buffer_total[0], T.int64(1)), 'float32')
	235	+ else:
	236	+ buffer_out[block] = total[0]
	237	+
	238	+ return reduce_block
	239	+
	240	+
	241	+@cache
	242	+def _reduce_kernel(cores, capacity=ROW_BLOCK):
	243	+ (n, blocks) = (T.dynamic('n', dtype='int64'), T.dynamic('blocks', dtype='int64'))
	244	+ reduce_block = _make_reduce_rows_block(n, capacity)
	245	+
	246	+ @T.prim_func
	247	+ def reduce_rows(
	248	+ buffer_loss: T.Tensor((n,), 'float32'),
	249	+ buffer_out: T.Tensor((blocks,), 'float32'),
	250	+ buffer_total: T.Tensor((1,), 'int64'),
	251	+ mean: T.int32,
	252	+ ):
	253	+ with T.Kernel(cores) as core, T.SimtVF(threads=256):
	254	+ for wave in T.serial(T.ceildiv(blocks, cores)):
	255	+ block = wave * cores + core
	256	+ if block < blocks:
	257	+ reduce_block(buffer_loss, buffer_out, buffer_total, mean, block)
	258	+
	259	+ return _compile(f'reduce_rows_{cores}_{capacity}', reduce_rows)
	260	+
	261	+
	262	+def _make_sum_count_block(n, capacity, first=False, final=False, joint=False):
	263	+ if joint:
	264	+
	265	+ @T.macro
	266	+ def update_count(buffer_counts, counts, row, ignore):
	267	+ T.reducer_update(counts[0], buffer_counts[row])
	268	+
	269	+ @T.macro
	270	+ def store_result(buffer_out, buffer_outcounts, total, valid, block, original_n, final_step, scaled):
	271	+ buffer_outcounts[block] = valid[0]
	272	+ if final_step != 0:
	273	+ if scaled != 0:
	274	+ buffer_out[block] = total[0] * (T.cast(original_n, 'float32') / T.cast(valid[0], 'float32'))
	275	+ else:
	276	+ buffer_out[block] = total[0] / T.cast(valid[0], 'float32')
	277	+ else:
	278	+ buffer_out[block] = total[0]
	279	+ else:
	280	+
	281	+ @T.macro
	282	+ def update_count(buffer_counts, counts, row, ignore):
	283	+ if first:
	284	+ T.reducer_update(counts[0], T.cast(buffer_counts[row] != ignore, 'int32'))
	285	+ else:
	286	+ T.reducer_update(counts[0], T.cast(buffer_counts[row], 'int32'))
	287	+
	288	+ @T.macro
	289	+ def store_result(buffer_out, buffer_outcounts, total, valid, block, original_n, final_step, scaled):
	290	+ buffer_outcounts[block] = T.cast(valid[0], 'int64')
	291	+ if final:
	292	+ buffer_out[block] = total[0] / T.cast(valid[0], 'float32')
	293	+ else:
	294	+ buffer_out[block] = total[0]
	295	+
	296	+ @T.macro
	297	+ def reduce_block(
	298	+ buffer_values, buffer_counts, buffer_out, buffer_outcounts, block, ignore, original_n, final_step, scaled
	299	+ ):
	300	+ sums = T.alloc_reducer((1,), 'float32', op='sum')
	301	+ counts = T.alloc_reducer((1,), 'int32', op='sum')
	302	+ T.reducer_init(sums)
	303	+ T.reducer_init(counts)
	304	+ for i in T.Parallel(capacity):
	305	+ row = block * capacity + i
	306	+ if row < n:
	307	+ T.reducer_update(sums[0], buffer_values[row])
	308	+ update_count(buffer_counts, counts, row, ignore)
	309	+ total = T.alloc_fragment((1,), 'float32')
	310	+ valid = T.alloc_fragment((1,), 'int32')
	311	+ T.finalize_reducer(sums, total)
	312	+ T.finalize_reducer(counts, valid)
	313	+ for _ in T.Parallel(1):
	314	+ store_result(buffer_out, buffer_outcounts, total, valid, block, original_n, final_step, scaled)
	315	+
	316	+ return reduce_block
	317	+
	318	+
	319	+@cache
	320	+def _reduce_sum_count_kernel(cores, capacity=ROW_BLOCK, first=False, final=False):
	321	+ (n, blocks) = (T.dynamic('n', dtype='int64'), T.dynamic('blocks', dtype='int64'))
	322	+ reduce_block = _make_sum_count_block(n, capacity, first=first, final=final)
	323	+
	324	+ @T.prim_func
	325	+ def reduce_sum_count(
	326	+ buffer_values: T.Tensor((n,), 'float32'),
	327	+ buffer_counts: T.Tensor((n,), 'int64'),
	328	+ buffer_out: T.Tensor((blocks,), 'float32'),
	329	+ buffer_outcounts: T.Tensor((blocks,), 'int64'),
	330	+ ignore: T.int64,
	331	+ ):
	332	+ with T.Kernel(cores) as core, T.SimtVF(threads=256):
	333	+ for wave in T.serial(T.ceildiv(blocks, cores)):
	334	+ block = wave * cores + core
	335	+ if block < blocks:
	336	+ reduce_block(
	337	+ buffer_values,
	338	+ buffer_counts,
	339	+ buffer_out,
	340	+ buffer_outcounts,
	341	+ block,
	342	+ ignore,
	343	+ T.int64(0),
	344	+ T.int32(0),
	345	+ T.int32(0),
	346	+ )
	347	+
	348	+ return _compile(f'reduce_sum_count_{cores}_{capacity}_{int(first)}_{int(final)}', reduce_sum_count)
	349	+
	350	+
	351	+@cache
	352	+def _joint_reduce_kernel(cores, capacity=ROW_BLOCK):
	353	+ (n, original_n, blocks) = (T.dynamic(name, dtype='int64') for name in ('n', 'original_n', 'blocks'))
	354	+ reduce_block = _make_sum_count_block(n, capacity, joint=True)
	355	+
	356	+ @T.prim_func
	357	+ def joint_reduce(
	358	+ buffer_loss: T.Tensor((n,), 'float32'),
	359	+ buffer_counts: T.Tensor((n,), 'int32'),
	360	+ buffer_out: T.Tensor((blocks,), 'float32'),
	361	+ buffer_outcount: T.Tensor((blocks,), 'int32'),
	362	+ original_n: T.int64,
	363	+ final: T.int32,
	364	+ scaled: T.int32,
	365	+ ):
	366	+ with T.Kernel(cores) as core, T.SimtVF(threads=256):
	367	+ for wave in T.serial(T.ceildiv(blocks, cores)):
	368	+ block = wave * cores + core
	369	+ if block < blocks:
	370	+ reduce_block(
	371	+ buffer_loss,
	372	+ buffer_counts,
	373	+ buffer_out,
	374	+ buffer_outcount,
	375	+ block,
	376	+ T.int64(0),
	377	+ original_n,
	378	+ final,
	379	+ scaled,
	380	+ )
	381	+
	382	+ return _compile(f'joint_reduce_{cores}_{capacity}', joint_reduce)

Asrc/cann_ops_tilelang/loss/_cross_entropy_row.py+443-0

@@ -0,0 +1,443 @@
	1	+# Copyright (c) 2026 ops-tilelang contributors
	2	+# SPDX-License-Identifier: MIT
	3	+# See LICENSE in the repository root for the full license text.
	4	+
	5	+from functools import cache
	6	+
	7	+import tilelang.ascend.language as T
	8	+
	9	+from ._cross_entropy_common import _compile
	10	+from ._cross_entropy_simd import _block_maximum, _plain_exponential_lanes, _row_block_statistics_update
	11	+
	12	+_ROW_PIPELINE_STAGES = 2
	13	+
	14	+
	15	+def _row_block_accumulator(dtype, block, v, single_block, stages):
	16	+ update_row_statistics = _row_block_statistics_update(dtype, single_block)
	17	+
	18	+ @T.macro
	19	+ def accumulate_block(buffer_x, row, x_ub, maxbuf, sumbuf, wave, sidx):
	20	+ for part in T.Pipelined(T.ceildiv(v, block), num_stages=stages):
	21	+ valid = T.min(T.int64(block), v - part * block)
	22	+ T.copy(buffer_x[row, part * block], x_ub[:valid])
	23	+ with T.SimdVF():
	24	+ update_row_statistics(x_ub, valid, maxbuf, sumbuf, wave, sidx)
	25	+
	26	+ return accumulate_block
	27	+
	28	+
	29	+@T.macro
	30	+def _simt_mean_half_gap(buffer_x, buffer_y, row, maxbuf, midx):
	31	+ return maxbuf[midx] * T.float32(0.5) - T.cast(buffer_x[row, buffer_y[row]], 'float32') * T.float32(0.5)
	32	+
	33	+
	34	+def _simt_mean_inverse(n, prescale):
	35	+ @T.macro
	36	+ def inverse(buffer_total):
	37	+ if prescale:
	38	+ return T.float32(1) / T.cast(n, 'float32')
	39	+ else:
	40	+ return T.float32(1) / T.cast(buffer_total[0], 'float32')
	41	+
	42	+ return inverse
	43	+
	44	+
	45	+def _batched_simt_finisher(n, batch_slots, prescale):
	46	+ mean_inverse = _simt_mean_inverse(n, prescale)
	47	+
	48	+ @T.macro
	49	+ def store_ignored(buffer_loss, owned_row, mean):
	50	+ if mean != 0 and n == 1:
	51	+ buffer_loss[owned_row] = T.float32(float('nan'))
	52	+ else:
	53	+ buffer_loss[owned_row] = T.float32(0)
	54	+
	55	+ @T.macro
	56	+ def prepare_inverse(buffer_total, inv_count):
	57	+ if n == 1:
	58	+ inv_count[0] = T.float32(1)
	59	+ else:
	60	+ inv_count[0] = mean_inverse(buffer_total)
	61	+
	62	+ @T.macro
	63	+ def finish(buffer_x, buffer_y, buffer_loss, buffer_total, owned_row, slot, ignore, mean, stats):
	64	+ if owned_row < n:
	65	+ if buffer_y[owned_row] == ignore:
	66	+ if prescale:
	67	+ buffer_total[owned_row] = 0
	68	+ store_ignored(buffer_loss, owned_row, mean)
	69	+ else:
	70	+ if prescale:
	71	+ buffer_total[owned_row] = 1
	72	+ if mean != 0:
	73	+ half_gap = _simt_mean_half_gap(buffer_x, buffer_y, owned_row, stats, slot)
	74	+ inv_count = T.alloc_fragment((1,), 'float32')
	75	+ prepare_inverse(buffer_total, inv_count)
	76	+ buffer_loss[owned_row] = T.log(stats[batch_slots + slot]) * inv_count[0] + half_gap * (
	77	+ T.float32(2) * inv_count[0]
	78	+ )
	79	+ else:
	80	+ buffer_loss[owned_row] = T.log(stats[batch_slots + slot]) - (
	81	+ T.cast(buffer_x[owned_row, buffer_y[owned_row]], 'float32') - stats[slot]
	82	+ )
	83	+
	84	+ return finish
	85	+
	86	+
	87	+def _wave_simt_finisher(n, prescale):
	88	+ mean_inverse = _simt_mean_inverse(n, prescale)
	89	+ sidx = 8 if prescale else 0
	90	+
	91	+ @T.macro
	92	+ def finish_row(buffer_x, buffer_y, buffer_loss, buffer_total, row, ignore, mean, row_max, row_sum):
	93	+ if buffer_y[row] == ignore:
	94	+ if mean != 0 and n == 1:
	95	+ buffer_loss[row] = T.float32(float('nan'))
	96	+ else:
	97	+ buffer_loss[row] = T.float32(0)
	98	+ elif mean != 0:
	99	+ half_gap = _simt_mean_half_gap(buffer_x, buffer_y, row, row_max, 0)
	100	+ inv_count = T.alloc_fragment((1,), 'float32')
	101	+ if n == 1:
	102	+ inv_count[0] = T.float32(1)
	103	+ else:
	104	+ inv_count[0] = mean_inverse(buffer_total)
	105	+ buffer_loss[row] = T.log(row_sum[sidx]) * inv_count[0] + half_gap * (T.float32(2) * inv_count[0])
	106	+ else:
	107	+ buffer_loss[row] = T.log(row_sum[sidx]) - (T.cast(buffer_x[row, buffer_y[row]], 'float32') - row_max[0])
	108	+
	109	+ @T.macro
	110	+ def finish(buffer_x, buffer_y, buffer_loss, buffer_total, row, ignore, mean, row_max, row_sum):
	111	+ with T.SimtVF(threads=32):
	112	+ for _ in T.Parallel(1):
	113	+ finish_row(buffer_x, buffer_y, buffer_loss, buffer_total, row, ignore, mean, row_max, row_sum)
	114	+
	115	+ return finish
	116	+
	117	+
	118	+def _row_factory_setup(dtype, block, single_block, batch_slots, static_v):
	119	+ if static_v is not None and not (single_block and 0 < static_v <= block):
	120	+ raise ValueError('static_v requires a single block and 0 < static_v <= block')
	121	+ n = T.dynamic('n', dtype='int64')
	122	+ v = T.dynamic('v', dtype='int64') if static_v is None else static_v
	123	+ stride = T.dynamic('stride', dtype='int64')
	124	+
	125	+ if batch_slots and not single_block:
	126	+ raise ValueError('batch_slots requires a single block')
	127	+ stats_size = 2 * batch_slots if batch_slots else 16
	128	+
	129	+ accumulate_row_blocks = _row_block_accumulator(dtype, block, v, single_block, _ROW_PIPELINE_STAGES)
	130	+ return n, v, stride, stats_size, accumulate_row_blocks
	131	+
	132	+
	133	+@cache
	134	+def _row_kernel(dtype, block, cores, single_block=False, batch_slots=0, static_v=None):
	135	+ n, v, stride, stats_size, accumulate_row_blocks = _row_factory_setup(
	136	+ dtype, block, single_block, batch_slots, static_v
	137	+ )
	138	+
	139	+ @T.macro
	140	+ def initialize_row_statistics(row_max, row_sum):
	141	+ with T.SimdVF():
	142	+ full_init = T.simd.pset(32)
	143	+ T.simd.vsts(
	144	+ row_max[0],
	145	+ T.simd.vdup(-T.infinity('float32'), T.float32, full_init),
	146	+ full_init,
	147	+ 'ONEPT_B32',
	148	+ )
	149	+ T.simd.vsts(row_sum[0], T.simd.vdup(0.0, T.float32, full_init), full_init, 'ONEPT_B32')
	150	+
	151	+ store_row_loss = _wave_simt_finisher(n, False)
	152	+ store_batched_row_loss = _batched_simt_finisher(n, batch_slots, False)
	153	+
	154	+ @T.macro
	155	+ def _row_loss_body(buffer_x, buffer_y, buffer_loss, buffer_total, ignore, mean, core):
	156	+ x_ub = T.alloc_shared((block + 128,), dtype)
	157	+ if batch_slots:
	158	+ stats = T.alloc_shared((stats_size,), 'float32')
	159	+ T.annotate_buffer_versions({x_ub: _ROW_PIPELINE_STAGES, stats: 1})
	160	+ else:
	161	+ row_max = T.alloc_shared((1,), 'float32')
	162	+ row_sum = T.alloc_shared((1,), 'float32')
	163	+ T.annotate_buffer_versions({x_ub: _ROW_PIPELINE_STAGES, row_max: 1, row_sum: 1})
	164	+ for wave in T.serial(T.ceildiv(n, cores)):
	165	+ row = wave * cores + core
	166	+ if row < n:
	167	+ if not single_block:
	168	+ initialize_row_statistics(row_max, row_sum)
	169	+ if batch_slots:
	170	+ accumulate_row_blocks(buffer_x, row, x_ub, stats, stats, wave, batch_slots + wave)
	171	+ else:
	172	+ accumulate_row_blocks(buffer_x, row, x_ub, row_max, row_sum, 0, 0)
	173	+ if not batch_slots:
	174	+ store_row_loss(buffer_x, buffer_y, buffer_loss, buffer_total, row, ignore, mean, row_max, row_sum)
	175	+ if batch_slots:
	176	+ T.ascend_pipe_barrier('PIPE_ALL')
	177	+ with T.SimtVF(threads=max(32, batch_slots)):
	178	+ slot = T.get_thread_binding()
	179	+ if slot < batch_slots:
	180	+ owned_row = slot * cores + core
	181	+ store_batched_row_loss(
	182	+ buffer_x, buffer_y, buffer_loss, buffer_total, owned_row, slot, ignore, mean, stats
	183	+ )
	184	+
	185	+ @T.prim_func
	186	+ def row_loss(
	187	+ buffer_x: T.StridedTensor((n, v), (stride, 1), dtype),
	188	+ buffer_y: T.Tensor((n,), 'int64'),
	189	+ buffer_loss: T.Tensor((n,), 'float32'),
	190	+ buffer_total: T.Tensor((1,), 'int64'),
	191	+ ignore: T.int64,
	192	+ mean: T.int32,
	193	+ ):
	194	+ with T.Kernel(cores) as core:
	195	+ _row_loss_body(buffer_x, buffer_y, buffer_loss, buffer_total, ignore, mean, core)
	196	+
	197	+ return _compile(f'row_{dtype}_{block}_{cores}_{int(single_block)}_{batch_slots}_{static_v}', row_loss)
	198	+
	199	+
	200	+@cache
	201	+def _simd_row_kernel(dtype, block, cores, inline_count=False):
	202	+ (n, v, k) = (T.dynamic(name, dtype='int64') for name in ('n', 'v', 'k'))
	203	+ stride = T.dynamic('stride', dtype='int64')
	204	+ narrow = dtype != 'float32'
	205	+ stages = 2
	206	+
	207	+ block_maximum = _block_maximum(dtype)
	208	+ exponential_lanes = _plain_exponential_lanes(narrow)
	209	+
	210	+ @T.macro
	211	+ def count_nonignored_pair(buffer_y, i, ignore, count):
	212	+ if buffer_y[2 * i] != T.cast(ignore, 'int32') or buffer_y[2 * i + 1] != T.cast(ignore >> 32, 'int32'):
	213	+ T.reducer_update(count[0], T.int32(1))
	214	+
	215	+ @T.macro
	216	+ def reduce_inline_count(buffer_y, buffer_total, count_ub, ignore, core):
	217	+ with T.SimtVF(threads=32):
	218	+ count = T.alloc_reducer((1,), 'int32', op='sum')
	219	+ T.reducer_init(count)
	220	+ for i in T.Parallel(32):
	221	+ if i < n:
	222	+ count_nonignored_pair(buffer_y, i, ignore, count)
	223	+ result = T.alloc_fragment((1,), 'int32')
	224	+ T.finalize_reducer(count, result)
	225	+ for _ in T.Parallel(1):
	226	+ count_ub[0] = result[0]
	227	+ if core == 0:
	228	+ buffer_total[0] = result[0]
	229	+ buffer_total[1] = T.int32(0)
	230	+
	231	+ @T.macro
	232	+ def initialize_simd_row_statistics(row_max, row_sum, row_answer):
	233	+ with T.SimdVF():
	234	+ full_init = T.simd.pset(32)
	235	+ T.simd.vsts(
	236	+ row_max[0],
	237	+ T.simd.vdup(-T.infinity('float32'), T.float32, full_init),
	238	+ full_init,
	239	+ 'ONEPT_B32',
	240	+ )
	241	+ T.simd.vsts(row_sum[0], T.simd.vdup(0.0, T.float32, full_init), full_init, 'ONEPT_B32')
	242	+ T.simd.vsts(row_answer[0], T.simd.vdup(0.0, T.float32, full_init), full_init, 'ONEPT_B32')
	243	+
	244	+ @T.macro
	245	+ def compute_online_statistics(x_ub, valid, row_max, row_sum, full):
	246	+ neg = T.simd.vdup(-T.infinity('float32'), T.float32, full)
	247	+ block_max = block_maximum(x_ub, valid, full, neg)
	248	+ run_max = T.simd.vld(row_max[0], dist='BRC_B32')
	249	+ run_sum = T.simd.vld(row_sum[0], dist='BRC_B32')
	250	+ m = T.simd.vmax(block_max, run_max, full)
	251	+ ls = exponential_lanes(x_ub, valid, m, full, neg)
	252	+ block_sum = T.simd.vdupv(T.simd.vcadd(ls, full), full)
	253	+ s = T.simd.vadd(
	254	+ T.simd.vmul(run_sum, T.simd.vexp(T.simd.vsub(run_max, m, full), full), full),
	255	+ block_sum,
	256	+ full,
	257	+ )
	258	+ return m, s
	259	+
	260	+ @T.macro
	261	+ def store_final_row_result(m, s, answer, ignored, mean, count_ub, output_ub, wave, full):
	262	+ result = T.simd.alloc_var(T.float32)
	263	+ if mean != 0:
	264	+ inverse = T.simd.alloc_var(T.float32)
	265	+ if n == 1:
	266	+ inverse = T.simd.vdup(1.0, T.float32, full)
	267	+ else:
	268	+ count = T.simd.vcvt(T.simd.vld(count_ub[0], dist='BRC_B32'), T.float32, full)
	269	+ inverse = T.simd.vdiv(T.simd.vdup(1.0, T.float32, full), count, full)
	270	+ half = T.simd.vdup(0.5, T.float32, full)
	271	+ gap = T.simd.vsub(T.simd.vmul(m, half, full), T.simd.vmul(answer, half, full), full)
	272	+ result = T.simd.vadd(
	273	+ T.simd.vmul(T.simd.vln(s, full), inverse, full),
	274	+ T.simd.vmul(gap, T.simd.vmul(T.simd.vdup(2.0, T.float32, full), inverse, full), full),
	275	+ full,
	276	+ )
	277	+ else:
	278	+ result = T.simd.vsub(T.simd.vln(s, full), T.simd.vsub(answer, m, full), full)
	279	+ if mean != 0 and n == 1:
	280	+ result = T.simd.vsel(T.simd.vdup(T.float32(float('nan')), T.float32, full), result, ignored)
	281	+ else:
	282	+ result = T.simd.vsel(T.simd.vdup(0.0, T.float32, full), result, ignored)
	283	+ T.simd.vsts(output_ub[wave], result, full, 'ONEPT_B32')
	284	+
	285	+ @T.macro
	286	+ def update_simd_block_and_store(x_ub, valid, row_stats, label_ub, ignore, part, mean, count_ub, output_ub, wave):
	287	+ row_max, row_sum, row_answer = row_stats
	288	+ with T.SimdVF():
	289	+ full = T.simd.pset(32)
	290	+ m, s = compute_online_statistics(x_ub, valid, row_max, row_sum, full)
	291	+ low = T.simd.vld(label_ub[0], dist='BRC_B32')
	292	+ high = T.simd.vld(label_ub[1], dist='BRC_B32')
	293	+ ignored = T.simd.pand(
	294	+ T.simd.vcmps(low, T.cast(ignore, 'int32'), full, 'eq'),
	295	+ T.simd.vcmps(high, T.cast(ignore >> 32, 'int32'), full, 'eq'),
	296	+ full,
	297	+ )
	298	+ in_block = T.simd.pand(
	299	+ T.simd.vcmps(low, T.cast(part * block, 'int32'), full, 'ge'),
	300	+ T.simd.vcmps(low, T.cast(part * block + valid, 'int32'), full, 'lt'),
	301	+ full,
	302	+ )
	303	+ in_block = T.simd.pand(in_block, T.simd.vcmps(high, 0, full, 'eq'), full)
	304	+ offset = T.simd.vsub(low, T.simd.vdup(T.cast(part * block, 'int32'), T.int32, full), full)
	305	+ safe_offset = T.simd.vsel(offset, T.simd.vdup(0, T.int32, full), in_block)
	306	+ answer = T.simd.vld(row_answer[0], dist='BRC_B32')
	307	+ if narrow:
	308	+ gather_full = T.simd.pset(16)
	309	+ gather_index = T.simd.vdupv(T.simd.vcvt(safe_offset, T.uint16, full), gather_full)
	310	+ gathered = T.simd.vcvt(
	311	+ T.simd.vgather2(T.access_ptr(x_ub[0], 'r', extent=block), gather_index, gather_full),
	312	+ T.float32,
	313	+ part=0,
	314	+ )
	315	+ else:
	316	+ gathered = T.simd.vgather2(
	317	+ T.access_ptr(x_ub[0], 'r', extent=block),
	318	+ T.reinterpret(safe_offset, 'uint32x64'),
	319	+ full,
	320	+ )
	321	+ answer = T.simd.vsel(gathered, answer, in_block)
	322	+ if part + 1 == T.ceildiv(v, block):
	323	+ store_final_row_result(m, s, answer, ignored, mean, count_ub, output_ub, wave, full)
	324	+ else:
	325	+ T.simd.vsts(row_max[0], m, full, 'ONEPT_B32')
	326	+ T.simd.vsts(row_sum[0], s, full, 'ONEPT_B32')
	327	+ T.simd.vsts(row_answer[0], answer, full, 'ONEPT_B32')
	328	+
	329	+ @T.macro
	330	+ def _simd_row_loss_body(buffer_x, buffer_y, buffer_loss, buffer_total, ignore, mean, core):
	331	+ x_ub = T.alloc_shared((block + 128,), dtype)
	332	+ row_max = T.alloc_shared((1,), 'float32')
	333	+ row_sum = T.alloc_shared((1,), 'float32')
	334	+ row_answer = T.alloc_shared((1,), 'float32')
	335	+ output_ub = T.alloc_shared((8,), 'float32')
	336	+ label_ub = T.alloc_shared((8,), 'int32')
	337	+ count_ub = T.alloc_shared((1 if inline_count else 8,), 'int32')
	338	+ T.annotate_buffer_versions({x_ub: stages, row_max: 1, row_sum: 1, row_answer: 1, output_ub: 1})
	339	+ if inline_count:
	340	+ reduce_inline_count(buffer_y, buffer_total, count_ub, ignore, core)
	341	+ T.ascend_pipe_barrier('PIPE_ALL')
	342	+ elif mean != 0 and n > 1:
	343	+ T.copy(buffer_total[0:2], count_ub[0:2])
	344	+ row_begin = n * core // cores
	345	+ row_end = n * (core + 1) // cores
	346	+ for wave in T.Pipelined(row_end - row_begin, num_stages=2):
	347	+ row = row_begin + wave
	348	+ if row < row_end:
	349	+ T.copy(buffer_y[2 * row], label_ub[0:2])
	350	+ initialize_simd_row_statistics(row_max, row_sum, row_answer)
	351	+ for part in T.Pipelined(T.ceildiv(v, block), num_stages=stages):
	352	+ valid = T.min(T.int64(block), v - part * block)
	353	+ T.copy(buffer_x[row, part * block], x_ub[:valid])
	354	+ update_simd_block_and_store(
	355	+ x_ub,
	356	+ valid,
	357	+ (row_max, row_sum, row_answer),
	358	+ label_ub,
	359	+ ignore,
	360	+ part,
	361	+ mean,
	362	+ count_ub,
	363	+ output_ub,
	364	+ wave,
	365	+ )
	366	+ if row_begin < row_end:
	367	+ T.copy(output_ub[: row_end - row_begin], buffer_loss[row_begin:row_end], l2_cache_ctrl=0)
	368	+
	369	+ @T.prim_func
	370	+ def row_loss(
	371	+ buffer_x: T.StridedTensor((n, v), (stride, 1), dtype),
	372	+ buffer_y: T.Tensor((2 * n,), 'int32'),
	373	+ buffer_loss: T.Tensor((n,), 'float32'),
	374	+ buffer_total: T.Tensor((2,), 'int32'),
	375	+ ignore: T.int64,
	376	+ mean: T.int32,
	377	+ ):
	378	+ with T.Kernel(cores) as core:
	379	+ _simd_row_loss_body(buffer_x, buffer_y, buffer_loss, buffer_total, ignore, mean, core)
	380	+
	381	+ return _compile(f'row_fullsimd_{dtype}_{block}_{cores}_{int(inline_count)}', row_loss)
	382	+
	383	+
	384	+@cache
	385	+def _scaled_row_kernel(dtype, block, cores, single_block=False, batch_slots=0, static_v=None):
	386	+ n, v, stride, stats_size, accumulate_row_blocks = _row_factory_setup(
	387	+ dtype, block, single_block, batch_slots, static_v
	388	+ )
	389	+
	390	+ @T.macro
	391	+ def initialize_scaled_statistics(stats):
	392	+ with T.SimdVF():
	393	+ full_init = T.simd.pset(32)
	394	+ T.simd.vsts(
	395	+ stats[0],
	396	+ T.simd.vdup(-T.infinity('float32'), T.float32, full_init),
	397	+ full_init,
	398	+ 'ONEPT_B32',
	399	+ )
	400	+ T.simd.vsts(stats[8], T.simd.vdup(0.0, T.float32, full_init), full_init, 'ONEPT_B32')
	401	+
	402	+ store_scaled_row_loss = _wave_simt_finisher(n, True)
	403	+ store_batched_scaled_loss = _batched_simt_finisher(n, batch_slots, True)
	404	+
	405	+ @T.macro
	406	+ def _scaled_row_loss_body(buffer_x, buffer_y, buffer_loss, buffer_total, ignore, mean, core):
	407	+ x_ub = T.alloc_shared((block + 128,), dtype)
	408	+ stats = T.alloc_shared((stats_size,), 'float32')
	409	+ T.annotate_buffer_versions({x_ub: _ROW_PIPELINE_STAGES, stats: 1})
	410	+ for wave in T.serial(T.ceildiv(n, cores)):
	411	+ row = wave * cores + core
	412	+ if row < n:
	413	+ if not single_block:
	414	+ initialize_scaled_statistics(stats)
	415	+ if batch_slots:
	416	+ accumulate_row_blocks(buffer_x, row, x_ub, stats, stats, wave, batch_slots + wave)
	417	+ else:
	418	+ accumulate_row_blocks(buffer_x, row, x_ub, stats, stats, 0, 8)
	419	+ if not batch_slots:
	420	+ store_scaled_row_loss(buffer_x, buffer_y, buffer_loss, stats, row, ignore, mean, stats, stats)
	421	+ if batch_slots:
	422	+ T.ascend_pipe_barrier('PIPE_ALL')
	423	+ with T.SimtVF(threads=max(32, batch_slots)):
	424	+ slot = T.get_thread_binding()
	425	+ if slot < batch_slots:
	426	+ owned_row = slot * cores + core
	427	+ store_batched_scaled_loss(
	428	+ buffer_x, buffer_y, buffer_loss, buffer_total, owned_row, slot, ignore, mean, stats
	429	+ )
	430	+
	431	+ @T.prim_func
	432	+ def row_loss(
	433	+ buffer_x: T.StridedTensor((n, v), (stride, 1), dtype),
	434	+ buffer_y: T.Tensor((n,), 'int64'),
	435	+ buffer_loss: T.Tensor((n,), 'float32'),
	436	+ buffer_total: T.Tensor((n,), 'int32'),
	437	+ ignore: T.int64,
	438	+ mean: T.int32,
	439	+ ):
	440	+ with T.Kernel(cores) as core:
	441	+ _scaled_row_loss_body(buffer_x, buffer_y, buffer_loss, buffer_total, ignore, mean, core)
	442	+
	443	+ return _compile(f'row_{dtype}_{block}_{cores}_{int(single_block)}_{batch_slots}_{static_v}', row_loss)

Asrc/cann_ops_tilelang/loss/_cross_entropy_simd.py+210-0

@@ -0,0 +1,210 @@
	1	+# Copyright (c) 2026 ops-tilelang contributors
	2	+# SPDX-License-Identifier: MIT
	3	+# See LICENSE in the repository root for the full license text.
	4	+
	5	+import tilelang.ascend.language as T
	6	+
	7	+
	8	+def _half_block_maximum(dtype):
	9	+ @T.macro
	10	+ def block_maximum(x_ub, valid, full, neg):
	11	+ block_max = T.simd.alloc_var(T.float32)
	12	+ half_full = T.simd.pset(16)
	13	+ half_neg = T.simd.vdup(-T.infinity(dtype), dtype, half_full)
	14	+ h0 = T.simd.alloc_var(dtype)
	15	+ h1 = T.simd.alloc_var(dtype)
	16	+ h2 = T.simd.alloc_var(dtype)
	17	+ h3 = T.simd.alloc_var(dtype)
	18	+ h0 = half_neg
	19	+ h1 = half_neg
	20	+ h2 = half_neg
	21	+ h3 = half_neg
	22	+ if valid >= 512:
	23	+ for vec in T.serial(valid // 512):
	24	+ h0 = T.simd.vmax(h0, T.simd.vld(x_ub[vec * 512]), half_full)
	25	+ h1 = T.simd.vmax(h1, T.simd.vld(x_ub[vec * 512 + 128]), half_full)
	26	+ h2 = T.simd.vmax(h2, T.simd.vld(x_ub[vec * 512 + 256]), half_full)
	27	+ h3 = T.simd.vmax(h3, T.simd.vld(x_ub[vec * 512 + 384]), half_full)
	28	+ h0 = T.simd.vmax(h0, h1, half_full)
	29	+ h0 = T.simd.vmax(h0, h2, half_full)
	30	+ h0 = T.simd.vmax(h0, h3, half_full)
	31	+ if valid // 128 % 4 > 0:
	32	+ for vec in T.serial(valid // 128 % 4):
	33	+ h0 = T.simd.vmax(h0, T.simd.vld(x_ub[(valid // 512 * 4 + vec) * 128]), half_full)
	34	+ if valid % 128 != 0:
	35	+ half_tail = T.simd.update_mask(valid % 128, 16)
	36	+ half_value = T.simd.vld(x_ub[valid // 128 * 128])
	37	+ h0 = T.simd.vmax(h0, T.simd.vsel(half_value, half_neg, half_tail), half_full)
	38	+ max_even = T.simd.vcvt(h0, T.float32, part=0)
	39	+ max_odd = T.simd.vcvt(h0, T.float32, part=1)
	40	+ block_max = T.simd.vdupv(T.simd.vcmax(T.simd.vmax(max_even, max_odd, full), full), full)
	41	+ return block_max
	42	+
	43	+ return block_maximum
	44	+
	45	+
	46	+def _float_block_maximum():
	47	+ @T.macro
	48	+ def block_maximum(x_ub, valid, full, neg):
	49	+ block_max = T.simd.alloc_var(T.float32)
	50	+ lm = T.simd.alloc_var(T.float32)
	51	+ lm = neg
	52	+ lm1 = T.simd.alloc_var(T.float32)
	53	+ lm1 = neg
	54	+ lm2 = T.simd.alloc_var(T.float32)
	55	+ lm2 = neg
	56	+ lm3 = T.simd.alloc_var(T.float32)
	57	+ lm3 = neg
	58	+ if valid >= 256:
	59	+ for vec in T.serial(valid // 256):
	60	+ group_value0 = T.simd.vld(x_ub[(vec * 4 + 0) * 64])
	61	+ lm = T.simd.vmax(lm, group_value0, full)
	62	+ group_value1 = T.simd.vld(x_ub[(vec * 4 + 1) * 64])
	63	+ lm1 = T.simd.vmax(lm1, group_value1, full)
	64	+ group_value2 = T.simd.vld(x_ub[(vec * 4 + 2) * 64])
	65	+ lm2 = T.simd.vmax(lm2, group_value2, full)
	66	+ group_value3 = T.simd.vld(x_ub[(vec * 4 + 3) * 64])
	67	+ lm3 = T.simd.vmax(lm3, group_value3, full)
	68	+ lm = T.simd.vmax(lm, lm1, full)
	69	+ lm = T.simd.vmax(lm, lm2, full)
	70	+ lm = T.simd.vmax(lm, lm3, full)
	71	+ if valid // 64 % 4 > 0:
	72	+ for vec in T.serial(valid // 64 % 4):
	73	+ value = T.simd.vld(x_ub[(valid // 256 * 4 + vec) * 64])
	74	+ lm = T.simd.vmax(lm, value, full)
	75	+ if valid % 64 != 0:
	76	+ tail_value = T.simd.vld(x_ub[valid // 64 * 64])
	77	+ tail = T.simd.update_mask(valid % 64, 32)
	78	+ lm = T.simd.vmax(lm, T.simd.vsel(tail_value, neg, tail), full)
	79	+ block_max = T.simd.vdupv(T.simd.vcmax(lm, full), full)
	80	+ return block_max
	81	+
	82	+ return block_maximum
	83	+
	84	+
	85	+def _block_maximum(dtype):
	86	+ if dtype == 'float32':
	87	+ return _float_block_maximum()
	88	+ return _half_block_maximum(dtype)
	89	+
	90	+
	91	+def _exponential_lanes(narrow, fused_exp, broadcast):
	92	+ if narrow:
	93	+
	94	+ @T.macro
	95	+ def load_lane(x_ub, offset):
	96	+ return T.simd.vcvt(T.simd.vld(x_ub[offset], dist='UNPK_B16'), T.float32, part=0)
	97	+ else:
	98	+
	99	+ @T.macro
	100	+ def load_lane(x_ub, offset):
	101	+ return T.simd.vld(x_ub[offset])
	102	+
	103	+ if fused_exp:
	104	+
	105	+ @T.macro
	106	+ def selected_exp(value, m, full):
	107	+ return T.simd.vexpdif(value, m, full)
	108	+ else:
	109	+
	110	+ @T.macro
	111	+ def selected_exp(value, m, full):
	112	+ return T.simd.vexp(T.simd.vsub(value, m, full), full)
	113	+
	114	+ @T.macro
	115	+ def exponential_lanes(x_ub, valid, m, full, neg):
	116	+ ls = T.simd.alloc_var(T.float32)
	117	+ ls = T.simd.vdup(0.0, T.float32, full)
	118	+ ls1 = T.simd.alloc_var(T.float32)
	119	+ ls1 = T.simd.vdup(0.0, T.float32, full)
	120	+ ls2 = T.simd.alloc_var(T.float32)
	121	+ ls2 = T.simd.vdup(0.0, T.float32, full)
	122	+ ls3 = T.simd.alloc_var(T.float32)
	123	+ ls3 = T.simd.vdup(0.0, T.float32, full)
	124	+ if valid >= 256:
	125	+ for vec in T.serial(valid // 256):
	126	+ if narrow:
	127	+ packed_group0 = T.simd.vld(x_ub[(vec * 4 + 0) * 64])
	128	+ group_value0 = T.simd.vcvt(packed_group0, T.float32, part=0)
	129	+ else:
	130	+ group_value0 = T.simd.vld(x_ub[(vec * 4 + 0) * 64])
	131	+ ls = T.simd.vadd(ls, selected_exp(group_value0, m, full), full)
	132	+ if narrow:
	133	+ group_value1 = T.simd.vcvt(packed_group0, T.float32, part=1)
	134	+ else:
	135	+ group_value1 = T.simd.vld(x_ub[(vec * 4 + 1) * 64])
	136	+ ls1 = T.simd.vadd(ls1, selected_exp(group_value1, m, full), full)
	137	+ if narrow:
	138	+ packed_group1 = T.simd.vld(x_ub[(vec * 4 + 2) * 64])
	139	+ group_value2 = T.simd.vcvt(packed_group1, T.float32, part=0)
	140	+ else:
	141	+ group_value2 = T.simd.vld(x_ub[(vec * 4 + 2) * 64])
	142	+ ls2 = T.simd.vadd(ls2, selected_exp(group_value2, m, full), full)
	143	+ if narrow:
	144	+ group_value3 = T.simd.vcvt(packed_group1, T.float32, part=1)
	145	+ else:
	146	+ group_value3 = T.simd.vld(x_ub[(vec * 4 + 3) * 64])
	147	+ ls3 = T.simd.vadd(ls3, selected_exp(group_value3, m, full), full)
	148	+ ls = T.simd.vadd(ls, ls1, full)
	149	+ ls = T.simd.vadd(ls, ls2, full)
	150	+ ls = T.simd.vadd(ls, ls3, full)
	151	+ if valid // 64 % 4 > 0:
	152	+ for vec in T.serial(valid // 64 % 4):
	153	+ value2 = load_lane(x_ub, (valid // 256 * 4 + vec) * 64)
	154	+ ls = T.simd.vadd(ls, selected_exp(value2, m, full), full)
	155	+ if valid % 64 != 0:
	156	+ tail_value2 = load_lane(x_ub, valid // 64 * 64)
	157	+ tail2 = T.simd.update_mask(valid % 64, 32)
	158	+ tail_value2 = T.simd.vsel(tail_value2, neg, tail2)
	159	+ ls = T.simd.vadd(ls, selected_exp(tail_value2, m, full), full)
	160	+ block_sum = T.simd.vdupv(T.simd.vcadd(ls, full), full) if broadcast else ls
	161	+ return block_sum
	162	+
	163	+ return exponential_lanes
	164	+
	165	+
	166	+def _block_exponential_lanes(dtype):
	167	+ return _exponential_lanes(narrow=dtype != 'float32', fused_exp=dtype != 'float32', broadcast=False)
	168	+
	169	+
	170	+def _block_exponential_sum(dtype):
	171	+ return _exponential_lanes(narrow=dtype != 'float32', fused_exp=dtype != 'float32', broadcast=True)
	172	+
	173	+
	174	+def _plain_exponential_lanes(narrow):
	175	+ return _exponential_lanes(narrow, fused_exp=False, broadcast=False)
	176	+
	177	+
	178	+def _row_block_statistics_update(dtype, single_block):
	179	+ narrow = dtype != 'float32'
	180	+ block_maximum = _block_maximum(dtype)
	181	+ block_exponential_sum = _block_exponential_sum(dtype)
	182	+
	183	+ @T.macro
	184	+ def update_row_statistics(x_ub, valid, max_buffer, sum_buffer, wave, sum_offset):
	185	+ full = T.simd.pset(32)
	186	+ neg = T.simd.vdup(-T.infinity('float32'), T.float32, full)
	187	+ block_max = block_maximum(x_ub, valid, full, neg)
	188	+ if single_block:
	189	+ m = block_max
	190	+ else:
	191	+ run_max = T.simd.vld(max_buffer[wave], dist='BRC_B32')
	192	+ run_sum = T.simd.vld(sum_buffer[sum_offset], dist='BRC_B32')
	193	+ m = T.simd.vmax(block_max, run_max, full)
	194	+ block_sum = block_exponential_sum(x_ub, valid, m, full, neg)
	195	+ if single_block:
	196	+ s = block_sum
	197	+ else:
	198	+ s = T.simd.vadd(
	199	+ T.simd.vmul(
	200	+ run_sum,
	201	+ T.simd.vexpdif(run_max, m, full) if narrow else T.simd.vexp(T.simd.vsub(run_max, m, full), full),
	202	+ full,
	203	+ ),
	204	+ block_sum,
	205	+ full,
	206	+ )
	207	+ T.simd.vsts(max_buffer[wave], m, full, 'ONEPT_B32')
	208	+ T.simd.vsts(sum_buffer[sum_offset], s, full, 'ONEPT_B32')
	209	+
	210	+ return update_row_statistics

Asrc/cann_ops_tilelang/loss/cross_entropy.py+351-0

@@ -0,0 +1,351 @@
	1	+# Copyright (c) 2026 ops-tilelang contributors
	2	+# SPDX-License-Identifier: MIT
	3	+# See LICENSE in the repository root for the full license text.
	4	+
	5	+import tilelang
	6	+import torch
	7	+
	8	+from cann_ops_tilelang.config import get_num_vector_cores
	9	+
	10	+from .cross_entropy_kernel import (
	11	+ FUSED_UB_BUDGET,
	12	+ ROW_BLOCK,
	13	+ SIMD_ROW_MAX_WAVES,
	14	+ SMALL_MAX_LANE_WORK,
	15	+ VOCAB_BYTES,
	16	+ _batched_small_kernel,
	17	+ _chunk_kernel,
	18	+ _counted_row_kernel,
	19	+ _fully_simd_fused_kernel,
	20	+ _joint_reduce_kernel,
	21	+ _merge_kernel,
	22	+ _partial_launch_kernel,
	23	+ _reduce_kernel,
	24	+ _reduce_sum_count_kernel,
	25	+ _row_kernel,
	26	+ _scaled_row_kernel,
	27	+ _simd_row_kernel,
	28	+ _valid_count_kernel,
	29	+ _valid_count_merge_kernel,
	30	+)
	31	+
	32	+
	33	+def _validate(logits, targets, ignore_index, reduction):
	34	+ if not isinstance(logits, torch.Tensor) or not isinstance(targets, torch.Tensor):
	35	+ raise TypeError('logits and targets must be tensors')
	36	+ if logits.ndim != 2 or targets.ndim != 1:
	37	+ raise ValueError('expected logits [N, V] and targets [N]')
	38	+ if logits.shape[0] != targets.shape[0] or logits.shape[1] < 1:
	39	+ raise ValueError('matching N and V >= 1 are required')
	40	+ if logits.dtype not in (torch.float32, torch.bfloat16, torch.float16) or targets.dtype != torch.int64:
	41	+ raise TypeError('expected FP32/BF16/FP16 logits and int64 targets')
	42	+ if logits.device.type != 'npu' or targets.device != logits.device:
	43	+ raise ValueError('inputs must be on the same NPU')
	44	+ if logits.stride(1) != 1 or (logits.shape[0] > 1 and logits.stride(0) < logits.shape[1]):
	45	+ raise ValueError('logits require contiguous columns and nonoverlapping rows')
	46	+ if not targets.is_contiguous():
	47	+ raise ValueError('targets must be contiguous')
	48	+ if type(ignore_index) is not int:
	49	+ raise TypeError('ignore_index must be a Python int, not bool')
	50	+ if not -(263) <= ignore_index < 263:
	51	+ raise ValueError('ignore_index must fit int64')
	52	+ if reduction not in ('none', 'sum', 'mean'):
	53	+ raise ValueError('reduction must be none, sum, or mean')
	54	+
	55	+
	56	+def _active_cores(tasks, available):
	57	+ return min(available, tilelang.next_power_of_2(tasks))
	58	+
	59	+
	60	+@torch.no_grad()
	61	+def cross_entropy_forward(logits, targets, *, ignore_index=-100, reduction='none'):
	62	+ _validate(logits, targets, ignore_index, reduction)
	63	+ n, v = logits.shape
	64	+ if n == 0:
	65	+ return _empty_loss(logits, reduction)
	66	+ dtype = str(logits.dtype).removeprefix('torch.')
	67	+ element_size = logits.element_size()
	68	+ small_width = max(32, tilelang.next_power_of_2(v))
	69	+ block = min(VOCAB_BYTES // element_size, small_width)
	70	+ k = (v + block - 1) // block
	71	+ cores = get_num_vector_cores()
	72	+ waves = (n + cores - 1) // cores
	73	+ route = _whole_row_route(n, v, block, cores, reduction, waves)
	74	+ if route:
	75	+ if reduction == 'sum':
	76	+ return _partial_sum(logits, targets, ignore_index, dtype, block, cores)
	77	+ return _whole_row_mean(logits, targets, ignore_index, dtype, block, cores, route)
	78	+ row_work = waves * v
	79	+ small_rows = min(8, 2048 // small_width)
	80	+ lane_work = small_width // 32 * (1 + int(dtype != 'float32'))
	81	+ uses_small = lane_work <= SMALL_MAX_LANE_WORK and small_rows >= 4 and n >= 2 * cores
	82	+ if k > 1:
	83	+ fused, fused_block, slots = _fused_route(v, k, n, cores, row_work, element_size)
	84	+ if fused:
	85	+ return _launch_fused(logits, targets, ignore_index, reduction, dtype, n, cores, fused_block, slots)
	86	+ inline_count = post_mean = False
	87	+ if reduction == 'mean':
	88	+ inline_count, post_mean, row_cores, row_slots = _mean_plan(
	89	+ logits, dtype, block, k, cores, waves, uses_small, n, v, element_size
	90	+ )
	91	+ if row_slots:
	92	+ return _single_block_scaled_mean(logits, targets, ignore_index, dtype, block, row_cores, row_slots)
	93	+ loss = torch.empty((n,), device=logits.device, dtype=torch.float32)
	94	+ total_count = torch.empty((1,), device=logits.device, dtype=torch.int64)
	95	+ if reduction == 'mean' and not (inline_count or post_mean):
	96	+ _prepare_count(logits, targets, total_count, ignore_index, n, cores)
	97	+ kernel_mean = int(reduction == 'mean' and not post_mean)
	98	+ if uses_small:
	99	+ _batched_small_kernel(dtype, small_width, small_rows, _active_cores((n + small_rows - 1) // small_rows, cores))(
	100	+ logits, targets, loss, total_count, ignore_index, kernel_mean
	101	+ )
	102	+ elif n >= cores or k == 1 or row_work <= (n * k + cores - 1) // cores * block + block:
	103	+ _launch_rows(logits, targets, ignore_index, kernel_mean, dtype, block, cores, inline_count, loss, total_count)
	104	+ else:
	105	+ _launch_chunks(logits, targets, ignore_index, kernel_mean, dtype, block, cores, loss, total_count)
	106	+ if reduction == 'none':
	107	+ return loss
	108	+ if n == 1:
	109	+ return loss.reshape(())
	110	+ return _finish_loss(logits, loss, targets, total_count, ignore_index, reduction, n, cores, post_mean)
	111	+
	112	+
	113	+def _empty_loss(logits, reduction):
	114	+ if reduction == 'none':
	115	+ return torch.empty((0,), device=logits.device, dtype=torch.float32)
	116	+ return torch.full((), float('nan') if reduction == 'mean' else 0.0, device=logits.device, dtype=torch.float32)
	117	+
	118	+
	119	+def _inline_mean(n, v, k, cores, waves):
	120	+ if k != 1 or not 1 < n <= 32:
	121	+ return False
	122	+ return n < 2 * cores and waves <= SIMD_ROW_MAX_WAVES and v <= 2147483647
	123	+
	124	+
	125	+def _post_mean(logits, count_tiles, cores, n, v):
	126	+ if not 8 < count_tiles <= cores:
	127	+ return False
	128	+ return n * (2 * torch.finfo(logits.dtype).max + v + 1) < torch.finfo(torch.float32).max / 4
	129	+
	130	+
	131	+def _count_cores(count_tiles, cores):
	132	+ return _active_cores(count_tiles, cores) if count_tiles > 8 else 1
	133	+
	134	+
	135	+def _mean_plan(logits, dtype, block, k, cores, waves, uses_small, n, v, element_size):
	136	+ inline_count = _inline_mean(n, v, k, cores, waves)
	137	+ if inline_count:
	138	+ return inline_count, False, 0, 0
	139	+ if k != 1 or n > 2**24 or dtype == 'bfloat16':
	140	+ return inline_count, False, 0, 0
	141	+ count_tiles = (n + ROW_BLOCK - 1) // ROW_BLOCK
	142	+ if dtype == 'float16':
	143	+ return inline_count, _post_mean(logits, count_tiles, cores, n, v), 0, 0
	144	+ if uses_small or waves <= SIMD_ROW_MAX_WAVES:
	145	+ return inline_count, False, 0, 0
	146	+ if _count_cores(count_tiles, cores) == 1:
	147	+ return inline_count, False, 0, 0
	148	+ row_cores = _active_cores(n, cores)
	149	+ row_slots = tilelang.next_power_of_2((n + row_cores - 1) // row_cores)
	150	+ row_ub_bytes = 2 * (block + 128) * element_size + 2 * max(8, row_slots) * 4
	151	+ if row_slots > 256 or row_ub_bytes > FUSED_UB_BUDGET:
	152	+ return inline_count, False, 0, 0
	153	+ return inline_count, False, row_cores, row_slots
	154	+
	155	+
	156	+def _launch_fused(logits, targets, ignore_index, reduction, dtype, n, cores, fused_block, slots):
	157	+ workspace = torch.empty((cores, slots * 8), device=logits.device, dtype=torch.float32)
	158	+ output = torch.empty((n,) if reduction == 'none' else (1,), device=logits.device, dtype=torch.float32)
	159	+ _fully_simd_fused_kernel(dtype, fused_block, cores, slots)(
	160	+ logits,
	161	+ targets.view(torch.int32),
	162	+ workspace,
	163	+ output,
	164	+ ignore_index,
	165	+ {'none': 0, 'sum': 1, 'mean': 2}[reduction],
	166	+ )
	167	+ return output if reduction == 'none' else output.reshape(())
	168	+
	169	+
	170	+def _launch_rows(logits, targets, ignore_index, kernel_mean, dtype, block, cores, inline_count, loss, total_count):
	171	+ n, v = logits.shape
	172	+ single_block = v <= block
	173	+ if n <= 2147483647 and v <= 2147483647 and ((n + cores - 1) // cores <= SIMD_ROW_MAX_WAVES):
	174	+ _simd_row_kernel(dtype, block, _active_cores(n, cores), inline_count)(
	175	+ logits,
	176	+ targets.view(torch.int32),
	177	+ loss,
	178	+ total_count.view(torch.int32),
	179	+ ignore_index,
	180	+ kernel_mean,
	181	+ )
	182	+ else:
	183	+ row_cores = _active_cores(n, cores)
	184	+ row_slots = tilelang.next_power_of_2((n + row_cores - 1) // row_cores)
	185	+ row_ub_bytes = 2 * (block + 128) * logits.element_size() + 2 * max(8, row_slots) * 4
	186	+ batch_slots = row_slots if single_block and row_slots <= 256 and (row_ub_bytes <= FUSED_UB_BUDGET) else 0
	187	+ _row_kernel(
	188	+ dtype, block, row_cores, single_block, batch_slots, static_v=v if single_block and batch_slots else None
	189	+ )(logits, targets, loss, total_count, ignore_index, kernel_mean)
	190	+
	191	+
	192	+def _launch_chunks(logits, targets, ignore_index, kernel_mean, dtype, block, cores, loss, total_count):
	193	+ n, v = logits.shape
	194	+ k = (v + block - 1) // block
	195	+ maxima = torch.empty((n, k), device=logits.device, dtype=torch.float32)
	196	+ sums = torch.empty_like(maxima)
	197	+ _chunk_kernel(dtype, block, _active_cores(n * k, cores))(logits, maxima, sums)
	198	+ _merge_kernel(dtype, max(32, tilelang.next_power_of_2(k)), _active_cores(n, cores))(
	199	+ logits, targets, maxima, sums, loss, ignore_index, total_count, kernel_mean
	200	+ )
	201	+
	202	+
	203	+def _whole_row_route(n, v, block, cores, reduction, waves):
	204	+ if reduction == 'none' or n > 2**24:
	205	+ return 0
	206	+ if n <= SIMD_ROW_MAX_WAVES * cores or v <= block:
	207	+ return 0
	208	+ if reduction == 'sum':
	209	+ return 1
	210	+ return 2 if v > 2 * block and waves >= ROW_BLOCK else 1
	211	+
	212	+
	213	+def _whole_row_mean(logits, targets, ignore_index, dtype, block, cores, route):
	214	+ n = logits.shape[0]
	215	+ active = _active_cores(n, cores)
	216	+ if route == 2:
	217	+ values = torch.empty((active,), device=logits.device, dtype=torch.float32)
	218	+ counts = torch.empty((active,), device=logits.device, dtype=torch.int32)
	219	+ output = torch.empty((1,), device=logits.device, dtype=torch.float32)
	220	+ _partial_launch_kernel(dtype, block, active, True)(logits, targets, values, counts, output, ignore_index)
	221	+ return output.reshape(())
	222	+ values = torch.empty((n,), device=logits.device, dtype=torch.float32)
	223	+ counts = torch.empty((n,), device=logits.device, dtype=torch.int32)
	224	+ _counted_row_kernel(dtype, block, active, True)(logits, targets, values, counts, ignore_index)
	225	+ remaining = n
	226	+ while True:
	227	+ capacity = 2 * ROW_BLOCK if ROW_BLOCK < remaining <= 2 * ROW_BLOCK else ROW_BLOCK
	228	+ blocks = (remaining + capacity - 1) // capacity
	229	+ output = torch.empty((blocks,), device=logits.device, dtype=torch.float32)
	230	+ output_count = torch.empty((blocks,), device=logits.device, dtype=torch.int32)
	231	+ _joint_reduce_kernel(_active_cores(blocks, cores), capacity)(
	232	+ values, counts, output, output_count, n, int(blocks == 1), 1
	233	+ )
	234	+ if blocks == 1:
	235	+ return output.reshape(())
	236	+ (values, counts, remaining) = (output, output_count, blocks)
	237	+
	238	+
	239	+def _single_block_scaled_mean(logits, targets, ignore_index, dtype, block, cores, slots):
	240	+ n = logits.shape[0]
	241	+ values = torch.empty((n,), device=logits.device, dtype=torch.float32)
	242	+ counts = torch.empty((n,), device=logits.device, dtype=torch.int32)
	243	+ _scaled_row_kernel(dtype, block, cores, True, slots, static_v=logits.shape[1])(
	244	+ logits, targets, values, counts, ignore_index, 1
	245	+ )
	246	+ remaining = n
	247	+ while True:
	248	+ blocks = (remaining + ROW_BLOCK - 1) // ROW_BLOCK
	249	+ output = torch.empty((blocks,), device=logits.device, dtype=torch.float32)
	250	+ output_count = torch.empty((blocks,), device=logits.device, dtype=torch.int32)
	251	+ _joint_reduce_kernel(_active_cores(blocks, cores), ROW_BLOCK)(
	252	+ values, counts, output, output_count, n, int(blocks == 1), 1
	253	+ )
	254	+ if blocks == 1:
	255	+ return output.reshape(())
	256	+ (values, counts, remaining) = (output, output_count, blocks)
	257	+
	258	+
	259	+def _prepare_count(logits, targets, total_count, ignore_index, n, cores):
	260	+ if n <= 1:
	261	+ return
	262	+ count_tiles = (n + ROW_BLOCK - 1) // ROW_BLOCK
	263	+ count_cores = _count_cores(count_tiles, cores)
	264	+ if count_cores == 1:
	265	+ _valid_count_kernel(n > 2147483647)(targets, total_count, ignore_index)
	266	+ else:
	267	+ partial_count = torch.empty((count_cores,), device=logits.device, dtype=torch.int64)
	268	+ _valid_count_kernel(n > 2147483647, count_cores)(targets, partial_count, ignore_index)
	269	+ _valid_count_merge_kernel()(partial_count, total_count)
	270	+
	271	+
	272	+def _finish_loss(logits, loss, targets, total_count, ignore_index, reduction, n, cores, post_mean):
	273	+ if post_mean:
	274	+ counts = targets
	275	+ first = True
	276	+ while True:
	277	+ capacity = 2 * ROW_BLOCK if ROW_BLOCK < n <= 2 * ROW_BLOCK else ROW_BLOCK
	278	+ blocks = (n + capacity - 1) // capacity
	279	+ output = torch.empty((blocks,), device=logits.device, dtype=torch.float32)
	280	+ output_counts = torch.empty((blocks,), device=logits.device, dtype=torch.int64)
	281	+ _reduce_sum_count_kernel(_active_cores(blocks, cores), capacity, first, blocks == 1)(
	282	+ loss, counts, output, output_counts, ignore_index
	283	+ )
	284	+ if blocks == 1:
	285	+ return output.reshape(())
	286	+ (loss, counts, n, first) = (output, output_counts, blocks, False)
	287	+ while True:
	288	+ capacity = 2 * ROW_BLOCK if ROW_BLOCK < n <= 2 * ROW_BLOCK else ROW_BLOCK
	289	+ blocks = (n + capacity - 1) // capacity
	290	+ output = torch.empty((blocks,), device=logits.device, dtype=torch.float32)
	291	+ _reduce_kernel(_active_cores(blocks, cores), capacity)(
	292	+ loss, output, total_count, int(reduction == 'mean' and blocks == 1)
	293	+ )
	294	+ if blocks == 1:
	295	+ return output.reshape(())
	296	+ (loss, n) = (output, blocks)
	297	+
	298	+
	299	+def _partial_sum(logits, targets, ignore_index, dtype, block, route_cores):
	300	+ n = logits.shape[0]
	301	+ active = _active_cores(n, route_cores)
	302	+ values = torch.empty((active,), device=logits.device, dtype=torch.float32)
	303	+ counts = torch.empty((active,), device=logits.device, dtype=torch.int32)
	304	+ output = torch.empty((1,), device=logits.device, dtype=torch.float32)
	305	+ _partial_launch_kernel(dtype, block, active, False)(logits, targets, values, counts, output, ignore_index)
	306	+ return output.reshape(())
	307	+
	308	+
	309	+def _fused_route(v, k, n, cores, row_work, element_size):
	310	+ fused_block = ((v + k - 1) // k + 127) // 128 * 128
	311	+ fused_k = (v + fused_block - 1) // fused_block
	312	+ slots = (n * fused_k + cores - 1) // cores
	313	+ fused = (
	314	+ cores <= n <= 2 * cores
	315	+ and n % cores != 0
	316	+ and (2 <= fused_k <= 32)
	317	+ and (slots <= 16)
	318	+ and (slots * fused_block * 8 <= row_work * 7)
	319	+ )
	320	+ fused_ub = (
	321	+ 2 * (fused_block + 128) * element_size
	322	+ + (slots * 8 + 4 * cores + 64 + cores * slots * 8 + (2 * cores + 63) // 64 * 64) * 4
	323	+ )
	324	+ fused = fused and fused_ub <= FUSED_UB_BUDGET
	325	+ return fused, fused_block, slots
	326	+
	327	+
	328	+@torch.library.custom_op('cann_ops_tilelang::cross_entropy_forward', mutates_args=())
	329	+def _cross_entropy_forward_custom(
	330	+ logits: torch.Tensor,
	331	+ targets: torch.Tensor,
	332	+ *,
	333	+ ignore_index: int = -100,
	334	+ reduction: str = 'none',
	335	+) -> torch.Tensor:
	336	+ return cross_entropy_forward(logits, targets, ignore_index=ignore_index, reduction=reduction)
	337	+
	338	+
	339	+@_cross_entropy_forward_custom.register_fake
	340	+def _cross_entropy_forward_fake(
	341	+ logits: torch.Tensor,
	342	+ targets: torch.Tensor,
	343	+ *,
	344	+ ignore_index: int = -100,
	345	+ reduction: str = 'none',
	346	+) -> torch.Tensor:
	347	+ _validate(logits, targets, ignore_index=ignore_index, reduction=reduction)
	348	+ return logits.new_empty(
	349	+ (logits.shape[0],) if reduction == 'none' else (),
	350	+ dtype=torch.float32,
	351	+ )

Asrc/cann_ops_tilelang/loss/cross_entropy_kernel.py+27-0

@@ -0,0 +1,27 @@
	1	+# Copyright (c) 2026 ops-tilelang contributors
	2	+# SPDX-License-Identifier: MIT
	3	+# See LICENSE in the repository root for the full license text.
	4	+
	5	+from ._cross_entropy_chunk import _chunk_kernel as _chunk_kernel
	6	+from ._cross_entropy_chunk import _fully_simd_fused_kernel as _fully_simd_fused_kernel
	7	+from ._cross_entropy_common import BALANCED_FUSED_BLOCK as BALANCED_FUSED_BLOCK
	8	+from ._cross_entropy_common import COMPILED_KERNELS as COMPILED_KERNELS
	9	+from ._cross_entropy_common import FUSED_UB_BUDGET as FUSED_UB_BUDGET
	10	+from ._cross_entropy_common import GRID_CORES as GRID_CORES
	11	+from ._cross_entropy_common import ROW_BLOCK as ROW_BLOCK
	12	+from ._cross_entropy_common import SIMD_ROW_MAX_WAVES as SIMD_ROW_MAX_WAVES
	13	+from ._cross_entropy_common import SMALL_MAX_LANE_WORK as SMALL_MAX_LANE_WORK
	14	+from ._cross_entropy_common import VOCAB_BLOCK as VOCAB_BLOCK
	15	+from ._cross_entropy_common import VOCAB_BYTES as VOCAB_BYTES
	16	+from ._cross_entropy_counted import _counted_row_kernel as _counted_row_kernel
	17	+from ._cross_entropy_counted import _partial_launch_kernel as _partial_launch_kernel
	18	+from ._cross_entropy_reduce import _batched_small_kernel as _batched_small_kernel
	19	+from ._cross_entropy_reduce import _joint_reduce_kernel as _joint_reduce_kernel
	20	+from ._cross_entropy_reduce import _merge_kernel as _merge_kernel
	21	+from ._cross_entropy_reduce import _reduce_kernel as _reduce_kernel
	22	+from ._cross_entropy_reduce import _reduce_sum_count_kernel as _reduce_sum_count_kernel
	23	+from ._cross_entropy_reduce import _valid_count_kernel as _valid_count_kernel
	24	+from ._cross_entropy_reduce import _valid_count_merge_kernel as _valid_count_merge_kernel
	25	+from ._cross_entropy_row import _row_kernel as _row_kernel
	26	+from ._cross_entropy_row import _scaled_row_kernel as _scaled_row_kernel
	27	+from ._cross_entropy_row import _simd_row_kernel as _simd_row_kernel

Atests/loss/__init__.py+5-0

@@ -0,0 +1,5 @@
	1	+# Copyright (c) 2026 ops-tilelang contributors
	2	+# SPDX-License-Identifier: MIT
	3	+# See LICENSE in the repository root for the full license text.
	4	+
	5	+__all__ = []

Atests/loss/cases.py+301-0

@@ -0,0 +1,301 @@
	1	+# Copyright (c) 2026 ops-tilelang contributors
	2	+# SPDX-License-Identifier: MIT
	3	+# See LICENSE in the repository root for the full license text.
	4	+
	5	+import os
	6	+from dataclasses import dataclass, replace
	7	+from itertools import product
	8	+
	9	+import torch
	10	+
	11	+VOCAB_BLOCK = 8192
	12	+ROW_BLOCK = 256
	13	+GRID_CORES = 64
	14	+DTYPES = ('float32', 'bfloat16', 'float16')
	15	+REDUCTIONS = ('none', 'sum', 'mean')
	16	+
	17	+
	18	+@dataclass(frozen=True)
	19	+class Case:
	20	+ name: str
	21	+ n: int
	22	+ v: int
	23	+ dtype: str = 'float32'
	24	+ reduction: str = 'none'
	25	+ pattern: str = 'random'
	26	+ ignore: str = 'none'
	27	+ ignore_index: int = -100
	28	+ padding: int = 0
	29	+ offset: int = 0
	30	+ seed: int = 20261026
	31	+
	32	+ @property
	33	+ def id(self):
	34	+ return f'{self.name}-N{self.n}-V{self.v}-{self.dtype}-{self.reduction}'
	35	+
	36	+
	37	+def get_test_level():
	38	+ level = int(os.environ.get('OPS_TILELANG_TEST_LEVEL', '1'))
	39	+ if level not in (0, 1, 2):
	40	+ raise ValueError('OPS_TILELANG_TEST_LEVEL must be 0, 1, or 2')
	41	+ return level
	42	+
	43	+
	44	+def generate_test_params(level=1, *, block_v=VOCAB_BLOCK, block_rows=ROW_BLOCK, grid_cores=GRID_CORES):
	45	+ if level not in (0, 1, 2):
	46	+ raise ValueError('invalid test level or planning tile size')
	47	+ if block_v < 2 or block_rows < 2 or grid_cores < 1:
	48	+ raise ValueError('invalid test level or planning tile size')
	49	+ geometries = _test_geometries(block_v, block_rows, grid_cores)
	50	+ if level == 0:
	51	+ return [
	52	+ geometries[2],
	53	+ replace(geometries[11], dtype='bfloat16', reduction='mean'),
	54	+ replace(geometries[18], dtype='float16', reduction='sum'),
	55	+ ]
	56	+ cases = [
	57	+ replace(case, dtype=dtype, reduction=reduction)
	58	+ for (case, dtype, reduction) in product(geometries, DTYPES, REDUCTIONS)
	59	+ ]
	60	+ cases.extend(Case('fp32-overflow', 3, 2, reduction=r, pattern='overflow') for r in REDUCTIONS)
	61	+ cases.extend(Case('fp32-sum-overflow', 3, 2, reduction=r, pattern='sum_overflow') for r in REDUCTIONS)
	62	+ for dtype, reduction in product(('float32', 'bfloat16'), REDUCTIONS):
	63	+ cases.append(
	64	+ Case('row-overflow-finite-mean', 3, 2, dtype=dtype, reduction=reduction, pattern='row_overflow_mixed')
	65	+ )
	66	+ cases.append(
	67	+ Case(
	68	+ 'sum-overflow-with-ignore',
	69	+ 7,
	70	+ 2,
	71	+ dtype=dtype,
	72	+ reduction=reduction,
	73	+ pattern='sum_overflow',
	74	+ ignore='partial',
	75	+ )
	76	+ )
	77	+ if level == 2:
	78	+ _add_interaction_cases(cases, block_v, block_rows)
	79	+ return cases
	80	+
	81	+
	82	+def make_inputs(case, *, block_v=VOCAB_BLOCK, return_storage=False):
	83	+ rng = torch.Generator().manual_seed(case.seed)
	84	+ x = torch.randn(case.n, case.v, generator=rng)
	85	+ positions = [
	86	+ 0,
	87	+ case.v - 1,
	88	+ case.v // 2,
	89	+ min(block_v - 1, case.v - 1),
	90	+ min(block_v, case.v - 1),
	91	+ min(2 * block_v - 1, case.v - 1),
	92	+ (case.v - 1) // block_v * block_v,
	93	+ ]
	94	+ y = torch.tensor([positions[i % len(positions)] for i in range(case.n)], dtype=torch.int64)
	95	+ _apply_pattern(case, x, y, positions)
	96	+ if case.ignore == 'partial':
	97	+ y[::2] = case.ignore_index
	98	+ elif case.ignore == 'all':
	99	+ y.fill_(case.ignore_index)
	100	+ elif case.ignore == 'one_valid':
	101	+ y.fill_(case.ignore_index)
	102	+ y[-1] = case.v - 1
	103	+ elif case.ignore != 'none':
	104	+ raise ValueError(f'unknown ignore mode: {case.ignore}')
	105	+ dtype = getattr(torch, case.dtype)
	106	+ stride = case.v + case.padding
	107	+ storage = torch.full((case.offset + case.n * stride,), 777, dtype=dtype)
	108	+ logits = storage.as_strided((case.n, case.v), (stride, 1), case.offset)
	109	+ logits.copy_(x.to(dtype))
	110	+ target_storage = torch.full((case.n + 1,), 999999, dtype=torch.int64)
	111	+ targets = target_storage[1:]
	112	+ targets.copy_(y)
	113	+ if return_storage:
	114	+ return logits, targets, storage, target_storage
	115	+ return (logits, targets)
	116	+
	117	+
	118	+def coverage_tags(case, logits, targets, *, block_v=VOCAB_BLOCK, block_rows=ROW_BLOCK, grid_cores=GRID_CORES):
	119	+ valid = targets != case.ignore_index
	120	+ tags = {f'dtype:{case.dtype}', f'reduction:{case.reduction}', f'pattern:{case.pattern}'}
	121	+ tags.add('rows:empty' if case.n == 0 else 'rows:nonempty')
	122	+ tags.add('vocab:one-class' if case.v == 1 else 'vocab:many-classes')
	123	+ tags.add('chunks:single' if case.v <= block_v else 'chunks:multiple')
	124	+ tags.add('vocab:full-block' if case.v % block_v == 0 else 'vocab:tail')
	125	+ tags.add('rows:reduction-tail' if case.n % block_rows else 'rows:reduction-aligned')
	126	+ if case.n > block_rows:
	127	+ tags.add('rows:multi-reduction-block')
	128	+ if case.n > grid_cores:
	129	+ tags.add('rows:persistent-multi-wave')
	130	+ tags.add('layout:row-strided' if case.padding else 'layout:compact')
	131	+ if logits.storage_offset():
	132	+ tags.add('layout:offset')
	133	+ if case.n:
	134	+ _add_target_tags(tags, case, targets, valid, block_v, grid_cores)
	135	+ if case.ignore_index != -100:
	136	+ tags.add('ignore:custom-in-range' if 0 <= case.ignore_index < case.v else 'ignore:custom-out-of-range')
	137	+ return tags
	138	+
	139	+
	140	+REQUIRED_TAGS = {
	141	+ *(f'dtype:{d}' for d in DTYPES),
	142	+ *(f'reduction:{r}' for r in REDUCTIONS),
	143	+ 'rows:empty',
	144	+ 'rows:nonempty',
	145	+ 'vocab:one-class',
	146	+ 'vocab:many-classes',
	147	+ 'chunks:single',
	148	+ 'chunks:multiple',
	149	+ 'vocab:full-block',
	150	+ 'vocab:tail',
	151	+ 'rows:reduction-tail',
	152	+ 'rows:reduction-aligned',
	153	+ 'rows:multi-reduction-block',
	154	+ 'rows:persistent-multi-wave',
	155	+ 'layout:row-strided',
	156	+ 'layout:compact',
	157	+ 'layout:offset',
	158	+ 'tasks:one-wave',
	159	+ 'tasks:multi-wave',
	160	+ 'tasks:full-wave',
	161	+ 'tasks:tail-wave',
	162	+ 'ignore:none',
	163	+ 'ignore:partial',
	164	+ 'ignore:all',
	165	+ 'ignore:one-valid',
	166	+ 'ignore:custom-in-range',
	167	+ 'ignore:custom-out-of-range',
	168	+ 'target:first',
	169	+ 'target:last',
	170	+ 'target:before-boundary',
	171	+ 'target:after-boundary',
	172	+ 'pattern:example',
	173	+ 'pattern:boundary',
	174	+ 'pattern:zeros',
	175	+ 'pattern:negative',
	176	+ 'pattern:large_constant',
	177	+ 'pattern:large_range',
	178	+ 'pattern:dominant',
	179	+ 'pattern:wide_random',
	180	+ 'pattern:overflow',
	181	+ 'pattern:sum_overflow',
	182	+ 'pattern:row_overflow_mixed',
	183	+ 'pattern:random',
	184	+}
	185	+
	186	+
	187	+def _test_geometries(block_v, block_rows, grid_cores):
	188	+ geometries = [
	189	+ Case('empty', 0, 7),
	190	+ Case('one-class', 3, 1, ignore='partial'),
	191	+ Case('cat-dog-bird', 3, 3, pattern='example'),
	192	+ Case('tiny-tail', 7, 31),
	193	+ Case('vector-full', 7, 32),
	194	+ Case('vector-tail', 7, 33),
	195	+ Case('block-minus-one', 7, block_v - 1, pattern='boundary'),
	196	+ Case('block-exact', 7, block_v, pattern='boundary'),
	197	+ Case('block-plus-one', 7, block_v + 1, pattern='boundary'),
	198	+ Case('two-full-blocks', 7, 2 * block_v, pattern='boundary'),
	199	+ Case('two-blocks-plus-one', 7, 2 * block_v + 1, pattern='boundary'),
	200	+ Case('model-vocab', 7, 129280, pattern='boundary', ignore='partial'),
	201	+ Case('model-minus-one', 2, 129279),
	202	+ Case('model-plus-one', 2, 129281),
	203	+ Case('row-minus-one', block_rows - 1, 33),
	204	+ Case('row-exact', block_rows, 33),
	205	+ Case('row-plus-one', block_rows + 1, 33),
	206	+ Case('persistent-tail', 2 * grid_cores + 1, 127),
	207	+ Case('strided-offset', 7, block_v + 1, padding=7, offset=3, ignore='partial'),
	208	+ Case('label-first-last', 7, 17, pattern='boundary'),
	209	+ Case('all-ignore', 7, block_v + 1, ignore='all'),
	210	+ Case('only-one-valid', 7, 33, ignore='one_valid'),
	211	+ Case('custom-ignore-positive', 7, 33, ignore='partial', ignore_index=40),
	212	+ Case('custom-ignore-class', 7, 33, ignore='partial', ignore_index=0),
	213	+ Case('custom-ignore-negative', 7, 33, ignore='partial', ignore_index=-7),
	214	+ Case('zeros', 3, 33, pattern='zeros'),
	215	+ Case('negative-padding', 3, block_v + 1, pattern='negative'),
	216	+ Case('large-common-offset', 3, 33, pattern='large_constant'),
	217	+ Case('large-range', 3, 33, pattern='large_range'),
	218	+ Case('confident-right-wrong', 7, 33, pattern='dominant'),
	219	+ Case('wide-random', 7, block_v + 1, pattern='wide_random'),
	220	+ Case('single-row', 1, 7),
	221	+ Case('single-row-model', 1, 129280),
	222	+ Case('many-row-blocks', 16 * block_rows + 1, 33, ignore='partial'),
	223	+ ]
	224	+ return geometries
	225	+
	226	+
	227	+def _add_interaction_cases(cases, block_v, block_rows):
	228	+ interaction_shapes = [
	229	+ (1, 1),
	230	+ (7, block_v - 1),
	231	+ (7, block_v),
	232	+ (7, block_v + 1),
	233	+ (7, 129280),
	234	+ (block_rows + 1, 33),
	235	+ ]
	236	+ for (n, v), dtype, reduction, ignore, padding in product(
	237	+ interaction_shapes, DTYPES, REDUCTIONS, ('none', 'partial', 'all'), (0, 7)
	238	+ ):
	239	+ cases.append(
	240	+ Case(
	241	+ f'interaction-{ignore}-pad{padding}',
	242	+ n,
	243	+ v,
	244	+ dtype,
	245	+ reduction,
	246	+ ignore=ignore,
	247	+ padding=padding,
	248	+ offset=1,
	249	+ seed=20261027,
	250	+ )
	251	+ )
	252	+
	253	+
	254	+def _apply_pattern(case, x, y, positions):
	255	+ if case.pattern == 'example':
	256	+ x[:] = torch.tensor([2.0, 1.0, 0.0])
	257	+ y[:] = torch.tensor([0, 1, 2])
	258	+ elif case.pattern == 'boundary':
	259	+ x.fill_(-7)
	260	+ for i in range(case.n):
	261	+ x[i, positions[(i + 1) % len(positions)]] = 9
	262	+ elif case.pattern == 'zeros':
	263	+ x.zero_()
	264	+ elif case.pattern == 'negative':
	265	+ x.fill_(-10000)
	266	+ elif case.pattern == 'large_constant':
	267	+ x.fill_(10000 if case.dtype == 'float16' else 100000000.0)
	268	+ elif case.pattern == 'large_range':
	269	+ x.fill_(-10000)
	270	+ x[:, :2] = torch.tensor([10000.0, 9999.0])
	271	+ elif case.pattern == 'dominant':
	272	+ x.fill_(-100)
	273	+ x[:, 0] = 100
	274	+ elif case.pattern == 'wide_random':
	275	+ x *= 20
	276	+ elif case.pattern in ('overflow', 'sum_overflow', 'row_overflow_mixed'):
	277	+ magnitude = 1e38 if case.pattern == 'sum_overflow' else 2.5e38
	278	+ x[:] = torch.tensor([magnitude, -magnitude])
	279	+ y.fill_(1)
	280	+ if case.pattern == 'row_overflow_mixed':
	281	+ y[1:] = 0
	282	+ elif case.pattern != 'random':
	283	+ raise ValueError(f'unknown pattern: {case.pattern}')
	284	+
	285	+
	286	+def _add_target_tags(tags, case, targets, valid, block_v, grid_cores):
	287	+ tasks = case.n * ((case.v + block_v - 1) // block_v)
	288	+ tags.add('tasks:one-wave' if tasks <= grid_cores else 'tasks:multi-wave')
	289	+ tags.add('tasks:full-wave' if tasks % grid_cores == 0 else 'tasks:tail-wave')
	290	+ count = int(valid.sum())
	291	+ tags.add('ignore:all' if count == 0 else 'ignore:none' if count == case.n else 'ignore:partial')
	292	+ if count == 1:
	293	+ tags.add('ignore:one-valid')
	294	+ for label, value in (
	295	+ ('first', 0),
	296	+ ('last', case.v - 1),
	297	+ ('before-boundary', block_v - 1),
	298	+ ('after-boundary', block_v),
	299	+ ):
	300	+ if bool(((targets == value) & valid).any()):
	301	+ tags.add(f'target:{label}')

Atests/loss/checks.py+81-0

@@ -0,0 +1,81 @@
	1	+# Copyright (c) 2026 ops-tilelang contributors
	2	+# SPDX-License-Identifier: MIT
	3	+# See LICENSE in the repository root for the full license text.
	4	+
	5	+import torch
	6	+
	7	+from .cases import make_inputs
	8	+from .reference import cross_entropy_golden
	9	+
	10	+RTOL = 1e-05
	11	+ATOL = 1e-05
	12	+
	13	+
	14	+def assert_forward_matches(
	15	+ actual, logits_cpu, targets_cpu, *, ignore_index=-100, reduction='none', expected_device='cpu'
	16	+):
	17	+ expected = cross_entropy_golden(logits_cpu, targets_cpu, ignore_index=ignore_index, reduction=reduction)
	18	+ assert_expected_matches(
	19	+ actual, expected, targets_cpu, ignore_index=ignore_index, reduction=reduction, expected_device=expected_device
	20	+ )
	21	+
	22	+
	23	+def assert_expected_matches(
	24	+ actual, expected, targets_cpu, *, ignore_index=-100, reduction='none', expected_device='cpu'
	25	+):
	26	+ if not (isinstance(actual, torch.Tensor)):
	27	+ raise AssertionError('output must be a tensor')
	28	+ if not (actual.dtype == torch.float32):
	29	+ raise AssertionError('loss must be FP32 even for BF16/FP16 input')
	30	+ if not (actual.shape == expected.shape):
	31	+ raise AssertionError('loss shape does not match reduction')
	32	+ if not (actual.device == torch.device(expected_device)):
	33	+ raise AssertionError('loss is on the wrong device')
	34	+ if not (actual.is_contiguous()):
	35	+ raise AssertionError('loss must be contiguous')
	36	+ actual = actual.detach().cpu()
	37	+ valid_count = int((targets_cpu != ignore_index).sum())
	38	+ allowed_nan = reduction == 'mean' and valid_count == 0
	39	+ if not (bool(torch.isnan(expected).any()) == allowed_nan):
	40	+ raise AssertionError('precision case invariant failed')
	41	+ if not (torch.equal(torch.isnan(actual), torch.isnan(expected))):
	42	+ raise AssertionError('NaN mask differs')
	43	+ if not (torch.equal(torch.isposinf(actual), torch.isposinf(expected))):
	44	+ raise AssertionError('+Inf mask differs')
	45	+ if not (torch.equal(torch.isneginf(actual), torch.isneginf(expected))):
	46	+ raise AssertionError('-Inf mask differs')
	47	+ finite = torch.isfinite(expected)
	48	+ atol = ATOL * max(valid_count, 1) if reduction == 'sum' else ATOL
	49	+ torch.testing.assert_close(actual[finite], expected[finite], rtol=RTOL, atol=atol)
	50	+ if reduction == 'none':
	51	+ ignored = targets_cpu == ignore_index
	52	+ if not (torch.equal(actual[ignored], torch.zeros_like(actual[ignored]))):
	53	+ raise AssertionError('ignored loss must be exactly zero')
	54	+ elif valid_count == 0 and reduction == 'sum':
	55	+ if not (actual.item() == 0.0):
	56	+ raise AssertionError('empty/all-ignored sum must be exactly zero')
	57	+
	58	+
	59	+def run_precision_case(implementation, case, *, device='cpu'):
	60	+ (x_cpu, y_cpu, x_base, y_base) = make_inputs(case, return_storage=True)
	61	+ x_storage = x_base.to(device, copy=True)
	62	+ y_storage = y_base.to(device, copy=True)
	63	+ x = x_storage.as_strided(x_cpu.shape, x_cpu.stride(), x_cpu.storage_offset())
	64	+ y = y_storage.as_strided(y_cpu.shape, y_cpu.stride(), y_cpu.storage_offset())
	65	+ if not (x.stride() == x_cpu.stride() and x.storage_offset() == x_cpu.storage_offset()):
	66	+ raise AssertionError('precision case invariant failed')
	67	+ if not (y.stride() == y_cpu.stride() and y.storage_offset() == y_cpu.storage_offset()):
	68	+ raise AssertionError('precision case invariant failed')
	69	+ result = implementation(x, y, ignore_index=case.ignore_index, reduction=case.reduction)
	70	+ assert_forward_matches(
	71	+ result, x_cpu, y_cpu, ignore_index=case.ignore_index, reduction=case.reduction, expected_device=x.device
	72	+ )
	73	+ if not (torch.equal(x_storage.cpu(), x_base)):
	74	+ raise AssertionError('logits or padding modified')
	75	+ if not (torch.equal(y_storage.cpu(), y_base)):
	76	+ raise AssertionError('targets or guard storage modified')
	77	+ if result.numel():
	78	+ pointer = result.untyped_storage().data_ptr()
	79	+ if not (pointer not in (x_storage.untyped_storage().data_ptr(), y_storage.untyped_storage().data_ptr())):
	80	+ raise AssertionError('output aliases an input')
	81	+ return result

Atests/loss/reference.py+61-0

@@ -0,0 +1,61 @@
	1	+# Copyright (c) 2026 ops-tilelang contributors
	2	+# SPDX-License-Identifier: MIT
	3	+# See LICENSE in the repository root for the full license text.
	4	+
	5	+import torch
	6	+
	7	+INPUT_DTYPES = (torch.float32, torch.bfloat16, torch.float16)
	8	+REDUCTIONS = ('none', 'sum', 'mean')
	9	+
	10	+
	11	+def validate_reference_inputs(logits, targets, *, ignore_index=-100, reduction='none'):
	12	+ if not isinstance(logits, torch.Tensor) or not isinstance(targets, torch.Tensor):
	13	+ raise TypeError('logits and targets must be tensors')
	14	+ if logits.ndim != 2 or targets.ndim != 1:
	15	+ raise ValueError('expected logits [N, V] and targets [N]')
	16	+ if logits.shape[0] != targets.shape[0] or logits.shape[1] < 1:
	17	+ raise ValueError('matching N and V >= 1 are required')
	18	+ if logits.dtype not in INPUT_DTYPES or targets.dtype != torch.int64:
	19	+ raise TypeError('expected FP32/BF16/FP16 logits and int64 targets')
	20	+ if logits.device.type != 'cpu' or targets.device.type != 'cpu':
	21	+ raise ValueError('references require CPU inputs')
	22	+ if logits.stride(1) != 1 or (logits.shape[0] > 1 and logits.stride(0) < logits.shape[1]):
	23	+ raise ValueError('logits require contiguous columns and nonoverlapping rows')
	24	+ if not targets.is_contiguous():
	25	+ raise ValueError('targets must be contiguous')
	26	+ if type(ignore_index) is not int:
	27	+ raise TypeError('ignore_index must be a Python int, not bool')
	28	+ if not -(263) <= ignore_index < 263:
	29	+ raise ValueError('ignore_index must fit int64')
	30	+ if reduction not in REDUCTIONS:
	31	+ raise ValueError('reduction must be none, sum, or mean')
	32	+ if not bool(torch.isfinite(logits).all()):
	33	+ raise ValueError('input logits must be finite')
	34	+ invalid = (targets != ignore_index) & ((targets < 0) \| (targets >= logits.shape[1]))
	35	+ if bool(invalid.any()):
	36	+ raise ValueError('target must be in [0, V) or equal ignore_index')
	37	+
	38	+
	39	+@torch.no_grad()
	40	+def row_losses_fp64(logits, targets, *, ignore_index=-100):
	41	+ validate_reference_inputs(logits, targets, ignore_index=ignore_index)
	42	+ valid = targets != ignore_index
	43	+ x = logits.double()[valid]
	44	+ shifted = x - x.max(dim=1, keepdim=True).values
	45	+ selected = shifted.gather(1, targets[valid, None]).squeeze(1)
	46	+ rows = torch.zeros(logits.shape[0], dtype=torch.float64)
	47	+ rows[valid] = shifted.exp().sum(dim=1).log() - selected
	48	+ return rows
	49	+
	50	+
	51	+@torch.no_grad()
	52	+def cross_entropy_golden(logits, targets, *, ignore_index=-100, reduction='none'):
	53	+ if reduction not in REDUCTIONS:
	54	+ raise ValueError('reduction must be none, sum, or mean')
	55	+ rows = row_losses_fp64(logits, targets, ignore_index=ignore_index)
	56	+ valid = targets != ignore_index
	57	+ if reduction == 'none':
	58	+ return rows.float()
	59	+ if reduction == 'sum':
	60	+ return rows.sum().float()
	61	+ return (rows.sum() / valid.sum()).float()

Atests/loss/test_cross_entropy.py+75-0

@@ -0,0 +1,75 @@
	1	+# Copyright (c) 2026 ops-tilelang contributors
	2	+# SPDX-License-Identifier: MIT
	3	+# See LICENSE in the repository root for the full license text.
	4	+
	5	+import pytest
	6	+import torch
	7	+
	8	+from cann_ops_tilelang.loss import cross_entropy_forward
	9	+
	10	+from .cases import generate_test_params, get_test_level
	11	+from .checks import run_precision_case
	12	+from .reference import cross_entropy_golden
	13	+
	14	+
	15	+@pytest.mark.npu
	16	+@pytest.mark.parametrize('case', generate_test_params(get_test_level()), ids=lambda case: case.id)
	17	+def test_cross_entropy_precision(case):
	18	+ run_precision_case(cross_entropy_forward, case, device='npu:0')
	19	+
	20	+
	21	+@pytest.mark.parametrize('reduction', ['none', 'sum', 'mean'])
	22	+def test_cross_entropy_reference(reduction):
	23	+ x = torch.tensor([[2.0, 1.0, 0.0], [1.0, 3.0, 2.0]])
	24	+ y = torch.tensor([0, 1])
	25	+ actual = cross_entropy_golden(x, y, reduction=reduction)
	26	+ expected = torch.nn.functional.cross_entropy(x, y, reduction=reduction)
	27	+ torch.testing.assert_close(actual, expected)
	28	+
	29	+
	30	+@pytest.mark.npu
	31	+@pytest.mark.parametrize(
	32	+ 'kind', ['dtype_x', 'dtype_y', 'rank_x', 'rank_y', 'length', 'stride', 'reduction', 'ignore_bool']
	33	+)
	34	+def test_cross_entropy_validation(kind):
	35	+ x = torch.zeros((3, 7), device='npu:0')
	36	+ y = torch.zeros((3,), device='npu:0', dtype=torch.int64)
	37	+ kwargs = {}
	38	+ if kind == 'dtype_x':
	39	+ x = x.to(torch.int32)
	40	+ if kind == 'dtype_y':
	41	+ y = y.to(torch.int32)
	42	+ if kind == 'rank_x':
	43	+ x = x.flatten()
	44	+ if kind == 'rank_y':
	45	+ y = y[:, None]
	46	+ if kind == 'length':
	47	+ y = y[:2]
	48	+ if kind == 'stride':
	49	+ x = x[:, ::2]
	50	+ if kind == 'reduction':
	51	+ kwargs['reduction'] = 'invalid'
	52	+ if kind == 'ignore_bool':
	53	+ kwargs['ignore_index'] = True
	54	+ with pytest.raises((TypeError, ValueError)):
	55	+ cross_entropy_forward(x, y, **kwargs)
	56	+
	57	+
	58	+@pytest.mark.npu
	59	+@pytest.mark.benchmark
	60	+@pytest.mark.parametrize('dtype', ['float32', 'float16', 'bfloat16'])
	61	+@pytest.mark.parametrize('reduction', ['none', 'sum', 'mean'])
	62	+@pytest.mark.parametrize('n', [64, 1024, 2048, 4096, 8192, 16384, 32768, 65536])
	63	+def test_cross_entropy_benchmark(benchmark_timer, benchmark_record, n, dtype, reduction):
	64	+ x = torch.randn((n, 129280), device='npu:0', dtype=getattr(torch, dtype))
	65	+ y = torch.randint(129280, (n,), device='npu:0')
	66	+ result = cross_entropy_forward(x, y, reduction=reduction)
	67	+ time_us = benchmark_timer(lambda: cross_entropy_forward(x, y, reduction=reduction))
	68	+ count = x.numel() * x.element_size() + y.numel() * y.element_size() + result.numel() * result.element_size()
	69	+ benchmark_record(
	70	+ kernel='cross_entropy_forward',
	71	+ operation='fwd',
	72	+ params=dict(n=n, v=129280, dtype=dtype, reduction=reduction),
	73	+ time_us=time_us,
	74	+ bandwidth_gbs=count / time_us / 1000.0,
	75	+ )

Atests/loss/test_cross_entropy_compile.py+204-0

@@ -0,0 +1,204 @@
	1	+# Copyright (c) 2026 ops-tilelang contributors
	2	+# SPDX-License-Identifier: MIT
	3	+# See LICENSE in the repository root for the full license text.
	4	+
	5	+from importlib import import_module
	6	+
	7	+import pytest
	8	+import torch
	9	+from torch._subclasses.fake_tensor import FakeTensorMode
	10	+
	11	+from cann_ops_tilelang import cross_entropy_forward
	12	+
	13	+from .reference import cross_entropy_golden
	14	+
	15	+pytestmark = pytest.mark.npu
	16	+
	17	+
	18	+@pytest.fixture
	19	+def isolated_graph_rng():
	20	+ generator = torch.npu.default_generators[torch.npu.current_device()]
	21	+ original = generator.graphsafe_get_state()
	22	+ generator.graphsafe_set_state(generator.clone_state())
	23	+ try:
	24	+ yield
	25	+ finally:
	26	+ generator.graphsafe_set_state(original)
	27	+
	28	+
	29	+@pytest.fixture(autouse=True)
	30	+def reset_compile_cache():
	31	+ torch.compiler.reset()
	32	+ yield
	33	+ torch.compiler.reset()
	34	+
	35	+
	36	+def _inputs(dtype, n=7):
	37	+ logits = torch.randn((n, 41), device='npu', dtype=dtype)[:, :33]
	38	+ targets = torch.arange(n, device=logits.device, dtype=torch.int64) % 33
	39	+ targets[::3] = -100
	40	+ return logits, targets
	41	+
	42	+
	43	+def _check_forward(forward, logits, targets, reduction):
	44	+ caller_mode = torch.is_grad_enabled()
	45	+ before_logits = logits.detach().cpu().clone()
	46	+ before_targets = targets.cpu().clone()
	47	+ expected = cross_entropy_golden(before_logits, before_targets, ignore_index=-100, reduction=reduction)
	48	+ with torch.no_grad():
	49	+ output = forward(logits, targets, ignore_index=-100, reduction=reduction)
	50	+ assert not torch.is_grad_enabled()
	51	+ assert output.shape == expected.shape
	52	+ assert output.dtype == torch.float32
	53	+ assert output.device == logits.device
	54	+ assert not output.requires_grad
	55	+ assert output.grad_fn is None
	56	+ assert output.untyped_storage() is not logits.untyped_storage()
	57	+ assert output.untyped_storage() is not targets.untyped_storage()
	58	+ torch.testing.assert_close(output.cpu(), expected, rtol=1e-4, atol=1e-5, equal_nan=True)
	59	+ torch.testing.assert_close(logits.detach().cpu(), before_logits, rtol=0, atol=0)
	60	+ torch.testing.assert_close(targets.cpu(), before_targets, rtol=0, atol=0)
	61	+ assert torch.is_grad_enabled() == caller_mode
	62	+
	63	+
	64	+@pytest.mark.usefixtures('isolated_graph_rng')
	65	+@pytest.mark.parametrize('dtype', (torch.float32, torch.bfloat16, torch.float16))
	66	+@pytest.mark.parametrize('reduction', ('none', 'sum', 'mean'))
	67	+def test_compile_replay(dtype, reduction):
	68	+ logits, targets = _inputs(dtype)
	69	+ logits.requires_grad_()
	70	+ direct = torch.ops.cann_ops_tilelang.cross_entropy_forward.default
	71	+ compiled = torch.compile(direct, backend='npugraph_ex', fullgraph=True, dynamic=False)
	72	+ original_logits = logits
	73	+ original_targets = targets
	74	+ for iteration in range(3):
	75	+ if iteration == 1:
	76	+ with torch.no_grad():
	77	+ logits.mul_(0.5).add_(torch.randn_like(logits))
	78	+ targets.copy_(
	79	+ torch.tensor(
	80	+ [32, -100, 5, 4, -100, 9, 1],
	81	+ device=targets.device,
	82	+ dtype=torch.int64,
	83	+ )
	84	+ )
	85	+ assert logits is original_logits
	86	+ assert targets is original_targets
	87	+ elif iteration == 2:
	88	+ logits, targets = _inputs(dtype)
	89	+ logits.requires_grad_()
	90	+ assert logits is not original_logits
	91	+ assert targets is not original_targets
	92	+ assert logits.untyped_storage() is not original_logits.untyped_storage()
	93	+ assert targets.untyped_storage() is not original_targets.untyped_storage()
	94	+ assert logits.shape == (7, 33)
	95	+ assert logits.stride() == (41, 1)
	96	+ assert targets.shape == (7,)
	97	+ assert targets.stride() == (1,)
	98	+ for forward in (cross_entropy_forward, direct, compiled):
	99	+ _check_forward(forward, logits, targets, reduction)
	100	+
	101	+
	102	+@pytest.mark.usefixtures('isolated_graph_rng')
	103	+@pytest.mark.parametrize('reduction', ('none', 'sum', 'mean'))
	104	+@pytest.mark.parametrize('n', (0, 7), ids=('empty', 'all_ignored'))
	105	+def test_compile_empty_and_all_ignored(n, reduction):
	106	+ logits, targets = _inputs(torch.float32, n)
	107	+ targets.fill_(-100)
	108	+ direct = torch.ops.cann_ops_tilelang.cross_entropy_forward.default
	109	+ compiled = torch.compile(
	110	+ direct,
	111	+ backend='eager' if n == 0 else 'npugraph_ex',
	112	+ fullgraph=True,
	113	+ dynamic=False,
	114	+ )
	115	+ for forward in (cross_entropy_forward, direct, compiled):
	116	+ _check_forward(forward, logits, targets, reduction)
	117	+
	118	+
	119	+@pytest.mark.parametrize('dtype', (torch.float32, torch.bfloat16, torch.float16))
	120	+@pytest.mark.parametrize('reduction', ('none', 'sum', 'mean'))
	121	+@pytest.mark.parametrize('n', (0, 7))
	122	+def test_custom_op_fake_metadata(monkeypatch, dtype, reduction, n):
	123	+ host = import_module(cross_entropy_forward.__module__)
	124	+
	125	+ def fail(args, *kwargs):
	126	+ raise AssertionError('fake execution must not enter the host implementation, compiler, or device query')
	127	+
	128	+ monkeypatch.setattr(host, 'cross_entropy_forward', fail)
	129	+ monkeypatch.setattr(host, 'get_num_vector_cores', fail)
	130	+ monkeypatch.setattr(host.tilelang, 'compile', fail)
	131	+ caller_mode = torch.is_grad_enabled()
	132	+ with FakeTensorMode(), torch.no_grad():
	133	+ logits = torch.empty_strided((n, 33), (41, 1), device='npu', dtype=dtype)
	134	+ targets = torch.empty((n,), device=logits.device, dtype=torch.int64)
	135	+ output = torch.ops.cann_ops_tilelang.cross_entropy_forward.default(
	136	+ logits, targets, ignore_index=-100, reduction=reduction
	137	+ )
	138	+ second = torch.ops.cann_ops_tilelang.cross_entropy_forward.default(
	139	+ logits, targets, ignore_index=-100, reduction=reduction
	140	+ )
	141	+ assert not torch.is_grad_enabled()
	142	+ for result in (output, second):
	143	+ assert result.shape == ((n,) if reduction == 'none' else ())
	144	+ assert result.dtype == torch.float32
	145	+ assert result.device == logits.device
	146	+ assert not result.requires_grad
	147	+ assert result.grad_fn is None
	148	+ assert result is not logits
	149	+ assert result is not targets
	150	+ assert result.untyped_storage() is not logits.untyped_storage()
	151	+ assert result.untyped_storage() is not targets.untyped_storage()
	152	+ assert output is not second
	153	+ assert output.untyped_storage() is not second.untyped_storage()
	154	+ invalid_targets = torch.empty((n,), device=logits.device, dtype=torch.int32)
	155	+ with pytest.raises(TypeError, match='int64 targets'):
	156	+ torch.ops.cann_ops_tilelang.cross_entropy_forward.default(logits, invalid_targets, reduction=reduction)
	157	+ assert torch.is_grad_enabled() == caller_mode
	158	+
	159	+
	160	+@pytest.mark.parametrize('reduction', ('none', 'sum', 'mean'))
	161	+def test_custom_op_opcheck(reduction):
	162	+ logits, targets = _inputs(torch.float32)
	163	+ caller_mode = torch.is_grad_enabled()
	164	+ with torch.no_grad():
	165	+ torch.library.opcheck(
	166	+ torch.ops.cann_ops_tilelang.cross_entropy_forward.default,
	167	+ (logits, targets),
	168	+ kwargs={'ignore_index': -100, 'reduction': reduction},
	169	+ test_utils=('test_schema', 'test_faketensor'),
	170	+ )
	171	+ assert not torch.is_grad_enabled()
	172	+ assert torch.is_grad_enabled() == caller_mode
	173	+
	174	+
	175	+def test_public_eager_rejects_bool_ignore_index():
	176	+ logits, targets = _inputs(torch.float32)
	177	+ _check_forward(cross_entropy_forward, logits, targets, 'none')
	178	+ caller_mode = torch.is_grad_enabled()
	179	+ with torch.no_grad():
	180	+ before_logits = logits.cpu().clone()
	181	+ before_targets = targets.cpu().clone()
	182	+ with pytest.raises(TypeError, match='Python int, not bool'):
	183	+ cross_entropy_forward(logits, targets, ignore_index=True)
	184	+ assert not torch.is_grad_enabled()
	185	+ torch.testing.assert_close(logits.cpu(), before_logits, rtol=0, atol=0)
	186	+ torch.testing.assert_close(targets.cpu(), before_targets, rtol=0, atol=0)
	187	+ assert torch.is_grad_enabled() == caller_mode
	188	+
	189	+
	190	+@pytest.mark.usefixtures('isolated_graph_rng')
	191	+@pytest.mark.parametrize('grad_enabled', (False, True))
	192	+@pytest.mark.parametrize('requires_grad', (False, True))
	193	+def test_compile_preserves_grad_mode(grad_enabled, requires_grad):
	194	+ logits, targets = _inputs(torch.float32)
	195	+ logits.requires_grad_(requires_grad)
	196	+ direct = torch.ops.cann_ops_tilelang.cross_entropy_forward.default
	197	+ compiled = torch.compile(direct, backend='npugraph_ex', fullgraph=True, dynamic=False)
	198	+ caller_mode = torch.is_grad_enabled()
	199	+ with torch.set_grad_enabled(grad_enabled):
	200	+ for forward in (cross_entropy_forward, direct, compiled, compiled):
	201	+ _check_forward(forward, logits, targets, 'mean')
	202	+ assert torch.is_grad_enabled() == grad_enabled
	203	+ assert logits.requires_grad == requires_grad
	204	+ assert torch.is_grad_enabled() == caller_mode

Atests/loss/test_cross_entropy_routes.py+29-0

@@ -0,0 +1,29 @@
	1	+# Copyright (c) 2026 ops-tilelang contributors
	2	+# SPDX-License-Identifier: MIT
	3	+# See LICENSE in the repository root for the full license text.
	4	+
	5	+from itertools import product
	6	+
	7	+import pytest
	8	+
	9	+from cann_ops_tilelang.config import get_num_vector_cores
	10	+from cann_ops_tilelang.loss import cross_entropy_forward
	11	+
	12	+from .cases import Case
	13	+from .checks import run_precision_case
	14	+
	15	+pytestmark = pytest.mark.npu
	16	+
	17	+
	18	+@pytest.mark.parametrize('dtype,reduction', list(product(['float32', 'float16', 'bfloat16'], ['none', 'sum', 'mean'])))
	19	+@pytest.mark.parametrize('route', ['remainder', 'whole_row', 'single_block'])
	20	+def test_cross_entropy_routes(dtype, reduction, route):
	21	+ cores = get_num_vector_cores()
	22	+ if route == 'remainder':
	23	+ n, v = cores + max(1, cores // 7), 129280
	24	+ elif route == 'whole_row':
	25	+ n, v = 2 * cores + 1, 129280
	26	+ else:
	27	+ n, v = 16 * 256 + 1, 8193
	28	+ case = Case(route, n, v, dtype=dtype, reduction=reduction, ignore='partial')
	29	+ run_precision_case(cross_entropy_forward, case, device='npu:0')