ACCImplicitRoutine.cpp

Go to the documentation of this file.

//===- ACCImplicitRoutine.cpp - OpenACC Implicit Routine Transform -------===//
//
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
// See https://llvm.org/LICENSE.txt for license information.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
//
//===----------------------------------------------------------------------===//
//
// This pass implements the implicit rules described in OpenACC specification
// for `Routine Directive` (OpenACC 3.4 spec, section 2.15.1).
//
// "If no explicit routine directive applies to a procedure whose definition
// appears in the program unit being compiled, then the implementation applies
// an implicit routine directive to that procedure if any of the following
// conditions holds:
// - The procedure is called or its address is accessed in a compute region."
//
// The specification further states:
// "When the implementation applies an implicit routine directive to a
// procedure, it must recursively apply implicit routine directives to other
// procedures for which the above rules specify relevant dependencies. Such
// dependencies can form a cycle, so the implementation must take care to avoid
// infinite recursion."
//
// This pass implements these requirements by:
// 1. Walking through all OpenACC compute constructs and functions already
//    marked with `acc routine` in the module and identifying function calls
//    within these regions.
// 2. Creating implicit `acc.routine` operations for functions that don't
//    already have routine declarations.
// 3. Recursively walking through all existing `acc routine` and creating
//    implicit routine operations for function calls within these routines,
//    while avoiding infinite recursion through proper tracking.
//
// Calls in a branch that does not execute for `the-device-types` are skipped.
// That list holds the runtime `acc_device_t` values for which `acc.on_device`
// is true on the target.
//
// Requirements:
// -------------
// To use this pass in a pipeline, the following requirements must be met:
//
// 1. Operation Interface Implementation: Operations that define functions
//    or call functions should implement `mlir::FunctionOpInterface` and
//    `mlir::CallOpInterface` respectively.
//
// 2. Analysis Registration (Optional): If custom behavior is needed for
//    determining if a symbol use is valid within GPU regions, the dialect
//    should pre-register the `acc::OpenACCSupport` analysis.
//===----------------------------------------------------------------------===//
 
#include "mlir/Dialect/OpenACC/Extensions/AllExtensions.h"
#include "mlir/Dialect/OpenACC/Transforms/Passes.h"
 
#include "mlir/Dialect/OpenACC/Analysis/OpenACCSupport.h"
#include "mlir/Dialect/OpenACC/OpenACC.h"
#include "mlir/Dialect/OpenACC/OpenACCUtils.h"
#include "mlir/IR/Builders.h"
#include "mlir/IR/BuiltinAttributes.h"
#include "mlir/IR/BuiltinOps.h"
#include "mlir/IR/Operation.h"
#include "mlir/IR/Value.h"
#include "mlir/Interfaces/CallInterfaces.h"
#include "mlir/Interfaces/FunctionInterfaces.h"
#include <queue>
 
#define DEBUG_TYPE "acc-implicit-routine"
 
namespace mlir {
namespace acc {
#define GEN_PASS_DEF_ACCIMPLICITROUTINE
#include "mlir/Dialect/OpenACC/Transforms/Passes.h.inc"
} // namespace acc
} // namespace mlir
 
namespace {
using namespace mlir;
 
class ACCImplicitRoutine
    : public acc::impl::ACCImplicitRoutineBase<ACCImplicitRoutine> {
private:
  unsigned routineCounter = 0;
  static constexpr llvm::StringRef accRoutinePrefix = "acc_routine_";
 
  // Count existing routine operations and update counter
  void initRoutineCounter(ModuleOp module) {
    module.walk([&](acc::RoutineOp routineOp) { routineCounter++; });
  }
 
  // Generate a unique name for the routine and create the routine operation
  acc::RoutineOp createRoutineOp(OpBuilder &builder, Location loc,
                                 FunctionOpInterface &callee) {
    std::string routineName =
        (accRoutinePrefix + std::to_string(routineCounter++)).str();
    auto routineOp = acc::RoutineOp::create(
        builder, loc,
        /* sym_name=*/builder.getStringAttr(routineName),
        /* sym_visibility=*/nullptr,
        /* func_name=*/
        mlir::SymbolRefAttr::get(builder.getContext(),
                                 builder.getStringAttr(callee.getName())),
        /* bindIdName=*/nullptr,
        /* bindStrName=*/nullptr,
        /* bindIdNameDeviceType=*/nullptr,
        /* bindStrNameDeviceType=*/nullptr,
        /* worker=*/nullptr,
        /* vector=*/nullptr,
        /* seq=*/nullptr,
        /* nohost=*/nullptr,
        /* implicit=*/builder.getUnitAttr(),
        /* gang=*/nullptr,
        /* gangDim=*/nullptr,
        /* gangDimDeviceType=*/nullptr);
 
    // Assert that the callee does not already have routine info attribute
    assert(!callee->hasDiscardableAttr(acc::getRoutineInfoAttrName()) &&
           "function is already associated with a routine");
 
    callee->setDiscardableAttr(
        acc::getRoutineInfoAttrName(),
        mlir::acc::RoutineInfoAttr::get(
            builder.getContext(),
            {mlir::SymbolRefAttr::get(builder.getContext(),
                                      builder.getStringAttr(routineName))}));
    return routineOp;
  }
 
  // Used to walk through a compute region looking for function calls.
  LogicalResult
  implicitRoutineForCallsInComputeRegions(Operation *op, SymbolTable &symTab,
                                          mlir::OpBuilder &builder,
                                          acc::OpenACCSupport &accSupport) {
    LogicalResult result = success();
    op->walk([&](CallOpInterface callOp) {
      if (acc::isInOffTargetBranch(callOp.getOperation(), *theDeviceTypes))
        return;
      if (!callOp.getCallableForCallee())
        return;
 
      auto calleeSymbolRef =
          dyn_cast<SymbolRefAttr>(callOp.getCallableForCallee());
      // When call is done through ssa value, the callee is not a symbol.
      // Skip it because we don't know the call target.
      if (!calleeSymbolRef)
        return;
 
      auto callee = symTab.lookup<FunctionOpInterface>(
          calleeSymbolRef.getLeafReference().str());
      // If the callee does not exist or is already a valid symbol for GPU
      // regions, skip it
      if (!callee)
        return;
      // Already a valid symbol for GPU regions (e.g. an existing routine or an
      // LLVM intrinsic), skip it.
      if (accSupport.isValidSymbolUse(callOp.getOperation(), calleeSymbolRef))
        return;
      // Without a definition in this compilation unit an implicit routine
      // cannot be applied, so the call target needs explicit routine
      // information.
      if (callee.isExternal()) {
        callOp->emitError() << "Calls in an acc compute region must be marked "
                               "with acc routine: "
                            << calleeSymbolRef.getLeafReference();
        result = failure();
        return;
      }
      builder.setInsertionPoint(callee);
      createRoutineOp(builder, callee.getLoc(), callee);
    });
    return result;
  }
 
  // Recursively handle calls within a routine operation
  LogicalResult implicitRoutineForCallsInRoutine(
      acc::RoutineOp routineOp, mlir::OpBuilder &builder,
      acc::OpenACCSupport &accSupport, acc::DeviceType targetDeviceType) {
    // When bind clause is used, it means that the target is different than the
    // function to which the `acc routine` is used with. Skip this case to
    // avoid implicitly recursively marking calls that would not end up on
    // device.
    if (routineOp.getEffectiveBindName(targetDeviceType))
      return success();
 
    SymbolTable symTab(routineOp->getParentOfType<ModuleOp>());
    std::queue<acc::RoutineOp> routineQueue;
    routineQueue.push(routineOp);
    LogicalResult result = success();
    while (!routineQueue.empty()) {
      auto currentRoutine = routineQueue.front();
      routineQueue.pop();
      auto func = symTab.lookup<FunctionOpInterface>(
          currentRoutine.getFuncName().getLeafReference());
      func.walk([&](CallOpInterface callOp) {
        if (acc::isInOffTargetBranch(callOp.getOperation(), *theDeviceTypes))
          return;
        if (!callOp.getCallableForCallee())
          return;
 
        auto calleeSymbolRef =
            dyn_cast<SymbolRefAttr>(callOp.getCallableForCallee());
        // When call is done through ssa value, the callee is not a symbol.
        // Skip it because we don't know the call target.
        if (!calleeSymbolRef)
          return;
 
        auto callee = symTab.lookup<FunctionOpInterface>(
            calleeSymbolRef.getLeafReference().str());
        // If the callee does not exist or is already a valid symbol for GPU
        // regions, skip it
        if (!callee)
          return;
        // Already a valid symbol for GPU regions (e.g. an existing routine or
        // an LLVM intrinsic), skip it.
        if (accSupport.isValidSymbolUse(callOp.getOperation(), calleeSymbolRef))
          return;
        // Without a definition in this compilation unit an implicit routine
        // cannot be applied, so the call target needs explicit routine
        // information.
        if (callee.isExternal()) {
          callOp->emitError()
              << "Calls in acc routine must also be marked with acc routine: "
              << calleeSymbolRef.getLeafReference();
          result = failure();
          return;
        }
        builder.setInsertionPoint(callee);
        auto newRoutineOp = createRoutineOp(builder, callee.getLoc(), callee);
        routineQueue.push(newRoutineOp);
      });
    }
    return result;
  }
 
public:
  using ACCImplicitRoutineBase<ACCImplicitRoutine>::ACCImplicitRoutineBase;
 
  void getDependentDialects(DialectRegistry &registry) const override {
    ACCImplicitRoutineBase<ACCImplicitRoutine>::getDependentDialects(registry);
    acc::registerAllExtensions(registry);
  }
 
  void runOnOperation() override {
    auto module = getOperation();
    mlir::OpBuilder builder(module.getContext());
    SymbolTable symTab(module);
    initRoutineCounter(module);
 
    acc::OpenACCSupport &accSupport = getAnalysis<acc::OpenACCSupport>();
 
    LogicalResult result = success();
 
    // Handle compute regions
    module.walk([&](Operation *op) {
      if (isa<ACC_COMPUTE_CONSTRUCT_OPS>(op))
        if (failed(implicitRoutineForCallsInComputeRegions(op, symTab, builder,
                                                           accSupport)))
          result = failure();
    });
 
    // Use the device type option from the pass options.
    acc::DeviceType targetDeviceType = deviceType;
 
    // Handle existing routines
    module.walk([&](acc::RoutineOp routineOp) {
      if (failed(implicitRoutineForCallsInRoutine(
              routineOp, builder, accSupport, targetDeviceType)))
        result = failure();
    });
 
    if (failed(result))
      return signalPassFailure();
  }
};
} // namespace