What's changed

Added TT implementation of Selective Super Resolution (SSR) model in tt-metal.

Results

PCC Score

Hardware: Wormhole n150:

Note

Code refactoring and optimisation are still in progress.

Torch Fallback usage in code

• Tensor unpadding (window_attn_tr.py, OCAB.py):

  • ttnn.transformer operation requirement: Head size must be a multiple of 32.
  • Weight and Bias padding to the initial linear layer, ensures compatibility with ttnn.transformer operations.
  • resulting Tensor output from ttnn.transformer operations is unpaded.

• Pixel Shuffle Optimization (upsample.py):

  • Used torch.nn.functional.pixel_shuffle instead of manual reshapes and permute operations.
  • Provides significant performance improvements during upsampling.

• Bicubic Upsample (ssr.py):

  • TTNN currently does not support bicubic interpolation mode for the ttnn.upsample operation.
    Approximation with ttnn upsample with mode as bilinear was resulting in lower pcc.

• Window Reverse (test_ssr.py):

  • Window reverse was performed on the host for SSR without convolution, as it is the final layer and can be treated as post-processing.

Checklist

• All post commit CI passes
• Blackhole Post commit CI with demo tests passes (if applicable)
• Model regression CI passes (if applicable)
• Device performance regression CI passes (if applicable)
• (For models and ops writers) Single-card demo tests CI passes (if applicable) See recommended dev flow.
• Galaxy quick CI passes (if applicable)
• Galaxy demo tests, for Llama CI passes, if applicable, because of current Llama work
• (For runtime and ops writers) T3000 unit tests CI passes (if applicable, since this is run on push to main)
• (For models and ops writers) T3000 demo tests CI passes (if applicable, since this is required for release)
• New/Existing tests provide coverage for changes