1 /* Copyright 2021 The TensorFlow Authors. All Rights Reserved.
2
3 Licensed under the Apache License, Version 2.0 (the "License");
4 you may not use this file except in compliance with the License.
5 You may obtain a copy of the License at
6
7 http://www.apache.org/licenses/LICENSE-2.0
8
9 Unless required by applicable law or agreed to in writing, software
10 distributed under the License is distributed on an "AS IS" BASIS,
11 WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 See the License for the specific language governing permissions and
13 limitations under the License.
14 ==============================================================================*/
15 #include "tensorflow/lite/profiling/memory_usage_monitor.h"
16
17 #include <memory>
18
19 #include <gtest/gtest.h>
20 #include "absl/time/clock.h"
21 #include "absl/time/time.h"
22 #include "tensorflow/lite/profiling/memory_info.h"
23
24 namespace tflite {
25 namespace profiling {
26 namespace memory {
27
28 class MemoryUsageNotSupportedSampler : public MemoryUsageMonitor::Sampler {
29 public:
IsSupported()30 bool IsSupported() override { return false; }
31 };
32
TEST(MemoryUsageMonitor,NotSupported)33 TEST(MemoryUsageMonitor, NotSupported) {
34 MemoryUsageMonitor monitor1(50, std::unique_ptr<MemoryUsageMonitor::Sampler>(
35 new MemoryUsageNotSupportedSampler()));
36 EXPECT_FLOAT_EQ(MemoryUsageMonitor::kInvalidMemUsageMB,
37 monitor1.GetPeakMemUsageInMB());
38
39 MemoryUsageMonitor monitor2(50, nullptr);
40 EXPECT_FLOAT_EQ(MemoryUsageMonitor::kInvalidMemUsageMB,
41 monitor2.GetPeakMemUsageInMB());
42 }
43
44 // Just smoke tests for different call combinations.
45 class MemoryUsageMonitorTest : public ::testing::Test {
46 protected:
47 class FakeMemoryUsageSampler : public MemoryUsageMonitor::Sampler {
48 public:
FakeMemoryUsageSampler(int64_t * num_sleeps)49 explicit FakeMemoryUsageSampler(int64_t* num_sleeps)
50 : sleep_cnt_(num_sleeps) {}
IsSupported()51 bool IsSupported() override { return true; }
GetMemoryUsage()52 MemoryUsage GetMemoryUsage() override {
53 MemoryUsage result;
54 result.max_rss_kb = 5 * ((*sleep_cnt_) + 1) * 1024;
55 return result;
56 }
SleepFor(const absl::Duration & duration)57 void SleepFor(const absl::Duration& duration) override {
58 (*sleep_cnt_)++;
59 absl::SleepFor(duration);
60 }
61
62 private:
63 int64_t* const sleep_cnt_ = nullptr;
64 };
65
SetUp()66 void SetUp() override {
67 monitor_ = std::make_unique<MemoryUsageMonitor>(
68 50, std::unique_ptr<MemoryUsageMonitor::Sampler>(
69 new FakeMemoryUsageSampler(&num_sleeps_)));
70 }
71
72 int64_t num_sleeps_ = 0;
73 std::unique_ptr<MemoryUsageMonitor> monitor_ = nullptr;
74 };
75
TEST_F(MemoryUsageMonitorTest,StartAndStop)76 TEST_F(MemoryUsageMonitorTest, StartAndStop) {
77 monitor_->Start();
78 monitor_->Stop();
79 EXPECT_FLOAT_EQ(5.0 * (num_sleeps_ + 1), monitor_->GetPeakMemUsageInMB());
80 }
81
TEST_F(MemoryUsageMonitorTest,NoStartAndStop)82 TEST_F(MemoryUsageMonitorTest, NoStartAndStop) {
83 monitor_->Stop();
84 EXPECT_FLOAT_EQ(MemoryUsageMonitor::kInvalidMemUsageMB,
85 monitor_->GetPeakMemUsageInMB());
86 }
87
TEST_F(MemoryUsageMonitorTest,StartAndNoStop)88 TEST_F(MemoryUsageMonitorTest, StartAndNoStop) {
89 monitor_->Start();
90 EXPECT_FLOAT_EQ(MemoryUsageMonitor::kInvalidMemUsageMB,
91 monitor_->GetPeakMemUsageInMB());
92 }
93
TEST_F(MemoryUsageMonitorTest,StopFirst)94 TEST_F(MemoryUsageMonitorTest, StopFirst) {
95 monitor_->Stop();
96 EXPECT_FLOAT_EQ(MemoryUsageMonitor::kInvalidMemUsageMB,
97 monitor_->GetPeakMemUsageInMB());
98 monitor_->Start();
99 EXPECT_FLOAT_EQ(MemoryUsageMonitor::kInvalidMemUsageMB,
100 monitor_->GetPeakMemUsageInMB());
101 }
102
TEST_F(MemoryUsageMonitorTest,MultiStartAndStops)103 TEST_F(MemoryUsageMonitorTest, MultiStartAndStops) {
104 monitor_->Start();
105 monitor_->Start();
106 monitor_->Stop();
107 monitor_->Stop();
108 EXPECT_FLOAT_EQ(5.0 * (num_sleeps_ + 1), monitor_->GetPeakMemUsageInMB());
109 }
110
TEST_F(MemoryUsageMonitorTest,StartStopPairs)111 TEST_F(MemoryUsageMonitorTest, StartStopPairs) {
112 monitor_->Start();
113 monitor_->Stop();
114 EXPECT_FLOAT_EQ(5.0 * (num_sleeps_ + 1), monitor_->GetPeakMemUsageInMB());
115
116 monitor_->Start();
117 // Sleep for at least for a duration that's longer than the sampling interval
118 // passed to 'monitor_' (i.e. 50 ms) to simulate the memory usage increase.
119 absl::SleepFor(absl::Milliseconds(100));
120 monitor_->Stop();
121 EXPECT_GE(num_sleeps_, 1);
122 EXPECT_FLOAT_EQ(5.0 * (num_sleeps_ + 1), monitor_->GetPeakMemUsageInMB());
123 }
124
125 } // namespace memory
126 } // namespace profiling
127 } // namespace tflite
128