[AMD] Fix DSV4 JIT build on rocm (#31131)

Co-authored-by: kangwangamd <kangwang@amd.com>
This commit is contained in:
YC Yen-Ching Tseng
2026-07-15 09:58:12 -07:00
committed by GitHub
co-authored by kangwangamd
parent c879f3da5c
commit a8b60433c2
2 changed files with 15 additions and 8 deletions
@@ -26,6 +26,9 @@
#ifndef cudaDevAttrComputeCapabilityMajor
#define cudaDevAttrComputeCapabilityMajor hipDeviceAttributeComputeCapabilityMajor
#endif
#ifndef cudaDevAttrComputeCapabilityMinor
#define cudaDevAttrComputeCapabilityMinor hipDeviceAttributeComputeCapabilityMinor
#endif
#ifndef cudaRuntimeGetVersion
#define cudaRuntimeGetVersion hipRuntimeGetVersion
#endif
@@ -68,6 +71,18 @@ inline auto get_cc_major(int device_id) -> int {
return cc_major;
}
// Return the Minor compute capability for the given device
inline auto get_cc_minor(int device_id) -> int {
int cc_minor;
RuntimeDeviceCheck(cudaDeviceGetAttribute(&cc_minor, cudaDevAttrComputeCapabilityMinor, device_id));
return cc_minor;
}
// Return the SM version (major * 10 + minor) for the given device
inline auto get_sm_version(int device_id) -> int {
return get_cc_major(device_id) * 10 + get_cc_minor(device_id);
}
// Return the runtime version
inline auto get_runtime_version() -> int {
int runtime_version;
@@ -242,14 +242,6 @@ inline void RuntimeDeviceCheck(DebugInfo location = {}) {
return RuntimeDeviceCheck(::cudaGetLastError(), location);
}
inline int getSMVersion(int device_id) {
int sm_major = 0;
int sm_minor = 0;
RuntimeDeviceCheck(cudaDeviceGetAttribute(&sm_major, cudaDevAttrComputeCapabilityMajor, device_id));
RuntimeDeviceCheck(cudaDeviceGetAttribute(&sm_minor, cudaDevAttrComputeCapabilityMinor, device_id));
return sm_major * 10 + sm_minor;
}
inline auto alloc_workspace_tensor(size_t required_bytes, DLDevice device) -> tvm::ffi::Tensor {
if (required_bytes == 0) return {};
DLDataType u8 = {kDLUInt, 8, 1};