[MLIR][TORCH] Fix tensor literal int elem type to be signless

The element type of tensor literal should be signless when converted to
builtin tensor types.

Author: cathyzhyi

e2e_testing/torchscript/basic.py

@@ -602,7 +602,7 @@ def forward(self, x):
 @register_test_case(module_factory=lambda: ContiguousModule())
 def ContiguousModule_basic(module, tu: TestUtils):
     module.forward(tu.rand(3, 1))
-    
+
 class TensorToInt(torch.nn.Module):
     def __init__(self):
         super().__init__()
@@ -699,7 +699,7 @@ def __init__(self):
 
     @export
     @annotate_args([
-        None,   
+        None,
         ([-1, -1], torch.float32, True),
         ([-1, -1], torch.float32, True),
         ([-1, -1], torch.float32, True),
@@ -718,7 +718,7 @@ def __init__(self):
 
     @export
     @annotate_args([
-        None,   
+        None,
         ([-1, -1], torch.float32, True),
         ([-1, -1], torch.float32, True),
         ([-1, -1], torch.float32, True),
@@ -739,7 +739,7 @@ def __init__(self):
 
     @export
     @annotate_args([
-        None,   
+        None,
     ])
 
     def forward(self):
@@ -756,7 +756,7 @@ def __init__(self):
 
     @export
     @annotate_args([
-        None,   
+        None,
     ])
 
     def forward(self):
@@ -963,3 +963,38 @@ def forward(self, lhs):
 def TModuleRank0_basic(module, tu: TestUtils):
     module.forward(torch.tensor(7, dtype=torch.float32))
 
+class TensorLiteralModule(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+        torch.manual_seed(0)
+        self.t = torch.randint(-5, 5, (2, 3))
+
+    @export
+    @annotate_args([
+        None,
+    ])
+    def forward(self):
+        return torch.add(self.t, self.t)
+
+@register_test_case(module_factory=lambda: TensorLiteralModule())
+def TensorLiteralModule_basic(module, tu: TestUtils):
+    module.forward()
+
+
+class TensorOpaqueLiteralModule(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+        torch.manual_seed(0)
+        self.t = torch.randint(-5, 5, (256, 1024))
+
+    @export
+    @annotate_args([
+        None,
+    ])
+    def forward(self):
+        return torch.add(self.t, self.t)
+
+@register_test_case(module_factory=lambda: TensorOpaqueLiteralModule())
+def TensorOpaqueLiteralModule_basic(module, tu: TestUtils):
+    module.forward()
+

lib/Conversion/TorchToStd/TorchToStd.cpp

@@ -87,6 +87,49 @@ class ConvertAtenGtIntOp : public OpConversionPattern<AtenGtIntOp> {
 };
 } // namespace
 
+// Tensors with integer types need to be converted to signless integer
+// element type. All tensors with element types other than integer can reuse
+// existing elements attribute.
+namespace {
+class ConvertTorchTensorLiteralOp
+    : public OpConversionPattern<ValueTensorLiteralOp> {
+public:
+  using OpConversionPattern<ValueTensorLiteralOp>::OpConversionPattern;
+  using OpAdaptor = ValueTensorLiteralOp::Adaptor;
+  LogicalResult
+  matchAndRewrite(ValueTensorLiteralOp op, OpAdaptor adaptor,
+                  ConversionPatternRewriter &rewriter) const override {
+    MLIRContext *context = op->getContext();
+    if (auto elements = op.valueAttr().dyn_cast<DenseIntElementsAttr>()) {
+      Type elemTy = op.valueAttr().getElementType();
+      unsigned bitWidth = elemTy.getIntOrFloatBitWidth();
+      Type builtinTensorElemTy = IntegerType::get(context, bitWidth);
+      rewriter.replaceOpWithNewOp<arith::ConstantOp>(
+          op, elements.mapValues(builtinTensorElemTy, [&](const APInt &v) {
+            return APInt(bitWidth, v.getSExtValue());
+          }));
+      return success();
+    }
+    if (auto elements = op.valueAttr().dyn_cast<OpaqueElementsAttr>()) {
+      if (auto type = elements.getType().dyn_cast<RankedTensorType>()) {
+        if (auto intType = type.getElementType().dyn_cast<IntegerType>()) {
+          Type builtinTensorElemTy =
+              IntegerType::get(context, intType.getIntOrFloatBitWidth());
+          auto shapedType =
+              RankedTensorType::get(type.getShape(), builtinTensorElemTy);
+          rewriter.replaceOpWithNewOp<arith::ConstantOp>(
+              op, OpaqueElementsAttr::get(elements.getDialect(), shapedType,
+                                          elements.getValue()));
+          return success();
+        }
+      }
+    }
+    rewriter.replaceOpWithNewOp<arith::ConstantOp>(op, op.valueAttr());
+    return success();
+  }
+};
+} // namespace
+
 namespace {
 template <typename OpTy>
 class ConvertTorchConstantOp : public OpConversionPattern<OpTy> {
@@ -133,8 +176,8 @@ class ConvertTorchToStd : public ConvertTorchToStdBase<ConvertTorchToStd> {
     target.addIllegalOp<AtenGtIntOp>();
     patterns.add<ConvertAtenGtIntOp>(typeConverter, context);
     target.addIllegalOp<ValueTensorLiteralOp>();
-    patterns.add<ConvertTorchConstantOp<ValueTensorLiteralOp>>(typeConverter,
-                                                               context);
+    patterns.add<ConvertTorchTensorLiteralOp>(typeConverter, context);
+
     target.addIllegalOp<ConstantBoolOp>();
     patterns.add<ConvertTorchConstantOp<ConstantBoolOp>>(typeConverter,
                                                          context);

lib/Dialect/TorchConversion/Transforms/VerifyLinalgOnTensorsBackendContract.cpp

@@ -64,6 +64,7 @@ class VerifyLinalgOnTensorsBackendContractPass
     // Tensor operations should go through linalg and the tensor dialect.
     target.addDynamicallyLegalDialect<linalg::LinalgDialect>(opHasLegalTypes);
     target.addDynamicallyLegalDialect<tensor::TensorDialect>(opHasLegalTypes);
+    target.addDynamicallyLegalDialect<AffineDialect>(opHasLegalTypes);
 
     // AssertOp is used to terminate the program for error guards.
     target.addLegalOp<AssertOp>();