fix sgl-kernel build failure on AMD (#3352)

This commit is contained in:
Yineng Zhang
2025-02-07 00:35:59 +08:00
committed by GitHub
parent 6792411e7f
commit adeee15204

View File

@@ -16,7 +16,9 @@ limitations under the License.
#pragma once #pragma once
#include <cuda_runtime.h> #include <cuda_runtime.h>
#ifndef USE_ROCM
#include <pytorch_extension_utils.h> #include <pytorch_extension_utils.h>
#endif
#include <torch/extension.h> #include <torch/extension.h>
#include <sstream> #include <sstream>
@@ -63,6 +65,7 @@ inline int getSMVersion() {
return sm_major * 10 + sm_minor; return sm_major * 10 + sm_minor;
} }
#ifndef USE_ROCM
#define DISPATCH_PYTORCH_DTYPE_TO_CTYPE_FLOAT_FP16(pytorch_dtype, c_type, ...) \ #define DISPATCH_PYTORCH_DTYPE_TO_CTYPE_FLOAT_FP16(pytorch_dtype, c_type, ...) \
[&]() -> bool { \ [&]() -> bool { \
switch (pytorch_dtype) { \ switch (pytorch_dtype) { \
@@ -79,6 +82,7 @@ inline int getSMVersion() {
return false; \ return false; \
} \ } \
}() }()
#endif
#define DISPATCH_CASE_INTEGRAL_TYPES(...) \ #define DISPATCH_CASE_INTEGRAL_TYPES(...) \
AT_DISPATCH_CASE(at::ScalarType::Byte, __VA_ARGS__) \ AT_DISPATCH_CASE(at::ScalarType::Byte, __VA_ARGS__) \