[TensorExpr] CudaCodegen: restart counter for function names unique ID inside each codegen instantiation. (#47254)
Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/47254
CUDA codegen used a static global counter for picking names for
functions, but the functions only need to be unique in the scope of the
given codegen. This PR fixes that.
Differential Revision: D24698271
Test Plan: Imported from OSS
Reviewed By: bertmaher
Pulled By: ZolotukhinM
fbshipit-source-id: 516c0087b86b35bbb6ea7c71bb0ed9c3daaca2b8
diff --git a/torch/csrc/jit/tensorexpr/cuda_codegen.cpp b/torch/csrc/jit/tensorexpr/cuda_codegen.cpp
index f100aae..8aca5db 100644
--- a/torch/csrc/jit/tensorexpr/cuda_codegen.cpp
+++ b/torch/csrc/jit/tensorexpr/cuda_codegen.cpp
@@ -658,12 +658,14 @@
};
std::string CudaCodeGen::GetUniqueFuncName(const std::string& func_prefix) {
- // We are using a global counter here to make sure difference instances
- // within CudaCodeGen have different names.
- static int64_t counter = 0;
- ++counter;
- int64_t value = counter;
- return func_prefix + "_" + std::to_string(value);
+ int64_t counter = 0;
+ std::string name = func_prefix;
+ while (taken_func_names.count(name)) {
+ name = func_prefix + "_" + std::to_string(counter++);
+ }
+
+ taken_func_names.insert(name);
+ return name;
}
bool GPUMetaVarRewriter::isFullExtent() {
diff --git a/torch/csrc/jit/tensorexpr/cuda_codegen.h b/torch/csrc/jit/tensorexpr/cuda_codegen.h
index b23f33f..2820b24 100644
--- a/torch/csrc/jit/tensorexpr/cuda_codegen.h
+++ b/torch/csrc/jit/tensorexpr/cuda_codegen.h
@@ -239,6 +239,7 @@
std::unique_ptr<CudaPrinter> printer_;
std::unique_ptr<CudaAnalysis> cuda_analysis_;
std::unique_ptr<GPUMetaVarRewriter> metavar_rewriter_;
+ std::unordered_set<std::string> taken_func_names;
CUfunction function_;
bool has_random_ = false;
