Fix 02-fused-softmax tutorial on BMG

Signed-off-by: Anatoly Myachev <[email protected]>

Author: anmyachev

python/tutorials/02-fused-softmax.py

@@ -126,10 +126,10 @@ def softmax_kernel(output_ptr, input_ptr, input_row_stride, output_row_stride, n
 max_num_resident_warps = NUM_SM * warps_per_sm
 kernels = {}
 # Possible SLM allocation sizes in kB
-tg_slm_sizes = [i * 2**i for i in [0, 1, 2, 4, 8, 16, 24, 32, 48, 64, 96, 128]]  # TODO: Get from properties
+tg_slm_sizes = [2**i for i in [0, 1, 2, 4, 8, 16, 24, 32, 48, 64, 96, 128]]  # TODO: Get from properties
 
 
-def softmax(x):
+def softmax(x, y):
 
     def occupancy(num_warps, size_smem):
 
@@ -155,9 +155,6 @@ def allocated_slm_size(size_smem):
     # way so you don't have to come up with manual heuristics yourself.
     num_warps = min(max_num_warps, max(1, BLOCK_SIZE // (WARP_SIZE * 4)))
 
-    # Allocate output
-    y = torch.empty_like(x)
-
     # pre-compile kernel to get register usage and compute thread occupancy.
     kernel = softmax_kernel.warmup(y, x, x.stride(0), y.stride(0), n_rows, n_cols, num_warps=num_warps,
                                    threads_per_warp=WARP_SIZE, BLOCK_SIZE=BLOCK_SIZE, grid=(1, ))
@@ -189,7 +186,8 @@ def allocated_slm_size(size_smem):
 
 torch.manual_seed(0)
 x = torch.randn(1823, 781, device=DEVICE)
-y_triton = softmax(x)
+y = torch.empty_like(x)
+y_triton = softmax(x, y)
 y_torch = torch.softmax(x, axis=1)
 assert torch.allclose(y_triton, y_torch), (y_triton, y_torch)
 
@@ -226,7 +224,7 @@ def benchmark(M, N, provider):
     if provider == 'torch':
         ms = triton.testing.do_bench(lambda: torch.softmax(x, axis=-1))
     if provider == 'triton':
-        ms = triton.testing.do_bench(lambda: softmax(x))
+        ms = triton.testing.do_bench(lambda: softmax(x, y))
     gbps = lambda ms: 2 * x.numel() * x.element_size() * 1e-9 / (ms * 1e-3)
     return gbps(ms)