[TORCH] Modified fx_importer to support hop_while_loop (#4338)

This PR adds support for emitting graphs for Pytorch HOPs, beginning
with `torch._higher_order_ops.while_loop`.

The proposed change is to modify the `import_program` to call function
`_import_all_child_modules`, which recursively imports the stateless
graph for all the children modules.

Since HOP operator graphs are stateless graphs with no mutation, it is
correct to import them as stateless graphs, although the method
`import_stateless_graph` is marked as "deprecated".

---------

Signed-off-by: Keshav Vinayak Jha <[email protected]>

Author: keshavvinayak01

lib/Conversion/TorchToSCF/TorchToSCF.cpp

@@ -150,6 +150,10 @@ class ConvertTorchPrimLoopWhileLikeOp : public OpConversionPattern<PrimLoopOp> {
           targetType = Torch::IntType::get(op->getContext());
         torchArg = typeConverter->materializeSourceConversion(
             rewriter, scfWhileOp.getLoc(), targetType, {to});
+      } else if (auto tty = dyn_cast<RankedTensorType>(targetType)) {
+        targetType = op.getIterArgsInit()[barg.index()].getType();
+        torchArg = typeConverter->materializeSourceConversion(
+            rewriter, scfWhileOp.getLoc(), targetType, {to});
       }
       if (!torchArg)
         return rewriter.notifyMatchFailure(op,
@@ -173,14 +177,6 @@ class ConvertTorchPrimLoopWhileLikeOp : public OpConversionPattern<PrimLoopOp> {
                                              "unsupported type of the operand");
         loopConditionIterArgs.push_back(shouldContinue);
         for (auto torchArg : primLoopConditionOp.getIterArgs()) {
-          Type torchType = torchArg.getType();
-
-          // If the argument is a torch tensor, directly add it in the list of
-          // iter args.
-          if (isa<Torch::BaseTensorType>(torchType)) {
-            loopConditionIterArgs.push_back(torchArg);
-            continue;
-          }
           Value arg = typeConverter->materializeTargetConversion(
               rewriter, scfWhileOp->getLoc(),
               typeConverter->convertType(torchArg.getType()), {torchArg});

lib/Dialect/Torch/Transforms/Passes.cpp

@@ -70,6 +70,9 @@ void mlir::torch::Torch::createTorchScriptModuleToTorchBackendPipeline(
 
 void mlir::torch::Torch::createTorchDynamoExportToTorchBackendPipeline(
     OpPassManager &pm, const TorchLoweringPipelineOptions &options) {
+  // Inline func.call operations created by higher-order ops like while_loop
+  // to conform to the linalg-on-tensors backend contract.
+  pm.addPass(createInlinerPass());
   pm.addNestedPass<func::FuncOp>(
       createReduceOpVariantsPass(options.extraLibrary));
   pm.addNestedPass<func::FuncOp>(createCanonicalizerPass());

projects/pt1/e2e_testing/xfail_sets.py

@@ -246,6 +246,8 @@
     "IsFloatingPointInt_False",
     "TorchPrimLoopForLikeModule_basic",
     "TorchPrimLoopWhileLikeModule_basic",
+    # torch._dynamo.exc.BackendCompilerFailed: Unsupported op: get_attr
+    "TorchPrimLoopWhileLikeHOPModule_basic",
     "ScalarConstantTupleModule_basic",
     # END tests failing due to: empty graph in dynamo
     # ERROR due to: backend never runs because of empty frame
@@ -481,6 +483,7 @@
     "TensorToBoolZeroRank_basic",
     "TensorToBool_basic",
     "ThresholdBackward2dMixedModule_basic",
+    "TorchPrimLoopWhileLikeHOPModule_basic",  # Compilation error: failed to legalize operation 'func.call'
     "UnsafeViewCollapseDynamicWithAtenSizeIntModule_basic",
     "UpSampleNearest2dDynamicFactor_basic",
     "ViewCollapseDynamicWithAtenSizeIntModule_basic",
@@ -982,6 +985,8 @@
     "ElementwiseClampMinModule_bfloat16",
     "ElementwiseClampModule_bfloat16",
     "ElementwiseReluModule_bfloat16",
+    # Runtime error: failed to legalize operation 'torch.constant.int'
+    "TorchPrimLoopWhileLikeHOPModule_basic",
 }
 
 FX_IMPORTER_STABLEHLO_CRASHING_SET = {
@@ -2564,6 +2569,7 @@
 
 LTC_XFAIL_SET = {
     "TorchPrimLoopForLikeTensorArgModule_basic" "CollapseAllDimensionsModule_basic",
+    "TorchPrimLoopWhileLikeHOPModule_basic",
     "CollapseRank1DynamicModule_basic",
     "CollapseStaticModule_basic",
     "CollapsePartialDynamicModule_basic",
@@ -3253,6 +3259,8 @@
     "ToCopyWithDTypeModule_basic",
     "TorchPrimLoopForLikeModule_basic",
     "TorchPrimLoopWhileLikeModule_basic",
+    # RuntimeError: Detected that you are using FX to torch.jit.trace a dynamo-optimized function
+    "TorchPrimLoopWhileLikeHOPModule_basic",
     "TraceModule_basic",
     "TraceModule_empty",
     "TraceModule_nonsquare",
@@ -3952,6 +3960,8 @@
     "ThresholdBackward2dMixedModule_basic",
     "TorchPrimLoopForLikeModule_basic",
     "TorchPrimLoopWhileLikeModule_basic",
+    # Runtime error: failed to legalize operation 'torch.aten.Bool.Tensor'
+    "TorchPrimLoopWhileLikeHOPModule_basic",
     "TraceModule_empty",
     "TraceUnsignedIntModule_empty",
     "TransposedConv1dNegativePadding_basic",
@@ -5024,6 +5034,8 @@
     "ToDtypeFloatFromIntModule_basic",
     "TorchPrimLoopForLikeModule_basic",
     "TorchPrimLoopWhileLikeModule_basic",
+    # RuntimeError: Detected that you are using FX to torch.jit.trace a dynamo-optimized function
+    "TorchPrimLoopWhileLikeHOPModule_basic",
     "TraceModule_basic",
     "TraceModule_empty",
     "TraceModule_nonsquare",

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

@@ -10,6 +10,7 @@
 from torch_mlir_e2e_test.framework import TestUtils
 from torch_mlir_e2e_test.registry import register_test_case
 from torch_mlir_e2e_test.annotations import annotate_args, export
+from torch._higher_order_ops.while_loop import while_loop
 
 # ==============================================================================
 
@@ -78,3 +79,36 @@ def TorchPrimLoopForLikeTensorArgModule_basic(module, tu: TestUtils):
     x_test = torch.zeros([7, 9]).float()
 
     module.forward(x_test)
+
+
+# ==============================================================================
+
+
+class TorchPrimLoopWhileLikeHOPModule(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+
+    def body_fn(self, i, x):
+        return i + 1, x + 1
+
+    def cond_fn(self, i, x):
+        return i < 3
+
+    @export
+    @annotate_args(
+        [
+            None,
+            ([7, 9], torch.float32, True),
+        ]
+    )
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        i0 = torch.tensor(0)
+        out_i, out_x = while_loop(self.cond_fn, self.body_fn, (i0, x))
+        return out_i, out_x
+
+
+@register_test_case(module_factory=lambda: TorchPrimLoopWhileLikeHOPModule())
+def TorchPrimLoopWhileLikeHOPModule_basic(module, tu: TestUtils):
+    x_test = torch.zeros([7, 9]).float()
+
+    module.forward(x_test)