Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add support for new logical sharding + alignment in TensorLayout and tensor creation #14771

Open
wants to merge 2 commits into
base: main
Choose a base branch
from

Commits on Nov 14, 2024

  1. #13127: Add support for new logical sharding + alignment in TensorLay…

    …out and tensor creation
    
    - Add ShardMode enum to specify shard shape in shard spec as either physical or logical
      * ShardMode::PHYSICAL: This is current behaviour that we will deprecate!
        ** It is less expressive than using shard shape as logical (ie. must be tile aligned for TILE layout etc...)
        ** It fundamentally operates on padded shape and is confusing and incompatible with logical shape
      * ShardMode::LOGICAL: Shard shape cuts 2D logical shape and each shard is aligned after
        ** Without alignment restrictions, you can cut 2D logical shape more arbitrarily
        ** Existing sharding can be switched over to this entirely (just need codeowners to help out and flip...)
      * Default everywhere will be ShardMode::PHYSICAL with a warning message
    - Switch tests/ttnn/unit_tests/operations/test_paged_update_cache.py to use logical shard shape as an example
      * Introduce tensor.logical_volume() (as opposed to tensor.volume() which returns physical volume based on padded shape)
      * TODO: Rename volume() -> physical_volume() and logical_volume() -> volume()
    - Add new c++ tests to test tensor creation with logical shard shape + alignment
      * IMPORTANT: Need to update host data manipulation to be aware of new logical sharding for use from python!
    
    To support these changes, some changes to TensorLayout:
    - Make private TensorLayout constructor with alignment public with these changes:
      * legacyShapeToAlignment will try to return 2D alignment if possible (ie. only padding on height/width)
        ** Goal is to transition alignment to be 2D only if we remove poor use cases of padding on non-height/width dims
      * legacyShapeToAlignment is only expected to be used for ShardMode::PHYSICAL and uses default alignment for sharded tensors
        ** Before interleaved or sharded will just use padded shape for alignment
        ** One exception is for row major sharded tensors where we use shard width if shape is padded;
           otherwise, we only take shard width for BLOCK/WIDTH sharded cases and original physical shape for HEIGHT sharded
      * legacyShapeToAlignment (and alignment in general) will work iff there is only padding on height and/or width
        ** IMPORTANT: This means we are expecting tensors with arbitrary padding along non-height/width to be interleaved only!
    - If ShardMode::LOGICAL:
      * In TensorLayout::compute_shard_spec_buffer, calculate physical shard shape based on shard shape + alignment
      * In TensorLayout::compute_physical_shape, calculate physical shape based on number of logical shards
    - Clean up handling of sharded tensors and error messages in ttnn/cpp/ttnn/tensor/layout/page_config.cpp
    - Add Size constructor for std::array<uint32_t, 2>
    TT-BrianLiu committed Nov 14, 2024
    Configuration menu
    Copy the full SHA
    f7354bf View commit details
    Browse the repository at this point in the history

Commits on Nov 15, 2024

  1. #0: clean up

    TT-BrianLiu committed Nov 15, 2024
    Configuration menu
    Copy the full SHA
    17af0d5 View commit details
    Browse the repository at this point in the history