SSR functional full network

[ign-krishnanand]

What's changed

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

Results

PCC Score

Hardware: Wormhole n150:

Output	Random weights and input data
SR output	0.9999265025206074
Features output	0.9897763897274661

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

[dvartaniansTT]

Please add a README that includes a brief intro to the model.
instructions to set up the proper env to run pytests and the command to run them and what each test does.
also include the device and end to perf numbers in the README. use this as an example

[dvartaniansTT]

LGTM

[ign-saurav]

Hi @mbezuljTT, Could you please review the PR?