github.com/johnnyeven/libtools@v0.0.0-20191126065708-61829c1adf46/third_party/mlir/lib/Conversion/GPUToCUDA/GenerateCubinAccessors.cpp (about)

     1  //===- GenerateCubinAccessors.cpp - MLIR GPU lowering passes --------------===//
     2  //
     3  // Copyright 2019 The MLIR Authors.
     4  //
     5  // Licensed under the Apache License, Version 2.0 (the "License");
     6  // you may not use this file except in compliance with the License.
     7  // You may obtain a copy of the License at
     8  //
     9  //   http://www.apache.org/licenses/LICENSE-2.0
    10  //
    11  // Unless required by applicable law or agreed to in writing, software
    12  // distributed under the License is distributed on an "AS IS" BASIS,
    13  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    14  // See the License for the specific language governing permissions and
    15  // limitations under the License.
    16  // =============================================================================
    17  //
    18  // This file implements a pass to generate LLVMIR functions that return the
    19  // data stored in nvvm.cubin char* blob.
    20  //
    21  //===----------------------------------------------------------------------===//
    22  
    23  #include "mlir/Conversion/GPUToCUDA/GPUToCUDAPass.h"
    24  #include "mlir/Dialect/GPU/GPUDialect.h"
    25  #include "mlir/Dialect/LLVMIR/LLVMDialect.h"
    26  #include "mlir/IR/Attributes.h"
    27  #include "mlir/IR/Builders.h"
    28  #include "mlir/IR/Function.h"
    29  #include "mlir/IR/Identifier.h"
    30  #include "mlir/IR/Module.h"
    31  #include "mlir/IR/StandardTypes.h"
    32  #include "mlir/Pass/Pass.h"
    33  #include "mlir/Pass/PassRegistry.h"
    34  
    35  #include "llvm/ADT/STLExtras.h"
    36  
    37  namespace mlir {
    38  namespace {
    39  
    40  // TODO(herhut): Move to shared location.
    41  constexpr const char *kCubinAnnotation = "nvvm.cubin";
    42  constexpr const char *kCubinGetterAnnotation = "nvvm.cubingetter";
    43  constexpr const char *kCubinGetterSuffix = "_cubin";
    44  constexpr const char *kCubinStorageSuffix = "_cubin_cst";
    45  
    46  /// A pass generating global strings and getter functions for all cubin blobs
    47  /// annotated on functions via the nvvm.cubin attribute.
    48  class GpuGenerateCubinAccessorsPass
    49      : public ModulePass<GpuGenerateCubinAccessorsPass> {
    50  private:
    51    LLVM::LLVMType getIndexType() {
    52      unsigned bits =
    53          llvmDialect->getLLVMModule().getDataLayout().getPointerSizeInBits();
    54      return LLVM::LLVMType::getIntNTy(llvmDialect, bits);
    55    }
    56  
    57    // Inserts a global constant string containing `blob` into the parent module
    58    // of `orig` and generates the function that returns the address of the first
    59    // character of this string.
    60    // TODO(herhut): consider fusing this pass with launch-func-to-cuda.
    61    void generate(FuncOp orig, StringAttr blob) {
    62      Location loc = orig.getLoc();
    63      SmallString<128> nameBuffer(orig.getName());
    64      auto module = orig.getParentOfType<ModuleOp>();
    65      assert(module && "function must belong to a module");
    66  
    67      // Insert the getter function just after the original function.
    68      OpBuilder moduleBuilder(module.getBody(), module.getBody()->begin());
    69      moduleBuilder.setInsertionPoint(orig.getOperation()->getNextNode());
    70      auto getterType = moduleBuilder.getFunctionType(
    71          llvm::None, LLVM::LLVMType::getInt8PtrTy(llvmDialect));
    72      nameBuffer.append(kCubinGetterSuffix);
    73      auto result = moduleBuilder.create<FuncOp>(
    74          loc, StringRef(nameBuffer), getterType, ArrayRef<NamedAttribute>());
    75      Block *entryBlock = result.addEntryBlock();
    76  
    77      // Drop the getter suffix before appending the storage suffix.
    78      nameBuffer.resize(orig.getName().size());
    79      nameBuffer.append(kCubinStorageSuffix);
    80  
    81      // Obtain the address of the first character of the global string containing
    82      // the cubin and return from the getter.
    83      OpBuilder builder(entryBlock);
    84      Value *startPtr = LLVM::createGlobalString(
    85          loc, builder, StringRef(nameBuffer), blob.getValue(), llvmDialect);
    86      builder.create<LLVM::ReturnOp>(loc, startPtr);
    87  
    88      // Store the name of the getter on the function for easier lookup.
    89      orig.setAttr(kCubinGetterAnnotation, builder.getSymbolRefAttr(result));
    90    }
    91  
    92  public:
    93    // Perform the conversion on the module.  This may insert globals, so it
    94    // cannot be done on multiple functions in parallel.
    95    void runOnModule() override {
    96      llvmDialect =
    97          getModule().getContext()->getRegisteredDialect<LLVM::LLVMDialect>();
    98  
    99      for (auto func : getModule().getOps<FuncOp>()) {
   100        StringAttr cubinBlob = func.getAttrOfType<StringAttr>(kCubinAnnotation);
   101        if (!cubinBlob)
   102          continue;
   103        generate(func, cubinBlob);
   104      }
   105    }
   106  
   107  private:
   108    LLVM::LLVMDialect *llvmDialect;
   109  };
   110  
   111  } // anonymous namespace
   112  
   113  std::unique_ptr<ModulePassBase> createGenerateCubinAccessorPass() {
   114    return std::make_unique<GpuGenerateCubinAccessorsPass>();
   115  }
   116  
   117  static PassRegistration<GpuGenerateCubinAccessorsPass>
   118      pass("generate-cubin-accessors",
   119           "Generate LLVMIR functions that give access to cubin data");
   120  
   121  } // namespace mlir