EmulateAtomics.cpp

Go to the documentation of this file.

//===- EmulateAtomics.cpp - Emulate unsupported AMDGPU atomics ------===//
//
// 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
//
//===----------------------------------------------------------------------===//
 
#include "mlir/Dialect/AMDGPU/Transforms/Passes.h"
 
#include "mlir/Dialect/AMDGPU/IR/AMDGPUDialect.h"
#include "mlir/Dialect/Arith/IR/Arith.h"
#include "mlir/Dialect/ControlFlow/IR/ControlFlowOps.h"
#include "mlir/Dialect/LLVMIR/ROCDLTargetInfo.h"
#include "mlir/Dialect/Vector/IR/VectorOps.h"
#include "mlir/IR/BuiltinAttributes.h"
#include "mlir/IR/TypeUtilities.h"
#include "mlir/Transforms/DialectConversion.h"
 
namespace mlir::amdgpu {
#define GEN_PASS_DEF_AMDGPUEMULATEATOMICSPASS
#include "mlir/Dialect/AMDGPU/Transforms/Passes.h.inc"
} // namespace mlir::amdgpu
 
using namespace mlir;
using namespace mlir::amdgpu;
 
namespace {
struct AmdgpuEmulateAtomicsPass
    : public amdgpu::impl::AmdgpuEmulateAtomicsPassBase<
          AmdgpuEmulateAtomicsPass> {
  using AmdgpuEmulateAtomicsPassBase<
      AmdgpuEmulateAtomicsPass>::AmdgpuEmulateAtomicsPassBase;
  void runOnOperation() override;
};
 
template <typename AtomicOp, typename ArithOp>
struct RawBufferAtomicByCasPattern : public OpConversionPattern<AtomicOp> {
  using OpConversionPattern<AtomicOp>::OpConversionPattern;
  using Adaptor = typename AtomicOp::Adaptor;
 
  LogicalResult
  matchAndRewrite(AtomicOp atomicOp, Adaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override;
};
} // namespace
 
namespace {
enum class DataArgAction : unsigned char {
  Duplicate,
  Drop,
};
} // namespace
 
// Fix up the fact that, when we're migrating from a general bugffer atomic
// to a load or to a CAS, the number of openrands, and thus the number of
// entries needed in operandSegmentSizes, needs to change. We use this method
// because we'd like to preserve unknown attributes on the atomic instead of
// discarding them.
static void patchOperandSegmentSizes(ArrayRef<NamedAttribute> attrs,
                                     SmallVectorImpl<NamedAttribute> &newAttrs,
                                     DataArgAction action) {
  newAttrs.reserve(attrs.size());
  for (NamedAttribute attr : attrs) {
    if (attr.getName().getValue() != "operandSegmentSizes") {
      newAttrs.push_back(attr);
      continue;
    }
    auto segmentAttr = cast<DenseI32ArrayAttr>(attr.getValue());
    MLIRContext *context = segmentAttr.getContext();
    DenseI32ArrayAttr newSegments;
    switch (action) {
    case DataArgAction::Drop:
      newSegments = DenseI32ArrayAttr::get(
          context, segmentAttr.asArrayRef().drop_front());
      break;
    case DataArgAction::Duplicate: {
      SmallVector<int32_t> newVals;
      ArrayRef<int32_t> oldVals = segmentAttr.asArrayRef();
      newVals.push_back(oldVals[0]);
      newVals.append(oldVals.begin(), oldVals.end());
      newSegments = DenseI32ArrayAttr::get(context, newVals);
      break;
    }
    }
    newAttrs.push_back(NamedAttribute(attr.getName(), newSegments));
  }
}
 
template <typename OpTy>
static typename OpTy::Properties
getPropertiesFromAttrs(OpBuilder &builder, Location loc,
                       ArrayRef<NamedAttribute> attrs) {
  typename OpTy::Properties properties{};
  OpTy::populateDefaultProperties(
      OperationName(OpTy::getOperationName(), builder.getContext()),
      properties);
  LogicalResult result =
      OpTy::setPropertiesFromAttr(properties, builder.getDictionaryAttr(attrs),
                                  [&]() { return emitError(loc); });
  assert(succeeded(result) && "failed to convert operation properties");
  (void)result;
  return properties;
}
 
// A helper function to flatten a vector value to a scalar containing its bits,
// returning the value itself if othetwise.
static Value flattenVecToBits(ConversionPatternRewriter &rewriter, Location loc,
                              Value val) {
  auto vectorType = dyn_cast<VectorType>(val.getType());
  if (!vectorType)
    return val;
 
  int64_t bitwidth =
      vectorType.getElementTypeBitWidth() * vectorType.getNumElements();
  Type allBitsType = rewriter.getIntegerType(bitwidth);
  auto allBitsVecType = VectorType::get({1}, allBitsType);
  Value bitcast = vector::BitCastOp::create(rewriter, loc, allBitsVecType, val);
  Value scalar = vector::ExtractOp::create(rewriter, loc, bitcast, 0);
  return scalar;
}
 
template <typename AtomicOp, typename ArithOp>
LogicalResult RawBufferAtomicByCasPattern<AtomicOp, ArithOp>::matchAndRewrite(
    AtomicOp atomicOp, Adaptor adaptor,
    ConversionPatternRewriter &rewriter) const {
  Location loc = atomicOp.getLoc();
 
  NamedAttrList origProperties;
  atomicOp->getName().walkInherentAttrs(atomicOp,
                                        [&](StringRef name, Attribute &attr) {
                                          origProperties.append(name, attr);
                                        });
  ArrayRef<NamedAttribute> discardableAttrs =
      atomicOp->getDiscardableAttrDictionary().getValue();
  ValueRange operands = adaptor.getOperands();
  Value data = operands.take_front()[0];
  ValueRange invariantArgs = operands.drop_front();
  Type dataType = data.getType();
 
  SmallVector<NamedAttribute> loadAttrs;
  patchOperandSegmentSizes(origProperties, loadAttrs, DataArgAction::Drop);
  auto loadProperties =
      getPropertiesFromAttrs<RawBufferLoadOp>(rewriter, loc, loadAttrs);
  Value initialLoad =
      RawBufferLoadOp::create(rewriter, loc, TypeRange{dataType}, invariantArgs,
                              loadProperties, discardableAttrs);
  Block *currentBlock = rewriter.getInsertionBlock();
  Block *afterAtomic =
      rewriter.splitBlock(currentBlock, rewriter.getInsertionPoint());
  afterAtomic->addArgument(dataType, loc);
  Block *loopBlock = rewriter.createBlock(afterAtomic, {dataType}, {loc});
 
  rewriter.setInsertionPointToEnd(currentBlock);
  cf::BranchOp::create(rewriter, loc, loopBlock, initialLoad);
 
  rewriter.setInsertionPointToEnd(loopBlock);
  Value prevLoad = loopBlock->getArgument(0);
  Value operated = ArithOp::create(rewriter, loc, data, prevLoad);
  dataType = operated.getType();
 
  SmallVector<NamedAttribute> cmpswapAttrs;
  patchOperandSegmentSizes(origProperties, cmpswapAttrs,
                           DataArgAction::Duplicate);
  SmallVector<Value> cmpswapArgs = {operated, prevLoad};
  cmpswapArgs.append(invariantArgs.begin(), invariantArgs.end());
  auto cmpswapProperties = getPropertiesFromAttrs<RawBufferAtomicCmpswapOp>(
      rewriter, loc, cmpswapAttrs);
  Value atomicRes = RawBufferAtomicCmpswapOp::create(
      rewriter, loc, TypeRange{dataType}, cmpswapArgs, cmpswapProperties,
      discardableAttrs);
 
  // We care about exact bitwise equality here, so do some bitcasts.
  // These will fold away during lowering to the ROCDL dialect, where
  // an int->float bitcast is introduced to account for the fact that cmpswap
  // only takes integer arguments.
 
  Value prevLoadForCompare = flattenVecToBits(rewriter, loc, prevLoad);
  Value atomicResForCompare = flattenVecToBits(rewriter, loc, atomicRes);
  if (auto floatDataTy = dyn_cast<FloatType>(dataType)) {
    Type equivInt = rewriter.getIntegerType(floatDataTy.getWidth());
    prevLoadForCompare =
        arith::BitcastOp::create(rewriter, loc, equivInt, prevLoad);
    atomicResForCompare =
        arith::BitcastOp::create(rewriter, loc, equivInt, atomicRes);
  }
  Value canLeave =
      arith::CmpIOp::create(rewriter, loc, arith::CmpIPredicate::eq,
                            atomicResForCompare, prevLoadForCompare);
  cf::CondBranchOp::create(rewriter, loc, canLeave, afterAtomic,
                           ValueRange{prevLoad}, loopBlock, atomicRes);
  rewriter.replaceOp(atomicOp, ValueRange{afterAtomic->getArgument(0)});
  return success();
}
 
/// Returns whether \p op can be lowered to a native buffer atomic fadd on
/// \p target.
static bool isFaddNativelySupported(const ROCDL::TargetInfo &target,
                                    RawBufferAtomicFaddOp op) {
  namespace AMDGPU = ::llvm::AMDGPU;
 
  // A target with only the no-return form can still perform the atomic, it
  // just cannot report the old value, so it suffices when the result is dead.
  bool hasFadd = target.has(AMDGPU::FEAT_ATOMIC_FADD_RTN_INSTS) ||
                 (target.has(AMDGPU::FEAT_ATOMIC_FADD_NO_RTN_INSTS) &&
                  op.getOldValue().use_empty());
  if (!hasFadd)
    return false;
 
  // The packed 16-bit forms are separate instructions with their own features.
  Type elemType = getElementTypeOrSelf(op.getValue().getType());
  if (isa<Float16Type>(elemType))
    return target.has(AMDGPU::FEAT_ATOMIC_BUFFER_GLOBAL_PK_ADD_F16_INSTS);
  if (isa<BFloat16Type>(elemType))
    return target.has(AMDGPU::FEAT_ATOMIC_BUFFER_PK_ADD_BF16_INST);
  return true;
}
 
void mlir::amdgpu::populateAmdgpuEmulateAtomicsPatterns(
    ConversionTarget &target, RewritePatternSet &patterns,
    const ROCDL::TargetInfo &targetInfo, PatternBenefit benefit) {
  namespace AMDGPU = ::llvm::AMDGPU;
 
  target.addDynamicallyLegalOp<RawBufferAtomicFaddOp>(
      [targetInfo](RawBufferAtomicFaddOp op) -> bool {
        return isFaddNativelySupported(targetInfo, op);
      });
 
  // Floating-point min and max are only emulated on gfx9; later generations
  // have them for every type this op accepts. f64 is the one gfx9 case that
  // may be native.
  if (targetInfo.isGeneration(9)) {
    target.addDynamicallyLegalOp<RawBufferAtomicFmaxOp>(
        [targetInfo](RawBufferAtomicFmaxOp op) -> bool {
          return op.getValue().getType().isF64() &&
                 targetInfo.has(AMDGPU::FEAT_ATOMIC_FMIN_FMAX_GLOBAL_F64);
        });
  }
  patterns.add<
      RawBufferAtomicByCasPattern<RawBufferAtomicFaddOp, arith::AddFOp>,
      RawBufferAtomicByCasPattern<RawBufferAtomicFmaxOp, arith::MaximumFOp>,
      RawBufferAtomicByCasPattern<RawBufferAtomicSmaxOp, arith::MaxSIOp>,
      RawBufferAtomicByCasPattern<RawBufferAtomicUminOp, arith::MinUIOp>>(
      patterns.getContext(), benefit);
}
 
void AmdgpuEmulateAtomicsPass::runOnOperation() {
  Operation *op = getOperation();
  FailureOr<ROCDL::TargetInfo> targetInfo =
      ROCDL::TargetInfo::get(ROCDL::resolveArchOption(arch, chipset),
                             /*waveSize=*/0, [&] { return op->emitError(); });
  if (failed(targetInfo))
    return signalPassFailure();
 
  MLIRContext &ctx = getContext();
  ConversionTarget target(ctx);
  RewritePatternSet patterns(&ctx);
  target.markUnknownOpDynamicallyLegal(
      [](Operation *op) -> bool { return true; });
 
  populateAmdgpuEmulateAtomicsPatterns(target, patterns, *targetInfo);
  if (failed(applyPartialConversion(op, target, std::move(patterns))))
    return signalPassFailure();
}