fix gm size overflow (#62940)

lj970926 · web-flow · commit 69217ad9e881 · 2024-03-22T13:51:08.000+08:00
diff --git a/paddle/phi/backends/xpu/xpu_context.cc b/paddle/phi/backends/xpu/xpu_context.cc
@@ -31,7 +31,7 @@ namespace xpu = baidu::xpu::api;
 namespace phi {
 
 struct XPUContext::Impl {
-  void SetL3Cache(int l3_size = 1024) {
+  void SetL3Cache(int64_t l3_size = 1024) {
     PADDLE_ENFORCE_XPU_SUCCESS(xpu_wait(context_->xpu_stream));
     context_->_l3_mgr.set(nullptr, 0, true);  // free origin l3
     void* l3_ptr = nullptr;
@@ -130,7 +130,7 @@ struct XPUContext::Impl {
     }
   }
 
-  void Init(int gm_default_size = 1024, int l3_default_size = 1024) {
+  void Init(int64_t gm_default_size = 1024, int64_t l3_default_size = 1024) {
     owned_ = true;
     backends::xpu::XPUDeviceGuard guard(place_.GetDeviceId());
     LOG_FIRST_N(WARNING, 1)
@@ -222,26 +222,26 @@ struct XPUContext::Impl {
   xpu::BKCLContext_t bkcl_context_{nullptr};
 };
 
-static int get_gm_size(int i) {
-  int default_size = 1024;
+static int64_t get_gm_size(int i) {
+  int64_t default_size = 1024;
   if (std::getenv("XPUAPI_DEFAULT_SIZE") != nullptr) {
-    default_size = atoi(std::getenv("XPUAPI_DEFAULT_SIZE"));
+    default_size = std::atoll(std::getenv("XPUAPI_DEFAULT_SIZE"));
   }
   std::string cur_env = std::string("XPUAPI_DEFAULT_SIZE") + std::to_string(i);
   if (std::getenv(cur_env.c_str()) != nullptr) {
-    default_size = atoi(std::getenv(cur_env.c_str()));
+    default_size = std::atoll(std::getenv(cur_env.c_str()));
   }
   return default_size;
 }
 
-static int get_l3_size(int i) {
-  int default_size = 1024;
+static int64_t get_l3_size(int i) {
+  int64_t default_size = 1024;
   if (std::getenv("XPU_PADDLE_L3_SIZE") != nullptr) {
-    default_size = atoi(std::getenv("XPU_PADDLE_L3_SIZE"));
+    default_size = std::atoll(std::getenv("XPU_PADDLE_L3_SIZE"));
   }
   std::string cur_env = std::string("XPU_PADDLE_L3_SIZE") + std::to_string(i);
   if (std::getenv(cur_env.c_str()) != nullptr) {
-    default_size = atoi(std::getenv(cur_env.c_str()));
+    default_size = std::atoll(std::getenv(cur_env.c_str()));
   }
   return default_size;
 }
@@ -324,7 +324,7 @@ void XPUContext::SetXContext(xpu::Context* context, int i) {
   impls_[i]->SetXContext(context);
 }
 
-void XPUContext::SetL3Cache(int l3_size, int i) {
+void XPUContext::SetL3Cache(int64_t l3_size, int i) {
   impls_[i]->SetL3Cache(l3_size);
 }
 
diff --git a/paddle/phi/backends/xpu/xpu_context.h b/paddle/phi/backends/xpu/xpu_context.h
@@ -71,7 +71,7 @@ class XPUContext : public DeviceContext,
   // resource as external, and will not delete any resource when destructing.
   void SetXContext(xpu::Context*, int i = 0);
 
-  void SetL3Cache(int l3_size = 1024, int i = 0);
+  void SetL3Cache(int64_t l3_size = 1024, int i = 0);
 
   void SetXpuVersion(int version);
 

Original file line number	Diff line number	Diff line change
`@@ -31,7 +31,7 @@ namespace xpu = baidu::xpu::api;`
`31`	`31`	`namespace phi {`
`32`	`32`
`33`	`33`	`struct XPUContext::Impl {`
`34`		`- void SetL3Cache(int l3_size = 1024) {`
	`34`	`+ void SetL3Cache(int64_t l3_size = 1024) {`
`35`	`35`	`PADDLE_ENFORCE_XPU_SUCCESS(xpu_wait(context_->xpu_stream));`
`36`	`36`	`context_->_l3_mgr.set(nullptr, 0, true); // free origin l3`
`37`	`37`	`void* l3_ptr = nullptr;`
`@@ -130,7 +130,7 @@ struct XPUContext::Impl {`
`130`	`130`	`}`
`131`	`131`	`}`
`132`	`132`
`133`		`- void Init(int gm_default_size = 1024, int l3_default_size = 1024) {`
	`133`	`+ void Init(int64_t gm_default_size = 1024, int64_t l3_default_size = 1024) {`
`134`	`134`	`owned_ = true;`
`135`	`135`	`backends::xpu::XPUDeviceGuard guard(place_.GetDeviceId());`
`136`	`136`	`LOG_FIRST_N(WARNING, 1)`
`@@ -222,26 +222,26 @@ struct XPUContext::Impl {`
`222`	`222`	`xpu::BKCLContext_t bkcl_context_{nullptr};`
`223`	`223`	`};`
`224`	`224`
`225`		`-static int get_gm_size(int i) {`
`226`		`- int default_size = 1024;`
	`225`	`+static int64_t get_gm_size(int i) {`
	`226`	`+ int64_t default_size = 1024;`
`227`	`227`	`if (std::getenv("XPUAPI_DEFAULT_SIZE") != nullptr) {`
`228`		`- default_size = atoi(std::getenv("XPUAPI_DEFAULT_SIZE"));`
	`228`	`+ default_size = std::atoll(std::getenv("XPUAPI_DEFAULT_SIZE"));`
`229`	`229`	`}`
`230`	`230`	`std::string cur_env = std::string("XPUAPI_DEFAULT_SIZE") + std::to_string(i);`
`231`	`231`	`if (std::getenv(cur_env.c_str()) != nullptr) {`
`232`		`- default_size = atoi(std::getenv(cur_env.c_str()));`
	`232`	`+ default_size = std::atoll(std::getenv(cur_env.c_str()));`
`233`	`233`	`}`
`234`	`234`	`return default_size;`
`235`	`235`	`}`
`236`	`236`
`237`		`-static int get_l3_size(int i) {`
`238`		`- int default_size = 1024;`
	`237`	`+static int64_t get_l3_size(int i) {`
	`238`	`+ int64_t default_size = 1024;`
`239`	`239`	`if (std::getenv("XPU_PADDLE_L3_SIZE") != nullptr) {`
`240`		`- default_size = atoi(std::getenv("XPU_PADDLE_L3_SIZE"));`
	`240`	`+ default_size = std::atoll(std::getenv("XPU_PADDLE_L3_SIZE"));`
`241`	`241`	`}`
`242`	`242`	`std::string cur_env = std::string("XPU_PADDLE_L3_SIZE") + std::to_string(i);`
`243`	`243`	`if (std::getenv(cur_env.c_str()) != nullptr) {`
`244`		`- default_size = atoi(std::getenv(cur_env.c_str()));`
	`244`	`+ default_size = std::atoll(std::getenv(cur_env.c_str()));`
`245`	`245`	`}`
`246`	`246`	`return default_size;`
`247`	`247`	`}`
`@@ -324,7 +324,7 @@ void XPUContext::SetXContext(xpu::Context* context, int i) {`
`324`	`324`	`impls_[i]->SetXContext(context);`
`325`	`325`	`}`
`326`	`326`
`327`		`-void XPUContext::SetL3Cache(int l3_size, int i) {`
	`327`	`+void XPUContext::SetL3Cache(int64_t l3_size, int i) {`
`328`	`328`	`impls_[i]->SetL3Cache(l3_size);`
`329`	`329`	`}`
`330`	`330`