Skip to content

Commit

Permalink
add scater to tcp custom op
Browse files Browse the repository at this point in the history
  • Loading branch information
matthewfl committed Sep 5, 2024
1 parent 29062f8 commit ba92d9f
Showing 1 changed file with 23 additions and 0 deletions.
23 changes: 23 additions & 0 deletions lib/Conversion/TorchToTcp/TcpCustomOp.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -353,6 +353,28 @@ class ConvertAtenViewOp : public OpConversionPattern<AtenViewOp> {
}
};

class ConvertAtenSliceScatterOp
: public OpConversionPattern<AtenSliceScatterOp> {
using OpConversionPattern::OpConversionPattern;

LogicalResult
matchAndRewrite(AtenSliceScatterOp op, OpAdaptor adaptor,
ConversionPatternRewriter &rewriter) const override {
// this should really have some tcp op to reduce to. So going to CustomOp
// is more of a placeholder than a serious implementation
torch_to_tcp::TorchToTcpCustomOpConversionHelper helper{op, rewriter,
getTypeConverter()};
helper.addOperand("self", adaptor.getSelf());
helper.addOperand("src", adaptor.getSrc());
helper.addIntAttr("dim", op.getDim());
helper.addIntAttr("start", op.getStart());
helper.addIntAttr("end", op.getEnd());
helper.addIntAttr("step", op.getStep());

return helper.replace();
}
};

} // namespace

void torch_to_tcp::populateTcpCustomOpPatternsAndLegality(
Expand All @@ -373,6 +395,7 @@ void torch_to_tcp::populateTcpCustomOpPatternsAndLegality(
INSERT_ATEN_TO_TCP_CUSTOM_OP_PATTERN(AtenSortOp);
INSERT_ATEN_TO_TCP_CUSTOM_OP_PATTERN(AtenCumsumOp);
INSERT_ATEN_TO_TCP_CUSTOM_OP_PATTERN(AtenMinDimOp);
INSERT_ATEN_TO_TCP_CUSTOM_OP_PATTERN(AtenSliceScatterOp);
// AtenViewOp can still live after torch-to-tcp conversion
patterns.add<ConvertAtenViewOp>(typeConverter, patterns.getContext());
#undef INSERT_ATEN_TO_TCP_CUSTOM_OP_PATTERN
Expand Down

0 comments on commit ba92d9f

Please sign in to comment.