cinn(backend): support run rope graph from llama (#61115)

6clc · web-flow · commit 2baf472821cd · 2024-01-27T15:15:16.000+08:00
support run rope graph from llama
diff --git a/paddle/cinn/hlir/op/contrib/gather_nd.cc b/paddle/cinn/hlir/op/contrib/gather_nd.cc
@@ -86,6 +86,48 @@ ir::Tensor GatherNd(const ir::Tensor &x,
   return res;
 }
 
+ir::Tensor GatherNdSymbolic(const ir::Tensor &x,
+                            const ir::Tensor &index,
+                            const std::string &name) {
+  std::vector<Expr> x_shape = x->shape;
+  std::vector<Expr> index_shape = index->shape;
+  size_t x_shape_size = x_shape.size();
+  size_t index_shape_size = index_shape.size();
+  std::vector<Expr> out_shape;
+  out_shape.insert(out_shape.end(), index_shape.begin(), index_shape.end() - 1);
+  out_shape.insert(out_shape.end(),
+                   x_shape.begin() + index_shape.back().as_int64(),
+                   x_shape.end());
+  auto res = Compute(
+      out_shape,
+      [=](const std::vector<Expr> &indices) {
+        std::vector<Expr> indices_position;
+        for (size_t i = 0; i < index_shape_size - 1; ++i) {
+          indices_position.push_back(
+              ir::Cast::Make(cinn::common::Int(64), indices[i]));
+        }
+        indices_position.push_back(
+            ir::Cast::Make(cinn::common::Int(64), Expr(0)));
+        size_t indices_position_size = indices_position.size();
+        std::vector<Expr> real_indices;
+        for (size_t i = 0; i < index_shape.back().as_int64(); ++i) {
+          indices_position[indices_position_size - 1] =
+              ir::Cast::Make(cinn::common::Int(64), Expr(i));
+          real_indices.push_back(
+              ir::Cast::Make(cinn::common::Int(64), index(indices_position)));
+        }
+        if (real_indices.size() == x_shape_size) {
+          return x(real_indices);
+        }
+        for (size_t i = index_shape_size - 1; i < indices.size(); ++i) {
+          real_indices.push_back(indices[i]);
+        }
+        return x(real_indices);
+      },
+      name);
+  return res;
+}
+
 std::shared_ptr<framework::OpStrategy> StrategyForGatherNd(
     const framework::NodeAttr &attrs,
     const std::vector<ir::Tensor> &inputs,
@@ -190,7 +232,7 @@ std::shared_ptr<framework::OpStrategy> StrategyForGatherNdSymbolic(
                 << ", output_shapes: " << utils::Join(output_shapes[0], ", ");
         CHECK_EQ(pack_args.size(), 3U);
         std::string tensor_name = pack_args[2].operator std::string();
-        ir::Tensor out = GatherNd(tensor_x, tensor_index, tensor_name);
+        ir::Tensor out = GatherNdSymbolic(tensor_x, tensor_index, tensor_name);
         std::vector<CINNValue> res;
         stages->InsertLazily(out);
         res.push_back(CINNValue(out));
diff --git a/paddle/cinn/hlir/op/transform.cc b/paddle/cinn/hlir/op/transform.cc
@@ -1616,7 +1616,7 @@ std::shared_ptr<OpStrategy> StrategyForSliceSymbolic(
 
   std::vector<Expr> output_shape;
   for (auto &i : output_shapes[0]) {
-    output_shape.push_back(Expr(i));
+    output_shape.push_back(i->dim_expr);
     LOG(INFO) << "output_shape: " << output_shape.back();
     CHECK(output_shape.back().type().valid());
   }
@@ -1636,7 +1636,7 @@ std::shared_ptr<OpStrategy> StrategyForSliceSymbolic(
         CHECK(arg_pack[1].is_string());
         std::string tensor_name = arg_pack[1].operator std::string();
 
-        auto out = pe::Slice(
+        auto out = pe::SliceSymbolic(
             A, starts, axes, strides, decrease_axis, output_shape, tensor_name);
         LOG(INFO) << "out: " << out;
         auto stages = CreateStages({out});
diff --git a/paddle/cinn/hlir/pe/transform.cc b/paddle/cinn/hlir/pe/transform.cc
@@ -1058,6 +1058,64 @@ ir::Tensor Slice(const ir::Tensor& A,
       output_name);
 }
 
+ir::Tensor SliceSymbolic(const ir::Tensor& A,
+                         const std::vector<int>& starts,
+                         const std::vector<int>& axes,
+                         const std::vector<int>& strides,
+                         const std::vector<int>& decrease_axis,
+                         const std::vector<Expr>& output_shape,
+                         const std::string& output_name) {
+  std::vector<Expr> input_shape;
+  for (const auto& shape : A->shape) {
+    input_shape.emplace_back(shape);
+  }
+
+  std::vector<int> new_starts(starts);
+  for (int i = 0; i < axes.size(); i++) {
+    CHECK(input_shape[axes[i]].is_constant())
+        << "Not supported Slice in dynamic dimensions, because the "
+           "relationship between slice range and symbol size cannot be "
+           "determined at compile time";
+    if (new_starts[i] < -input_shape[axes[i]].as_int64()) {
+      new_starts[i] = 0;
+    } else if (new_starts[i] < 0) {
+      new_starts[i] = input_shape[axes[i]].as_int64() + new_starts[i];
+    } else if (new_starts[i] > input_shape[axes[i]].as_int64()) {
+      new_starts[i] = input_shape[axes[i]].as_int64() - 1;
+    }
+  }
+
+  // output = input[starts:ends:strides]
+  // Note that when strides < 0, the output reverse:
+  // data=[[1,2,3,4],[5,6,7,8],]
+  // axes=[0,1]
+  // starts=[1,3]
+  // ends=[2,0]
+  // strides=[1,-1]
+  // ==> result=[[8,7,6],]
+  return Compute(
+      output_shape,
+      [=](const std::vector<Expr>& indice) {
+        std::vector<Expr> temp;
+        int indice_i = 0;
+        for (int i = 0; i < input_shape.size(); ++i) {
+          if (std::find(decrease_axis.cbegin(), decrease_axis.cend(), i) !=
+              decrease_axis.cend()) {
+            temp.emplace_back(0);
+          } else {
+            temp.emplace_back(indice[indice_i]);
+            indice_i++;
+          }
+        }
+        for (int i = 0; i < axes.size(); i++) {
+          temp[axes[i]] =
+              temp[axes[i]] * Expr(strides[i]) + Expr(new_starts[i]);
+        }
+        return A(temp);
+      },
+      output_name);
+}
+
 ir::Tensor SliceAssign(const ir::Tensor& input,
                        const ir::Tensor& assign,
                        const std::vector<int>& axes,
diff --git a/paddle/cinn/hlir/pe/transform.h b/paddle/cinn/hlir/pe/transform.h
@@ -184,6 +184,14 @@ ir::Tensor Slice(const ir::Tensor& A,
                  const std::vector<Expr>& output_shape,
                  const std::string& output_name);
 
+ir::Tensor SliceSymbolic(const ir::Tensor& A,
+                         const std::vector<int>& starts,
+                         const std::vector<int>& axes,
+                         const std::vector<int>& strides,
+                         const std::vector<int>& decrease_axis,
+                         const std::vector<Expr>& output_shape,
+                         const std::string& output_name);
+
 /**
  * @brief Perform meta op SliceAssign
  * @param input The input tensor
diff --git a/test/ir/pir/cinn/CMakeLists.txt b/test/ir/pir/cinn/CMakeLists.txt
@@ -44,5 +44,14 @@ if(WITH_GPU)
       ${CMAKE_CURRENT_SOURCE_DIR}/test_rms_norm.py
     WORKING_DIRECTORY ${CMAKE_BINARY_DIR})
   set_tests_properties(test_rms_norm PROPERTIES LABELS "RUN_TYPE=CINN")
+  add_test(
+    NAME test_rope_symbolic
+    COMMAND
+      ${CMAKE_COMMAND} -E env
+      PYTHONPATH=${CMAKE_BINARY_DIR}:${CMAKE_BINARY_DIR}/python/:$ENV{PYTHONPATH}
+      FLAGS_cinn_bucket_compile=True FLAGS_enable_pir_api=1 ${PYTHON_EXECUTABLE}
+      ${CMAKE_CURRENT_SOURCE_DIR}/test_rope.py
+    WORKING_DIRECTORY ${CMAKE_BINARY_DIR})
+  set_tests_properties(test_rope_symbolic PROPERTIES LABELS "RUN_TYPE=CINN")
 
 endif()
diff --git a/test/ir/pir/cinn/test_rope.py b/test/ir/pir/cinn/test_rope.py
@@ -0,0 +1,82 @@
+# 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
+
+from test_cinn_sub_graph import TestCinnSubGraphBase, apply_to_static
+
+import paddle
+from paddle import nn
+
+
+class RotaryPosEmb(nn.Layer):
+    def __init__(self):
+        super().__init__()
+
+    def forward(self, q, k, cos, sin, position_ids):
+        cos = cos.squeeze(axis=[0, 2])  # [seq_len, dim]
+        sin = sin.squeeze(axis=[0, 2])  # [seq_len, dim]
+
+        cos = cos[position_ids].unsqueeze(2)  # [bs, seq_len, 1, dim]
+        sin = sin[position_ids].unsqueeze(2)  # [bs, seq_len, 1, dim]
+        q_embed = (q * cos) + (self.rotate_half(q) * sin)
+        k_embed = (k * cos) + (self.rotate_half(k) * sin)
+        return q_embed, k_embed
+
+    def rotate_half(self, x):
+        """Rotates half the hidden dims of the input."""
+        x1 = x[..., : x.shape[-1] // 2]
+        x2 = x[..., x.shape[-1] // 2 :]
+        return paddle.concat([-x2, x1], axis=-1)  # shape is the same as x
+
+
+class TestRotaryPosEmb(TestCinnSubGraphBase):
+    def prepare_data(self):
+        self.q = paddle.randn([1, 2048, 8, 96], dtype="float32")
+        self.q.stop_gradient = False
+
+        self.k = paddle.randn([1, 2048, 8, 96], dtype="float32")
+        self.k.stop_gradient = False
+
+        self.cos = paddle.randn([1, 2048, 1, 96], dtype="float32")
+        self.cos.stop_gradient = False
+
+        self.sin = paddle.randn([1, 2048, 1, 96], dtype="float32")
+        self.sin.stop_gradient = False
+
+        self.position_ids = paddle.arange(end=2048, dtype="int64").unsqueeze(0)
+        self.position_ids.stop_gradient = False
+
+    def eval(self, use_cinn):
+        paddle.seed(2022)
+        net = RotaryPosEmb()
+        net.eval()
+        if use_cinn:
+            net = apply_to_static(net, use_cinn)
+
+        out = net(self.q, self.k, self.cos, self.sin, self.position_ids)
+        return out
+
+    def test_eval(self):
+        cinn_outs = self.eval(use_cinn=True)
+        # dy_outs = self.eval(use_cinn=False)
+
+        # TODO(phlrain): Need to check result
+        # for cinn_out, dy_out in zip(cinn_outs, dy_outs):
+        #     np.testing.assert_allclose(
+        #         cinn_out.numpy(), dy_out.numpy(), atol=1e-8
+        #     )
+
+
+if __name__ == '__main__':
+    unittest.main()
