offload/unittests/OffloadAPI/common/Fixtures.hpp

//===------- Offload API tests - gtest fixtures --==-----------------------===//
//
// 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
//
//===----------------------------------------------------------------------===//

#include <OffloadAPI.h>
#include <OffloadPrint.hpp>
#include <gtest/gtest.h>
#include <optional>
#include <string>
#include <thread>

#include "Environment.hpp"
#include "Shared/EnvironmentVar.h"

#pragma once

#ifndef ASSERT_SUCCESS
#define ASSERT_SUCCESS(ACTUAL)                                                 \
  do {                                                                         \
    ol_result_t Res = ACTUAL;                                                  \
    if (Res && Res->Code != OL_ERRC_SUCCESS) {                                 \
      GTEST_FAIL() << #ACTUAL " returned " << Res->Code << ": "                \
                   << Res->Details;                                            \
    }                                                                          \
  } while (0)
#endif

#ifndef ASSERT_SUCCESS_OR_UNSUPPORTED
#define ASSERT_SUCCESS_OR_UNSUPPORTED(ACTUAL)                                  \
  do {                                                                         \
    ol_result_t Res = ACTUAL;                                                  \
    if (Res && Res->Code == OL_ERRC_UNSUPPORTED) {                             \
      GTEST_SKIP() << #ACTUAL " returned unsupported; skipping test";          \
      return;                                                                  \
    } else if (Res && Res->Code != OL_ERRC_SUCCESS) {                          \
      GTEST_FAIL() << #ACTUAL " returned " << Res->Code << ": "                \
                   << Res->Details;                                            \
    }                                                                          \
  } while (0)
#endif

#ifndef ASSERT_ERROR
#define ASSERT_ERROR(EXPECTED, ACTUAL)                                         \
  do {                                                                         \
    ol_result_t Res = ACTUAL;                                                  \
    if (!Res)                                                                  \
      GTEST_FAIL() << #ACTUAL " succeeded when we expected it to fail";        \
    if (Res->Code != EXPECTED)                                                 \
      GTEST_FAIL() << #ACTUAL " was expected to return "                       \
                   << #EXPECTED " but instead returned " << Res->Code << ": "  \
                   << Res->Details;                                            \
  } while (0)
#endif

#ifndef ASSERT_ANY_ERROR
#define ASSERT_ANY_ERROR(ACTUAL)                                               \
  do {                                                                         \
    ol_result_t Res = ACTUAL;                                                  \
    ASSERT_TRUE(Res);                                                          \
  } while (0)
#endif

struct BackendMatcher {
  ol_platform_backend_t Backend;
  std::string Message;

  BackendMatcher(ol_platform_backend_t B, std::string M = {})
      : Backend(B), Message(std::move(M)) {}
};

struct LevelZero : BackendMatcher {
  LevelZero(std::string M = {})
      : BackendMatcher(OL_PLATFORM_BACKEND_LEVEL_ZERO, std::move(M)) {}
};

struct CUDA : BackendMatcher {
  CUDA(std::string M = {})
      : BackendMatcher(OL_PLATFORM_BACKEND_CUDA, std::move(M)) {}
};

struct AMDGPU : BackendMatcher {
  AMDGPU(std::string M = {})
      : BackendMatcher(OL_PLATFORM_BACKEND_AMDGPU, std::move(M)) {}
};

inline std::string knownFailureMessage(const BackendMatcher &M) {
  std::string Msg;
  llvm::raw_string_ostream OS(Msg);
  OS << "Known failure on " << M.Backend;
  if (!M.Message.empty())
    OS << ": " << M.Message;
  return Msg;
}

inline std::optional<std::string>
findKnownFailure(ol_platform_backend_t CurBackend,
                 std::initializer_list<BackendMatcher> Matchers) {
  for (const auto &M : Matchers) {
    if (M.Backend == CurBackend)
      return knownFailureMessage(M);
  }
  return std::nullopt;
}

#define SKIP_KNOWN_FAILURE(...)                                                \
  if (auto KFMsg =                                                             \
          ::findKnownFailure(this->getPlatformBackend(), {__VA_ARGS__}))       \
  GTEST_SKIP() << *KFMsg

#define RETURN_ON_FATAL_FAILURE(...)                                           \
  __VA_ARGS__;                                                                 \
  if (this->HasFatalFailure() || this->IsSkipped()) {                          \
    return;                                                                    \
  }                                                                            \
  (void)0

inline std::string SanitizeString(const std::string &Str) {
  auto NewStr = Str;
  std::replace_if(
      NewStr.begin(), NewStr.end(), [](char C) { return !std::isalnum(C); },
      '_');
  return NewStr;
}

template <typename Fn> inline void threadify(Fn body) {
  std::vector<std::thread> Threads;
  for (size_t I = 0; I < 20; I++) {
    Threads.emplace_back(
        [&body](size_t I) {
          std::string ScopeMsg{"Thread #"};
          ScopeMsg.append(std::to_string(I));
          SCOPED_TRACE(ScopeMsg);
          body(I);
        },
        I);
  }
  for (auto &T : Threads) {
    T.join();
  }
}

/// Skip the current test when OFFLOAD_FORCE_SYNC_OPS is enabled. Tests using
/// ManuallyTriggeredTask enqueue a host task that blocks until a later
/// `trigger`; forcing operations synchronous makes the enqueue wait for that
/// task inline, so it can never be triggered and times out.
#define SKIP_IF_FORCE_SYNC_OPS()                                               \
  do {                                                                         \
    if (BoolEnvar("OFFLOAD_FORCE_SYNC_OPS", false))                            \
      GTEST_SKIP() << "incompatible with OFFLOAD_FORCE_SYNC_OPS";              \
  } while (0)

/// Enqueues a task to the queue that can be manually resolved.
// It will block until `trigger` is called.
struct ManuallyTriggeredTask {
  std::mutex M;
  std::condition_variable CV;
  bool Flag = false;
  ol_event_handle_t CompleteEvent;

  ol_result_t enqueue(ol_queue_handle_t Queue) {
    if (auto Err = olLaunchHostFunction(
            Queue,
            [](void *That) {
              static_cast<ManuallyTriggeredTask *>(That)->wait();
            },
            this))
      return Err;

    return olCreateEvent(Queue, OL_EVENT_FLAGS_NONE, &CompleteEvent);
  }

  void wait() {
    std::unique_lock<std::mutex> lk(M);
    CV.wait_for(lk, std::chrono::milliseconds(1000), [&] { return Flag; });
    EXPECT_TRUE(Flag);
  }

  ol_result_t trigger() {
    Flag = true;
    CV.notify_one();

    return olSyncEvent(CompleteEvent);
  }
};

struct OffloadTest : ::testing::Test {
  ol_device_handle_t Host = TestEnvironment::getHostDevice();
};

template <class T> using OffloadParam = std::tuple<TestEnvironment::Device, T>;

template <class T>
struct OffloadDeviceTestWithParam
    : OffloadTest,
      ::testing::WithParamInterface<OffloadParam<T>> {
  void SetUp() override {
    RETURN_ON_FATAL_FAILURE(OffloadTest::SetUp());

    auto &DeviceParam = std::get<0>(this->GetParam());
    Device = DeviceParam.Handle;
    if (Device == nullptr)
      GTEST_SKIP() << "No available devices.";

    ASSERT_SUCCESS(olCreateContext(1, &Device, &Context));
  }

  void TearDown() override {
    if (Context)
      olDestroyContext(Context);
    RETURN_ON_FATAL_FAILURE(OffloadTest::TearDown());
  }

  ol_platform_backend_t getPlatformBackend() const {
    ol_platform_handle_t Platform = nullptr;
    if (olGetDeviceInfo(Device, OL_DEVICE_INFO_PLATFORM,
                        sizeof(ol_platform_handle_t), &Platform))
      return OL_PLATFORM_BACKEND_UNKNOWN;
    ol_platform_backend_t Backend;
    if (olGetPlatformInfo(Platform, OL_PLATFORM_INFO_BACKEND,
                          sizeof(ol_platform_backend_t), &Backend))
      return OL_PLATFORM_BACKEND_UNKNOWN;
    return Backend;
  }

  const OffloadParam<T> &getParamTuple() const { return this->GetParam(); }

  const T &getTestParam() { return std::get<1>(getParamTuple()); }

  ol_device_handle_t Device = nullptr;
  ol_context_handle_t Context = nullptr;
};

// In order to avoid code duplication, the unparameterized versions of fixtures
// are aliases for parameterized fixtures, with `int` type chosen arbitrarily as
// an ignored parameter type. The single mock parameter of value `0` is combined
// with the devices in the provided macros, yielding tuples
// `std::tuple<TestEnvironment::Device, int>`. The hidden `int` parameter is not
// used, but it enables users to instantiate unparameterized tests without the
// knowledge about the details related to the implementation of fixtures.
// Moreover, it allows for modifying only one version of the fixture, without
// the need to also change the other version: either parameterized or
// unparameterized.
using OffloadDeviceTest = OffloadDeviceTestWithParam<int>;

template <typename T>
struct OffloadPlatformTestWithParam : OffloadDeviceTestWithParam<T> {
  void SetUp() override {
    RETURN_ON_FATAL_FAILURE(OffloadDeviceTestWithParam<T>::SetUp());

    ASSERT_SUCCESS(olGetDeviceInfo(this->Device, OL_DEVICE_INFO_PLATFORM,
                                   sizeof(Platform), &Platform));
    ASSERT_NE(Platform, nullptr);
  }

  ol_platform_handle_t Platform = nullptr;
};

using OffloadPlatformTest = OffloadPlatformTestWithParam<int>;

// Fixture for a generic program test. If you want a different program, use
// offloadQueueTest and create your own program handle with the binary you want.
template <typename T>
struct OffloadProgramTestWithParam : OffloadDeviceTestWithParam<T> {
  void SetUp() override { SetUpWith("foo"); }

  void SetUpWith(const char *ProgramName) {
    RETURN_ON_FATAL_FAILURE(OffloadDeviceTestWithParam<T>::SetUp());
    ASSERT_TRUE(TestEnvironment::loadDeviceBinary(ProgramName, this->Device,
                                                  DeviceBin));
    ASSERT_GE(DeviceBin->getBufferSize(), 0lu);
    ASSERT_SUCCESS(olCreateProgram(this->Context, this->Device,
                                   DeviceBin->getBufferStart(),
                                   DeviceBin->getBufferSize(), &Program));
  }

  void TearDown() override {
    if (Program) {
      olDestroyProgram(Program);
    }
    RETURN_ON_FATAL_FAILURE(OffloadDeviceTestWithParam<T>::TearDown());
  }

  ol_program_handle_t Program = nullptr;
  std::unique_ptr<llvm::MemoryBuffer> DeviceBin;
};

using OffloadProgramTest = OffloadProgramTestWithParam<int>;

struct OffloadKernelTest : OffloadProgramTest {
  void SetUp() override {
    RETURN_ON_FATAL_FAILURE(OffloadProgramTest::SetUp());
    ASSERT_SUCCESS(olGetSymbol(Program, "foo", OL_SYMBOL_KIND_KERNEL, &Kernel));
  }

  void TearDown() override {
    RETURN_ON_FATAL_FAILURE(OffloadProgramTest::TearDown());
  }

  ol_symbol_handle_t Kernel = nullptr;
};

template <typename T>
struct OffloadGlobalTestWithParam : OffloadProgramTestWithParam<T> {
  void SetUp() override {
    RETURN_ON_FATAL_FAILURE(
        OffloadProgramTestWithParam<T>::SetUpWith("global"));
    ASSERT_SUCCESS(olGetSymbol(this->Program, "global",
                               OL_SYMBOL_KIND_GLOBAL_VARIABLE, &Global));
  }

  void TearDown() override {
    RETURN_ON_FATAL_FAILURE(OffloadProgramTestWithParam<T>::TearDown());
  }

  ol_symbol_handle_t Global = nullptr;
};

using OffloadGlobalTest = OffloadGlobalTestWithParam<int>;

struct OffloadQueueTest : OffloadDeviceTest {
  void SetUp() override {
    RETURN_ON_FATAL_FAILURE(OffloadDeviceTest::SetUp());
    ASSERT_SUCCESS(olCreateQueue(Context, Device, &Queue));
  }

  void TearDown() override {
    if (Queue) {
      olDestroyQueue(Queue);
    }
    RETURN_ON_FATAL_FAILURE(OffloadDeviceTest::TearDown());
  }

  ol_queue_handle_t Queue = nullptr;
};

struct OffloadEventTest : OffloadQueueTest {
  void SetUp() override {
    RETURN_ON_FATAL_FAILURE(OffloadQueueTest::SetUp());
    ASSERT_SUCCESS(olCreateEvent(Queue, OL_EVENT_FLAGS_NONE, &Event));
    ASSERT_SUCCESS(olSyncQueue(Queue));
  }

  void TearDown() override {
    if (Event)
      olDestroyEvent(Event);
    RETURN_ON_FATAL_FAILURE(OffloadQueueTest::TearDown());
  }

  ol_event_handle_t Event = nullptr;
};

struct LaunchKernelTestBase : OffloadQueueTest {
  void SetUpProgram(const char *program) {
    RETURN_ON_FATAL_FAILURE(OffloadQueueTest::SetUp());
    ASSERT_TRUE(TestEnvironment::loadDeviceBinary(program, Device, DeviceBin));
    ASSERT_GE(DeviceBin->getBufferSize(), 0lu);
    ASSERT_SUCCESS(olCreateProgram(Context, Device, DeviceBin->getBufferStart(),
                                   DeviceBin->getBufferSize(), &Program));

    LaunchArgs.Dimensions = 1;
    LaunchArgs.GroupSize = {64, 1, 1};
    LaunchArgs.NumGroups = {1, 1, 1};
    LaunchArgs.DynSharedMemory = 0;
  }

  void TearDown() override {
    if (Program)
      olDestroyProgram(Program);
    RETURN_ON_FATAL_FAILURE(OffloadQueueTest::TearDown());
  }

  std::unique_ptr<llvm::MemoryBuffer> DeviceBin;
  ol_program_handle_t Program = nullptr;
  ol_kernel_launch_size_args_t LaunchArgs{};
};

struct LaunchSingleKernelTestBase : LaunchKernelTestBase {
  void SetUpKernel(const char *kernel) {
    RETURN_ON_FATAL_FAILURE(SetUpProgram(kernel));
    ASSERT_SUCCESS(
        olGetSymbol(Program, kernel, OL_SYMBOL_KIND_KERNEL, &Kernel));
  }

  ol_symbol_handle_t Kernel = nullptr;
};

using DevicesVec = std::vector<TestEnvironment::Device>;

inline DevicesVec getDevicesAndHost() {
  DevicesVec Res(TestEnvironment::getDevices());
  TestEnvironment::Device Host{TestEnvironment::getHostDevice(), "HOST"};

  Res.push_back(Host);

  return Res;
}

template <class T>
inline std::string
defaultPrinterWithParam(const ::testing::TestParamInfo<OffloadParam<T>> &info) {
  auto device = std::get<0>(info.param);
  auto param = std::get<1>(info.param);

  std::string placeholder;
  llvm::raw_string_ostream ss(placeholder);

  ss << device.Name << "__" << param;

  return SanitizeString(ss.str());
}

inline std::string
defaultPrinter(const ::testing::TestParamInfo<OffloadParam<int>> &info) {
  auto device = std::get<0>(info.param);

  return SanitizeString(device.Name);
}

// Devices might not be available for offload testing, so allow uninstantiated
// tests (as the device list will be empty). This means that all tests requiring
// a device will be silently skipped.
#define OFFLOAD_TESTS_INSTANTIATE_WITH_DEVICES(FIXTURE, DEVICES)               \
  INSTANTIATE_TEST_SUITE_P(                                                    \
      , FIXTURE,                                                               \
      testing::Combine(::testing::ValuesIn(DEVICES), testing::ValuesIn({0})),  \
      defaultPrinter);                                                         \
  GTEST_ALLOW_UNINSTANTIATED_PARAMETERIZED_TEST(FIXTURE)

#define OFFLOAD_TESTS_INSTANTIATE_DEVICE_FIXTURE(FIXTURE)                      \
  OFFLOAD_TESTS_INSTANTIATE_WITH_DEVICES(FIXTURE, TestEnvironment::getDevices())

#define OFFLOAD_TESTS_INSTANTIATE_HOST_DEVICE_FIXTURE(FIXTURE)                 \
  OFFLOAD_TESTS_INSTANTIATE_WITH_DEVICES(FIXTURE, getDevicesAndHost())

#define OFFLOAD_TESTS_INSTANTIATE_WITH_DEVICES_WITH_PARAM(FIXTURE, VALUES,     \
                                                          DEVICES, PRINTER)    \
  INSTANTIATE_TEST_SUITE_P(                                                    \
      , FIXTURE,                                                               \
      testing::Combine(::testing::ValuesIn(TestEnvironment::getDevices()),     \
                       ::testing::ValuesIn(VALUES)),                           \
      PRINTER);                                                                \
  GTEST_ALLOW_UNINSTANTIATED_PARAMETERIZED_TEST(FIXTURE)

#define OFFLOAD_TESTS_INSTANTIATE_DEVICE_FIXTURE_WITH_PARAM(FIXTURE, VALUES,   \
                                                            PRINTER)           \
  OFFLOAD_TESTS_INSTANTIATE_WITH_DEVICES_WITH_PARAM(                           \
      FIXTURE, VALUES, TestEnvironment::getDevices(), PRINTER)

#define OFFLOAD_TESTS_INSTANTIATE_HOST_DEVICE_FIXTURE_WITH_PARAM(              \
    FIXTURE, VALUES, PRINTER)                                                  \
  OFFLOAD_TESTS_INSTANTIATE_WITH_DEVICES_WITH_PARAM(                           \
      FIXTURE, VALUES, getDevicesAndHost(), PRINTER)