171 lines · cpp
1//===- MarkDeclareTarget.cpp -------------------------------------------===//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// Mark functions called from explicit target code as implicitly declare target.10//11//===----------------------------------------------------------------------===//12 13#include "flang/Optimizer/OpenMP/Passes.h"14#include "mlir/Dialect/Func/IR/FuncOps.h"15#include "mlir/Dialect/LLVMIR/LLVMDialect.h"16#include "mlir/Dialect/OpenMP/OpenMPDialect.h"17#include "mlir/IR/BuiltinDialect.h"18#include "mlir/IR/BuiltinOps.h"19#include "mlir/IR/Operation.h"20#include "mlir/IR/SymbolTable.h"21#include "mlir/Pass/Pass.h"22#include "mlir/Support/LLVM.h"23#include "llvm/ADT/SmallPtrSet.h"24#include "llvm/ADT/TypeSwitch.h"25 26namespace flangomp {27#define GEN_PASS_DEF_MARKDECLARETARGETPASS28#include "flang/Optimizer/OpenMP/Passes.h.inc"29} // namespace flangomp30 31namespace {32class MarkDeclareTargetPass33 : public flangomp::impl::MarkDeclareTargetPassBase<MarkDeclareTargetPass> {34 35 struct ParentInfo {36 mlir::omp::DeclareTargetDeviceType devTy;37 mlir::omp::DeclareTargetCaptureClause capClause;38 bool automap;39 };40 41 void processSymbolRef(mlir::SymbolRefAttr symRef, ParentInfo parentInfo,42 llvm::SmallPtrSet<mlir::Operation *, 16> visited) {43 if (auto currFOp =44 getOperation().lookupSymbol<mlir::func::FuncOp>(symRef)) {45 auto current = llvm::dyn_cast<mlir::omp::DeclareTargetInterface>(46 currFOp.getOperation());47 48 if (current.isDeclareTarget()) {49 auto currentDt = current.getDeclareTargetDeviceType();50 51 // Found the same function twice, with different device_types,52 // mark as Any as it belongs to both53 if (currentDt != parentInfo.devTy &&54 currentDt != mlir::omp::DeclareTargetDeviceType::any) {55 current.setDeclareTarget(mlir::omp::DeclareTargetDeviceType::any,56 current.getDeclareTargetCaptureClause(),57 current.getDeclareTargetAutomap());58 }59 } else {60 current.setDeclareTarget(parentInfo.devTy, parentInfo.capClause,61 parentInfo.automap);62 }63 64 markNestedFuncs(parentInfo, currFOp, visited);65 }66 }67 68 void processReductionRefs(std::optional<mlir::ArrayAttr> symRefs,69 ParentInfo parentInfo,70 llvm::SmallPtrSet<mlir::Operation *, 16> visited) {71 if (!symRefs)72 return;73 74 for (auto symRef : symRefs->getAsRange<mlir::SymbolRefAttr>()) {75 if (auto declareReductionOp =76 getOperation().lookupSymbol<mlir::omp::DeclareReductionOp>(77 symRef)) {78 markNestedFuncs(parentInfo, declareReductionOp, visited);79 }80 }81 }82 83 void84 processReductionClauses(mlir::Operation *op, ParentInfo parentInfo,85 llvm::SmallPtrSet<mlir::Operation *, 16> visited) {86 llvm::TypeSwitch<mlir::Operation &>(*op)87 .Case([&](mlir::omp::LoopOp op) {88 processReductionRefs(op.getReductionSyms(), parentInfo, visited);89 })90 .Case([&](mlir::omp::ParallelOp op) {91 processReductionRefs(op.getReductionSyms(), parentInfo, visited);92 })93 .Case([&](mlir::omp::SectionsOp op) {94 processReductionRefs(op.getReductionSyms(), parentInfo, visited);95 })96 .Case([&](mlir::omp::SimdOp op) {97 processReductionRefs(op.getReductionSyms(), parentInfo, visited);98 })99 .Case([&](mlir::omp::TargetOp op) {100 processReductionRefs(op.getInReductionSyms(), parentInfo, visited);101 })102 .Case([&](mlir::omp::TaskgroupOp op) {103 processReductionRefs(op.getTaskReductionSyms(), parentInfo, visited);104 })105 .Case([&](mlir::omp::TaskloopOp op) {106 processReductionRefs(op.getReductionSyms(), parentInfo, visited);107 processReductionRefs(op.getInReductionSyms(), parentInfo, visited);108 })109 .Case([&](mlir::omp::TaskOp op) {110 processReductionRefs(op.getInReductionSyms(), parentInfo, visited);111 })112 .Case([&](mlir::omp::TeamsOp op) {113 processReductionRefs(op.getReductionSyms(), parentInfo, visited);114 })115 .Case([&](mlir::omp::WsloopOp op) {116 processReductionRefs(op.getReductionSyms(), parentInfo, visited);117 })118 .Default([](mlir::Operation &) {});119 }120 121 void markNestedFuncs(ParentInfo parentInfo, mlir::Operation *currOp,122 llvm::SmallPtrSet<mlir::Operation *, 16> visited) {123 if (visited.contains(currOp))124 return;125 visited.insert(currOp);126 127 currOp->walk([&, this](mlir::Operation *op) {128 if (auto callOp = llvm::dyn_cast<mlir::CallOpInterface>(op)) {129 if (auto symRef = llvm::dyn_cast_if_present<mlir::SymbolRefAttr>(130 callOp.getCallableForCallee())) {131 processSymbolRef(symRef, parentInfo, visited);132 }133 }134 processReductionClauses(op, parentInfo, visited);135 });136 }137 138 // This pass executes on mlir::ModuleOp's marking functions contained within139 // as implicitly declare target if they are called from within an explicitly140 // marked declare target function or a target region (TargetOp)141 void runOnOperation() override {142 for (auto functionOp : getOperation().getOps<mlir::func::FuncOp>()) {143 auto declareTargetOp = llvm::dyn_cast<mlir::omp::DeclareTargetInterface>(144 functionOp.getOperation());145 if (declareTargetOp.isDeclareTarget()) {146 llvm::SmallPtrSet<mlir::Operation *, 16> visited;147 ParentInfo parentInfo{declareTargetOp.getDeclareTargetDeviceType(),148 declareTargetOp.getDeclareTargetCaptureClause(),149 declareTargetOp.getDeclareTargetAutomap()};150 markNestedFuncs(parentInfo, functionOp, visited);151 }152 }153 154 // TODO: Extend to work with reverse-offloading, this shouldn't155 // require too much effort, just need to check the device clause156 // when it's lowering has been implemented and change the157 // DeclareTargetDeviceType argument from nohost to host depending on158 // the contents of the device clause159 getOperation()->walk([&](mlir::omp::TargetOp tarOp) {160 llvm::SmallPtrSet<mlir::Operation *, 16> visited;161 ParentInfo parentInfo = {162 /*devTy=*/mlir::omp::DeclareTargetDeviceType::nohost,163 /*capClause=*/mlir::omp::DeclareTargetCaptureClause::to,164 /*automap=*/false,165 };166 markNestedFuncs(parentInfo, tarOp, visited);167 });168 }169};170} // namespace171