dpaintrin.h

Go to the documentation of this file.

/*
 * NVIDIA_COPYRIGHT_BEGIN
 *
 * Copyright (c) 2023-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 *
 * NVIDIA CORPORATION, its affiliates and licensors retain all intellectual
 * property and proprietary rights in and to this material, related
 * documentation and any modifications thereto. Any use, reproduction,
 * disclosure or distribution of this material and related documentation
 * without an express license agreement from NVIDIA CORPORATION or
 * its affiliates is strictly prohibited.
 *
 * NVIDIA_COPYRIGHT_END
 */
 
//===--------- dpaintrin.h - Header file for all DPA intrinsics -----------===
#ifndef __DPAINTRIN_H
#define __DPAINTRIN_H
 
/*
 * Users need to define following macro before including this header file to
 * use a specific version of DPA intrinsics
 */
#ifndef DPA_INTRIN_VERSION_USED
#define DPA_INTRIN_VERSION_USED (DPA_INTRIN_VERSION(1, 3))
#endif
 
#if (DPA_INTRIN_VERSION_USED == (DPA_INTRIN_VERSION(1, 3)))
#define __DPA_HEAP __MSPACE_HEAP
#define __DPA_MEMORY __MSPACE_MEMORY
#define __DPA_MMIO __MSPACE_MMIO
#define __DPA_SYSTEM __MSPACE_SYSTEM
 
#define __DPA_R __MOP_R
#define __DPA_W __MOP_W
#define __DPA_RW __MOP_RW
 
#define __dpa_thread_fence(MEMORY_SPACE, PRED_OP, SUCC_OP)                     \
  __dpa_thread_fence_internal_1_3(MEMORY_SPACE, PRED_OP, SUCC_OP);
 
#define __dpa_thread_memory_fence(OP1, OP2)                                    \
  __dpa_thread_fence(__DPA_MEMORY, OP1, OP2)
 
#define __dpa_thread_outbox_fence(OP1, OP2)                                    \
  __dpa_thread_fence(__DPA_MMIO, OP1, OP2)
 
#define __dpa_thread_window_fence(OP1, OP2)                                    \
  __dpa_thread_fence(__DPA_MMIO, OP1, OP2)
 
#define __dpa_thread_system_fence()                                            \
  __dpa_thread_fence(__DPA_SYSTEM, __DPA_RW, __DPA_RW)
 
#define __dpa_thread_window_read_inv()                                         \
  __dpa_thread_fence(__DPA_MMIO, __DPA_R, __DPA_R)
 
#define __dpa_thread_window_writeback()                                        \
  __dpa_thread_fence(__DPA_MMIO, __DPA_W, __DPA_W)
 
#if __NV_DPA >= NV_DPA_CX8
#define __dpa_thread_l1_flush()                                                 \
  __dpa_thread_l1_flush_internal_1_3()
#endif
 
#if __NV_DPA == __NV_DPA_BF3
#define __dpa_thread_memory_writeback()                                        \
  __dpa_thread_fence(__DPA_MEMORY, __DPA_W, __DPA_W)
#else
#define __dpa_thread_memory_writeback() __dpa_compiler_barrier()
#endif
 
#if defined(__riscv_xrpfxp)
#define __dpa_fxp_rcp(N) __dpa_fxp_rcp_internal_1_3(N)
#define __dpa_fxp_pow2(N) __dpa_fxp_pow2_internal_1_3(N)
#define __dpa_fxp_log2(N) __dpa_fxp_log2_internal_1_3(N)
#endif // __riscv_xrpfxp
 
#if defined(__riscv_xnvcc)
#define __dpa_data_ignore(ADDR) __dpa_data_ignore_internal_1_3(ADDR)
#endif // __riscv_xnvcc
 
#define __dpa_thread_cycles() __dpa_thread_cycles_internal_1_3()
#define __dpa_thread_inst_ret() __dpa_thread_inst_ret_internal_1_3()
#define __dpa_thread_time() __dpa_thread_time_internal_1_3()
 
#if __NV_DPA >= NV_DPA_CX8
#define __dpa_remote_atomic_load(PTR, MEMORDER) \
    __dpa_remote_atomic_load_internal_1_3(PTR, MEMORDER)
 
#define __dpa_remote_atomic_exchange(PTR, VAL, MEMORDER) \
    __dpa_remote_atomic_exchange_internal_1_3(PTR, VAL, MEMORDER)
 
#define __dpa_remote_atomic_add_fetch(PTR, VAL, MEMORDER) \
    __dpa_remote_atomic_add_fetch_internal_1_3(PTR, VAL, MEMORDER)
#define __dpa_remote_atomic_sub_fetch(PTR, VAL, MEMORDER) \
    __dpa_remote_atomic_sub_fetch_internal_1_3(PTR, VAL, MEMORDER)
 
#define __dpa_remote_atomic_fetch_add(PTR, VAL, MEMORDER) \
    __dpa_remote_atomic_fetch_add_internal_1_3(PTR, VAL, MEMORDER)
#define __dpa_remote_atomic_fetch_sub(PTR, VAL, MEMORDER) \
    __dpa_remote_atomic_fetch_sub_internal_1_3(PTR, VAL, MEMORDER)
 
#if defined(__riscv_xcas)
#define __dpa_remote_atomic_compare_exchange(PTR, EXPECTED, DESIRED, \
    SUCC_MEMORDER, FAIL_MEMORDER) \
    __dpa_remote_atomic_compare_exchange_internal_1_3(PTR, EXPECTED, DESIRED, \
    SUCC_MEMORDER, FAIL_MEMORDER)
#endif
 
#endif
 
#include <stdint.h>
#define __extract_fields(dst, src, ...) \
  __extract_fields_internal(dst, src, \
      sizeof((int []) {__VA_ARGS__}) / sizeof(int) / 3, \
      (int []) {__VA_ARGS__})
 
static inline __attribute__((always_inline)) 
uint64_t __extract_64bits(uint64_t arr[], int bitoffset, int bitlength)
{
  uint64_t result = arr[bitoffset/64] >> (bitoffset % 64);
  // read the remaining part from the next index
  if (bitoffset % 64 + bitlength > 64)
    result |= arr[bitoffset/64+1] << (64 - bitoffset % 64);
  return result;
}
 
static inline __attribute__((always_inline))
void __extract_fields_internal (uint64_t *restrict dst, uint64_t *restrict src,
                                int extract_count, int varg[]) 
{
  #pragma unroll
  for (int i = 0; i < extract_count; i++) {
    int reverse = 0;
    // triplet in bits
    int dst_offset = varg[i*3] * 8;
    int src_offset = varg[i*3+1] * 8;
    int length = varg[i*3+2] * 8;
    if (length < 0) {
      length = -length;
      reverse = 1;
    } else if (length == 0)
      continue;
 
    // initial suffix to copy to dst if it starts in the middle of an 8-byte boundary
    if (dst_offset % 64) {
      int segment_offset = dst_offset % 64;
      uint64_t srcmask = -1l;
      if (segment_offset + length < 64) // whole extraction is within the 8 byte
        srcmask >>= 64 - length;
      srcmask <<= segment_offset;
      uint64_t dstmask = ~srcmask;
      uint64_t srcval = __extract_64bits(src, src_offset,  64 - segment_offset) << segment_offset;
      uint64_t oldval = dst[dst_offset/64];
      dst[dst_offset/64] = (srcval & srcmask) | (oldval & dstmask);
      dst_offset += 64 - segment_offset;
      src_offset += 64 - segment_offset;
      length -= 64 - segment_offset;
      if (length < 0)
        length = 0;
    }
 
    // copy whole 8-byte segments
    while (length / 64) {
      dst[dst_offset/64] = __extract_64bits(src, src_offset, 64);
      dst_offset += 64;
      src_offset += 64;
      length -= 64;
    }
 
    // remaining prefix 
    if (length) {
      uint64_t srcval =
        __extract_64bits(src, src_offset, length) & ((1l<<length)-1);
      uint64_t oldvalmasked = dst[dst_offset/64] & (-1l << length);
      dst[dst_offset/64] = srcval | oldvalmasked;
    }
 
    if (reverse) {
      int num_bytes = -varg[i*3+2];
      int offset = varg[i*3];
      uint64_t val = dst[offset/8];
      uint64_t swapped = val >> (offset % 8)*8;
      switch (num_bytes)
      {
      case 2:
        swapped = __builtin_bswap16(swapped);
        break;
      case 4:
        swapped = __builtin_bswap32(swapped);
        break;
      case 8:
        swapped = __builtin_bswap64(swapped);
        break;
      }
      swapped <<= (offset % 8)*8;
      uint64_t mask = (1l << (offset % 8)*8) - 1;
      int upper_part = offset % 8 + num_bytes;
      if (upper_part < 8)
        mask |= -1l << upper_part * 8;
      dst[offset/8] = swapped | (val & mask);
    }
  }
}
 
#else
 
#error Bad value for DPA_INTRIN_VERSION_USED
 
#endif // DPA_INTRIN_VERSION_USED
 
#endif // __DPAINTRIN_H