[Runtime] Fix out of threads error to raise on every call (#7857)

Currently the attention tutorial is failing on gb200. This happens
because we only check for too many threads on the first `_init_modules`
call. This changes it so that when we raise `OutOfResources` we also set
`self._run` to a function that will raise the error. This means we will
still raise the error, even if the module has already been loaded.

Author: peterbell10

python/triton/compiler/compiler.py

@@ -403,6 +403,10 @@ def __missing__(self, key):
         return value
 
 
+def _raise_error(err, *args, **kwargs):
+    raise err
+
+
 class CompiledKernel:
 
     def __init__(self, src, metadata_group, hash):
@@ -439,32 +443,38 @@ def __init__(self, src, metadata_group, hash):
     def _init_handles(self):
         if self.module is not None:
             return
+
+        def raise_(err):
+            self._run = functools.partial(_raise_error, err)
+            raise err
+
         device = driver.active.get_current_device()
         # create launcher
         self._run = driver.active.launcher_cls(self.src, self.metadata)
         # not enough shared memory to run the kernel
         max_shared = max_shared_mem(device)
         if self.metadata.shared > max_shared:
-            raise OutOfResources(self.metadata.shared, max_shared, "shared memory")
+            raise_(OutOfResources(self.metadata.shared, max_shared, "shared memory"))
         if hasattr(self.metadata, "tmem_size") and self.metadata.tmem_size is not None:
             # Use blackwell max tmem size for now, this should be moved in device properties
             max_tmem_size = 512  # tmem size in number of columns
             if self.metadata.tmem_size > max_tmem_size:
-                raise OutOfResources(self.metadata.tmem_size, max_tmem_size, "tensor memory")
+                raise_(OutOfResources(self.metadata.tmem_size, max_tmem_size, "tensor memory"))
         if knobs.runtime.kernel_load_start_hook is not None:
             knobs.runtime.kernel_load_start_hook(self.module, self.function, self.name, self.metadata_group, self.hash)
         # TODO: n_regs, n_spills should be metadata generated when calling `ptxas`
         self.module, self.function, self.n_regs, self.n_spills, self.n_max_threads = driver.active.utils.load_binary(
             self.name, self.kernel, self.metadata.shared, device)
         warp_size = driver.active.get_current_target().warp_size
         if self.metadata.num_warps * warp_size > self.n_max_threads:
-            raise OutOfResources(self.metadata.num_warps * warp_size, self.n_max_threads, "threads")
+            raise_(OutOfResources(self.metadata.num_warps * warp_size, self.n_max_threads, "threads"))
         if knobs.runtime.kernel_load_end_hook is not None:
             knobs.runtime.kernel_load_end_hook(self.module, self.function, self.name, self.metadata_group, self.hash)
 
     @property
     def run(self):
-        self._init_handles()
+        if self._run is None:
+            self._init_handles()
         return self._run
 
     def launch_metadata(self, grid, stream, *args):