blob: da3f028814cd1625b774a839958ebaed07c860f5 [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>
David Netoc6f3ab22018-04-06 18:02:31 -040028#include <cstring>
David Netodd992212017-06-23 17:47:55 -040029
David Netoc6f3ab22018-04-06 18:02:31 -040030#include "llvm/IR/Constants.h"
31#include "llvm/IR/DerivedTypes.h"
32#include "llvm/IR/Function.h"
David Netoc6f3ab22018-04-06 18:02:31 -040033#include "llvm/IR/IRBuilder.h"
Diego Novillo3cc8d7a2019-04-10 13:30:34 -040034#include "llvm/IR/Instructions.h"
David Netoc6f3ab22018-04-06 18:02:31 -040035#include "llvm/IR/Metadata.h"
36#include "llvm/IR/Module.h"
37#include "llvm/Pass.h"
38#include "llvm/Support/CommandLine.h"
39#include "llvm/Support/raw_ostream.h"
40#include "llvm/Transforms/Utils/Cloning.h"
David Netodd992212017-06-23 17:47:55 -040041
David Neto4feb7a42017-10-06 17:29:42 -040042#include "ArgKind.h"
David Netodd992212017-06-23 17:47:55 -040043
44using namespace llvm;
45
46#define DEBUG_TYPE "clusterpodkernelargs"
47
48namespace {
49struct ClusterPodKernelArgumentsPass : public ModulePass {
50 static char ID;
51 ClusterPodKernelArgumentsPass() : ModulePass(ID) {}
52
53 bool runOnModule(Module &M) override;
54};
David Neto48f56a42017-10-06 16:44:25 -040055
David Netodd992212017-06-23 17:47:55 -040056} // namespace
57
58char ClusterPodKernelArgumentsPass::ID = 0;
59static RegisterPass<ClusterPodKernelArgumentsPass>
Diego Novillo3cc8d7a2019-04-10 13:30:34 -040060 X("ClusterPodKernelArgumentsPass", "Cluster POD Kernel Arguments Pass");
David Netodd992212017-06-23 17:47:55 -040061
62namespace clspv {
63llvm::ModulePass *createClusterPodKernelArgumentsPass() {
64 return new ClusterPodKernelArgumentsPass();
65}
66} // namespace clspv
67
68bool ClusterPodKernelArgumentsPass::runOnModule(Module &M) {
69 bool Changed = false;
70 LLVMContext &Context = M.getContext();
71
72 SmallVector<Function *, 8> WorkList;
73
74 for (Function &F : M) {
75 if (F.isDeclaration() || F.getCallingConv() != CallingConv::SPIR_KERNEL) {
76 continue;
77 }
78 for (Argument &Arg : F.args()) {
79 if (!isa<PointerType>(Arg.getType())) {
80 WorkList.push_back(&F);
81 break;
82 }
83 }
84 }
85
Diego Novillo3cc8d7a2019-04-10 13:30:34 -040086 SmallVector<CallInst *, 8> CallList;
David Netod5b3f982017-09-28 14:49:49 -040087
David Netoc6f3ab22018-04-06 18:02:31 -040088 // Note: The transformation done in this pass preserves the pointer-to-local
89 // arg to spec-id mapping.
90 clspv::ArgIdMapType arg_spec_id_map = clspv::AllocateArgSpecIds(M);
91
Diego Novillo3cc8d7a2019-04-10 13:30:34 -040092 for (Function *F : WorkList) {
David Netodd992212017-06-23 17:47:55 -040093 Changed = true;
94
David Neto156783e2017-07-05 15:39:41 -040095 // An ArgMapping describes how a kernel argument is remapped.
96 struct ArgMapping {
97 std::string name;
98 // 0-based argument index in the old kernel function.
99 unsigned old_index;
100 // 0-based argument index in the new kernel function.
David Netoc6f3ab22018-04-06 18:02:31 -0400101 int new_index;
David Neto156783e2017-07-05 15:39:41 -0400102 // Offset of the argument value within the new kernel argument.
103 // This is always zero for non-POD arguments. For a POD argument,
104 // this is the byte offset within the POD arguments struct.
105 unsigned offset;
Kévin PETITa353c832018-03-20 23:21:21 +0000106 // Size of the argument
107 unsigned arg_size;
Kévin Petit8bea15e2019-04-09 14:05:17 +0100108 // Argument type.
109 clspv::ArgKind arg_kind;
David Netoc6f3ab22018-04-06 18:02:31 -0400110 // If non-negative, this argument is a pointer-to-local, and the value
111 // here is the specialization constant id for the array size.
112 int spec_id;
David Neto156783e2017-07-05 15:39:41 -0400113 };
114
David Netodd992212017-06-23 17:47:55 -0400115 // In OpenCL, kernel arguments are either pointers or POD. A composite with
Kévin Petit921c1ab2019-03-19 21:25:44 +0000116 // an element or member that is a pointer is not allowed. So we'll use POD
David Netodd992212017-06-23 17:47:55 -0400117 // as a shorthand for non-pointer.
118
119 SmallVector<Type *, 8> PtrArgTys;
120 SmallVector<Type *, 8> PodArgTys;
David Neto156783e2017-07-05 15:39:41 -0400121 SmallVector<ArgMapping, 8> RemapInfo;
122 unsigned arg_index = 0;
David Netoc6f3ab22018-04-06 18:02:31 -0400123 int new_index = 0;
David Netodd992212017-06-23 17:47:55 -0400124 for (Argument &Arg : F->args()) {
125 Type *ArgTy = Arg.getType();
126 if (isa<PointerType>(ArgTy)) {
127 PtrArgTys.push_back(ArgTy);
David Neto862b7d82018-06-14 18:48:37 -0400128 const auto kind = clspv::GetArgKindForType(ArgTy);
David Netoc6f3ab22018-04-06 18:02:31 -0400129 int spec_id = -1;
David Neto862b7d82018-06-14 18:48:37 -0400130 if (kind == clspv::ArgKind::Local) {
David Netoc6f3ab22018-04-06 18:02:31 -0400131 spec_id = arg_spec_id_map[&Arg];
132 assert(spec_id > 0);
133 }
David Neto862b7d82018-06-14 18:48:37 -0400134 RemapInfo.push_back({std::string(Arg.getName()), arg_index, new_index++,
Kévin Petit8bea15e2019-04-09 14:05:17 +0100135 0u, 0u, kind, spec_id});
David Netodd992212017-06-23 17:47:55 -0400136 } else {
137 PodArgTys.push_back(ArgTy);
138 }
David Neto156783e2017-07-05 15:39:41 -0400139 arg_index++;
David Netodd992212017-06-23 17:47:55 -0400140 }
141
David Netodd992212017-06-23 17:47:55 -0400142 // Put the pointer arguments first, and then POD arguments struct last.
David Neto2ded02e2017-10-23 15:30:59 -0400143 // Use StructType::get so we reuse types where possible.
144 auto PodArgsStructTy = StructType::get(Context, PodArgTys);
David Netodd992212017-06-23 17:47:55 -0400145 SmallVector<Type *, 8> NewFuncParamTys(PtrArgTys);
146 NewFuncParamTys.push_back(PodArgsStructTy);
147
David Neto156783e2017-07-05 15:39:41 -0400148 // We've recorded the remapping for pointer arguments. Now record the
149 // remapping for POD arguments.
150 {
Kévin PETITa353c832018-03-20 23:21:21 +0000151 const DataLayout DL(&M);
152 const auto StructLayout = DL.getStructLayout(PodArgsStructTy);
David Neto156783e2017-07-05 15:39:41 -0400153 arg_index = 0;
154 int pod_index = 0;
David Neto156783e2017-07-05 15:39:41 -0400155 for (Argument &Arg : F->args()) {
156 Type *ArgTy = Arg.getType();
157 if (!isa<PointerType>(ArgTy)) {
Kévin PETITa353c832018-03-20 23:21:21 +0000158 unsigned arg_size = DL.getTypeStoreSize(ArgTy);
David Neto156783e2017-07-05 15:39:41 -0400159 RemapInfo.push_back(
David Netoc6f3ab22018-04-06 18:02:31 -0400160 {std::string(Arg.getName()), arg_index, new_index,
Diego Novillo3cc8d7a2019-04-10 13:30:34 -0400161 unsigned(StructLayout->getElementOffset(pod_index++)), arg_size,
162 clspv::GetArgKindForType(ArgTy), -1});
David Neto156783e2017-07-05 15:39:41 -0400163 }
164 arg_index++;
165 }
166 }
167
David Netodd992212017-06-23 17:47:55 -0400168 FunctionType *NewFuncTy =
169 FunctionType::get(F->getReturnType(), NewFuncParamTys, false);
170
171 // Create the new function and set key properties.
172 auto NewFunc = Function::Create(NewFuncTy, F->getLinkage());
173 // The new function adopts the real name so that linkage to the outside
174 // world remains the same.
175 NewFunc->setName(F->getName());
176 F->setName(NewFunc->getName().str() + ".inner");
177
178 NewFunc->setCallingConv(F->getCallingConv());
179 F->setCallingConv(CallingConv::SPIR_FUNC);
180
Kévin Petit921c1ab2019-03-19 21:25:44 +0000181 // Transfer attributes that don't apply to the POD arguments
182 // to the new functions.
183 auto Attributes = F->getAttributes();
184 SmallVector<std::pair<unsigned, AttributeSet>, 8> AttrBuildInfo;
185
186 // Return attributes have to come first
187 if (Attributes.hasAttributes(AttributeList::ReturnIndex)) {
188 auto idx = AttributeList::ReturnIndex;
189 auto attrs = Attributes.getRetAttributes();
190 AttrBuildInfo.push_back(std::make_pair(idx, attrs));
191 }
192
Kévin Petit8bea15e2019-04-09 14:05:17 +0100193 // Then attributes for non-POD parameters
Kévin Petit921c1ab2019-03-19 21:25:44 +0000194 for (auto &rinfo : RemapInfo) {
Kévin Petit8bea15e2019-04-09 14:05:17 +0100195 bool argIsPod = rinfo.arg_kind == clspv::ArgKind::Pod ||
196 rinfo.arg_kind == clspv::ArgKind::PodUBO;
197 if (!argIsPod && Attributes.hasParamAttrs(rinfo.old_index)) {
Kévin Petit921c1ab2019-03-19 21:25:44 +0000198 auto idx = rinfo.new_index + AttributeList::FirstArgIndex;
199 auto attrs = Attributes.getParamAttributes(rinfo.old_index);
200 AttrBuildInfo.push_back(std::make_pair(idx, attrs));
201 }
202 }
203
alan-bakerbccf62c2019-03-29 10:32:41 -0400204 // And finally function attributes.
Kévin Petit921c1ab2019-03-19 21:25:44 +0000205 if (Attributes.hasAttributes(AttributeList::FunctionIndex)) {
Kévin Petit921c1ab2019-03-19 21:25:44 +0000206 auto idx = AttributeList::FunctionIndex;
alan-bakerbccf62c2019-03-29 10:32:41 -0400207 auto attrs = Attributes.getFnAttributes();
208 AttrBuildInfo.push_back(std::make_pair(idx, attrs));
Kévin Petit921c1ab2019-03-19 21:25:44 +0000209 }
alan-bakerbccf62c2019-03-29 10:32:41 -0400210 auto newAttributes = AttributeList::get(M.getContext(), AttrBuildInfo);
Kévin Petit921c1ab2019-03-19 21:25:44 +0000211 NewFunc->setAttributes(newAttributes);
212
David Netodd992212017-06-23 17:47:55 -0400213 // Move OpenCL kernel named attributes.
214 // TODO(dneto): Attributes starting with kernel_arg_* should be rewritten
215 // to reflect change in the argument shape.
216 std::vector<const char *> Metadatas{
217 "reqd_work_group_size", "kernel_arg_addr_space",
218 "kernel_arg_access_qual", "kernel_arg_type",
219 "kernel_arg_base_type", "kernel_arg_type_qual"};
220 for (auto name : Metadatas) {
221 NewFunc->setMetadata(name, F->getMetadata(name));
222 F->setMetadata(name, nullptr);
223 }
224
David Neto156783e2017-07-05 15:39:41 -0400225 IRBuilder<> Builder(BasicBlock::Create(Context, "entry", NewFunc));
226
227 // Set kernel argument mapping metadata.
228 {
229 // Attach a metadata node named "kernel_arg_map" to the new kernel
230 // function. It is a tuple of nodes, each of which is a tuple for
231 // each argument, with members:
232 // - Argument name
233 // - Ordinal index in the original kernel function
234 // - Ordinal index in the new kernel function
235 // - Byte offset within the argument. This is always 0 for pointer
236 // arguments. For POD arguments this is the offest within the POD
237 // argument struct.
David Neto48f56a42017-10-06 16:44:25 -0400238 // - Argument type
Diego Novillo3cc8d7a2019-04-10 13:30:34 -0400239 LLVMContext &Context = M.getContext();
240 SmallVector<Metadata *, 8> mappings;
David Neto156783e2017-07-05 15:39:41 -0400241 for (auto &arg_mapping : RemapInfo) {
242 auto *name_md = MDString::get(Context, arg_mapping.name);
243 auto *old_index_md =
244 ConstantAsMetadata::get(Builder.getInt32(arg_mapping.old_index));
Diego Novillo3cc8d7a2019-04-10 13:30:34 -0400245 auto *new_index_md =
246 ConstantAsMetadata::get(Builder.getInt32(arg_mapping.new_index));
David Netoc6f3ab22018-04-06 18:02:31 -0400247 auto *offset_md =
David Neto156783e2017-07-05 15:39:41 -0400248 ConstantAsMetadata::get(Builder.getInt32(arg_mapping.offset));
Kévin PETITa353c832018-03-20 23:21:21 +0000249 auto *arg_size_md =
250 ConstantAsMetadata::get(Builder.getInt32(arg_mapping.arg_size));
Kévin Petit8bea15e2019-04-09 14:05:17 +0100251 auto argKindName = GetArgKindName(arg_mapping.arg_kind);
252 auto *argtype_md = MDString::get(Context, argKindName);
Diego Novillo3cc8d7a2019-04-10 13:30:34 -0400253 auto *spec_id_md =
254 ConstantAsMetadata::get(Builder.getInt32(arg_mapping.spec_id));
255 auto *arg_md = MDNode::get(
256 Context, {name_md, old_index_md, new_index_md, offset_md,
257 arg_size_md, argtype_md, spec_id_md});
David Neto156783e2017-07-05 15:39:41 -0400258 mappings.push_back(arg_md);
259 }
260
261 NewFunc->setMetadata("kernel_arg_map", MDNode::get(Context, mappings));
262 }
263
David Netodd992212017-06-23 17:47:55 -0400264 // Insert the function after the original, to preserve ordering
265 // in the module as much as possible.
266 auto &FunctionList = M.getFunctionList();
267 for (auto Iter = FunctionList.begin(), IterEnd = FunctionList.end();
268 Iter != IterEnd; ++Iter) {
269 if (&*Iter == F) {
270 FunctionList.insertAfter(Iter, NewFunc);
271 break;
272 }
273 }
274
275 // The body of the wrapper is essentially a call to the original function,
276 // but we have to unwrap the non-pointer arguments from the struct.
David Netodd992212017-06-23 17:47:55 -0400277
278 // Map the wrapper's arguments to the callee's arguments.
279 SmallVector<Argument *, 8> CallerArgs;
280 for (Argument &Arg : NewFunc->args()) {
281 CallerArgs.push_back(&Arg);
282 }
283 Argument *PodArg = CallerArgs.back();
284 PodArg->setName("podargs");
285
286 SmallVector<Value *, 8> CalleeArgs;
287 unsigned podIndex = 0;
288 unsigned ptrIndex = 0;
289 for (const Argument &Arg : F->args()) {
290 if (isa<PointerType>(Arg.getType())) {
291 CalleeArgs.push_back(CallerArgs[ptrIndex++]);
292 } else {
293 CalleeArgs.push_back(Builder.CreateExtractValue(PodArg, {podIndex++}));
294 }
295 CalleeArgs.back()->setName(Arg.getName());
296 }
297 assert(ptrIndex + podIndex == F->arg_size());
Kévin Petit98d9c332019-03-13 15:03:40 +0000298 assert(ptrIndex == PtrArgTys.size());
299 assert(podIndex != 0);
300 assert(podIndex == PodArgTys.size());
David Netodd992212017-06-23 17:47:55 -0400301
302 auto Call = Builder.CreateCall(F, CalleeArgs);
303 Call->setCallingConv(F->getCallingConv());
David Netod5b3f982017-09-28 14:49:49 -0400304 CallList.push_back(Call);
David Netodd992212017-06-23 17:47:55 -0400305
306 Builder.CreateRetVoid();
307 }
308
David Neto482550a2018-03-24 05:21:07 -0700309 // Inline the inner function. It's cleaner to do this.
310 for (CallInst *C : CallList) {
311 InlineFunctionInfo info;
312 Changed |= InlineFunction(C, info);
David Netod5b3f982017-09-28 14:49:49 -0400313 }
314
David Netodd992212017-06-23 17:47:55 -0400315 return Changed;
316}