[BugFix][TIRx] Legalize BF16 conversions inside Let and Bind values

[anerli]

Problem

BF16ComputeLegalize leaves a BF16 conversion in place when it sits inside the value of a Let expression or a Bind statement. BF16StorageLegalize then retypes the BF16 buffer to uint16, and the surviving Cast("float32", A[i]) becomes an integer-to-float conversion of the raw bits. A kernel that writes

b: T.float32 = T.Cast("float32", A[i])
C[i] = b * T.float32(2)

reads BF16 1/1024 as 14976.0. The same kernel with the cast written inline at the use site is legalized correctly.

Root cause

ComputeLegalizer::VisitExpr_(const LetNode*) and VisitStmt_(const BindNode*) call PromoteToTarget(op->value) on the unvisited value. Every other handler in the class (Cast, Select, Broadcast, Shuffle, BufferStore, the binary operators) calls PromoteToTarget(this->VisitExpr(...)). Nothing in TileLang runs this pass on the CPU pipeline today, and the CUDA gate skips it, so the two handlers had not processed real IR.

Change

• src/tirx/transform/unsupported_dtype_legalize.cc: visit the bound value before promoting it, in both handlers.

Tests

tests/python/tirx-transform/test_tir_transform_bf16_legalize.py:

• test_bf16_bind_value_will_legalize: a typed assignment in script (a Bind) whose value casts a BF16 load; structural equality against the legalized form.
• test_bf16_let_value_will_legalize: a Let expression built directly; asserts the Let survives and no cast to or from bfloat16 remains.

Validation

Apple M-series, macOS 15, built inside TileLang (85fd8fc2) with USE_METAL=ON USE_LLVM=ON (LLVM 20.1.8).

pytest tests/python/tirx-transform/test_tir_transform_bf16_legalize.py    7 passed

With TileLang's CPU pipeline running BF16ComputeLegalize (companion TileLang PR), testing/python/cpu plus testing/python/llvm: 120 passed, 0 failed. A downstream BF16 pointwise kernel that binds its casts computes correctly on the LLVM target with this change and reads raw bits without it.

Related

• TileLang PR [BugFix][CPU] Legalize BF16 arithmetic in the CPU pipeline adds BF16ComputeLegalize to CPUPassPipelineBody. It is independent of this fix and can merge in either order: it repairs the compile failures on its own, and this fix repairs the remaining case of a conversion bound to a variable, which TileLang's eager frontend emits for every assignment.
• Follow-up in TileLang once this merges: bump 3rdparty/tvm, and extend testing/python/cpu/test_tilelang_cpu_bf16_legalize.py with the bound-value kernel (a, b = A[i].astype("float32"), B[i].astype("float32")) in both the device-free lowering check and the LLVM execution check. Those cases fail against the current 3rdparty/tvm and pass with this change; they were left out of the TileLang PR for that reason.