How to run 02-fused-softmax.py on a CPU?

[banach-space]

Hi folks,

First off, thank you for triton-shared—it's fantastic work!

I’ve successfully built the project and run the examples on my AArch64 machine. However, I’m running into an issue with 02-fused-softmax.py:

AssertionError: Torch not compiled with CUDA enabled

Clearly, something is trying to target CUDA, but I’m not sure what. I modified the example to explicitly select the CPU backend:

diff --git a/python/tutorials/02-fused-softmax.py b/python/tutorials/02-fused-softmax.py
index 57f027277..8b10eb3d3 100644
--- a/python/tutorials/02-fused-softmax.py
+++ b/python/tutorials/02-fused-softmax.py
@@ -27,6 +27,8 @@ import triton
 import triton.language as tl
 from triton.runtime import driver

+from triton.backends.triton_shared.driver import CPUDriver
+

 def is_hip():
     return triton.runtime.driver.active.get_current_target().backend == "hip"
@@ -110,12 +112,9 @@ def softmax_kernel(output_ptr, input_ptr, input_row_stride, output_row_stride, n
 # %%
 # We can create a helper function that enqueues the kernel and its (meta-)arguments for any given input tensor.

-device = torch.cuda.current_device()
-properties = driver.active.utils.get_device_properties(device)
-NUM_SM = properties["multiprocessor_count"]
-NUM_REGS = properties["max_num_regs"]
-SIZE_SMEM = properties["max_shared_mem"]
-WARP_SIZE = properties["warpSize"]
+triton.runtime.driver.set_active(CPUDriver())
+device = 'cpu'
+# device = torch.cuda.current_device()
 target = triton.runtime.driver.active.get_current_target()
 kernels = {}

@@ -146,25 +145,7 @@ def softmax(x):
         kernel._init_handles()
         n_regs = kernel.n_regs
         size_smem = kernel.metadata.shared
-        if is_hip():
-            # NUM_REGS represents the number of regular purpose registers. On CDNA architectures this is half of all registers available.
-            # However, this is not always the case. In most cases all registers can be used as regular purpose registers.
-            # ISA SECTION (3.6.4 for CDNA3)
-            # VGPRs are allocated out of two pools: regular VGPRs and accumulation VGPRs. Accumulation VGPRs are used
-            # with matrix VALU instructions, and can also be loaded directly from memory. A wave may have up to 512 total
-            # VGPRs, 256 of each type. When a wave has fewer than 512 total VGPRs, the number of each type is flexible - it is
-            # not required to be equal numbers of both types.
-            if is_cdna():
-                NUM_GPRS = NUM_REGS * 2
-
-            # MAX_NUM_THREADS represents maximum number of resident threads per multi-processor.
-            # When we divide this number with WARP_SIZE we get maximum number of waves that can
-            # execute on a CU (multi-processor)  in parallel.
-            MAX_NUM_THREADS = properties["max_threads_per_sm"]
-            max_num_waves = MAX_NUM_THREADS // WARP_SIZE
-            occupancy = min(NUM_GPRS // WARP_SIZE // n_regs, max_num_waves) // num_warps
-        else:
-            occupancy = NUM_REGS // (n_regs * WARP_SIZE * num_warps)
+        occupancy = NUM_REGS // (n_regs * WARP_SIZE * num_warps)
         occupancy = min(occupancy, SIZE_SMEM // size_smem)
         num_programs = NUM_SM * occupancy
         kernels[BLOCK_SIZE] = (kernel, num_programs)

Unfortunately, I’m still getting the same error, so I suspect my changes might not be sufficient. I’m very new to this and probably missing something obvious - please bear with me 😅

Any guidance would be greatly appreciated. Let me know if you need additional logs or details from my setup!

Thanks,
Andrzej

[parsifal-47]

Hi Andrzej,
the perf part internally very specific to GPU, that is why in my PR here #163 I do not use it to benchmark triton on CPU. I would recommend to remove "perf" part entirely and retry.

Thank you!

[banach-space]

Hi Renat, thanks for getting back to me so quickly and for the link. Yes, removing "perf" the perf part fixes the issue :)

Let me close this and just follow what you did in #163. It would be nice to have it merged in :)

-Andrzej