Move handling of integer signedness to the backend conversions

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):