| 1 | // Inspired by |
| 2 | // https://github.com/NVIDIA/DALI/blob/main/include/dali/core/static_switch.h |
| 3 | // and https://github.com/pytorch/pytorch/blob/master/aten/src/ATen/Dispatch.h |
| 4 | |
| 5 | // clang-format off |
| 6 | // adapted from https://github.com/state-spaces/mamba/blob/main/csrc/selective_scan/static_switch.h |
| 7 | #pragma once |
| 8 | |
| 9 | /// @param COND - a boolean expression to switch by |
| 10 | /// @param CONST_NAME - a name given for the constexpr bool variable. |
| 11 | /// @param ... - code to execute for true and false |
| 12 | /// |
| 13 | /// Usage: |
| 14 | /// ``` |
| 15 | /// BOOL_SWITCH(flag, BoolConst, [&] { |
| 16 | /// some_function<BoolConst>(...); |
| 17 | /// }); |
| 18 | /// ``` |
| 19 | #define BOOL_SWITCH(COND, CONST_NAME, ...) \ |
| 20 | [&] { \ |
| 21 | if (COND) { \ |
| 22 | constexpr bool CONST_NAME = true; \ |
| 23 | return __VA_ARGS__(); \ |
| 24 | } else { \ |
| 25 | constexpr bool CONST_NAME = false; \ |
| 26 | return __VA_ARGS__(); \ |
| 27 | } \ |
| 28 | }() |