Fix L2 cache size calculation to handle symbolic expressions and ensure float...

Fix L2 cache size calculation to handle symbolic expressions and ensure float conversion of hit ratios in annotation (#576)

tilelang/jit/adapter/wrapper.py

@@ -353,7 +353,11 @@ class TLCUDASourceWrapper(object):
             # get persisting_l2_cache_max_size
             from tilelang.carver.arch.driver import get_persisting_l2_cache_max_size
             persisting_l2_cache_max_size = get_persisting_l2_cache_max_size()
-            num_bytes = min(size_in_bytes, persisting_l2_cache_max_size)
+            try:
+                num_bytes = min(size_in_bytes, persisting_l2_cache_max_size)
+            except Exception:
+                # as size_in_bytes maybe a symbolic expression
+                num_bytes = persisting_l2_cache_max_size
 
             init_l2_persistent_map += L2_PERSISTENT_MAP_INIT_FUNC.format(
                 buffer_name, float(hit_ratio), size_in_bytes, num_bytes)

tilelang/language/__init__.py

@@ -172,7 +172,7 @@ def annotate_l2_hit_ratio(l2_hit_ratio_map: Dict):
     _l2_hit_ratio_map = {}
     for buffer, hit_ratio in l2_hit_ratio_map.items():
         assert buffer.scope() == "global", "persistent L2 can only be applied to global buffers"
-        _l2_hit_ratio_map[buffer.data] = hit_ratio
+        _l2_hit_ratio_map[buffer.data] = float(hit_ratio)
     return block_attr({"l2_hit_ratio_map": _l2_hit_ratio_map})