Optimized CPU tiling and detiling for weight resharding.

tpu_raiden/weight_sync/tiling_utils.cc

@@ -14,16 +14,16 @@
 
 #include "tpu_raiden/weight_sync/tiling_utils.h"
 
+#include <algorithm>
 #include <cstdint>
 #include <cstring>
 #include <utility>
 #include <vector>
 
-#include "absl/types/span.h"
-#include "xla/layout.h"
-
 #include "absl/status/status.h"
+#include "absl/types/span.h"
 #include "xla/index_util.h"
+#include "xla/layout.h"
 #include "xla/layout_util.h"
 #include "xla/shape.h"
 #include "xla/shape_util.h"
@@ -82,6 +82,144 @@ int64_t GetTiledBufferElements(const xla::Shape& shape) {
   return total_elements;
 }
 
+bool IsStandardRowMajorTiled(const xla::Shape& shape,
+                             const xla::Layout& layout) {
+  if (layout.tiles().size() != 1) return false;
+  if (layout.tiles(0).dimensions().size() != 2) return false;
+
+  const int R = shape.dimensions().size();
+  if (R < 2) return false;
+
+  for (int i = 0; i < R; ++i) {
+    if (layout.minor_to_major(i) != R - 1 - i) {
+      return false;
+    }
+  }
+  return true;
+}
+
+absl::Status TileBufferNDOptimized(const uint8_t* src_linear,
+                                   uint8_t* dst_tiled, const xla::Shape& shape,
+                                   const xla::Layout& layout) {
+  const int R = shape.dimensions().size();
+  int64_t H = shape.dimensions(layout.minor_to_major(1));
+  int64_t W = shape.dimensions(layout.minor_to_major(0));
+  int64_t itemsize =
+      xla::ShapeUtil::ByteSizeOfPrimitiveType(shape.element_type());
+
+  const xla::Tile& tile = layout.tiles(0);
+  int64_t tile_H = tile.dimension(0);
+  int64_t tile_W = tile.dimension(1);
+
+  int64_t num_tiles_0 = xla::CeilOfRatio(H, tile_H);
+  int64_t num_tiles_1 = xla::CeilOfRatio(W, tile_W);
+  int64_t tile_size_bytes = tile_H * tile_W * itemsize;
+
+  int64_t batch_size = 1;
+  for (int i = 2; i < R; ++i) {
+    batch_size *= shape.dimensions(layout.minor_to_major(i));
+  }
+
+  int64_t matrix_size_bytes = H * W * itemsize;
+  int64_t tiled_matrix_size_bytes = num_tiles_0 * num_tiles_1 * tile_size_bytes;
+
+  // Zero-initialize the destination buffer to handle padding automatically
+  int64_t total_physical_elements = GetTiledBufferElements(shape);
+  std::memset(dst_tiled, 0, total_physical_elements * itemsize);
+
+  for (int64_t b = 0; b < batch_size; ++b) {
+    const uint8_t* src_batch_ptr = src_linear + b * matrix_size_bytes;
+    uint8_t* dst_batch_ptr = dst_tiled + b * tiled_matrix_size_bytes;
+
+    for (int64_t tile_row = 0; tile_row < num_tiles_0; ++tile_row) {
+      for (int64_t tile_col = 0; tile_col < num_tiles_1; ++tile_col) {
+        int64_t tile_index = tile_row * num_tiles_1 + tile_col;
+        uint8_t* dst_tile_ptr = dst_batch_ptr + tile_index * tile_size_bytes;
+
+        for (int64_t r = 0; r < tile_H; ++r) {
+          int64_t logical_row = tile_row * tile_H + r;
+          if (logical_row >= H) {
+            continue;
+          }
+
+          int64_t logical_col_start = tile_col * tile_W;
+          int64_t valid_elements = std::min(tile_W, W - logical_col_start);
+          if (valid_elements <= 0) {
+            continue;
+          }
+
+          const uint8_t* src_row_ptr =
+              src_batch_ptr + (logical_row * W + logical_col_start) * itemsize;
+          uint8_t* dst_row_ptr = dst_tile_ptr + (r * tile_W) * itemsize;
+
+          std::memcpy(dst_row_ptr, src_row_ptr, valid_elements * itemsize);
+        }
+      }
+    }
+  }
+  return absl::OkStatus();
+}
+
+absl::Status DetileBufferNDOptimized(const uint8_t* src_tiled,
+                                     uint8_t* dst_linear,
+                                     const xla::Shape& shape,
+                                     const xla::Layout& layout) {
+  const int R = shape.dimensions().size();
+  int64_t H = shape.dimensions(layout.minor_to_major(1));
+  int64_t W = shape.dimensions(layout.minor_to_major(0));
+  int64_t itemsize =
+      xla::ShapeUtil::ByteSizeOfPrimitiveType(shape.element_type());
+
+  const xla::Tile& tile = layout.tiles(0);
+  int64_t tile_H = tile.dimension(0);
+  int64_t tile_W = tile.dimension(1);
+
+  int64_t num_tiles_0 = xla::CeilOfRatio(H, tile_H);
+  int64_t num_tiles_1 = xla::CeilOfRatio(W, tile_W);
+  int64_t tile_size_bytes = tile_H * tile_W * itemsize;
+
+  int64_t batch_size = 1;
+  for (int i = 2; i < R; ++i) {
+    batch_size *= shape.dimensions(layout.minor_to_major(i));
+  }
+
+  int64_t matrix_size_bytes = H * W * itemsize;
+  int64_t tiled_matrix_size_bytes = num_tiles_0 * num_tiles_1 * tile_size_bytes;
+
+  for (int64_t b = 0; b < batch_size; ++b) {
+    const uint8_t* src_batch_ptr = src_tiled + b * tiled_matrix_size_bytes;
+    uint8_t* dst_batch_ptr = dst_linear + b * matrix_size_bytes;
+
+    for (int64_t tile_row = 0; tile_row < num_tiles_0; ++tile_row) {
+      for (int64_t tile_col = 0; tile_col < num_tiles_1; ++tile_col) {
+        int64_t tile_index = tile_row * num_tiles_1 + tile_col;
+        const uint8_t* src_tile_ptr =
+            src_batch_ptr + tile_index * tile_size_bytes;
+
+        for (int64_t r = 0; r < tile_H; ++r) {
+          int64_t logical_row = tile_row * tile_H + r;
+          if (logical_row >= H) {
+            continue;
+          }
+
+          int64_t logical_col_start = tile_col * tile_W;
+          int64_t valid_elements = std::min(tile_W, W - logical_col_start);
+          if (valid_elements <= 0) {
+            continue;
+          }
+
+          uint8_t* dst_row_ptr =
+              dst_batch_ptr + (logical_row * W + logical_col_start) * itemsize;
+          const uint8_t* src_row_ptr = src_tile_ptr + (r * tile_W) * itemsize;
+
+          std::memcpy(dst_row_ptr, src_row_ptr, valid_elements * itemsize);
+        }
+      }
+    }
+  }
+  return absl::OkStatus();
+}
+
 }  // namespace
 
 absl::Status DetileBuffer(const uint8_t* src_tiled, uint8_t* dst_linear,
@@ -90,6 +228,10 @@ absl::Status DetileBuffer(const uint8_t* src_tiled, uint8_t* dst_linear,
     return absl::InternalError("Buffer is not tiled");
   }
 
+  if (IsStandardRowMajorTiled(shape, layout)) {
+    return DetileBufferNDOptimized(src_tiled, dst_linear, shape, layout);
+  }
+
   int64_t itemsize =
       xla::ShapeUtil::ByteSizeOfPrimitiveType(shape.element_type());
 
@@ -119,6 +261,10 @@ absl::Status TileBuffer(const uint8_t* src_linear, uint8_t* dst_tiled,
     return absl::InternalError("Buffer is not tiled");
   }
 
+  if (IsStandardRowMajorTiled(shape, layout)) {
+    return TileBufferNDOptimized(src_linear, dst_tiled, shape, layout);
+  }
+
   int64_t itemsize =
       xla::ShapeUtil::ByteSizeOfPrimitiveType(shape.element_type());
 

tpu_raiden/weight_sync/tiling_utils_test.cc

@@ -226,5 +226,114 @@ TEST(TilingUtilsTest, Bf16SubTiling) {
   }
 }
 
+TEST(TilingUtilsTest, Standard3D) {
+  // 3D matrix of shape 2x8x8, element type float (4 bytes).
+  // Layout has minor_to_major={2, 1, 0} (standard row-major), and tiling with
+  // tile dimensions 4x4.
+  const int64_t D0 = 2;
+  const int64_t H = 8;
+  const int64_t W = 8;
+  const int64_t tH = 4;
+  const int64_t tW = 4;
+
+  xla::Shape shape = xla::ShapeUtil::MakeShapeWithDenseLayout(
+      xla::PrimitiveType::F32, {D0, H, W}, {2, 1, 0}, {xla::Tile({tH, tW})});
+
+  const int64_t num_elements = D0 * H * W;
+  std::vector<float> src_linear(num_elements);
+  for (int i = 0; i < num_elements; ++i) {
+    src_linear[i] = static_cast<float>(i);
+  }
+
+  const int64_t tiled_size_bytes = num_elements * sizeof(float);
+  std::vector<uint8_t> dst_tiled(tiled_size_bytes);
+
+  absl::Status tile_status =
+      TileBuffer(reinterpret_cast<const uint8_t*>(src_linear.data()),
+                 dst_tiled.data(), shape, shape.layout());
+  EXPECT_TRUE(tile_status.ok()) << tile_status.ToString();
+
+  // Verify tiled structure.
+  // Block 0 (d0=0) covers indices 0-63.
+  // Block 1 (d0=1) covers indices 64-127.
+  // Within Block 1, the 2D matrix is 8x8 tiled with 4x4.
+  // Let's check d0=1, row 0, col 4: linear index is 1*64 + 0*8 + 4 = 68. Value
+  // is 68.0. Physically, Block 1 starts at 64 * 4 = 256 bytes. Within Block 1,
+  // row 0, col 4 is in Tile(0,1), offset 0. Tile(0,1) index is 1. Offset within
+  // Block 1: 1 * 16 * 4 = 64 bytes. Total physical offset: 256 + 64 = 320 bytes
+  // (index 80 in float array).
+  float* dst_tiled_float = reinterpret_cast<float*>(dst_tiled.data());
+  EXPECT_EQ(dst_tiled_float[80], 68.0f);
+
+  // Detile back.
+  std::vector<float> dst_linear(num_elements, 0.0f);
+  absl::Status detile_status = DetileBuffer(
+      dst_tiled.data(), reinterpret_cast<uint8_t*>(dst_linear.data()), shape,
+      shape.layout());
+  EXPECT_TRUE(detile_status.ok()) << detile_status.ToString();
+
+  for (int i = 0; i < num_elements; ++i) {
+    EXPECT_EQ(dst_linear[i], src_linear[i]) << "Mismatch at index " << i;
+  }
+}
+
+TEST(TilingUtilsTest, Standard4D) {
+  // 4D matrix of shape 2x3x8x8, element type float (4 bytes).
+  // Layout has minor_to_major={3, 2, 1, 0} (standard row-major), and tiling
+  // with tile dimensions 4x4.
+  const int64_t D0 = 2;
+  const int64_t D1 = 3;
+  const int64_t H = 8;
+  const int64_t W = 8;
+  const int64_t tH = 4;
+  const int64_t tW = 4;
+
+  xla::Shape shape = xla::ShapeUtil::MakeShapeWithDenseLayout(
+      xla::PrimitiveType::F32, {D0, D1, H, W}, {3, 2, 1, 0},
+      {xla::Tile({tH, tW})});
+
+  const int64_t num_elements = D0 * D1 * H * W;
+  std::vector<float> src_linear(num_elements);
+  for (int i = 0; i < num_elements; ++i) {
+    src_linear[i] = static_cast<float>(i);
+  }
+
+  const int64_t tiled_size_bytes = num_elements * sizeof(float);
+  std::vector<uint8_t> dst_tiled(tiled_size_bytes);
+
+  absl::Status tile_status =
+      TileBuffer(reinterpret_cast<const uint8_t*>(src_linear.data()),
+                 dst_tiled.data(), shape, shape.layout());
+  EXPECT_TRUE(tile_status.ok()) << tile_status.ToString();
+
+  // Verify tiled structure.
+  // We have D0 * D1 = 6 batches.
+  // Each batch is 8x8 tiled with 4x4.
+  // Let's check d0=1, d1=1, row 0, col 4:
+  // Linear index: d0*(D1*H*W) + d1*(H*W) + row*W + col
+  // = 1*(3*8*8) + 1*(8*8) + 0*8 + 4
+  // = 192 + 64 + 4 = 260. Value is 260.0.
+  // Physically, each batch has size H * W * sizeof(float) = 64 * 4 = 256 bytes.
+  // Batch index is d0 * D1 + d1 = 1 * 3 + 1 = 4.
+  // Batch offset: 4 * 256 = 1024 bytes.
+  // Within batch 4, row 0, col 4 is in Tile(0,1), offset 0.
+  // Tile(0,1) index is 1.
+  // Offset within batch: 1 * 16 * 4 = 64 bytes.
+  // Total physical offset: 1024 + 64 = 1088 bytes (index 272 in float array).
+  float* dst_tiled_float = reinterpret_cast<float*>(dst_tiled.data());
+  EXPECT_EQ(dst_tiled_float[272], 260.0f);
+
+  // Detile back.
+  std::vector<float> dst_linear(num_elements, 0.0f);
+  absl::Status detile_status = DetileBuffer(
+      dst_tiled.data(), reinterpret_cast<uint8_t*>(dst_linear.data()), shape,
+      shape.layout());
+  EXPECT_TRUE(detile_status.ok()) << detile_status.ToString();
+
+  for (int i = 0; i < num_elements; ++i) {
+    EXPECT_EQ(dst_linear[i], src_linear[i]) << "Mismatch at index " << i;
+  }
+}
+
 }  // namespace
 }  // namespace tpu_raiden::weight_sync