From 935a5f67d1d5f8f1bd95150f389680e4db0b3a59 Mon Sep 17 00:00:00 2001 From: Benjamin Kramer <benny.kra@googlemail.com> Date: Sat, 12 Feb 2022 15:23:41 +0100 Subject: [PATCH] [AffineMap] Move result exprs into trailing storage. NFCI. --- mlir/lib/IR/AffineMap.cpp | 11 +++-------- mlir/lib/IR/AffineMapDetail.h | 27 +++++++++++++++++++-------- 2 files changed, 22 insertions(+), 16 deletions(-) diff --git a/mlir/lib/IR/AffineMap.cpp b/mlir/lib/IR/AffineMap.cpp index 73cbd3bf336f..5d0ffd6fb5c6 100644 --- a/mlir/lib/IR/AffineMap.cpp +++ b/mlir/lib/IR/AffineMap.cpp @@ -299,22 +299,17 @@ unsigned AffineMap::getNumSymbols() const { assert(map && "uninitialized map storage"); return map->numSymbols; } -unsigned AffineMap::getNumResults() const { - assert(map && "uninitialized map storage"); - return map->results.size(); -} +unsigned AffineMap::getNumResults() const { return getResults().size(); } unsigned AffineMap::getNumInputs() const { assert(map && "uninitialized map storage"); return map->numDims + map->numSymbols; } - ArrayRef<AffineExpr> AffineMap::getResults() const { assert(map && "uninitialized map storage"); - return map->results; + return map->results(); } AffineExpr AffineMap::getResult(unsigned idx) const { - assert(map && "uninitialized map storage"); - return map->results[idx]; + return getResults()[idx]; } unsigned AffineMap::getDimPosition(unsigned idx) const { diff --git a/mlir/lib/IR/AffineMapDetail.h b/mlir/lib/IR/AffineMapDetail.h index c0a906b8365c..732c7fd1d3a1 100644 --- a/mlir/lib/IR/AffineMapDetail.h +++ b/mlir/lib/IR/AffineMapDetail.h @@ -17,36 +17,47 @@ #include "mlir/IR/AffineMap.h" #include "mlir/Support/StorageUniquer.h" #include "llvm/ADT/ArrayRef.h" +#include "llvm/Support/TrailingObjects.h" namespace mlir { namespace detail { -struct AffineMapStorage : public StorageUniquer::BaseStorage { +struct AffineMapStorage final + : public StorageUniquer::BaseStorage, + public llvm::TrailingObjects<AffineMapStorage, AffineExpr> { /// The hash key used for uniquing. using KeyTy = std::tuple<unsigned, unsigned, ArrayRef<AffineExpr>>; unsigned numDims; unsigned numSymbols; - - /// The affine expressions for this (multi-dimensional) map. - /// TODO: use trailing objects for this. - ArrayRef<AffineExpr> results; + unsigned numResults; MLIRContext *context; + /// The affine expressions for this (multi-dimensional) map. + ArrayRef<AffineExpr> results() const { + return {getTrailingObjects<AffineExpr>(), numResults}; + } + bool operator==(const KeyTy &key) const { return std::get<0>(key) == numDims && std::get<1>(key) == numSymbols && - std::get<2>(key) == results; + std::get<2>(key) == results(); } // Constructs an AffineMapStorage from a key. The context must be set by the // caller. static AffineMapStorage * construct(StorageUniquer::StorageAllocator &allocator, const KeyTy &key) { - auto *res = new (allocator.allocate<AffineMapStorage>()) AffineMapStorage(); + auto results = std::get<2>(key); + auto byteSize = + AffineMapStorage::totalSizeToAlloc<AffineExpr>(results.size()); + auto *rawMem = allocator.allocate(byteSize, alignof(AffineMapStorage)); + auto *res = new (rawMem) AffineMapStorage(); res->numDims = std::get<0>(key); res->numSymbols = std::get<1>(key); - res->results = allocator.copyInto(std::get<2>(key)); + res->numResults = results.size(); + std::uninitialized_copy(results.begin(), results.end(), + res->getTrailingObjects<AffineExpr>()); return res; } };