#sdy #debug support debugging tool to save the edge source shardings.

shardy/dialect/sdy/ir/constants.h

@@ -35,16 +35,30 @@ inline constexpr StringRef kShardingRuleAttr = "sdy.sharding_rule";
 // caused a value to be sharded a certain way.
 inline constexpr StringRef kShardingOriginsAttr = "sdy.sharding_origins";
 
+// Attribute name for saving which operand/result sharding of an op caused its
+// value to be sharded a certain way.
+inline constexpr StringRef kPropagationEdgesAttr = "sdy.propagation_edges";
+
 // Attribute name like `kShardingOriginsAttr` but for
 // `ShardableDataFlowOpInterface` op block arguments.
 inline constexpr StringRef kBlockArgShardingOriginsAttr =
     "sdy.block_arg_sharding_origins";
 
+// Attribute name like `kPropagationEdgesAttr` but for
+// `ShardableDataFlowOpInterface` op block arguments.
+inline constexpr StringRef kBlockArgPropagationEdgesAttr =
+    "sdy.block_arg_propagation_edges";
+
 // Attribute name like `kShardingOriginsAttr` but for
 // `ShardableDataFlowOpInterface` op results.
 inline constexpr StringRef kResultShardingOriginsAttr =
     "sdy.result_sharding_origins";
 
+// Attribute name like `kPropagationEdgesAttr` but for
+// `ShardableDataFlowOpInterface` op results.
+inline constexpr StringRef kResultPropagationEdgesAttr =
+    "sdy.result_propagation_edges";
+
 // Attribute name for the unique name of a sharding origin. Is either an
 // `sdy.sharding_constraint`, or `sdy.ManualComputationOp` input/output.
 inline constexpr StringRef kShardingOriginNameAttr = "sdy.sharding_origin_name";

shardy/dialect/sdy/transforms/export/BUILD

@@ -58,6 +58,7 @@ cc_library(
         "//shardy/dialect/sdy/transforms/propagation:op_sharding_rule_registry",
         "//shardy/dialect/sdy/transforms/propagation:sharding_projection",
         "//shardy/dialect/sdy/transforms/propagation:utils",
+        "//shardy/dialect/sdy/transforms/propagation/debugging:source_sharding",
         "@llvm-project//llvm:Support",
         "@llvm-project//mlir:FuncDialect",
         "@llvm-project//mlir:IR",

shardy/dialect/sdy/transforms/export/passes.td

@@ -28,6 +28,19 @@ def SinkDataFlowEdgesPass : Pass<"sdy-sink-data-flow-edges", "func::FuncOp"> {
   }];
   let dependentDialects = ["mlir::sdy::SdyDialect"];
   //TODO(tomnatan): consider moving the sharding to all targets that can have a sharding attached.
+
+  let options = [
+    Option<"sinkDebugShardingOrigins", "sink-debug-sharding-origins", "bool",
+           /*default=*/"false",
+           "Whether to sink the debug sharding origins info. See "
+           "`debug-sharding-origins` option in propagation for more info.">,
+    Option<"sinkDebugPropagationEdgeSharding",
+           "sink-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 UpdateNonDivisibleInputOutputShardingsPass : Pass<"sdy-update-non-divisible-input-output-shardings", "func::FuncOp"> {

shardy/dialect/sdy/transforms/export/sink_data_flow_edges.cc

@@ -30,9 +30,10 @@ limitations under the License.
 #include "mlir/IR/Visitors.h"
 #include "mlir/Pass/Pass.h"  // IWYU pragma: keep
 #include "mlir/Support/LLVM.h"
-#include "shardy/dialect/sdy/ir/constants.h"
 #include "shardy/dialect/sdy/ir/dialect.h"
 #include "shardy/dialect/sdy/ir/utils.h"
+#include "shardy/dialect/sdy/transforms/export/passes.h"  // IWYU pragma: keep
+#include "shardy/dialect/sdy/transforms/propagation/debugging/source_sharding.h"
 
 namespace mlir {
 namespace sdy {
@@ -81,44 +82,6 @@ SmallVector<TensorShardingAttr> getShardingsFromDataFlowEdges(
   return shardings;
 }
 
-// Saves an array of all the origin sharding dictionaries for the given
-// `edgeOwners` on `op`. If non exist, nothing is saved.
-//
-// For debugging the origin shardings, we want to preserve the origin sharding
-// dictionaries from the `DataFlowEdgeOp`s on the owning op so that they are
-// preserved after the propagation pipeline.
-//
-// See the `debug-sharding-origins` config on propagation for more details.
-//
-// TODO(b/388458831): add `saveDebugPropagationInfo` to the pass and pass it in
-// here. Can then reserve the right size for `originShardingDicts` and not need
-// the `exists` boolean.
-void buildOriginShardingDictsFromDataFlowEdges(ValueRange edgeOwners,
-                                               Operation* op,
-                                               StringRef attrName,
-                                               IRRewriter& rewriter) {
-  SmallVector<Attribute> originShardingDicts;
-  // TODO(b/388458831): pass through a boolean indicating whether the origin
-  // sharding debug information is enabled.
-  bool exists = false;
-  for (Value edgeOwner : edgeOwners) {
-    DictionaryAttr dict;
-    if (auto dataFlowEdgeOp = DataFlowEdgeOp::lookup(edgeOwner)) {
-      dict =
-          dataFlowEdgeOp->getAttrOfType<DictionaryAttr>(kShardingOriginsAttr);
-    }
-    if (!dict) {
-      dict = rewriter.getDictionaryAttr({});
-    } else {
-      exists = true;
-    }
-    originShardingDicts.push_back(dict);
-  }
-  if (exists) {
-    op->setAttr(attrName, rewriter.getArrayAttr(originShardingDicts));
-  }
-}
-
 struct SinkDataFlowEdgesPass
     : public impl::SinkDataFlowEdgesPassBase<SinkDataFlowEdgesPass> {
   using SinkDataFlowEdgesPassBase::SinkDataFlowEdgesPassBase;
@@ -152,17 +115,19 @@ struct SinkDataFlowEdgesPass
         shardableDataFlowOp.setBlockArgumentEdgeOwnerShardings(
             blockArgShardings);
       }
-      buildOriginShardingDictsFromDataFlowEdges(
-          blockArgOwners, op, kBlockArgShardingOriginsAttr, rewriter);
+      saveDebugInfoDictsFromDataFlowEdges(
+          blockArgOwners, op, sinkDebugShardingOrigins,
+          sinkDebugPropagationEdgeSharding, EdgeNodeType::OPERAND, rewriter);
 
       ResultRange resultOwners = shardableDataFlowOp.getOpResultEdgeOwners();
       if (SmallVector<TensorShardingAttr> resultShardings =
               getShardingsFromDataFlowEdges(resultOwners);
           !resultShardings.empty()) {
         shardableDataFlowOp.setOpResultEdgeOwnerShardings(resultShardings);
       }
-      buildOriginShardingDictsFromDataFlowEdges(
-          resultOwners, op, kResultShardingOriginsAttr, rewriter);
+      saveDebugInfoDictsFromDataFlowEdges(
+          resultOwners, op, sinkDebugShardingOrigins,
+          sinkDebugPropagationEdgeSharding, EdgeNodeType::RESULT, rewriter);
       return WalkResult::advance();
     });
   }

shardy/dialect/sdy/transforms/propagation/basic_propagation.cc

@@ -300,8 +300,8 @@ LogicalResult propagateTensorShardings(
   if (context->hasActionHandler()) {
     context->executeAction<SourceShardingAction>(
         updateShardings,
-        /*IRUnits=*/{op}, operandsParams.tensors, resultsParams.tensors, mesh,
-        shardingRule, shardingProjection);
+        /*IRUnits=*/{op}, op, operandsParams.tensors, resultsParams.tensors,
+        mesh, shardingRule, shardingProjection, anyUpdated);
   } else {
     updateShardings();
   }
@@ -649,7 +649,8 @@ void BasicPropagationPassImpl::runOnOperation() {
   MLIRContext& context = getContext();
 
   // Prepare debugging handler for sharding origins and edge sources.
-  ShardingDebugMappings mappings(debugShardingOrigins, debugEdgeSourceSharding);
+  ShardingDebugMappings mappings(debugShardingOrigins,
+                                 debugPropagationEdgeSharding);
   SourceShardingHandler handler(&mappings);
   handler.prepareHandler(moduleOp);
 
@@ -683,7 +684,7 @@ void BasicPropagationPassImpl::setPropagationOptions(
   dumpDirectory = options.dumpDirectory.str();
   conservativePropagation = options.conservativePropagation;
   debugShardingOrigins = options.debugShardingOrigins;
-  debugEdgeSourceSharding = options.debugEdgeSourceSharding;
+  debugPropagationEdgeSharding = options.debugPropagationEdgeSharding;
 }
 
 std::unique_ptr<Pass> createBasicPropagationPass(

shardy/dialect/sdy/transforms/propagation/basic_propagation.h

@@ -126,12 +126,12 @@ class BasicPropagationPassImpl : public OperationPass<ModuleOp> {
           "before propagation."),
       llvm::cl::init(false)};
 
-  Option<bool> debugEdgeSourceSharding{
-      *this, "debug-edge-source-sharding",
+  Option<bool> debugPropagationEdgeSharding{
+      *this, "debug-propagation-edge-sharding",
       llvm::cl::desc(
-          "whether to save information about the edge source of a sharding "
-          "on the MLIR module. These are from which operand/result a sharding "
-          "was propagated."),
+          "whether to save information about the SSA value edges of how a "
+          "sharding on the MLIR module propagated around. These are from which "
+          "operand/result a sharding was propagated to a given op."),
       llvm::cl::init(false)};
 
  private: