Introduce isFullyReplicated() for TensorShardingAttr

Google-ML-Automation · copybara-github · commit a0d337eecb49 · 2024-07-31T12:10:06.000-07:00
PiperOrigin-RevId: 658095168
diff --git a/shardy/dialect/sdy/ir/attrs.td b/shardy/dialect/sdy/ir/attrs.td
@@ -370,6 +370,13 @@ def Sdy_TensorSharding : AttrDef<Sdy_Dialect, "TensorSharding"> {
       return getDimSharding(dim).getIsClosed();
     }
 
+    bool isFullyReplicated() const {
+      return llvm::all_of(getDimShardings(),
+                      [](const DimensionShardingAttr dimSharding) {
+                         return dimSharding.emptyAxes();
+                      });
+    }
+
     StringRef getMeshName() const {
       return getMeshSymName().getValue();
     }

Original file line number	Diff line number	Diff line change
`@@ -370,6 +370,13 @@ def Sdy_TensorSharding : AttrDef<Sdy_Dialect, "TensorSharding"> {`
`370`	`370`	`return getDimSharding(dim).getIsClosed();`
`371`	`371`	`}`
`372`	`372`
	`373`	`+ bool isFullyReplicated() const {`
	`374`	`+ return llvm::all_of(getDimShardings(),`
	`375`	`+ [](const DimensionShardingAttr dimSharding) {`
	`376`	`+ return dimSharding.emptyAxes();`
	`377`	`+ });`
	`378`	`+ }`
	`379`	`+`
`373`	`380`	`StringRef getMeshName() const {`
`374`	`381`	`return getMeshSymName().getValue();`
`375`	`382`	`}`