layout.rs

// SPDX-FileCopyrightText: Copyright (c) 2024-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
// SPDX-License-Identifier: Apache-2.0

#![deny(missing_docs)]

//! # Block Layout Management 🧱
//!
//! This module is responsible for defining and managing the memory layout of data blocks.
//! It provides the foundational traits and concrete implementations for how blocks,
//! composed of multiple layers and pages, are arranged within a given [`Storage`].
//! The primary goal is to abstract the complexities of memory organization, including
//! contiguity, strides, and alignment, to ensure efficient data access and manipulation.
//!
//! ## Core Concepts
//!
//! ### 1. Layout Traits
//! The module defines a set of traits to ensure a consistent interface across different layout strategies:
//! - [`BlockLayout`]: The central trait that combines configuration and lookup capabilities. It specifies the
//!   associated [`StorageType`].
//! - [`BlockLayoutConfig`]: Provides metadata about the layout, such as the number of blocks, layers, page size,
//!   and data type.
//! - [`BlockLayoutLookup`]: Offers methods to retrieve the memory address and size of a specific memory region
//!   (page) within the layout.
//!
//! ### 2. Layout Configuration
//! The [`LayoutConfig`] struct is used to define the parameters of a block layout, including:
//! - `num_blocks`: Total number of blocks.
//! - `num_layers`: Number of layers per block.
//! - `page_size`: Size of each page (often corresponds to a dimension like sequence length or number of tokens).
//! - `inner_dim`: The inner dimension of the data (e.g., hidden size).
//! - `alignment`: Required memory alignment for certain operations or hardware. Must be a power of 2.
//! - `dtype`: The data type ([`DType`]) of the elements stored.
//!
//! This configuration is validated to ensure consistency and correctness (e.g., alignment must be a power of 2).
//!
//! ### 3. Concrete Layouts
//! Currently, the primary implemented layout is:
//! - [`FullyContiguous<S>`]: Represents a layout where all blocks and their constituent layers are stored sequentially
//!   in a single contiguous memory region provided by the generic storage `S`. It handles potential alignment
//!   requirements by calculating a `base_offset` within the provided storage and adjusting strides between blocks if
//!   necessary.
//!
//! ### 4. Strides and Alignment
//! The layout calculations meticulously handle strides between layers and blocks. For instance, in [`FullyContiguousConfig`]:
//! - `layer_stride_in_bytes`: The size of one memory region (page).
//! - `natural_block_stride`: The size of one block if there were no additional alignment padding between blocks.
//! - `block_stride_in_bytes`: The actual stride between the start of consecutive blocks, potentially larger than
//!   `natural_block_stride` to meet `alignment` requirements.
//! - `base_offset`: An offset applied from the start of the allocated [`Storage`] to ensure the first block's
//!   data begins at an aligned address.
//!
//! The function `align_up` is a utility to ensure values are aligned to the nearest multiple of a power-of-2 alignment.
//!
//! ### 5. Storage Interaction
//! Layouts are tightly coupled with the [`Storage`] trait from the `super::storage` module.
//! The [`BlockLayout::allocate`] method uses a [`StorageAllocator`] to obtain the necessary memory,
//! calculating the required size including any padding for alignment.
//!
//! ### 6. Error Handling
//! Operations within this module can result in [`LayoutError`], which covers issues like invalid configuration, validation errors, or out-of-bounds indexing.
//!
//! ## Usage Example
//!
//! ```rust
//! use dynamo_llm::block_manager::layout::{
//!     LayoutConfig, FullyContiguous, BlockLayout, BlockLayoutLookup, BlockLayoutConfig,
//! };
//! use dynamo_llm::block_manager::storage::{SystemAllocator, StorageType};
//! use dynamo_llm::common::dtype::DType;
//!
//! // Define the layout configuration
//! let config = LayoutConfig::builder()
//!     .num_blocks(10)
//!     .num_layers(4)
//!     .outer_dim(1)
//!     .page_size(16)
//!     .inner_dim(128)
//!     .dtype(DType::FP16)
//!     .build()
//!     .unwrap();
//!
//!
//! // Allocate a FullyContiguous layout using a SystemAllocator
//! let allocator = SystemAllocator;
//! let layout = FullyContiguous::allocate(config, &allocator).unwrap();
//!
//! // Access layout properties
//! assert_eq!(layout.num_blocks(), 10);
//! assert_eq!(layout.storage_type(), StorageType::System);
//!
//! // Get the address of a specific page
//! let addr = layout.memory_region_addr(0, 0).unwrap();
//! println!("Address of block 0, layer 0: {}", addr);
//! ```
//!
//! ## NIXL Integration
//! This module also includes a submodule `nixl` ([`crate::block_manager::layout::nixl`])
//! which extends these layout concepts for NIXL (NVIDIA Interface eXchange Layer), enabling
//! layouts to be registered and serialized for use in distributed environments.

// todo: coming soon...
// pub mod distributed;

pub mod nixl;
mod utils;

use utils::*;

use derive_getters::Getters;
use thiserror::Error;

use crate::block_manager::storage::{Storage, StorageAllocator};
use derive_builder::Builder;
use serde::{Deserialize, Serialize};
use tracing::instrument;
use validator::Validate;

use super::storage::StorageType;

/// Errors that can occur during layout operations
#[derive(Debug, Error)]
#[allow(missing_docs)]
pub enum LayoutError {
    #[error("Invalid configuration: {0}")]
    InvalidConfig(String),

    #[error("Validation failed: {0}")]
    ValidationError(#[from] validator::ValidationErrors),

    #[error("Invalid block index: {0}")]
    InvalidBlockIndex(usize),

    #[error("Invalid layer index: {0}")]
    InvalidLayerIndex(usize),

    #[error("Invalid outer index: {0}")]
    InvalidOuterIndex(usize),

    #[error("Operation failed: {0}")]
    OperationFailed(String),

    #[error("Serialization error: {0}")]
    SerdeError(#[from] serde_json::Error),
}

/// Storage pattern for layers
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub enum LayoutType {
    /// All layers are contiguous in memory [n_blocks, n_layers, outer_dim, ...]
    FullyContiguous,

    /// All layers are stored separately.
    /// If outer_contiguous is true, for each layer: [outer_dim, n_blocks, ...]
    /// If outer_contiguous is false, for each layer: [n_blocks, outer_dim, ...]
    /// When outer_dim is 1, these two modes are equivalent.
    LayerSeparate {
        /// If true, the outer dimension is contiguous. Otherwise, the block dimension is contiguous.
        outer_contiguous: bool,
    },
}

/// Local Memory Region
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, Getters)]
pub struct LocalMemoryRegion {
    #[getter(copy)]
    addr: usize,

    #[getter(copy)]
    size: usize,

    #[getter(copy)]
    storage_type: StorageType,
}

/// Core trait for block layouts
pub trait BlockLayout: GenericBlockLayout {
    /// The type of storage this layout uses
    type StorageType: Storage;

    /// Returns the layout type
    fn layout_type(&self) -> LayoutType;

    /// Get the memory regions for all blocks and layers
    fn storage(&self) -> Vec<&Self::StorageType>;

    /// Get the mutable memory regions for all blocks and layers
    fn storage_mut(&mut self) -> Vec<&mut Self::StorageType>;
}

/// Generic trait for block layouts - type-erased on the [Storage] object.
pub trait GenericBlockLayout: BlockLayoutConfig + Send + Sync {
    /// Storage type for the layout
    fn storage_type(&self) -> &StorageType;

    /// Full configuration for the layout
    fn config(&self) -> &LayoutConfig;

    /// Get the memory region for a specific page [page_size, inner_dim]
    ///
    /// # Arguments
    ///
    /// * `block_idx` - The index of the block
    /// * `layer_idx` - The index of the layer
    /// * `outer_idx` - The index of the outer dimension, e.g. if
    ///
    fn memory_region(
        &self,
        block_idx: usize,
        layer_idx: usize,
        outer_idx: usize,
    ) -> Result<LocalMemoryRegion, LayoutError>;
}

/// Configuration for block layouts
pub trait BlockLayoutConfig: std::fmt::Debug {
    /// Returns the layout config
    fn layout_config(&self) -> LayoutConfig;

    /// Returns the total number of blocks this layout manages
    fn num_blocks(&self) -> usize {
        self.layout_config().num_blocks
    }

    /// Returns the number of layers per block
    fn num_layers(&self) -> usize {
        self.layout_config().num_layers
    }

    /// Returns the number of outer dimensions per block
    /// In some cases, K and V might be indexed separately, so in that example one might have 2 outer dimensions
    /// For MLA, this is 1.
    /// The location of the outer dimension in the shape of the tensor layout is defined by the layout type.
    fn outer_dim(&self) -> usize {
        self.layout_config().outer_dim
    }

    /// Returns the size of each block in bytes
    fn page_size(&self) -> usize {
        self.layout_config().page_size
    }

    /// Returns the inner dimension size
    fn inner_dim(&self) -> usize {
        self.layout_config().inner_dim
    }

    /// The size of the data for a layout (pre base_offset)
    fn layout_data_bytes(&self) -> usize;
}

/// Configuration for block layouts
#[derive(Debug, Clone, Builder, Validate, Serialize, Deserialize, PartialEq, Eq)]
pub struct LayoutConfig {
    /// Number of blocks
    #[validate(range(min = 1))]
    pub num_blocks: usize,

    /// Number of layers
    #[validate(range(min = 1))]
    pub num_layers: usize,

    /// Number of outer dimensions
    #[validate(range(min = 1, max = 2))]
    pub outer_dim: usize,

    /// Page size
    #[validate(range(min = 1))]
    pub page_size: usize,

    /// Inner dimension
    #[validate(range(min = 1))]
    pub inner_dim: usize,

    /// Alignment
    #[validate(custom(function = "validate_power_of_2"))]
    #[builder(default = "1")]
    pub alignment: usize,

    /// Data type
    #[builder(default = "2")]
    pub dtype_width_bytes: usize,
}

impl LayoutConfig {
    /// Builder for LayoutConfig
    pub fn builder() -> LayoutConfigBuilder {
        LayoutConfigBuilder::default()
    }
}

/// Internal struct to hold calculated layout dimensions specific to FullyContiguous.
// Module-level, but only used internally by FullyContiguous
#[derive(Debug, Clone, Serialize, Deserialize)]
pub(crate) struct FullyContiguousConfig {
    inner: LayoutConfig,

    /// Minimum contiguous memory region size
    /// Inner dimension * page size * dtype size
    memory_region_size: usize,

    /// Stride between outer dimensions
    outer_dim_stride_in_bytes: usize,

    /// Stride between layers
    layer_stride_in_bytes: usize,

    /// Natural block stride
    natural_block_stride: usize,

    /// Block stride in bytes
    block_stride_in_bytes: usize, // Aligned if necessary

    /// Size of the layout data itself (post base offset)
    layout_data_bytes: usize, // Size of the layout data itself (post base offset)
}

impl FullyContiguousConfig {
    /// Calculates the core dimensions based on the configuration.
    /// Returns an error if the configuration is invalid.
    fn new(config: LayoutConfig) -> Result<Self, LayoutError> {
        // Validate first, propagating errors via `?`
        config.validate()?;

        let alignment = config.alignment;
        let outer_dim_stride_in_bytes =
            config.page_size * config.inner_dim * config.dtype_width_bytes;
        let layer_stride_in_bytes = outer_dim_stride_in_bytes * config.outer_dim;
        let memory_region_size = layer_stride_in_bytes;
        let natural_block_stride = config.num_layers * layer_stride_in_bytes;

        let block_stride_in_bytes = if alignment > 1 {
            align_up(natural_block_stride, alignment)
        } else {
            natural_block_stride
        };

        let layout_data_bytes =
            (config.num_blocks - 1) * block_stride_in_bytes + natural_block_stride;

        Ok(Self {
            inner: config,
            memory_region_size,
            outer_dim_stride_in_bytes,
            layer_stride_in_bytes,
            natural_block_stride,
            block_stride_in_bytes,
            layout_data_bytes,
        })
    }

    /// Calculate the total number of bytes required for allocation, including initial alignment padding.
    /// Panics if the provided configuration is invalid.
    pub fn required_allocation_size(&self) -> usize {
        let initial_padding = self.inner.alignment.saturating_sub(1);
        self.layout_data_bytes + initial_padding
    }
}

impl BlockLayoutConfig for FullyContiguousConfig {
    fn layout_config(&self) -> LayoutConfig {
        self.inner.clone()
    }

    fn layout_data_bytes(&self) -> usize {
        self.layout_data_bytes
    }
}

/// Contiguous memory layout where all blocks and layers are sequential
#[derive(Debug)]
pub struct FullyContiguous<S: Storage> {
    /// Configuration for the layout
    config: FullyContiguousConfig,

    /// Storage for the layoutk
    storage: S,

    /// Storage type for the layout
    storage_type: StorageType,

    // Offset from storage.addr() to the aligned start of block 0
    base_offset: usize,
}

impl<S: Storage> FullyContiguous<S> {
    /// Create a new contiguous layout using the provided configuration and pre-allocated storage.
    /// Performs validation and calculates strides/offsets.
    #[instrument(level = "debug", skip(storage), fields(config = ?config))]
    pub fn new(config: LayoutConfig, mut storage: Vec<S>) -> Result<Self, LayoutError> {
        // Calculate dimensions, which includes validation.
        let config = FullyContiguousConfig::new(config)?;

        if storage.len() != 1 {
            return Err(LayoutError::InvalidConfig(
                "FullyContiguous layout requires exactly one storage region".to_string(),
            ));
        }
        let storage = storage.remove(0);
        let storage_type = storage.storage_type();

        let base_offset = validate_storage(&storage, &config)?;

        tracing::debug!(
            config.memory_region_size,
            config.layer_stride_in_bytes,
            config.block_stride_in_bytes,
            config.natural_block_stride,
            alignment = config.inner.alignment,
            base_offset,
            "Calculated layout strides (aligned)"
        );

        Ok(Self {
            config,
            storage,
            storage_type,
            base_offset,
        })
    }

    /// Internal constructor used for reconstruction from serialized parts.
    /// Assumes the provided config, storage, and base_offset are consistent
    /// and skips size/alignment validation against the storage.
    pub(crate) fn new_internal(
        config: FullyContiguousConfig,
        storage: S,
        storage_type: StorageType,
        base_offset: usize,
    ) -> Result<Self, LayoutError> {
        // Basic check: Ensure the storage address matches expectations based on offset if possible?
        // Maybe not strictly necessary if we trust the serialized data.
        Ok(Self {
            config,
            storage,
            storage_type,
            base_offset,
        })
    }

    /// Allocate storage using the provided allocator and create a new FullyContiguous layout.
    ///
    /// Calculates the required size based on the configuration, allocates the storage
    /// (including potential padding for initial alignment), and then constructs the
    /// `FullyContiguous` layout instance.
    ///
    /// # Type Parameters
    ///
    /// * `A`: The type of the storage allocator, implementing `StorageAllocator<S>`.
    ///
    /// # Arguments
    ///
    /// * `config` - The layout configuration.
    /// * `allocator` - A reference to the storage allocator.
    ///
    /// # Returns
    ///
    /// A `Result` containing the new `FullyContiguous<S>` instance or an error if allocation
    /// or layout creation fails.
    #[instrument(level = "debug", skip(allocator), fields(config = ?config))]
    pub fn allocate(
        config: LayoutConfig,
        allocator: &dyn StorageAllocator<S>,
    ) -> Result<Self, LayoutError> {
        // Calculate total bytes needed. Propagate error if config is invalid.
        let config = FullyContiguousConfig::new(config)?;
        let bytes_to_allocate = config.required_allocation_size();

        tracing::debug!(
            bytes_to_allocate,
            alignment = config.inner.alignment,
            "Calculated storage size for allocation (with alignment padding)"
        );

        let storage = allocator.allocate(bytes_to_allocate).map_err(|e| {
            LayoutError::OperationFailed(format!("Storage allocation failed: {}", e))
        })?;
        tracing::debug!(
            allocated_size = storage.size(),
            allocated_addr = storage.addr(),
            "Storage allocated successfully"
        );

        // Pass the config by value as Self::new takes ownership
        Self::new(config.inner, vec![storage])
    }
}

impl<S: Storage> BlockLayout for FullyContiguous<S> {
    type StorageType = S;

    fn layout_type(&self) -> LayoutType {
        LayoutType::FullyContiguous
    }

    fn storage(&self) -> Vec<&Self::StorageType> {
        vec![&self.storage]
    }

    fn storage_mut(&mut self) -> Vec<&mut Self::StorageType> {
        vec![&mut self.storage]
    }
}

impl<S: Storage> GenericBlockLayout for FullyContiguous<S> {
    fn storage_type(&self) -> &StorageType {
        &self.storage_type
    }

    fn config(&self) -> &LayoutConfig {
        &self.config.inner
    }

    fn memory_region(
        &self,
        block_idx: usize,
        layer_idx: usize,
        outer_idx: usize,
    ) -> Result<LocalMemoryRegion, LayoutError> {
        validate_indices(&self.config, block_idx, layer_idx, outer_idx)?;

        // Start from the aligned base address
        let aligned_start_addr = self.storage.addr() as usize + self.base_offset;

        // Calculate offset relative to the aligned start using stored config
        let block_offset = block_idx * self.config.block_stride_in_bytes;
        let layer_offset = layer_idx * self.config.layer_stride_in_bytes;
        let outer_offset = outer_idx * self.config.outer_dim_stride_in_bytes;
        let final_addr = aligned_start_addr + block_offset + layer_offset + outer_offset;

        Ok(LocalMemoryRegion {
            addr: final_addr,
            size: self.config.memory_region_size,
            storage_type: self.storage_type,
        })
    }
}

impl<S: Storage> BlockLayoutConfig for FullyContiguous<S> {
    fn layout_config(&self) -> LayoutConfig {
        self.config.inner.clone()
    }

    fn layout_data_bytes(&self) -> usize {
        self.config.layout_data_bytes
    }
}

/// Configuration for layer-separated layouts.
/// This is used in vLLM, where every layer has its own allocation.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub(crate) struct LayerSeparateConfig {
    inner: LayoutConfig,

    /// Size of each contiguous memory region
    memory_region_size: usize,

    /// Stride between outer dimensions
    outer_dim_stride_in_bytes: usize,

    /// Block stride in bytes
    block_stride_in_bytes: usize,

    /// Size of the layout data itself (post base offset)
    layout_data_bytes: usize,

    /// Indicator for outer contiguous or block contiguous
    is_outer_contiguous: bool,
}

impl LayerSeparateConfig {
    fn new(config: LayoutConfig, is_outer_contiguous: bool) -> Result<Self, LayoutError> {
        config.validate()?;

        let alignment = config.alignment;
        let memory_region_size = config.page_size * config.inner_dim * config.dtype_width_bytes;

        let outer_dim_stride_in_bytes;
        let block_stride_in_bytes;
        let layout_data_bytes;

        if is_outer_contiguous {
            block_stride_in_bytes = if alignment > 1 {
                align_up(memory_region_size, alignment)
            } else {
                memory_region_size
            };
            outer_dim_stride_in_bytes = block_stride_in_bytes * config.num_blocks;
            layout_data_bytes = outer_dim_stride_in_bytes * config.outer_dim;
        } else {
            outer_dim_stride_in_bytes = memory_region_size;
            let natural_block_stride = outer_dim_stride_in_bytes * config.outer_dim;
            block_stride_in_bytes = if alignment > 1 {
                align_up(natural_block_stride, alignment)
            } else {
                natural_block_stride
            };
            layout_data_bytes = block_stride_in_bytes * config.num_blocks;
        }

        Ok(Self {
            inner: config,
            memory_region_size,
            outer_dim_stride_in_bytes,
            block_stride_in_bytes,
            layout_data_bytes,
            is_outer_contiguous,
        })
    }

    pub fn required_allocation_size(&self) -> usize {
        let initial_padding = self.inner.alignment.saturating_sub(1);
        self.layout_data_bytes + initial_padding
    }
}

impl BlockLayoutConfig for LayerSeparateConfig {
    fn layout_config(&self) -> LayoutConfig {
        self.inner.clone()
    }

    fn layout_data_bytes(&self) -> usize {
        self.layout_data_bytes
    }
}

/// Layer-separated layout where each layer has its own allocation.
#[derive(Debug)]
pub struct LayerSeparate<S: Storage> {
    /// Configuration for the layout
    config: LayerSeparateConfig,

    /// Storage for the layout
    storages: Vec<S>,

    /// Storage type for the layout
    storage_type: StorageType,

    /// Base offset from storage.addr() to the aligned start of block 0
    base_offsets: Vec<usize>,
}

impl<S: Storage> LayerSeparate<S> {
    /// Create a new LayerSeparate layout.
    #[instrument(level = "debug", skip(storages), fields(config = ?config))]
    pub fn new(
        config: LayoutConfig,
        storages: Vec<S>,
        is_outer_contiguous: bool,
    ) -> Result<Self, LayoutError> {
        if storages.len() != config.num_layers {
            return Err(LayoutError::InvalidConfig(
                "LayerSeparate layout requires exactly one storage region per layer".to_string(),
            ));
        }

        let config = LayerSeparateConfig::new(config, is_outer_contiguous)?;

        let storage_type = storages[0].storage_type();
        let mut base_offsets = Vec::new();
        for storage in &storages {
            let base_offset = validate_storage(storage, &config)?;

            tracing::debug!(
                config.memory_region_size,
                config.block_stride_in_bytes,
                config.outer_dim_stride_in_bytes,
                alignment = config.inner.alignment,
                base_offset,
                "Calculated layout strides (aligned)"
            );

            base_offsets.push(base_offset);
        }

        Ok(Self {
            config,
            storages,
            storage_type,
            base_offsets,
        })
    }

    pub(crate) fn new_internal(
        config: LayerSeparateConfig,
        storages: Vec<S>,
        storage_type: StorageType,
        base_offsets: Vec<usize>,
    ) -> Result<Self, LayoutError> {
        Ok(Self {
            config,
            storages,
            storage_type,
            base_offsets,
        })
    }

    /// Allocate a new LayerSeparate layout.
    /// `is_outer_contiguous` determines whether the outer dimension or the block dimension is contiguous.
    /// The amount of [`Storage`]s allocated is equal to the number of layers in the config.
    pub fn allocate(
        config: LayoutConfig,
        allocator: &dyn StorageAllocator<S>,
        is_outer_contiguous: bool,
    ) -> Result<Self, LayoutError> {
        // Calculate total bytes needed. Propagate error if config is invalid.
        let config = LayerSeparateConfig::new(config, is_outer_contiguous)?;
        let bytes_to_allocate = config.required_allocation_size();

        tracing::debug!(
            bytes_to_allocate,
            alignment = config.inner.alignment,
            "Calculated storage size for allocation (with alignment padding)"
        );

        let mut storages = Vec::new();

        for _ in 0..config.inner.num_layers {
            let storage = allocator.allocate(bytes_to_allocate).map_err(|e| {
                LayoutError::OperationFailed(format!("Storage allocation failed: {}", e))
            })?;
            storages.push(storage);
        }

        tracing::debug!(
            allocated_size = storages[0].size(),
            allocated_addr = storages[0].addr(),
            "Storage allocated successfully"
        );

        // Pass the config by value as Self::new takes ownership
        Self::new(config.inner, storages, is_outer_contiguous)
    }
}

impl<S: Storage> GenericBlockLayout for LayerSeparate<S> {
    fn storage_type(&self) -> &StorageType {
        &self.storage_type
    }

    fn config(&self) -> &LayoutConfig {
        &self.config.inner
    }

    fn memory_region(
        &self,
        block_idx: usize,
        layer_idx: usize,
        outer_idx: usize,
    ) -> Result<LocalMemoryRegion, LayoutError> {
        validate_indices(&self.config, block_idx, layer_idx, outer_idx)?;

        // Start from the aligned base address
        let aligned_start_addr =
            self.storages[layer_idx].addr() as usize + self.base_offsets[layer_idx];

        // Calculate offset relative to the aligned start using stored config
        let block_offset = block_idx * self.config.block_stride_in_bytes;
        let outer_offset = outer_idx * self.config.outer_dim_stride_in_bytes;
        let final_addr = aligned_start_addr + block_offset + outer_offset;

        Ok(LocalMemoryRegion {
            addr: final_addr,
            size: self.config.memory_region_size,
            storage_type: self.storages[layer_idx].storage_type(),
        })
    }
}

impl<S: Storage> BlockLayout for LayerSeparate<S> {
    type StorageType = S;

    fn layout_type(&self) -> LayoutType {
        LayoutType::LayerSeparate {
            outer_contiguous: self.config.is_outer_contiguous,
        }
    }

    fn storage(&self) -> Vec<&Self::StorageType> {
        self.storages.iter().collect()
    }

    fn storage_mut(&mut self) -> Vec<&mut Self::StorageType> {
        self.storages.iter_mut().collect()
    }
}

impl<S: Storage> BlockLayoutConfig for LayerSeparate<S> {
    fn layout_config(&self) -> LayoutConfig {
        self.config.inner.clone()
    }

    fn layout_data_bytes(&self) -> usize {
        self.config.layout_data_bytes
    }
}

#[allow(missing_docs)]
#[cfg(test)]
pub mod tests {
    use super::*;
    use crate::block_manager::storage::tests::{NullDeviceAllocator, NullDeviceStorage};
    use crate::block_manager::storage::{StorageType, SystemAllocator};
    use dynamo_runtime::logging::init as init_logging;

    const NUM_BLOCKS: usize = 7;
    const NUM_LAYERS: usize = 5;
    const OUTER_DIM: usize = 2;
    const PAGE_SIZE: usize = 4;
    const INNER_DIM: usize = 13;
    const DTYPE_WIDTH_BYTES: usize = 4;

    /// Helper function to calculate expected memory offset
    fn calculate_expected_offset(
        base_addr: u64,
        block_idx: usize,
        layer_idx: usize,
        block_stride: usize,
        layer_stride: usize,
    ) -> u64 {
        base_addr + (block_idx * block_stride + layer_idx * layer_stride) as u64
    }

    // Updated setup_layout: Calculates size internally, uses default alignment for simplicity in non-alignment tests.
    pub fn setup_layout(
        alignment: Option<usize>, // Option to override default alignment
    ) -> Result<FullyContiguous<NullDeviceStorage>, LayoutError> {
        let config = LayoutConfig {
            num_blocks: NUM_BLOCKS,
            num_layers: NUM_LAYERS,
            outer_dim: OUTER_DIM,
            page_size: PAGE_SIZE,
            inner_dim: INNER_DIM,
            alignment: alignment.unwrap_or(1),
            dtype_width_bytes: DTYPE_WIDTH_BYTES,
        };

        FullyContiguous::allocate(config, &NullDeviceAllocator)
    }

    #[test]
    fn test_fc_creation_invalid_alignment() {
        let config = LayoutConfig::builder()
            .num_blocks(NUM_BLOCKS)
            .num_layers(NUM_LAYERS)
            .outer_dim(OUTER_DIM)
            .page_size(PAGE_SIZE)
            .inner_dim(INNER_DIM)
            .alignment(3)
            .build()
            .unwrap();

        assert!(config.validate().is_err());
    }

    #[test]
    fn test_fc_creation_success() {
        // Setup with default (None) alignment
        let layout_result = setup_layout(None);
        assert!(
            layout_result.is_ok(),
            "Layout creation failed: {:?}",
            layout_result.err()
        );
    }

    #[test]
    fn test_fc_creation_insufficient_storage() {
        init_logging();
        let config = LayoutConfig {
            num_blocks: NUM_BLOCKS,
            num_layers: NUM_LAYERS,
            outer_dim: OUTER_DIM,
            page_size: PAGE_SIZE,
            inner_dim: INNER_DIM,
            alignment: 1,
            dtype_width_bytes: DTYPE_WIDTH_BYTES,
        };
        // Calculate correct size needed
        let fc_config = FullyContiguousConfig::new(config.clone()).unwrap();
        let required_size = fc_config.required_allocation_size();
        let storage = NullDeviceStorage::new((required_size - 1) as u64);
        let layout_result = FullyContiguous::new(config, vec![storage]);

        assert!(layout_result.is_err());
        match layout_result.err().unwrap() {
            LayoutError::InvalidConfig(_) => {} // Expected error
            e => panic!("Expected InvalidConfig error, got {:?}", e),
        }
    }

    #[test]
    fn test_fc_accessor_methods() {
        let layout = setup_layout(None).expect("Layout setup failed");

        assert_eq!(layout.num_blocks(), NUM_BLOCKS);
        assert_eq!(layout.num_layers(), NUM_LAYERS);
        assert_eq!(layout.outer_dim(), OUTER_DIM);
        assert_eq!(layout.page_size(), PAGE_SIZE);
        assert_eq!(layout.inner_dim(), INNER_DIM);
    }

    #[test]
    fn test_fc_offset_calculation() {
        let layout = setup_layout(None).expect("Layout setup failed");

        let dims = layout.config.clone();
        let block_stride = dims.block_stride_in_bytes;
        let layer_stride = dims.layer_stride_in_bytes;
        let base_addr = layout.storage.addr() + layout.base_offset as u64;

        // Test first block, first layer
        let expected_offset_0_0 =
            calculate_expected_offset(base_addr, 0, 0, block_stride, layer_stride);
        assert_eq!(
            layout.memory_region(0, 0, 0).unwrap().addr as u64,
            expected_offset_0_0
        );

        // Test first block, last layer
        let last_layer_idx = NUM_LAYERS - 1;
        let expected_offset_0_last =
            calculate_expected_offset(base_addr, 0, last_layer_idx, block_stride, layer_stride);
        assert_eq!(
            layout.memory_region(0, last_layer_idx, 0).unwrap().addr as u64,
            expected_offset_0_last
        );

        // Test last block, first layer
        let last_block_idx = NUM_BLOCKS - 1;
        let expected_offset_last_0 =
            calculate_expected_offset(base_addr, last_block_idx, 0, block_stride, layer_stride);
        assert_eq!(
            layout.memory_region(last_block_idx, 0, 0).unwrap().addr as u64,
            expected_offset_last_0
        );

        // Test last block, last layer
        let expected_offset_last_last = calculate_expected_offset(
            base_addr,
            last_block_idx,
            last_layer_idx,
            block_stride,
            layer_stride,
        );
        assert_eq!(
            layout
                .memory_region(last_block_idx, last_layer_idx, 0)
                .unwrap()
                .addr as u64,
            expected_offset_last_last
        );

        // Test intermediate block/layer
        let mid_block_idx = NUM_BLOCKS / 2;
        let mid_layer_idx = NUM_LAYERS / 2;
        let expected_offset_mid_mid = calculate_expected_offset(
            base_addr,
            mid_block_idx,
            mid_layer_idx,
            block_stride,
            layer_stride,
        );
        assert_eq!(
            layout
                .memory_region(mid_block_idx, mid_layer_idx, 0)
                .unwrap()
                .addr as u64,
            expected_offset_mid_mid
        );
    }

    #[test]
    fn test_fc_invalid_block_index() {
        let layout = setup_layout(None).expect("Layout setup failed");
        let result = layout.memory_region(NUM_BLOCKS, 0, 0); // Index == num_blocks (out of bounds)
        assert!(result.is_err());
        assert!(matches!(
            result.err().unwrap(),
            LayoutError::InvalidBlockIndex(NUM_BLOCKS)
        ));
    }

    #[test]
    fn test_fc_invalid_layer_index() {
        let layout = setup_layout(None).expect("Layout setup failed");
        let result = layout.memory_region(0, NUM_LAYERS, 0); // Index == num_layers (out of bounds)
        assert!(result.is_err());
        assert!(matches!(
            result.err().unwrap(),
            LayoutError::InvalidLayerIndex(NUM_LAYERS)
        ));
    }

    #[test]
    fn test_fc_invalid_outer_index() {
        let layout = setup_layout(None).expect("Layout setup failed");
        let result = layout.memory_region(0, 0, OUTER_DIM); // Index == num_outer_dims (out of bounds)
        assert!(result.is_err());
        assert!(matches!(
            result.err().unwrap(),
            LayoutError::InvalidOuterIndex(OUTER_DIM)
        ));
    }

    #[test]
    fn test_fc_allocation_system() {
        init_logging();
        let config = LayoutConfig {
            num_blocks: NUM_BLOCKS,
            num_layers: NUM_LAYERS,
            outer_dim: OUTER_DIM,
            page_size: PAGE_SIZE,
            inner_dim: INNER_DIM,
            alignment: 1,
            dtype_width_bytes: DTYPE_WIDTH_BYTES,
        };

        let allocator = SystemAllocator;
        let layout_result = FullyContiguous::allocate(config, &allocator);

        assert!(layout_result.is_ok());
        let layout = layout_result.unwrap();

        // Basic checks on the allocated layout
        assert_eq!(layout.num_blocks(), NUM_BLOCKS);
        assert_eq!(layout.num_layers(), NUM_LAYERS);
        assert_eq!(layout.page_size(), PAGE_SIZE);
        assert_eq!(layout.inner_dim(), INNER_DIM);
        assert_eq!(layout.storage.storage_type(), StorageType::System);
        assert_eq!(
            layout.storage.size(),
            layout.config.required_allocation_size()
        );

        assert_eq!(
            layout.storage.size(),
            NUM_BLOCKS * NUM_LAYERS * OUTER_DIM * PAGE_SIZE * INNER_DIM * DTYPE_WIDTH_BYTES
        );
    }

    #[test]
    fn test_fc_alignment() {
        init_logging();
        const ALIGNMENT: usize = 256; // Must be power of 2

        let config = LayoutConfig {
            num_blocks: NUM_BLOCKS,
            num_layers: NUM_LAYERS,
            outer_dim: OUTER_DIM,
            page_size: PAGE_SIZE,
            inner_dim: INNER_DIM,
            alignment: ALIGNMENT,
            dtype_width_bytes: DTYPE_WIDTH_BYTES,
        };

        // Calculate expected size needed *for the data layout itself*
        let memory_region_size = PAGE_SIZE * INNER_DIM * DTYPE_WIDTH_BYTES;
        assert_eq!(memory_region_size, 208);

        let natural_block_stride = OUTER_DIM * NUM_LAYERS * memory_region_size;
        assert_eq!(natural_block_stride, 2080);

        let aligned_block_stride = align_up(natural_block_stride, ALIGNMENT);
        assert_eq!(aligned_block_stride, 2304);

        // Calculate the expected *allocated* size (data + initial padding)
        let fc_config = FullyContiguousConfig::new(config.clone()).unwrap();
        let expected_allocated_size = fc_config.required_allocation_size();

        // Use allocate method
        let allocator = SystemAllocator;
        let layout_result = FullyContiguous::allocate(config.clone(), &allocator);

        assert!(
            layout_result.is_ok(),
            "Allocation failed: {:?}",
            layout_result.err()
        );
        let layout = layout_result.unwrap();

        // Verify total *allocated* size matches expectation
        assert_eq!(
            layout.storage.size(),
            expected_allocated_size,
            "Allocated storage size mismatch"
        );
        assert_eq!(
            layout.config.block_stride_in_bytes, aligned_block_stride,
            "Stored block stride mismatch"
        );

        // Check alignment of block starts
        let addr_block_0 = layout
            .memory_region(0, 0, 0)
            .expect("Failed to get addr block 0");
        let addr_block_1 = layout
            .memory_region(1, 0, 0)
            .expect("Failed to get addr block 1");
        let addr_block_2 = layout
            .memory_region(2, 0, 0)
            .expect("Failed to get addr block 2");

        // All blocks should now be aligned due to base_offset adjustment
        assert_eq!(
            addr_block_0.addr as u64 % ALIGNMENT as u64,
            0,
            "Block 0 start address is not aligned"
        );
        assert_eq!(
            addr_block_1.addr as u64 % ALIGNMENT as u64,
            0,
            "Block 1 start address is not aligned"
        );
        assert_eq!(
            addr_block_2.addr as u64 % ALIGNMENT as u64,
            0,
            "Block 2 start address is not aligned"
        );

        // Verify the difference matches the aligned stride
        assert_eq!(
            addr_block_1.addr as u64 - addr_block_0.addr as u64,
            aligned_block_stride as u64,
            "Stride between block 0 and 1 mismatch"
        );
        assert_eq!(
            addr_block_2.addr as u64 - addr_block_1.addr as u64,
            aligned_block_stride as u64,
            "Stride between block 1 and 2 mismatch"
        );
    }

    // LayerSeparate Tests

    /// Helper function to setup LayerSeparate layout with specified configuration
    pub fn setup_layer_separate_layout(
        alignment: Option<usize>,
        is_outer_contiguous: bool,
    ) -> Result<LayerSeparate<NullDeviceStorage>, LayoutError> {
        let config = LayoutConfig {
            num_blocks: NUM_BLOCKS,
            num_layers: NUM_LAYERS,
            outer_dim: OUTER_DIM,
            page_size: PAGE_SIZE,
            inner_dim: INNER_DIM,
            alignment: alignment.unwrap_or(1),
            dtype_width_bytes: DTYPE_WIDTH_BYTES,
        };

        // Create one storage per layer
        let ls_config = LayerSeparateConfig::new(config.clone(), is_outer_contiguous)?;
        let required_size = ls_config.required_allocation_size();
        let mut storages = Vec::new();
        for _ in 0..NUM_LAYERS {
            storages.push(NullDeviceStorage::new(required_size as u64));
        }

        LayerSeparate::new(config, storages, is_outer_contiguous)
    }

    #[test]
    fn test_ls_creation_success_outer_contiguous() {
        let layout_result = setup_layer_separate_layout(None, true);
        assert!(
            layout_result.is_ok(),
            "LayerSeparate creation failed: {:?}",
            layout_result.err()
        );

        let layout = layout_result.unwrap();
        assert_eq!(
            layout.layout_type(),
            LayoutType::LayerSeparate {
                outer_contiguous: true
            }
        );
    }

    #[test]
    fn test_ls_creation_success_block_contiguous() {
        let layout_result = setup_layer_separate_layout(None, false);
        assert!(
            layout_result.is_ok(),
            "LayerSeparate creation failed: {:?}",
            layout_result.err()
        );

        let layout = layout_result.unwrap();
        assert_eq!(
            layout.layout_type(),
            LayoutType::LayerSeparate {
                outer_contiguous: false
            }
        );
    }

    #[test]
    fn test_ls_creation_wrong_storage_count() {
        let config = LayoutConfig {
            num_blocks: NUM_BLOCKS,
            num_layers: NUM_LAYERS,
            outer_dim: OUTER_DIM,
            page_size: PAGE_SIZE,
            inner_dim: INNER_DIM,
            alignment: 1,
            dtype_width_bytes: DTYPE_WIDTH_BYTES,
        };

        // Create wrong number of storages (should be NUM_LAYERS, but provide NUM_LAYERS - 1)
        let mut storages = Vec::new();
        for _ in 0..(NUM_LAYERS - 1) {
            storages.push(NullDeviceStorage::new(1000));
        }

        let layout_result = LayerSeparate::new(config, storages, true);
        assert!(layout_result.is_err());
        match layout_result.err().unwrap() {
            LayoutError::InvalidConfig(_) => {} // Expected error
            e => panic!("Expected InvalidConfig error, got {:?}", e),
        }
    }

    #[test]
    fn test_ls_accessor_methods() {
        let layout = setup_layer_separate_layout(None, true).expect("Layout setup failed");

        assert_eq!(layout.num_blocks(), NUM_BLOCKS);
        assert_eq!(layout.num_layers(), NUM_LAYERS);
        assert_eq!(layout.outer_dim(), OUTER_DIM);
        assert_eq!(layout.page_size(), PAGE_SIZE);
        assert_eq!(layout.inner_dim(), INNER_DIM);
        assert_eq!(layout.storage().len(), NUM_LAYERS);
        assert_eq!(layout.storage_type(), &StorageType::Null);
    }

    #[test]
    fn test_ls_memory_region_outer_contiguous() {
        let layout = setup_layer_separate_layout(None, true).expect("Layout setup failed");

        // Test accessing different blocks within the same layer
        let region_0_0_0 = layout.memory_region(0, 0, 0).unwrap();
        let region_1_0_0 = layout.memory_region(1, 0, 0).unwrap();

        // In outer_contiguous mode, blocks are sequential within each layer
        let expected_block_stride = layout.config.block_stride_in_bytes;
        assert_eq!(
            region_1_0_0.addr - region_0_0_0.addr,
            expected_block_stride,
            "Block stride mismatch in outer_contiguous mode"
        );

        // Test accessing different outer dimensions
        let region_0_0_1 = layout.memory_region(0, 0, 1).unwrap();
        let expected_outer_stride = layout.config.outer_dim_stride_in_bytes;
        assert_eq!(
            region_0_0_1.addr - region_0_0_0.addr,
            expected_outer_stride,
            "Outer dimension stride mismatch"
        );

        // Test accessing different layers (should be in different storage)
        let region_0_1_0 = layout.memory_region(0, 1, 0).unwrap();
        let region_0_0_0_storage_addr = layout.storages[0].addr() as usize + layout.base_offsets[0];
        let region_0_1_0_storage_addr = layout.storages[1].addr() as usize + layout.base_offsets[1];

        assert_eq!(region_0_0_0.addr, region_0_0_0_storage_addr);
        assert_eq!(region_0_1_0.addr, region_0_1_0_storage_addr);
    }

    #[test]
    fn test_ls_memory_region_block_contiguous() {
        let layout = setup_layer_separate_layout(None, false).expect("Layout setup failed");

        // Test accessing different blocks within the same layer
        let region_0_0_0 = layout.memory_region(0, 0, 0).unwrap();
        let region_1_0_0 = layout.memory_region(1, 0, 0).unwrap();

        // In block_contiguous mode, blocks have different stride calculation
        let expected_block_stride = layout.config.block_stride_in_bytes;
        assert_eq!(
            region_1_0_0.addr - region_0_0_0.addr,
            expected_block_stride,
            "Block stride mismatch in block_contiguous mode"
        );

        // Test accessing different outer dimensions within same block
        let region_0_0_1 = layout.memory_region(0, 0, 1).unwrap();
        let expected_outer_stride = layout.config.outer_dim_stride_in_bytes;
        assert_eq!(
            region_0_0_1.addr - region_0_0_0.addr,
            expected_outer_stride,
            "Outer dimension stride mismatch in block_contiguous mode"
        );
    }

    #[test]
    fn test_ls_invalid_indices() {
        let layout = setup_layer_separate_layout(None, true).expect("Layout setup failed");

        // Test invalid block index
        let result = layout.memory_region(NUM_BLOCKS, 0, 0);
        assert!(result.is_err());
        assert!(matches!(
            result.err().unwrap(),
            LayoutError::InvalidBlockIndex(NUM_BLOCKS)
        ));

        // Test invalid layer index
        let result = layout.memory_region(0, NUM_LAYERS, 0);
        assert!(result.is_err());
        assert!(matches!(
            result.err().unwrap(),
            LayoutError::InvalidLayerIndex(NUM_LAYERS)
        ));

        // Test invalid outer index
        let result = layout.memory_region(0, 0, OUTER_DIM);
        assert!(result.is_err());
        assert!(matches!(
            result.err().unwrap(),
            LayoutError::InvalidOuterIndex(OUTER_DIM)
        ));
    }

    #[test]
    fn test_ls_memory_region_size() {
        let layout = setup_layer_separate_layout(None, true).expect("Layout setup failed");

        let region = layout.memory_region(0, 0, 0).unwrap();
        let expected_size = PAGE_SIZE * INNER_DIM * DTYPE_WIDTH_BYTES;

        assert_eq!(region.size, expected_size);
    }

    #[test]
    fn test_ls_all_blocks_layers_accessible() {
        let layout = setup_layer_separate_layout(None, true).expect("Layout setup failed");

        // Test that we can access all valid combinations of indices
        for block_idx in 0..NUM_BLOCKS {
            for layer_idx in 0..NUM_LAYERS {
                for outer_idx in 0..OUTER_DIM {
                    let result = layout.memory_region(block_idx, layer_idx, outer_idx);
                    assert!(
                        result.is_ok(),
                        "Failed to access block {}, layer {}, outer {}: {:?}",
                        block_idx,
                        layer_idx,
                        outer_idx,
                        result.err()
                    );
                }
            }
        }
    }

    #[test]
    fn test_ls_storage_mutability() {
        let mut layout = setup_layer_separate_layout(None, true).expect("Layout setup failed");

        // Test that we can get mutable references to storage
        let mut_storages = layout.storage_mut();
        assert_eq!(mut_storages.len(), NUM_LAYERS);

        // Verify each storage is accessible
        for (i, storage) in mut_storages.iter().enumerate() {
            assert!(storage.size() > 0, "Storage {} has zero size", i);
        }
    }

    #[test]
    fn test_ls_alignment() {
        init_logging();
        const ALIGNMENT: usize = 128; // Must be power of 2

        let config = LayoutConfig {
            num_blocks: NUM_BLOCKS,
            num_layers: NUM_LAYERS,
            outer_dim: OUTER_DIM,
            page_size: PAGE_SIZE,
            inner_dim: INNER_DIM,
            alignment: ALIGNMENT,
            dtype_width_bytes: DTYPE_WIDTH_BYTES,
        };

        // Create storages with sufficient size
        let ls_config = LayerSeparateConfig::new(config.clone(), true).unwrap();
        let required_size = ls_config.required_allocation_size();
        let mut storages = Vec::new();
        for _ in 0..NUM_LAYERS {
            storages.push(NullDeviceStorage::new(required_size as u64));
        }

        let layout_result = LayerSeparate::new(config, storages, true);
        assert!(
            layout_result.is_ok(),
            "Layout creation with alignment failed"
        );

        let layout = layout_result.unwrap();

        // Check that block addresses are properly aligned within each layer
        for layer_idx in 0..NUM_LAYERS {
            let addr_block_0 = layout.memory_region(0, layer_idx, 0).unwrap();
            let addr_block_1 = layout.memory_region(1, layer_idx, 0).unwrap();

            // First block should be aligned
            assert_eq!(
                addr_block_0.addr % ALIGNMENT,
                0,
                "Block 0 in layer {} is not aligned",
                layer_idx
            );

            // Subsequent blocks should maintain alignment
            assert_eq!(
                addr_block_1.addr % ALIGNMENT,
                0,
                "Block 1 in layer {} is not aligned",
                layer_idx
            );
        }
    }

    #[test]
    fn test_ls_stride_calculations_outer_contiguous() {
        let layout = setup_layer_separate_layout(None, true).expect("Layout setup failed");

        let memory_region_size = PAGE_SIZE * INNER_DIM * DTYPE_WIDTH_BYTES;

        // In outer_contiguous mode:
        // outer_dim_stride = block_stride * num_blocks
        // block_stride = memory_region_size (aligned)
        assert_eq!(layout.config.memory_region_size, memory_region_size);
        assert_eq!(layout.config.block_stride_in_bytes, memory_region_size); // No alignment needed
        assert_eq!(
            layout.config.outer_dim_stride_in_bytes,
            layout.config.block_stride_in_bytes * NUM_BLOCKS
        );
    }

    #[test]
    fn test_ls_stride_calculations_block_contiguous() {
        let layout = setup_layer_separate_layout(None, false).expect("Layout setup failed");

        let memory_region_size = PAGE_SIZE * INNER_DIM * DTYPE_WIDTH_BYTES;

        // In block_contiguous mode:
        // outer_dim_stride = memory_region_size
        // block_stride = outer_dim_stride * outer_dim (aligned)
        assert_eq!(layout.config.memory_region_size, memory_region_size);
        assert_eq!(layout.config.outer_dim_stride_in_bytes, memory_region_size);
        assert_eq!(
            layout.config.block_stride_in_bytes,
            memory_region_size * OUTER_DIM
        );
    }

    #[test]
    fn test_ls_layout_data_bytes() {
        let layout_outer = setup_layer_separate_layout(None, true).expect("Layout setup failed");
        let layout_block = setup_layer_separate_layout(None, false).expect("Layout setup failed");

        // For outer_contiguous: layout_data_bytes = outer_dim_stride * outer_dim
        let expected_outer = layout_outer.config.outer_dim_stride_in_bytes * OUTER_DIM;
        assert_eq!(layout_outer.layout_data_bytes(), expected_outer);

        // For block_contiguous: layout_data_bytes = block_stride * num_blocks
        let expected_block = layout_block.config.block_stride_in_bytes * NUM_BLOCKS;
        assert_eq!(layout_block.layout_data_bytes(), expected_block);
    }

    #[test]
    fn test_ls_allocate() {
        let config = LayoutConfig {
            num_blocks: NUM_BLOCKS,
            num_layers: NUM_LAYERS,
            outer_dim: OUTER_DIM,
            page_size: PAGE_SIZE,
            inner_dim: INNER_DIM,
            alignment: 1,
            dtype_width_bytes: DTYPE_WIDTH_BYTES,
        };

        LayerSeparate::allocate(config, &NullDeviceAllocator, true)
            .expect("Layout allocation failed");
    }
}