AtomicOps.cpp

Go to the documentation of this file.

//===- AtomicOps.cpp - MLIR SPIR-V Atomic Ops  ----------------------------===//
//
// 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
//
//===----------------------------------------------------------------------===//
//
// Defines the atomic operations in the SPIR-V dialect.
//
//===----------------------------------------------------------------------===//
 
#include "mlir/Dialect/SPIRV/IR/SPIRVOps.h"
 
#include "SPIRVOpUtils.h"
#include "SPIRVParsingUtils.h"
 
using namespace mlir::spirv::AttrNames;
 
namespace mlir::spirv {
 
// Parses an atomic update op. If the update op does not take a value (like
// AtomicIIncrement) `hasValue` must be false.
static ParseResult parseAtomicUpdateOp(OpAsmParser &parser,
                                       OperationState &state, bool hasValue) {
  spirv::Scope scope;
  spirv::MemorySemantics memoryScope;
  SmallVector<OpAsmParser::UnresolvedOperand, 2> operandInfo;
  OpAsmParser::UnresolvedOperand ptrInfo, valueInfo;
  Type type;
  SMLoc loc;
  if (parseEnumStrAttr<spirv::ScopeAttr>(scope, parser, state,
                                         kMemoryScopeAttrName) ||
      parseEnumStrAttr<spirv::MemorySemanticsAttr>(memoryScope, parser, state,
                                                   kSemanticsAttrName) ||
      parser.parseOperandList(operandInfo, (hasValue ? 2 : 1)) ||
      parser.getCurrentLocation(&loc) || parser.parseColonType(type))
    return failure();
 
  auto ptrType = llvm::dyn_cast<spirv::PointerType>(type);
  if (!ptrType)
    return parser.emitError(loc, "expected pointer type");
 
  SmallVector<Type, 2> operandTypes;
  operandTypes.push_back(ptrType);
  if (hasValue)
    operandTypes.push_back(ptrType.getPointeeType());
  if (parser.resolveOperands(operandInfo, operandTypes, parser.getNameLoc(),
                             state.operands))
    return failure();
  return parser.addTypeToList(ptrType.getPointeeType(), state.types);
}
 
// Prints an atomic update op.
static void printAtomicUpdateOp(Operation *op, OpAsmPrinter &printer) {
  printer << " \"";
  auto scopeAttr = op->getAttrOfType<spirv::ScopeAttr>(kMemoryScopeAttrName);
  printer << spirv::stringifyScope(scopeAttr.getValue()) << "\" \"";
  auto memorySemanticsAttr =
      op->getAttrOfType<spirv::MemorySemanticsAttr>(kSemanticsAttrName);
  printer << spirv::stringifyMemorySemantics(memorySemanticsAttr.getValue())
          << "\" " << op->getOperands() << " : " << op->getOperand(0).getType();
}
 
template <typename T>
static StringRef stringifyTypeName();
 
template <>
StringRef stringifyTypeName<IntegerType>() {
  return "integer";
}
 
template <>
StringRef stringifyTypeName<FloatType>() {
  return "float";
}
 
// Verifies an atomic update op.
template <typename ExpectedElementType>
static LogicalResult verifyAtomicUpdateOp(Operation *op) {
  auto ptrType = llvm::cast<spirv::PointerType>(op->getOperand(0).getType());
  auto elementType = ptrType.getPointeeType();
  if (!llvm::isa<ExpectedElementType>(elementType))
    return op->emitOpError() << "pointer operand must point to an "
                             << stringifyTypeName<ExpectedElementType>()
                             << " value, found " << elementType;
 
  if (op->getNumOperands() > 1) {
    auto valueType = op->getOperand(1).getType();
    if (valueType != elementType)
      return op->emitOpError("expected value to have the same type as the "
                             "pointer operand's pointee type ")
             << elementType << ", but found " << valueType;
  }
  auto memorySemantics =
      op->getAttrOfType<spirv::MemorySemanticsAttr>(kSemanticsAttrName)
          .getValue();
  if (failed(verifyMemorySemantics(op, memorySemantics))) {
    return failure();
  }
  return success();
}
 
template <typename T>
static void printAtomicCompareExchangeImpl(T atomOp, OpAsmPrinter &printer) {
  printer << " \"" << stringifyScope(atomOp.getMemoryScope()) << "\" \""
          << stringifyMemorySemantics(atomOp.getEqualSemantics()) << "\" \""
          << stringifyMemorySemantics(atomOp.getUnequalSemantics()) << "\" "
          << atomOp.getOperands() << " : " << atomOp.getPointer().getType();
}
 
static ParseResult parseAtomicCompareExchangeImpl(OpAsmParser &parser,
                                                  OperationState &state) {
  spirv::Scope memoryScope;
  spirv::MemorySemantics equalSemantics, unequalSemantics;
  SmallVector<OpAsmParser::UnresolvedOperand, 3> operandInfo;
  Type type;
  if (parseEnumStrAttr<spirv::ScopeAttr>(memoryScope, parser, state,
                                         kMemoryScopeAttrName) ||
      parseEnumStrAttr<spirv::MemorySemanticsAttr>(
          equalSemantics, parser, state, kEqualSemanticsAttrName) ||
      parseEnumStrAttr<spirv::MemorySemanticsAttr>(
          unequalSemantics, parser, state, kUnequalSemanticsAttrName) ||
      parser.parseOperandList(operandInfo, 3))
    return failure();
 
  auto loc = parser.getCurrentLocation();
  if (parser.parseColonType(type))
    return failure();
 
  auto ptrType = llvm::dyn_cast<spirv::PointerType>(type);
  if (!ptrType)
    return parser.emitError(loc, "expected pointer type");
 
  if (parser.resolveOperands(
          operandInfo,
          {ptrType, ptrType.getPointeeType(), ptrType.getPointeeType()},
          parser.getNameLoc(), state.operands))
    return failure();
 
  return parser.addTypeToList(ptrType.getPointeeType(), state.types);
}
 
template <typename T>
static LogicalResult verifyAtomicCompareExchangeImpl(T atomOp) {
  // According to the spec:
  // "The type of Value must be the same as Result Type. The type of the value
  // pointed to by Pointer must be the same as Result Type. This type must also
  // match the type of Comparator."
  if (atomOp.getType() != atomOp.getValue().getType())
    return atomOp.emitOpError("value operand must have the same type as the op "
                              "result, but found ")
           << atomOp.getValue().getType() << " vs " << atomOp.getType();
 
  if (atomOp.getType() != atomOp.getComparator().getType())
    return atomOp.emitOpError(
               "comparator operand must have the same type as the op "
               "result, but found ")
           << atomOp.getComparator().getType() << " vs " << atomOp.getType();
 
  Type pointeeType =
      llvm::cast<spirv::PointerType>(atomOp.getPointer().getType())
          .getPointeeType();
  if (atomOp.getType() != pointeeType)
    return atomOp.emitOpError(
               "pointer operand's pointee type must have the same "
               "as the op result type, but found ")
           << pointeeType << " vs " << atomOp.getType();
 
  // TODO: Unequal cannot be set to Release or Acquire and Release.
  // In addition, Unequal cannot be set to a stronger memory-order then Equal.
 
  return success();
}
 
//===----------------------------------------------------------------------===//
// spirv.AtomicAndOp
//===----------------------------------------------------------------------===//
 
LogicalResult AtomicAndOp::verify() {
  return verifyAtomicUpdateOp<IntegerType>(getOperation());
}
 
ParseResult AtomicAndOp::parse(OpAsmParser &parser, OperationState &result) {
  return parseAtomicUpdateOp(parser, result, true);
}
 
void AtomicAndOp::print(OpAsmPrinter &p) { printAtomicUpdateOp(*this, p); }
 
//===----------------------------------------------------------------------===//
// spirv.AtomicCompareExchangeOp
//===----------------------------------------------------------------------===//
 
LogicalResult AtomicCompareExchangeOp::verify() {
  return verifyAtomicCompareExchangeImpl(*this);
}
 
ParseResult AtomicCompareExchangeOp::parse(OpAsmParser &parser,
                                           OperationState &result) {
  return parseAtomicCompareExchangeImpl(parser, result);
}
 
void AtomicCompareExchangeOp::print(OpAsmPrinter &p) {
  printAtomicCompareExchangeImpl(*this, p);
}
 
//===----------------------------------------------------------------------===//
// spirv.AtomicCompareExchangeWeakOp
//===----------------------------------------------------------------------===//
 
LogicalResult AtomicCompareExchangeWeakOp::verify() {
  return verifyAtomicCompareExchangeImpl(*this);
}
 
ParseResult AtomicCompareExchangeWeakOp::parse(OpAsmParser &parser,
                                               OperationState &result) {
  return parseAtomicCompareExchangeImpl(parser, result);
}
 
void AtomicCompareExchangeWeakOp::print(OpAsmPrinter &p) {
  printAtomicCompareExchangeImpl(*this, p);
}
 
//===----------------------------------------------------------------------===//
// spirv.AtomicExchange
//===----------------------------------------------------------------------===//
 
void AtomicExchangeOp::print(OpAsmPrinter &printer) {
  printer << " \"" << stringifyScope(getMemoryScope()) << "\" \""
          << stringifyMemorySemantics(getSemantics()) << "\" " << getOperands()
          << " : " << getPointer().getType();
}
 
ParseResult AtomicExchangeOp::parse(OpAsmParser &parser,
                                    OperationState &result) {
  spirv::Scope memoryScope;
  spirv::MemorySemantics semantics;
  SmallVector<OpAsmParser::UnresolvedOperand, 2> operandInfo;
  Type type;
  if (parseEnumStrAttr<spirv::ScopeAttr>(memoryScope, parser, result,
                                         kMemoryScopeAttrName) ||
      parseEnumStrAttr<spirv::MemorySemanticsAttr>(semantics, parser, result,
                                                   kSemanticsAttrName) ||
      parser.parseOperandList(operandInfo, 2))
    return failure();
 
  auto loc = parser.getCurrentLocation();
  if (parser.parseColonType(type))
    return failure();
 
  auto ptrType = llvm::dyn_cast<spirv::PointerType>(type);
  if (!ptrType)
    return parser.emitError(loc, "expected pointer type");
 
  if (parser.resolveOperands(operandInfo, {ptrType, ptrType.getPointeeType()},
                             parser.getNameLoc(), result.operands))
    return failure();
 
  return parser.addTypeToList(ptrType.getPointeeType(), result.types);
}
 
LogicalResult AtomicExchangeOp::verify() {
  if (getType() != getValue().getType())
    return emitOpError("value operand must have the same type as the op "
                       "result, but found ")
           << getValue().getType() << " vs " << getType();
 
  Type pointeeType =
      llvm::cast<spirv::PointerType>(getPointer().getType()).getPointeeType();
  if (getType() != pointeeType)
    return emitOpError("pointer operand's pointee type must have the same "
                       "as the op result type, but found ")
           << pointeeType << " vs " << getType();
 
  return success();
}
 
//===----------------------------------------------------------------------===//
// spirv.AtomicIAddOp
//===----------------------------------------------------------------------===//
 
LogicalResult AtomicIAddOp::verify() {
  return verifyAtomicUpdateOp<IntegerType>(getOperation());
}
 
ParseResult AtomicIAddOp::parse(OpAsmParser &parser, OperationState &result) {
  return parseAtomicUpdateOp(parser, result, true);
}
 
void AtomicIAddOp::print(OpAsmPrinter &p) { printAtomicUpdateOp(*this, p); }
 
//===----------------------------------------------------------------------===//
// spirv.EXT.AtomicFAddOp
//===----------------------------------------------------------------------===//
 
LogicalResult EXTAtomicFAddOp::verify() {
  return verifyAtomicUpdateOp<FloatType>(getOperation());
}
 
ParseResult EXTAtomicFAddOp::parse(OpAsmParser &parser,
                                   OperationState &result) {
  return parseAtomicUpdateOp(parser, result, true);
}
 
void spirv::EXTAtomicFAddOp::print(OpAsmPrinter &p) {
  printAtomicUpdateOp(*this, p);
}
 
//===----------------------------------------------------------------------===//
// spirv.AtomicIDecrementOp
//===----------------------------------------------------------------------===//
 
LogicalResult AtomicIDecrementOp::verify() {
  return verifyAtomicUpdateOp<IntegerType>(getOperation());
}
 
ParseResult AtomicIDecrementOp::parse(OpAsmParser &parser,
                                      OperationState &result) {
  return parseAtomicUpdateOp(parser, result, false);
}
 
void AtomicIDecrementOp::print(OpAsmPrinter &p) {
  printAtomicUpdateOp(*this, p);
}
 
//===----------------------------------------------------------------------===//
// spirv.AtomicIIncrementOp
//===----------------------------------------------------------------------===//
 
LogicalResult AtomicIIncrementOp::verify() {
  return verifyAtomicUpdateOp<IntegerType>(getOperation());
}
 
ParseResult AtomicIIncrementOp::parse(OpAsmParser &parser,
                                      OperationState &result) {
  return parseAtomicUpdateOp(parser, result, false);
}
 
void AtomicIIncrementOp::print(OpAsmPrinter &p) {
  printAtomicUpdateOp(*this, p);
}
 
//===----------------------------------------------------------------------===//
// spirv.AtomicISubOp
//===----------------------------------------------------------------------===//
 
LogicalResult AtomicISubOp::verify() {
  return verifyAtomicUpdateOp<IntegerType>(getOperation());
}
 
ParseResult AtomicISubOp::parse(OpAsmParser &parser, OperationState &result) {
  return parseAtomicUpdateOp(parser, result, true);
}
 
void AtomicISubOp::print(OpAsmPrinter &p) { printAtomicUpdateOp(*this, p); }
 
//===----------------------------------------------------------------------===//
// spirv.AtomicOrOp
//===----------------------------------------------------------------------===//
 
LogicalResult AtomicOrOp::verify() {
  return verifyAtomicUpdateOp<IntegerType>(getOperation());
}
 
ParseResult AtomicOrOp::parse(OpAsmParser &parser, OperationState &result) {
  return parseAtomicUpdateOp(parser, result, true);
}
 
void AtomicOrOp::print(OpAsmPrinter &p) { printAtomicUpdateOp(*this, p); }
 
//===----------------------------------------------------------------------===//
// spirv.AtomicSMaxOp
//===----------------------------------------------------------------------===//
 
LogicalResult AtomicSMaxOp::verify() {
  return verifyAtomicUpdateOp<IntegerType>(getOperation());
}
 
ParseResult AtomicSMaxOp::parse(OpAsmParser &parser, OperationState &result) {
  return parseAtomicUpdateOp(parser, result, true);
}
 
void AtomicSMaxOp::print(OpAsmPrinter &p) { printAtomicUpdateOp(*this, p); }
 
//===----------------------------------------------------------------------===//
// spirv.AtomicSMinOp
//===----------------------------------------------------------------------===//
 
LogicalResult AtomicSMinOp::verify() {
  return verifyAtomicUpdateOp<IntegerType>(getOperation());
}
 
ParseResult AtomicSMinOp::parse(OpAsmParser &parser, OperationState &result) {
  return parseAtomicUpdateOp(parser, result, true);
}
 
void AtomicSMinOp::print(OpAsmPrinter &p) { printAtomicUpdateOp(*this, p); }
 
//===----------------------------------------------------------------------===//
// spirv.AtomicUMaxOp
//===----------------------------------------------------------------------===//
 
LogicalResult AtomicUMaxOp::verify() {
  return verifyAtomicUpdateOp<IntegerType>(getOperation());
}
 
ParseResult AtomicUMaxOp::parse(OpAsmParser &parser, OperationState &result) {
  return parseAtomicUpdateOp(parser, result, true);
}
 
void AtomicUMaxOp::print(OpAsmPrinter &p) { printAtomicUpdateOp(*this, p); }
 
//===----------------------------------------------------------------------===//
// spirv.AtomicUMinOp
//===----------------------------------------------------------------------===//
 
LogicalResult AtomicUMinOp::verify() {
  return verifyAtomicUpdateOp<IntegerType>(getOperation());
}
 
ParseResult AtomicUMinOp::parse(OpAsmParser &parser, OperationState &result) {
  return parseAtomicUpdateOp(parser, result, true);
}
 
void AtomicUMinOp::print(OpAsmPrinter &p) { printAtomicUpdateOp(*this, p); }
 
//===----------------------------------------------------------------------===//
// spirv.AtomicXorOp
//===----------------------------------------------------------------------===//
 
LogicalResult AtomicXorOp::verify() {
  return verifyAtomicUpdateOp<IntegerType>(getOperation());
}
 
ParseResult AtomicXorOp::parse(OpAsmParser &parser, OperationState &result) {
  return parseAtomicUpdateOp(parser, result, true);
}
 
void AtomicXorOp::print(OpAsmPrinter &p) { printAtomicUpdateOp(*this, p); }
 
} // namespace mlir::spirv