[PRIM][PIR]Migrate prim rules

paddle/fluid/primitive/codegen/gen.py

@@ -72,8 +72,20 @@
 ]
 
 
-PRIM_VJP = ['divide_grad', 'sum_grad']  # vjp list of primitive op
-CUSTOM_VJP = ['gelu_grad']  # custom vjp list of composite op
+PRIM_VJP = [
+    'divide_grad',
+    'sum_grad',
+    'cast_grad',
+    'add_grad',
+    'multiply_grad',
+    'elementwise_pow_grad',
+    'reshape_grad',
+    'split_grad',
+    'tanh_grad',
+    'transpose_grad',
+    'concat_grad',
+]  # vjp list of primitive op
+CUSTOM_VJP = ['gelu_grad', 'layer_norm_grad']  # custom vjp list of composite op
 VJP_COMPS = PRIM_VJP + CUSTOM_VJP
 
 BACKENDS = [
@@ -149,6 +161,8 @@
     'embedding_grad',
     'sqrt',
     'uniform',
+    'split',
+    'transpose',
 ]
 
 

paddle/fluid/primitive/codegen/templates/rule/vjp/generated/generated_vjp.cc.j2

@@ -106,7 +106,7 @@ paddle::Tensor* {{api.outputs[i].name}} = !stop_gradients[{{i}}][0] ? &vjp_res[{
     {% else %}
 std::vector<paddle::Tensor*> {{api.outputs[i].name}}(stop_gradients[{{i}}].size(), nullptr);
 for (size_t i=0; i< stop_gradients[{{i}}].size(); i++ ) {
-  {{api.outputs[i].name}} =  !stop_gradients[{{i}}][i] ?  &vjp_res[{{i}}][i] : nullptr;
+  {{api.outputs[i].name}}[i] =  !stop_gradients[{{i}}][i] ?  &vjp_res[{{i}}][i] : nullptr;
 }
     {% endif %}
   {% endfor %}

paddle/fluid/primitive/primitive.yaml

@@ -49,3 +49,4 @@
 - erf
 - tanh
 - full
+- cast