diff --git a/tpu_raiden/weight_sync/tiling_utils.cc b/tpu_raiden/weight_sync/tiling_utils.cc index 7cb85be4..b954bd17 100644 --- a/tpu_raiden/weight_sync/tiling_utils.cc +++ b/tpu_raiden/weight_sync/tiling_utils.cc @@ -14,16 +14,16 @@ #include "tpu_raiden/weight_sync/tiling_utils.h" +#include #include #include #include #include -#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()); diff --git a/tpu_raiden/weight_sync/tiling_utils_test.cc b/tpu_raiden/weight_sync/tiling_utils_test.cc index 6c07d113..5e9d5a2d 100644 --- a/tpu_raiden/weight_sync/tiling_utils_test.cc +++ b/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 src_linear(num_elements); + for (int i = 0; i < num_elements; ++i) { + src_linear[i] = static_cast(i); + } + + const int64_t tiled_size_bytes = num_elements * sizeof(float); + std::vector dst_tiled(tiled_size_bytes); + + absl::Status tile_status = + TileBuffer(reinterpret_cast(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(dst_tiled.data()); + EXPECT_EQ(dst_tiled_float[80], 68.0f); + + // Detile back. + std::vector dst_linear(num_elements, 0.0f); + absl::Status detile_status = DetileBuffer( + dst_tiled.data(), reinterpret_cast(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 src_linear(num_elements); + for (int i = 0; i < num_elements; ++i) { + src_linear[i] = static_cast(i); + } + + const int64_t tiled_size_bytes = num_elements * sizeof(float); + std::vector dst_tiled(tiled_size_bytes); + + absl::Status tile_status = + TileBuffer(reinterpret_cast(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(dst_tiled.data()); + EXPECT_EQ(dst_tiled_float[272], 260.0f); + + // Detile back. + std::vector dst_linear(num_elements, 0.0f); + absl::Status detile_status = DetileBuffer( + dst_tiled.data(), reinterpret_cast(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