[SDY][Cleanup] remove apply_sharding_constraints.h file by moving debug args td file options so that constructor is auto-generated via tblgen.

PiperOrigin-RevId: 810599911

Author: Varcho

shardy/dialect/sdy/transforms/import/BUILD

@@ -44,7 +44,6 @@ cc_library(
         "sharding_group_import.cc",
     ],
     hdrs = [
-        "apply_sharding_constraints.h",
         "passes.h",
     ],
     deps = [

shardy/dialect/sdy/transforms/import/apply_sharding_constraints.cc

@@ -13,25 +13,19 @@ See the License for the specific language governing permissions and
 limitations under the License.
 ==============================================================================*/
 
-#include "shardy/dialect/sdy/transforms/import/apply_sharding_constraints.h"
-
 #include <cassert>
 #include <functional>
-#include <memory>
 
 #include "llvm/ADT/STLExtras.h"
-#include "llvm/Support/CommandLine.h"
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/BuiltinAttributes.h"
 #include "mlir/IR/BuiltinOps.h"
 #include "mlir/IR/MLIRContext.h"
 #include "mlir/IR/Value.h"
 #include "mlir/IR/ValueRange.h"
-#include "mlir/Pass/Pass.h"
 #include "mlir/Support/LLVM.h"
 #include "shardy/dialect/sdy/ir/dialect.h"
 #include "shardy/dialect/sdy/ir/utils.h"
-#include "shardy/dialect/sdy/transforms/common/propagation_options.h"
 #include "shardy/dialect/sdy/transforms/import/passes.h"  // IWYU pragma: keep
 #include "shardy/dialect/sdy/transforms/propagation/debugging/source_sharding.h"
 #include "shardy/dialect/sdy/transforms/propagation/op_sharding_rule_registry.h"
@@ -239,14 +233,7 @@ struct ApplyShardingConstraintsPass
   ApplyShardingConstraintsPass(const ApplyShardingConstraintsPass& other)
       : ApplyShardingConstraintsPassBase<ApplyShardingConstraintsPass>(other) {
     debugShardingOrigins = other.debugShardingOrigins;
-    debugPropagationEdgeSharding = other.debugPropagationEdgeSharding;
-  }
-
-  // Constructor to be used when creating the pass programmatically with
-  // options.
-  explicit ApplyShardingConstraintsPass(const PropagationOptions& options) {
-    debugShardingOrigins = options.debugShardingOrigins;
-    debugPropagationEdgeSharding = options.debugPropagationEdgeSharding;
+    debugPropagationEdges = other.debugPropagationEdges;
   }
 
   void runOnOperation() final {
@@ -255,7 +242,7 @@ struct ApplyShardingConstraintsPass
 
     // Prepare debugging handler for sharding origins and edge sources.
     ShardingDebugMappings mappings(debugShardingOrigins,
-                                   debugPropagationEdgeSharding);
+                                   debugPropagationEdges);
     SourceShardingHandler handler(&mappings);
     // Prepare the handler and register it to the context.
     handler.prepareHandler(moduleOp);
@@ -315,26 +302,9 @@ struct ApplyShardingConstraintsPass
     context.registerActionHandler(nullptr);
     handler.saveOnModule(moduleOp);
   }
-
-  Option<bool> debugShardingOrigins{
-      *this, "debug-sharding-origins",
-      llvm::cl::desc("whether to save sharding origin information"),
-      llvm::cl::init(false)};
-
-  Option<bool> debugPropagationEdgeSharding{
-      *this, "debug-propagation-edge-sharding",
-      llvm::cl::desc("whether to save propagation edge information"),
-      llvm::cl::init(false)};
 };
 
 }  // namespace
 
-// This function can be used to create the pass with specific options
-// programmatically.
-std::unique_ptr<Pass> createApplyShardingConstraintsPass(
-    const PropagationOptions& options) {
-  return std::make_unique<ApplyShardingConstraintsPass>(options);
-}
-
 }  // namespace sdy
 }  // namespace mlir

shardy/dialect/sdy/transforms/import/apply_sharding_constraints.h

@@ -21,7 +21,6 @@ limitations under the License.
 #include "mlir/Transforms/Passes.h"
 #include "shardy/common/file_utils.h"
 #include "shardy/dialect/sdy/transforms/common/propagation_options.h"
-#include "shardy/dialect/sdy/transforms/import/apply_sharding_constraints.h"
 #include "shardy/dialect/sdy/transforms/import/passes.h"
 
 namespace mlir {
@@ -66,7 +65,9 @@ void addImportPipeline(OpPassManager& pm, int& dumpIndex,
       options.dumpDirectory, "before_propagation", dumpIndex++));
 
   pm.addNestedPass<func::FuncOp>(createAddDataFlowEdgesPass());
-  pm.addPass(createApplyShardingConstraintsPass(options));
+  pm.addPass(
+      createApplyShardingConstraintsPass(ApplyShardingConstraintsPassOptions{
+          options.debugShardingOrigins, options.debugPropagationEdgeSharding}));
   // The sharding group import pass must run after applying sharding
   // constraints. This ensures we can detect sharding conflicts between group
   // members which have pre-propagation shardings due to sharding constraints.

shardy/dialect/sdy/transforms/import/import_pipeline.cc

@@ -105,6 +105,19 @@ def ApplyShardingConstraintsPass : Pass<"sdy-apply-sharding-constraints", "Modul
     dangling case).
   }];
   let dependentDialects = ["mlir::sdy::SdyDialect"];
+
+  let options = [
+    Option<"debugShardingOrigins", "debug-sharding-origins", "bool",
+           /*default=*/"false",
+           "Whether to compute the debug origin shardings for constraints. See "
+           "`debug-sharding-origins` option in propagation for more info.">,
+    Option<"debugPropagationEdges",
+           "debug-propagation-edge-sharding", "bool",
+           /*default=*/"false",
+           "Whether to sink the debug propagation edge sharding info. See "
+           "`debug-propagation-edge-sharding` option in propagation for more "
+           "info.">
+  ];
 }
 
 def ConstantOrScalarSplitterPass : Pass<"sdy-constant-or-scalar-splitter", "func::FuncOp"> {