ACCSpecializeForDevice.cpp

Go to the documentation of this file.

//===- ACCSpecializeForDevice.cpp -----------------------------------------===//
//
// 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 strips OpenACC constructs that are invalid or unnecessary inside
// device code (specialized acc routines or compute construct regions).
//
// Overview:
// ---------
// In a specialized acc routine or compute construct, many OpenACC operations
// do not make sense because they are host-side constructs. This pass removes
// or transforms these operations appropriately:
//
// - Data operations that manage device memory from host perspective
// - Compute constructs that launch kernels (we're already on device)
// - Runtime operations like init/shutdown/set/wait
//
// Transformations:
// ----------------
// The pass applies the following transformations:
//
// 1. Data Entry Ops (replaced with var operand):
//    acc.attach, acc.copyin, acc.create, acc.declare_device_resident,
//    acc.declare_link, acc.deviceptr, acc.get_deviceptr, acc.nocreate,
//    acc.present, acc.update_device, acc.use_device
//
// 2. Data Exit Ops (erased):
//    acc.copyout, acc.delete, acc.detach, acc.update_host
//
// 3. Structured Data/Compute Constructs (region inlined):
//    acc.data, acc.host_data, acc.kernel_environment, acc.parallel,
//    acc.serial, acc.kernels
//
// 4. Unstructured Data Ops (erased):
//    acc.enter_data, acc.exit_data, acc.update, acc.declare_enter,
//    acc.declare_exit
//
// 5. Runtime Ops (erased):
//    acc.init, acc.shutdown, acc.set, acc.wait
//
// 6. acc.on_device (folded):
//    acc.on_device with constant device type is folded to a boolean constant.
//    With a device type only known at runtime it is replaced by comparisons
//    against the device types being specialized for.
//
// Scope of Application:
// ---------------------
// - For functions with `acc.specialized_routine` attribute: patterns are
//   applied to the entire function body.
// - For non-specialized functions: patterns are applied only to ACC
//   operations INSIDE compute constructs (parallel, serial, kernels),
//   not to the compute constructs themselves or their data operands.
//
// Note: acc.cache, acc.private, acc.reduction, acc.firstprivate are NOT
// transformed by this pass as they are valid in device code.
//
//===----------------------------------------------------------------------===//
 
#include "mlir/Dialect/OpenACC/Extensions/AllExtensions.h"
#include "mlir/Dialect/OpenACC/Transforms/Passes.h"
 
#include "mlir/Dialect/Arith/IR/Arith.h"
#include "mlir/Dialect/Func/IR/FuncOps.h"
#include "mlir/Dialect/OpenACC/OpenACC.h"
#include "mlir/Dialect/OpenACC/Transforms/ACCSpecializePatterns.h"
#include "mlir/Dialect/SCF/IR/SCF.h"
#include "mlir/IR/MLIRContext.h"
#include "mlir/IR/PatternMatch.h"
#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
 
namespace mlir {
namespace acc {
#define GEN_PASS_DEF_ACCSPECIALIZEFORDEVICE
#include "mlir/Dialect/OpenACC/Transforms/Passes.h.inc"
} // namespace acc
} // namespace mlir
 
using namespace mlir;
using namespace mlir::acc;
 
namespace {
/// Fold acc.on_device in device code. A constant device type folds to a
/// boolean constant, any other device type to a comparison against the device
/// types being specialized for.
class FoldAccOnDeviceOpConversion : public OpRewritePattern<acc::OnDeviceOp> {
  using OpRewritePattern<acc::OnDeviceOp>::OpRewritePattern;
 
  const TypesForDevice &theDeviceTypes;
 
public:
  FoldAccOnDeviceOpConversion(MLIRContext *context,
                              const TypesForDevice &theDeviceTypes)
      : OpRewritePattern<acc::OnDeviceOp>(context),
        theDeviceTypes(theDeviceTypes) {}
 
  LogicalResult matchAndRewrite(acc::OnDeviceOp op,
                                PatternRewriter &rewriter) const override {
    Value deviceType = op.getDeviceType();
    APInt constVal;
    if (matchPattern(deviceType, m_ConstantInt(&constVal))) {
      bool result = theDeviceTypes.contains(constVal.getSExtValue());
      rewriter.replaceOpWithNewOp<arith::ConstantOp>(
          op, rewriter.getI1Type(), rewriter.getBoolAttr(result));
      return success();
    }
 
    // The device type is only known at runtime, inline the definition of
    // `on_device`.
    Location loc = op.getLoc();
    Value result;
    for (int64_t type : theDeviceTypes) {
      Value typeValue = arith::ConstantOp::create(
          rewriter, loc, deviceType.getType(),
          rewriter.getIntegerAttr(deviceType.getType(), type));
      Value isType = arith::CmpIOp::create(
          rewriter, loc, arith::CmpIPredicate::eq, deviceType, typeValue);
      if (result)
        result = arith::OrIOp::create(rewriter, loc, result, isType);
      else
        result = isType;
    }
    assert(result && "At least one device type expected.");
    rewriter.replaceOp(op, result);
    return success();
  }
};
 
class ACCSpecializeForDevice
    : public acc::impl::ACCSpecializeForDeviceBase<ACCSpecializeForDevice> {
public:
  using ACCSpecializeForDeviceBase<
      ACCSpecializeForDevice>::ACCSpecializeForDeviceBase;
 
  void getDependentDialects(DialectRegistry &registry) const override {
    ACCSpecializeForDeviceBase<ACCSpecializeForDevice>::getDependentDialects(
        registry);
    acc::registerAllExtensions(registry);
  }
 
  void runOnOperation() override {
    func::FuncOp func = getOperation();
 
    TypesForDevice types{theDeviceTypes.begin(), theDeviceTypes.end()};
 
    RewritePatternSet patterns(&getContext());
    acc::populateACCSpecializeForDevicePatterns(patterns, types);
    GreedyRewriteConfig config;
    config.setUseTopDownTraversal(true);
 
    if (acc::isSpecializedAccRoutine(func)) {
      // For specialized acc routines, apply patterns to the entire function
      (void)applyPatternsGreedily(func, std::move(patterns), config);
    } else {
      // For non-specialized functions, apply patterns only to ACC operations
      // inside compute constructs (not to the compute constructs themselves).
      // Use ExistingOps strictness so the greedy driver does not expand the
      // worklist to parent ops, which would accidentally unwrap the compute
      // construct (e.g. after inlining acc routines with their own data
      // regions).
      config.setStrictness(GreedyRewriteStrictness::ExistingOps);
      SmallVector<Operation *> opsToTransform;
      func.walk([&](Operation *op) {
        if (isa<ACC_COMPUTE_CONSTRUCT_OPS>(op)) {
          // Walk inside the compute construct and collect ACC ops
          op->walk([&](Operation *innerOp) {
            // Skip the compute construct itself
            if (innerOp == op)
              return;
            if (isa<acc::OpenACCDialect>(innerOp->getDialect()))
              opsToTransform.push_back(innerOp);
          });
        }
      });
      if (!opsToTransform.empty())
        (void)applyOpPatternsGreedily(opsToTransform, std::move(patterns),
                                      config);
    }
  }
};
} // namespace
 
//===----------------------------------------------------------------------===//
// Pattern population functions
//===----------------------------------------------------------------------===//
 
void mlir::acc::populateACCSpecializeForDevicePatterns(
    RewritePatternSet &patterns, const TypesForDevice &theDeviceTypes) {
  MLIRContext *context = patterns.getContext();
 
  // Declare patterns - erase declare_enter and its associated declare_exit
  patterns.insert<ACCDeclareEnterOpConversion>(context);
 
  // Data entry ops - replaced with their var operand
  // Note: acc.cache, acc.private, acc.reduction, acc.firstprivate are NOT
  // included here - they are valid in device code
  patterns.insert<ACCOpReplaceWithVarConversion<acc::AttachOp>,
                  ACCOpReplaceWithVarConversion<acc::CopyinOp>,
                  ACCOpReplaceWithVarConversion<acc::CreateOp>,
                  ACCOpReplaceWithVarConversion<acc::DeclareDeviceResidentOp>,
                  ACCOpReplaceWithVarConversion<acc::DeclareLinkOp>,
                  ACCOpReplaceWithVarConversion<acc::DevicePtrOp>,
                  ACCOpReplaceWithVarConversion<acc::GetDevicePtrOp>,
                  ACCOpReplaceWithVarConversion<acc::NoCreateOp>,
                  ACCOpReplaceWithVarConversion<acc::PresentOp>,
                  ACCOpReplaceWithVarConversion<acc::UpdateDeviceOp>,
                  ACCOpReplaceWithVarConversion<acc::UseDeviceOp>>(context);
 
  // Data exit ops - simply erased (no results)
  patterns.insert<ACCOpEraseConversion<acc::CopyoutOp>,
                  ACCOpEraseConversion<acc::DeleteOp>,
                  ACCOpEraseConversion<acc::DetachOp>,
                  ACCOpEraseConversion<acc::UpdateHostOp>>(context);
 
  // Structured data constructs - unwrap their regions
  patterns.insert<ACCRegionUnwrapConversion<acc::DataOp>,
                  ACCRegionUnwrapConversion<acc::HostDataOp>,
                  ACCRegionUnwrapConversion<acc::KernelEnvironmentOp>>(context);
 
  // Compute constructs - unwrap their regions
  patterns.insert<ACCRegionUnwrapConversion<acc::ParallelOp>,
                  ACCRegionUnwrapConversion<acc::SerialOp>,
                  ACCRegionUnwrapConversion<acc::KernelsOp>>(context);
 
  // Unstructured data operations - erase them
  patterns.insert<ACCOpEraseConversion<acc::EnterDataOp>,
                  ACCOpEraseConversion<acc::ExitDataOp>,
                  ACCOpEraseConversion<acc::UpdateOp>>(context);
 
  // Runtime operations - erase them
  patterns.insert<
      ACCOpEraseConversion<acc::InitOp>, ACCOpEraseConversion<acc::ShutdownOp>,
      ACCOpEraseConversion<acc::SetOp>, ACCOpEraseConversion<acc::WaitOp>>(
      context);
 
  // Fold acc.on_device calls so dead-code elimination can remove host-only
  // code paths in device code.
  patterns.insert<FoldAccOnDeviceOpConversion>(context, theDeviceTypes);
}