XeVMDialect.cpp

Go to the documentation of this file.

//===-- XeVMDialect.cpp - XeVM dialect registration -------------*- C++ -*-===//
//
// This file is licensed 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/LLVMIR/XeVMDialect.h"
#include "mlir/Dialect/GPU/IR/CompilationInterfaces.h"
#include "mlir/Dialect/Utils/StaticValueUtils.h"
#include "mlir/IR/TypeUtilities.h"
#include "llvm/ADT/APFloat.h"
#include "llvm/ADT/SmallSet.h"
#include "llvm/Support/FileSystem.h"
#include "llvm/Support/MathExtras.h"
 
using namespace mlir;
using namespace mlir::xevm;
 
#include "mlir/Dialect/LLVMIR/XeVMOpsDialect.cpp.inc"
#include "mlir/Dialect/LLVMIR/XeVMOpsEnums.cpp.inc"
 
namespace {
static constexpr uint32_t subgroupSize = 16;
 
template <typename Op>
LogicalResult verifyMatrixInput(Op op) {
  static_assert(llvm::is_one_of<Op, BlockLoad2dOp, BlockStore2dOp,
                                BlockPrefetch2dOp>::value,
                "Unexpected template parameter");
 
  std::optional<int64_t> width = getConstantIntValue(op.getBaseWidth());
  std::optional<int64_t> pitch = getConstantIntValue(op.getBasePitch());
  if (pitch && width && *pitch < *width)
    return op->emitOpError(
        "4th operand (base pitch) should be >= 2nd operand (base width)");
 
  uint32_t elemSize = op.getElemSizeInBits();
  if (elemSize < 8 || !llvm::isPowerOf2_32(elemSize) || elemSize > 32)
    return op->emitOpError("expecting 'elem_size_in_bits' to be 8, 16, or 32");
 
  uint32_t tileHeight = op.getTileHeight();
  if (tileHeight > 32 || !llvm::isPowerOf2_32(tileHeight))
    return op->emitOpError("expecting tile_height to be 1, 2, 4, 8, 16, or 32");
 
  uint32_t vBlocks = op.getVBlocks();
  if (vBlocks > 8 || !llvm::isPowerOf2_32(vBlocks))
    return op->emitOpError("expecting v_blocks to be 1, 2, 4, or 8");
 
  return success();
}
 
LogicalResult verify2DBlockLoadRestriction(BlockLoad2dOp op) {
  VectorType resTy = op.getRes().getType();
  if (!resTy.getElementType().isIntOrFloat())
    return op.emitOpError()
           << "expecting result element type to be int or float";
  unsigned resElemTySize = resTy.getElementType().getIntOrFloatBitWidth();
  unsigned resSize = resTy.getNumElements() * resElemTySize;
  unsigned expectedSize = op.getElemSizeInBits() * op.getTileHeight() *
                          op.getTileWidth() * op.getVBlocks() / subgroupSize;
  if (resSize != expectedSize)
    return op.emitOpError() << "result size of " << resSize
                            << " bits does not match the expected size of "
                            << expectedSize << " bits";
 
  if (op.getTranspose() && op.getPackRegister())
    return op.emitOpError("transpose and pack_register are mutually exclusive");
 
  if (!op.getTranspose() && !op.getPackRegister()) {
    uint32_t tileHeight = op.getTileHeight();
    if (tileHeight < 1 || tileHeight > 32)
      return op.emitOpError("expecting tile_height to be between 1 and 32");
 
    uint32_t tileWidth = op.getTileWidth();
    uint32_t vBlocks = op.getVBlocks();
    switch (op.getElemSizeInBits()) {
    case 8:
      if (tileWidth < 4 || tileWidth > 64)
        return op.emitOpError("expecting tile_width to be between 4 and 64");
      if (vBlocks != 1 && vBlocks != 2 && vBlocks != 4)
        return op.emitOpError("expecting v_blocks to be 1, 2, or 4");
      if (tileWidth * vBlocks > 64)
        return op.emitOpError(
            "tile_width * v_blocks should be less than or equal "
            "to 64 for 8 bit elements");
      break;
    case 16:
      if (tileWidth < 2 || tileWidth > 32)
        return op.emitOpError("expecting tile_width to be between 2 and 32");
      if (vBlocks != 1 && vBlocks != 2 && vBlocks != 4)
        return op.emitOpError("expecting v_blocks to be 1, 2, or 4");
      if (tileWidth * vBlocks > 32)
        return op.emitOpError(
            "tile_width * v_blocks should be less than or equal "
            "to 32 for 16 bit elements");
      break;
    case 32:
      if (tileWidth < 1 || tileWidth > 16)
        return op.emitOpError("expecting tile_width to be between 1 and 16");
      if (vBlocks != 1 && vBlocks != 2)
        return op.emitOpError("expecting v_blocks to be 1 or 2");
      if (tileWidth * vBlocks > 16)
        return op.emitOpError(
            "tile_width * v_blocks should be less than or equal "
            "to 16 for 32 bit elements");
      break;
    case 64:
      if (tileWidth < 1 || tileWidth > 8)
        return op.emitOpError("expecting tile_width to be between 1 and 8");
      if (vBlocks != 1)
        return op.emitOpError("expecting v_blocks to be 1");
      break;
    default:
      return op.emitOpError(
          "expecting elem_size_in_bits to be 8, 16, 32, or 64");
    }
 
    return success();
  }
 
  if (op.getTranspose()) {
    assert(!op.getPackRegister() && "Expecting pack_register should be false");
 
    uint32_t vBlocks = op.getVBlocks();
    if (vBlocks != 1)
      return op.emitOpError("expecting v_blocks to be 1");
 
    uint32_t tileHeight = op.getTileHeight();
    uint32_t tileWidth = op.getTileWidth();
    switch (op.getElemSizeInBits()) {
    case 32:
      if (tileHeight < 1 || tileHeight > 32)
        return op.emitOpError("expecting tile_height to be between 1 and 32");
      if (tileWidth < 1 || tileWidth > 8)
        return op.emitOpError("expecting tile_width to be between 1 and 8");
      break;
    case 64:
      if (tileHeight != 8)
        return op.emitOpError(
            "expecting tile_height to be 8 for 64 bit elements");
      if (tileWidth != 1 && tileWidth != 2 && tileWidth != 4)
        return op.emitOpError("expecting tile_width to be 1, 2, or 4");
      break;
    default:
      return op.emitOpError("transpose is only supported for 32 and 64 bit "
                            "elements");
    }
 
    return success();
  }
 
  assert(op.getPackRegister() && !op.getTranspose() &&
         "Expecting pack_register should be true and transpose should be "
         "false");
 
  uint32_t vBlocks = op.getVBlocks();
  if (vBlocks != 1 && vBlocks != 2 && vBlocks != 4)
    return op.emitOpError("expecting v_blocks to be 1, 2, or 4");
 
  uint32_t tileHeight = op.getTileHeight();
  uint32_t tileWidth = op.getTileWidth();
  switch (op.getElemSizeInBits()) {
  case 8:
    if (tileHeight < 4 || tileHeight > 32)
      return op.emitOpError("expecting tile_height to be between 4 and 32");
    if (tileWidth < 4 || tileWidth > 16)
      return op.emitOpError("expecting tile_width to be between 4 and 16");
    break;
  case 16:
    if (tileHeight < 2 || tileHeight > 32)
      return op.emitOpError("expecting tile_height to be between 2 and 32");
    if (tileWidth < 2 || tileWidth > 16)
      return op.emitOpError("expecting tile_width to be between 2 and 16");
    if (tileWidth * vBlocks > 32)
      return op.emitOpError(
          "tile_width * v_blocks should be less than or equal "
          "to 32 for 16 bit elements");
    break;
  default:
    return op.emitOpError("pack_register is only supported for 8 and 16 bit "
                          "elements");
  }
 
  return success();
}
 
static LogicalResult verify2DBlockStoreRestriction(BlockStore2dOp op) {
  uint32_t tileHeight = op.getTileHeight();
  if (tileHeight < 1 || tileHeight > 8)
    return op.emitOpError("expecting tile_height to be between 1 and 8");
 
  uint32_t tileWidth = op.getTileWidth();
  switch (op.getElemSizeInBits()) {
  case 8:
    if (tileWidth < 4 || tileWidth > 64)
      return op.emitOpError("expecting tile_width to be between 4 and 64");
    break;
  case 16:
    if (tileWidth < 2 || tileWidth > 32)
      return op.emitOpError("expecting tile_width to be between 2 and 32");
    break;
  case 32:
    if (tileWidth < 1 || tileWidth > 16)
      return op.emitOpError("expecting tile_width to be between 1 and 16");
    break;
  case 64:
    if (tileWidth < 1 || tileWidth > 8)
      return op.emitOpError("expecting tile_width to be between 1 and 8");
    break;
  default:
    return op.emitOpError("expecting elem_size_in_bits to be 8, 16, 32, or 64");
  }
 
  uint32_t vBlocks = op.getVBlocks();
  if (vBlocks != 1)
    return op.emitOpError("expecting v_blocks to be 1");
  return success();
}
 
} // namespace
 
LogicalResult BlockLoad2dOp::verify() {
  if (verify2DBlockLoadRestriction(*this).failed())
    return failure();
 
  if (verifyMatrixInput(*this).failed())
    return failure();
 
  VectorType resTy = getRes().getType();
  if (!resTy.getElementType().isIntOrFloat())
    return emitOpError() << "expecting result element type to be int of float";
  unsigned resElemTySize = resTy.getElementType().getIntOrFloatBitWidth();
  if (getElemSizeInBits() == 32 || getPackRegister()) {
    if (resElemTySize != 32)
      return emitOpError() << "expecting result element type to be 32 bits";
  }
 
  uint32_t tileWidth = getTileWidth();
  if (getPackRegister()) {
    if (tileWidth != 16)
      return emitOpError(
          "tile_width when pack_register is true should be equal "
          "to subgroup size (16 elements)");
    return success();
  }
 
  return success();
}
 
LogicalResult BlockStore2dOp::verify() {
  if (verify2DBlockStoreRestriction(*this).failed())
    return failure();
 
  if (verifyMatrixInput(*this).failed())
    return failure();
 
  uint32_t tileWidth = getTileWidth();
  switch (getElemSizeInBits()) {
  case 8:
    if (tileWidth != 16 && tileWidth != 32)
      return emitOpError("tile_width for 8 bit elements should be equal to "
                         "16 or 32");
    break;
  case 16:
    if (tileWidth != 16)
      return emitOpError("tile_width for 16 bit elements should be equal "
                         "to 16");
    break;
  case 32:
    if (tileWidth != 16)
      return emitOpError("tile_width for 32 bit elements should be equal "
                         "to 16");
    break;
  default:
    llvm_unreachable("unexpected element size");
  }
 
  return success();
}
 
LogicalResult BlockPrefetch2dOp::verify() {
  if (verifyMatrixInput(*this).failed())
    return failure();
 
  uint32_t tileWidth = getTileWidth();
  switch (getElemSizeInBits()) {
  case 8:
    if (tileWidth != 16 && tileWidth != 32)
      return emitOpError("tile_width for 8 bit elements should be equal to "
                         "16 or 32");
    break;
  case 16:
    if (tileWidth != 16)
      return emitOpError("tile_width for 16 bit elements should be equal "
                         "to 16");
    break;
  case 32:
    if (tileWidth != 8 && tileWidth != 16)
      return emitOpError(
          "tile_width for 32 bit elements should be equal to 8 or 16");
    break;
  default:
    llvm_unreachable("unexpected element size");
  }
 
  return success();
}
 
template <typename OpType, typename = std::enable_if_t<llvm::is_one_of<
                               OpType, BlockLoadOp, BlockStoreOp>::value>>
LogicalResult verify1DBlockArg(OpType op) {
  Type srcOrDstTy;
  if constexpr (std::is_same_v<OpType, BlockLoadOp>)
    srcOrDstTy = op.getResult().getType();
  else
    srcOrDstTy = op.getVal().getType();
  VectorType vTy = dyn_cast<VectorType>(srcOrDstTy);
  // scalar case is always valid
  if (!vTy)
    return success();
  int elemTySize = vTy.getElementType().getIntOrFloatBitWidth() / 8;
  if (elemTySize == 1) {
    llvm::SmallSet<int, 4> validSizes{2, 4, 8, 16};
    if (validSizes.contains(vTy.getNumElements()))
      return success();
    else
      return op.emitOpError(
          "vector size must be 2, 4, 8 or 16 for 8-bit element type");
  } else {
    llvm::SmallSet<int, 3> validSizes{2, 4, 8};
    if (validSizes.contains(vTy.getNumElements()))
      return success();
    else
      return op.emitOpError(
          "vector size must be 2, 4 or 8 for element type > 8 bits");
  }
}
 
LogicalResult BlockLoadOp::verify() { return verify1DBlockArg(*this); }
 
LogicalResult BlockStoreOp::verify() { return verify1DBlockArg(*this); }
 
LogicalResult MMAOp::verify() {
  if (getC()) {
    if (getResult().getType() != getC().getType())
      return emitOpError("type of C operand must match result type");
  }
  return success();
}
 
LogicalResult MMAMxOp::verify() {
  if (getC()) {
    if (getResult().getType() != getC().getType())
      return emitOpError("type of C operand must match result type");
  }
  return success();
}
 
/// Number of bits one narrow float value occupies. The narrow values of a
/// `xevm.truncf` destination, or a `xevm.extf` source, are packed into whole
/// bytes, so a sub-byte format fits several values per byte.
static int64_t getNarrowFloatBitWidth(TruncfDstElemTypes etype) {
  return etype == TruncfDstElemTypes::E2M1 ? 4 : 8;
}
static int64_t getNarrowFloatBitWidth(ExtfSrcElemTypes etype) {
  return etype == ExtfSrcElemTypes::E2M1 ? 4 : 8;
}
 
/// Number of values `ty` holds: its length if it is a vector, and one
/// otherwise. SPIR-V has no vector of length one and uses a scalar instead, so
/// a conversion of two fp4 values, which pack into a single byte, has a scalar
/// on its packed side.
static int64_t getNumValues(Type ty) {
  if (auto vecTy = dyn_cast<VectorType>(ty))
    return vecTy.getNumElements();
  return 1;
}
 
/// Total bit width of `ty`, which is a scalar or a vector of a scalar.
static int64_t getPackedBitWidth(Type ty) {
  if (auto vecTy = dyn_cast<VectorType>(ty))
    return vecTy.getNumElements() * vecTy.getElementTypeBitWidth();
  return ty.getIntOrFloatBitWidth();
}
 
/// Verifies that `packedTy` is exactly wide enough to hold `numValues` values
/// of `narrowBits` bits each, rounded up to whole bytes.
static LogicalResult verifyPackedWidth(Operation *op, StringRef packedName,
                                       Type packedTy, int64_t numValues,
                                       int64_t narrowBits) {
  int64_t expected = llvm::alignTo(numValues * narrowBits, 8);
  int64_t actual = getPackedBitWidth(packedTy);
  if (actual != expected)
    return op->emitOpError()
           << packedName << " should be " << expected << " bits wide to hold "
           << numValues << " value(s) of " << narrowBits << " bits, but it is "
           << actual;
  return success();
}
 
LogicalResult TruncfOp::verify() {
  Type srcTy = getSrc().getType();
  Type dstTy = getDst().getType();
  if (getElementTypeOrSelf(srcTy).getIntOrFloatBitWidth() <=
      getElementTypeOrSelf(dstTy).getIntOrFloatBitWidth())
    return emitError(
        "dst element bitwidth should be less than src element bitwidth");
  return verifyPackedWidth(*this, "dst", dstTy, getNumValues(srcTy),
                           getNarrowFloatBitWidth(getDstEtype().getEtype()));
}
 
LogicalResult ExtfOp::verify() {
  Type srcTy = getSrc().getType();
  Type dstTy = getDst().getType();
  if (getElementTypeOrSelf(srcTy).getIntOrFloatBitWidth() >=
      getElementTypeOrSelf(dstTy).getIntOrFloatBitWidth())
    return emitError(
        "dst element bitwidth should be greater than src element bitwidth");
  return verifyPackedWidth(*this, "src", srcTy, getNumValues(dstTy),
                           getNarrowFloatBitWidth(getSrcEtype().getEtype()));
}
 
/// Float semantics the element type attributes of `xevm.truncf` and `xevm.extf`
/// stand for. The narrow formats are the OCP FP8 and FP4 ones: `bf8` is
/// E5M2, `f8` is E4M3 and `e2m1` is FP4.
static const llvm::fltSemantics *getFloatSemantics(TruncfSrcElemTypes etype) {
  switch (etype) {
  case TruncfSrcElemTypes::F16:
    return &llvm::APFloat::IEEEhalf();
  case TruncfSrcElemTypes::BF16:
    return &llvm::APFloat::BFloat();
  }
  return nullptr;
}
 
static const llvm::fltSemantics *getFloatSemantics(ExtfDstElemTypes etype) {
  switch (etype) {
  case ExtfDstElemTypes::F16:
    return &llvm::APFloat::IEEEhalf();
  case ExtfDstElemTypes::BF16:
    return &llvm::APFloat::BFloat();
  }
  return nullptr;
}
 
static const llvm::fltSemantics *getFloatSemantics(TruncfDstElemTypes etype) {
  switch (etype) {
  case TruncfDstElemTypes::BF8:
    return &llvm::APFloat::Float8E5M2();
  case TruncfDstElemTypes::F8:
    return &llvm::APFloat::Float8E4M3FN();
  case TruncfDstElemTypes::E2M1:
    return &llvm::APFloat::Float4E2M1FN();
  }
  return nullptr;
}
 
static const llvm::fltSemantics *getFloatSemantics(ExtfSrcElemTypes etype) {
  switch (etype) {
  case ExtfSrcElemTypes::BF8:
    return &llvm::APFloat::Float8E5M2();
  case ExtfSrcElemTypes::F8:
    return &llvm::APFloat::Float8E4M3FN();
  case ExtfSrcElemTypes::E2M1:
    return &llvm::APFloat::Float4E2M1FN();
  }
  return nullptr;
}
 
/// truncf(extf(a)) -> a, when the two ops convert through the same pair of
/// formats and every narrow value survives the round trip. `bf8` does not
/// qualify: it is IEEE-like, and extending a signaling NaN quiets it, so the
/// original value cannot be recovered. This mirrors `arith.truncf`, which gates
/// the same fold on `APFloatBase::isLosslesslyConvertibleTo`.
OpFoldResult TruncfOp::fold(FoldAdaptor) {
  auto extfOp = getSrc().getDefiningOp<ExtfOp>();
  if (!extfOp)
    return {};
 
  const llvm::fltSemantics *narrowSem =
      getFloatSemantics(getDstEtype().getEtype());
  const llvm::fltSemantics *wideSem =
      getFloatSemantics(getSrcEtype().getEtype());
  if (narrowSem != getFloatSemantics(extfOp.getSrcEtype().getEtype()) ||
      wideSem != getFloatSemantics(extfOp.getDstEtype().getEtype()))
    return {};
 
  Value narrowSrc = extfOp.getSrc();
  if (narrowSrc.getType() != getDst().getType())
    return {};
 
  if (!llvm::APFloatBase::isLosslesslyConvertibleTo(*narrowSem, *wideSem))
    return {};
 
  return narrowSrc;
}
 
LogicalResult BitcastShuffleOp::verify() {
  Type srcTy = getSrc().getType();
  Type resTy = getRes().getType();
  auto srcVecTy = dyn_cast<VectorType>(srcTy);
  auto resVecTy = dyn_cast<VectorType>(resTy);
  // Only a pack (vector -> scalar) and an unpack (scalar -> vector) are
  // supported, so exactly one side is a vector.
  if (static_cast<bool>(srcVecTy) == static_cast<bool>(resVecTy))
    return emitOpError("expected exactly one of src and res to be a vector: a "
                       "pack takes a vector and returns a scalar, an unpack "
                       "takes a scalar and returns a vector");
 
  auto getTotalBitWidth = [](Type ty) -> unsigned {
    if (auto vecTy = dyn_cast<VectorType>(ty))
      return vecTy.getNumElements() * vecTy.getElementTypeBitWidth();
    return ty.getIntOrFloatBitWidth();
  };
  if (getTotalBitWidth(srcTy) != getTotalBitWidth(resTy))
    return emitOpError("src and res types must have the same total bit width");
  return success();
}
 
LogicalResult
XeVMTargetAttr::verify(function_ref<InFlightDiagnostic()> emitError, int O,
                       StringRef triple, StringRef chip, DictionaryAttr flags,
                       ArrayAttr linkFiles) {
  if (O < 0 || O > 3) {
    return emitError()
           << "The optimization level must be a number between 0 and 3.";
  }
  if (triple.empty()) {
    return emitError() << "The target triple cannot be empty.";
  }
  if (chip.empty()) {
    return emitError() << "The target chip cannot be empty.";
  }
  if (linkFiles) {
    for (Attribute fileAttr : linkFiles) {
      if (auto fileStrAttr = llvm::dyn_cast<StringAttr>(fileAttr)) {
        StringRef filePath = fileStrAttr.getValue();
        if (filePath.empty()) {
          return emitError() << "File paths in linkFiles cannot be empty.";
        }
        if (!llvm::sys::fs::exists(filePath)) {
          return emitError() << "File '" << filePath << "' does not exist.";
        }
      }
    }
  }
  return success();
}
 
void XeVMDialect::initialize() {
  addOperations<
#define GET_OP_LIST
#include "mlir/Dialect/LLVMIR/XeVMOps.cpp.inc"
      >();
 
  addAttributes<
#define GET_ATTRDEF_LIST
#include "mlir/Dialect/LLVMIR/XeVMOpsAttributes.cpp.inc"
      >();
  declarePromisedInterface<mlir::gpu::TargetAttrInterface,
                           mlir::xevm::XeVMTargetAttr>();
}
 
#define GET_OP_CLASSES
#include "mlir/Dialect/LLVMIR/XeVMOps.cpp.inc"
 
#define GET_ATTRDEF_CLASSES
#include "mlir/Dialect/LLVMIR/XeVMOpsAttributes.cpp.inc"