[PIR+CINN]Support data transformer for cinn_runtime.jit_kernel (#62742)

Aurelius84 · web-flow · commit cb3f79f54768 · 2024-03-18T10:28:49.000+08:00
* [PIR+CINN]Support data transformer for cinn_runtime.jit_kernel

* fix unused var
diff --git a/paddle/fluid/pir/transforms/pd_op_to_kernel_pass.cc b/paddle/fluid/pir/transforms/pd_op_to_kernel_pass.cc
@@ -1817,17 +1817,38 @@ void HandleForSpecialOp(
   }
 
   if (op_item->name() == "cinn_runtime.jit_kernel") {
-    if (op_item->num_operands() > 0) {
-      for (size_t i = 0; i < op_item->num_operands(); ++i) {
-        auto cur_in = op_item->operand_source(i);
-        if (!cur_in) {
-          vec_inputs.emplace_back();
-          continue;
+    for (size_t i = 0; i < op_item->num_operands(); ++i) {
+      auto cur_in = op_item->operand_source(i);
+      if (!cur_in) {
+        vec_inputs.emplace_back();
+        continue;
+      }
+      auto new_in = GetNewInput(
+          cur_in, *map_value_pair, static_cast<int>(i), op_item->name());
+      // For data transform
+      if (new_in.type().isa<AllocatedDenseTensorType>()) {
+        auto in_place =
+            new_in.type().dyn_cast<AllocatedDenseTensorType>().place();
+        auto dst_backend = phi::TransToPhiBackend(place);
+        bool need_trans =
+            (in_place.GetType() != phi::AllocationType::UNDEFINED) &&
+            (paddle::experimental::NeedTransformPlace(
+                in_place, dst_backend, {}));
+        if (need_trans) {
+          VLOG(6) << "need trans from " << in_place << " to " << dst_backend;
+          auto value_type =
+              op_item->operand_source(i).type().dyn_cast<DenseTensorType>();
+          auto out_place = phi::TransToPhiPlace(dst_backend);
+          auto out_type =
+              AllocatedDenseTensorType::get(ctx, out_place, value_type);
+          phi::KernelKey kernel_key(phi::Backend::GPU,
+                                    phi::DataLayout::ANY,
+                                    TransToPhiDataType(value_type.dtype()));
+          new_in = AddPlaceTransferOp(
+              new_in, out_type, in_place, out_place, kernel_key, block);
         }
-        auto new_in = GetNewInput(
-            cur_in, *map_value_pair, static_cast<int>(i), op_item->name());
-        vec_inputs.push_back(new_in);
       }
+      vec_inputs.push_back(new_in);
     }
 
     for (size_t i = 0; i < op_item->num_results(); ++i) {
diff --git a/test/ir/pir/cinn/test_data_transform.py b/test/ir/pir/cinn/test_data_transform.py
@@ -0,0 +1,66 @@
+# Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import unittest
+
+import numpy as np
+import utils
+from test_cinn_sub_graph import TestCinnSubGraphBase
+
+import paddle
+from paddle import nn
+
+
+class DataTransformNet(nn.Layer):
+    def __init__(self):
+        super().__init__()
+
+    def forward(self, x):
+        y = paddle.tensor.creation.fill_constant(
+            x.shape, 'float32', 1.0, force_cpu=True
+        )
+        y = paddle.static.Print(y)
+        z = paddle.nn.functional.relu(y)
+        return x + z
+
+
+class TestDataTransformNet(TestCinnSubGraphBase):
+    def prepare_data(self):
+        self.shape = [16, 16]
+        self.x = paddle.randn(self.shape, dtype="float32")
+
+    def check_jit_kernel_info(self, static_fn):
+        utils.check_jit_kernel_number(static_fn, 1)
+        utils.check_jit_kernel_structure(static_fn, {utils.JIT_KERNEL_NAME: 1})
+
+    def eval(self, use_cinn):
+        paddle.seed(2022)
+        net = DataTransformNet()
+        net = utils.apply_to_static(net, use_cinn)
+        net.eval()
+        out = net(self.x)
+        if use_cinn:
+            self.check_jit_kernel_info(net.forward)
+        return out
+
+    def test_eval(self):
+        cinn_out = self.eval(use_cinn=True)
+        dy_out = self.eval(use_cinn=False)
+        np.testing.assert_allclose(
+            cinn_out.numpy(), dy_out.numpy(), atol=1e-6, rtol=1e-6
+        )
+
+
+if __name__ == '__main__':
+    unittest.main()
