diff --git a/lib/PTO/IR/VPTO.cpp b/lib/PTO/IR/VPTO.cpp index a73ddf8059..a9a1e1a2e3 100644 --- a/lib/PTO/IR/VPTO.cpp +++ b/lib/PTO/IR/VPTO.cpp @@ -2566,6 +2566,105 @@ static bool isStructuredAccStoreInt32PreQuantMode(AccStoreQuantPreMode mode) { } } +enum class StructuredAccStoreDestinationFamily { + Any, + F32, + F16, + BF16, + I32, + I16, + I8, + I4, + FP8 +}; + +static StructuredAccStoreDestinationFamily +getStructuredAccStorePreQuantDestinationFamily(AccStoreQuantPreMode mode) { + switch (mode) { + case AccStoreQuantPreMode::F32F16: + case AccStoreQuantPreMode::QF322F16PreVec: + case AccStoreQuantPreMode::QF322F16PreScalar: + case AccStoreQuantPreMode::DEQF16Vec: + case AccStoreQuantPreMode::DEQF16Scalar: + return StructuredAccStoreDestinationFamily::F16; + case AccStoreQuantPreMode::F32BF16: + case AccStoreQuantPreMode::QF322BF16PreVec: + case AccStoreQuantPreMode::QF322BF16PreScalar: + case AccStoreQuantPreMode::QS322BF16PreVec: + case AccStoreQuantPreMode::QS322BF16PreScalar: + return StructuredAccStoreDestinationFamily::BF16; + case AccStoreQuantPreMode::QF322F32PreVec: + case AccStoreQuantPreMode::QF322F32PreScalar: + return StructuredAccStoreDestinationFamily::F32; + case AccStoreQuantPreMode::QF322HIF8PreVec: + case AccStoreQuantPreMode::QF322HIF8PreScalar: + case AccStoreQuantPreMode::QF322HIF8PreHybridVec: + case AccStoreQuantPreMode::QF322HIF8PreHybridScalar: + case AccStoreQuantPreMode::QF322FP8PreVec: + case AccStoreQuantPreMode::QF322FP8PreScalar: + return StructuredAccStoreDestinationFamily::FP8; + case AccStoreQuantPreMode::DEQS32IntVec: + case AccStoreQuantPreMode::DEQS32IntScalar: + return StructuredAccStoreDestinationFamily::I32; + case AccStoreQuantPreMode::QF162S16PreVec: + case AccStoreQuantPreMode::QF162S16PreScalar: + case AccStoreQuantPreMode::DEQS16Vec: + case AccStoreQuantPreMode::DEQS16Scalar: + return StructuredAccStoreDestinationFamily::I16; + case AccStoreQuantPreMode::QF162B8PreVec: + case AccStoreQuantPreMode::QF162B8PreScalar: + case AccStoreQuantPreMode::REQ8Vec: + case AccStoreQuantPreMode::REQ8Scalar: + case AccStoreQuantPreMode::QF322B8PreVec: + case AccStoreQuantPreMode::QF322B8PreScalar: + return StructuredAccStoreDestinationFamily::I8; + case AccStoreQuantPreMode::QF162S4PreVec: + case AccStoreQuantPreMode::QF162S4PreScalar: + case AccStoreQuantPreMode::REQ4Vec: + case AccStoreQuantPreMode::REQ4Scalar: + case AccStoreQuantPreMode::QF322S4PreVec: + case AccStoreQuantPreMode::QF322S4PreScalar: + return StructuredAccStoreDestinationFamily::I4; + case AccStoreQuantPreMode::NoConvert: + return StructuredAccStoreDestinationFamily::Any; + } + llvm_unreachable("unknown acc-store pre_quant mode"); +} + +static bool isStructuredAccStoreDestinationFamily( + Type type, StructuredAccStoreDestinationFamily family) { + switch (family) { + case StructuredAccStoreDestinationFamily::Any: + return true; + case StructuredAccStoreDestinationFamily::F32: + return type.isF32(); + case StructuredAccStoreDestinationFamily::F16: + return type.isF16(); + case StructuredAccStoreDestinationFamily::BF16: + return type.isBF16(); + case StructuredAccStoreDestinationFamily::I32: + if (auto intType = dyn_cast(type)) + return intType.getWidth() == 32; + return false; + case StructuredAccStoreDestinationFamily::I16: + if (auto intType = dyn_cast(type)) + return intType.getWidth() == 16 && !intType.isUnsigned(); + return false; + case StructuredAccStoreDestinationFamily::I8: + if (auto intType = dyn_cast(type)) + return intType.getWidth() == 8; + return false; + case StructuredAccStoreDestinationFamily::I4: + if (auto intType = dyn_cast(type)) + return intType.getWidth() == 4 && !intType.isUnsigned(); + return false; + case StructuredAccStoreDestinationFamily::FP8: + return pto::isPTOFloat8Type(type) || pto::isPTOHiFloat8Type(type) || + pto::isPTOHiFloat8x2Type(type); + } + llvm_unreachable("unknown acc-store destination family"); +} + static ParseResult parseStructuredAccStoreUnitFlag(OpAsmParser &parser, StructuredAccStoreAsmState &state) { if (state.unitFlag) @@ -2856,6 +2955,12 @@ static LogicalResult verifyStructuredAccStoreLike( << "pre_quant requires source element type to be f32 or i32, got " << sourceElementType; } + + StructuredAccStoreDestinationFamily destinationFamily = + getStructuredAccStorePreQuantDestinationFamily(*preQuantMode); + if (!isStructuredAccStoreDestinationFamily(destinationElementType, + destinationFamily)) + return emitIncompatibleQuantModeError(); } if (clipValue && !isStructuredAccStoreClipSupportedElementType(destinationElementType)) diff --git a/ptodsl/docs/user_guide/07-data-movement-ops.md b/ptodsl/docs/user_guide/07-data-movement-ops.md index cabe44abdf..308908a4a0 100644 --- a/ptodsl/docs/user_guide/07-data-movement-ops.md +++ b/ptodsl/docs/user_guide/07-data-movement-ops.md @@ -1046,6 +1046,69 @@ pto.mte_l0c_ub(acc, ub, 16, 32, 16, 32, split=pto.SplitMode.N) # split N --- +#### `pre_quant` modes and destination types + +`pre_quant=(payload, mode)` applies a conversion or quantization step before +`pre_relu`, `clip`, saturation, and the final store. The selected mode must +match both the accumulator source element type and the destination element +type. + +Payload rules: + +- `_vec` modes require a scaling pointer payload with `f16`, `bf16`, or `f32` + elements. +- `_scalar` modes require an `f16`, `bf16`, or `f32` scalar payload. +- `f32_f16` and `f32_bf16` also require an `f16`, `bf16`, or `f32` scalar + payload; the payload is required by the public syntax but does not select + per-channel scaling. + +Legal mode families: + +| Mode family | Accumulator source dtype | Legal destination dtype | +|-------------|--------------------------|--------------------------| +| `f32_f16`, `qf322f16_pre_*`, `deqf16_*` | `f32` for `f32_f16` / `qf322f16_pre_*`; `i32` for `deqf16_*` | `f16` | +| `f32_bf16`, `qf322bf16_pre_*`, `qs322bf16_pre_*` | `f32` for `f32_bf16` / `qf322bf16_pre_*`; `i32` for `qs322bf16_pre_*` | `bf16` | +| `qf322f32_pre_*` | `f32` | `f32` | +| `qf322hif8_pre_*`, `qf322fp8_pre_*` | `f32` | PTO FP8 / HiFloat8 family | +| `deqs32_int_*` | `i32` | 32-bit integer | +| `qf162s16_pre_*`, `deqs16_*` | `i32` | signed or signless 16-bit integer | +| `qf162b8_pre_*`, `req8_*`, `qf322b8_pre_*` | `f32` for `qf162b8_pre_*` / `qf322b8_pre_*`; `i32` for `req8_*` | 8-bit integer | +| `qf162s4_pre_*`, `req4_*`, `qf322s4_pre_*` | `f32` for `qf162s4_pre_*` / `qf322s4_pre_*`; `i32` for `req4_*` | signed or signless 4-bit integer | + +Examples: + +```python +# Legal: f32 accumulator values are converted to bf16 on writeback. +pto.mte_l0c_gm( + acc_f32, + dst_bf16, + 16, + 16, + 16, + 16, + 0, + 0, + pre_quant=(pto.bf16(1.0), "f32_bf16"), + layout="nz2nd", +) + +# Illegal: "f32_bf16" requires a bf16 destination, not f16. +pto.mte_l0c_gm( + acc_f32, + dst_f16, + 16, + 16, + 16, + 16, + 0, + 0, + pre_quant=(pto.bf16(1.0), "f32_bf16"), + layout="nz2nd", +) +``` + +--- + ### Cube data movement quick reference | Data Flow | Operation | Src Space | Dst Space | diff --git a/ptodsl/ptodsl/_ops.py b/ptodsl/ptodsl/_ops.py index 53dd3e2255..95774838d2 100644 --- a/ptodsl/ptodsl/_ops.py +++ b/ptodsl/ptodsl/_ops.py @@ -3716,6 +3716,74 @@ def _acc_store_unit_flag_attr(unit_flag): context="acc store unit_flag", ) +_ACC_STORE_PRE_QUANT_MODES = frozenset({ + "no_convert", + "f32_f16", + "qf322hif8_pre_vec", + "qf322hif8_pre_scalar", + "qf322hif8_pre_hybrid_vec", + "qf322hif8_pre_hybrid_scalar", + "deqs32_int_vec", + "deqs32_int_scalar", + "req8_vec", + "req8_scalar", + "deqf16_vec", + "deqf16_scalar", + "qf322fp8_pre_vec", + "qf322fp8_pre_scalar", + "qf322f32_pre_vec", + "qf322f32_pre_scalar", + "f32_bf16", + "qf162b8_pre_vec", + "qf162b8_pre_scalar", + "qf162s4_pre_vec", + "qf162s4_pre_scalar", + "req4_vec", + "req4_scalar", + "qf322b8_pre_vec", + "qf322b8_pre_scalar", + "qf322s4_pre_vec", + "qf322s4_pre_scalar", + "deqs16_vec", + "deqs16_scalar", + "qf162s16_pre_vec", + "qf162s16_pre_scalar", + "qf322f16_pre_vec", + "qf322f16_pre_scalar", + "qf322bf16_pre_vec", + "qf322bf16_pre_scalar", + "qs322bf16_pre_vec", + "qs322bf16_pre_scalar", +}) + + +def _is_acc_store_float_scalar_payload(value) -> bool: + raw_value = unwrap_surface_value(value) + return hasattr(raw_value, "type") and any( + cls.isinstance(raw_value.type) for cls in (F16Type, BF16Type, F32Type) + ) + + +def _is_acc_store_scaling_pointer_payload(value) -> bool: + raw_value = unwrap_surface_value(value) + if not hasattr(raw_value, "type"): + return False + try: + ptr_type = _pto.PtrType(raw_value.type) + except Exception: + return False + scaling_attr = _pto.AddressSpaceAttr.get(_pto.AddressSpace.SCALING) + memory_space = getattr(ptr_type, "memory_space", None) + if ( + memory_space != scaling_attr + and getattr(memory_space, "value", None) != _pto.AddressSpace.SCALING + ): + return False + return any( + cls.isinstance(ptr_type.element_type) + for cls in (F16Type, BF16Type, F32Type) + ) + def _acc_store_pre_quant(pre_quant): if pre_quant is None: @@ -3723,9 +3791,23 @@ def _acc_store_pre_quant(pre_quant): if not isinstance(pre_quant, tuple) or len(pre_quant) != 2: raise TypeError("acc store pre_quant expects (payload, mode)") payload, mode = pre_quant + normalized_mode = _normalize_token(mode, context="acc store pre_quant mode") + if normalized_mode not in _ACC_STORE_PRE_QUANT_MODES: + raise ValueError(f"unsupported acc store pre_quant mode: {normalized_mode}") + if normalized_mode.endswith("_vec"): + if not _is_acc_store_scaling_pointer_payload(payload): + raise TypeError( + "acc store vector pre_quant payload must be a scaling pointer " + "with f16, bf16, or f32 elements" + ) + else: + if not _is_acc_store_float_scalar_payload(payload): + raise TypeError( + "acc store scalar pre_quant payload must be an f16, bf16, or f32 scalar" + ) return ( unwrap_surface_value(payload), - Attribute.parse(f"#pto"), + Attribute.parse(f"#pto"), ) diff --git a/ptodsl/tests/test_jit_compile.py b/ptodsl/tests/test_jit_compile.py index 557c4b3a2a..9852d69c33 100644 --- a/ptodsl/tests/test_jit_compile.py +++ b/ptodsl/tests/test_jit_compile.py @@ -1793,6 +1793,96 @@ def public_cube_surface_probe( pto.mte_l0c_ub(acc_tile.as_ptr(), out_tile.as_ptr(), m, n, n, n, split=pto.SplitMode.M, layout="nz2nd") +@pto.jit(target="a5", mode="explicit") +def acc_store_pre_quant_surface_probe( + dst: pto.ptr(pto.bf16, "gm"), +): + size = pto.const(16) + zero = pto.const(0) + src = pto.alloc_tile( + shape=[16, 16], + dtype=pto.f32, + memory_space=pto.MemorySpace.ACC, + valid_shape=[16, 16], + blayout="ColMajor", + slayout="RowMajor", + ) + pto.mte_l0c_gm( + src.as_ptr(), + dst, + size, + size, + size, + size, + zero, + zero, + pre_quant=(pto.bf16(1.0), "f32_bf16"), + layout="nz2nd", + ) + + +@pto.jit(target="a5", mode="explicit") +def acc_store_bad_pre_quant_mode_probe( + dst: pto.ptr(pto.bf16, "gm"), +): + size = pto.const(16) + zero = pto.const(0) + src = pto.alloc_tile( + shape=[16, 16], + dtype=pto.f32, + memory_space=pto.MemorySpace.ACC, + valid_shape=[16, 16], + blayout="ColMajor", + slayout="RowMajor", + ) + pto.mte_l0c_gm( + src.as_ptr(), + dst, + size, + size, + size, + size, + zero, + zero, + pre_quant=(pto.bf16(1.0), "not_a_mode"), + layout="nz2nd", + ) + + +@pto.jit(target="a5", mode="explicit") +def acc_store_bad_pre_quant_payload_probe( + dst: pto.ptr(pto.bf16, "gm"), +): + size = pto.const(16) + zero = pto.const(0) + src = pto.alloc_tile( + shape=[16, 16], + dtype=pto.f32, + memory_space=pto.MemorySpace.ACC, + valid_shape=[16, 16], + blayout="ColMajor", + slayout="RowMajor", + ) + payload = pto.alloc_tile( + shape=[16], + dtype=pto.bf16, + memory_space=pto.MemorySpace.SCALING, + valid_shape=[16], + ) + pto.mte_l0c_gm( + src.as_ptr(), + dst, + size, + size, + size, + size, + zero, + zero, + pre_quant=(payload.as_ptr(), "f32_bf16"), + layout="nz2nd", + ) + + @pto.cube def public_cube_tile_mx_probe( mat_lhs: pto.Tile, @@ -4982,6 +5072,22 @@ def _enter_inline_simt_with_resource_attr(): public_surface_text = public_surface_exports_probe.compile().mlir_text() expect_parse_roundtrip_and_verify(public_surface_text, "public surface export specialization") + acc_store_pre_quant_text = acc_store_pre_quant_surface_probe.compile().mlir_text() + expect_parse_roundtrip_and_verify(acc_store_pre_quant_text, "acc-store pre_quant public mode specialization") + expect( + "pre_quant(%" in acc_store_pre_quant_text and "mode = f32_bf16" in acc_store_pre_quant_text, + "mte_l0c_gm pre_quant should preserve the f32_bf16 public mode", + ) + expect_raises( + ValueError, + lambda: acc_store_bad_pre_quant_mode_probe.compile().mlir_text(), + "unsupported acc store pre_quant mode", + ) + expect_raises( + TypeError, + lambda: acc_store_bad_pre_quant_payload_probe.compile().mlir_text(), + "acc store scalar pre_quant payload must be an f16, bf16, or f32 scalar", + ) compile_time_query_text = compile_time_query_probe.compile().mlir_text() expect_parse_roundtrip_and_verify(compile_time_query_text, "compile-time query specialization") eager_scalar_text = eager_scalar_constructor_probe.compile().mlir_text() diff --git a/test/lit/vpto/acc_store_gm_f32_bf16.pto b/test/lit/vpto/acc_store_gm_f32_bf16.pto new file mode 100644 index 0000000000..378f09d673 --- /dev/null +++ b/test/lit/vpto/acc_store_gm_f32_bf16.pto @@ -0,0 +1,36 @@ +// Copyright (c) 2026 Huawei Technologies Co., Ltd. +// This program is free software, you can redistribute it and/or modify it under the terms and conditions of +// CANN Open Software License Agreement Version 2.0 (the "License"). +// Please refer to the License for details. You may not use this file except in compliance with the License. +// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +// INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +// See LICENSE in the root of the software repository for the full text of the License. + +// RUN: ptoas --pto-arch=a5 --pto-backend=vpto --emit-vpto --mlir-print-ir-before=vpto-expand-wrapper-ops %s -o /dev/null 2>&1 | FileCheck %s --check-prefix=ROUNDTRIP +// RUN: ptoas --pto-arch=a5 --pto-backend=vpto --emit-vpto --mlir-print-ir-after=vpto-expand-wrapper-ops %s -o /dev/null 2>&1 | FileCheck %s --check-prefix=EXPAND + +module attributes {"pto.target_arch" = "a5", pto.kernel_kind = #pto.kernel_kind} { + func.func @acc_store_gm_f32_bf16(%src: !pto.ptr, + %dst: !pto.ptr) { + %c0_i64 = arith.constant 0 : i64 + %c16_i64 = arith.constant 16 : i64 + %c32_i64 = arith.constant 32 : i64 + %c1_bf16 = arith.constant 1.000000e+00 : bf16 + + pto.mte_l0c_gm %src, %dst, %c16_i64, %c16_i64, %c16_i64, %c32_i64, + %c0_i64, %c0_i64, + pre_quant(%c1_bf16, mode = f32_bf16), + nz2nd + : !pto.ptr, !pto.ptr, i64, i64, i64, i64, i64, i64, + bf16 + + return + } +} + +// ROUNDTRIP-LABEL: func.func @acc_store_gm_f32_bf16( +// ROUNDTRIP: pto.mte_l0c_gm %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, pre_quant(%{{.*}}, mode = f32_bf16), nz2nd + +// EXPAND-LABEL: func.func @acc_store_gm_f32_bf16( +// EXPAND: pto.set_quant_pre +// EXPAND: pto.copy_matrix_cc_to_gm diff --git a/test/lit/vpto/acc_store_verify_invalid_f32_bf16_dst.pto b/test/lit/vpto/acc_store_verify_invalid_f32_bf16_dst.pto new file mode 100644 index 0000000000..2a882380dd --- /dev/null +++ b/test/lit/vpto/acc_store_verify_invalid_f32_bf16_dst.pto @@ -0,0 +1,30 @@ +// Copyright (c) 2026 Huawei Technologies Co., Ltd. +// This program is free software, you can redistribute it and/or modify it under the terms and conditions of +// CANN Open Software License Agreement Version 2.0 (the "License"). +// Please refer to the License for details. You may not use this file except in compliance with the License. +// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +// INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +// See LICENSE in the root of the software repository for the full text of the License. + +// RUN: not ptoas --pto-arch=a5 %s -o - 2>&1 | FileCheck %s + +module attributes {"pto.target_arch" = "a5"} { + func.func @acc_store_verify_invalid_f32_bf16_dst( + %src: !pto.ptr, %dst: !pto.ptr) { + %c0_i64 = arith.constant 0 : i64 + %c16_i64 = arith.constant 16 : i64 + %c32_i64 = arith.constant 32 : i64 + %c1_bf16 = arith.constant 1.000000e+00 : bf16 + + pto.mte_l0c_gm %src, %dst, %c16_i64, %c16_i64, %c16_i64, %c32_i64, + %c0_i64, %c0_i64, + pre_quant(%c1_bf16, mode = f32_bf16), + nz2nd + : !pto.ptr, !pto.ptr, i64, i64, i64, i64, i64, i64, + bf16 + + return + } +} + +// CHECK: pre_quant mode f32_bf16 is incompatible with source element type 'f32' and destination element type 'f16' diff --git a/test/lit/vpto/acc_store_verify_invalid_f32_bf16_vec_payload.pto b/test/lit/vpto/acc_store_verify_invalid_f32_bf16_vec_payload.pto new file mode 100644 index 0000000000..f0b46d6ca2 --- /dev/null +++ b/test/lit/vpto/acc_store_verify_invalid_f32_bf16_vec_payload.pto @@ -0,0 +1,30 @@ +// Copyright (c) 2026 Huawei Technologies Co., Ltd. +// This program is free software, you can redistribute it and/or modify it under the terms and conditions of +// CANN Open Software License Agreement Version 2.0 (the "License"). +// Please refer to the License for details. You may not use this file except in compliance with the License. +// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +// INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +// See LICENSE in the root of the software repository for the full text of the License. + +// RUN: not ptoas --pto-arch=a5 %s -o - 2>&1 | FileCheck %s + +module attributes {"pto.target_arch" = "a5"} { + func.func @acc_store_verify_invalid_f32_bf16_vec_payload( + %src: !pto.ptr, %dst: !pto.ptr, + %payload: !pto.ptr) { + %c0_i64 = arith.constant 0 : i64 + %c16_i64 = arith.constant 16 : i64 + %c32_i64 = arith.constant 32 : i64 + + pto.mte_l0c_gm %src, %dst, %c16_i64, %c16_i64, %c16_i64, %c32_i64, + %c0_i64, %c0_i64, + pre_quant(%payload, mode = f32_bf16), + nz2nd + : !pto.ptr, !pto.ptr, i64, i64, i64, i64, i64, i64, + !pto.ptr + + return + } +} + +// CHECK: scalar pre_quant mode requires f16/bf16/f32 payload