offload/include/Shared/APITypes.h

//===-- Shared/APITypes.h - Offload and plugin API types --------*- C++ -*-===//
//
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
// See https://llvm.org/LICENSE.txt for license information.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
//
//===----------------------------------------------------------------------===//
//
// This file defines types used in the interface between the user code, the
// target independent offload runtime library, and target dependent plugins.
//
//===----------------------------------------------------------------------===//

#ifndef OMPTARGET_SHARED_API_TYPES_H
#define OMPTARGET_SHARED_API_TYPES_H

#include "Environment.h"

#include "llvm/ADT/SmallString.h"
#include "llvm/ADT/SmallVector.h"
#include "llvm/Frontend/Offloading/Utility.h"

#include <cstddef>
#include <cstdint>
#include <mutex>
#include <utility>

enum TargetAllocTy : int32_t;

extern "C" {

/// This struct is a record of the device image information
struct __tgt_device_image {
  void *ImageStart; // Pointer to the target code start
  void *ImageEnd;   // Pointer to the target code end
  llvm::offloading::EntryTy
      *EntriesBegin; // Begin of table with all target entries
  llvm::offloading::EntryTy *EntriesEnd; // End of table (non inclusive)
};

struct __tgt_device_info {
  void *Context = nullptr;
  void *Device = nullptr;
  void *Platform = nullptr;
};

/// This struct is a record of all the host code that may be offloaded to a
/// target.
struct __tgt_bin_desc {
  int32_t NumDeviceImages;          // Number of device types supported
  __tgt_device_image *DeviceImages; // Array of device images (1 per dev. type)
  llvm::offloading::EntryTy
      *HostEntriesBegin; // Begin of table with all host entries
  llvm::offloading::EntryTy *HostEntriesEnd; // End of table (non inclusive)
};

/// This struct contains the offload entries identified by the target runtime
struct __tgt_target_table {
  llvm::offloading::EntryTy
      *EntriesBegin; // Begin of the table with all the entries
  llvm::offloading::EntryTy
      *EntriesEnd; // End of the table with all the entries (non inclusive)
};

/// This struct contains a handle to a loaded binary in the plugin device.
struct __tgt_device_binary {
  uintptr_t handle;
};

// clang-format on

/// This struct contains information exchanged between different asynchronous
/// operations for device-dependent optimization and potential synchronization
struct __tgt_async_info {
  // A pointer to a queue-like structure where offloading operations are issued.
  // We assume to use this structure to do synchronization. In CUDA backend, it
  // is CUstream.
  void *Queue = nullptr;

  /// A collection of allocations that are associated with this stream and that
  /// should be freed after finalization. Also store the type of the allocation.
  llvm::SmallVector<std::pair<void *, TargetAllocTy>, 2> AssociatedAllocations;

  /// Mutex to guard access to AssociatedAllocations and the Queue.
  std::mutex Mutex;

  /// The kernel launch environment used to issue a kernel. Stored here to
  /// ensure it is a valid location while the transfer to the device is
  /// happening.
  KernelLaunchEnvironmentTy KernelLaunchEnvironment;

  /// Opaque handle for profiler-specific data (e.g., OMPT trace record info).
  /// Owned by the profiler; the runtime threads this pointer through the plugin
  /// layer to associate async operations with trace records.
  void *ProfilerData = nullptr;
};

/// This struct contains all of the arguments to a target kernel region launch.
struct KernelArgsTy {
  uint32_t Version = 0; // Version of this struct for ABI compatibility.
  uint32_t NumArgs = 0; // Number of arguments in each input pointer.
  void **ArgBasePtrs =
      nullptr;                 // Base pointer of each argument (e.g. a struct).
  void **ArgPtrs = nullptr;    // Pointer to the argument data.
  int64_t *ArgSizes = nullptr; // Size of the argument data in bytes.
  int64_t *ArgTypes = nullptr; // Type of the data (e.g. to / from).
  void **ArgNames = nullptr;   // Name of the data for debugging, possibly null.
  void **ArgMappers = nullptr; // User-defined mappers, possibly null.
  uint64_t Tripcount =
      0; // Tripcount for the teams / distribute loop, 0 otherwise.
  struct {
    uint64_t NoWait : 1; // Was this kernel spawned with a `nowait` clause.
    uint64_t IsCUDA : 1; // Was this kernel spawned via CUDA.
    uint64_t DynCGroupMemFallback : 2; // The fallback for dynamic cgroup mem.
    uint64_t Cooperative : 1; // Was this kernel spawned as cooperative.
    uint64_t IsPtrArgs : 1;   // Arguments are laid out as an array of pointers.
    uint64_t StrictBlocks : 1; // The user-requested number of blocks is strict.
    uint64_t StrictThreads
        : 1; // The user-requested number of threads is strict.
    uint64_t Unused : 56;
  } Flags = {0, 0, 0, 0, 0, 0, 0, 0};
  // User-requested number of blocks (for x,y,z dimension).
  uint32_t UserNumBlocks[3] = {0, 0, 0};
  // User-requested number of threads (for x,y,z dimension).
  uint32_t UserThreadLimit[3] = {0, 0, 0};
  uint32_t DynCGroupMem = 0; // Amount of dynamic cgroup memory requested.
};
static_assert(sizeof(KernelArgsTy().Flags) == sizeof(uint64_t),
              "Invalid struct size");
static_assert(sizeof(KernelArgsTy) ==
                  (8 * sizeof(int32_t) + 3 * sizeof(int64_t) +
                   4 * sizeof(void **) + 2 * sizeof(int64_t *)),
              "Invalid struct size");

/// Array of pointers to kernel launch arguments and the size of that array.
struct KernelLaunchParamsTy {
  /// Number of kernel arguments in \p Args.
  uint32_t NumArgs = 0;
  /// Array of \p NumArgs pointers, each pointing at one argument's value.
  void **Args = nullptr;
};

/// The outcome of a kernel replay.
struct KernelReplayOutcomeTy {
  /// The path to the file that stores the output memory snapshot after the
  /// kernel has been replayed.
  llvm::SmallString<128> OutputFilepath;
  /// The execution time of the kernel replay in nanoseconds. This time includes
  /// the the kernel launch and synchronization time. Replay I/O is excluded.
  uint64_t KernelReplayTimeNs = 0;
  /// The pointer to the device memory allocation used to replay. This can be
  /// reused for future replays of the same kernel.
  void *ReplayDeviceAlloc = nullptr;
};
}

#endif // OMPTARGET_SHARED_API_TYPES_H