[Infer Symbolic Shape No.245,247] Add symbol_infer_interface for weight_{dequantize.quantize} (#68979)

ooooo-create · web-flow · commit e97b2e1d2d8f · 2024-11-13T14:11:48.000+08:00
* Add symbol_infer_interface for weight_{dequantize.quantize}

* Fix

* Fix

* Refine logic

* revert irrelevant change
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
@@ -2374,12 +2374,60 @@ bool UnpoolOpInferSymbolicShape(pir::Operation *op,
   return true;
 }
 
-// bool WeightDequantizeOpInferSymbolicShape(pir::Operation *op,
-//                                           pir::InferSymbolicShapeContext
-//                                           *infer_context) {
-//   // pass
-//   return true;
-// }
+bool WeightDequantizeOpInferSymbolicShape(
+    pir::Operation *op, pir::InferSymbolicShapeContext *infer_context) {
+  const auto &x_shape =
+      infer_context->GetShapeOrDataForValue(op->operand_source(0)).shape();
+  const auto &scale_shape =
+      infer_context->GetShapeOrDataForValue(op->operand_source(1)).shape();
+  PADDLE_ENFORCE_EQ(x_shape.size(),
+                    2UL,
+                    common::errors::InvalidArgument(
+                        "The x tensor of dequantize op must be 2D, but got[%d]",
+                        x_shape.size()));
+  int group_size = op->attribute<pir::Int32Attribute>("group_size").data();
+  std::string algo = op->attribute<pir::StrAttribute>("algo").AsString();
+  PADDLE_ENFORCE_EQ(
+      (group_size == -1 || group_size == 64 || group_size == 128),
+      true,
+      common::errors::InvalidArgument("group_size must be -1, 64 or 128."));
+  symbol::DimExpr real_channel_shape;
+  if (algo == "weight_only_int8") {
+    real_channel_shape = x_shape[0];
+  } else if (algo == "weight_only_int4") {
+    real_channel_shape = x_shape[0] * 2;
+  } else {
+    PADDLE_THROW(common::errors::InvalidArgument(
+        "The algo must be in ['weight_only_int8', 'weight_only_int4'], "
+        "but got[%s]",
+        algo));
+  }
+  if (group_size == -1) {
+    PADDLE_ENFORCE_EQ(scale_shape.size(),
+                      1UL,
+                      common::errors::InvalidArgument(
+                          "The scale tensor of dequantize op must "
+                          "be 1D in per-channel mode, but got[%d]",
+                          scale_shape.size()));
+    infer_context->AddEqualCstr(scale_shape[0], real_channel_shape);
+  } else {
+    PADDLE_ENFORCE_EQ(scale_shape.size(),
+                      2UL,
+                      common::errors::InvalidArgument(
+                          "The scale tensor of dequantize op must "
+                          "be 2D in group-wise mode, but got[%d]",
+                          scale_shape.size()));
+    infer_context->AddEqualCstr(scale_shape[0],
+                                (x_shape[1] + (group_size - 1)) / group_size);
+    infer_context->AddEqualCstr(scale_shape[1], real_channel_shape);
+  }
+  std::vector<symbol::DimExpr> out_shape{x_shape[1], real_channel_shape};
+  infer_context->SetShapeOrDataForValue(
+      op->result(0),
+      symbol::ShapeOrDataDimExprs{
+          symbol::TensorShapeOrDataDimExprs(out_shape)});
+  return true;
+}
 
 bool YoloBoxOpInferSymbolicShape(
     pir::Operation *op, pir::InferSymbolicShapeContext *infer_context) {
diff --git a/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/binary_infer_sym.h b/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/binary_infer_sym.h
@@ -98,7 +98,7 @@ OP_DECLARE_INFER_SYMBOLIC_SHAPE(TdmChild)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(TriangularSolve)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Unpool3d)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Unpool)
-// OP_DECLARE_INFER_SYMBOLIC_SHAPE(WeightDequantize)
+OP_DECLARE_INFER_SYMBOLIC_SHAPE(WeightDequantize)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(YoloBox)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(YoloBoxHead)
 
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
@@ -4426,10 +4426,70 @@ bool UnstackOpInferSymbolicShape(
   return true;
 }
 
-// bool WeightQuantizeOpInferSymbolicShape(
-//     pir::Operation *op, pir::InferSymbolicShapeContext *infer_context) {
-//   // pass
-//   return true;
-// }
+bool WeightQuantizeOpInferSymbolicShape(
+    pir::Operation *op, pir::InferSymbolicShapeContext *infer_context) {
+  const auto &x_shape =
+      infer_context->GetShapeOrDataForValue(op->operand_source(0)).shape();
+  PADDLE_ENFORCE_EQ(
+      x_shape.size(),
+      2UL,
+      common::errors::InvalidArgument(
+          "The x tensor of quant op must be 2D, but got[%d]", x_shape.size()));
+  if (x_shape[0].isa<int64_t>()) {
+    int64_t x_shape_0 = x_shape[0].dyn_cast<int64_t>();
+    PADDLE_ENFORCE_EQ(
+        x_shape_0 % 64,
+        0,
+        common::errors::InvalidArgument(
+            "The first dimension of input must be divisible by 64, but got[%d]",
+            x_shape_0));
+  }
+  if (x_shape[1].isa<int64_t>()) {
+    int64_t x_shape_1 = x_shape[1].dyn_cast<int64_t>();
+    PADDLE_ENFORCE_EQ(
+        x_shape_1 % 16,
+        0,
+        common::errors::InvalidArgument("The second dimension of input must be "
+                                        "divisible by 16, but got[%d]",
+                                        x_shape_1));
+  }
+
+  int group_size = op->attribute<pir::Int32Attribute>("group_size").data();
+  std::string algo = op->attribute<pir::StrAttribute>("algo").AsString();
+  PADDLE_ENFORCE_EQ(
+      ((group_size == -1) || (group_size == 64) || (group_size == 128)),
+      true,
+      common::errors::InvalidArgument(
+          "Currently, group_size only support -1, 64 or 128."));
+  std::vector<symbol::DimExpr> scale_shape;
+  std::vector<symbol::DimExpr> out_shape;
+  if (group_size != -1) {
+    symbol::DimExpr scale_shape_0 =
+        (x_shape[0] + (group_size - 1)) / group_size;
+    scale_shape = {scale_shape_0, x_shape[1]};
+  } else {
+    scale_shape = {x_shape[1]};
+  }
+  if (algo == "weight_only_int8" || algo == "llm.int8") {
+    out_shape = {x_shape[1], x_shape[0]};
+  } else if (algo == "weight_only_int4") {
+    out_shape = {x_shape[1] / 2, x_shape[0]};
+  } else {
+    PADDLE_THROW(common::errors::InvalidArgument(
+        "The algo must be in ['weight_only_int8', 'weight_only_int4', "
+        "'llm.int8'], but got[%s]",
+        algo));
+  }
+  infer_context->SetShapeOrDataForValue(
+      op->result(0),
+      symbol::ShapeOrDataDimExprs{
+          symbol::TensorShapeOrDataDimExprs(out_shape)});
+
+  infer_context->SetShapeOrDataForValue(
+      op->result(1),
+      symbol::ShapeOrDataDimExprs{
+          symbol::TensorShapeOrDataDimExprs(scale_shape)});
+  return true;
+}
 
 }  // namespace paddle::dialect
diff --git a/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/unary_infer_sym.h b/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/unary_infer_sym.h
@@ -164,6 +164,6 @@ OP_DECLARE_INFER_SYMBOLIC_SHAPE(Unsqueeze)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Unsqueeze_)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Unfold)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Unstack)
-// OP_DECLARE_INFER_SYMBOLIC_SHAPE(WeightQuantize)
+OP_DECLARE_INFER_SYMBOLIC_SHAPE(WeightQuantize)
 
 }  // namespace paddle::dialect
diff --git a/paddle/phi/infermeta/unary.cc b/paddle/phi/infermeta/unary.cc
@@ -5961,10 +5961,10 @@ void WeightQuantizeInferMeta(const MetaTensor& x,
   } else if (algo == "weight_only_int4") {
     dim_out = std::vector<int64_t>({x_dims[1] / 2, x_dims[0]});
   } else {
-    common::errors::InvalidArgument(
+    PADDLE_THROW(common::errors::InvalidArgument(
         "The algo must be in ['weight_only_int8', 'weight_only_int4', "
         "'llm.int8'], but got[%s]",
-        algo);
+        algo));
   }
   out->set_dims(common::make_ddim(dim_out));
 
diff --git a/paddle/phi/ops/yaml/ops.yaml b/paddle/phi/ops/yaml/ops.yaml
@@ -5332,7 +5332,7 @@
     func : weight_dequantize
     data_type : out_dtype
   traits : paddle::dialect::ForwardOnlyTrait
-  # interfaces : paddle::dialect::InferSymbolicShapeInterface
+  interfaces : paddle::dialect::InferSymbolicShapeInterface
 
 - op : weight_only_linear
   args : (Tensor x, Tensor weight, Tensor bias, Tensor weight_scale, str weight_dtype, int arch = 80, int group_size = -1)
@@ -5356,7 +5356,7 @@
     data_type : x
     backend : x
   traits : paddle::dialect::ForwardOnlyTrait
-  # interfaces : paddle::dialect::InferSymbolicShapeInterface
+  interfaces : paddle::dialect::InferSymbolicShapeInterface
 
 - op : weighted_sample_neighbors
   args : (Tensor row, Tensor colptr, Tensor edge_weight, Tensor input_nodes, Tensor eids, int sample_size, bool return_eids)
