fix(ptodsl): preserve signless arithmetic operand types

[liuzidi]

Fix runtime integer arithmetic lowering after 5eb87c2. Integer literals are emitted signless, and mixed signed/signless operands are reconciled in signless arithmetic types. Explicit signed/unsigned scalar and VMI storage semantics remain unchanged. Regression: PTODSL jit compile suite and topk_gate E=256/384/512/72 K=8/9/8/6 num_sms=2.

[liuzidi]

Background

This PR addresses the first PTOAS-related compilation failure observed while compiling TileLang's VMI topk_gate kernel with num_sms=2. The failure is not a kernel algorithm error and is not caused by the explicit VMI signed/unsigned types used by the kernel. It is a frontend type-adaptation bug that is exposed by runtime integer expressions such as:

(num_tokens + 3) // 4

Here, num_tokens is a runtime si32 argument and 3/4 are Python integer literals generated during tracing.

Observed failure

Before this change, PTOAS could emit malformed builtin.unrealized_conversion_cast operations after adapting integer operands for arith operations. A representative form was:

%c3_i32 = arith.constant 3 : i32
%5 = builtin.unrealized_conversion_cast %c3_i32 : si32 to i32

The SSA value %c3_i32 is defined as signless i32, but the cast declares its source operand as signed si32. In other traces the same issue appeared through a reused runtime value, for example:

%1 = builtin.unrealized_conversion_cast %arg2 : si32 to i32
%4 = builtin.unrealized_conversion_cast %1 : si32 to i32

The second cast again declares an si32 source even though %1 is already i32. PTOAS therefore rejects the module while parsing/verifying MLIR:

use of value '%c3_i32' expects different type than prior uses:
'si32' vs 'i32'
Error: Failed to parse MLIR.

Root cause

MLIR arithmetic operations use signless integer types (iN) for their operands and results. PTOAS's previous adaptation path materialized runtime literals as signed siN, then stripped signedness while preparing operands for arith. Under tracing/type reuse, the adapted signless value could be reused while the conversion builder still described the original signed type. This made the operation's declared source type disagree with the SSA value's actual type.

The important issue is therefore not that si32 is intrinsically wrong, nor that all integer types should be globally replaced. The issue is that a value's declared MLIR type must remain consistent across every use, especially when a traced value is reused.

Fix

The fix is intentionally limited to the arithmetic adaptation boundary in:

• ptodsl/ptodsl/_scalar_adaptation.py

Specifically:

1. Runtime integer literals used in arithmetic are materialized directly as signless iN constants, which is the type expected by arith.
2. When mixed signed/unsigned and signless integer operands are reconciled for an arithmetic operation, the common arithmetic operand type is signless.
3. Values that already have the required integer type no longer go through redundant signedness strip/restore conversions.
4. Result signedness restoration remains in place where the surrounding PTOAS/VMI contract requires it.

This keeps the change at the exact boundary where arith requires signless operands. It does not change explicit pto.siN/pto.uiN construction, VMI type semantics, memory ABI behavior, or the signedness of values stored by the generated kernel. An earlier experiment that removed result signedness restoration was reverted because it broke signed stores; the final patch does not include that behavior change.

Regression coverage

The new regression test is:

• ptodsl/tests/test_issue_1405_signed_runtime_arithmetic.py

It covers a signed runtime argument, literal arithmetic, floor division, index arithmetic, and a signed store. The test also checks that no malformed literal conversion of the following form is generated:

%c\\w+ = builtin\\.unrealized_conversion_cast .*: si32 to i32

Validation performed with PTOAS rebuilt from this PR branch:

• Full third_party/ptoas/ptodsl/tests/test_jit_compile.py
• test_issue_1405_signed_runtime_arithmetic.py
• test_scalar_cast.py
• test_ptoas_frontend_verify.py
• Explicit mixed signed/unsigned probes, including preservation of arith.maxsi, arith.minui, and signed stores
• Width probes for si8/ui8, si16/ui16, si32/ui32, and si64/ui64
• Fresh TileLang + PTOAS compilation of representative topk_gate configurations (num_sms=2), including:
  • E=256, K=8
  • E=384, K=9
  • E=512, K=8
  • E=72, K=6
• The complete relevant num_groups=1 set from data/vmi/moe.jsonl, including token-count cases such as 4, 512, and 4001 where applicable

All of the above compilation and frontend regression checks pass with the patched PTOAS.

Scope relative to the other failure class

The separate second failure class concerns a non-32-byte GM-to-UB MTE destination row stride in packed UE8M0 cast_back. That is a physical-access legalization issue and is addressed by the PTOAS main change at 5eb87c21ab9479d834e66968f63f0b1def292764; it is intentionally not solved by adding padding or changing the kernel in this PR.

Hardware launch/numerical validation was not possible in this environment because NPU runtime initialization fails (507899/107002). The successful checks reported here are frontend parsing, verification, JIT compilation, and fresh TileLang-to-PTOAS compilation.

[liuzidi]

这个 PR 解决的是一个 PTOAS 编译器在处理整数运算时的类型对不上问题。

通俗地说，TileLang 的 topk_gate kernel 里有类似这样的计算：

(num_tokens + 3) // 4

其中：

• num_tokens 是运行时传进来的整数；
• 3 和 4 是编译时生成的整数常量。

旧版本 PTOAS 在处理这类表达式时，会把运行时参数当成一种“带符号整数”（例如 si32），但把常量或中间结果处理成另一种“普通整数”（例如 i32）。之后它又尝试在两者之间插入转换，但转换声明的输入类型和实际值类型不一致，最终生成了非法的 MLIR：

%c3_i32 = arith.constant 3 : i32
%x = builtin.unrealized_conversion_cast %c3_i32 : si32 to i32

这里 %c3_i32 实际是 i32，但转换却说它是 si32。PTOAS 在解析和校验 MLIR 时就会报错，导致 topk_gate 无法完成 lowering。

这个 PR 的解决方式

修复集中在 PTOAS 的标量类型适配逻辑：

1. 参与 arith 运算的整数常量直接生成 signless iN 类型，和 MLIR arithmetic 的要求保持一致。
2. 如果运算中同时出现 signed、unsigned 和 signless 整数，就先统一成一个合法的 signless arithmetic 类型。
3. 如果一个值本来已经是目标类型，就不再重复插入 signedness 转换，避免类型被错误复用。
4. 显式的 pto.si32、pto.ui32 类型和最终结果的 signedness 仍然保留，不改变 kernel 原本的数值语义。

所以，这个 PR 不是把所有整数都粗暴替换成 i32，也不是修改 TileLang kernel，而是修正 PTOAS 在“进入 MLIR arithmetic 运算之前”对整数类型的适配方式，确保每个 SSA 值在所有使用位置上的类型声明始终一致。

修复后，原先 topk_gate 中因为 si32/i32 不一致导致的编译错误可以正常通过。

[Zhendong404]

Overall the direction is right: keeping arith operands signless from materialization through the op removes the "si32 identity + i32 value" mismatch that broke topk_gate. A few issues inline — the main one is a silent unsigned semantics regression introduced by the signless-reconcile rule.

[Zhendong404]

This rule silently changes unsigned semantics. Since _materialize_runtime_literal now anchors literals as signless, any uiN runtime value combined with a literal lands here and the result type becomes signless iN. Downstream op selection reads signedness from the reconciled type (_runtime_scalar_ops.py: emit_runtime_binary_op/emit_runtime_compare/emit_runtime_min/emit_runtime_max), where signless defaults to signed — so e.g. (n_ui32 + 3) // 4 now lowers to floordivsi and n_ui32 + 3 < k to slt instead of divui/ult. For values >= 2^31 this is a silent numeric wrong-result, not a compile error. Suggest preserving the non-signless side's authored signedness on the result (compute in the signless domain, restore the uiN/siN identity afterwards), or at minimum capture the operand signedness before reconciliation for op selection. Either way, please add a ui32 regression test.

[liuzidi]

Fixed in 0331b1e. Reconciliation now returns the authored integer type separately from the signless arithmetic operands. Runtime op selection uses that authored type, so ui32 with a literal lowers to divui/remui and unsigned comparison predicates, while MLIR arithmetic operands remain signless. Results are restored to the authored uiN/siN type. A dedicated ui32 regression now covers add, floor division, remainder, comparison, min, and max.

[Zhendong404]

This comment says "restore the authored signedness only on the result", but the implementation does not actually restore the authored signedness: emit_runtime_binary_op restores to lhs.type, which is the reconciled signless type at this point, not the authored si32/ui32. Either restore the authored type (preferred — see my other comment on the unsigned case) or reword the comment so it does not mislead future maintenance.

[liuzidi]

Fixed in 0331b1e. The result is now restored using the authored type captured before signless reconciliation, rather than the reconciled lhs.type. The comment is accurate with the updated implementation.

[Zhendong404]

int(value) narrows the accepted literal forms: the previous path went through _materialize_integer_literal → _parse_integer_value, which also accepts integer strings (e.g. "0x10"). If string literals are still part of the surface, this is a small behavior regression; suggest raw = _parse_integer_value(value, target_type=anchor_type) here to keep parity.

[liuzidi]

Fixed in 0331b1e. _materialize_runtime_literal now uses _parse_integer_value(value, target_type=anchor_type) before creating the signless constant, preserving integer strings such as "0x10".

[Zhendong404]

Test coverage suggestions: (1) add an unsigned (pto.ui32) case covering add/floordiv/mod/compare/min/max — that is where the reconcile rule can silently flip op selection to signed (divui→floordivsi, ult→slt); (2) this negative regex is tied to one specific malformed shape (%c* name, si32→i32); a more robust check would catch any unrealized_conversion_cast whose declared source type mismatches the defining constant, or at least generalize to si\d+ to i\d+; (3) import re belongs at module top; (4) assert "si32" in text / "i32" in text are nearly tautological and can be dropped or tightened.

[liuzidi]

The unsigned regression was added in 0331b1e and checks divui, remui, cmpi ult, maxui, and minui. The malformed-cast assertion remains focused on the exact siN-to-iN constant-cast shape that triggered this PR; authored-type propagation now covers the broader semantic issue. The top-level import re and assertion cleanup are minor follow-ups, and can be included in a subsequent cleanup if preferred.