From 02de9f22b5d3b7af401eca845ee8b4be514300a9 Mon Sep 17 00:00:00 2001
From: Tom Natan <tomnatan@google.com>
Date: Mon, 10 Feb 2025 02:13:35 -0800
Subject: [PATCH] test structured bindings

PiperOrigin-RevId: 725127187
---
 .../sdy/transforms/propagation/basic_propagation.cc  | 10 +++-------
 .../propagation/debugging/source_sharding.cc         | 12 +++++++-----
 2 files changed, 10 insertions(+), 12 deletions(-)

diff --git a/shardy/dialect/sdy/transforms/propagation/basic_propagation.cc b/shardy/dialect/sdy/transforms/propagation/basic_propagation.cc
index e375bce4..61356310 100644
--- a/shardy/dialect/sdy/transforms/propagation/basic_propagation.cc
+++ b/shardy/dialect/sdy/transforms/propagation/basic_propagation.cc
@@ -614,11 +614,7 @@ LogicalResult BasicPropagationPassImpl::propagate(
       conservativePropagation, shardingGroupMap);
   // We only need a single iteration (and another to confirm convergence), since
   // we make sure ops whose sharding changes are added back to the worklist.
-  GreedyRewriteConfig config;
-  config.useTopDownTraversal = true;
-  config.enableRegionSimplification = mlir::GreedySimplifyRegionLevel::Disabled;
-  config.fold = false;
-  config.cseConstants = false;
+  GreedyRewriteConfig config{.useTopDownTraversal = true};
   if (failed(applyPatternsGreedily(moduleOp, std::move(patterns), config))) {
     // We should always converge in 2 iterations, if we don't, something is
     // wrong.
@@ -628,8 +624,8 @@ LogicalResult BasicPropagationPassImpl::propagate(
     return failure();
   }
 
-  // Pushes any shardings from the values returned in the terminator of the body
-  // of `funcOp` to the corresponding `funcOp` result type attrs.
+  // Pushes any shardings from tha values returned in the terminator of the body
+  // of `funcOp` to the coresponding `funcOp` result type attrs.
   if (failed(propagateFuncResults(moduleOp, symbolTable, factorPropagation,
                                   shardingGroupMap))) {
     return failure();
diff --git a/shardy/dialect/sdy/transforms/propagation/debugging/source_sharding.cc b/shardy/dialect/sdy/transforms/propagation/debugging/source_sharding.cc
index 6461e0fc..e990bc23 100644
--- a/shardy/dialect/sdy/transforms/propagation/debugging/source_sharding.cc
+++ b/shardy/dialect/sdy/transforms/propagation/debugging/source_sharding.cc
@@ -304,7 +304,7 @@ void saveShardingOriginsOnModule(
 
     func::FuncOp funcOp = getEnclosingOfType<func::FuncOp>(owningOp);
 
-    // TODO(bartchr): Swap the map to store `ValueOrFuncResult` to avoid having
+    // TODO(bartchr): Swop the map to sture `ValueOrFuncResult` to avoid having
     // to do this terminator finding logic just to set the func result attr.
     OpOperand* terminatorOperand = getTerminatorOperand(value, funcOp);
 
@@ -317,10 +317,12 @@ void saveShardingOriginsOnModule(
                            builder.getDictionaryAttr(entries));
     }
     TypeSwitch<Operation*, void>(owningOp)
-        .Case<func::FuncOp>([&, value = value](func::FuncOp funcOp) {
-          funcOp.setArgAttr(cast<BlockArgument>(value).getArgNumber(),
-                            kShardingOriginsAttr,
-                            builder.getDictionaryAttr(entries));
+        .Case<func::FuncOp>([&](func::FuncOp funcOp) {
+          if (value) {
+            funcOp.setArgAttr(cast<BlockArgument>(value).getArgNumber(),
+                              kShardingOriginsAttr,
+                              builder.getDictionaryAttr(entries));
+          }
         })
         .Case<ShardingConstraintOp, DataFlowEdgeOp>([&](Operation* op) {
           op->setAttr(kShardingOriginsAttr, builder.getDictionaryAttr(entries));