Skip to content

fix: add pto.tfillpad_inplace; stop tfillpad from auto-lowering to TFILLPAD_INPLACE #473

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

Merged
zhangstevenunity merged 1 commit into from
Apr 15, 2026
Merged

fix: add pto.tfillpad_inplace; stop tfillpad from auto-lowering to TFILLPAD_INPLACE #473

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
28 changes: 28 additions & 0 deletions include/PTO/IR/PTOOps.td
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3640,6 +3640,34 @@ def TFillPadExpandOp : PTO_TOp<"tfillpad_expand", [
}];
}

def TFillPadInplaceOp : PTO_TOp<"tfillpad_inplace", [
PTO_DpsInitOpInterface,
OpPipeInterface,
DeclareOpInterfaceMethods<MemoryEffectsOpInterface>
]> {
let summary = "In-place fill padding on shared backing storage: src supplies valid-region bounds, dst supplies target pad bounds (tilebuf, DPS)";

let arguments = (ins
PTODpsType:$src,
PTODpsType:$dst
);

let results = (outs);

let hasVerifier = 1;

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

let extraClassDeclaration = [{
::mlir::pto::PIPE getPipe() { return ::mlir::pto::PIPE::PIPE_V; }
::mlir::MutableOperandRange getDpsInitsMutable() { return getDstMutable(); }
}];
}

def TGatherOp : PTO_TOp<"tgather", [
AttrSizedOperandSegments,
PTO_DpsInitOpInterface,
Expand Down
6 changes: 6 additions & 0 deletions lib/PTO/IR/PTO.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4537,6 +4537,11 @@ mlir::LogicalResult mlir::pto::TFillPadExpandOp::verify() {
/*allowDstExpand=*/true, "tfillpad_expand");
}

mlir::LogicalResult mlir::pto::TFillPadInplaceOp::verify() {
return verifyTFillPadLike(getOperation(), getSrc().getType(), getDst().getType(),
/*allowDstExpand=*/false, "tfillpad_inplace");
}


llvm::LogicalResult mlir::pto::TGatherOp::verify() {
auto isSupportedGatherElemTypeA5Index = [&](Type ty) -> bool {
Expand Down Expand Up @@ -9788,6 +9793,7 @@ void TInsertFPOp::getEffects(

PTO_DEFINE_UNARY_EFFECTS(TFillPadOp, getSrcMutable(), getDstMutable())
PTO_DEFINE_UNARY_EFFECTS(TFillPadExpandOp, getSrcMutable(), getDstMutable())
PTO_DEFINE_UNARY_EFFECTS(TFillPadInplaceOp, getSrcMutable(), getDstMutable())

void TGatherOp::getEffects(
SmallVectorImpl<SideEffects::EffectInstance<MemoryEffects::Effect>> &effects) {
Expand Down
32 changes: 27 additions & 5 deletions lib/PTO/Transforms/PTOToEmitC.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6696,12 +6696,33 @@ struct PTOFillPadToEmitC : public OpConversionPattern<pto::TFillPadOp> {

Value src = peelUnrealized(adaptor.getSrc());
Value dst = peelUnrealized(adaptor.getDst());
llvm::StringRef callee =
(op.getSrc() == op.getDst() || src == dst) ? "TFILLPAD_INPLACE"
: "TFILLPAD";

rewriter.create<emitc::CallOpaqueOp>(
loc, TypeRange{}, callee,
loc, TypeRange{}, "TFILLPAD",
/*args=*/ArrayAttr{}, /*templateArgs=*/ArrayAttr{},
/*operands=*/ValueRange{dst, src});

rewriter.eraseOp(op);
return success();
Comment on lines 6700 to +6706
Copy link
Copy Markdown

@gemini-code-assist gemini-code-assist bot Apr 13, 2026

Choose a reason for hiding this comment

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

medium

Using replaceOpWithNewOp is more idiomatic in MLIR conversion patterns than manually creating an operation and then erasing the original one. It ensures that the conversion driver is correctly notified of the replacement.

Suggested change
rewriter.create<emitc::CallOpaqueOp>(
loc, TypeRange{}, callee,
loc, TypeRange{}, "TFILLPAD",
/*args=*/ArrayAttr{}, /*templateArgs=*/ArrayAttr{},
/*operands=*/ValueRange{dst, src});
rewriter.eraseOp(op);
return success();
rewriter.replaceOpWithNewOp<emitc::CallOpaqueOp>(
op, TypeRange{}, "TFILLPAD",
/*args=*/ArrayAttr{}, /*templateArgs=*/ArrayAttr{},
/*operands=*/ValueRange{dst, src});
return success();

}
};
//===----------------------------------------------------------------------===//
// pto.tfillpad_inplace lowering -> TFILLPAD_INPLACE(dst, src)
//===----------------------------------------------------------------------===//

struct PTOFillPadInplaceToEmitC
: public OpConversionPattern<pto::TFillPadInplaceOp> {
using OpConversionPattern<pto::TFillPadInplaceOp>::OpConversionPattern;

LogicalResult matchAndRewrite(pto::TFillPadInplaceOp op, OpAdaptor adaptor,
ConversionPatternRewriter &rewriter) const override {
auto loc = op.getLoc();

Value src = peelUnrealized(adaptor.getSrc());
Value dst = peelUnrealized(adaptor.getDst());

rewriter.create<emitc::CallOpaqueOp>(
loc, TypeRange{}, "TFILLPAD_INPLACE",
/*args=*/ArrayAttr{}, /*templateArgs=*/ArrayAttr{},
/*operands=*/ValueRange{dst, src});

Comment on lines +6717 to 6728
Copy link
Copy Markdown

@gemini-code-assist gemini-code-assist bot Apr 13, 2026

Choose a reason for hiding this comment

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

critical

The matchAndRewrite implementation for PTOFillPadInplaceToEmitC is incomplete. It is missing the return statement and the closing braces for both the function and the struct. This will cause a compilation error. Additionally, using replaceOpWithNewOp is preferred.

  LogicalResult matchAndRewrite(pto::TFillPadInplaceOp op, OpAdaptor adaptor,
                                ConversionPatternRewriter &rewriter) const override {
    auto loc = op.getLoc();

    Value src = peelUnrealized(adaptor.getSrc());
    Value dst = peelUnrealized(adaptor.getDst());

    rewriter.replaceOpWithNewOp<emitc::CallOpaqueOp>(
        op, TypeRange{}, "TFILLPAD_INPLACE",
        /*args=*/ArrayAttr{}, /*templateArgs=*/ArrayAttr{},
        /*operands=*/ValueRange{dst, src});

    return success();
  }
};

Expand Down Expand Up @@ -9852,7 +9873,8 @@ static void populatePTOToEmitCPatterns(RewritePatternSet &patterns,
patterns.add<PTOPartAddToEmitC>(typeConverter, ctx);
patterns.add<PTOExtractToEmitC, PTOExtractFPToEmitC, PTOInsertToEmitC,
PTOInsertFPToEmitC>(typeConverter, ctx);
patterns.add<PTOFillPadToEmitC, PTOFillPadExpandToEmitC>(typeConverter, ctx);
patterns.add<PTOFillPadToEmitC, PTOFillPadInplaceToEmitC, PTOFillPadExpandToEmitC>(
typeConverter, ctx);
patterns.add<PTOGatherToEmitC>(typeConverter, ctx);
patterns.add<PTOGatherbToEmitC>(typeConverter, ctx);
patterns.add<PTOMovFPToEmitC>(typeConverter, ctx);
Expand Down
26 changes: 26 additions & 0 deletions lib/PTO/Transforms/PTOViewToMemref.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2314,6 +2314,32 @@ struct PTOViewToMemrefPass
dst);
}

SmallVector<mlir::pto::TFillPadInplaceOp, 8> fillpadInplaceOps;
func.walk(
[&](mlir::pto::TFillPadInplaceOp op) { fillpadInplaceOps.push_back(op); });

for (auto op : fillpadInplaceOps) {
IRRewriter rewriter(ctx);
Comment on lines +2321 to +2322
Copy link
Copy Markdown

@gemini-code-assist gemini-code-assist bot Apr 13, 2026

Choose a reason for hiding this comment

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

medium

Instantiating an IRRewriter inside a loop is inefficient. It is better to create the rewriter once outside the loop and update its insertion point for each operation.

Suggested change
for (auto op : fillpadInplaceOps) {
IRRewriter rewriter(ctx);
IRRewriter rewriter(ctx);
for (auto op : fillpadInplaceOps) {

Copy link
Copy Markdown
Contributor Author

@FangRui0 FangRui0 Apr 13, 2026

Choose a reason for hiding this comment

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

fixed

rewriter.setInsertionPoint(op);

Value src = op.getSrc();
Value dst = op.getDst();

auto srcTy = dyn_cast<MemRefType>(src.getType());
auto dstTy = dyn_cast<MemRefType>(dst.getType());
if (!srcTy || !dstTy) {
op.emitError("ins/outs are not memref yet");
signalPassFailure();
return;
}

rewriter.replaceOpWithNewOp<pto::TFillPadInplaceOp>(
op,
TypeRange{},
src,
dst);
}

// --- TSetValOp [Dst, Offset, Val] ---
// Lower tile-world scalar write to memref-world SETVAL DPS op.
SmallVector<mlir::pto::TSetValOp, 8> tsetvalops;
Expand Down
4 changes: 2 additions & 2 deletions test/basic/tfillpad_inplace_alias_lowering.pto
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3,8 +3,8 @@
module {
func.func @tfillpad_inplace_alias() {
%tile = pto.alloc_tile : !pto.tile_buf<loc=vec, dtype=f32, rows=32, cols=32, v_row=32, v_col=32, blayout=row_major, slayout=none_box, fractal=512, pad=1>
pto.tfillpad ins(%tile : !pto.tile_buf<loc=vec, dtype=f32, rows=32, cols=32, v_row=32, v_col=32, blayout=row_major, slayout=none_box, fractal=512, pad=1>)
outs(%tile : !pto.tile_buf<loc=vec, dtype=f32, rows=32, cols=32, v_row=32, v_col=32, blayout=row_major, slayout=none_box, fractal=512, pad=1>)
pto.tfillpad_inplace ins(%tile : !pto.tile_buf<loc=vec, dtype=f32, rows=32, cols=32, v_row=32, v_col=32, blayout=row_major, slayout=none_box, fractal=512, pad=1>)
outs(%tile : !pto.tile_buf<loc=vec, dtype=f32, rows=32, cols=32, v_row=32, v_col=32, blayout=row_major, slayout=none_box, fractal=512, pad=1>)
return
}
}
Expand Down
14 changes: 14 additions & 0 deletions test/basic/tfillpad_same_ssa_lowers_to_tfillpad.pto
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,14 @@
// RUN: ptoas %s | FileCheck %s

module {
func.func @tfillpad_same_ssa() {
%tile = pto.alloc_tile : !pto.tile_buf<loc=vec, dtype=f32, rows=32, cols=32, v_row=32, v_col=32, blayout=row_major, slayout=none_box, fractal=512, pad=1>
pto.tfillpad ins(%tile : !pto.tile_buf<loc=vec, dtype=f32, rows=32, cols=32, v_row=32, v_col=32, blayout=row_major, slayout=none_box, fractal=512, pad=1>)
outs(%tile : !pto.tile_buf<loc=vec, dtype=f32, rows=32, cols=32, v_row=32, v_col=32, blayout=row_major, slayout=none_box, fractal=512, pad=1>)
return
}
}

// CHECK-LABEL: AICORE void tfillpad_same_ssa(
// CHECK: TFILLPAD(
// CHECK-NOT: TFILLPAD_INPLACE(
2 changes: 1 addition & 1 deletion test/samples/Fillpad/fillpad_inplace.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -49,7 +49,7 @@ def build():

tile = pto.AllocTileOp(tile_ty).result
pto.TLoadOp(None, sv0, tile)
pto.TFillPadOp(tile, tile)
pto.TFillPadInplaceOp(tile, tile)
pto.TStoreOp(None, tile, sv1)
func.ReturnOp([])

Expand Down
4 changes: 2 additions & 2 deletions test/samples/runop.sh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -847,12 +847,12 @@ PY

if [[ "$base" == "fillpad_inplace" ]]; then
if ! grep -Fq "TFILLPAD_INPLACE(" "$cpp"; then
echo -e "${A}(${base}.py)\tFAIL\tmissing TFILLPAD_INPLACE() lowering for aliased pto.tfillpad"
echo -e "${A}(${base}.py)\tFAIL\tmissing TFILLPAD_INPLACE() lowering for pto.tfillpad_inplace"
overall=1
continue
fi
if grep -Fq "TFILLPAD_EXPAND(" "$cpp"; then
echo -e "${A}(${base}.py)\tFAIL\tpto.tfillpad alias path should not lower via TFILLPAD_EXPAND()"
echo -e "${A}(${base}.py)\tFAIL\tpto.tfillpad_inplace should not lower via TFILLPAD_EXPAND()"
overall=1
continue
fi
Expand Down
Loading