87 lines · cpp
1//===- Utils.cpp - MLIR ROCDL target utils ----------------------*- C++ -*-===//2//3// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.4// See https://llvm.org/LICENSE.txt for license information.5// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception6//7//===----------------------------------------------------------------------===//8//9// This files defines ROCDL target related utility classes and functions.10//11//===----------------------------------------------------------------------===//12 13#include "mlir/Target/LLVM/ROCDL/Utils.h"14#include "mlir/Dialect/GPU/IR/GPUDialect.h"15 16#include "llvm/ADT/StringMap.h"17#include "llvm/Frontend/Offloading/Utility.h"18 19using namespace mlir;20using namespace mlir::ROCDL;21 22std::optional<DenseMap<StringAttr, NamedAttrList>>23mlir::ROCDL::getAMDHSAKernelsELFMetadata(Builder &builder,24 ArrayRef<char> elfData) {25 uint16_t elfABIVersion;26 llvm::StringMap<llvm::offloading::amdgpu::AMDGPUKernelMetaData> kernels;27 llvm::MemoryBufferRef buffer(StringRef(elfData.data(), elfData.size()),28 "buffer");29 // Get the metadata.30 llvm::Error error = llvm::offloading::amdgpu::getAMDGPUMetaDataFromImage(31 buffer, kernels, elfABIVersion);32 // Return `nullopt` if the metadata couldn't be retrieved.33 if (error) {34 llvm::consumeError(std::move(error));35 return std::nullopt;36 }37 // Helper lambda for converting values.38 auto getI32Array = [&builder](const uint32_t *array) {39 return builder.getDenseI32ArrayAttr({static_cast<int32_t>(array[0]),40 static_cast<int32_t>(array[1]),41 static_cast<int32_t>(array[2])});42 };43 DenseMap<StringAttr, NamedAttrList> kernelMD;44 for (const auto &[name, kernel] : kernels) {45 NamedAttrList attrs;46 // Add kernel metadata.47 attrs.append("agpr_count", builder.getI64IntegerAttr(kernel.AGPRCount));48 attrs.append("sgpr_count", builder.getI64IntegerAttr(kernel.SGPRCount));49 attrs.append("vgpr_count", builder.getI64IntegerAttr(kernel.VGPRCount));50 attrs.append("sgpr_spill_count",51 builder.getI64IntegerAttr(kernel.SGPRSpillCount));52 attrs.append("vgpr_spill_count",53 builder.getI64IntegerAttr(kernel.VGPRSpillCount));54 attrs.append("wavefront_size",55 builder.getI64IntegerAttr(kernel.WavefrontSize));56 attrs.append("max_flat_workgroup_size",57 builder.getI64IntegerAttr(kernel.MaxFlatWorkgroupSize));58 attrs.append("group_segment_fixed_size",59 builder.getI64IntegerAttr(kernel.GroupSegmentList));60 attrs.append("private_segment_fixed_size",61 builder.getI64IntegerAttr(kernel.PrivateSegmentSize));62 attrs.append("reqd_workgroup_size",63 getI32Array(kernel.RequestedWorkgroupSize));64 attrs.append("workgroup_size_hint", getI32Array(kernel.WorkgroupSizeHint));65 kernelMD[builder.getStringAttr(name)] = std::move(attrs);66 }67 return std::move(kernelMD);68}69 70gpu::KernelTableAttr mlir::ROCDL::getKernelMetadata(Operation *gpuModule,71 ArrayRef<char> elfData) {72 auto module = cast<gpu::GPUModuleOp>(gpuModule);73 Builder builder(module.getContext());74 SmallVector<gpu::KernelMetadataAttr> kernels;75 std::optional<DenseMap<StringAttr, NamedAttrList>> mdMapOrNull =76 getAMDHSAKernelsELFMetadata(builder, elfData);77 for (auto funcOp : module.getBody()->getOps<LLVM::LLVMFuncOp>()) {78 if (!funcOp->getDiscardableAttr("rocdl.kernel"))79 continue;80 kernels.push_back(gpu::KernelMetadataAttr::get(81 funcOp, mdMapOrNull ? builder.getDictionaryAttr(82 mdMapOrNull->lookup(funcOp.getNameAttr()))83 : nullptr));84 }85 return gpu::KernelTableAttr::get(gpuModule->getContext(), kernels);86}87