KOKKOS_IF_ON_HOST and KOKKOS_IF_ON_DEVICE

Overview

The KOKKOS_IF_ON_HOST and KOKKOS_IF_ON_DEVICE macros are a pair of function-like macros introduced in Kokkos 3.6 that enable portable conditional compilation within a single ``KOKKOS_FUNCTION`` body. They allow you to select which code is compiled and executed based on whether the function is running on the host (CPU) or a device (GPU, etc.). These macros provide an alternative to non-portable preprocessor idioms like #ifdef __CUDA_ARCH__.

Motivation

Traditional preprocessor directives like #ifdef __CUDA_ARCH__ rely on a split compilation model, where host and device code are compiled in separate passes. While this model is supported by some compilers (like nvcc), it is not universally portable. Other modern compilers for GPU-accelerated code, such as those that support OpenACC or OpenMPTarget, use a unified compilation approach where both host and device code are compiled in a single pass. As a result, code written with backend-specific macros is not portable across different compilers and programming models

The KOKKOS_IF_ON_HOST and KOKKOS_IF_ON_DEVICE macros solve this portability problem by allowing the compiler to conditionally compile code within a single pass, enabling a single code base to be used with a wider range of compilers and backends.

Usage

These macros are designed to be used within a function decorated with KOKKOS_FUNCTION. They accept a single argument, which is a block of code enclosed in double parentheses. The code inside the macro’s parentheses will only be compiled and executed on the specified architecture

Signature

KOKKOS_IF_ON_HOST(( /* code to be compiled on the host */ ))
KOKKOS_IF_ON_DEVICE(( /* code to be compiled on the device */ ))

Example: Host/Device Overloading

A common use case is to provide different implementations of a function for host and device execution.

struct MyS { int i; };

KOKKOS_FUNCTION MyS MakeStruct() {
  // This return statement is only compiled for the host target.
  KOKKOS_IF_ON_HOST((
    return MyS{0};
  ))

  // This return statement is only compiled for the device target.
  KOKKOS_IF_ON_DEVICE((
    return MyS{1};
  ))
}

Important Considerations

Scope

Each KOKKOS_IF_ON_* macro introduces a new scope. Any variables declared within the macro’s parentheses are local to that scope and will not be accessible outside of it.

KOKKOS_IF_ON_HOST((
  int x = 0; // 'x' is only visible within this scope
  std::cout << x << '\n';
)) // The scope of 'x' ends here.

constexpr Context

These macros cannot be used in a context that requires a constexpr (constant expression).

Best Practices

Avoid using these macros whenever possible.

KOKKOS_IF_ON_HOST and KOKKOS_IF_ON_DEVICE should be considered a last resort for code differentiation. The primary goal of Kokkos is to achieve high-performance portability through a unified code base. Relying on these macros can hinder this goal by introducing host/device-specific logic.

Before using these macros, consider alternative approaches like partial template specialization on execution spaces or using Kokkos’s built-in functionalities, which are designed to be portable across all backends. Using these macros should be limited to situations where a fundamental difference between host and device APIs necessitates separate code paths, such as for I/O operations or specific backend intrinsics.