[Runtime] Device API to query L2 cache size

include/tvm/runtime/device_api.h

@@ -48,7 +48,8 @@ enum DeviceAttrKind : int {
   kMaxRegistersPerBlock = 9,
   kGcnArch = 10,
   kApiVersion = 11,
-  kDriverVersion = 12
+  kDriverVersion = 12,
+  kL2CacheSizeBytes = 13,
 };
 
 #ifdef TVM_KALLOC_ALIGNMENT

python/tvm/_ffi/runtime_ctypes.py

@@ -488,6 +488,24 @@ def driver_version(self):
         """
         return self._GetDeviceAttr(self.device_type, self.device_id, 12)
 
+    @property
+    def l2_cache_size_bytes(self):
+        """Return the size of the device L2 cache in bytes
+
+        Supported devices include CUDA/ROCM/OpenCL.
+
+        Returns
+        -------
+        l2_cache_size_bytes : int or None
+            The size of the device L2 cache in bytes returned by device runtime API.
+            Return None if the device does not support this feature.
+
+        Note
+        ----
+        The value returned by opencl's API is smaller than actual device L2 cache size.
+        """
+        return self._GetDeviceAttr(self.device_type, self.device_id, 13)
+
     def texture_spatial_limit(self):
         """Returns limits for textures by spatial dimensions
 

python/tvm/target/target.py

@@ -244,6 +244,10 @@ def supports_cooperative_matrix(self):
     def features(self):
         return TargetFeatures(self)
 
+    @property
+    def l2_cache_size_bytes(self):
+        return int(self.attrs.get("l2_cache_size_bytes", 0))
+
     def get_kind_attr(self, attr_name):
         """Get additional attribute about the target kind.
 

src/runtime/cuda/cuda_device_api.cc

@@ -105,6 +105,12 @@ class CUDADeviceAPI final : public DeviceAPI {
       }
       case kDriverVersion:
         return;
+      case kL2CacheSizeBytes:
+        // Get size of device l2 cache size in bytes.
+        int l2_size = 0;
+        CUDA_CALL(cudaDeviceGetAttribute(&l2_size, cudaDevAttrL2CacheSize, dev.device_id));
+        *rv = l2_size;
+        return;
     }
     *rv = value;
   }

src/runtime/metal/metal_device_api.mm

@@ -81,6 +81,8 @@
         return;
       case kDriverVersion:
         return;
+      case kL2CacheSizeBytes:
+        return;
     }
   };
 }

src/runtime/opencl/opencl_device_api.cc

@@ -198,6 +198,13 @@ void OpenCLWorkspace::GetAttr(Device dev, DeviceAttrKind kind, TVMRetValue* rv)
       *rv = std::string(value);
       break;
     }
+    case kL2CacheSizeBytes:
+      // NOTE(Zihao): this API cannot reflect the real L2 cache size in both CUDA/AMD GPUs.
+      cl_ulong value;
+      OPENCL_CALL(clGetDeviceInfo(device_id, CL_DEVICE_GLOBAL_MEM_CACHE_SIZE, sizeof(value), &value,
+                                  nullptr));
+      *rv = static_cast<int64_t>(value);
+      break;
   }
 }
 

src/runtime/rocm/rocm_device_api.cc

@@ -122,6 +122,11 @@ class ROCMDeviceAPI final : public DeviceAPI {
       }
       case kDriverVersion:
         return;
+      case kL2CacheSizeBytes:
+        // Get size of device l2 cache size in bytes.
+        int l2_size;
+        ROCM_CALL(hipDeviceGetAttribute(&l2_size, hipDeviceAttributeL2CacheSize, device.device_id));
+        *rv = l2_size;
     }
     *rv = value;
   }

src/runtime/vulkan/vulkan_device_api.cc

@@ -160,6 +160,9 @@ void VulkanDeviceAPI::GetAttr(Device dev, DeviceAttrKind kind, TVMRetValue* rv)
       *rv = os.str();
       break;
     }
+
+    case kL2CacheSizeBytes:
+      break;
   }
 }