aboutsummaryrefslogtreecommitdiff
path: root/llvm/tools/llvm-offload-wrapper/llvm-offload-wrapper.cpp
blob: d65b402571ae8c23b16f3746dd4edacba9f0d507 (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
//===- llvm-offload-wrapper: Create runtime registration code for devices -===//
//
// 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
//
//===----------------------------------------------------------------------===//
//
// Provides a utility for generating runtime registration code for device code.
// We take a binary image (CUDA fatbinary, HIP offload bundle, LLVM binary) and
// create a new IR module that calls the respective runtime to load it on the
// device.
//
//===----------------------------------------------------------------------===//

#include "llvm/Bitcode/BitcodeWriter.h"
#include "llvm/Frontend/Offloading/OffloadWrapper.h"
#include "llvm/Frontend/Offloading/Utility.h"
#include "llvm/Object/OffloadBinary.h"
#include "llvm/Support/CommandLine.h"
#include "llvm/Support/FileOutputBuffer.h"
#include "llvm/Support/FileSystem.h"
#include "llvm/Support/InitLLVM.h"
#include "llvm/Support/MemoryBuffer.h"
#include "llvm/Support/Path.h"
#include "llvm/Support/Signals.h"
#include "llvm/Support/StringSaver.h"
#include "llvm/Support/WithColor.h"
#include "llvm/TargetParser/Host.h"

using namespace llvm;

static cl::opt<bool> Help("h", cl::desc("Alias for -help"), cl::Hidden);

static cl::OptionCategory
    OffloadWrapeprCategory("llvm-offload-wrapper options");

static cl::opt<object::OffloadKind> Kind(
    "kind", cl::desc("Wrap for offload kind:"), cl::cat(OffloadWrapeprCategory),
    cl::Required,
    cl::values(clEnumValN(object::OFK_OpenMP, "openmp", "Wrap OpenMP binaries"),
               clEnumValN(object::OFK_Cuda, "cuda", "Wrap CUDA binaries"),
               clEnumValN(object::OFK_HIP, "hip", "Wrap HIP binaries")));

static cl::opt<std::string> OutputFile("o", cl::desc("Write output to <file>."),
                                       cl::value_desc("file"),
                                       cl::cat(OffloadWrapeprCategory));

static cl::list<std::string> InputFiles(cl::Positional,
                                        cl::desc("Wrap input from <file>"),
                                        cl::value_desc("file"), cl::OneOrMore,
                                        cl::cat(OffloadWrapeprCategory));

static cl::opt<std::string>
    TheTriple("triple", cl::desc("Target triple for the wrapper module"),
              cl::init(sys::getDefaultTargetTriple()),
              cl::cat(OffloadWrapeprCategory));

static Error wrapImages(ArrayRef<ArrayRef<char>> BuffersToWrap) {
  if (BuffersToWrap.size() > 1 &&
      (Kind == llvm::object::OFK_Cuda || Kind == llvm::object::OFK_HIP))
    return createStringError(
        "CUDA / HIP offloading uses a single fatbinary or offload bundle");

  LLVMContext Context;
  Module M("offload.wrapper.module", Context);
  M.setTargetTriple(Triple());

  switch (Kind) {
  case llvm::object::OFK_OpenMP:
    if (Error Err = offloading::wrapOpenMPBinaries(
            M, BuffersToWrap, offloading::getOffloadEntryArray(M),
            /*Suffix=*/"", /*Relocatable=*/false))
      return Err;
    break;
  case llvm::object::OFK_Cuda:
    if (Error Err = offloading::wrapCudaBinary(
            M, BuffersToWrap.front(), offloading::getOffloadEntryArray(M),
            /*Suffix=*/"", /*EmitSurfacesAndTextures=*/false))
      return Err;
    break;
  case llvm::object::OFK_HIP:
    if (Error Err = offloading::wrapHIPBinary(
            M, BuffersToWrap.front(), offloading::getOffloadEntryArray(M)))
      return Err;
    break;
  case llvm::object::OFK_SYCL:
    if (Error Err = offloading::wrapSYCLBinaries(M, BuffersToWrap.front()))
      return Err;
    break;
  default:
    return createStringError(getOffloadKindName(Kind) +
                             " wrapping is not supported");
  }

  int FD = -1;
  if (std::error_code EC = sys::fs::openFileForWrite(OutputFile, FD))
    return errorCodeToError(EC);
  llvm::raw_fd_ostream OS(FD, true);
  WriteBitcodeToFile(M, OS);

  return Error::success();
}

int main(int argc, char **argv) {
  InitLLVM X(argc, argv);
  cl::HideUnrelatedOptions(OffloadWrapeprCategory);
  cl::ParseCommandLineOptions(
      argc, argv,
      "Generate runtime registration code for a device binary image\n");

  if (Help) {
    cl::PrintHelpMessage();
    return EXIT_SUCCESS;
  }

  auto ReportError = [argv](Error E) {
    logAllUnhandledErrors(std::move(E), WithColor::error(errs(), argv[0]));
    exit(EXIT_FAILURE);
  };

  SmallVector<std::unique_ptr<MemoryBuffer>> Buffers;
  SmallVector<ArrayRef<char>> BuffersToWrap;
  for (StringRef Input : InputFiles) {
    ErrorOr<std::unique_ptr<MemoryBuffer>> BufferOrErr =
        MemoryBuffer::getFileOrSTDIN(Input);
    if (std::error_code EC = BufferOrErr.getError())
      ReportError(createFileError(Input, EC));
    std::unique_ptr<MemoryBuffer> &Buffer =
        Buffers.emplace_back(std::move(*BufferOrErr));
    BuffersToWrap.emplace_back(
        ArrayRef<char>(Buffer->getBufferStart(), Buffer->getBufferSize()));
  }

  if (Error Err = wrapImages(BuffersToWrap))
    ReportError(std::move(Err));

  return EXIT_SUCCESS;
}