[TorchToLinalg] Support aten.mm with mixed sign extensions (#4722)

This lowering emits a matmul in its generic form when mixed sign
extensions are used.

Alternative would be to shift the unsigned side to signed, add a
zeropoint of -128 and emit a linalg.quantized_matmul. However i think
generally we should not do sign adjustments at the level of torch-mlir
since its very much target dependent how the signs should be handled.
Additionally undoing shifting the sign is a lot more difficult then
adding the shifting afterwards if desired.

---------

Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>

Author: ziereis

lib/Conversion/TorchToLinalg/Linear.cpp

@@ -116,18 +116,32 @@ class ConvertAtenMmOp : public OpConversionPattern<AtenMmOp> {
           op, "unsupported: aten.mm with mixed quantization");
     }
 
+    bool isUnsignedLhs = torch_to_linalg::isUnsignedTorchType(lhsTorchType);
+    bool isUnsignedRhs = torch_to_linalg::isUnsignedTorchType(rhsTorchType);
+
+    auto lhsIntType = dyn_cast<mlir::IntegerType>(lhsType.getElementType());
+    auto rhsIntType = dyn_cast<mlir::IntegerType>(rhsType.getElementType());
+
+    bool bothInt = lhsIntType && rhsIntType;
+
+    // Mixing integer widths would require extending the narrower operand to the
+    // wider one before the contraction, which none of the paths below do.
+    if (bothInt && lhsIntType.getWidth() != rhsIntType.getWidth()) {
+      return rewriter.notifyMatchFailure(
+          op, "unsupported: aten.mm with mixed integer widths");
+    }
+
+    bool isMixedSignedness = bothInt && isUnsignedLhs != isUnsignedRhs;
+
     if (lhsTorchType.getDtype() != rhsTorchType.getDtype()) {
-      if (!lhsZeroPoint) {
+      if (!lhsZeroPoint && !isMixedSignedness) {
         return rewriter.notifyMatchFailure(
             op, "unsupported: aten.mm with different input element types");
       }
       // Allows quantized types to mismatch since they will be cast to the same
       // type.
     }
 
-    bool isUnsigned = torch_to_linalg::isUnsignedTorchType(lhsTorchType);
-    bool isUnsignedR = torch_to_linalg::isUnsignedTorchType(rhsTorchType);
-
     Value lhsDim0 = tensor::DimOp::create(rewriter, loc, lhs, 0);
     Value rhsDim1 = tensor::DimOp::create(rewriter, loc, rhs, 1);
 
@@ -171,15 +185,43 @@ class ConvertAtenMmOp : public OpConversionPattern<AtenMmOp> {
       // change uint8 quantization -> int8 quantization
       int64_t numBits =
           cast<mlir::IntegerType>(lhsType.getElementType()).getWidth();
-      signShift(rewriter, loc, lhs, lhsZeroPoint, isUnsigned, numBits);
+      signShift(rewriter, loc, lhs, lhsZeroPoint, isUnsignedLhs, numBits);
       numBits = cast<mlir::IntegerType>(rhsType.getElementType()).getWidth();
-      signShift(rewriter, loc, rhs, rhsZeroPoint, isUnsignedR, numBits);
+      signShift(rewriter, loc, rhs, rhsZeroPoint, isUnsignedRhs, numBits);
 
       matmul = linalg::QuantizedMatmulOp::create(
                    rewriter, loc, zeroFill.getType(),
                    ValueRange{lhs, rhs, lhsZeroPoint, rhsZeroPoint}, zeroFill)
                    .getResult(0);
-    } else if (isUnsigned) {
+    } else if (isMixedSignedness) {
+      // `linalg.matmul` extends both operands with the same type function, so
+      // the contraction is written as a generic that extends each operand
+      // according to its own signedness instead.
+      MLIRContext *context = op.getContext();
+      AffineExpr m, n, k;
+      bindDims(context, m, n, k);
+      SmallVector<AffineMap> indexingMaps = {
+          AffineMap::get(3, 0, {m, k}, context),
+          AffineMap::get(3, 0, {k, n}, context),
+          AffineMap::get(3, 0, {m, n}, context)};
+      SmallVector<utils::IteratorType> iteratorTypes = {
+          utils::IteratorType::parallel, utils::IteratorType::parallel,
+          utils::IteratorType::reduction};
+      matmul =
+          linalg::GenericOp::create(
+              rewriter, loc, zeroFill.getType(), ValueRange{lhs, rhs}, zeroFill,
+              indexingMaps, iteratorTypes,
+              [&](OpBuilder &b, Location loc, ValueRange args) {
+                Value lhsElem = convertScalarToDtype(
+                    b, loc, args[0], elementType, lhsTorchType.getDtype());
+                Value rhsElem = convertScalarToDtype(
+                    b, loc, args[1], elementType, rhsTorchType.getDtype());
+                Value product = arith::MulIOp::create(b, loc, lhsElem, rhsElem);
+                Value sum = arith::AddIOp::create(b, loc, args[2], product);
+                linalg::YieldOp::create(b, loc, sum);
+              })
+              .getResult(0);
+    } else if (isUnsignedLhs && isUnsignedRhs) {
       auto matmulOp = linalg::MatmulOp::create(
           rewriter, loc, zeroFill.getType(), ValueRange{lhs, rhs}, zeroFill);
       matmulOp.setCast(linalg::TypeFn::cast_unsigned);

projects/pt1/e2e_testing/xfail_sets.py

@@ -34,8 +34,6 @@
     "ElementwiseClampInt16Module_basic",
     # TODO: The values are extremely close to the golden values, but the test fails because of strict rtol/atol.
     "AtenInstanceNormModuleFp16_basic",
-    # mixed uint8 x int8 aten.mm has no lowering yet
-    "AtenIntMMMixedSigni8_basic",
     # unimplemented lowering torch -> linalg for torchvision.deform_conv2d
     # this is added to check the torch.onnx.export -> import_onnx -> torch path
     "DeformConv2D_basic",
@@ -417,8 +415,6 @@
     "AtenFloatScalarModule_basic",
     # TODO: The values are extremely close to the golden values, but the test fails because of strict rtol/atol.
     "AtenInstanceNormModuleFp16_basic",
-    # mixed uint8 x int8 aten.mm has no lowering yet
-    "AtenIntMMMixedSigni8_basic",
     "AtenIntBoolOpConstFalseModule_basic",
     "AtenIntBoolOpConstTrueModule_basic",
     "AtenIntBoolOpModule_basic",

test/Conversion/TorchToLinalg/basic.mlir

@@ -201,6 +201,35 @@ func.func @torch.aten.mm$basic_unsigned(%arg0: !torch.vtensor<[?,?],ui32>, %arg1
 
 // -----
 
+// CHECK-LABEL: func.func @torch.aten.mm$mixed_signedness_lhs_unsigned(
+// CHECK: linalg.generic
+// CHECK-SAME: iterator_types = ["parallel", "parallel", "reduction"]
+// CHECK-SAME: ins(%{{.*}}, %{{.*}} : tensor<3x4xi8>, tensor<4x3xi8>)
+// CHECK: ^bb0(%[[LHS:.*]]: i8, %[[RHS:.*]]: i8, %[[ACC:.*]]: i32):
+// CHECK-DAG: %[[LHS_EXT:.*]] = arith.extui %[[LHS]] : i8 to i32
+// CHECK-DAG: %[[RHS_EXT:.*]] = arith.extsi %[[RHS]] : i8 to i32
+// CHECK: %[[MUL:.*]] = arith.muli %[[LHS_EXT]], %[[RHS_EXT]] : i32
+// CHECK: %[[SUM:.*]] = arith.addi %[[ACC]], %[[MUL]] : i32
+// CHECK: linalg.yield %[[SUM]] : i32
+func.func @torch.aten.mm$mixed_signedness_lhs_unsigned(%arg0: !torch.vtensor<[3,4],ui8>, %arg1: !torch.vtensor<[4,3],si8>) -> !torch.vtensor<[3,3],si32> {
+  %0 = torch.aten.mm %arg0, %arg1 : !torch.vtensor<[3,4],ui8>, !torch.vtensor<[4,3],si8> -> !torch.vtensor<[3,3],si32>
+  return %0 : !torch.vtensor<[3,3],si32>
+}
+
+// -----
+
+// CHECK-LABEL: func.func @torch.aten.mm$mixed_signedness_rhs_unsigned(
+// CHECK: ^bb0(%[[LHS:.*]]: i8, %[[RHS:.*]]: i8, %[[ACC:.*]]: i32):
+// CHECK-DAG: %[[LHS_EXT:.*]] = arith.extsi %[[LHS]] : i8 to i32
+// CHECK-DAG: %[[RHS_EXT:.*]] = arith.extui %[[RHS]] : i8 to i32
+// CHECK: arith.muli %[[LHS_EXT]], %[[RHS_EXT]] : i32
+func.func @torch.aten.mm$mixed_signedness_rhs_unsigned(%arg0: !torch.vtensor<[3,4],si8>, %arg1: !torch.vtensor<[4,3],ui8>) -> !torch.vtensor<[3,3],si32> {
+  %0 = torch.aten.mm %arg0, %arg1 : !torch.vtensor<[3,4],si8>, !torch.vtensor<[4,3],ui8> -> !torch.vtensor<[3,3],si32>
+  return %0 : !torch.vtensor<[3,3],si32>
+}
+
+// -----
+
 // If the operands are missing dtype, we cannot lower it.
 func.func @torch.aten.mm$no_convert$missing_dtype(%arg0: !torch.vtensor, %arg1: !torch.vtensor) -> !torch.vtensor {
   // expected-error@+1 {{failed to legalize}}