【Infer Symbolic Shape No.242】【BUAA】warpctc (#67539)

tlxd · web-flow · commit 7b4908029c84 · 2024-08-23T17:01:17.000+08:00
* Update ops.yaml multiary_infer_sym.cc multiary_infer_sym.h Warpctc Op + unary_infer_sym.cc OneHot Op + binary_infer_sym.cc Unpool3d Op

* Update multiary_infer_sym.cc binary_infer_sym.cc

* Update multiary_infer_sym.cc

* Update multiary_infer_sym.cc

* Update unary_infer_sym.cc unary.cc

* Update unary_infer_sym.cc unary.cc

* Update unary_infer_sym.cc

* Update multiary.cc + multiary_infer_sym.cc

* Update multiary.cc

* Update multiary.cc

* Update multiary.cc

* Update unary_infer_sym.cc one_hot op
diff --git a/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/binary_infer_sym.cc b/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/binary_infer_sym.cc
@@ -1474,7 +1474,7 @@ bool Unpool3dOpInferSymbolicShape(
       details::GetVectorAttr<int>(op, "output_size");
 
   std::vector<symbol::DimExpr> output_shape = {x_shape[0], x_shape[1]};
-  for (int i = 0; i < static_cast<int>(ksize.size()); ++i) {
+  for (size_t i = 0; i < ksize.size(); ++i) {
     output_shape.emplace_back(symbol::DimExpr(output_size[i]));
   }
 
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
@@ -2449,12 +2449,43 @@ bool ViterbiDecodeOpInferSymbolicShape(
   return true;
 }
 
-// bool WarpctcOpInferSymbolicShape(pir::Operation *op,
-//                                  pir::InferSymbolicShapeContext
-//                                  *infer_context) {
-//   // pass
-//   return true;
-// }
+bool WarpctcOpInferSymbolicShape(
+    pir::Operation *op, pir::InferSymbolicShapeContext *infer_context) {
+  const auto &logits_shape_or_data =
+      infer_context->GetShapeOrDataForValue(op->operand_source(0));
+  const std::vector<symbol::DimExpr> &logits_shape =
+      logits_shape_or_data.shape();
+
+  symbol::DimExpr max_sequence_length, num_sequences;
+  symbol::DimExpr sequence_width = symbol::DimExpr(1);
+
+  if (op->operand_source(2) && op->operand_source(3)) {
+    max_sequence_length = logits_shape[0];
+    num_sequences = logits_shape[1];
+    sequence_width = logits_shape[2];
+  } else {
+    max_sequence_length = infer_context->GetNextSymName();
+    num_sequences = infer_context->GetNextSymName();
+    for (size_t i = 1; i < logits_shape.size(); ++i) {
+      sequence_width = sequence_width * logits_shape[i];
+    }
+  }
+
+  std::vector<symbol::DimExpr> loss_shape = {num_sequences, symbol::DimExpr(1)};
+  std::vector<symbol::DimExpr> warpctcgrad_shape = {
+      max_sequence_length, num_sequences, sequence_width};
+
+  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 WarprnntOpInferSymbolicShape(
     pir::Operation *op, pir::InferSymbolicShapeContext *infer_context) {
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
@@ -114,7 +114,7 @@ OP_DECLARE_INFER_SYMBOLIC_SHAPE(TdmSampler)
 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(Warpctc)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Warprnnt)
 // OP_DECLARE_INFER_SYMBOLIC_SHAPE(WeightOnlyLinear)
 // OP_DECLARE_INFER_SYMBOLIC_SHAPE(WeightedSampleNeighbors)
diff --git a/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/unary_infer_sym.cc b/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/unary_infer_sym.cc
@@ -1915,28 +1915,25 @@ bool OneHotOpInferSymbolicShape(pir::Operation *op,
   const auto &x_shape_or_data =
       infer_context->GetShapeOrDataForValue(op->operand_source(0));
   const std::vector<symbol::DimExpr> &x_shape = x_shape_or_data.shape();
+  const auto &num_classes_shape_or_date =
+      infer_context->GetShapeOrDataForValue(op->operand_source(1));
 
   const auto &attributes = op->attributes();
   int64_t num_classes;
-  if (op->operand_source(1)) {
-    const auto &num_classes_shape_or_date =
-        infer_context->GetShapeOrDataForValue(op->operand_source(1));
-    if (attributes.find("num_classes") != attributes.end()) {
-      num_classes = op->attribute<pir::Int64Attribute>("num_classes").data();
-    } else if (num_classes_shape_or_date.data().has_value()) {
-      num_classes =
-          num_classes_shape_or_date.data().value().at(0).Get<int64_t>();
-    } else {
-      PADDLE_ENFORCE_EQ(!num_classes_shape_or_date.data().has_value(),
-                        true,
-                        common::errors::InvalidArgument(
-                            "The depth should have data! Please check."));
-    }
+  symbol::DimExpr num_classes_expr;
+
+  if (attributes.find("num_classes") != attributes.end()) {
+    num_classes = op->attribute<pir::Int64Attribute>("num_classes").data();
+    num_classes_expr = symbol::DimExpr(num_classes);
+  } else if (num_classes_shape_or_date.data().has_value()) {
+    num_classes_expr = num_classes_shape_or_date.data().value().at(0);
+  } else {
+    num_classes_expr = infer_context->GetNextSymName();
   }
 
   const std::vector<symbol::DimExpr> &out_shape = [&] {
     std::vector<symbol::DimExpr> out_shape = x_shape;
-    out_shape.push_back(symbol::DimExpr(num_classes));
+    out_shape.push_back(num_classes_expr);
     return out_shape;
   }();
 
diff --git a/paddle/phi/infermeta/multiary.cc b/paddle/phi/infermeta/multiary.cc
@@ -5408,11 +5408,15 @@ void WarpctcInferMeta(const MetaTensor& logits,
                       MetaTensor* loss,
                       MetaTensor* warpctcgrad) {
   auto logits_dims = logits.dims();
-  int sequence_width = 0;
+  int num_sequences, sequence_width, max_sequence_length;
 
-  if (logits_length) {
+  if (logits_length && labels_length) {
+    max_sequence_length = static_cast<int>(logits_dims[0]);
+    num_sequences = static_cast<int>(logits_dims[1]);
     sequence_width = static_cast<int>(logits_dims[2]);
   } else {
+    max_sequence_length = -1;
+    num_sequences = -1;
     sequence_width =
         static_cast<int>(common::product(logits_dims) / logits_dims[0]);
   }
@@ -5432,8 +5436,10 @@ void WarpctcInferMeta(const MetaTensor& logits,
           "but received %d",
           blank));
 
-  loss->set_dims({-1, 1});
+  loss->set_dims({num_sequences, 1});
   loss->set_dtype(logits.dtype());
+  warpctcgrad->set_dims({max_sequence_length, num_sequences, sequence_width});
+  warpctcgrad->set_dtype(logits.dtype());
 }
 
 void WarprnntInferMeta(const MetaTensor& input,
diff --git a/paddle/phi/infermeta/unary.cc b/paddle/phi/infermeta/unary.cc
@@ -2914,21 +2914,21 @@ void OneHotRawInferMeta(const MetaTensor& x,
 }
 
 void OneHotInferMeta(const MetaTensor& x,
-                     const Scalar& depth_t,
+                     const Scalar& num_classes,
                      MetaTensor* out) {
   auto x_dims = x.dims();
   PADDLE_ENFORCE_GE(x_dims.size(),
                     0,
                     common::errors::InvalidArgument(
                         "Rank of Input(X) should be at least 0."));
 
-  int depth = depth_t.to<int>();
+  int num_classes_int = num_classes.to<int>();
   auto out_dims_vec = common::vectorize(x_dims);
-  out_dims_vec.push_back(depth);
+  out_dims_vec.push_back(num_classes_int);
   auto out_dims = common::make_ddim(out_dims_vec);
+
   out->set_dims(out_dims);
   out->share_lod(x);
-
   out->set_dtype(phi::DataType::FLOAT32);
 }
 
diff --git a/paddle/phi/ops/yaml/ops.yaml b/paddle/phi/ops/yaml/ops.yaml
@@ -5062,7 +5062,7 @@
   optional : logits_length, labels_length
   intermediate : warpctcgrad
   backward : warpctc_grad
-  # interfaces : paddle::dialect::InferSymbolicShapeInterface
+  interfaces : paddle::dialect::InferSymbolicShapeInterface
 
 - op : warprnnt
   args : (Tensor input, Tensor label, Tensor input_lengths, Tensor label_lengths, int blank = 0, float fastemit_lambda = 0.0)

Original file line number	Diff line number	Diff line change
`@@ -1474,7 +1474,7 @@ bool Unpool3dOpInferSymbolicShape(`
`1474`	`1474`	`details::GetVectorAttr<int>(op, "output_size");`
`1475`	`1475`
`1476`	`1476`	`std::vector<symbol::DimExpr> output_shape = {x_shape[0], x_shape[1]};`
`1477`		`- for (int i = 0; i < static_cast<int>(ksize.size()); ++i) {`
	`1477`	`+ for (size_t i = 0; i < ksize.size(); ++i) {`
`1478`	`1478`	`output_shape.emplace_back(symbol::DimExpr(output_size[i]));`
`1479`	`1479`	`}`
`1480`	`1480`