Add dtype functions for aten.ge.Tensor and aten.le.Tensor

ramiro050 · May 5, 2023 · 899d8bc · 899d8bc
1 parent f58f9c2
commit 899d8bc
Show file tree

Hide file tree

Showing 19 changed files with 4,139 additions and 2,560 deletions.
diff --git a/build_tools/python_deploy/build_linux_packages.sh b/build_tools/python_deploy/build_linux_packages.sh
@@ -288,7 +288,7 @@ function test_in_tree() {
   python -m e2e_testing.main --config=lazy_tensor_core -v
 
   echo ":::: Run TorchDynamo e2e integration tests"
-  python -m e2e_testing.main --config=torchdynamo -v --crashing_tests_to_not_attempt_to_run_and_a_bug_is_filed RandnDtypeDeviceModule_basic
+  python -m e2e_testing.main --config=torchdynamo -v
 }
 
 function setup_venv() {

diff --git a/e2e_testing/xfail_sets.py b/e2e_testing/xfail_sets.py
@@ -93,11 +93,6 @@
     "ElementwiseAddScalar_NumToTensorFloat_Module_basic",
     # ERROR: assert isinstance(e, FakeTensor)
     "RsubInt0d_NumToTensor_Module_basic",
-
-    # Dtype function transition failures
-    "MobilenetV3Module_basic",
-    "ResNet18Module_basic",
-    "ResNet18StaticModule_basic",
 }
 
 STABLEHLO_PASS_SET = {
@@ -706,7 +701,8 @@
     "FullLikeModuleInt2DStatic_basic",
     "FullModuleInt3D_basic",
     "FullModuleFloat2D_basic",
-    "RepeatModule_basic"
+    "RepeatModule_basic",
+    "ResNet18StaticModule_basic",
 }
 
 LTC_XFAIL_SET = {

diff --git a/lib/Conversion/TorchToLinalg/TensorConstructors.cpp b/lib/Conversion/TorchToLinalg/TensorConstructors.cpp
@@ -226,7 +226,8 @@ class ConvertAtenEmptyMemoryFormatOp
         typeConverter->convertType(op.getType()).cast<RankedTensorType>();
     Type resultElementType;
     if (op.getDtype().getType().isa<Torch::NoneType>()) {
-      resultElementType = resultType.getElementType();
+      resultElementType = getDefaultDtypeForTorchScalar(
+          Torch::FloatType::get(op->getContext()));
     } else {
       int64_t dtypeInt;
       if (!matchPattern(op.getDtype(), m_TorchConstantInt(&dtypeInt)))

diff --git a/lib/Dialect/Torch/Transforms/AbstractInterpLibrary.cpp b/lib/Dialect/Torch/Transforms/AbstractInterpLibrary.cpp
diff --git a/lib/Dialect/Torch/Transforms/DecomposeComplexOps.cpp b/lib/Dialect/Torch/Transforms/DecomposeComplexOps.cpp
@@ -3716,8 +3716,14 @@ class DecomposeAtenRandnGeneratorOp
   LogicalResult matchAndRewrite(AtenRandnGeneratorOp op,
                                 PatternRewriter &rewriter) const override {
     Location loc = op.getLoc();
-    Type resultType = op.getType();
+    auto resultType = op.getType().cast<BaseTensorType>();
+
+    if (!resultType.hasDtype()) {
+      return rewriter.notifyMatchFailure(
+          op, "expected result type to have a dtype");
+    }
 
+    Value dtype = getDtypeIntValueForType(rewriter, loc, resultType.getDtype());
     Value none = rewriter.create<ConstantNoneOp>(loc);
     Value low = rewriter.create<Torch::ConstantFloatOp>(
         loc, rewriter.getF64FloatAttr((double)0.0));
@@ -3729,11 +3735,13 @@ class DecomposeAtenRandnGeneratorOp
         loc, rewriter.getF64FloatAttr((double)(2.0 * 3.14159)));
 
     Value emptyTensorA = rewriter.create<AtenEmptyMemoryFormatOp>(
-        loc, resultType, op.getSize(), /*dtype=*/none, /*layout=*/op.getLayout(),
+        loc, resultType, op.getSize(), /*dtype=*/dtype,
+        /*layout=*/op.getLayout(),
         /*device=*/op.getDevice(), /*pin_memory=*/op.getPinMemory(),
         /*memory_format=*/none);
     Value emptyTensorB = rewriter.create<AtenEmptyMemoryFormatOp>(
-        loc, resultType, op.getSize(), /*dtype=*/none, /*layout=*/op.getLayout(),
+        loc, resultType, op.getSize(), /*dtype=*/dtype,
+        /*layout=*/op.getLayout(),
         /*device=*/op.getDevice(), /*pin_memory=*/op.getPinMemory(),
         /*memory_format=*/none);