Pass some decompose-complex-ops options in torch-to-iree (iree-org#19076)

Some problematic decompositions in `torch-decompose-complex-ops` are now
disabled by default in the `torch-to-iree` pipeline. Also allows turning
off all decompositions with an option.

---------

Signed-off-by: zjgarvey <[email protected]>

Author: zjgarvey

compiler/plugins/input/Torch/InputConversion/Passes.cpp

@@ -34,8 +34,6 @@ void createTorchToIREEPipeline(
   // backends. We do this first as it tends to involve pattern-matching against
   // constants, (e.g. dimensions which must be constant in a ranked programming
   // model) and those constants get somewhat obscured by TorchToArith.
-  llvm::ArrayRef<std::string> emptyArrayRef;
-
   // Dynamic shape bindings add a lot of structure to the IR which we prefer to
   // leverage and eliminate prior to any other activity, so do this first.
   pm.addNestedPass<func::FuncOp>(createBindSymbolicShapesPass());
@@ -51,8 +49,9 @@ void createTorchToIREEPipeline(
       torch::Torch::createReduceOpVariantsPass(llvm::StringRef()));
   pm.addNestedPass<func::FuncOp>(
       mlir::torch::TorchConversion::createConvertCustomQuantOpPass());
-  pm.addNestedPass<func::FuncOp>(
-      torch::Torch::createDecomposeComplexOpsPass(emptyArrayRef));
+  if (options.decompose)
+    pm.addNestedPass<func::FuncOp>(
+        torch::Torch::createDecomposeComplexOpsPass(BackendLegalOps::get()));
   pm.addNestedPass<func::FuncOp>(torch::Torch::createFuseQuantizedOpsPass());
   pm.addNestedPass<func::FuncOp>(createCanonicalizerPass());
   pm.addNestedPass<func::FuncOp>(torch::Torch::createScalarizeShapesPass());

compiler/plugins/input/Torch/InputConversion/Passes.h

@@ -13,11 +13,28 @@
 
 namespace mlir::iree_compiler::TorchInput {
 
+// The following is a hard-coded list of ops we don't want to decompose in the
+// torch dialect, since they have disadvantageous decompositons for the
+// torch-to-linalg path. For example, decomposing `aten.flatten.using_ints` to
+// `aten.view` simply destroys useful information about what kind of reshape is
+// being performed, and hinders our ability, in some cases, to lower this to a
+// collapse instead of a generic reshape.
+struct BackendLegalOps {
+  static const llvm::SmallVector<std::string> get() {
+    return {"aten.flatten.using_ints", "aten.unflatten.int",
+            "aten.adaptive_avg_pool1d", "aten.adaptive_avg_pool2d",
+            "aten.adaptive_max_pool1d"};
+  };
+};
+
 struct TorchToIREELoweringPipelineOptions
     : public PassPipelineOptions<TorchToIREELoweringPipelineOptions> {
   Option<bool> strictSymbolicShapes{
       *this, "strict-symbolic-shapes",
       llvm::cl::desc("Use strict symbolic shapes."), llvm::cl::init(true)};
+  Option<bool> decompose{*this, "decompose",
+                         llvm::cl::desc("Decompose complex torch operations."),
+                         llvm::cl::init(true)};
 };
 
 // Creates a pipeline that lowers from the torch backend contract to IREE.

compiler/plugins/input/Torch/PluginRegistration.cpp

@@ -23,12 +23,16 @@ namespace {
 
 struct TorchOptions {
   bool strictSymbolicShapes = true;
+  bool decompose = true;
   void bindOptions(OptionsBinder &binder) {
     static llvm::cl::OptionCategory category("Torch Input");
     binder.opt<bool>(
         "iree-torch-use-strict-symbolic-shapes", strictSymbolicShapes,
         llvm::cl::cat(category),
         llvm::cl::desc("Forces dynamic shapes to be treated as strict"));
+    binder.opt<bool>("iree-torch-decompose-complex-ops", decompose,
+                     llvm::cl::cat(category),
+                     llvm::cl::desc("Decompose complex torch operations."));
   }
 };
 
@@ -58,23 +62,17 @@ struct TorchSession
     if (typeMnemonic == "onnx") {
       // ONNX input is a pre-processing step to torch.
       mlir::torch::Torch::TorchLoweringPipelineOptions torchOnnxPipelineOptions;
-      // The `aten.flatten.using_ints` and `aten.unflatten.int` are added to the
-      // list of backend legal ops so that they are not decomposed into the
-      // `aten.view` op during the run of `DecomposeComplexOps` pass. The issue
-      // with this is that the `aten.view` op eventually lowers to
-      // `tensor.reshape` op while there exists a direct torch->linalg lowering
-      // for both the flatten/unflatten ops which lowers to
-      // `tensor.collapse_shape/expand_shape` op, and this is a more preferred
-      // path for the downstream pipeline.
-      torchOnnxPipelineOptions.backendLegalOps = {"aten.flatten.using_ints",
-                                                  "aten.unflatten.int"};
+      torchOnnxPipelineOptions.decompose = options.decompose;
+      torchOnnxPipelineOptions.backendLegalOps =
+          TorchInput::BackendLegalOps::get();
       mlir::torch::Torch::createTorchOnnxToTorchBackendPipeline(
           passManager, torchOnnxPipelineOptions);
     }
 
     if (typeMnemonic == "torch" || typeMnemonic == "onnx") {
       TorchInput::TorchToIREELoweringPipelineOptions torchOptions;
       torchOptions.strictSymbolicShapes = options.strictSymbolicShapes;
+      torchOptions.decompose = options.decompose;
       TorchInput::createTorchToIREEPipeline(passManager, torchOptions);
       return true;
     }