@torch.compile materializes broadcasted tensors, resulting in huge allocations on GPU

[cbuchner1]

Hi,

after running a memory trace on the End 2 End example notebook with the techniques here, I am finding that the pyTorch graph compiler is a lot less likely to eliminate broadcast tensors from being materialized by fusing expressions than TensorFlow. In other words, the size of simulation scenarios in the SLS is suffering from restrictions that I haven't observed with TensorFlow and Sionna 1.x

In that End2End downlink example, I am seeing two places in the code in this notebook that results in huge allocation spikes, using torch 2.10.0

1. src/sionna/phy/channel/utils.py line 297
   h_f = h * e

This broadcast h * e expands to P times the size of the reduced product h_f = h_f.sum(dim=-3) by being materialized on the GPU with P representing the number of paths used by the channel model (e.g. often 24 for UMa).

2. src/sionna/phy/ofdm/precoding.py line 352
   h_eff = h @ g

For num_ofdm_symbols = 14 in that End 2 End example, the input and output tensors have these shapes

  h shape:     (1, 210, 21, 14, 128,  1,  12) ~0.707 GiB
  g shape:     (1,   1, 21, 14, 128, 12,  10) ~0.034 GiB
  h_eff shape: (1, 210, 21, 14,  128, 1,  10) ~0.589 GiB
  h_eff shape: (1, 210,  1, 21,  10, 14, 128) ~0.589 GiB  (after permute)

The broadcast performed during the multiplication is fully materialized on the GPU even in code running under @torch.compile and requires 210 times the size of the input tensor g (because it's being broadcast 210 times), making a ~7GiB allocation.

Would the Sionna team consider deploying workarounds to prevent the materialization of such broadcasts, or do you consider this a problem of pyTorch which should be fixed on their end?

[cbuchner1]

Possible workarounds to reduce the memory spikes are

1. This sequence of statements in cir_to_ofdm_channel()
   h = a.unsqueeze(-1)
   h_f = h * e
   h_f = h_f.sum(dim=-3)
   can possibly be replaced by a single statement
   h_f = torch.einsum("...pn,...pf->...nf", a, e.squeeze(-2))

The above modification is a suggestion by gpt-5.4-high, and it claims that "I verified the new version is lint-clean, matches the old implementation numerically for both tau shapes supported by cir_to_ofdm_channel(), and still works under torch.compile.".

2. In compute_effective_channel() of PrecodedChannel as well as _compute_effective_channel() of RZFPrecoder
   the sequence of statements
   g = g.unsqueeze(1)
   h_eff = h @ g
   can possibly be replaced by a single statement
   h_eff = torch.einsum("brtofxy,btofys->brtofxs", h, g)
   completely avoiding materialization of the broadcasted version of g

It does greatly reduce the output of torch.cuda.max_memory_allocated() for me, but I hate how much less readable this version is.

I do get a minor speedup and a nearly 4x reduction in peak GPU memory usage with the E2E sample, for 14 OFDM symbols and 100 slots simulation duration. Results are very similar (different rounding, but no major changes)

Before patch:
torch.cuda.max_memory_allocated(): 18.3148 GiB user 1m32.461s
After patch:
torch.cuda.max_memory_allocated(): 4.9795 GiB user 1m27.398s

Note that TensorFlow also has a tf.einsum() function and when running Sionna v1.x code in Eager mode this also materializes such large broadcast tensors. It may be worthwhile to also apply this workaround to Sionna v1.x.