diff --git a/paddle/fluid/imperative/tests/test_tracer.cc b/paddle/fluid/imperative/tests/test_tracer.cc index f112b9fc1e90bf06dc78cd73b9315d279f4cc723..9ef13a5d45b2e8894f77c698c036a6820d377c12 100644 --- a/paddle/fluid/imperative/tests/test_tracer.cc +++ b/paddle/fluid/imperative/tests/test_tracer.cc @@ -187,6 +187,16 @@ TEST(test_tracer, test_trace_op_with_multi_device_inputs) { } } #endif + +TEST(test_tracer, test_unique_name_generator) { + // generate two unique names + imperative::Tracer tracer; + auto fc_1 = tracer.GenerateUniqueName("fc"); + auto fc_2 = tracer.GenerateUniqueName("fc"); + ASSERT_STREQ("fc_1", fc_1.c_str()); + ASSERT_STREQ("fc_2", fc_2.c_str()); +} + } // namespace imperative } // namespace paddle diff --git a/paddle/fluid/imperative/tracer.h b/paddle/fluid/imperative/tracer.h index 5a7351c6dbe38b4d5977bd1115a1dddf4eb99471..7bddea53f15a15ebcbb7c3d1a2fb36a87c8a015f 100644 --- a/paddle/fluid/imperative/tracer.h +++ b/paddle/fluid/imperative/tracer.h @@ -29,13 +29,26 @@ namespace paddle { namespace imperative { +class UniqueNameGenerator { + public: + explicit UniqueNameGenerator(std::string prefix = "") : prefix_(prefix) {} + std::string Generate(std::string key = "tmp") { + return prefix_ + key + "_" + std::to_string(++id_); + } + + private: + std::atomic id_{0}; + std::string prefix_; +}; + class Tracer { DISABLE_COPY_AND_ASSIGN(Tracer); public: Tracer() : engine_(new BasicEngine()), - program_desc_tracer_(new jit::ProgramDescTracer()) {} + program_desc_tracer_(new jit::ProgramDescTracer()), + generator_(new UniqueNameGenerator()) {} ~Tracer() = default; @@ -64,6 +77,10 @@ class Tracer { return program_desc_tracer_.get(); } + std::string GenerateUniqueName(std::string key = "tmp") { + return generator_->Generate(key); + } + private: static size_t GenerateUniqueId() { static std::atomic id{0}; @@ -74,6 +91,7 @@ class Tracer { std::unique_ptr engine_; std::unique_ptr program_desc_tracer_; bool enable_program_desc_tracing_{false}; + std::unique_ptr generator_; }; } // namespace imperative