[NFC][Mosaic TPU] Separate out retiling from relayout

PiperOrigin-RevId: 658335679
jax-ml · Aug 1, 2024 · 0307438 · 0307438
1 parent 0734345
commit 0307438
Show file tree

Hide file tree

Showing 3 changed files with 211 additions and 183 deletions.
diff --git a/jaxlib/mosaic/dialect/tpu/layout.cc b/jaxlib/mosaic/dialect/tpu/layout.cc
@@ -713,7 +713,8 @@ llvm::hash_code hash_value(const VectorLayout& layout) {
   return llvm::hash_value(layout.as_tuple());
 }
 
-std::ostream &operator<<(std::ostream &os, VectorLayout::ImplicitDim dim) {
+template <typename Stream>
+Stream& printImplicitDim(Stream& os, VectorLayout::ImplicitDim dim) {
   switch (dim) {
     case VectorLayout::ImplicitDim::kNone:
       os << "none";
@@ -728,6 +729,15 @@ std::ostream &operator<<(std::ostream &os, VectorLayout::ImplicitDim dim) {
   return os;
 }
 
+std::ostream& operator<<(std::ostream& os, VectorLayout::ImplicitDim dim) {
+  return printImplicitDim(os, dim);
+}
+
+mlir::Diagnostic& operator<<(mlir::Diagnostic& diag,
+                             VectorLayout::ImplicitDim dim) {
+  return printImplicitDim(diag, dim);
+}
+
 std::optional<Layout> parseLayout(mlir::AsmParser& parser) {
   std::string layout_str;
   if (failed(parser.parseString(&layout_str))) {

diff --git a/jaxlib/mosaic/dialect/tpu/layout.h b/jaxlib/mosaic/dialect/tpu/layout.h
@@ -513,6 +513,8 @@ llvm::raw_ostream &operator<<(llvm::raw_ostream &os, const Layout &v);
 llvm::hash_code hash_value(const VectorLayout &layout);
 mlir::Diagnostic &operator<<(mlir::Diagnostic &diag, const Layout &v);
 std::ostream &operator<<(std::ostream &os, VectorLayout::ImplicitDim dim);
+mlir::Diagnostic &operator<<(mlir::Diagnostic &diag,
+                             VectorLayout::ImplicitDim dim);
 
 std::optional<Layout> parseLayout(mlir::AsmParser &parser);