Move handling of integer signedness to the backend conversions (#2597)

The function `getTypeForScalarType` currently takes an argument to
specify the signedness of integer types. This is leakage of backend
specific requirements into the torch dialect world. Because
`getTypeForScalarType` is a utility function for the torch dialect, it
should only produce types that match the sign conventions used by
PyTorch (regular integers are signed and unsigned integers are
unsigned).

This commit removes the signedness argument from
`getTypeForScalarType`, and moves the backend specific handling of
integer types to the backend code.

Author: ramiro050

include/torch-mlir/Dialect/Torch/Utils/Utils.h

@@ -26,9 +26,8 @@ bool getListConstructElements(Value v, SmallVectorImpl<Value> &elems);
 std::optional<int64_t> matchLegalConstantIndexIntoListOfSize(Value v,
                                                              int64_t length);
 torch_upstream::ScalarType getScalarTypeForType(Type type);
-FailureOr<Type> getTypeForScalarType(
-    MLIRContext *context, torch_upstream::ScalarType dtypeInt,
-    mlir::IntegerType::SignednessSemantics signedness = IntegerType::Signed);
+FailureOr<Type> getTypeForScalarType(MLIRContext *context,
+                                     torch_upstream::ScalarType dtypeInt);
 
 Type getTypeForTorchType(
     MLIRContext *context, Type type,

lib/Conversion/TorchToLinalg/TensorConstructors.cpp

@@ -127,9 +127,9 @@ class ConvertConstantTensorAllocOp : public OpConversionPattern<OpTy> {
       if (!matchPattern(op.getDtype(), m_TorchConstantInt(&dtypeInt)))
         return rewriter.notifyMatchFailure(
             op, "unimplemented: dtype must be a constant integer or none");
-      FailureOr<Type> maybeResultElementType = getTypeForScalarType(
-          op->getContext(), (torch_upstream::ScalarType)dtypeInt,
-          IntegerType::Signless);
+      FailureOr<Type> maybeResultElementType =
+          torch_to_linalg::getBackendTypeForScalarType(
+              op->getContext(), (torch_upstream::ScalarType)dtypeInt);
       if (failed(maybeResultElementType)) {
         return rewriter.notifyMatchFailure(
             op, "unable to convert `dtypeInt` to builtin type");
@@ -233,9 +233,9 @@ class ConvertAtenEmptyMemoryFormatOp
       if (!matchPattern(op.getDtype(), m_TorchConstantInt(&dtypeInt)))
         return rewriter.notifyMatchFailure(
             op, "unimplemented: dtype must be a constant integer or none");
-      FailureOr<Type> maybeResultElementType = getTypeForScalarType(
-          op->getContext(), (torch_upstream::ScalarType)dtypeInt,
-          IntegerType::Signless);
+      FailureOr<Type> maybeResultElementType =
+          torch_to_linalg::getBackendTypeForScalarType(
+              op->getContext(), (torch_upstream::ScalarType)dtypeInt);
       if (failed(maybeResultElementType)) {
         return rewriter.notifyMatchFailure(
             op, "unable to convert `dtypeInt` to builtin type");

lib/Conversion/TorchToLinalg/Uncategorized.cpp

@@ -1057,9 +1057,9 @@ static Value createLinalgPayloadCalculationForElementwiseOp(
       atenToDtype.emitError("unimplemented: dtype must be a constant integer");
       return nullptr;
     }
-    FailureOr<Type> maybeResultElementType = getTypeForScalarType(
-        atenToDtype->getContext(), (torch_upstream::ScalarType)dtypeInt,
-        IntegerType::Signless);
+    FailureOr<Type> maybeResultElementType =
+        torch_to_linalg::getBackendTypeForScalarType(
+            atenToDtype->getContext(), (torch_upstream::ScalarType)dtypeInt);
     if (failed(maybeResultElementType)) {
       atenToDtype.emitError("unable to convert `dtypeInt` to builtin type");
       return nullptr;

lib/Conversion/TorchToLinalg/Utils.cpp

@@ -20,7 +20,6 @@
 #include "torch-mlir/Conversion/Utils/Utils.h"
 #include "torch-mlir/Dialect/Torch/IR/TorchDialect.h"
 #include "torch-mlir/Dialect/Torch/IR/TorchOps.h"
-#include "torch-mlir/Dialect/Torch/Utils/TorchUpstream.h"
 #include "torch-mlir/Dialect/Torch/Utils/Utils.h"
 
 using namespace mlir;
@@ -546,3 +545,18 @@ Value torch_to_linalg::convertTensorToElementType(OpBuilder &b, Location loc,
   return torch_to_linalg::createElementwiseLinalgGeneric(
       b, loc, {tensor}, elementType, dtypePromoteBody);
 }
+
+FailureOr<Type> torch_to_linalg::getBackendTypeForScalarType(
+    MLIRContext *context, torch_upstream::ScalarType dtypeInt) {
+  FailureOr<Type> maybeType =
+      getTypeForScalarType(context, (torch_upstream::ScalarType)dtypeInt);
+  if (failed(maybeType)) {
+    return failure();
+  }
+  Type type = *maybeType;
+  // The linalg-on-tensors backend currently expects integers to be signless.
+  if (auto intType = type.dyn_cast<IntegerType>()) {
+    type = IntegerType::get(context, intType.getWidth(), IntegerType::Signless);
+  }
+  return type;
+}

lib/Conversion/TorchToLinalg/Utils.h

@@ -8,6 +8,7 @@
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Transforms/DialectConversion.h"
+#include "torch-mlir/Dialect/Torch/Utils/TorchUpstream.h"
 
 namespace mlir {
 namespace torch {
@@ -88,6 +89,12 @@ Value removeSizeInformation(OpBuilder &b, Location loc, Value tensor);
 Value convertTensorToElementType(OpBuilder &b, Location loc, Value tensor,
                                  Type elementType);
 
+// Convert a scalar type to the corresponding builtin type in the
+// linalg-on-tensors backend.
+FailureOr<Type>
+getBackendTypeForScalarType(MLIRContext *context,
+                            torch_upstream::ScalarType dtypeInt);
+
 } // namespace torch_to_linalg
 } // namespace torch
 } // namespace mlir

lib/Conversion/TorchToStablehlo/Basic.cpp

@@ -1672,13 +1672,18 @@ LogicalResult ConvertAtenOp<AtenEmptyMemoryFormatOp>::matchAndRewrite(
       return rewriter.notifyMatchFailure(
           op, "unimplemented: dtype must be a constant integer or none");
     FailureOr<Type> maybeResultElementType = getTypeForScalarType(
-        op->getContext(), (torch_upstream::ScalarType)dtypeInt,
-        IntegerType::Signless);
+        op->getContext(), (torch_upstream::ScalarType)dtypeInt);
     if (failed(maybeResultElementType)) {
       return rewriter.notifyMatchFailure(
           op, "unable to convert `dtypeInt` to builtin type");
     }
     resultElementType = *maybeResultElementType;
+    // The stablehlo backend expects signed integers to be signless.
+    if (resultElementType.isSignedInteger()) {
+      resultElementType = IntegerType::get(
+          op->getContext(), resultElementType.getIntOrFloatBitWidth(),
+          IntegerType::Signless);
+    }
   }
 
   // Create an uninitialized tensor of `resultSize` shape.

lib/Dialect/Torch/Utils/Utils.cpp

@@ -85,17 +85,16 @@ Type Torch::getTypeForTorchType(
 
 FailureOr<Type>
 Torch::getTypeForScalarType(MLIRContext *context,
-                            torch_upstream::ScalarType dtypeInt,
-                            mlir::IntegerType::SignednessSemantics signedness) {
+                            torch_upstream::ScalarType dtypeInt) {
   switch (dtypeInt) {
   case torch_upstream::ScalarType::Float:
     return Float32Type::get(context);
   case torch_upstream::ScalarType::Double:
     return Float64Type::get(context);
   case torch_upstream::ScalarType::Long:
-    return IntegerType::get(context, 64, signedness);
+    return IntegerType::get(context, 64, mlir::IntegerType::Signed);
   case torch_upstream::ScalarType::Int:
-    return IntegerType::get(context, 32, signedness);
+    return IntegerType::get(context, 32, mlir::IntegerType::Signed);
   case torch_upstream::ScalarType::Bool:
     return IntegerType::get(context, 1);
   case torch_upstream::ScalarType::BFloat16:
@@ -105,7 +104,7 @@ Torch::getTypeForScalarType(MLIRContext *context,
   case torch_upstream::ScalarType::Byte:
     return mlir::IntegerType::get(context, 8, mlir::IntegerType::Unsigned);
   case torch_upstream::ScalarType::Char:
-    return mlir::IntegerType::get(context, 8, signedness);
+    return mlir::IntegerType::get(context, 8, mlir::IntegerType::Signed);
   case torch_upstream::ScalarType::ComplexHalf:
     return mlir::ComplexType::get(Float16Type::get(context));
   case torch_upstream::ScalarType::ComplexFloat:

projects/pt1/python/torch_mlir_e2e_test/test_suite/constant_alloc.py

@@ -451,6 +451,25 @@ def EmptyModule_int(module, tu: TestUtils):
     module.forward()
 
 
+class EmptyUInt8Module(torch.nn.Module):
+
+    def __init__(self):
+        super().__init__()
+
+    @export
+    @annotate_args([
+        None,
+    ])
+    def forward(self):
+        empty = torch.ops.aten.empty([1], dtype=torch.uint8)
+        return torch.ops.aten.zeros_like(empty).to(torch.int8)
+
+
+@register_test_case(module_factory=lambda: EmptyUInt8Module())
+def EmptyModule_uint8(module, tu: TestUtils):
+    module.forward()
+
+
 class EmptyFloatModule(torch.nn.Module):
 
     def __init__(self):