blob: 230096d3f0b7fe733596c4e0df799017e2e10e31 [file] [log] [blame]
David Netodd992212017-06-23 17:47:55 -04001// Copyright 2017 The Clspv Authors. All rights reserved.
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7// http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
David Neto156783e2017-07-05 15:39:41 -040015// Cluster POD kernel arguments.
16//
17// Collect plain-old-data kernel arguments and place them into a single
18// struct argument, at the end. Other arguments are pointers, and retain
19// their relative order.
20//
21// We will create a kernel function as the new entry point, and change
22// the original kernel function into a regular SPIR function. Key
23// kernel metadata is moved from the old function to the wrapper.
24// We also attach a "kernel_arg_map" metadata node to the function to
25// encode the mapping from old kernel argument to new kernel argument.
26
David Netodd992212017-06-23 17:47:55 -040027#include <cassert>
28
29#include <llvm/IR/Constants.h>
30#include <llvm/IR/DerivedTypes.h>
31#include <llvm/IR/Function.h>
32#include <llvm/IR/Instructions.h>
33#include <llvm/IR/IRBuilder.h>
David Neto156783e2017-07-05 15:39:41 -040034#include <llvm/IR/Metadata.h>
David Netodd992212017-06-23 17:47:55 -040035#include <llvm/IR/Module.h>
36#include <llvm/Pass.h>
David Netod5b3f982017-09-28 14:49:49 -040037#include <llvm/Support/CommandLine.h>
David Netodd992212017-06-23 17:47:55 -040038#include <llvm/Support/raw_ostream.h>
David Netod5b3f982017-09-28 14:49:49 -040039#include <llvm/Transforms/Utils/Cloning.h>
David Netodd992212017-06-23 17:47:55 -040040
David Neto4feb7a42017-10-06 17:29:42 -040041#include "ArgKind.h"
David Netodd992212017-06-23 17:47:55 -040042
43using namespace llvm;
44
45#define DEBUG_TYPE "clusterpodkernelargs"
46
David Netod5b3f982017-09-28 14:49:49 -040047// TODO(dneto): Remove this after experimentation.
48static llvm::cl::opt<bool> no_inline_pod_fn(
49 "no-inline-pod-inner-function", llvm::cl::init(false),
50 llvm::cl::desc("DEPRECATED. Avoid inlining the inner function created by "
51 "clustering pod kernel args"));
52
David Netodd992212017-06-23 17:47:55 -040053namespace {
54struct ClusterPodKernelArgumentsPass : public ModulePass {
55 static char ID;
56 ClusterPodKernelArgumentsPass() : ModulePass(ID) {}
57
58 bool runOnModule(Module &M) override;
59};
David Neto48f56a42017-10-06 16:44:25 -040060
David Netodd992212017-06-23 17:47:55 -040061} // namespace
62
63char ClusterPodKernelArgumentsPass::ID = 0;
64static RegisterPass<ClusterPodKernelArgumentsPass>
65 X("ClusterPodKernelArgumentsPass",
66 "Cluster POD Kernel Arguments Pass");
67
68namespace clspv {
69llvm::ModulePass *createClusterPodKernelArgumentsPass() {
70 return new ClusterPodKernelArgumentsPass();
71}
72} // namespace clspv
73
74bool ClusterPodKernelArgumentsPass::runOnModule(Module &M) {
75 bool Changed = false;
76 LLVMContext &Context = M.getContext();
77
78 SmallVector<Function *, 8> WorkList;
79
80 for (Function &F : M) {
81 if (F.isDeclaration() || F.getCallingConv() != CallingConv::SPIR_KERNEL) {
82 continue;
83 }
84 for (Argument &Arg : F.args()) {
85 if (!isa<PointerType>(Arg.getType())) {
86 WorkList.push_back(&F);
87 break;
88 }
89 }
90 }
91
David Netod5b3f982017-09-28 14:49:49 -040092 SmallVector<CallInst*, 8> CallList;
93
David Netodd992212017-06-23 17:47:55 -040094 for (Function* F : WorkList) {
95 Changed = true;
96
David Neto156783e2017-07-05 15:39:41 -040097 // An ArgMapping describes how a kernel argument is remapped.
98 struct ArgMapping {
99 std::string name;
100 // 0-based argument index in the old kernel function.
101 unsigned old_index;
102 // 0-based argument index in the new kernel function.
103 unsigned new_index;
104 // Offset of the argument value within the new kernel argument.
105 // This is always zero for non-POD arguments. For a POD argument,
106 // this is the byte offset within the POD arguments struct.
107 unsigned offset;
David Neto48f56a42017-10-06 16:44:25 -0400108 // Argument type. Same range of values as the result of
David Neto4feb7a42017-10-06 17:29:42 -0400109 // clspv::GetArgKindForType.
110 const char* arg_kind;
David Neto156783e2017-07-05 15:39:41 -0400111 };
112
David Netodd992212017-06-23 17:47:55 -0400113 // In OpenCL, kernel arguments are either pointers or POD. A composite with
114 // an element or memeber that is a pointer is not allowed. So we'll use POD
115 // as a shorthand for non-pointer.
116
117 SmallVector<Type *, 8> PtrArgTys;
118 SmallVector<Type *, 8> PodArgTys;
David Neto156783e2017-07-05 15:39:41 -0400119 SmallVector<ArgMapping, 8> RemapInfo;
120 unsigned arg_index = 0;
David Netodd992212017-06-23 17:47:55 -0400121 for (Argument &Arg : F->args()) {
122 Type *ArgTy = Arg.getType();
123 if (isa<PointerType>(ArgTy)) {
124 PtrArgTys.push_back(ArgTy);
David Neto156783e2017-07-05 15:39:41 -0400125 RemapInfo.push_back({std::string(Arg.getName()), arg_index,
David Neto48f56a42017-10-06 16:44:25 -0400126 unsigned(RemapInfo.size()), 0u,
David Neto4feb7a42017-10-06 17:29:42 -0400127 clspv::GetArgKindForType(ArgTy)});
David Netodd992212017-06-23 17:47:55 -0400128 } else {
129 PodArgTys.push_back(ArgTy);
130 }
David Neto156783e2017-07-05 15:39:41 -0400131 arg_index++;
David Netodd992212017-06-23 17:47:55 -0400132 }
133
David Netodd992212017-06-23 17:47:55 -0400134 // Put the pointer arguments first, and then POD arguments struct last.
135 auto PodArgsStructTy =
136 StructType::create(PodArgTys, F->getName().str() + ".podargs");
137 SmallVector<Type *, 8> NewFuncParamTys(PtrArgTys);
138 NewFuncParamTys.push_back(PodArgsStructTy);
139
David Neto156783e2017-07-05 15:39:41 -0400140 // We've recorded the remapping for pointer arguments. Now record the
141 // remapping for POD arguments.
142 {
143 const auto StructLayout =
144 M.getDataLayout().getStructLayout(PodArgsStructTy);
145 arg_index = 0;
146 int pod_index = 0;
147 const unsigned num_pointer_args = unsigned(RemapInfo.size());
148 for (Argument &Arg : F->args()) {
149 Type *ArgTy = Arg.getType();
150 if (!isa<PointerType>(ArgTy)) {
151 RemapInfo.push_back(
152 {std::string(Arg.getName()), arg_index, num_pointer_args,
David Neto48f56a42017-10-06 16:44:25 -0400153 unsigned(StructLayout->getElementOffset(pod_index++)),
David Neto4feb7a42017-10-06 17:29:42 -0400154 clspv::GetArgKindForType(ArgTy)});
David Neto156783e2017-07-05 15:39:41 -0400155 }
156 arg_index++;
157 }
158 }
159
David Netodd992212017-06-23 17:47:55 -0400160 FunctionType *NewFuncTy =
161 FunctionType::get(F->getReturnType(), NewFuncParamTys, false);
162
163 // Create the new function and set key properties.
164 auto NewFunc = Function::Create(NewFuncTy, F->getLinkage());
165 // The new function adopts the real name so that linkage to the outside
166 // world remains the same.
167 NewFunc->setName(F->getName());
168 F->setName(NewFunc->getName().str() + ".inner");
169
170 NewFunc->setCallingConv(F->getCallingConv());
171 F->setCallingConv(CallingConv::SPIR_FUNC);
172
173 NewFunc->setAttributes(F->getAttributes());
174 // Move OpenCL kernel named attributes.
175 // TODO(dneto): Attributes starting with kernel_arg_* should be rewritten
176 // to reflect change in the argument shape.
177 std::vector<const char *> Metadatas{
178 "reqd_work_group_size", "kernel_arg_addr_space",
179 "kernel_arg_access_qual", "kernel_arg_type",
180 "kernel_arg_base_type", "kernel_arg_type_qual"};
181 for (auto name : Metadatas) {
182 NewFunc->setMetadata(name, F->getMetadata(name));
183 F->setMetadata(name, nullptr);
184 }
185
David Neto156783e2017-07-05 15:39:41 -0400186 IRBuilder<> Builder(BasicBlock::Create(Context, "entry", NewFunc));
187
188 // Set kernel argument mapping metadata.
189 {
190 // Attach a metadata node named "kernel_arg_map" to the new kernel
191 // function. It is a tuple of nodes, each of which is a tuple for
192 // each argument, with members:
193 // - Argument name
194 // - Ordinal index in the original kernel function
195 // - Ordinal index in the new kernel function
196 // - Byte offset within the argument. This is always 0 for pointer
197 // arguments. For POD arguments this is the offest within the POD
198 // argument struct.
David Neto48f56a42017-10-06 16:44:25 -0400199 // - Argument type
David Neto156783e2017-07-05 15:39:41 -0400200 LLVMContext& Context = M.getContext();
201 SmallVector<Metadata*, 8> mappings;
202 for (auto &arg_mapping : RemapInfo) {
203 auto *name_md = MDString::get(Context, arg_mapping.name);
204 auto *old_index_md =
205 ConstantAsMetadata::get(Builder.getInt32(arg_mapping.old_index));
206 auto *new_index_md =
207 ConstantAsMetadata::get(Builder.getInt32(arg_mapping.new_index));
208 auto *offset =
209 ConstantAsMetadata::get(Builder.getInt32(arg_mapping.offset));
David Neto4feb7a42017-10-06 17:29:42 -0400210 auto *argtype_md = MDString::get(Context, arg_mapping.arg_kind);
David Neto48f56a42017-10-06 16:44:25 -0400211 auto *arg_md = MDNode::get(
212 Context, {name_md, old_index_md, new_index_md, offset, argtype_md});
David Neto156783e2017-07-05 15:39:41 -0400213 mappings.push_back(arg_md);
214 }
215
216 NewFunc->setMetadata("kernel_arg_map", MDNode::get(Context, mappings));
217 }
218
David Netodd992212017-06-23 17:47:55 -0400219 // Insert the function after the original, to preserve ordering
220 // in the module as much as possible.
221 auto &FunctionList = M.getFunctionList();
222 for (auto Iter = FunctionList.begin(), IterEnd = FunctionList.end();
223 Iter != IterEnd; ++Iter) {
224 if (&*Iter == F) {
225 FunctionList.insertAfter(Iter, NewFunc);
226 break;
227 }
228 }
229
230 // The body of the wrapper is essentially a call to the original function,
231 // but we have to unwrap the non-pointer arguments from the struct.
David Netodd992212017-06-23 17:47:55 -0400232
233 // Map the wrapper's arguments to the callee's arguments.
234 SmallVector<Argument *, 8> CallerArgs;
235 for (Argument &Arg : NewFunc->args()) {
236 CallerArgs.push_back(&Arg);
237 }
238 Argument *PodArg = CallerArgs.back();
239 PodArg->setName("podargs");
240
241 SmallVector<Value *, 8> CalleeArgs;
242 unsigned podIndex = 0;
243 unsigned ptrIndex = 0;
244 for (const Argument &Arg : F->args()) {
245 if (isa<PointerType>(Arg.getType())) {
246 CalleeArgs.push_back(CallerArgs[ptrIndex++]);
247 } else {
248 CalleeArgs.push_back(Builder.CreateExtractValue(PodArg, {podIndex++}));
249 }
250 CalleeArgs.back()->setName(Arg.getName());
251 }
252 assert(ptrIndex + podIndex == F->arg_size());
253 assert(ptrIndex = PtrArgTys.size());
254 assert(podIndex = PodArgTys.size());
255
256 auto Call = Builder.CreateCall(F, CalleeArgs);
257 Call->setCallingConv(F->getCallingConv());
David Netod5b3f982017-09-28 14:49:49 -0400258 CallList.push_back(Call);
David Netodd992212017-06-23 17:47:55 -0400259
260 Builder.CreateRetVoid();
261 }
262
David Netod5b3f982017-09-28 14:49:49 -0400263 if (!no_inline_pod_fn) {
264 for (CallInst *C : CallList) {
265 InlineFunctionInfo info;
266 Changed |= InlineFunction(C, info);
267 }
268 }
269
David Netodd992212017-06-23 17:47:55 -0400270 return Changed;
271}