From 870ebb82be387049902d3928180259206ab63efb Mon Sep 17 00:00:00 2001 From: victor-eds Date: Tue, 26 Nov 2024 11:51:29 +0000 Subject: [PATCH] [NIT] Revert `Allocation` analysis changes https://github.com/intel/intel-xpu-backend-for-triton/pull/2605 introduced allocation analysis specialization downstream, but upstreaming lead to a different approach (https://github.com/triton-lang/triton/pull/5070). This commit reverts residual changes to the analysis matching upstream's https://github.com/triton-lang/triton/commit/deee78f4dfd46737b28c8d0aee78dd37cba00873. Signed-off-by: victor-eds --- include/triton/Analysis/Allocation.h | 88 ++++++++++++++-------------- lib/Analysis/Allocation.cpp | 14 ++--- 2 files changed, 50 insertions(+), 52 deletions(-) diff --git a/include/triton/Analysis/Allocation.h b/include/triton/Analysis/Allocation.h index f4de4c6b53..9d0d6684da 100644 --- a/include/triton/Analysis/Allocation.h +++ b/include/triton/Analysis/Allocation.h @@ -99,45 +99,6 @@ class Allocation { using BufferIdSetT = DenseSet; using FuncAllocMapT = CallGraph::FuncDataMapT; - /// A class that represents a shared memory buffer - struct BufferT { - /// Explicit: triton_gpu.local_alloc - /// Scratch: triton_gpu.convert_layout - /// Virtual: triton.call - enum class BufferKind { Explicit, Scratch, Virtual }; - - /// MT: thread-safe - inline static std::atomic nextId = 0; - - BufferKind kind; - BufferId id; - size_t size; - size_t alignment; - size_t offset; - - bool operator==(const BufferT &other) const { return id == other.id; } - bool operator<(const BufferT &other) const { return id < other.id; } - - BufferT() : BufferT(BufferKind::Explicit, 0) {} - BufferT(BufferKind kind, size_t size, size_t alignment = 4, - size_t offset = 0) - : kind(kind), id(nextId++), size(size), alignment(alignment), - offset(offset) {} - - size_t setOffsetAligned(size_t newOffset) { - return offset = llvm::alignTo(newOffset, alignment); - } - }; - - /// Op -> Scratch Buffer - using OpScratchMapT = DenseMap; - /// Value -> Explicit Buffer - using ValueBufferMapT = llvm::MapVector; - /// Value -> Alias Buffer - using AliasBufferMapT = llvm::MapVector>; - /// BufferId -> Buffer - using BufferSetT = std::map; - static constexpr BufferId InvalidBufferId = std::numeric_limits::max(); @@ -153,12 +114,6 @@ class Allocation { /// Returns the operation this analysis was constructed from. Operation *getOperation() const { return operation; } - const OpScratchMapT &getOpScratch() const { return opScratch; } - const OpScratchMapT &getOpVirtual() const { return opVirtual; } - const ValueBufferMapT &getValueBuffer() const { return valueBuffer; } - const AliasBufferMapT &getAliasBuffer() const { return aliasBuffer; } - void setSharedMemorySize(size_t size) { sharedMemorySize = size; } - /// Returns the offset of the given buffer in the shared memory. size_t getOffset(BufferId bufferId) const { return bufferSet.at(bufferId).offset; @@ -222,6 +177,47 @@ class Allocation { /// Returns mapping from operation to list of live LDS buffers std::map> getLiveBuffers(); +private: + /// A class that represents a shared memory buffer + struct BufferT { + /// Explicit: triton_gpu.local_alloc + /// Scratch: triton_gpu.convert_layout + /// Virtual: triton.call + enum class BufferKind { Explicit, Scratch, Virtual }; + + /// MT: thread-safe + inline static std::atomic nextId = 0; + + BufferKind kind; + BufferId id; + size_t size; + size_t alignment; + size_t offset; + + bool operator==(const BufferT &other) const { return id == other.id; } + bool operator<(const BufferT &other) const { return id < other.id; } + + BufferT() : BufferT(BufferKind::Explicit, 0) {} + BufferT(BufferKind kind, size_t size, size_t alignment = 4, + size_t offset = 0) + : kind(kind), id(nextId++), size(size), alignment(alignment), + offset(offset) {} + + size_t setOffsetAligned(size_t newOffset) { + return offset = llvm::alignTo(newOffset, alignment); + } + }; + + /// Op -> Scratch Buffer + using OpScratchMapT = DenseMap; + /// Value -> Explicit Buffer + using ValueBufferMapT = llvm::MapVector; + /// Value -> Alias Buffer + using AliasBufferMapT = llvm::MapVector>; + /// BufferId -> Buffer + using BufferSetT = std::map; + +private: template void addBuffer(KeyType &key, Args &&...args) { auto buffer = BufferT(Kind, std::forward(args)...); @@ -247,6 +243,8 @@ class Allocation { AliasBufferMapT aliasBuffer; BufferSetT bufferSet; size_t sharedMemorySize = 0; + + friend class triton::AllocationAnalysis; }; /// Static analysis that computes the allocation of shared memory buffers diff --git a/lib/Analysis/Allocation.cpp b/lib/Analysis/Allocation.cpp index 3403674e02..53897578aa 100644 --- a/lib/Analysis/Allocation.cpp +++ b/lib/Analysis/Allocation.cpp @@ -289,7 +289,7 @@ class AllocationAnalysis { /// Each buffer is allocated only once. void resolveExplicitBufferLiveness( function_ref(Value value)> getLiveness) { - for (auto valueBufferIter : allocation->getValueBuffer()) { + for (auto valueBufferIter : allocation->valueBuffer) { auto value = valueBufferIter.first; auto *buffer = valueBufferIter.second; bufferRange[buffer] = getLiveness(value); @@ -301,7 +301,7 @@ class AllocationAnalysis { /// arguments are involved. void resolveAliasBufferLiveness( function_ref(Value value)> getLiveness) { - for (const auto &aliasBufferIter : allocation->getAliasBuffer()) { + for (const auto &aliasBufferIter : allocation->aliasBuffer) { auto value = aliasBufferIter.first; auto buffers = aliasBufferIter.second; auto range = getLiveness(value); @@ -334,8 +334,8 @@ class AllocationAnalysis { operationId.lookup(op) + 1)}); } }; - processScratchMemory(allocation->getOpScratch()); - processScratchMemory(allocation->getOpVirtual()); + processScratchMemory(allocation->opScratch); + processScratchMemory(allocation->opVirtual); } /// Resolves liveness of all values involved under the root operation. @@ -499,7 +499,7 @@ class AllocationAnalysis { void allocate(const SmallVector &buffers, const GraphT &interference) { // Reset shared memory size - allocation->setSharedMemorySize(0); + allocation->sharedMemorySize = 0; // First-fit graph coloring // Neighbors are nodes that interfere with each other. // We color a node by finding the index of the first available @@ -534,8 +534,8 @@ class AllocationAnalysis { } if (colors.lookup(x) != 0) x->setOffsetAligned(newOffset); - allocation->setSharedMemorySize( - std::max(allocation->getSharedMemorySize(), x->offset + x->size)); + allocation->sharedMemorySize = + std::max(allocation->sharedMemorySize, x->offset + x->size); } }