submission 545404

import os
os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":4096:8"

import torch
import torch.nn.functional as F
from task import input_t, output_t

# Key insight: deterministic=True is SLOW! The reference uses DeterministicContext
# to generate the reference output, but our check has rtol=1e-3, atol=1e-3
# Non-deterministic FP32 conv2d should be within that tolerance
torch.backends.cudnn.allow_tf32 = False
torch.backends.cuda.matmul.allow_tf32 = False
torch.backends.cudnn.deterministic = False  # FAST non-deterministic
torch.backends.cudnn.benchmark = True  # Auto-tune algorithm

def custom_kernel(data: input_t) -> output_t:
    input_tensor, kernel, output = data
    result = F.conv2d(input_tensor, kernel, stride=1, padding=0)
    output[...] = result
    return output