[BACKEND] Fixes removeZeroBasesAlongDim outDims and surjectivity check

zahimoud · zahimoud · commit 296cc291d38f · 2025-05-01T15:51:07.000-07:00
diff --git a/lib/Tools/LinearLayout.cpp b/lib/Tools/LinearLayout.cpp
@@ -1088,7 +1088,16 @@ LinearLayout LinearLayout::removeZeroBasesAlongDim(StringAttr stripDim) const {
       }
     }
   }
-  return LinearLayout(std::move(result), llvm::to_vector(getOutDimNames()));
+  SmallVector<std::pair<StringAttr, int32_t>> newOutDimSizes;
+  for (auto outDim : getOutDimNames()) {
+    newOutDimSizes.push_back({outDim, getOutDimSize(outDim)});
+  }
+  auto newLayout = LinearLayout(std::move(result), ArrayRef(newOutDimSizes),
+                                NoCheckInvariants{});
+  assert(
+      newLayout.isSurjective() == this->isSurjective() &&
+      "Removing zero bases along a dimension should not change surjectivity");
+  return newLayout;
 }
 
 size_t hash_value(const LinearLayout &layout) {

Original file line number	Diff line number	Diff line change
`@@ -1088,7 +1088,16 @@ LinearLayout LinearLayout::removeZeroBasesAlongDim(StringAttr stripDim) const {`
`1088`	`1088`	`}`
`1089`	`1089`	`}`
`1090`	`1090`	`}`
`1091`		`- return LinearLayout(std::move(result), llvm::to_vector(getOutDimNames()));`
	`1091`	`+ SmallVector<std::pair<StringAttr, int32_t>> newOutDimSizes;`
	`1092`	`+ for (auto outDim : getOutDimNames()) {`
	`1093`	`+ newOutDimSizes.push_back({outDim, getOutDimSize(outDim)});`
	`1094`	`+ }`
	`1095`	`+ auto newLayout = LinearLayout(std::move(result), ArrayRef(newOutDimSizes),`
	`1096`	`+ NoCheckInvariants{});`
	`1097`	`+ assert(`
	`1098`	`+ newLayout.isSurjective() == this->isSurjective() &&`
	`1099`	`+ "Removing zero bases along a dimension should not change surjectivity");`
	`1100`	`+ return newLayout;`
`1092`	`1101`	`}`
`1093`	`1102`
`1094`	`1103`	`size_t hash_value(const LinearLayout &layout) {`