| Add the int8 x int8 -> int32 data type route to tla.mmad Co-authored-by: Alessio<alessio.russo@h-partners.com> # message auto-generated for no-merge-commit merge: !1120 merge mmad-int8 into master Add the int8 x int8 -> int32 data type route to tla.mmad Created-by: arusso Commit-by: Alessio Merged-by: cann-robot Description: ## 描述 tla.mmad previously routed only f16, bf16 and f32 operands, all accumulating into an fp32 L0C. This adds i8 x i8 -> i32, which the cube supports through the same non-MX mad intrinsic the float routes already use. The device wrapper is a single Catlass::Gemm::Mmad<int8_t, int8_t, int32_t> instantiation alongside the existing three, exposed as mmad_int8_int8_int32. TlaCubeRegionPass gains one more branch in the element-type routing to select it, and core_api.py adds ("i8", "i8", "i32") to the tla.mmad element-type contract. The bc layer registers the int8 operand copies the route needs: GM->L1 in zN and nZ, and L1->L0A / L1->L0B in both orientations. The one structural difference from every existing route is the accumulator: int8 accumulates into an **i32** L0C rather than fp32. Every copy out of L0C therefore had to stop assuming an fp32 accumulator. That contract is now expressed once and applied to all four L0C routes (GM, UB row-major, UB column-major, L1 zN): cpp // Whether fixpipe can carry this L0C element type out to the given destination // type. An fp32 accumulator (the float MMAD routes) may be narrowed to f16/bf16 // on the way out; an i32 accumulator (the int8 MMAD route) has no narrowing path // and must land as i32. static bool isLegalFixpipeElementType(StringRef srcElementType, StringRef dstElementType); Each of those exits is registered for int32_t -> int32_t, so the integer route has the same shape as the float ones: L0C -> GM, L0C -> UB (NO_SPLIT / SPLIT_M / SPLIT_N, both layouts) and L0C -> L1 (zN). The frontend enforces the matching rule for callers: an i32 accumulator may only be copied out as i32, while an fp32 one keeps its existing [f32, f16, bf16] choice. basic_matmul.py gains the route as a data-type choice: --dtype-a i8 --dtype-b i8 --dtype-c i32. The kernel takes an acc_is_int Constexpr and allocates the L0C tile as tla.Int32 when set; the host generates small integer operands, builds an exact int32 reference, and compares with element-wise equality instead of a tolerance. ## 关联的Issue N/A ## 原因 The DSL exposed only three of the operand formats the cube supports, which blocked int8-quantised workloads from being written in it. int8 is the cheapest of the missing formats to add: it needs no microscaling machinery, no sub-byte packing and no new element type, since i8 and i32 already exist in the DSL. CATLASS already selects an int32_t accumulator for this operand pair, so the device side is a single template instantiation. ## 测试 lit, in mmad-runtime-from-python.test, with both inputs emitted from legal @tla.kernel functions rather than hand-written IR: - an I8 case asserting the lowering reaches @mmad_int8_int8_int32 and that no tla.mmad survives; - an I8FIX case exercising every exit from the i32 accumulator in one kernel, asserting each copy resolves to an int32_t -> int32_t callee — copy_l0c_to_gm_row_major_int32_t, copy_l0c_to_ub_row_major_splitm_int32_t, copy_l0c_to_l1_zN_int32_t — and that no tla.mmad / tla.copy survives. This is what guards the shared fixpipe predicate: an i32 accumulator must never pick one of the fp32 callees. On-device, basic_matmul.py --dtype-a i8 --dtype-b i8 --dtype-c i32 runs the route end to end and checks the result with torch.equal against an exact int32 reference; integer MMAD is exact, so any tolerance would mask a defect. The example rejects a partial integer selection (--dtype-a i8 without i8/i32 on the other two), so the route cannot be entered half-configured. Full suites: pytest 1494 passed / 7 skipped, lit 124/124. ## 文档更新 docs/en/kernel_api_reference.md is regenerated, since core_api.py changed and the generated reference embeds source line numbers; the diff is limited to those mechanical Source: links. No new API surface is introduced, so there is no new entry: i8 operands with an i32 accumulator are a data-type choice on the existing tla.mmad, not a new op or parameter. The one thing a caller must know — that the integer route accumulates into an i32 L0C rather than fp32 — is visible in basic_matmul.py, which selects the L0C element type accordingly, and in the tla.mmad error message, which now names the supported element-type combinations including i8,i8 -> i32. ## 类型标签 - [ ] Bug修复 - [x] 新特性 - [ ] 性能优化 - [ ] 文档更新 - [ ] 其他，请描述： See merge request: cann/catlass!1120 | 1 个月前 |
