Skip to content

[PTO][IR] align textract tinsert and tquant advanced parameters #729

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
FangRui0 wants to merge 1 commit into
base: main
Choose a base branch
from
Open

[PTO][IR] align textract tinsert and tquant advanced parameters #729

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
29 changes: 28 additions & 1 deletion docs/PTO_IR_manual.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7273,10 +7273,18 @@ dst[i + indexRow, j + indexCol] = src[i, j]
| Name | Type | Description |
|------|------|-------------|
| `src` | `pto.tile_buf` | Source tile |
| `preQuantScalar` | `i64` (optional) | Optional pre-quant scalar for acc-source extraction forms |
| `indexRow` | `Index` | Destination row offset |
| `indexCol` | `Index` | Destination column offset |
| `dst` | `pto.tile_buf` | Destination tile |

**Attributes:**

| Name | Type | Description |
|------|------|-------------|
| `accToVecMode` | `#pto.acc_to_vec_mode<...>` (optional) | Acc-to-vec mode for acc-source vec destinations |
| `reluPreMode` | `#pto.relu_pre_mode<...>` | Optional relu-pre mode for acc-source forms |

**Results:** None. Writes into `dst` via DPS pattern.

**Constraints & Verification:**
Expand All @@ -7293,6 +7301,8 @@ dst[i + indexRow, j + indexCol] = src[i, j]

```mlir
pto.tinsert ins(%src, %row, %col : !pto.tile_buf<...>, index, index) outs(%dst : !pto.tile_buf<...>)
pto.tinsert ins(%src, %preq, %row, %col : !pto.tile_buf<...>, i64, index, index)
outs(%dst : !pto.tile_buf<...>) {reluPreMode = #pto.relu_pre_mode<normal_relu>}
```

---
Expand All @@ -7312,10 +7322,18 @@ dst[i, j] = src[i + indexRow, j + indexCol]
| Name | Type | Description |
|------|------|-------------|
| `src` | `pto.tile_buf` | Source tile |
| `preQuantScalar` | `i64` (optional) | Optional pre-quant scalar for acc-source extraction forms |
| `indexRow` | `Index` | Starting row |
| `indexCol` | `Index` | Starting column |
| `dst` | `pto.tile_buf` | Destination tile |

**Attributes:**

| Name | Type | Description |
|------|------|-------------|
| `accToVecMode` | `#pto.acc_to_vec_mode<...>` (optional) | Acc-to-vec mode for acc-source vec destinations |
| `reluPreMode` | `#pto.relu_pre_mode<...>` | Optional relu-pre mode for acc-source forms |

**Results:** None. Writes into `dst` via DPS pattern.

**Constraints & Verification:**
Expand All @@ -7341,6 +7359,8 @@ dst[i, j] = src[i + indexRow, j + indexCol]

```mlir
pto.textract ins(%src[%row, %col] : !pto.tile_buf<...>) outs(%dst : !pto.tile_buf<...>)
pto.textract ins(%src, %preq, %row, %col : !pto.tile_buf<...>, i64, index, index)
outs(%dst : !pto.tile_buf<...>) {reluPreMode = #pto.relu_pre_mode<normal_relu>}
```

---
Expand Down Expand Up @@ -7932,14 +7952,19 @@ dst[i, j] = Quantize(src[i, j]; fp, quant_type)
| Name | Type | Description |
|------|------|-------------|
| `src` | `pto.tile_buf` | Source tile (`f32`) |
| `fp` | `pto.tile_buf` | Scaling parameter tile |
| `fp` | `pto.tile_buf` (optional) | Legacy scaling parameter tile |
| `offset` | `pto.tile_buf` (optional) | Legacy asymmetric offset tile |
| `exp` / `max` / `scaling` | `pto.tile_buf` (optional) | A5 statistic/scaling helper tiles |
| `expZZ` | `pto.tile_buf` (optional) | A5 vec-store helper tile |
| `dst` | `pto.tile_buf` | Destination tile (`i8` for SYM, `ui8` for ASYM) |

**Attributes:**

| Name | Type | Description |
|------|------|-------------|
| `quant_type` | `#pto.quant_type` | `INT8_SYM` or `INT8_ASYM` |
| `quantScaleAlg` | `#pto.quant_scale_alg<...>` (optional) | A5 scale-algorithm selector for `exp/max/scaling` form |
| `vecStoreMode` | `#pto.vec_store_mode<...>` (optional) | A5 vec-store mode selector for `expZZ` form |

**Results:** None. Writes into `dst` via DPS pattern.

Expand All @@ -7948,6 +7973,8 @@ dst[i, j] = Quantize(src[i, j]; fp, quant_type)
- `src` element type must be `f32`.
- `dst` element type must be `i8` (`INT8_SYM`) or `ui8` (`INT8_ASYM`).
- A2/A3: `src` and `dst` must use row-major layout.
- Legacy scale-tile form uses `fp` and optional `offset`.
- A5 extended forms use `exp/max/scaling` or `exp/max/scaling/expZZ`; they are mutually exclusive with the legacy `fp/offset` form.

**Hardware Mapping:**

Expand Down
33 changes: 33 additions & 0 deletions include/PTO/IR/PTOAttrs.td
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -880,6 +880,39 @@ def PTO_AccToVecModeAttr : PTO_Attr<"AccToVecMode", "acc_to_vec_mode"> {
let parameters = (ins EnumParameter<PTO_AccToVecMode_Enum>:$value);
let assemblyFormat = "`<` params `>`";
}

def PTO_TInsertModeEnum : PTO_I32Enum<
"TInsertMode", "PTO A5 tinsert vec->mat mode", [
I32EnumAttrCase<"SPLIT2", 2, "split2">,
I32EnumAttrCase<"SPLIT4", 3, "split4">
]>;

def PTO_TInsertModeAttr : PTO_Attr<"TInsertMode", "tinsert_mode"> {
let parameters = (ins EnumParameter<PTO_TInsertModeEnum>:$value);
let assemblyFormat = "`<` params `>`";
}

def PTO_QuantScaleAlgEnum : PTO_I32Enum<
"QuantScaleAlg", "PTO TQUANT quant scale algorithm", [
I32EnumAttrCase<"OCP", 0, "ocp">,
I32EnumAttrCase<"NV", 1, "nv">
]>;

def PTO_QuantScaleAlgAttr : PTO_Attr<"QuantScaleAlg", "quant_scale_alg"> {
let parameters = (ins EnumParameter<PTO_QuantScaleAlgEnum>:$value);
let assemblyFormat = "`<` params `>`";
}

def PTO_VecStoreModeEnum : PTO_I32Enum<
"VecStoreMode", "PTO TQUANT vec store mode", [
I32EnumAttrCase<"ND", 0, "nd">,
I32EnumAttrCase<"NZ", 1, "nz">
]>;

def PTO_VecStoreModeAttr : PTO_Attr<"VecStoreMode", "vec_store_mode"> {
let parameters = (ins EnumParameter<PTO_VecStoreModeEnum>:$value);
let assemblyFormat = "`<` params `>`";
}
//===----------------------------------------------------------------------===//
// MaskPattern
//===----------------------------------------------------------------------===//
Expand Down
51 changes: 32 additions & 19 deletions include/PTO/IR/PTOOps.td
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4119,20 +4119,18 @@ def TExtractOp : PTO_TOp<"textract", [

let arguments = (ins
PTODpsType:$src,
Optional<I64>:$preQuantScalar,
Index:$indexRow,
Index:$indexCol,
PTODpsType:$dst
PTODpsType:$dst,
OptionalAttr<PTO_AccToVecModeAttr>:$accToVecMode,
DefaultValuedAttr<PTO_ReluPreModeAttr, "::mlir::pto::ReluPreMode::NoRelu">:$reluPreMode
);

let results = (outs);

let hasVerifier = 1;

let assemblyFormat = [{
`ins` `(` $src `,` $indexRow `,` $indexCol `:` qualified(type($src)) `,` type($indexRow) `,` type($indexCol) `)`
`outs` `(` $dst `:` qualified(type($dst) ) `)`
attr-dict
}];
let hasCustomAssemblyFormat = 1;

let extraClassDeclaration = [{
::mlir::pto::PIPE getPipe() {
Expand Down Expand Up @@ -4207,7 +4205,9 @@ def TExtractFPOp : PTO_TOp<"textract_fp", [
PTODpsType:$fp,
Index:$indexRow,
Index:$indexCol,
PTODpsType:$dst
PTODpsType:$dst,
OptionalAttr<PTO_AccToVecModeAttr>:$accToVecMode,
DefaultValuedAttr<PTO_ReluPreModeAttr, "::mlir::pto::ReluPreMode::NoRelu">:$reluPreMode
);

let results = (outs);
Expand Down Expand Up @@ -4235,20 +4235,19 @@ def TInsertOp : PTO_TOp<"tinsert", [

let arguments = (ins
PTODpsType:$src,
Optional<I64>:$preQuantScalar,
Index:$indexRow,
Index:$indexCol,
PTODpsType:$dst
PTODpsType:$dst,
OptionalAttr<PTO_AccToVecModeAttr>:$accToVecMode,
DefaultValuedAttr<PTO_ReluPreModeAttr, "::mlir::pto::ReluPreMode::NoRelu">:$reluPreMode,
OptionalAttr<PTO_TInsertModeAttr>:$tinsertMode
);

let results = (outs);

let hasVerifier = 1;

let assemblyFormat = [{
`ins` `(` $src `,` $indexRow `,` $indexCol `:` qualified(type($src)) `,` type($indexRow) `,` type($indexCol) `)`
`outs` `(` $dst `:` qualified(type($dst) ) `)`
attr-dict
}];
let hasCustomAssemblyFormat = 1;

let extraClassDeclaration = [{
// TINSERT runs on different pipes depending on target/path:
Expand Down Expand Up @@ -4324,7 +4323,9 @@ def TInsertFPOp : PTO_TOp<"tinsert_fp", [
PTODpsType:$fp,
Index:$indexRow,
Index:$indexCol,
PTODpsType:$dst
PTODpsType:$dst,
OptionalAttr<PTO_AccToVecModeAttr>:$accToVecMode,
DefaultValuedAttr<PTO_ReluPreModeAttr, "::mlir::pto::ReluPreMode::NoRelu">:$reluPreMode
);

let results = (outs);
Expand Down Expand Up @@ -5161,6 +5162,7 @@ def TPReluOp: PTO_TOp<"tprelu", [
//===----------------------------------------------------------------------===//

def TQuantOp : PTO_TOp<"tquant", [
AttrSizedOperandSegments,
PTO_DpsInitOpInterface,
OpPipeInterface,
DeclareOpInterfaceMethods<MemoryEffectsOpInterface>
Expand All @@ -5183,18 +5185,29 @@ def TQuantOp : PTO_TOp<"tquant", [

let arguments = (ins
PTODpsType:$src,
PTODpsType:$fp,
Optional<PTODpsType>:$fp,
Optional<PTODpsType>:$offset,
Optional<PTODpsType>:$exp,
Optional<PTODpsType>:$max,
Optional<PTODpsType>:$scaling,
Optional<PTODpsType>:$expZZ,
PTODpsType:$dst,
PTO_QuantTypeAttr:$quant_type
PTO_QuantTypeAttr:$quant_type,
OptionalAttr<PTO_QuantScaleAlgAttr>:$quantScaleAlg,
OptionalAttr<PTO_VecStoreModeAttr>:$vecStoreMode
);

let results = (outs);

let hasVerifier = 1;

let assemblyFormat = [{
`ins` `(` $src `,` $fp (`,` $offset^)? `:` qualified(type($src)) `,` qualified(type($fp)) (`,` qualified(type($offset))^)? `)`
`ins` `(` $src
(`,` $fp^)? (`,` $offset^)? (`,` $exp^)? (`,` $max^)? (`,` $scaling^)? (`,` $expZZ^)?
`:` qualified(type($src))
(`,` qualified(type($fp))^)? (`,` qualified(type($offset))^)? (`,` qualified(type($exp))^)?
(`,` qualified(type($max))^)? (`,` qualified(type($scaling))^)? (`,` qualified(type($expZZ))^)?
`)`
`outs` `(` $dst `:` qualified(type($dst)) `)`
attr-dict
}];
Comment on lines 5204 to 5213
Copy link
Copy Markdown

@gemini-code-assist gemini-code-assist Bot May 30, 2026

Choose a reason for hiding this comment

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

critical

The declarative assembly format for TQuantOp has multiple optional operands ($fp, $offset, $exp, $max, $scaling, $expZZ) parsed sequentially without any distinguishing keywords or markers.

Because the parser cannot distinguish SSA values by name, if any of the leading optional operands are omitted (for example, in the A5 form where fp and offset are omitted but exp, max, and scaling are provided), the parser will sequentially assign the first parsed optional operand to $fp, the second to $offset, and so on. This results in completely incorrect operand mapping and silent correctness bugs.

To resolve this, you should either:

  1. Define a custom parser and printer (let hasCustomAssemblyFormat = 1;) to handle the two mutually exclusive forms unambiguously.
  2. Or introduce distinguishing keywords in the declarative assembly format, e.g., (, fp = $fp^)? (, offset = $offset^)? etc.

Expand Down
Loading
Loading