Class DenseElementsAttrCache

lib/Bytecode/Reader/BytecodeReader.cpp:664–723

A cache for deduplicating constant attributes during parsing.

namespace {
///  A cache for deduplicating constant attributes during parsing.
class DenseElementsAttrCache {
public:
  FailureOr<DenseElementsAttr> getOrCreate(Type type, ArrayRef<uint8_t> data,
                                           MLIRContext &context) {
    // The key is a combination of the expected type and the raw data blob.
    std::pair<Type, ArrayRef<uint8_t>> key = {type, data};
    auto it = cache.find(key);
    if (it != cache.end())
      return it->second;

    // Create a reader for the constant data blob.
    EncodingReader reader(data, context);

    // Cast to TileType to get element type and shape info.
    if (!type)
      return reader.emitError() << "provided type is null";
    auto tileType = mlir::dyn_cast<cuda_tile::TileType>(type);
    if (!tileType || !tileType.getElementType().isIntOrFloat())
      return reader.emitError()
             << "expect Cuda Tile integer or float type but got: " << tileType;

    // Read the size of the raw data buffer.
    uint64_t rawDataSize;
    if (failed(reader.readVarInt(rawDataSize)))
      return reader.emitError() << "failed to read the size of the data buffer";

    // Read the raw byte data.
    ArrayRef<uint8_t> rawUint8Data;
    if (failed(reader.readBytes(rawDataSize, rawUint8Data)))
      return reader.emitError() << "failed to read the raw byte data";

    // Convert ArrayRef<uint8_t> to ArrayRef<char>.
    ArrayRef<char> rawData(reinterpret_cast<const char *>(rawUint8Data.data()),
                           rawUint8Data.size());
    // Validate the buffer size and format.
    if (!DenseElementsAttr::isValidRawBuffer(tileType, rawData)) {
      return reader.emitError() << "failed to validate buffer size and format";
    }

    DenseElementsAttr attr = nullptr;
    // Handle endianness conversion.
    if (llvm::endianness::native == llvm::endianness::big) {
      // Convert endianess.
      SmallVector<char, 64> outDataVec(rawData.size());
      MutableArrayRef<char> convRawData(outDataVec);
      DenseTypedElementsAttr::convertEndianOfArrayRefForBEmachine(
          rawData, convRawData, tileType);
      attr = DenseElementsAttr::getFromRawBuffer(tileType, convRawData);
    } else {
      attr = DenseElementsAttr::getFromRawBuffer(tileType, rawData);
    }

    if (attr)
      cache.insert({key, attr});
    return attr;
  }

private: