submission 524332

+ import os
+ os.environ['CUBLAS_WORKSPACE_CONFIG'] = ':4096:8'
+  
  import torch
  import torch.nn.functional as F
  from task import input_t, output_t
  ⋯ 12 unchanged lines
          n += 1
   
   
- # FFT is fast when W_fft tensor is small. W_fft size = Co*Ci * fH * fW2 * 8 bytes.
- # For C=128, k=16: 16384 * 288 * 145 * 8 = 5.1 GB -- memory bandwidth kills the gain.
- # For C=64,  k=16: 4096  * 144 * 73  * 8 = 344 MB -- fast.
- # Threshold: only use FFT when Co*Ci*fH*fW2 * 8 bytes < ~512 MB.
- _FFT_MAX_KSIZE = 16
- _FFT_MAX_CHANNELS = 64  # above this, W_fft memory cost dominates
-  
-  
  def _fft_conv2d(input_tensor: torch.Tensor, kernel: torch.Tensor, output: torch.Tensor) -> torch.Tensor:
      B, Ci, H, W = input_tensor.shape
      Co, _, kH, kW = kernel.shape
      H_out, W_out = H - kH + 1, W - kW + 1
   
-     # FFT size must be >= H+kH-1 to compute linear (non-circular) convolution
      fH = _next_fast_len(H + kH - 1)
      fW = _next_fast_len(W + kW - 1)
      fW2 = fW // 2 + 1
   
-     # Input FFT: (B, Ci, H, W) -> (B, Ci, fH, fW2) complex64
      X_fft = torch.fft.rfft2(input_tensor, s=(fH, fW))
-  
-     # Kernel FFT: batch all Co*Ci kernels in one cuFFT call
-     # (Co, Ci, kH, kW) -> (Co*Ci, kH, kW) -> rfft2 -> (Co, Ci, fH, fW2) complex64
      W_fft = torch.fft.rfft2(
          kernel.reshape(Co * Ci, kH, kW), s=(fH, fW)
      ).reshape(Co, Ci, fH, fW2)
   
-     # Cross-correlation in frequency domain:
-     # out[b, o, h, w] = sum_i X[b, i, h, w] * conj(W[o, i, h, w])
      out_fft = torch.einsum('bihw,oihw->bohw', X_fft, W_fft.conj())
-  
-     # IFFT and slice valid region
-     result = torch.fft.irfft2(out_fft, s=(fH, fW))
-     output[...] = result[:, :, :H_out, :W_out]
+     output[...] = torch.fft.irfft2(out_fft, s=(fH, fW))[:, :, :H_out, :W_out]
      return output
   
   
  ⋯ 1 unchanged lines
      input_tensor, kernel, output = data
      Co, Ci, kH, kW = kernel.shape
   
-     use_fft = (kW <= _FFT_MAX_KSIZE) and (Ci <= _FFT_MAX_CHANNELS)
-     if use_fft:
-         # FFT convolution: fast when W_fft fits in memory (~344 MB for C=64,k=16)
+     if kW <= 16:
+         # FFT convolution. For k<=16, FFT error is ~2e-5 (well within 1e-3 tolerance).
+         # W_fft is at most (128^2) * 288 * 145 * 8 = 5.1 GB which fits in A100 80GB.
          return _fft_conv2d(input_tensor, kernel, output)
-     elif kW <= _FFT_MAX_KSIZE:
-         # Large channels (C=128, k=16): W_fft would be 5.1 GB.
-         # Plain cuDNN is faster here.
-         output[...] = F.conv2d(input_tensor, kernel, stride=1, padding=0)
-         return output
      else:
-         # Large kernels (k=32): FFT float32 error too high; must match reference.
+         # k=32: FFT float32 error (~4e-3) exceeds atol=1e-3 for near-zero elements.
+         # Must use DeterministicContext to reproduce the reference's exact float32 result.
          with DeterministicContext():
              output[...] = F.conv2d(input_tensor, kernel, stride=1, padding=0)
          return output