【Infer Symbolic Shape No.243】【BUAA】warprnnt (#67567)

tlxd · web-flow · commit 88efb38d3a0d · 2024-08-22T17:15:43.000+08:00
* Update ops.yaml multiary_infer_sym.h multiary_infer_sym.cc

* Update multiary_infer_sym.cc

* Update multiary_infer_sym.cc

* Update multiary.cc
diff --git a/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/multiary_infer_sym.cc b/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/multiary_infer_sym.cc
@@ -2317,12 +2317,25 @@ bool ViterbiDecodeOpInferSymbolicShape(
 //   return true;
 // }
 
-// bool WarprnntOpInferSymbolicShape(pir::Operation *op,
-//                                   pir::InferSymbolicShapeContext
-//                                   *infer_context) {
-//   // pass
-//   return true;
-// }
+bool WarprnntOpInferSymbolicShape(
+    pir::Operation *op, pir::InferSymbolicShapeContext *infer_context) {
+  const auto &input_shape_or_data =
+      infer_context->GetShapeOrDataForValue(op->operand_source(0));
+  const std::vector<symbol::DimExpr> &input_shape = input_shape_or_data.shape();
+
+  std::vector<symbol::DimExpr> loss_shape = {input_shape[0]};
+  std::vector<symbol::DimExpr> warpctcgrad_shape = input_shape;
+
+  infer_context->SetShapeOrDataForValue(
+      op->result(0),
+      symbol::ShapeOrDataDimExprs{
+          symbol::TensorShapeOrDataDimExprs(loss_shape)});
+  infer_context->SetShapeOrDataForValue(
+      op->result(1),
+      symbol::ShapeOrDataDimExprs{
+          symbol::TensorShapeOrDataDimExprs(warpctcgrad_shape)});
+  return true;
+}
 
 // bool WeightOnlyLinearOpInferSymbolicShape(pir::Operation *op,
 //                                   pir::InferSymbolicShapeContext
diff --git a/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/multiary_infer_sym.h b/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/multiary_infer_sym.h
@@ -115,7 +115,7 @@ OP_DECLARE_INFER_SYMBOLIC_SHAPE(TrilinearInterp)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(HsigmoidLoss)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(ViterbiDecode)
 // OP_DECLARE_INFER_SYMBOLIC_SHAPE(Warpctc)
-// OP_DECLARE_INFER_SYMBOLIC_SHAPE(Warprnnt)
+OP_DECLARE_INFER_SYMBOLIC_SHAPE(Warprnnt)
 // OP_DECLARE_INFER_SYMBOLIC_SHAPE(WeightOnlyLinear)
 // OP_DECLARE_INFER_SYMBOLIC_SHAPE(WeightedSampleNeighbors)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Where)
diff --git a/paddle/phi/infermeta/multiary.cc b/paddle/phi/infermeta/multiary.cc
@@ -5444,8 +5444,8 @@ void WarprnntInferMeta(const MetaTensor& input,
                        float fastemit_lambda,
                        MetaTensor* loss,
                        MetaTensor* warpctcgrad) {
-  auto acts_dims = input.dims();
-  int D = static_cast<int>(acts_dims[3]);
+  auto input_dims = input.dims();
+  int D = static_cast<int>(input_dims[3]);
 
   PADDLE_ENFORCE_GE(
       blank,
@@ -5462,8 +5462,10 @@ void WarprnntInferMeta(const MetaTensor& input,
           "but received %d",
           blank));
 
-  loss->set_dims({-1});
+  loss->set_dims({input_dims[0]});
   loss->set_dtype(input.dtype());
+  warpctcgrad->set_dims(input_dims);
+  warpctcgrad->set_dtype(input.dtype());
 }
 
 void WeightOnlyLinearInferMeta(const MetaTensor& x,
diff --git a/paddle/phi/ops/yaml/ops.yaml b/paddle/phi/ops/yaml/ops.yaml
@@ -5071,7 +5071,7 @@
     data_type : input
   intermediate : warprnntgrad
   backward : warprnnt_grad
-  # interfaces : paddle::dialect::InferSymbolicShapeInterface
+  interfaces : paddle::dialect::InferSymbolicShapeInterface
 
 - op : weight_dequantize
   args : (Tensor x, Tensor scale, str algo = "weight_only_int8", DataType out_dtype = DataType::FLOAT16, int group_size = -1)
