From a74d7fb6ebdf65801f8d64263483491318aa8e16 Mon Sep 17 00:00:00 2001
From: houj04 <35131887+houj04@users.noreply.github.com>
Date: Wed, 15 Sep 2021 14:36:57 +0800
Subject: [PATCH] add set-xpu-device-id function for inference config. (#35572)

---
 paddle/fluid/inference/api/analysis_config.cc         |  8 ++++++++
 .../fluid/inference/api/analysis_predictor_tester.cc  | 11 +++++++++++
 paddle/fluid/inference/api/paddle_analysis_config.h   |  6 ++++++
 paddle/fluid/pybind/inference_api.cc                  |  2 ++
 4 files changed, 27 insertions(+)

diff --git a/paddle/fluid/inference/api/analysis_config.cc b/paddle/fluid/inference/api/analysis_config.cc
index f9c7be9cd4..ac540c7551 100644
--- a/paddle/fluid/inference/api/analysis_config.cc
+++ b/paddle/fluid/inference/api/analysis_config.cc
@@ -114,6 +114,14 @@ void AnalysisConfig::EnableXpu(int l3_workspace_size, bool locked,
   Update();
 }
 
+void AnalysisConfig::SetXpuDeviceId(int device_id) {
+  PADDLE_ENFORCE_EQ(use_xpu_, true,
+                    platform::errors::PreconditionNotMet(
+                        "Should call EnableXpu before SetXpuDeviceId."));
+  xpu_device_id_ = device_id;
+  Update();
+}
+
 void AnalysisConfig::EnableNpu(int device_id) {
 #ifdef PADDLE_WITH_ASCEND_CL
   use_npu_ = true;
diff --git a/paddle/fluid/inference/api/analysis_predictor_tester.cc b/paddle/fluid/inference/api/analysis_predictor_tester.cc
index 87af94a88d..513f3669a1 100644
--- a/paddle/fluid/inference/api/analysis_predictor_tester.cc
+++ b/paddle/fluid/inference/api/analysis_predictor_tester.cc
@@ -323,6 +323,17 @@ TEST(AnalysisPredictor, bf16_pass_strategy) {
   passStrategy.EnableMkldnnBfloat16();
 }
 
+#ifdef PADDLE_WITH_XPU
+TEST(AnalysisPredictor, set_xpu_device_id) {
+  AnalysisConfig config;
+  config.EnableXpu();
+  config.SetXpuDeviceId(0);
+  ASSERT_EQ(config.xpu_device_id(), 0);
+  config.SetXpuDeviceId(1);
+  ASSERT_EQ(config.xpu_device_id(), 1);
+}
+#endif
+
 }  // namespace paddle
 
 namespace paddle_infer {
diff --git a/paddle/fluid/inference/api/paddle_analysis_config.h b/paddle/fluid/inference/api/paddle_analysis_config.h
index dbdd0983b5..a64377f80f 100644
--- a/paddle/fluid/inference/api/paddle_analysis_config.h
+++ b/paddle/fluid/inference/api/paddle_analysis_config.h
@@ -203,6 +203,12 @@ struct PD_INFER_DECL AnalysisConfig {
                  const std::string& precision = "int16",
                  bool adaptive_seqlen = false);
   ///
+  /// \brief Set XPU device id.
+  ///
+  /// \param device_id the XPU card to use (default is 0).
+  ///
+  void SetXpuDeviceId(int device_id = 0);
+  ///
   /// \brief Turn on NPU.
   ///
   /// \param device_id device_id the NPU card to use (default is 0).
diff --git a/paddle/fluid/pybind/inference_api.cc b/paddle/fluid/pybind/inference_api.cc
index 87986aebe0..b1a91cd302 100644
--- a/paddle/fluid/pybind/inference_api.cc
+++ b/paddle/fluid/pybind/inference_api.cc
@@ -523,6 +523,8 @@ void BindAnalysisConfig(py::module *m) {
            py::arg("locked") = false, py::arg("autotune") = true,
            py::arg("autotune_file") = "", py::arg("precision") = "int16",
            py::arg("adaptive_seqlen") = false)
+      .def("set_xpu_device_id", &AnalysisConfig::SetXpuDeviceId,
+           py::arg("device_id") = 0)
       .def("enable_npu", &AnalysisConfig::EnableNpu, py::arg("device_id") = 0)
       .def("disable_gpu", &AnalysisConfig::DisableGpu)
       .def("use_gpu", &AnalysisConfig::use_gpu)
-- 
GitLab