feat: add B09 partial arg and histogram ops

docs/en/dev/codegen/00-pto_codegen.md

@@ -142,6 +142,9 @@ print(pto_code)
 | `tile.mul(lhs, rhs)` | `pto.tmul` |
 | `tile.add(a, b, c)` | `pto.taddc` (3-operand add) |
 | `tile.adds(tile, scalar)` | `pto.tadds` (tile + scalar) |
+| `tile.part_argmax(src0, src1, src0_idx, src1_idx)` | `pto.tpartargmax` (selected value and paired index outputs) |
+| `tile.part_argmin(src0, src1, src0_idx, src1_idx)` | `pto.tpartargmin` (selected value and paired index outputs) |
+| `tile.histogram(src, idx, byte=...)` | `pto.thistogram` (A5 cumulative 256-bin histogram) |
 | `tile.fillpad_expand(src, shape)` | `pto.tfillpad_expand ins(%src) outs(%dst)` (the `shape` tuple is type-deduction only; the larger `dst` and its pad come from the result type) |
 
 **`tile.slice` / `tile.assemble` lowering details.**  Both ops are lowered

docs/en/dev/ptoas-op-status.md

@@ -97,8 +97,8 @@ for lowering/compiler plumbing, plus other dialects such as VPTO, VMI, and SIMT.
 | pto.tpartadd | TPARTADD | tile+tensor | ✅ | ✅ | ✅ | ✅ | — |  |
 | pto.tpartmax | TPARTMAX | tile+tensor | ✅ | ✅ | ✅ | ✅ | — |  |
 | pto.tpartmin | TPARTMIN | tile+tensor | ✅ | ✅ | ✅ | ✅ | — |  |
-| pto.tpartargmax | TPARTARGMAX | tile | ✅ | ❌ | ❌ | ❌ | — | MISSING: lacks a complete frontend/codegen/ST path |
-| pto.tpartargmin | TPARTARGMIN | tile | ✅ | ❌ | ❌ | ❌ | — | MISSING: lacks a complete frontend/codegen/ST path |
+| pto.tpartargmax | TPARTARGMAX | tile | ✅ | ✅ | ✅ | ✅ | — | verified on A2/A3 hardware |
+| pto.tpartargmin | TPARTARGMIN | tile | ✅ | ✅ | ✅ | ✅ | — | verified on A2/A3 hardware |
 | pto.tpartmul | TPARTMUL | tile+tensor | ✅ | ✅ | ✅ | ✅ | — |  |
 | pto.tprelu | TPRELU | tile | ✅ | ✅ | ❌ | ❌ | — | path exists; historical ISA/semantic issue requires revalidation against the current pin |
 | pto.tadds | TADDS | tile+tensor | ✅ | ✅ | ✅ | ✅ | — |  |
@@ -136,7 +136,7 @@ for lowering/compiler plumbing, plus other dialects such as VPTO, VMI, and SIMT.
 | pto.trowmin | TROWMIN | tile+tensor | ✅ | ✅ | ✅ | ✅ | — | verified on A2/A3 hardware; A5 hardware verification pending |
 | pto.trowargmin | TROWARGMIN | tile+tensor | ✅ | ✅ | ✅ | ✅ | — |  |
 | pto.trowprod | TROWPROD | tile+tensor | ✅ | ✅ | ✅ | ✅ | — |  |
-| pto.thistogram | THISTOGRAM | tile | ✅ | ❌ | ❌ | ❌ | — | MISSING: lacks a complete frontend/codegen/ST path |
+| pto.thistogram | THISTOGRAM | tile | ✅ | ✅ | ✅ | ❌ | — | A5-only; same-name A5sim PTOAS compilation coverage added, A5 execution pending |
 | pto.tcolsum | TCOLSUM | tile+tensor | ✅ | ✅ | ✅ | ✅ | — |  |
 | pto.tcolmax | TCOLMAX | tile+tensor | ✅ | ✅ | ✅ | ✅ | — |  |
 | pto.tcolargmax | TCOLARGMAX | tile+tensor | ✅ | ✅ | ✅ | ✅ | — |  |
@@ -273,6 +273,6 @@ for lowering/compiler plumbing, plus other dialects such as VPTO, VMI, and SIMT.
 | **Source Compatibility / Manual Mode (1)** |  |  |  |  |  |  |  |  |
 | pto.tassign | TASSIGN | internal | ✅ | — | — | — | — | inactive backend hook; no standalone ST |
 
-**Stats**: 204 public/compatibility PTOAS ops; 113 have a pypto tile frontend and 75 have a tensor frontend;
-110 have same-name ST coverage (106 regular STs and 4 distributed STs); 62 lack same-name ST coverage
-(52 regular and 10 distributed); within these 204, another 32 ops are not suitable for standalone STs.
+**Stats**: 204 public/compatibility PTOAS ops; 116 have a pypto tile frontend and 75 have a tensor frontend;
+112 have same-name ST coverage (108 regular STs and 4 distributed STs); 60 lack same-name ST coverage
+(50 regular and 10 distributed); within these 204, another 32 ops are not suitable for standalone STs.

docs/zh/dev/codegen/00-pto_codegen.md

@@ -140,6 +140,9 @@ print(pto_code)
 | `tile.mul(lhs, rhs)` | `pto.tmul` |
 | `tile.add(a, b, c)` | `pto.taddc` (三操作数加法) |
 | `tile.adds(tile, scalar)` | `pto.tadds` (Tile + 标量) |
+| `tile.part_argmax(src0, src1, src0_idx, src1_idx)` | `pto.tpartargmax`（输出选中的值及配对索引） |
+| `tile.part_argmin(src0, src1, src0_idx, src1_idx)` | `pto.tpartargmin`（输出选中的值及配对索引） |
+| `tile.histogram(src, idx, byte=...)` | `pto.thistogram`（A5 累计 256-bin 直方图） |
 | `tile.fillpad_expand(src, shape)` | `pto.tfillpad_expand ins(%src) outs(%dst)`（`shape` 元组仅用于类型推导；更大的 `dst` 及其 pad 来自结果类型） |
 
 **`tile.slice` / `tile.assemble` 下沉细节。** 两个 op 都通过 `pto.subview`

docs/zh/dev/ptoas-op-status.md

@@ -83,8 +83,8 @@ lowering/compiler plumbing 使用的额外内部 op 未纳入，也不列 VPTO
 | pto.tpartadd | TPARTADD | tile+tensor | ✅ | ✅ | ✅ | ✅ | — |  |
 | pto.tpartmax | TPARTMAX | tile+tensor | ✅ | ✅ | ✅ | ✅ | — |  |
 | pto.tpartmin | TPARTMIN | tile+tensor | ✅ | ✅ | ✅ | ✅ | — |  |
-| pto.tpartargmax | TPARTARGMAX | tile | ✅ | ❌ | ❌ | ❌ | — | MISSING：缺完整前端/codegen/ST 链路 |
-| pto.tpartargmin | TPARTARGMIN | tile | ✅ | ❌ | ❌ | ❌ | — | MISSING：缺完整前端/codegen/ST 链路 |
+| pto.tpartargmax | TPARTARGMAX | tile | ✅ | ✅ | ✅ | ✅ | — | A2/A3 真机已验证 |
+| pto.tpartargmin | TPARTARGMIN | tile | ✅ | ✅ | ✅ | ✅ | — | A2/A3 真机已验证 |
 | pto.tpartmul | TPARTMUL | tile+tensor | ✅ | ✅ | ✅ | ✅ | — |  |
 | pto.tprelu | TPRELU | tile | ✅ | ✅ | ❌ | ❌ | — | 已有链路；历史 ISA/语义问题，需按当前 pin 复验 |
 | pto.tadds | TADDS | tile+tensor | ✅ | ✅ | ✅ | ✅ | — |  |
@@ -122,7 +122,7 @@ lowering/compiler plumbing 使用的额外内部 op 未纳入，也不列 VPTO
 | pto.trowmin | TROWMIN | tile+tensor | ✅ | ✅ | ✅ | ✅ | — | A2/A3 真机已验证；A5 真机待验证 |
 | pto.trowargmin | TROWARGMIN | tile+tensor | ✅ | ✅ | ✅ | ✅ | — |  |
 | pto.trowprod | TROWPROD | tile+tensor | ✅ | ✅ | ✅ | ✅ | — |  |
-| pto.thistogram | THISTOGRAM | tile | ✅ | ❌ | ❌ | ❌ | — | MISSING：缺完整前端/codegen/ST 链路 |
+| pto.thistogram | THISTOGRAM | tile | ✅ | ✅ | ✅ | ❌ | — | 仅 A5；已增加同名 A5sim PTOAS 编译覆盖，A5 执行待验证 |
 | pto.tcolsum | TCOLSUM | tile+tensor | ✅ | ✅ | ✅ | ✅ | — |  |
 | pto.tcolmax | TCOLMAX | tile+tensor | ✅ | ✅ | ✅ | ✅ | — |  |
 | pto.tcolargmax | TCOLARGMAX | tile+tensor | ✅ | ✅ | ✅ | ✅ | — |  |
@@ -259,6 +259,6 @@ lowering/compiler plumbing 使用的额外内部 op 未纳入，也不列 VPTO
 | **源码兼容 / 手动模式（1）** |  |  |  |  |  |  |  |  |
 | pto.tassign | TASSIGN | internal | ✅ | — | — | — | — | 失活 backend hook，不独立建 ST |
 
-**统计**：共 204 个 PTOAS 公共/兼容 op；pypto tile 前端 113 个，tensor 前端 75 个；
-同名 ST 覆盖 110 个（普通 ST 106，distributed ST 4）；无同名 ST 62 个
-（普通 52，distributed 10）；这 204 个中另有 32 个 op 不适合独立 ST。
+**统计**：共 204 个 PTOAS 公共/兼容 op；pypto tile 前端 116 个，tensor 前端 75 个；
+同名 ST 覆盖 112 个（普通 ST 108，distributed ST 4）；无同名 ST 60 个
+（普通 50，distributed 10）；这 204 个中另有 32 个 op 不适合独立 ST。

python/pypto/ir/op/tile_ops.py

@@ -935,6 +935,38 @@ def part_min(src0: Expr, src1: Expr, span: Span | None = None) -> Call:
     return _ir_core.create_op_call("tile.part_min", [src0, src1], {}, actual_span)
 
 
+def part_argmax(
+    src0: Expr,
+    src1: Expr,
+    src0_idx: Expr,
+    src1_idx: Expr,
+    span: Span | None = None,
+) -> Call:
+    """Partial element-wise maximum returning selected values and indices."""
+    actual_span = _get_span_or_capture(span)
+    return _ir_core.create_op_call("tile.part_argmax", [src0, src1, src0_idx, src1_idx], {}, actual_span)
+
+
+def part_argmin(
+    src0: Expr,
+    src1: Expr,
+    src0_idx: Expr,
+    src1_idx: Expr,
+    span: Span | None = None,
+) -> Call:
+    """Partial element-wise minimum returning selected values and indices."""
+    actual_span = _get_span_or_capture(span)
+    return _ir_core.create_op_call("tile.part_argmin", [src0, src1, src0_idx, src1_idx], {}, actual_span)
+
+
+def histogram(src: Expr, idx: Expr, byte: int = 1, span: Span | None = None) -> Call:
+    """A5 per-row cumulative histogram for one byte of a UINT16/UINT32 tile."""
+    if not isinstance(byte, int) or isinstance(byte, bool) or not 0 <= byte <= 3:
+        raise ValueError(f"histogram byte must be an integer in [0, 3], got {byte!r}")
+    actual_span = _get_span_or_capture(span)
+    return _ir_core.create_op_call("tile.histogram", [src, idx], {"byte": byte}, actual_span)
+
+
 def fmod(lhs: Expr, rhs: Expr, span: Span | None = None) -> Call:
     """Element-wise floating-point remainder of two tiles.
 

python/pypto/language/__init__.py

@@ -127,6 +127,7 @@ def scalar_func(x: pl.Scalar[pl.FP32]) -> pl.Scalar[pl.FP32]:
     gemv,
     gemv_acc,
     gemv_bias,
+    histogram,
     load,
     lrelu,
     matmul_bias,
@@ -138,6 +139,8 @@ def scalar_func(x: pl.Scalar[pl.FP32]) -> pl.Scalar[pl.FP32]:
     not_,
     or_,
     ors,
+    part_argmax,
+    part_argmin,
     prelu,
     relu,
     rem,
@@ -320,6 +323,9 @@ def scalar_func(x: pl.Scalar[pl.FP32]) -> pl.Scalar[pl.FP32]:
     "part_mul",
     "part_max",
     "part_min",
+    "part_argmax",
+    "part_argmin",
+    "histogram",
     "maximum",
     "exp",
     "log",

python/pypto/language/op/__init__.py

@@ -79,6 +79,7 @@
     gemv,
     gemv_acc,
     gemv_bias,
+    histogram,
     load,
     lrelu,
     matmul_bias,
@@ -91,6 +92,8 @@
     not_,
     or_,
     ors,
+    part_argmax,
+    part_argmin,
     prelu,
     relu,
     rem,
@@ -193,6 +196,9 @@
     "part_mul",
     "part_max",
     "part_min",
+    "part_argmax",
+    "part_argmin",
+    "histogram",
     "maximum",
     "min",
     "max",

python/pypto/language/op/tile_ops.py

@@ -120,6 +120,9 @@
     "part_mul",
     "part_max",
     "part_min",
+    "part_argmax",
+    "part_argmin",
+    "histogram",
     "fmod",
     "fmods",
     "and_",
@@ -2039,6 +2042,31 @@ def part_min(src0: Tile, src1: Tile) -> Tile:
     return Tile(expr=call_expr)
 
 
+def part_argmax(src0: Tile, src1: Tile, src0_idx: Tile, src1_idx: Tile) -> tuple[Tile, Tile]:
+    """Partial element-wise maximum returning ``(value, index)`` tiles."""
+    call_expr = _ir_ops.part_argmax(src0.unwrap(), src1.unwrap(), src0_idx.unwrap(), src1_idx.unwrap())
+    span = call_expr.span
+    return (
+        Tile(expr=_ir_core.TupleGetItemExpr(call_expr, 0, span)),
+        Tile(expr=_ir_core.TupleGetItemExpr(call_expr, 1, span)),
+    )
+
+
+def part_argmin(src0: Tile, src1: Tile, src0_idx: Tile, src1_idx: Tile) -> tuple[Tile, Tile]:
+    """Partial element-wise minimum returning ``(value, index)`` tiles."""
+    call_expr = _ir_ops.part_argmin(src0.unwrap(), src1.unwrap(), src0_idx.unwrap(), src1_idx.unwrap())
+    span = call_expr.span
+    return (
+        Tile(expr=_ir_core.TupleGetItemExpr(call_expr, 0, span)),
+        Tile(expr=_ir_core.TupleGetItemExpr(call_expr, 1, span)),
+    )
+
+
+def histogram(src: Tile, idx: Tile, byte: int = 1) -> Tile:
+    """A5 per-row cumulative histogram for a selected source byte."""
+    return Tile(expr=_ir_ops.histogram(src.unwrap(), idx.unwrap(), byte=byte))
+
+
 def fmod(lhs: Tile, rhs: Tile) -> Tile:
     """Element-wise floating-point remainder of two tiles.
 

src/backend/common/pto_ops_elementwise.cpp

@@ -14,7 +14,10 @@
  * @brief PTO codegen registration for elementwise / compute tile ops.
  */
 
+#include <algorithm>
+#include <array>
 #include <cstddef>
+#include <memory>
 #include <optional>
 #include <sstream>
 #include <string>
@@ -31,6 +34,7 @@
 #include "pypto/ir/kind_traits.h"
 #include "pypto/ir/scalar_expr.h"
 #include "pypto/ir/tile_view_semantics.h"
+#include "pypto/ir/transforms/utils/memref_utils.h"
 #include "pypto/ir/type.h"
 #include "src/backend/common/pto_ops_internal.h"
 
@@ -348,6 +352,88 @@ static std::string MakePrintCodegenPTO(const std::string& pto_op_name, const Cal
   return "";
 }
 
+static std::string MakePartArgCodegenPTO(const std::string& pto_op_name, const CallPtr& op,
+                                         codegen::CodegenBase& codegen_base) {
+  auto& codegen = AsPto(codegen_base);
+  CHECK(op->args_.size() == 4) << op->op_->name_
+                               << " requires 4 arguments (src0, src1, src0_idx, src1_idx), but got "
+                               << op->args_.size();
+
+  auto tuple_var = codegen.GetCurrentResultVar();
+  INTERNAL_CHECK_SPAN(tuple_var, op->span_)
+      << "Internal error: " << op->op_->name_ << " codegen requires current_result_var";
+  auto element_vars = codegen.ResolveTupleResultElements(tuple_var, 2);
+  INTERNAL_CHECK_SPAN(element_vars[0] && element_vars[1], op->span_)
+      << "Internal error: " << op->op_->name_ << " requires two TupleGetItemExpr consumers";
+
+  std::array<std::shared_ptr<const ir::TileType>, 2> result_types;
+  for (size_t i = 0; i < result_types.size(); ++i) {
+    result_types[i] = ir::GetTileTypeWithMemRef(element_vars[i]->GetType());
+    INTERNAL_CHECK_SPAN(result_types[i], element_vars[i]->span_)
+        << "Internal error: " << op->op_->name_ << " result " << i
+        << " must have TileType with MemRef set by InitMemRef";
+    codegen.EmitAllocTileForVar(element_vars[i], result_types[i]);
+  }
+
+  std::array<std::string, 4> inputs;
+  std::array<std::string, 4> input_types;
+  for (size_t i = 0; i < inputs.size(); ++i) {
+    inputs[i] = codegen.GetExprAsCode(op->args_[i]);
+    input_types[i] = codegen.GetExprTypeAnnotation(op->args_[i]);
+    INTERNAL_CHECK_SPAN(!inputs[i].empty(), op->args_[i]->span_)
+        << "Internal error: " << op->op_->name_ << " input " << i << " requires an SSA value";
+  }
+  const std::string dst = codegen.GetVarName(element_vars[0]);
+  const std::string dst_idx = codegen.GetVarName(element_vars[1]);
+  const std::string dst_type = codegen.GetTileBufTypeStringFromTileType(result_types[0]);
+  const std::string dst_idx_type = codegen.GetTileBufTypeStringFromTileType(result_types[1]);
+  INTERNAL_CHECK_SPAN(!dst.empty() && !dst_idx.empty(), op->span_)
+      << "Internal error: " << op->op_->name_ << " requires value and index outputs";
+
+  std::ostringstream oss;
+  oss << pto_op_name << " ins(" << inputs[0] << ", " << inputs[1] << ", " << inputs[2] << ", " << inputs[3];
+  if (std::all_of(input_types.begin(), input_types.end(),
+                  [](const std::string& type) { return !type.empty(); })) {
+    oss << " : " << input_types[0] << ", " << input_types[1] << ", " << input_types[2] << ", "
+        << input_types[3];
+  }
+  oss << ") outs(" << dst << ", " << dst_idx;
+  if (!dst_type.empty() && !dst_idx_type.empty()) {
+    oss << " : " << dst_type << ", " << dst_idx_type;
+  }
+  oss << ")";
+  codegen.Emit(oss.str());
+  return "";
+}
+
+static std::string MakeHistogramCodegenPTO(const CallPtr& op, codegen::CodegenBase& codegen_base) {
+  auto& codegen = AsPto(codegen_base);
+  CHECK(op->args_.size() == 2) << "tile.histogram requires 2 arguments (src, idx), but got "
+                               << op->args_.size();
+  const int byte = op->GetKwarg<int>("byte", 1);
+  std::string src = codegen.GetExprAsCode(op->args_[0]);
+  std::string idx = codegen.GetExprAsCode(op->args_[1]);
+  std::string src_type = codegen.GetExprTypeAnnotation(op->args_[0]);
+  std::string idx_type = codegen.GetExprTypeAnnotation(op->args_[1]);
+  std::string dst = codegen.GetCurrentResultTarget();
+  std::string dst_type = codegen.GetCurrentResultTileBufTypeString();
+  INTERNAL_CHECK_SPAN(!src.empty() && !idx.empty() && !dst.empty(), op->span_)
+      << "Internal error: tile.histogram requires input and output SSA values";
+
+  std::ostringstream oss;
+  oss << "pto.thistogram ins(" << src << ", " << idx;
+  if (!src_type.empty() && !idx_type.empty()) {
+    oss << " : " << src_type << ", " << idx_type;
+  }
+  oss << ") outs(" << dst;
+  if (!dst_type.empty()) {
+    oss << " : " << dst_type;
+  }
+  oss << ") {byte = " << byte << " : i32}";
+  codegen.Emit(oss.str());
+  return "";
+}
+
 struct SimpleOpEntry {
   const char* op_name;
   const char* pto_op_name;
@@ -497,6 +583,22 @@ void RegisterElementwiseOps(Backend& backend, const std::unordered_set<std::stri
     backend.RegisterOp(op_name).f_codegen(std::move(fn));
   };
 
+  auto register_part_arg = [&](const char* op_name, const char* pto_op_name) {
+    if (exclude_ops.count(op_name) > 0) return;
+    auto entry = backend.RegisterOp(op_name);
+    entry.f_codegen([pto_op = std::string(pto_op_name)](const CallPtr& op, codegen::CodegenBase& codegen) {
+      return MakePartArgCodegenPTO(pto_op, op, codegen);
+    });
+    for (size_t i = 0; i < 4; ++i) {
+      entry.set_input_layout(i, ir::TileLayout::row_major);
+    }
+    entry.set_output_layout(ir::TileLayout::row_major);
+  };
+  register_part_arg("tile.part_argmax", "pto.tpartargmax");
+  register_part_arg("tile.part_argmin", "pto.tpartargmin");
+  reg("tile.histogram",
+      [](const CallPtr& op, codegen::CodegenBase& codegen) { return MakeHistogramCodegenPTO(op, codegen); });
+
   auto register_precision_op = [&](const char* op_name, const char* pto_op_name, size_t arity,
                                    const char* attr_kind) {
     if (exclude_ops.count(op_name) > 0) return;