[python][compiler] Memoize device max shared memory per device (triton-lang#6503)

Similar to triton-lang#6000 this patch is
an upstreamed internal patch at Meta with the goal of reducing our
internal patches, cc @jamesjwu the original author.

When running various benchmarks with small kernels we see a non-trivial
amount of time spent fetching this property, and memoizing helped.

It might be worth looking into memoizing `get_device_properties`, but I
think that'd need a more careful treatment in the driver package in
order to properly handle arbitrary backends.

Author: danzimm

python/triton/compiler/compiler.py

@@ -164,6 +164,11 @@ def triton_key():
     return f'{__version__}' + '-'.join(contents)
 
 
+@functools.lru_cache()
+def max_shared_mem(device):
+    return driver.active.utils.get_device_properties(device)["max_shared_mem"]
+
+
 def parse(full_name, ext, context):
     if ext == "ttir" or ext == "ttgir":
         module = ir.parse_mlir_module(full_name, context)
@@ -400,7 +405,7 @@ def _init_handles(self):
         # create launcher
         self.run = driver.active.launcher_cls(self.src, self.metadata)
         # not enough shared memory to run the kernel
-        max_shared = driver.active.utils.get_device_properties(device)["max_shared_mem"]
+        max_shared = max_shared_mem(device)
         if self.metadata.shared > max_shared:
             raise OutOfResources(self.metadata.shared, max_shared, "shared memory")
         if hasattr(self.metadata, "tmem_size") and self.metadata.tmem_size is not None: