Skip to content

added lowering for torch._int_mm thereby supporting int8 matmul. #4573

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
NMJ0 wants to merge 3 commits into
base: main
Choose a base branch
from
+16 −12
Open

added lowering for torch._int_mm thereby supporting int8 matmul. #4573

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
2abb89d
added support for integer matrix multiplication
NMJ0 May 15, 2026
d9f0944
test
NMJ0 Jun 4, 2026
39ad3a5
test
NMJ0 Jun 4, 2026
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
24 changes: 14 additions & 10 deletions lib/Conversion/TorchToLinalg/Linear.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -71,12 +71,14 @@ static Value transposeValue(Location loc, Value value, ArrayRef<int64_t> perms,
return transpose;
}

class ConvertAtenMmOp : public OpConversionPattern<AtenMmOp> {
template <typename AtenOpT>
class ConvertAtenMmOp : public OpConversionPattern<AtenOpT> {
public:
using OpConversionPattern::OpConversionPattern;
using OpConversionPattern<AtenOpT>::OpConversionPattern;
LogicalResult
matchAndRewrite(AtenMmOp op, OpAdaptor adaptor,
ConversionPatternRewriter &rewriter) const override {
matchAndRewrite(
AtenOpT op, typename OpConversionPattern<AtenOpT>::OpAdaptor adaptor,
ConversionPatternRewriter &rewriter) const override {
Location loc = op->getLoc();
Value lhs = adaptor.getSelf();
Value rhs = adaptor.getMat2();
Expand Down Expand Up @@ -143,7 +145,7 @@ class ConvertAtenMmOp : public OpConversionPattern<AtenMmOp> {
}

TensorType resultType =
cast<TensorType>(getTypeConverter()->convertType(op.getType()));
cast<TensorType>(this->getTypeConverter()->convertType(op.getType()));
Type elementType = resultType.getElementType();
auto accumulatorDType =
getDefaultAccType(rewriter, lhsType.getElementType());
Expand All @@ -155,13 +157,13 @@ class ConvertAtenMmOp : public OpConversionPattern<AtenMmOp> {

Value matmul;
if (lhsZeroPoint) {
lhsZeroPoint = typeConverter->materializeTargetConversion(
lhsZeroPoint = this->typeConverter->materializeTargetConversion(
rewriter, loc,
getTypeConverter()->convertType(lhsZeroPoint.getType()),
this->getTypeConverter()->convertType(lhsZeroPoint.getType()),
lhsZeroPoint);
rhsZeroPoint = typeConverter->materializeTargetConversion(
rhsZeroPoint = this->typeConverter->materializeTargetConversion(
rewriter, loc,
getTypeConverter()->convertType(rhsZeroPoint.getType()),
this->getTypeConverter()->convertType(rhsZeroPoint.getType()),
rhsZeroPoint);
lhsZeroPoint = arith::TruncIOp::create(
rewriter, loc, rewriter.getI32Type(), lhsZeroPoint);
Expand Down Expand Up @@ -2524,7 +2526,9 @@ void mlir::torch::torch_to_linalg::populateLinearPatternsAndLegality(
ConversionTarget &target) {
MLIRContext *context = patterns.getContext();
target.addIllegalOp<AtenMmOp>();
patterns.add<ConvertAtenMmOp>(typeConverter, context);
patterns.add<ConvertAtenMmOp<AtenMmOp>>(typeConverter, context);
target.addIllegalOp<Aten_IntMmOp>();
patterns.add<ConvertAtenMmOp<Aten_IntMmOp>>(typeConverter, context);
target.addIllegalOp<AtenFlipOp>();
patterns.add<ConvertAtenFlipOp>(typeConverter, context);
target.addIllegalOp<AtenMatmulOp>();
Expand Down
2 changes: 1 addition & 1 deletion pytorch-requirements.txt
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,3 +1,3 @@
-f https://download.pytorch.org/whl/nightly/cpu/torch/
--pre
torch==2.12.0.dev20260323+cpu
torch==2.13.0.dev20260602+cpu
2 changes: 1 addition & 1 deletion torchvision-requirements.txt
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,3 +1,3 @@
-f https://download.pytorch.org/whl/nightly/cpu/torchvision/
--pre
torchvision==0.26.0.dev20260323+cpu
torchvision==0.28.0.dev20260603+cpu