squall/flash-attention
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Merge pull request #343 from danthe3rd/if_constexpr

Browse Source 
Fix compile error with `BOOL_SWITCH`
This commit is contained in:
Tri Dao 2023-07-19 04:27:07 -04:00 committed by GitHub
commit 31ae2488e6
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 50 additions and 49 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

99
csrc/flash_attn/src/static_switch.h
View File

@ -1,4 +1,5 @@
// Inspired by https://github.com/NVIDIA/DALI/blob/main/include/dali/core/static_switch.h // Inspired by
// https://github.com/NVIDIA/DALI/blob/main/include/dali/core/static_switch.h
// and https://github.com/pytorch/pytorch/blob/master/aten/src/ATen/Dispatch.h // and https://github.com/pytorch/pytorch/blob/master/aten/src/ATen/Dispatch.h
#pragma once #pragma once
@ -13,53 +14,53 @@
/// some_function<BoolConst>(...); /// some_function<BoolConst>(...);
/// }); /// });
/// ``` /// ```
#define BOOL_SWITCH(COND, CONST_NAME, ...) \ #define BOOL_SWITCH(COND, CONST_NAME, ...) \
[&] { \ [&] { \
if (COND) { \ if (COND) { \
constexpr bool CONST_NAME = true; \ constexpr static bool CONST_NAME = true; \
return __VA_ARGS__(); \ return __VA_ARGS__(); \
} else { \ } else { \
constexpr bool CONST_NAME = false; \ constexpr static bool CONST_NAME = false; \
return __VA_ARGS__(); \ return __VA_ARGS__(); \
} \ } \
}() }()
#define FP16_SWITCH(COND, ...) \ #define FP16_SWITCH(COND, ...) \
[&] { \ [&] { \
if (COND) { \ if (COND) { \
using elem_type = cutlass::half_t; \ using elem_type = cutlass::half_t; \
return __VA_ARGS__(); \ return __VA_ARGS__(); \
} else { \ } else { \
using elem_type = cutlass::bfloat16_t; \ using elem_type = cutlass::bfloat16_t; \
return __VA_ARGS__(); \ return __VA_ARGS__(); \
} \ } \
}() }()
#define FWD_HEADDIM_SWITCH(HEADDIM, ...) \ #define FWD_HEADDIM_SWITCH(HEADDIM, ...) \
[&] { \ [&] { \
if (HEADDIM <= 32) { \ if (HEADDIM <= 32) { \
constexpr int kHeadDim = 32; \ constexpr static int kHeadDim = 32; \
return __VA_ARGS__(); \ return __VA_ARGS__(); \
} else if (HEADDIM <= 64) { \ } else if (HEADDIM <= 64) { \
constexpr int kHeadDim = 64; \ constexpr static int kHeadDim = 64; \
return __VA_ARGS__(); \ return __VA_ARGS__(); \
} else if (HEADDIM <= 96) { \ } else if (HEADDIM <= 96) { \
constexpr int kHeadDim = 96; \ constexpr static int kHeadDim = 96; \
return __VA_ARGS__(); \ return __VA_ARGS__(); \
} else if (HEADDIM <= 128) { \ } else if (HEADDIM <= 128) { \
constexpr int kHeadDim = 128; \ constexpr static int kHeadDim = 128; \
return __VA_ARGS__(); \ return __VA_ARGS__(); \
} else if (HEADDIM <= 160) { \ } else if (HEADDIM <= 160) { \
constexpr int kHeadDim = 160; \ constexpr static int kHeadDim = 160; \
return __VA_ARGS__(); \ return __VA_ARGS__(); \
} else if (HEADDIM <= 192) { \ } else if (HEADDIM <= 192) { \
constexpr int kHeadDim = 192; \ constexpr static int kHeadDim = 192; \
return __VA_ARGS__(); \ return __VA_ARGS__(); \
} else if (HEADDIM <= 224) { \ } else if (HEADDIM <= 224) { \
constexpr int kHeadDim = 224; \ constexpr static int kHeadDim = 224; \
return __VA_ARGS__(); \ return __VA_ARGS__(); \
} else if (HEADDIM <= 256) { \ } else if (HEADDIM <= 256) { \
constexpr int kHeadDim = 256; \ constexpr static int kHeadDim = 256; \
return __VA_ARGS__(); \ return __VA_ARGS__(); \
} \ } \
}() }()