Refactor to simplify mesh extraction in getMeshOrDefault.

Google-ML-Automation · copybara-github · commit 984aecf29efc · 2025-09-04T05:59:57.000-07:00
It is supposed to be pure.

PiperOrigin-RevId: 800907631
diff --git a/shardy/dialect/sdy/transforms/export/explicit_reshards_util.cc b/shardy/dialect/sdy/transforms/export/explicit_reshards_util.cc
@@ -751,10 +751,9 @@ Mesh getMeshOrDefault(TensorShardingAttr sharding,
   if (!sharding) {
     return defaultMesh;
   }
-  StringRef meshName =
-      getCommonMeshName({sharding}, {}, symbolTable, /*ignoreDeviceIds=*/false)
-          .value();
-  return Mesh(getMeshAttr(symbolTable, meshName), meshName);
+  // NOTE: sharding always has a meshOrRef because it is a required parameter.
+  return Mesh(sharding.getMesh(symbolTable),
+              cast<FlatSymbolRefAttr>(sharding.getMeshOrRef()).getValue());
 }
 
 // Assumes that:

Original file line number	Diff line number	Diff line change
`@@ -751,10 +751,9 @@ Mesh getMeshOrDefault(TensorShardingAttr sharding,`
`751`	`751`	`if (!sharding) {`
`752`	`752`	`return defaultMesh;`
`753`	`753`	`}`
`754`		`- StringRef meshName =`
`755`		`- getCommonMeshName({sharding}, {}, symbolTable, /ignoreDeviceIds=/false)`
`756`		`- .value();`
`757`		`- return Mesh(getMeshAttr(symbolTable, meshName), meshName);`
	`754`	`+ // NOTE: sharding always has a meshOrRef because it is a required parameter.`
	`755`	`+ return Mesh(sharding.getMesh(symbolTable),`
	`756`	`+ cast<FlatSymbolRefAttr>(sharding.getMeshOrRef()).getValue());`
`758`	`757`	`}`
`759`	`758`
`760`	`759`	`// Assumes that:`