Skip to content

Feature/a8w8 matmul dequant fusion #1920

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

Closed
vegetabledoww wants to merge 1 commit into from
Closed

Feature/a8w8 matmul dequant fusion #1920

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
2 changes: 2 additions & 0 deletions include/pypto/codegen/pto/pto_codegen.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -370,6 +370,8 @@ class PTOCodegen : public CodegenBase {
void SetCurrentResultBuf(const std::string& buf);
void RegisterTileBufType(const std::string& ssa_name, const std::string& type_string);
std::string GetSSATileBufType(const std::string& ssa_name) const;
void AliasTileVarToExistingBuf(const ir::VarPtr& var, const std::string& ssa_name,
const std::string& type_string = "");
struct SubviewMaterializationInfo {
std::string source_ssa;
std::string source_type;
Expand Down
2 changes: 1 addition & 1 deletion include/pypto/ir/transforms/utils/cross_core_pipe.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -80,7 +80,7 @@ CallPtr CreateImportPeerBuffer(const std::string& buffer_name, const std::string
// (otherwise PTOAS derives it from `dir_mask`).
CallPtr CreateInitializePipe(core_affinity::CoreSide side, int dir_mask, int slot_size_bytes,
const ExprPtr& c2v_consumer_buf, const ExprPtr& v2c_consumer_buf,
std::optional<int> slot_num, const Span& span);
std::optional<int> slot_num, int pipe_id, const Span& span);

void CollectCrossCorePipeMetadata(const std::vector<StmtPtr>& stmts, CrossCorePipeMetadata& metadata);
CrossCorePipeMetadata CollectDominatingPipeSetupMetadata(const std::vector<StmtPtr>& stmts);
Expand Down
34 changes: 28 additions & 6 deletions src/backend/common/pto_ops_common.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2221,8 +2221,10 @@ static void EmitLogicalTpushValidShapeRestore(const CallPtr& op, codegen::PTOCod
static std::string FormatFrontendPipeAttrs(const CallPtr& op, int split) {
std::ostringstream oss;
oss << "{";
if (op->HasKwarg("id")) {
const int id = op->GetKwarg<int>("id", 0);
const bool has_explicit_id = op->HasKwarg("id");
const int id = has_explicit_id ? op->GetKwarg<int>("id", 0) : (split != 0 ? split : 0);
// Keep an explicit id=0 distinguishable from the omitted-id default.
if (id != 0 || has_explicit_id) {
CHECK(id >= 0) << "Frontend pipe 'id' attribute must be non-negative, got " << id;
oss << "id = " << id << ", ";
}
Expand Down Expand Up @@ -3506,7 +3508,22 @@ void RegisterPTOOps(Backend& backend, const std::unordered_set<std::string>& exc
if (src_space.has_value() && dst_space.has_value() && *src_space == *dst_space) {
auto src_offset = As<ir::ConstInt>((*src_tile->memref_)->byte_offset_);
auto dst_offset = As<ir::ConstInt>((*dst_tile->memref_)->byte_offset_);
if (src_offset && dst_offset && src_offset->value_ == dst_offset->value_) {
const bool same_dtype = src_tile->dtype_ == dst_tile->dtype_;
const bool same_tile_view = ir::tile_view_semantics::GetEffectiveTileView(*src_tile) ==
ir::tile_view_semantics::GetEffectiveTileView(*dst_tile);
bool same_shape = src_tile->shape_.size() == dst_tile->shape_.size();
for (size_t i = 0; same_shape && i < src_tile->shape_.size(); ++i) {
same_shape = IsSameDimExpr(src_tile->shape_[i], dst_tile->shape_[i]);
}
if (*src_space == ir::MemorySpace::Acc && same_dtype && same_shape) {
// PTOAS has no Acc->Acc tmov form. Treat same-shaped Acc moves as
// SSA aliases; these are inserted as layout/placement repairs before
// consumers such as tile.store, and copying would be illegal anyway.
codegen.SetCurrentResultBuf(codegen.GetExprAsCode(op->args_[0]));
return std::string("");
}
if (src_offset && dst_offset && src_offset->value_ == dst_offset->value_ && same_dtype &&
same_tile_view) {
Comment on lines +3511 to +3526

@coderabbitai coderabbitai Bot Jul 1, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

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

🩺 Stability & Availability | 🟠 Major | ⚡ Quick win

🧩 Analysis chain

🏁 Script executed:

#!/bin/bash
set -euo pipefail

echo "=== candidate reuse / allocation files ==="
fd -i 'memory' src/ir/transforms src/backend || true
fd -i 'memref' src || true

for f in $(fd -i 'memory' src/ir/transforms src/backend); do
  echo "=== AST outline: $f ==="
  ast-grep outline "$f" --view expanded || true
done

echo "=== tile.move + reuse/coalescing clues ==="
rg -nC3 --type=cpp 'tile\.move|byte_offset_|memref_|coalesc|reuse|target_type|blayout|slayout|valid_shape' src

Repository: hw-native-sys/pypto

Length of output: 50376

🏁 Script executed:

#!/bin/bash
set -euo pipefail

echo "=== pto_ops_common outline ==="
ast-grep outline src/backend/common/pto_ops_common.cpp --view expanded | sed -n '1,220p'

echo "=== tile.move / memory-reuse related definitions ==="
rg -nC4 'REGISTER_OP\("tile\.move"\)|set_output_reuses_input|not_inplace_safe|set_output_memory_inherit_input|MemoryReuse|InitMemRef|ReshapeTypeWithMemory|tile_view_semantics::GetEffectiveTileView|same_space|same-address|same address' src/ir src/backend | sed -n '1,260p'

echo "=== pto_ops_common around the cited lines ==="
sed -n '3410,3465p' src/backend/common/pto_ops_common.cpp

echo "=== tile op deducers around move/transpose/assemble/store ==="
sed -n '430,520p' src/ir/op/tile_ops/transform.cpp
sed -n '680,725p' src/ir/op/tile_ops/transform.cpp

Repository: hw-native-sys/pypto

Length of output: 35556

🏁 Script executed:

#!/bin/bash
set -euo pipefail

echo "=== tile.move registration / type deduction ==="
rg -nC6 'REGISTER_OP\("tile\.move"\)|DeduceTileMoveType|tile\.move' src/ir/op src/backend/common | sed -n '1,240p'

echo "=== memory reuse / memref init references ==="
rg -nC4 'MemoryReuse|InitMemRef|set_output_reuses_input|output_memory_inherit_input|reuses_input|memref_' src/ir/transforms src/backend/common src/ir/op | sed -n '1,260p'

echo "=== infer_tile_memory_space_pass around layout rewriting ==="
sed -n '410,470p' src/ir/transforms/infer_tile_memory_space_pass.cpp

echo "=== any explicit move/layout non-aliasing contract mentions ==="
rg -nC3 'not_inplace_safe|cannot reuse|must not reuse|same-space same-address|tile\.move.*layout|layout.*tile\.move|dtype.*tile\.move' src | sed -n '1,220p'

Repository: hw-native-sys/pypto

Length of output: 50375

🏁 Script executed:

#!/bin/bash
set -euo pipefail

echo "=== memory_reuse_pass: same-space / alias rules ==="
sed -n '380,460p' src/ir/transforms/memory_reuse_pass.cpp
echo
sed -n '1288,1388p' src/ir/transforms/memory_reuse_pass.cpp
echo
sed -n '1538,1565p' src/ir/transforms/memory_reuse_pass.cpp

echo "=== tile.move type deduction details ==="
sed -n '252,304p' src/ir/op/tile_ops/memory.cpp

Repository: hw-native-sys/pypto

Length of output: 13958

Keep layout-changing tile.move off shared buffers

tile.move is still eligible for MemoryReuse retargeting, so a dtype/layout-changing move can end up on the source buffer. This no-op elision only covers identical dtype + effective TileView; otherwise codegen can still emit an unsupported same-space pto.tmov. Mark these moves non-aliasing or force distinct MemRefs.

🤖 Prompt for AI Agents 
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@src/backend/common/pto_ops_common.cpp` around lines 3432 - 3436, The
tile.move no-op elision in the shared-buffer reuse path still allows
dtype/layout-changing moves to be retargeted onto the source buffer, which can
leave an unsupported same-space pto.tmov. Update the reuse/aliasing handling
around the src_offset/dst_offset check in src/backend/common/pto_ops_common.cpp
so that tile.move operations with differing dtype or effective TileView are
treated as non-aliasing, or otherwise force distinct MemRefs when the move
changes layout. Use the existing same_dtype and same_tile_view guard as the
location to prevent MemoryReuse retargeting for non-identical tiles.

// Alias the destination to the source SSA value so downstream
// references use the source's defined buffer, not the destination's
// alloc_tile (which would be unwritten after eliding the tmov).
Expand Down Expand Up @@ -3620,9 +3637,14 @@ void RegisterPTOOps(Backend& backend, const std::unordered_set<std::string>& exc
reg("tile.load", [](const ir::CallPtr& op, codegen::CodegenBase& codegen) {
return MakeTileLoadCodegenPTO(op, codegen);
});
reg("tile.store", [](const ir::CallPtr& op, codegen::CodegenBase& codegen) {
return MakeTileStoreCodegenPTO(op, codegen);
});
if (exclude_ops.count("tile.store") == 0) {
auto reg_entry = backend.RegisterOp("tile.store");
reg_entry
.f_codegen([](const ir::CallPtr& op, codegen::CodegenBase& codegen) {
return MakeTileStoreCodegenPTO(op, codegen);
})
.set_input_layout(0, ir::TileLayout::row_major);
}
// Distributed N6 ops — cross-rank tile load + per-rank signal notify/wait +
// synchronous bulk get/put. See MakeRemoteLoadCodegenPTO /
// MakeNotifyCodegenPTO / MakeWaitCodegenPTO / MakeGetCodegenPTO /
Expand Down
17 changes: 17 additions & 0 deletions src/codegen/pto/pto_codegen.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1255,6 +1255,16 @@ void PTOCodegen::RegisterTileBufType(const std::string& ssa_name, const std::str
fs_.ssa_to_tile_buf_type[ssa_name] = type_string;
}

void PTOCodegen::AliasTileVarToExistingBuf(const ir::VarPtr& var, const std::string& ssa_name,
const std::string& type_string) {
INTERNAL_CHECK(var != nullptr) << "Internal error: cannot alias null tile var";
INTERNAL_CHECK(!ssa_name.empty()) << "Internal error: cannot alias tile var to empty SSA";
BindVarToMlir(var, ssa_name);
if (!type_string.empty()) {
RegisterTileBufType(ssa_name, type_string);
}
}

std::string PTOCodegen::GetSSATileBufType(const std::string& ssa_name) const {
auto it = fs_.ssa_to_tile_buf_type.find(ssa_name);
return it != fs_.ssa_to_tile_buf_type.end() ? it->second : std::string{};
Expand Down Expand Up @@ -1383,6 +1393,13 @@ void PTOCodegen::VisitStmt_(const AssignStmtPtr& op) {
const bool alias_scatter_result_to_input = ShouldAliasScatterResultToInput(op);
const bool alias_array_update_to_input = ShouldAliasArrayUpdateResultToInput(op);

if (As<ir::TupleGetItemExpr>(op->value_)) {
auto key = GetVarKey(op->var_);
if (fs_.var_to_mlir.find(key) != fs_.var_to_mlir.end()) {
return;
}
}

if (auto tile_type = ir::GetTileTypeWithMemRef(op->var_->GetType())) {
if (!is_set_validshape && !alias_scatter_result_to_input) {
EmitAllocTileForVar(op->var_, tile_type);
Expand Down
6 changes: 6 additions & 0 deletions src/ir/op/tile_ops/memory.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -398,6 +398,12 @@ TypePtr DeduceTileCreateTileType(const std::vector<ExprPtr>& args,
std::optional<MemorySpace> creation_space = std::nullopt;
if (flat_layout) {
creation_space = MemorySpace::Mat;
} else if (target_memory_opt.has_value() && *target_memory_opt == MemorySpace::Left) {
tile_view.blayout = TileLayout::col_major;
tile_view.slayout = TileLayout::row_major;
} else if (target_memory_opt.has_value() && *target_memory_opt == MemorySpace::Right) {
tile_view.blayout = TileLayout::row_major;
tile_view.slayout = TileLayout::col_major;
} else if (target_memory_opt.has_value() && *target_memory_opt == MemorySpace::Acc) {
tile_view.blayout = TileLayout::col_major;
tile_view.slayout = TileLayout::row_major;
Expand Down
16 changes: 16 additions & 0 deletions src/ir/transforms/convert_tensor_to_tile_ops_pass.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -348,6 +348,17 @@ class TypePropagatingMutator : public IRMutator {
return std::make_shared<IterArg>(op->name_hint_, new_init->GetType(), new_init, op->span_);
}

ExprPtr VisitExpr_(const TupleGetItemExprPtr& op) override {
auto tuple = VisitExpr(op->tuple_);
if (auto make_tuple = As<MakeTuple>(tuple)) {
if (op->index_ >= 0 && static_cast<size_t>(op->index_) < make_tuple->elements_.size()) {
return VisitExpr(make_tuple->elements_[static_cast<size_t>(op->index_)]);
}
}
if (tuple.get() == op->tuple_.get()) return op;
return std::make_shared<TupleGetItemExpr>(tuple, op->index_, op->span_);
}

/// Override ForStmt to update return_vars types to match iter_arg types.
StmtPtr VisitStmt_(const ForStmtPtr& op) override {
auto result = IRMutator::VisitStmt_(op);
Expand Down Expand Up @@ -576,6 +587,11 @@ class TensorToTileMutator : public TypePropagatingMutator {
// remapped vars that the result expression references.
auto new_result = VisitExpr(conv_result.result);

if (As<MakeTuple>(new_result)) {
var_remap_[op->var_.get()] = new_result;
return SeqStmts::Flatten(std::move(stmts), op->span_);
}

auto tile_name = MakeTileValueName(op->var_->name_hint_);
auto tile_var = std::make_shared<Var>(tile_name, new_result->GetType(), op->var_->span_);
stmts.push_back(std::make_shared<AssignStmt>(tile_var, new_result, op->span_));
Expand Down
90 changes: 90 additions & 0 deletions src/ir/transforms/convert_to_ssa_pass.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -819,6 +819,47 @@ class SSAConverter {
auto cond = SubstExpr(op->condition_);
auto before = cur_;

// Frontend constants such as ``if True:`` can wrap manual-dependency
// producers that are consumed later in the parent scope. If we keep the
// artificial IfStmt until SSA conversion, those producer variables appear
// branch-local and the verifier rejects their later uses. Fold constants
// here, including explicit return_vars: the kept branch's trailing yield is
// stripped and each return_var is rebound to its yielded SSA value.
if (auto const_cond = StaticBoolValue(cond)) {
if (!*const_cond && !op->else_body_.has_value()) {
cur_ = before;
return std::make_shared<const SeqStmts>(std::vector<StmtPtr>{}, op->span_);
}
auto kept_body = ConvertStmt(*const_cond ? op->then_body_ : *op->else_body_);
if (op->return_vars_.empty()) {
return kept_body;
}
auto stripped = StripTrailingYield(kept_body, op->return_vars_.size());
std::vector<StmtPtr> extra_assigns;
for (size_t i = 0; i < op->return_vars_.size(); ++i) {
auto rv_key = op->return_vars_[i].get();
auto yielded = stripped.yielded_values[i];
if (auto yielded_var = As<Var>(yielded)) {
cur_[rv_key] = yielded_var;
continue;
}
int v = NextVersion(rv_key);
auto nrv = std::make_shared<Var>(BuildAutoNamedVersion(rv_key->name_hint_, "rv", v),
op->return_vars_[i]->GetType(), op->return_vars_[i]->span_);
extra_assigns.push_back(std::make_shared<AssignStmt>(nrv, yielded, op->span_));
cur_[rv_key] = nrv;
}
if (extra_assigns.empty()) return stripped.body;
std::vector<StmtPtr> stmts;
if (auto seq = As<SeqStmts>(stripped.body)) {
stmts = seq->stmts_;
} else if (stripped.body) {
stmts.push_back(stripped.body);
}
stmts.insert(stmts.end(), extra_assigns.begin(), extra_assigns.end());
return SeqStmts::Flatten(std::move(stmts), op->span_);
}

// Convert then branch
auto new_then = ConvertStmt(op->then_body_);
auto then_ver = cur_;
Expand Down Expand Up @@ -1224,6 +1265,55 @@ class SSAConverter {
return SeqStmts::Flatten({s, yield}, span);
}

struct StrippedYield {
StmtPtr body;
std::vector<ExprPtr> yielded_values;
};

static StmtPtr EmptyBody(const Span& span) {
return std::make_shared<const SeqStmts>(std::vector<StmtPtr>{}, span);
}

static std::optional<bool> StaticBoolValue(const ExprPtr& e) {
if (auto b = As<ConstBool>(e)) return b->value_;
if (auto i = As<ConstInt>(e); i && i->dtype() == DataType::BOOL) return i->value_ != 0;
return std::nullopt;
}

static StrippedYield StripTrailingYield(const StmtPtr& s, size_t return_var_count) {
AssertNoMidBodyYield(s);
if (auto scope = As<RuntimeScopeStmt>(s)) {
auto copy = MutableCopy(scope);
auto stripped = StripTrailingYield(scope->body_, return_var_count);
copy->body_ = stripped.body ? stripped.body : EmptyBody(scope->body_->span_);
return {copy, std::move(stripped.yielded_values)};
}
if (auto scope = As<SplitAivScopeStmt>(s)) {
auto copy = MutableCopy(scope);
auto stripped = StripTrailingYield(scope->body_, return_var_count);
copy->body_ = stripped.body ? stripped.body : EmptyBody(scope->body_->span_);
return {copy, std::move(stripped.yielded_values)};
}
if (auto seq = As<SeqStmts>(s)) {
INTERNAL_CHECK_SPAN(!seq->stmts_.empty(), seq->span_)
<< "ConvertToSSA: IfStmt with return_vars must end with YieldStmt";
auto stmts = seq->stmts_;
auto stripped = StripTrailingYield(stmts.back(), return_var_count);
if (stripped.body) {
stmts.back() = stripped.body;
} else {
stmts.pop_back();
}
return {SeqStmts::Flatten(std::move(stmts), seq->span_), std::move(stripped.yielded_values)};
}
auto yield = As<YieldStmt>(s);
INTERNAL_CHECK_SPAN(yield, s->span_) << "ConvertToSSA: IfStmt with return_vars must end with YieldStmt";
INTERNAL_CHECK_SPAN(yield->value_.size() == return_var_count, yield->span_)
<< "ConvertToSSA: yielded value count " << yield->value_.size()
<< " does not match return_vars count " << return_var_count;
return {nullptr, yield->value_};
}

// ── State ──────────────────────────────────────────────────────────

std::unordered_map<const Var*, VarPtr> cur_; // var pointer → latest version
Expand Down
Loading