[PIR] xpu support pir 1 (#67631)

wanghuancoder · web-flow · commit ab42ff7f5bd9 · 2024-08-23T11:06:00.000+08:00
* xpu support pir 1

* refine

* refine
diff --git a/paddle/fluid/pir/transforms/pd_op_to_kernel_pass.cc b/paddle/fluid/pir/transforms/pd_op_to_kernel_pass.cc
@@ -1631,9 +1631,17 @@ void AddShadowFeedForValue(
     pir::IrContext* ctx,
     std::unordered_map<pir::Operation*, pir::Operation*>* map_op_pair,
     std::unordered_map<pir::Value, pir::Value>* map_value_pair) {
+  phi::Backend backend = phi::Backend::GPU;
+#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
+  backend = phi::Backend::GPU;
+#endif
+#if defined(PADDLE_WITH_XPU)
+  backend = phi::Backend::XPU;
+#endif
+
   if (op_item->result(index).type().isa<DenseTensorType>()) {
     phi::KernelKey shadow_key{
-        phi::Backend::GPU,
+        backend,
         phi::DataLayout::ANY,
         TransToPhiDataType(
             op_item->result(index).type().dyn_cast<DenseTensorType>().dtype())};
@@ -1668,7 +1676,7 @@ void AddShadowFeedForValue(
     }
     // Add ShadowFeedTensors Op
     phi::KernelKey shadow_key{
-        phi::Backend::GPU,
+        backend,
         phi::DataLayout::ANY,
         TransToPhiDataType(vec_type[0].dyn_cast<DenseTensorType>().dtype())};
 
@@ -1723,7 +1731,8 @@ void AddShadowFeedForTuplePopOp(
   }
 
   // if value place not gpu, add shadow feed op
-  if (phi::is_gpu_place(place) && add_shadow_feed) {
+  if ((phi::is_gpu_place(place) || phi::is_xpu_place(place)) &&
+      add_shadow_feed) {
     for (size_t i = 0; i < op_item->num_results(); ++i) {
       AddShadowFeedForValue(i,
                             op_item,
@@ -2949,9 +2958,11 @@ void AddShadowFeedOpForDataOrFeed(
     std::unordered_map<pir::Operation*, pir::Operation*>* map_op_pair,
     std::unordered_map<pir::Value, pir::Value>* map_value_pair) {
   bool feed_op_add_shadow_feed =
-      (op_item->isa<FeedOp>()) && phi::is_gpu_place(place);
+      (op_item->isa<FeedOp>()) &&
+      (phi::is_gpu_place(place) || phi::is_xpu_place(place));
   bool data_op_add_shadow_feed =
-      (op_item->isa<DataOp>()) && phi::is_gpu_place(place) &&
+      (op_item->isa<DataOp>()) &&
+      (phi::is_gpu_place(place) || phi::is_xpu_place(place)) &&
       (kernel_op->attributes()
            .at("place")
            .dyn_cast<PlaceAttribute>()
@@ -3170,7 +3181,7 @@ void ProcessBlock(
           AllocatedDenseTensorType::get(ctx, phi::Place(), dense_tensor_type));
     }
   }
-  if (phi::is_gpu_place(place)) {
+  if (phi::is_gpu_place(place) || phi::is_xpu_place(place)) {
     for (auto& [keyword, arg] : block->kwargs()) {
       if (auto dense_tensor_type = arg.type().dyn_cast<DenseTensorType>()) {
         auto dtype = dense_tensor_type.dtype();
diff --git a/paddle/phi/backends/xpu/xpu1_op_list.cc b/paddle/phi/backends/xpu/xpu1_op_list.cc
@@ -38,6 +38,12 @@ XPUOpMap& get_kl1_ops() {
                      phi::DataType::INT32,
                      phi::DataType::INT64,
                      phi::DataType::BOOL})},
+      {"assign_raw",
+       XPUKernelSet({phi::DataType::FLOAT32,
+                     phi::DataType::FLOAT64,
+                     phi::DataType::INT32,
+                     phi::DataType::INT64,
+                     phi::DataType::BOOL})},
       {"assign_value",
        XPUKernelSet({phi::DataType::FLOAT32,
                      phi::DataType::INT32,
diff --git a/paddle/phi/backends/xpu/xpu2_op_list.cc b/paddle/phi/backends/xpu/xpu2_op_list.cc
@@ -70,6 +70,13 @@ XPUOpMap& get_kl2_ops() {
                      phi::DataType::FLOAT16,
                      phi::DataType::INT64,
                      phi::DataType::BOOL})},
+      {"assign_raw",
+       XPUKernelSet({phi::DataType::FLOAT32,
+                     phi::DataType::FLOAT64,
+                     phi::DataType::INT32,
+                     phi::DataType::FLOAT16,
+                     phi::DataType::INT64,
+                     phi::DataType::BOOL})},
       {"assign_value",
        XPUKernelSet({phi::DataType::FLOAT32,
                      phi::DataType::INT32,
diff --git a/paddle/phi/backends/xpu/xpu3_op_list.cc b/paddle/phi/backends/xpu/xpu3_op_list.cc
@@ -65,6 +65,14 @@ XPUOpMap& get_kl3_ops() {
                      phi::DataType::BFLOAT16,
                      phi::DataType::INT64,
                      phi::DataType::BOOL})},
+      {"assign_raw",
+       XPUKernelSet({phi::DataType::FLOAT32,
+                     phi::DataType::FLOAT64,
+                     phi::DataType::INT32,
+                     phi::DataType::FLOAT16,
+                     phi::DataType::BFLOAT16,
+                     phi::DataType::INT64,
+                     phi::DataType::BOOL})},
       {"assign_value",
        XPUKernelSet({phi::DataType::FLOAT32,
                      phi::DataType::INT32,
diff --git a/paddle/phi/core/kernel_factory.cc b/paddle/phi/core/kernel_factory.cc
@@ -195,6 +195,46 @@ bool KernelFactory::HasKernel(const std::string& kernel_name,
   if (kernel_iter == iter->second.end()) {
     return false;
   }
+
+  if (kernel_key.backend() == Backend::XPU) {
+#if defined(PADDLE_WITH_XPU_KP)
+    auto fluid_op_name = TransToFluidOpName(kernel_name);
+    bool has_kp_kernel = false;
+    VLOG(6) << "fluid_op_name: " << TransToFluidOpName(kernel_name);
+    bool is_xpu_kp_supported = phi::backends::xpu::is_xpu_kp_support_op(
+        fluid_op_name, kernel_key.dtype());
+    // Check in xpu_kp
+    if (is_xpu_kp_supported && FLAGS_run_kp_kernel) {
+      auto kernel_key_kp =
+          KernelKey(Backend::KPS, kernel_key.layout(), kernel_key.dtype());
+      auto kernel_iter_kp = iter->second.find(kernel_key_kp);
+      has_kp_kernel = (kernel_iter_kp != iter->second.end());
+      if (has_kp_kernel) {
+        kernel_key = kernel_key_kp;
+        kernel_iter = kernel_iter_kp;
+      }
+    }
+    // check in xpu
+    bool xpu_unsupport = !phi::backends::xpu::is_xpu_support_op(
+        fluid_op_name, kernel_key.dtype());
+    VLOG(6) << "Current KernelKey is " << kernel_key;
+    // Fall back to CPU, when FLAGS_enable_api_kernel_fallback is true and op
+    // was unregistered in xpu and kp
+    if (FLAGS_enable_api_kernel_fallback &&
+        (kernel_iter == iter->second.end() ||
+         (xpu_unsupport && !has_kp_kernel))) {
+      return false;
+    }
+#elif defined(PADDLE_WITH_XPU) && !defined(PADDLE_WITH_XPU_KP)
+    VLOG(6) << "fluid_op_name: " << TransToFluidOpName(kernel_name);
+    if ((FLAGS_enable_api_kernel_fallback &&
+         kernel_iter == iter->second.end()) ||
+        !phi::backends::xpu::is_xpu_support_op(TransToFluidOpName(kernel_name),
+                                               kernel_key.dtype())) {
+      return false;
+    }
+#endif
+  }
   return true;
 }
 
diff --git a/test/xpu/test_fleet_exe_dist_model_run_xpu.py b/test/xpu/test_fleet_exe_dist_model_run_xpu.py
@@ -34,62 +34,65 @@ def tearDown(self):
         self.temp_dir.cleanup()
 
     def test_dist_model_run(self):
-        # step 0: declare folder to save the model and params
-        path_prefix = os.path.join(
-            self.temp_dir.name, "dist_model_run_test/inf"
-        )
+        with paddle.pir_utils.OldIrGuard():
+            # step 0: declare folder to save the model and params
+            path_prefix = os.path.join(
+                self.temp_dir.name, "dist_model_run_test/inf"
+            )
 
-        # step 1: saving the inference model and params
-        x = paddle.static.data(name='x', shape=[28, 28], dtype='float32')
-        y = paddle.static.data(name='y', shape=[28, 1], dtype='int64')
-        predict = paddle.static.nn.fc(x, 10, activation='softmax')
-        loss = paddle.nn.functional.cross_entropy(predict, y)
-        avg_loss = paddle.tensor.stat.mean(loss)
-        exe = paddle.static.Executor(paddle.XPUPlace(0))
-        exe.run(paddle.static.default_startup_program())
-        x_data = np.random.randn(28, 28).astype('float32')
-        y_data = np.random.randint(0, 9, size=[28, 1]).astype('int64')
-        exe.run(
-            paddle.static.default_main_program(),
-            feed={'x': x_data, 'y': y_data},
-            fetch_list=[avg_loss],
-        )
-        paddle.static.save_inference_model(path_prefix, [x, y], [avg_loss], exe)
-        print('save model to', path_prefix)
+            # step 1: saving the inference model and params
+            x = paddle.static.data(name='x', shape=[28, 28], dtype='float32')
+            y = paddle.static.data(name='y', shape=[28, 1], dtype='int64')
+            predict = paddle.static.nn.fc(x, 10, activation='softmax')
+            loss = paddle.nn.functional.cross_entropy(predict, y)
+            avg_loss = paddle.tensor.stat.mean(loss)
+            exe = paddle.static.Executor(paddle.XPUPlace(0))
+            exe.run(paddle.static.default_startup_program())
+            x_data = np.random.randn(28, 28).astype('float32')
+            y_data = np.random.randint(0, 9, size=[28, 1]).astype('int64')
+            exe.run(
+                paddle.static.default_main_program(),
+                feed={'x': x_data, 'y': y_data},
+                fetch_list=[avg_loss],
+            )
+            paddle.static.save_inference_model(
+                path_prefix, [x, y], [avg_loss], exe
+            )
+            print('save model to', path_prefix)
 
-        # step 2: prepare fake data for the inference
-        x_tensor = np.random.randn(28, 28).astype('float32')
-        y_tensor = np.random.randint(0, 9, size=[28, 1]).astype('int64')
+            # step 2: prepare fake data for the inference
+            x_tensor = np.random.randn(28, 28).astype('float32')
+            y_tensor = np.random.randint(0, 9, size=[28, 1]).astype('int64')
 
-        # step 3: init the dist model to inference with fake data
-        config = core.DistModelConfig()
-        config.model_dir = path_prefix
-        config.place = 'XPU'
-        dist = core.DistModel(config)
-        dist.init()
-        dist_x = core.DistModelTensor(x_tensor, 'x')
-        dist_y = core.DistModelTensor(y_tensor, 'y')
-        input_data = [dist_x, dist_y]
-        output_rst = dist.run(input_data)
-        dist_model_rst = output_rst[0].as_ndarray().ravel().tolist()
-        print("dist model rst:", dist_model_rst)
+            # step 3: init the dist model to inference with fake data
+            config = core.DistModelConfig()
+            config.model_dir = path_prefix
+            config.place = 'XPU'
+            dist = core.DistModel(config)
+            dist.init()
+            dist_x = core.DistModelTensor(x_tensor, 'x')
+            dist_y = core.DistModelTensor(y_tensor, 'y')
+            input_data = [dist_x, dist_y]
+            output_rst = dist.run(input_data)
+            dist_model_rst = output_rst[0].as_ndarray().ravel().tolist()
+            print("dist model rst:", dist_model_rst)
 
-        # step 4: use framework's api to inference with fake data
-        [
-            inference_program,
-            feed_target_names,
-            fetch_targets,
-        ] = paddle.static.load_inference_model(path_prefix, exe)
-        results = exe.run(
-            inference_program,
-            feed={'x': x_tensor, 'y': y_tensor},
-            fetch_list=fetch_targets,
-        )
-        load_inference_model_rst = results[0]
-        print("load inference model api rst:", load_inference_model_rst)
+            # step 4: use framework's api to inference with fake data
+            [
+                inference_program,
+                feed_target_names,
+                fetch_targets,
+            ] = paddle.static.load_inference_model(path_prefix, exe)
+            results = exe.run(
+                inference_program,
+                feed={'x': x_tensor, 'y': y_tensor},
+                fetch_list=fetch_targets,
+            )
+            load_inference_model_rst = results[0]
+            print("load inference model api rst:", load_inference_model_rst)
 
-        # step 5: compare two results
-        np.testing.assert_allclose(dist_model_rst, load_inference_model_rst)
+            # step 5: compare two results
+            np.testing.assert_allclose(dist_model_rst, load_inference_model_rst)
 
 
 if __name__ == '__main__':
diff --git a/test/xpu/test_linspace_op_xpu.py b/test/xpu/test_linspace_op_xpu.py
@@ -20,7 +20,7 @@
     create_test_class,
     get_xpu_op_support_types,
 )
-from op_test_xpu import XPUOpTest, convert_np_dtype_to_dtype_
+from op_test_xpu import XPUOpTest
 
 import paddle
 
@@ -46,7 +46,13 @@ def setUp(self):
                 'Num': np.array([11]).astype('int32'),
             }
             self.outputs = {'Out': np.arange(0, 11).astype(self.dtype)}
-            self.attrs = {'dtype': int(convert_np_dtype_to_dtype_(self.dtype))}
+            self.attrs = {
+                'dtype': int(
+                    paddle.base.framework.convert_np_dtype_to_proto_type(
+                        self.dtype
+                    )
+                )
+            }
 
         def set_attrs(self):
             pass
