| //===- ACCAtomicPatterns.cpp - ACC atomic to LLVM patterns ------*- C++ -*-===// |
| // |
| // 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 |
| // |
| //===----------------------------------------------------------------------===// |
| // |
| // Lowers OpenACC atomic operations (read, write, update, capture) to LLVM |
| // dialect atomicrmw / cmpxchg sequences. |
| // |
| //===----------------------------------------------------------------------===// |
| |
| #include "mlir/Conversion/OpenACCToLLVM/ACCToLLVM.h" |
| |
| #include "mlir/Conversion/LLVMCommon/Pattern.h" |
| #include "mlir/Dialect/Arith/IR/Arith.h" |
| #include "mlir/Dialect/Complex/IR/Complex.h" |
| #include "mlir/Dialect/LLVMIR/LLVMDialect.h" |
| #include "mlir/Dialect/MemRef/IR/MemRef.h" |
| #include "mlir/Dialect/OpenACC/Analysis/OpenACCSupport.h" |
| #include "mlir/Dialect/OpenACC/OpenACC.h" |
| #include "mlir/Dialect/OpenACC/OpenACCUtilsType.h" |
| #include "mlir/IR/BuiltinOps.h" |
| #include "mlir/Interfaces/SideEffectInterfaces.h" |
| #include "llvm/ADT/TypeSwitch.h" |
| #include "llvm/Support/Debug.h" |
| |
| #include <cstdint> |
| #include <optional> |
| #include <set> |
| #include <utility> |
| |
| #define DEBUG_TYPE "acc-atomic-patterns" |
| |
| using namespace mlir; |
| using namespace mlir::acc; |
| |
| namespace { |
| |
| constexpr uint64_t kBitsInByte = 8; |
| |
| template <typename AtomicOpTy> |
| class ACCAtomicOpConversion : public ConvertOpToLLVMPattern<AtomicOpTy> { |
| using typename ConvertOpToLLVMPattern<AtomicOpTy>::OpAdaptor; |
| |
| public: |
| ACCAtomicOpConversion(const LLVMTypeConverter &typeConverter, |
| OpenACCSupport &accSupport, |
| const ACCAtomicLoadAddressCallback &getLoadAddress) |
| : ConvertOpToLLVMPattern<AtomicOpTy>(typeConverter), |
| accSupport(accSupport), getLoadAddress(getLoadAddress) {} |
| |
| LogicalResult |
| matchAndRewrite(AtomicOpTy op, OpAdaptor adaptor, |
| ConversionPatternRewriter &rewriter) const override; |
| |
| private: |
| OpenACCSupport &accSupport; |
| ACCAtomicLoadAddressCallback getLoadAddress; |
| |
| size_t getComplexStructElementSizeInBits(Type ty) const { |
| auto structType = dyn_cast<LLVM::LLVMStructType>(ty); |
| if (!structType || structType.getBody().size() != 2 || |
| structType.getBody()[0] != structType.getBody()[1] || |
| !structType.getBody()[0].isIntOrFloat()) |
| return 0; |
| size_t elementSizeInBits = structType.getBody()[0].getIntOrFloatBitWidth(); |
| if (elementSizeInBits == 0 || elementSizeInBits > 64) |
| llvm_unreachable("unexpected complex type"); |
| return elementSizeInBits; |
| } |
| |
| /// Serialize a complex value into an integer. |
| Value serializeExpr(Value expr, ConversionPatternRewriter &rewriter) const { |
| size_t elementSizeInBits = |
| getComplexStructElementSizeInBits(expr.getType()); |
| if (!elementSizeInBits) |
| return expr; |
| |
| Location loc = expr.getLoc(); |
| Value firstValue = LLVM::ExtractValueOp::create(rewriter, loc, expr, 0); |
| Value secondValue = LLVM::ExtractValueOp::create(rewriter, loc, expr, 1); |
| MLIRContext *context = rewriter.getContext(); |
| Type intCastTy = IntegerType::get(context, elementSizeInBits); |
| Type intTy = IntegerType::get(context, elementSizeInBits * 2); |
| firstValue = LLVM::BitcastOp::create(rewriter, loc, intCastTy, firstValue); |
| secondValue = |
| LLVM::BitcastOp::create(rewriter, loc, intCastTy, secondValue); |
| firstValue = LLVM::ZExtOp::create(rewriter, loc, intTy, firstValue); |
| secondValue = LLVM::ZExtOp::create(rewriter, loc, intTy, secondValue); |
| Value shlVal = |
| LLVM::ConstantOp::create(rewriter, loc, intTy, elementSizeInBits); |
| Value result = LLVM::ShlOp::create(rewriter, loc, secondValue, shlVal); |
| return LLVM::OrOp::create(rewriter, loc, result, firstValue); |
| } |
| |
| /// Deserialize an integer into a complex value. |
| Value deserializeExpr(Value expr, Type origTy, |
| ConversionPatternRewriter &rewriter) const { |
| size_t elementSizeInBits = getComplexStructElementSizeInBits(origTy); |
| if (!elementSizeInBits) |
| return expr; |
| |
| auto intTy = dyn_cast<IntegerType>(expr.getType()); |
| if (!intTy || intTy.getWidth() != elementSizeInBits * 2) |
| return expr; |
| |
| Location loc = expr.getLoc(); |
| MLIRContext *context = rewriter.getContext(); |
| Type elemTy = IntegerType::get(context, elementSizeInBits); |
| Value low = LLVM::TruncOp::create(rewriter, loc, elemTy, expr); |
| Value shiftAmount = |
| LLVM::ConstantOp::create(rewriter, loc, intTy, elementSizeInBits); |
| Value highFull = LLVM::LShrOp::create(rewriter, loc, expr, shiftAmount); |
| Value high = LLVM::TruncOp::create(rewriter, loc, elemTy, highFull); |
| Type origElemTy = cast<LLVM::LLVMStructType>(origTy).getBody()[0]; |
| if (origElemTy != elemTy) { |
| low = LLVM::BitcastOp::create(rewriter, loc, origElemTy, low); |
| high = LLVM::BitcastOp::create(rewriter, loc, origElemTy, high); |
| } |
| Value undefStruct = LLVM::UndefOp::create(rewriter, loc, origTy); |
| Value structWithLow = LLVM::InsertValueOp::create( |
| rewriter, loc, origTy, undefStruct, low, ArrayRef<int64_t>{0}); |
| return LLVM::InsertValueOp::create(rewriter, loc, origTy, structWithLow, |
| high, ArrayRef<int64_t>{1}); |
| } |
| |
| uint64_t getAtomicSizeInBytes(Type originalTy, Type convertedTy, |
| ModuleOp module) const { |
| std::optional<TypeSizeAndAlignment> sizeAndAlignment = |
| acc::getTypeSizeAndAlignment(originalTy, module, &accSupport); |
| if (!sizeAndAlignment) |
| sizeAndAlignment = |
| acc::getTypeSizeAndAlignment(convertedTy, module, &accSupport); |
| assert(sizeAndAlignment && "atomic type size is not computable"); |
| return sizeAndAlignment->first.getFixedValue(); |
| } |
| |
| Type getAtomicType(Type originalTy, Type convertedTy, ModuleOp module) const { |
| if (convertedTy.isIntOrFloat()) |
| return convertedTy; |
| return IntegerType::get( |
| convertedTy.getContext(), |
| getAtomicSizeInBytes(originalTy, convertedTy, module) * kBitsInByte); |
| } |
| |
| Type getReferencedElementType(Value ref, ModuleOp module = nullptr) const { |
| auto ptr = dyn_cast<PointerLikeType>(ref.getType()); |
| if (!ptr) |
| llvm_unreachable("unexpected type"); |
| Type elementTy = ptr.getElementType(); |
| Type convertedTy = this->getTypeConverter()->convertType(elementTy); |
| if (module) |
| return getAtomicType(elementTy, convertedTy, module); |
| return convertedTy; |
| } |
| |
| /// Memrefs are converted to descriptors rather than bare pointers, so the |
| /// element pointer must be recomputed from the descriptor. |
| Value getAtomicPointer(Value originalRef, Value convertedPtr, Location loc, |
| ConversionPatternRewriter &rewriter) const { |
| auto memrefTy = dyn_cast<MemRefType>(originalRef.getType()); |
| if (!memrefTy) |
| return convertedPtr; |
| |
| // Extract aligned base pointer (index 1) and offset (index 2). |
| Value alignedPtr = |
| LLVM::ExtractValueOp::create(rewriter, loc, convertedPtr, 1); |
| Value offset = LLVM::ExtractValueOp::create(rewriter, loc, convertedPtr, 2); |
| Type elemPtrType = LLVM::LLVMPointerType::get(rewriter.getContext()); |
| return LLVM::GEPOp::create( |
| rewriter, loc, elemPtrType, |
| this->getTypeConverter()->convertType(memrefTy.getElementType()), |
| alignedPtr, offset); |
| } |
| |
| Block *constructCmpxchgLoop(Value ptr, Type type, Value expr, |
| ConversionPatternRewriter &rewriter) const; |
| |
| /// Emit a single atomicrmw for a `x = x <binop> expr` capture, or fail so |
| /// the caller falls back to the cmpxchg loop. |
| LogicalResult |
| tryEmitCaptureAtomicRMW(AtomicCaptureOp capture, AtomicUpdateOp update, |
| AtomicReadOp read, |
| ConversionPatternRewriter &rewriter) const; |
| |
| Value genUpdateCmpxchgLoop(AtomicUpdateOp update, |
| ConversionPatternRewriter &rewriter) const; |
| }; |
| |
| template <> |
| LogicalResult ACCAtomicOpConversion<AtomicReadOp>::matchAndRewrite( |
| AtomicReadOp read, OpAdaptor adaptor, |
| ConversionPatternRewriter &rewriter) const { |
| Location loc = read.getLoc(); |
| Value xRef = read.getX(); |
| Value xPtr = getAtomicPointer(xRef, adaptor.getX(), loc, rewriter); |
| ModuleOp mod = read->getParentOfType<ModuleOp>(); |
| Type xType = getReferencedElementType(xRef, mod); |
| |
| auto ordering = LLVM::AtomicOrdering::monotonic; |
| Value storeVal; |
| if (xType.isSignlessInteger()) { |
| Value zero = LLVM::ConstantOp::create(rewriter, loc, xType, 0); |
| storeVal = LLVM::AtomicRMWOp::create(rewriter, loc, LLVM::AtomicBinOp::_or, |
| xPtr, zero, ordering); |
| } else { |
| unsigned bitWidth = xType.getIntOrFloatBitWidth(); |
| Type intType = IntegerType::get(rewriter.getContext(), bitWidth); |
| Value zero = LLVM::ConstantOp::create(rewriter, loc, intType, 0); |
| Value intVal = LLVM::AtomicRMWOp::create( |
| rewriter, loc, LLVM::AtomicBinOp::_or, xPtr, zero, ordering); |
| storeVal = LLVM::BitcastOp::create(rewriter, loc, xType, intVal); |
| } |
| |
| Value vRef = read.getV(); |
| Value vPtr = getAtomicPointer(vRef, adaptor.getV(), loc, rewriter); |
| Type vType = getReferencedElementType(vRef, mod); |
| |
| if (xType != vType) { |
| // Convert `x` if the types do not match. |
| auto vPtrType = cast<PointerLikeType>(vRef.getType()); |
| storeVal = vPtrType.genCast(rewriter, loc, storeVal, vType); |
| if (!storeVal) { |
| return rewriter.notifyMatchFailure( |
| read, "failed to convert the loaded value to the destination type"); |
| } |
| } |
| auto storeOp = LLVM::StoreOp::create(rewriter, loc, storeVal, vPtr); |
| rewriter.replaceOp(read, storeOp); |
| return success(); |
| } |
| |
| template <> |
| LogicalResult ACCAtomicOpConversion<AtomicWriteOp>::matchAndRewrite( |
| AtomicWriteOp write, OpAdaptor adaptor, |
| ConversionPatternRewriter &rewriter) const { |
| Location loc = write.getLoc(); |
| Value expr = serializeExpr(adaptor.getExpr(), rewriter); |
| Value xRef = write.getX(); |
| Value xPtr = getAtomicPointer(xRef, adaptor.getX(), loc, rewriter); |
| ModuleOp mod = write->getParentOfType<ModuleOp>(); |
| Type xType = getReferencedElementType(xRef, mod); |
| |
| auto ordering = LLVM::AtomicOrdering::monotonic; |
| if (!xType.isSignlessInteger()) { |
| unsigned bitWidth = xType.getIntOrFloatBitWidth(); |
| Type intType = IntegerType::get(rewriter.getContext(), bitWidth); |
| expr = LLVM::BitcastOp::create(rewriter, loc, intType, expr); |
| } |
| LLVM::AtomicRMWOp::create(rewriter, loc, LLVM::AtomicBinOp::xchg, xPtr, expr, |
| ordering); |
| rewriter.eraseOp(write); |
| return success(); |
| } |
| |
| // Generate a cmpxchg loop based on the GenericAtomicRMWOpLowering algorithm. |
| template <typename AtomicOpTy> |
| Block *ACCAtomicOpConversion<AtomicOpTy>::constructCmpxchgLoop( |
| Value ptr, Type type, Value expr, |
| ConversionPatternRewriter &rewriter) const { |
| // Split the block into initial, loop, and ending parts. |
| Location loc = rewriter.getInsertionPoint()->getLoc(); |
| Block *initBlock = rewriter.getInsertionBlock(); |
| Block *loopBlock = |
| rewriter.splitBlock(initBlock, rewriter.getInsertionPoint()); |
| loopBlock->addArgument(type, loc); |
| Block *endBlock = |
| rewriter.splitBlock(loopBlock, rewriter.getInsertionPoint()); |
| |
| // Compute the loaded value and branch to the loop block. |
| rewriter.setInsertionPointToEnd(initBlock); |
| Value init = LLVM::LoadOp::create(rewriter, loc, type, ptr); |
| LLVM::BrOp::create(rewriter, loc, init, loopBlock); |
| |
| // Prepare the body of the loop block. |
| rewriter.setInsertionPointToStart(loopBlock); |
| |
| Value loopArgument = loopBlock->getArgument(0); |
| Value result = serializeExpr(rewriter.getRemappedValue(expr), rewriter); |
| ModuleOp mod = initBlock->getParent()->getParentOfType<ModuleOp>(); |
| Type convertedExprType = |
| this->getTypeConverter()->convertType(expr.getType()); |
| Type exprType = getAtomicType(expr.getType(), convertedExprType, mod); |
| |
| // Cast to an integer type. |
| if (!exprType.isSignlessInteger()) { |
| Type tmpType = IntegerType::get( |
| rewriter.getContext(), |
| getAtomicSizeInBytes(expr.getType(), convertedExprType, mod) * |
| kBitsInByte); |
| result = LLVM::BitcastOp::create(rewriter, loc, tmpType, result); |
| loopArgument = |
| LLVM::BitcastOp::create(rewriter, loc, tmpType, loopArgument); |
| } |
| |
| // Prepare the epilog of the loop block. |
| // Append the cmpxchg op to the end of the loop block. |
| auto successOrdering = LLVM::AtomicOrdering::acq_rel; |
| auto failureOrdering = LLVM::AtomicOrdering::monotonic; |
| auto cmpxchg = |
| LLVM::AtomicCmpXchgOp::create(rewriter, loc, ptr, loopArgument, result, |
| successOrdering, failureOrdering); |
| // Extract the %new_loaded and %ok values from the pair. |
| Value newLoaded = LLVM::ExtractValueOp::create(rewriter, loc, cmpxchg, 0); |
| Value ok = LLVM::ExtractValueOp::create(rewriter, loc, cmpxchg, 1); |
| |
| // Cast back to the original type. |
| if (!exprType.isSignlessInteger()) |
| newLoaded = LLVM::BitcastOp::create(rewriter, loc, exprType, newLoaded); |
| |
| // Conditionally branch to the end or back to the loop depending on %ok. |
| LLVM::CondBrOp::create(rewriter, loc, ok, endBlock, ArrayRef<Value>(), |
| loopBlock, newLoaded); |
| |
| return loopBlock; |
| } |
| |
| static Value skipUnrealizedConversionOp(Value v) { |
| if (auto convOp = |
| dyn_cast_or_null<UnrealizedConversionCastOp>(v.getDefiningOp())) |
| return skipUnrealizedConversionOp(convOp.getOperand(0)); |
| return v; |
| } |
| |
| /// Obtain the foremost value of addr to indicate the origin of the storage. |
| static Value getBaseStorage(Value addr, ConversionPatternRewriter &rewriter) { |
| Operation *op = skipUnrealizedConversionOp(addr).getDefiningOp(); |
| if (auto gepOp = dyn_cast_or_null<LLVM::GEPOp>(op)) { |
| addr = gepOp.getBase(); |
| op = addr.getDefiningOp(); |
| } |
| if (auto extractOp = dyn_cast_or_null<LLVM::ExtractValueOp>(op)) { |
| // addr is in a struct. Get the inserted value corresponding to the |
| // extraction. |
| ArrayRef extractingPosition = extractOp.getPosition(); |
| Value container = skipUnrealizedConversionOp(extractOp.getContainer()); |
| Value remappedContainer = rewriter.getRemappedValue(container); |
| op = remappedContainer ? remappedContainer.getDefiningOp() : nullptr; |
| bool inserted = false; |
| while (auto insertValueOp = dyn_cast_or_null<LLVM::InsertValueOp>(op)) { |
| ArrayRef insertingPosition = insertValueOp.getPosition(); |
| if (insertingPosition == extractingPosition) { |
| addr = insertValueOp.getValue(); |
| inserted = true; |
| break; |
| } |
| op = insertValueOp.getContainer().getDefiningOp(); |
| } |
| // The aggregate holding addr is not built by insertions, so the value it |
| // is converted from identifies the storage. |
| if (!inserted) |
| addr = container; |
| } |
| // op might be alloca already, or addr might be an argument. |
| return skipUnrealizedConversionOp(addr); |
| } |
| |
| /// Include flow dependency (v -> expr) in the generated loop of |
| /// `{ atomic.read, atomic.write/update }`. |
| static void moveDependency(Value vRef, Value vPtr, Value expr, |
| Value loopArgument, Operation &loopHead, |
| ConversionPatternRewriter &rewriter, |
| const ACCAtomicLoadAddressCallback &getLoadAddress) { |
| Value vStorage = getBaseStorage(vPtr, rewriter); |
| // A dependency may be reached before it is itself converted, in which case |
| // its address is still expressed in terms of `vRef`. |
| Value vStorageRef = skipUnrealizedConversionOp(vRef); |
| LLVM_DEBUG({ |
| llvm::dbgs() << "[acc-atomic] moveDependency\n"; |
| llvm::dbgs() << " vPtr = " << vPtr << "\n"; |
| llvm::dbgs() << " vStorage = " << vStorage << "\n"; |
| llvm::dbgs() << " vStorageRef = " << vStorageRef << "\n"; |
| llvm::dbgs() << " expr = " << expr << "\n"; |
| }); |
| llvm::DenseMap<Value, Value> remappedToOriginal; |
| std::set<std::pair<Operation *, std::set<Operation *>>> worklist; |
| std::set<Operation *> included; |
| Value mappedExpr = rewriter.getRemappedValue(expr); |
| remappedToOriginal[mappedExpr] = expr; |
| if (Operation *exprDef = mappedExpr.getDefiningOp()) |
| worklist.insert(std::pair{exprDef, std::set<Operation *>{}}); |
| |
| while (!worklist.empty()) { |
| auto [dep, post] = worklist.extract(worklist.begin()).value(); |
| if (!dep) |
| continue; |
| LLVM_DEBUG(llvm::dbgs() << " visit dep = " << *dep << "\n"); |
| if (vPtr.getDefiningOp() && |
| !vPtr.getDefiningOp()->getParentOp()->isAncestor(dep)) { |
| // Outside the parental region. No load found in this flow. |
| LLVM_DEBUG(llvm::dbgs() << " -> skipped: not in parental region\n"); |
| continue; |
| } |
| |
| Value addr; |
| if (auto load = dyn_cast<LLVM::LoadOp>(dep)) { |
| addr = load.getAddr(); |
| } else if (auto load = dyn_cast<memref::LoadOp>(dep)) { |
| addr = load.getMemref(); |
| } else if (getLoadAddress) { |
| addr = getLoadAddress(dep); |
| } |
| if (addr) { |
| Value baseStorage = getBaseStorage(addr, rewriter); |
| LLVM_DEBUG({ |
| llvm::dbgs() << " addr = " << addr << "\n"; |
| llvm::dbgs() << " baseStorage = " << baseStorage << "\n"; |
| llvm::dbgs() << " remapped = " |
| << rewriter.getRemappedValue(baseStorage) << "\n"; |
| llvm::dbgs() << " matchRef=" << (baseStorage == vStorageRef) |
| << " matchConv=" |
| << (rewriter.getRemappedValue(baseStorage) == vStorage) |
| << "\n"; |
| }); |
| if (baseStorage == vStorageRef || |
| rewriter.getRemappedValue(baseStorage) == vStorage) { |
| // Found the load of `v`. Include this flow in dependency. |
| Value load = dep->getResult(0); |
| auto replaceUses = [&](Operation *op) { |
| rewriter.modifyOpInPlace(op, [&] { |
| op->replaceUsesOfWith(remappedToOriginal[load], loopArgument); |
| }); |
| }; |
| for (Operation *p : post) |
| replaceUses(p); |
| replaceUses(&loopHead); |
| included.insert(post.begin(), post.end()); |
| } |
| continue; |
| } |
| post.insert(dep); |
| for (Value operand : dep->getOperands()) { |
| Value mappedOperand = rewriter.getRemappedValue(operand); |
| if (auto *d = mappedOperand.getDefiningOp()) { |
| if (dep == d) { |
| if (auto *op = operand.getDefiningOp()) { |
| remappedToOriginal[op->getResult(0)] = operand; |
| worklist.insert(std::pair{op, post}); |
| } |
| } else { |
| remappedToOriginal[mappedOperand] = operand; |
| worklist.insert(std::pair{d, post}); |
| } |
| } |
| } |
| } |
| |
| // Include dependency. |
| SmallVector<Operation *> includedInOrder; |
| vPtr.getParentRegion()->walk([&](Operation *op) { |
| if (included.find(op) != included.end()) |
| includedInOrder.push_back(op); |
| }); |
| for (Operation *d : includedInOrder) |
| d->moveBefore(&loopHead); |
| } |
| |
| /// Generate a cmpxchg loop for update and return a stored value. |
| template <typename AtomicOpTy> |
| Value ACCAtomicOpConversion<AtomicOpTy>::genUpdateCmpxchgLoop( |
| AtomicUpdateOp update, ConversionPatternRewriter &rewriter) const { |
| Location loc = update.getLoc(); |
| Value xRef = update.getX(); |
| Value xPtr = |
| getAtomicPointer(xRef, rewriter.getRemappedValue(xRef), loc, rewriter); |
| ModuleOp mod = update->getParentOfType<ModuleOp>(); |
| Type xType = getReferencedElementType(xRef, mod); |
| Type xTypeOrig = getReferencedElementType(xRef); |
| |
| Block &updateBlock = update.getRegion().front(); |
| Value updateArgument = updateBlock.getArgument(0); |
| Operation *terminator = updateBlock.getTerminator(); |
| Value expr = terminator->getOperand(0); |
| |
| Block *loopBlock = constructCmpxchgLoop(xPtr, xType, expr, rewriter); |
| Value loopArgument = loopBlock->getArgument(0); |
| Operation &loopHead = loopBlock->front(); |
| |
| rewriter.setInsertionPointToStart(loopBlock); |
| loopArgument = deserializeExpr(loopArgument, xTypeOrig, rewriter); |
| |
| // Move in and out flow dependency (x -> expr). Some computation might be |
| // outside atomic regions. |
| moveDependency(xRef, xPtr, expr, loopArgument, loopHead, rewriter, |
| getLoadAddress); |
| // Move out the residue. |
| rewriter.replaceAllUsesWith(cast<BlockArgument>(updateArgument), |
| {loopArgument}); |
| |
| updateBlock.walk([&](Operation *op) { |
| if (!op->hasTrait<OpTrait::IsTerminator>()) |
| rewriter.moveOpBefore(op, &loopHead); |
| }); |
| |
| if (auto cmpxchg = dyn_cast<LLVM::AtomicCmpXchgOp>(loopHead)) |
| return cmpxchg.getVal(); |
| if (auto bitcast = dyn_cast<LLVM::BitcastOp>(loopHead)) |
| // Handling a non-integer type. |
| return bitcast.getArg(); |
| if (auto extract = dyn_cast<LLVM::ExtractValueOp>(loopHead)) |
| // Handling a complex type. |
| return extract.getContainer(); |
| llvm_unreachable("invalid cmpxchg loop"); |
| } |
| |
| static std::optional<LLVM::AtomicBinOp> getAtomicBinOp(Operation *op, |
| bool updateIsLhs) { |
| return TypeSwitch<Operation *, std::optional<LLVM::AtomicBinOp>>(op) |
| .Case<arith::AddFOp>([](auto) { return LLVM::AtomicBinOp::fadd; }) |
| .Case<arith::AddIOp>([](auto) { return LLVM::AtomicBinOp::add; }) |
| .Case<arith::SubFOp>( |
| [updateIsLhs](auto) -> std::optional<LLVM::AtomicBinOp> { |
| // atomicrmw fsub is always `*ptr = *ptr - val`. |
| if (!updateIsLhs) |
| return std::nullopt; |
| return LLVM::AtomicBinOp::fsub; |
| }) |
| .Case<arith::SubIOp>( |
| [updateIsLhs](auto) -> std::optional<LLVM::AtomicBinOp> { |
| // atomicrmw sub is always `*ptr = *ptr - val`. |
| if (!updateIsLhs) |
| return std::nullopt; |
| return LLVM::AtomicBinOp::sub; |
| }) |
| .Case<arith::AndIOp>([](auto) { return LLVM::AtomicBinOp::_and; }) |
| .Case<arith::OrIOp>([](auto) { return LLVM::AtomicBinOp::_or; }) |
| .Case<arith::XOrIOp>([](auto) { return LLVM::AtomicBinOp::_xor; }) |
| .Case<arith::MaxSIOp>([](auto) { return LLVM::AtomicBinOp::max; }) |
| .Case<arith::MinSIOp>([](auto) { return LLVM::AtomicBinOp::min; }) |
| .Case<arith::MaxUIOp>([](auto) { return LLVM::AtomicBinOp::umax; }) |
| .Case<arith::MinUIOp>([](auto) { return LLVM::AtomicBinOp::umin; }) |
| .Case<arith::MaximumFOp>([](auto) { return LLVM::AtomicBinOp::fmaximum; }) |
| .Case<arith::MinimumFOp>([](auto) { return LLVM::AtomicBinOp::fminimum; }) |
| .Case<arith::MaxNumFOp>( |
| [](auto) { return LLVM::AtomicBinOp::fmaximumnum; }) |
| .Case<arith::MinNumFOp>( |
| [](auto) { return LLVM::AtomicBinOp::fminimumnum; }) |
| .Default([](Operation *) { return std::nullopt; }); |
| } |
| |
| /// Match an update region computing `x = x <binop> expr`, and return the |
| /// atomicrmw kind with the binary operation. Shared by the update and capture |
| /// conversions so the same region cannot take atomicrmw in one and the cmpxchg |
| /// loop in the other. |
| static std::optional<std::pair<LLVM::AtomicBinOp, Operation *>> |
| matchAtomicBinOpUpdate(AtomicUpdateOp update) { |
| Block &block = update.getRegion().front(); |
| Value arg = block.getArgument(0); |
| Operation *yield = block.getTerminator(); |
| if (!yield || yield->getNumOperands() != 1) |
| return std::nullopt; |
| Operation *binOp = yield->getOperand(0).getDefiningOp(); |
| if (!binOp || binOp->getBlock() != &block || binOp->getNumOperands() != 2 || |
| binOp->getNumResults() != 1) |
| return std::nullopt; |
| // The updated value has to feed the binop and nothing else. |
| if (!arg.hasOneUse() || arg.use_begin()->getOwner() != binOp) |
| return std::nullopt; |
| bool updateIsLhs = binOp->getOperand(0) == arg; |
| std::optional<LLVM::AtomicBinOp> kind = getAtomicBinOp(binOp, updateIsLhs); |
| if (!kind) |
| return std::nullopt; |
| return std::make_pair(*kind, binOp); |
| } |
| |
| /// Generate llvm.atomicrmw or an llvm.cmpxchg loop. |
| template <> |
| LogicalResult ACCAtomicOpConversion<AtomicUpdateOp>::matchAndRewrite( |
| AtomicUpdateOp update, OpAdaptor adaptor, |
| ConversionPatternRewriter &rewriter) const { |
| Block &updateBlock = update.getRegion().front(); |
| Value updateArgument = updateBlock.getArgument(0); |
| |
| // Collect operations that depend on the update argument. |
| std::set<Operation *> dependents; |
| SmallVector<Value> worklist; |
| worklist.push_back(updateArgument); |
| while (!worklist.empty()) { |
| Value value = worklist.back(); |
| worklist.pop_back(); |
| for (OpOperand &use : value.getUses()) { |
| Operation *useOp = use.getOwner(); |
| dependents.insert(useOp); |
| if (useOp->getNumResults() == 1) |
| worklist.push_back(useOp->getResult(0)); |
| } |
| } |
| |
| // Move independent operations out of the update block. |
| SmallVector<Operation *> independent; |
| for (Operation &op : updateBlock.getOperations()) { |
| if (dependents.find(&op) == dependents.end()) { |
| if (op.hasTrait<OpTrait::IsTerminator>()) |
| llvm_unreachable("invalid update operation"); |
| independent.push_back(&op); |
| } |
| } |
| for (Operation *op : independent) |
| rewriter.moveOpBefore(op, update); |
| |
| // Map arith op to atomicrmw kind. |
| // |
| // In the current version of LLVM, these are the equivalences for float |
| // min/max across the different instructions intrinsics and operations - |
| // confusingly they have slightly different and not very descriptive names. |
| // |
| // | MLIR | atomicrmw inst | llvm intrinsic | |
| // |------------+----------------+-------------------| |
| // | - | fmax | llvm.maxnum.* | |
| // | MaximumFOp | fmaximum | llvm.maximum.* | |
| // | MaxNumFOp | fmaximumnum | llvm.maximumnum.* | |
| // |
| // Sources: |
| // https://llvm.org/docs/LangRef.html#id236 |
| // https://llvm.org/docs/LangRef.html#floating-point-min-max-intrinsics-comparison |
| // https://mlir.llvm.org/docs/Dialects/ArithOps/#arithmaximumf-arithmaximumfop |
| // https://mlir.llvm.org/docs/Dialects/ArithOps/#arithmaxnumf-arithmaxnumfop |
| |
| // Select the kind and the val of atomicrmw. |
| std::optional<Value> val = std::nullopt; |
| std::optional<LLVM::AtomicBinOp> kind = std::nullopt; |
| |
| if (auto matched = matchAtomicBinOpUpdate(update)) { |
| Operation *binOp = matched->second; |
| bool updateIsLhs = binOp->getOperand(0) == updateArgument; |
| kind = matched->first; |
| val = binOp->getOperand(updateIsLhs ? 1 : 0); |
| } |
| |
| // Per-component atomicrmw info for complex type decomposition. |
| // Decomposed complex ops (complex.re/im + arith binop + complex.create) |
| // produce per-component binary ops that each need a separate atomicrmw. |
| struct ComponentAtomic { |
| LLVM::AtomicBinOp kind; |
| Value val; |
| int32_t fieldIdx; |
| }; |
| SmallVector<ComponentAtomic, 2> componentAtomics; |
| |
| if (!val || !kind) { |
| Type convertedArgTy = |
| this->getTypeConverter()->convertType(updateArgument.getType()); |
| if (auto structTy = dyn_cast<LLVM::LLVMStructType>(convertedArgTy)) { |
| if (structTy.getBody().size() == 2 && |
| structTy.getBody()[0] == structTy.getBody()[1] && |
| structTy.getBody()[0].isIntOrFloat() && |
| structTy.getBody()[0].getIntOrFloatBitWidth() > 32) { |
| for (Operation &op : updateBlock.getOperations()) { |
| if (op.hasTrait<OpTrait::IsTerminator>() || op.getNumOperands() < 2) |
| continue; |
| int32_t fieldIdx = -1; |
| Value externalVal = nullptr; |
| bool updateIsLhs = false; |
| for (unsigned i = 0; i < 2; ++i) { |
| Value operand = op.getOperand(i); |
| if (auto reOp = operand.getDefiningOp<complex::ReOp>()) { |
| if (reOp.getOperand() == updateArgument) { |
| fieldIdx = 0; |
| externalVal = op.getOperand(1 - i); |
| updateIsLhs = i == 0; |
| } |
| } else if (auto imOp = operand.getDefiningOp<complex::ImOp>()) { |
| if (imOp.getOperand() == updateArgument) { |
| fieldIdx = 1; |
| externalVal = op.getOperand(1 - i); |
| updateIsLhs = i == 0; |
| } |
| } |
| } |
| if (fieldIdx < 0) |
| continue; |
| auto componentKind = getAtomicBinOp(&op, updateIsLhs); |
| if (!componentKind) |
| continue; |
| componentAtomics.push_back({*componentKind, externalVal, fieldIdx}); |
| } |
| } |
| } |
| } |
| |
| Location loc = update.getLoc(); |
| Value xPtr = getAtomicPointer(update.getX(), adaptor.getX(), loc, rewriter); |
| |
| // Require distinct real/imag lanes; duplicate fieldIdx values must fall back |
| // to cmpxchg rather than emitting two atomicrmw ops on the same component. |
| bool hasDistinctComplexLanes = false; |
| if (componentAtomics.size() == 2) { |
| unsigned lanes = 0; |
| for (const ComponentAtomic &ca : componentAtomics) { |
| if (ca.fieldIdx == 0 || ca.fieldIdx == 1) |
| lanes |= 1u << ca.fieldIdx; |
| } |
| hasDistinctComplexLanes = lanes == 0b11; |
| } |
| |
| if (val && kind) { |
| auto ordering = LLVM::AtomicOrdering::monotonic; |
| LLVM::AtomicRMWOp::create(rewriter, loc, *kind, xPtr, |
| rewriter.getRemappedValue(*val), ordering); |
| } else if (hasDistinctComplexLanes) { |
| auto structTy = cast<LLVM::LLVMStructType>( |
| this->getTypeConverter()->convertType(updateArgument.getType())); |
| Type ptrType = LLVM::LLVMPointerType::get(rewriter.getContext()); |
| auto ordering = LLVM::AtomicOrdering::monotonic; |
| for (ComponentAtomic &ca : componentAtomics) { |
| Value elemPtr = |
| LLVM::GEPOp::create(rewriter, loc, ptrType, structTy, xPtr, |
| ArrayRef<LLVM::GEPArg>{0, ca.fieldIdx}); |
| LLVM::AtomicRMWOp::create(rewriter, loc, ca.kind, elemPtr, |
| rewriter.getRemappedValue(ca.val), ordering); |
| } |
| } else { |
| // Fallback to the llvm.cmpxchg loop generation. |
| genUpdateCmpxchgLoop(update, rewriter); |
| } |
| rewriter.eraseOp(update); |
| return success(); |
| } |
| |
| /// The cmpxchg loop re-points reads of `v`/`x` inside `expr` at the atomically |
| /// loaded value (see moveDependency). An atomicrmw cannot do that, so any |
| /// `expr` that loads from memory keeps the loop: the load may be the captured |
| /// value itself, and the address is not reliably comparable here. |
| static bool exprReadsMemory(Value expr) { |
| SmallVector<Value> worklist{expr}; |
| llvm::DenseSet<Operation *> seen; |
| while (!worklist.empty()) { |
| Operation *def = worklist.pop_back_val().getDefiningOp(); |
| if (!def || !seen.insert(def).second) |
| continue; |
| if (!isMemoryEffectFree(def)) |
| return true; |
| worklist.append(def->getOperands().begin(), def->getOperands().end()); |
| } |
| return false; |
| } |
| |
| /// Emit a single atomicrmw for a `x = x <binop> expr` capture. |
| template <typename AtomicOpTy> |
| LogicalResult ACCAtomicOpConversion<AtomicOpTy>::tryEmitCaptureAtomicRMW( |
| AtomicCaptureOp capture, AtomicUpdateOp update, AtomicReadOp read, |
| ConversionPatternRewriter &rewriter) const { |
| if (read.getX() != update.getX()) |
| return failure(); |
| auto matched = matchAtomicBinOpUpdate(update); |
| if (!matched) |
| return failure(); |
| auto [kind, binOp] = *matched; |
| |
| Value arg = update.getRegion().front().getArgument(0); |
| bool updateIsLhs = binOp->getOperand(0) == arg; |
| Value expr = binOp->getOperand(updateIsLhs ? 1 : 0); |
| |
| // Keep serialized and aggregate types on the cmpxchg path. |
| Type argTy = arg.getType(); |
| if (!argTy.isIntOrFloat() || |
| this->getTypeConverter()->convertType(argTy) != argTy) |
| return failure(); |
| // The operand must already be available, and must not be the captured value. |
| Operation *exprDef = expr.getDefiningOp(); |
| if (exprDef && exprDef->getBlock() == &update.getRegion().front()) |
| return failure(); |
| if (exprReadsMemory(expr)) |
| return failure(); |
| |
| Location loc = capture.getLoc(); |
| Value xRef = update.getX(); |
| Value vRef = read.getV(); |
| Value xPtr = |
| getAtomicPointer(xRef, rewriter.getRemappedValue(xRef), loc, rewriter); |
| Value vPtr = |
| getAtomicPointer(vRef, rewriter.getRemappedValue(vRef), loc, rewriter); |
| |
| rewriter.setInsertionPoint(capture); |
| auto rmw = LLVM::AtomicRMWOp::create(rewriter, loc, kind, xPtr, |
| rewriter.getRemappedValue(expr), |
| LLVM::AtomicOrdering::monotonic); |
| // atomicrmw yields the old value; `{update, read}` captures the new one. |
| Value captured = rmw.getRes(); |
| if (capture.getFirstOp() == update.getOperation()) { |
| rewriter.moveOpAfter(binOp, rmw); |
| binOp->replaceUsesOfWith(arg, captured); |
| captured = binOp->getResult(0); |
| } |
| rewriter.replaceOpWithNewOp<LLVM::StoreOp>(capture, captured, vPtr); |
| return success(); |
| } |
| |
| /// Generate an llvm.cmpxchg loop. |
| template <> |
| LogicalResult ACCAtomicOpConversion<AtomicCaptureOp>::matchAndRewrite( |
| AtomicCaptureOp capture, OpAdaptor /*adaptor*/, |
| ConversionPatternRewriter &rewriter) const { |
| Operation *firstOp = capture.getFirstOp(); |
| Operation *secondOp = capture.getSecondOp(); |
| Value vPtr = nullptr; |
| Value storeVal = nullptr; |
| |
| // A single `x = x <binop> expr` capture becomes one atomicrmw. The cmpxchg |
| // loop below serializes retries and collapses under contention. |
| if (AtomicUpdateOp update = capture.getAtomicUpdateOp()) |
| if (AtomicReadOp read = capture.getAtomicReadOp()) |
| if (succeeded(tryEmitCaptureAtomicRMW(capture, update, read, rewriter))) |
| return success(); |
| |
| if (auto firstReadStmt = dyn_cast<AtomicReadOp>(firstOp)) { |
| Location loc = capture.getLoc(); |
| Value xRef = firstReadStmt.getX(); |
| Value xPtr = |
| getAtomicPointer(xRef, rewriter.getRemappedValue(xRef), loc, rewriter); |
| ModuleOp mod = capture->getParentOfType<ModuleOp>(); |
| Type xType = getReferencedElementType(xRef, mod); |
| Type xTypeOrig = getReferencedElementType(xRef); |
| Value vRef = firstReadStmt.getV(); |
| vPtr = |
| getAtomicPointer(vRef, rewriter.getRemappedValue(vRef), loc, rewriter); |
| |
| Value expr = nullptr; |
| if (auto secondWriteStmt = dyn_cast<AtomicWriteOp>(secondOp)) { |
| // 1. `{ atomic.read, atomic.write }` pattern |
| expr = secondWriteStmt.getExpr(); |
| } else if (auto secondUpdateStmt = dyn_cast<AtomicUpdateOp>(secondOp)) { |
| // 2. `{ atomic.read, atomic.update }` pattern |
| Block &updateBlock = secondUpdateStmt.getRegion().front(); |
| Operation *terminator = updateBlock.getTerminator(); |
| expr = terminator->getOperand(0); |
| } |
| |
| Block *loopBlock = constructCmpxchgLoop(xPtr, xType, expr, rewriter); |
| Value loopArgument = loopBlock->getArgument(0); |
| Operation &loopHead = loopBlock->front(); |
| auto condBr = cast<LLVM::CondBrOp>(loopBlock->back()); |
| storeVal = condBr.getFalseDestOperands()[0]; |
| |
| rewriter.setInsertionPointToStart(loopBlock); |
| loopArgument = deserializeExpr(loopArgument, xTypeOrig, rewriter); |
| |
| // Include flow dependency (v -> expr). |
| if (auto secondUpdateStmt = dyn_cast<AtomicUpdateOp>(secondOp)) { |
| Block &updateBlock = secondUpdateStmt.getRegion().front(); |
| Value updateArgument = updateBlock.getArgument(0); |
| updateArgument.replaceAllUsesWith(loopArgument); |
| updateBlock.walk([&](Operation *op) { |
| if (!op->hasTrait<OpTrait::IsTerminator>()) |
| rewriter.moveOpBefore(op, &loopHead); |
| }); |
| } |
| moveDependency(vRef, vPtr, expr, loopArgument, loopHead, rewriter, |
| getLoadAddress); |
| } else if (auto firstUpdateStmt = dyn_cast<AtomicUpdateOp>(firstOp)) { |
| if (auto secondReadStmt = dyn_cast<AtomicReadOp>(secondOp)) { |
| // 3. `{ atomic.update, atomic.read }` pattern |
| storeVal = genUpdateCmpxchgLoop(firstUpdateStmt, rewriter); |
| |
| Value vRef = secondReadStmt.getV(); |
| vPtr = getAtomicPointer(vRef, rewriter.getRemappedValue(vRef), |
| capture.getLoc(), rewriter); |
| } |
| } |
| // Generate `v = x`. |
| rewriter.setInsertionPoint(capture); |
| rewriter.replaceOpWithNewOp<LLVM::StoreOp>(capture, storeVal, vPtr); |
| return success(); |
| } |
| |
| } // namespace |
| |
| namespace mlir { |
| |
| void configureACCAtomicConversionLegality(ConversionTarget &target) { |
| target.addIllegalOp<AtomicReadOp, AtomicWriteOp, AtomicUpdateOp, |
| AtomicCaptureOp>(); |
| } |
| |
| void populateACCAtomicPatterns(const LLVMTypeConverter &converter, |
| RewritePatternSet &patterns, |
| acc::OpenACCSupport &accSupport, |
| ACCAtomicLoadAddressCallback getLoadAddress) { |
| patterns.add<ACCAtomicOpConversion<AtomicReadOp>, |
| ACCAtomicOpConversion<AtomicWriteOp>, |
| ACCAtomicOpConversion<AtomicUpdateOp>, |
| ACCAtomicOpConversion<AtomicCaptureOp>>(converter, accSupport, |
| getLoadAddress); |
| } |
| |
| } // namespace mlir |
