[compiler] Move StableHLO partitioning attribute handling to StablehloToPlan pass

Move the logic for reading StableHLO partitioning attributes (mhlo.num_replicas
and mhlo.num_partitions) and setting the executor process grid shape from the
plan-clustering pass to the StablehloToPlan conversion pass. This is a more
appropriate location since this conversion happens during the StableHLO to Plan
dialect conversion phase.

Also clean up unused includes in Clustering.cpp and remove unnecessary
dependent dialects declaration from Passes.td.

GitOrigin-RevId: 112e7c7386773df5cb57982ac1e3a831af5be621

Author: christopherbate

mlir-tensorrt/compiler/include/mlir-tensorrt/Dialect/Plan/Transforms/Passes.td

@@ -330,26 +330,18 @@ def ClusteringPass : Pass<"plan-clustering", "::mlir::ModuleOp"> {
     [CompilerBackendAttrInterface](../IR/PlanInterfaces.td).
   }];
 
-  let options = [
-    Option<"entrypoint", "entrypoint", "std::string", "\"\"",
-      "the name of the entrypoint function; if empty then the clustering runs"
-      " on all functions">,
-    Option<"forceEntrypointsReturnAllocs",
-      "force-entrypoints-return-allocs", "bool", "false",
-      "allow backend clusters to directly allocate outputs">,
-    Option<"disableCreateShapeFuncPass", "disable-create-shape-func-pass", "bool", "false",
-      "don't apply create shape to func pass in TensorRT clusters">,
-    InputKindOption
-  ];
-
-  let dependentDialects = [
-    // TODO: TensorRT and Tensor dialects needed since the different
-    // backends may create these ops. Add a way for backends to declare
-    // dialect dependencies.
-    "::mlir::tensor::TensorDialect",
-    "::mlir::tensorrt::TensorRTDialect",
-    "::mlir::plan::PlanDialect"
-  ];
+  let options =
+      [Option<"entrypoint", "entrypoint", "std::string", "\"\"",
+              "the name of the entrypoint function; if empty then the "
+              "clustering runs"
+              " on all functions">,
+       Option<"forceEntrypointsReturnAllocs", "force-entrypoints-return-allocs",
+              "bool", "false",
+              "allow backend clusters to directly allocate outputs">,
+       Option<"disableCreateShapeFuncPass", "disable-create-shape-func-pass",
+              "bool", "false",
+              "don't apply create shape to func pass in TensorRT clusters">,
+       InputKindOption];
 }
 
 //===----------------------------------------------------------------------===//

mlir-tensorrt/compiler/lib/Conversion/StablehloToPlan/StablehloToPlan.cpp

@@ -288,6 +288,15 @@ struct ShapeAssertionPattern
     return success();
   }
 };
+} // namespace
+
+static auto getIntegerAttrOrDefault(Operation *op, StringRef name,
+                                    int64_t defaultValue) {
+  if (auto attr = op->getAttrOfType<IntegerAttr>(name))
+    return attr.getInt();
+  return defaultValue;
+}
+namespace {
 
 struct ConvertStablehloToPlanPass
     : public impl::ConvertStablehloToPlanPassBase<ConvertStablehloToPlanPass> {
@@ -325,6 +334,22 @@ struct ConvertStablehloToPlanPass
       if (walkResult.wasInterrupted())
         return signalPassFailure();
     }
+
+    // Check for StableHLO partitioning attributes and attach the executor
+    // grid shape attribute.
+    ModuleOp module = getOperation();
+    {
+      auto numReplicas =
+          getIntegerAttrOrDefault(module, "mhlo.num_replicas", 1);
+      auto numPartitions =
+          getIntegerAttrOrDefault(module, "mhlo.num_partitions", 1);
+      if (failed(executor::setModuleProcessGridShape(
+              module, {numReplicas, numPartitions}))) {
+        emitError(module->getLoc())
+            << "failed to set the Executor process grid shape attribute";
+        return signalPassFailure();
+      }
+    }
   }
 };
 } // namespace

mlir-tensorrt/compiler/lib/Dialect/Plan/Transforms/Clustering.cpp

@@ -21,20 +21,16 @@
 /// Implementation of the `plan-clustering` pass.
 ///
 //===----------------------------------------------------------------------===//
-#include "mlir-executor/Executor/IR/Executor.h"
 #include "mlir-executor/Transforms/Clustering/Patterns.h"
 #include "mlir-tensorrt-common/Interfaces/TensorKindOpInterface.h"
 #include "mlir-tensorrt-dialect/Analysis/TensorKindAnalysis.h"
-#include "mlir-tensorrt-dialect/TensorRT/IR/TensorRTDialect.h"
 #include "mlir-tensorrt/Dialect/Plan/IR/Plan.h"
 #include "mlir-tensorrt/Dialect/Plan/IR/PlanInterfaces.h"
 #include "mlir-tensorrt/Dialect/Plan/Transforms/Passes.h"
 #include "mlir-tensorrt/Utils/DataFlowUtils.h"
 #include "mlir/Analysis/DataFlow/ConstantPropagationAnalysis.h"
 #include "mlir/Analysis/DataFlow/DeadCodeAnalysis.h"
 #include "mlir/Analysis/DataFlowFramework.h"
-#include "mlir/Dialect/Func/IR/FuncOps.h"
-#include "mlir/Dialect/Tensor/IR/Tensor.h"
 #include "mlir/IR/OpDefinition.h"
 #include "mlir/Transforms/DialectConversion.h"
 #include "mlir/Transforms/RegionUtils.h"
@@ -149,13 +145,6 @@ applyClusteringToFunc(RewriterBase &rewriter, FunctionOpInterface func,
   return success();
 }
 
-static auto getIntegerAttrOrDefault(Operation *op, StringRef name,
-                                    int64_t defaultValue) {
-  if (auto attr = op->getAttrOfType<IntegerAttr>(name))
-    return attr.getInt();
-  return defaultValue;
-}
-
 namespace {
 class ClusteringPass : public plan::impl::ClusteringPassBase<ClusteringPass> {
 public:
@@ -229,22 +218,6 @@ class ClusteringPass : public plan::impl::ClusteringPassBase<ClusteringPass> {
         return signalPassFailure();
     }
 
-    // Check for StableHLO partitioning attributes and attach the executor
-    // grid shape attribute.
-    /// TODO: move this logic into a standalone pass that handles partitioning.
-    {
-      auto numReplicas =
-          getIntegerAttrOrDefault(module, "mhlo.num_replicas", 1);
-      auto numPartitions =
-          getIntegerAttrOrDefault(module, "mhlo.num_partitions", 1);
-      if (failed(executor::setModuleProcessGridShape(
-              module, {numReplicas, numPartitions}))) {
-        emitError(module->getLoc())
-            << "failed to set the Executor process grid shape attribute";
-        return signalPassFailure();
-      }
-    }
-
     // Drop clustering attributes since they are no longer needed.
     module->removeAttr(plan::PlanDialect::kBackendsAttrName);
   }

mlir-tensorrt/compiler/test/Dialect/Plan/stablehlo-clustering.mlir

@@ -75,7 +75,7 @@ func.func @test(%arg0: tensor<4xi32>, %arg1: tensor<i32>)
 
 }
 
-// CHECK-LABEL: module @host_backends_with_values attributes
+// CHECK-LABEL: module @host_backends_with_values
 // CHECK-LABEL: func.func @test
 //  CHECK-SAME: (%[[arg0:.+]]: tensor<4xi32>, %[[arg1:.+]]: tensor<i32>)
 //   CHECK-DAG:     %[[cst:.+]] = stablehlo.constant dense<0.000000e+00> : tensor<f32>

mlir-tensorrt/executor/include/mlir-executor/Transforms/Clustering/Patterns.h

@@ -113,59 +113,6 @@ class ClusteringPatternSet {
   SmallVector<std::unique_ptr<RewriteType>> patterns;
 };
 
-/// Apply a set of clustering patterns to the function. Patterns are sorted and
-/// applied in decreasing order by benefit.
-LogicalResult
-applyClusteringPatterns(FunctionOpInterface mainFunc,
-                        ClusteringPatternSet<ClusteringRewriter> &patterns);
-
-/// A type of a function that can filter cluster region operations.
-using RegionOpFilterFn = std::function<bool(Operation *)>;
-
-/// Create a cluster region op filter using the specified parameters.
-RegionOpFilterFn getRegionOpFilter(
-    Attribute target,
-    unsigned operationCnt = std::numeric_limits<unsigned>::max(),
-    IsClusterableOpFn canOpCluster = [](Operation *op) { return true; });
-
-/// Given a `func.func` operation, this class describes a base
-/// pattern for doing a "scf::ExecuteRegion" based transformation
-class RegionOpFusionRewriter {
-public:
-  RegionOpFusionRewriter(const SmallVector<RegionOpFilterFn> &filters,
-                         Attribute newTarget,
-                         ClusterRegionOpBuilderFunc regionOpBuilderFunc)
-      : filters(filters), target(newTarget),
-        regionOpBuilderFunc(regionOpBuilderFunc) {}
-
-  /// This function walks on the mainFunc graph and  finds any matched
-  /// patterns according to filters. After it finds matched consecutive
-  /// Operation* in the graph and it will try to merge them into 1
-  /// single Operation* and rewrite it into the graph with a new
-  /// clustering target
-  void run(FunctionOpInterface mainFunc, RewriterBase &rewriter);
-
-private:
-  /// A list of filter functions that identify scf.execute_region operations of
-  /// interest
-  SmallVector<RegionOpFilterFn> filters;
-
-  /// the target of the merged region operation will be set to
-  Attribute target;
-
-  ClusterRegionOpBuilderFunc regionOpBuilderFunc;
-};
-
-/// Return the "target" of a particular cluster represented by the
-/// `scf.execute_region` operation. This currently returns the StringAttr
-/// named `__cluster_target__` if present, failure otherwise.
-FailureOr<Attribute> getClusterTarget(Operation *regionOp);
-
-/// Apply a set of region-op rewriter patterns to the function.
-LogicalResult applyRegionOpRewritePatterns(
-    FunctionOpInterface mainFunc,
-    ClusteringPatternSet<RegionOpFusionRewriter> &patterns);
-
 } // namespace mlir
 
 #endif // MLIR_TENSORRT_TRANSFORMS_CLUSTERING_PATTERNS_H