@@ -30,32 +30,32 @@ namespace tensorrt_llm::runtime
3030class MemoryCounters
3131{
3232public:
33- using SizeType32 = std::size_t ;
33+ using SizeType = std::size_t ;
3434 using DiffType = std::ptrdiff_t ;
3535
3636 MemoryCounters () = default ;
3737
38- [[nodiscard]] SizeType32 getGpu () const
38+ [[nodiscard]] SizeType getGpu () const
3939 {
4040 return mGpu ;
4141 }
4242
43- [[nodiscard]] SizeType32 getCpu () const
43+ [[nodiscard]] SizeType getCpu () const
4444 {
4545 return mCpu ;
4646 }
4747
48- [[nodiscard]] SizeType32 getPinned () const
48+ [[nodiscard]] SizeType getPinned () const
4949 {
5050 return mPinned ;
5151 }
5252
53- [[nodiscard]] SizeType32 getUVM () const
53+ [[nodiscard]] SizeType getUVM () const
5454 {
5555 return mUVM ;
5656 }
5757
58- [[nodiscard]] SizeType32 getPinnedPool () const
58+ [[nodiscard]] SizeType getPinnedPool () const
5959 {
6060 return mPinnedPool ;
6161 }
@@ -91,12 +91,8 @@ class MemoryCounters
9191 };
9292
9393 template <MemoryType T>
94- void allocate (SizeType32 size)
94+ void allocate (SizeType size)
9595 {
96- if (size > static_cast <SizeType32>(std::numeric_limits<DiffType>::max ()))
97- {
98- TLLM_THROW (" Memory size too large for diff type: %zu" , size);
99- }
10096 auto const sizeDiff = static_cast <DiffType>(size);
10197 if constexpr (T == MemoryType::kGPU )
10298 {
@@ -129,15 +125,11 @@ class MemoryCounters
129125 }
130126 }
131127
132- void allocate (MemoryType memoryType, SizeType32 size);
128+ void allocate (MemoryType memoryType, SizeType size);
133129
134130 template <MemoryType T>
135- void deallocate (SizeType32 size)
131+ void deallocate (SizeType size)
136132 {
137- if (size > static_cast <SizeType32>(std::numeric_limits<DiffType>::max ()))
138- {
139- TLLM_THROW (" Memory size too large for diff type: %zu" , size);
140- }
141133 auto const sizeDiff = -static_cast <DiffType>(size);
142134 if constexpr (T == MemoryType::kGPU )
143135 {
@@ -170,23 +162,18 @@ class MemoryCounters
170162 }
171163 }
172164
173- void deallocate (MemoryType memoryType, SizeType32 size);
165+ void deallocate (MemoryType memoryType, SizeType size);
174166
175167 static MemoryCounters& getInstance ();
176168
177- MemoryCounters (MemoryCounters const &) = delete ;
178- MemoryCounters& operator =(MemoryCounters const &) = delete ;
179- MemoryCounters (MemoryCounters&&) = delete ;
180- MemoryCounters& operator =(MemoryCounters&&) = delete ;
181-
182- static std::string bytesToString (SizeType32 bytes, int precision = 2 );
169+ static std::string bytesToString (SizeType bytes, int precision = 2 );
183170
184171 static std::string bytesToString (DiffType bytes, int precision = 2 );
185172
186173 [[nodiscard]] std::string toString () const ;
187174
188175private:
189- std::atomic<SizeType32 > mGpu {}, mCpu {}, mPinned {}, mUVM {}, mPinnedPool {};
176+ std::atomic<SizeType > mGpu {}, mCpu {}, mPinned {}, mUVM {}, mPinnedPool {};
190177 std::atomic<DiffType> mGpuDiff {}, mCpuDiff {}, mPinnedDiff {}, mUVMDiff {}, mPinnedPoolDiff {};
191178};
192179
0 commit comments