David Neto | dd99221 | 2017-06-23 17:47:55 -0400 | [diff] [blame] | 1 | // 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 Neto | 156783e | 2017-07-05 15:39:41 -0400 | [diff] [blame] | 15 | // 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 Neto | dd99221 | 2017-06-23 17:47:55 -0400 | [diff] [blame] | 27 | #include <cassert> |
David Neto | c6f3ab2 | 2018-04-06 18:02:31 -0400 | [diff] [blame] | 28 | #include <cstring> |
David Neto | dd99221 | 2017-06-23 17:47:55 -0400 | [diff] [blame] | 29 | |
David Neto | c6f3ab2 | 2018-04-06 18:02:31 -0400 | [diff] [blame] | 30 | #include "llvm/IR/Constants.h" |
| 31 | #include "llvm/IR/DerivedTypes.h" |
| 32 | #include "llvm/IR/Function.h" |
David Neto | c6f3ab2 | 2018-04-06 18:02:31 -0400 | [diff] [blame] | 33 | #include "llvm/IR/IRBuilder.h" |
Diego Novillo | 3cc8d7a | 2019-04-10 13:30:34 -0400 | [diff] [blame] | 34 | #include "llvm/IR/Instructions.h" |
David Neto | c6f3ab2 | 2018-04-06 18:02:31 -0400 | [diff] [blame] | 35 | #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 Neto | dd99221 | 2017-06-23 17:47:55 -0400 | [diff] [blame] | 41 | |
David Neto | 4feb7a4 | 2017-10-06 17:29:42 -0400 | [diff] [blame] | 42 | #include "ArgKind.h" |
David Neto | dd99221 | 2017-06-23 17:47:55 -0400 | [diff] [blame] | 43 | |
| 44 | using namespace llvm; |
| 45 | |
| 46 | #define DEBUG_TYPE "clusterpodkernelargs" |
| 47 | |
| 48 | namespace { |
| 49 | struct ClusterPodKernelArgumentsPass : public ModulePass { |
| 50 | static char ID; |
| 51 | ClusterPodKernelArgumentsPass() : ModulePass(ID) {} |
| 52 | |
| 53 | bool runOnModule(Module &M) override; |
| 54 | }; |
David Neto | 48f56a4 | 2017-10-06 16:44:25 -0400 | [diff] [blame] | 55 | |
David Neto | dd99221 | 2017-06-23 17:47:55 -0400 | [diff] [blame] | 56 | } // namespace |
| 57 | |
| 58 | char ClusterPodKernelArgumentsPass::ID = 0; |
| 59 | static RegisterPass<ClusterPodKernelArgumentsPass> |
Diego Novillo | 3cc8d7a | 2019-04-10 13:30:34 -0400 | [diff] [blame] | 60 | X("ClusterPodKernelArgumentsPass", "Cluster POD Kernel Arguments Pass"); |
David Neto | dd99221 | 2017-06-23 17:47:55 -0400 | [diff] [blame] | 61 | |
| 62 | namespace clspv { |
| 63 | llvm::ModulePass *createClusterPodKernelArgumentsPass() { |
| 64 | return new ClusterPodKernelArgumentsPass(); |
| 65 | } |
| 66 | } // namespace clspv |
| 67 | |
| 68 | bool 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 Novillo | 3cc8d7a | 2019-04-10 13:30:34 -0400 | [diff] [blame] | 86 | SmallVector<CallInst *, 8> CallList; |
David Neto | d5b3f98 | 2017-09-28 14:49:49 -0400 | [diff] [blame] | 87 | |
David Neto | c6f3ab2 | 2018-04-06 18:02:31 -0400 | [diff] [blame] | 88 | // 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 Novillo | 3cc8d7a | 2019-04-10 13:30:34 -0400 | [diff] [blame] | 92 | for (Function *F : WorkList) { |
David Neto | dd99221 | 2017-06-23 17:47:55 -0400 | [diff] [blame] | 93 | Changed = true; |
| 94 | |
David Neto | 156783e | 2017-07-05 15:39:41 -0400 | [diff] [blame] | 95 | // 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 Neto | c6f3ab2 | 2018-04-06 18:02:31 -0400 | [diff] [blame] | 101 | int new_index; |
David Neto | 156783e | 2017-07-05 15:39:41 -0400 | [diff] [blame] | 102 | // 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 PETIT | a353c83 | 2018-03-20 23:21:21 +0000 | [diff] [blame] | 106 | // Size of the argument |
| 107 | unsigned arg_size; |
Kévin Petit | 8bea15e | 2019-04-09 14:05:17 +0100 | [diff] [blame] | 108 | // Argument type. |
| 109 | clspv::ArgKind arg_kind; |
David Neto | c6f3ab2 | 2018-04-06 18:02:31 -0400 | [diff] [blame] | 110 | // 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 Neto | 156783e | 2017-07-05 15:39:41 -0400 | [diff] [blame] | 113 | }; |
| 114 | |
David Neto | dd99221 | 2017-06-23 17:47:55 -0400 | [diff] [blame] | 115 | // In OpenCL, kernel arguments are either pointers or POD. A composite with |
Kévin Petit | 921c1ab | 2019-03-19 21:25:44 +0000 | [diff] [blame] | 116 | // an element or member that is a pointer is not allowed. So we'll use POD |
David Neto | dd99221 | 2017-06-23 17:47:55 -0400 | [diff] [blame] | 117 | // as a shorthand for non-pointer. |
| 118 | |
| 119 | SmallVector<Type *, 8> PtrArgTys; |
| 120 | SmallVector<Type *, 8> PodArgTys; |
David Neto | 156783e | 2017-07-05 15:39:41 -0400 | [diff] [blame] | 121 | SmallVector<ArgMapping, 8> RemapInfo; |
| 122 | unsigned arg_index = 0; |
David Neto | c6f3ab2 | 2018-04-06 18:02:31 -0400 | [diff] [blame] | 123 | int new_index = 0; |
David Neto | dd99221 | 2017-06-23 17:47:55 -0400 | [diff] [blame] | 124 | for (Argument &Arg : F->args()) { |
| 125 | Type *ArgTy = Arg.getType(); |
| 126 | if (isa<PointerType>(ArgTy)) { |
| 127 | PtrArgTys.push_back(ArgTy); |
David Neto | 862b7d8 | 2018-06-14 18:48:37 -0400 | [diff] [blame] | 128 | const auto kind = clspv::GetArgKindForType(ArgTy); |
David Neto | c6f3ab2 | 2018-04-06 18:02:31 -0400 | [diff] [blame] | 129 | int spec_id = -1; |
David Neto | 862b7d8 | 2018-06-14 18:48:37 -0400 | [diff] [blame] | 130 | if (kind == clspv::ArgKind::Local) { |
David Neto | c6f3ab2 | 2018-04-06 18:02:31 -0400 | [diff] [blame] | 131 | spec_id = arg_spec_id_map[&Arg]; |
| 132 | assert(spec_id > 0); |
| 133 | } |
David Neto | 862b7d8 | 2018-06-14 18:48:37 -0400 | [diff] [blame] | 134 | RemapInfo.push_back({std::string(Arg.getName()), arg_index, new_index++, |
Kévin Petit | 8bea15e | 2019-04-09 14:05:17 +0100 | [diff] [blame] | 135 | 0u, 0u, kind, spec_id}); |
David Neto | dd99221 | 2017-06-23 17:47:55 -0400 | [diff] [blame] | 136 | } else { |
| 137 | PodArgTys.push_back(ArgTy); |
| 138 | } |
David Neto | 156783e | 2017-07-05 15:39:41 -0400 | [diff] [blame] | 139 | arg_index++; |
David Neto | dd99221 | 2017-06-23 17:47:55 -0400 | [diff] [blame] | 140 | } |
| 141 | |
David Neto | dd99221 | 2017-06-23 17:47:55 -0400 | [diff] [blame] | 142 | // Put the pointer arguments first, and then POD arguments struct last. |
David Neto | 2ded02e | 2017-10-23 15:30:59 -0400 | [diff] [blame] | 143 | // Use StructType::get so we reuse types where possible. |
| 144 | auto PodArgsStructTy = StructType::get(Context, PodArgTys); |
David Neto | dd99221 | 2017-06-23 17:47:55 -0400 | [diff] [blame] | 145 | SmallVector<Type *, 8> NewFuncParamTys(PtrArgTys); |
| 146 | NewFuncParamTys.push_back(PodArgsStructTy); |
| 147 | |
David Neto | 156783e | 2017-07-05 15:39:41 -0400 | [diff] [blame] | 148 | // We've recorded the remapping for pointer arguments. Now record the |
| 149 | // remapping for POD arguments. |
| 150 | { |
Kévin PETIT | a353c83 | 2018-03-20 23:21:21 +0000 | [diff] [blame] | 151 | const DataLayout DL(&M); |
| 152 | const auto StructLayout = DL.getStructLayout(PodArgsStructTy); |
David Neto | 156783e | 2017-07-05 15:39:41 -0400 | [diff] [blame] | 153 | arg_index = 0; |
| 154 | int pod_index = 0; |
David Neto | 156783e | 2017-07-05 15:39:41 -0400 | [diff] [blame] | 155 | for (Argument &Arg : F->args()) { |
| 156 | Type *ArgTy = Arg.getType(); |
| 157 | if (!isa<PointerType>(ArgTy)) { |
Kévin PETIT | a353c83 | 2018-03-20 23:21:21 +0000 | [diff] [blame] | 158 | unsigned arg_size = DL.getTypeStoreSize(ArgTy); |
David Neto | 156783e | 2017-07-05 15:39:41 -0400 | [diff] [blame] | 159 | RemapInfo.push_back( |
David Neto | c6f3ab2 | 2018-04-06 18:02:31 -0400 | [diff] [blame] | 160 | {std::string(Arg.getName()), arg_index, new_index, |
Diego Novillo | 3cc8d7a | 2019-04-10 13:30:34 -0400 | [diff] [blame] | 161 | unsigned(StructLayout->getElementOffset(pod_index++)), arg_size, |
| 162 | clspv::GetArgKindForType(ArgTy), -1}); |
David Neto | 156783e | 2017-07-05 15:39:41 -0400 | [diff] [blame] | 163 | } |
| 164 | arg_index++; |
| 165 | } |
| 166 | } |
| 167 | |
David Neto | dd99221 | 2017-06-23 17:47:55 -0400 | [diff] [blame] | 168 | 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 Petit | 921c1ab | 2019-03-19 21:25:44 +0000 | [diff] [blame] | 181 | // 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 Petit | 8bea15e | 2019-04-09 14:05:17 +0100 | [diff] [blame] | 193 | // Then attributes for non-POD parameters |
Kévin Petit | 921c1ab | 2019-03-19 21:25:44 +0000 | [diff] [blame] | 194 | for (auto &rinfo : RemapInfo) { |
Kévin Petit | 8bea15e | 2019-04-09 14:05:17 +0100 | [diff] [blame] | 195 | 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 Petit | 921c1ab | 2019-03-19 21:25:44 +0000 | [diff] [blame] | 198 | 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-baker | bccf62c | 2019-03-29 10:32:41 -0400 | [diff] [blame] | 204 | // And finally function attributes. |
Kévin Petit | 921c1ab | 2019-03-19 21:25:44 +0000 | [diff] [blame] | 205 | if (Attributes.hasAttributes(AttributeList::FunctionIndex)) { |
Kévin Petit | 921c1ab | 2019-03-19 21:25:44 +0000 | [diff] [blame] | 206 | auto idx = AttributeList::FunctionIndex; |
alan-baker | bccf62c | 2019-03-29 10:32:41 -0400 | [diff] [blame] | 207 | auto attrs = Attributes.getFnAttributes(); |
| 208 | AttrBuildInfo.push_back(std::make_pair(idx, attrs)); |
Kévin Petit | 921c1ab | 2019-03-19 21:25:44 +0000 | [diff] [blame] | 209 | } |
alan-baker | bccf62c | 2019-03-29 10:32:41 -0400 | [diff] [blame] | 210 | auto newAttributes = AttributeList::get(M.getContext(), AttrBuildInfo); |
Kévin Petit | 921c1ab | 2019-03-19 21:25:44 +0000 | [diff] [blame] | 211 | NewFunc->setAttributes(newAttributes); |
| 212 | |
David Neto | dd99221 | 2017-06-23 17:47:55 -0400 | [diff] [blame] | 213 | // 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 Neto | 156783e | 2017-07-05 15:39:41 -0400 | [diff] [blame] | 225 | 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 Neto | 48f56a4 | 2017-10-06 16:44:25 -0400 | [diff] [blame] | 238 | // - Argument type |
Diego Novillo | 3cc8d7a | 2019-04-10 13:30:34 -0400 | [diff] [blame] | 239 | LLVMContext &Context = M.getContext(); |
| 240 | SmallVector<Metadata *, 8> mappings; |
David Neto | 156783e | 2017-07-05 15:39:41 -0400 | [diff] [blame] | 241 | 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 Novillo | 3cc8d7a | 2019-04-10 13:30:34 -0400 | [diff] [blame] | 245 | auto *new_index_md = |
| 246 | ConstantAsMetadata::get(Builder.getInt32(arg_mapping.new_index)); |
David Neto | c6f3ab2 | 2018-04-06 18:02:31 -0400 | [diff] [blame] | 247 | auto *offset_md = |
David Neto | 156783e | 2017-07-05 15:39:41 -0400 | [diff] [blame] | 248 | ConstantAsMetadata::get(Builder.getInt32(arg_mapping.offset)); |
Kévin PETIT | a353c83 | 2018-03-20 23:21:21 +0000 | [diff] [blame] | 249 | auto *arg_size_md = |
| 250 | ConstantAsMetadata::get(Builder.getInt32(arg_mapping.arg_size)); |
Kévin Petit | 8bea15e | 2019-04-09 14:05:17 +0100 | [diff] [blame] | 251 | auto argKindName = GetArgKindName(arg_mapping.arg_kind); |
| 252 | auto *argtype_md = MDString::get(Context, argKindName); |
Diego Novillo | 3cc8d7a | 2019-04-10 13:30:34 -0400 | [diff] [blame] | 253 | 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 Neto | 156783e | 2017-07-05 15:39:41 -0400 | [diff] [blame] | 258 | mappings.push_back(arg_md); |
| 259 | } |
| 260 | |
| 261 | NewFunc->setMetadata("kernel_arg_map", MDNode::get(Context, mappings)); |
| 262 | } |
| 263 | |
David Neto | dd99221 | 2017-06-23 17:47:55 -0400 | [diff] [blame] | 264 | // 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 Neto | dd99221 | 2017-06-23 17:47:55 -0400 | [diff] [blame] | 277 | |
| 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 Petit | 98d9c33 | 2019-03-13 15:03:40 +0000 | [diff] [blame] | 298 | assert(ptrIndex == PtrArgTys.size()); |
| 299 | assert(podIndex != 0); |
| 300 | assert(podIndex == PodArgTys.size()); |
David Neto | dd99221 | 2017-06-23 17:47:55 -0400 | [diff] [blame] | 301 | |
| 302 | auto Call = Builder.CreateCall(F, CalleeArgs); |
| 303 | Call->setCallingConv(F->getCallingConv()); |
David Neto | d5b3f98 | 2017-09-28 14:49:49 -0400 | [diff] [blame] | 304 | CallList.push_back(Call); |
David Neto | dd99221 | 2017-06-23 17:47:55 -0400 | [diff] [blame] | 305 | |
| 306 | Builder.CreateRetVoid(); |
| 307 | } |
| 308 | |
David Neto | 482550a | 2018-03-24 05:21:07 -0700 | [diff] [blame] | 309 | // Inline the inner function. It's cleaner to do this. |
| 310 | for (CallInst *C : CallList) { |
| 311 | InlineFunctionInfo info; |
| 312 | Changed |= InlineFunction(C, info); |
David Neto | d5b3f98 | 2017-09-28 14:49:49 -0400 | [diff] [blame] | 313 | } |
| 314 | |
David Neto | dd99221 | 2017-06-23 17:47:55 -0400 | [diff] [blame] | 315 | return Changed; |
| 316 | } |