Skip to content

pre_quant doc & verify #922

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
and0d0 wants to merge 1 commit into
base: main
Choose a base branch
from
Open

pre_quant doc & verify #922

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
105 changes: 105 additions & 0 deletions lib/PTO/IR/VPTO.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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<IntegerType>(type))
return intType.getWidth() == 32;
return false;
case StructuredAccStoreDestinationFamily::I16:
if (auto intType = dyn_cast<IntegerType>(type))
return intType.getWidth() == 16 && !intType.isUnsigned();
return false;
case StructuredAccStoreDestinationFamily::I8:
if (auto intType = dyn_cast<IntegerType>(type))
return intType.getWidth() == 8;
return false;
case StructuredAccStoreDestinationFamily::I4:
if (auto intType = dyn_cast<IntegerType>(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)
Expand Down Expand Up @@ -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))
Expand Down
63 changes: 63 additions & 0 deletions ptodsl/docs/user_guide/07-data-movement-ops.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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 |
Expand Down
84 changes: 83 additions & 1 deletion ptodsl/ptodsl/_ops.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3716,16 +3716,98 @@ 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:
return None, None
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"
)
Comment on lines +3797 to +3807

@gemini-code-assist gemini-code-assist Bot Jul 9, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

对于 no_convert 模式，由于不需要进行量化或转换，因此不应该强制要求 payload 必须是 float 标量。在当前的逻辑中，如果 normalized_mode 为 no_convert，它会进入 else 分支并触发 TypeError。

建议对 no_convert 模式进行特殊处理，跳过 payload 的类型校验。

    if normalized_mode == "no_convert":
        pass
    elif 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<quant_pre_mode {_normalize_token(mode, context='acc store pre_quant mode')}>"),
Attribute.parse(f"#pto<quant_pre_mode {normalized_mode}>"),
)


Expand Down
Loading
Loading