shape dialect Broadcast

Author: jiahy0825

paddle/cinn/hlir/dialect/operator/transforms/CMakeLists.txt

@@ -29,4 +29,13 @@ if(NOT CINN_ONLY)
  cinn_op_dialect
  op_dialect_vjp)
 
+ cinn_cc_library(
+ fully_insert_broadcast_pass
+ SRCS
+ fully_insert_broadcast_pass.cc
+ DEPS
+ pir
+ cinn_op_dialect
+ op_dialect_vjp)
+
 endif()

paddle/fluid/pir/dialect/operator/ir/manual_op.cc

@@ -46,7 +46,6 @@ paddle::dialect::AddNOp, paddle::dialect::AddN_Op,
 #include "paddle/pir/core/builtin_op.h"
 #include "paddle/pir/core/builtin_type.h"
 #include "paddle/pir/core/ir_context.h"
-#include "paddle/phi/api/lib/data_type_set.h"
 
 namespace paddle {
 namespace dialect {
@@ -2685,89 +2684,6 @@ phi::DataType Increment_Op::GetKernelTypeForVar(
  return expected_kernel_dtype;
 }
 
-void ShapeBroadcastOp::Build(pir::Builder &builder, pir::OperationArgument &argument, pir::Value x_, pir::Value y_) {
- VLOG(4) << "Start build ShapeBroadcastOp";
-
-
-
- VLOG(4) << "Builder construction inputs";
- std::vector<pir::Value> argument_inputs = {x_, y_};
- argument.AddInputs(argument_inputs);
-
- VLOG(4) << "Builder construction attributes";
-
- VLOG(4) << "Builder construction outputs";
- paddle::dialect::DenseTensorType x = x_.type().dyn_cast<paddle::dialect::DenseTensorType>();
- paddle::dialect::DenseTensorType y = y_.type().dyn_cast<paddle::dialect::DenseTensorType>();
-
- VLOG(4) << "Builder construction dense_x";
- paddle::dialect::IrTensor ir_tensor_x(paddle::dialect::TransToPhiDataType(x.dtype()),
- x.dims(),
- x.data_layout(),
- x.lod(),
- x.offset());
- VLOG(4) << "Builder construction meta_x";
- paddle::dialect::IrMetaTensor meta_x(&ir_tensor_x);
-
- VLOG(4) << "Builder construction dense_y";
- paddle::dialect::IrTensor ir_tensor_y(paddle::dialect::TransToPhiDataType(y.dtype()),
- y.dims(),
- y.data_layout(),
- y.lod(),
- y.offset());
- VLOG(4) << "Builder construction meta_y";
- paddle::dialect::IrMetaTensor meta_y(&ir_tensor_y);
- paddle::dialect::IrTensor dense_out;
- paddle::dialect::IrMetaTensor meta_out(&dense_out);
-
- phi::ElementwiseInferMeta(meta_x, meta_y, &meta_out);
-
- std::vector<pir::Type> argument_outputs;
- pir::Type out_dense_tensor_type = paddle::dialect::DenseTensorType::get(pir::IrContext::Instance(), paddle::dialect::TransToIrDataType(dense_out.dtype()), dense_out.dims(), dense_out.layout(), dense_out.lod(), dense_out.offset());
- argument_outputs.push_back(out_dense_tensor_type);
- argument.AddOutputs(argument_outputs.begin(), argument_outputs.end());
- ::pir::PassStopGradientsDefaultly(argument);
-
-}
-
-namespace {
-
-void ShapeBroadcastOpInferMeta(const phi::MetaTensor& x,
- const phi::MetaTensor& y,
- phi::MetaTensor* out) {
- PADDLE_ENFORCE_EQ(x.dims().size(), 1);
- PADDLE_ENFORCE_EQ(y.dims().size(), 1);
- out->set_dims({std::max<int64_t>(x.dims().at(0), y.dims().at(0))});
- // dtype need promote when meet input dtype with more precision
- paddle::experimental::DataTypeSet dtype_set{x.dtype()};
- dtype_set = dtype_set | paddle::experimental::DataTypeSet(y.dtype());
- DataType promote_result = PromoteTypes(dtype_set);
- if (promote_result == DataType::UNDEFINED) {
- promote_result = x.dtype();
- }
- out->set_dtype(promote_result);
- out->set_layout(x.layout());
- out->share_lod(x);
-}
-
-}
-
-void ShapeBroadcastOp::InferMeta( phi::InferMetaContext *infer_meta ) {
- auto fn = PD_INFER_META(ShapeBroadcastOpInferMeta);
- fn(infer_meta);
-}
-
-
-phi::DataType ShapeBroadcastOp::GetKernelTypeForVar(
- const std::string& var_name,
- const phi::DataType& tensor_dtype,
- const phi::DataType& expected_kernel_dtype) {
- VLOG(4) << "Get KernelType for Var of op: ShapeBroadcastOp";
-
- return expected_kernel_dtype;
-}
-
-
 } // namespace dialect
 } // namespace paddle
 
@@ -2789,5 +2705,4 @@ IR_DEFINE_EXPLICIT_TYPE_ID(paddle::dialect::ExpandOp)
 IR_DEFINE_EXPLICIT_TYPE_ID(paddle::dialect::SelectInputOp)
 IR_DEFINE_EXPLICIT_TYPE_ID(paddle::dialect::IncrementOp)
 IR_DEFINE_EXPLICIT_TYPE_ID(paddle::dialect::Increment_Op)
-IR_DEFINE_EXPLICIT_TYPE_ID(paddle::dialect::ShapeBroadcastOp)
 #endif

paddle/fluid/pir/dialect/operator/ir/manual_op.h

@@ -515,30 +515,6 @@ class Increment_Op
  const std::vector<std::vector<bool>> &stop_gradients);
 };
 
-class IR_API ShapeBroadcastOp : public pir::Op<ShapeBroadcastOp,paddle::dialect::InferMetaInterface,paddle::dialect::GetKernelTypeForVarInterface> {
- public:
- using Op::Op;
- static const char *name() { return "pd_op.shape_broadcast"; }
- static constexpr const char **attributes_name = nullptr;
- static constexpr uint32_t attributes_num = 0;
- static void Build(pir::Builder &builder, pir::OperationArgument &argument, pir::Value x_, pir::Value y_);
- 
- void VerifySig() {}
-
-
- pir::Value x() { return operand_source(0); }
- pir::Value y() { return operand_source(1); }
- pir::OpResult out() { return result(0); }
-
- static void InferMeta(phi::InferMetaContext *infer_meta);
-
- static phi::DataType GetKernelTypeForVar(
- const std::string& var_name,
- const phi::DataType& tensor_dtype,
- const phi::DataType& expected_kernel_dtype);
-
-};
-
 } // namespace dialect
 } // namespace paddle
 
@@ -560,4 +536,3 @@ IR_DECLARE_EXPLICIT_TYPE_ID(paddle::dialect::ExpandOp)
 IR_DECLARE_EXPLICIT_TYPE_ID(paddle::dialect::SelectInputOp)
 IR_DECLARE_EXPLICIT_TYPE_ID(paddle::dialect::IncrementOp)
 IR_DECLARE_EXPLICIT_TYPE_ID(paddle::dialect::Increment_Op)
-IR_DECLARE_EXPLICIT_TYPE_ID(paddle::dialect::ShapeBroadcastOp)

paddle/pir/dialect/shape/ir/shape_op.cc

@@ -16,6 +16,8 @@
 #include "paddle/pir/core/builtin_attribute.h"
 #include "paddle/pir/core/builtin_op.h"
 #include "paddle/pir/core/builtin_type.h"
+#include "paddle/pir/dialect/shape/utils/dim_expr.h"
+#include "paddle/pir/dialect/shape/utils/shape_utils.h"
 
 namespace pir::shape {
 
@@ -363,6 +365,85 @@ void IndexCastOp::Build(Builder &builder, // NOLINT
  argument.output_types.emplace_back(out);
 }
 
+void ShapeBroadcastOp::Build(pir::Builder &builder,
+ pir::OperationArgument &argument,
+ pir::Value x_,
+ pir::Value y_) {
+ std::vector<pir::Value> argument_inputs = {x_, y_};
+ argument.AddInputs(argument_inputs);
+
+ IrContext *ctx = IrContext::Instance();
+ Type dtype = IndexType::get(ctx);
+ int64_t x_rank = x_.type()
+ .dyn_cast<DenseTensorType>()
+ .dyn_cast<ShapedTypeInterface>()
+ .GetRank();
+ int64_t y_rank = y_.type()
+ .dyn_cast<DenseTensorType>()
+ .dyn_cast<ShapedTypeInterface>()
+ .GetRank();
+ CHECK_EQ(x_rank, y_rank);
+ pir::DDim dims = {x_rank};
+ pir::DataLayout data_layout = pir::DataLayout::NCHW;
+ pir::LoD lod = {{0, 1, 2}};
+ size_t offset = 0;
+
+ argument.output_types.emplace_back(
+ DenseTensorType::get(ctx, dtype, dims, data_layout, lod, offset));
+}
+
+namespace {
+
+symbol::DimExpr GetBroadcastDimExpr(const symbol::DimExpr &lhs,
+ const symbol::DimExpr &rhs) {
+ if (lhs.isa<std::int64_t>() && rhs.isa<std::int64_t>()) {
+ CHECK_EQ(lhs.dyn_cast<std::int64_t>(), rhs.dyn_cast<std::int64_t>());
+ } else if (lhs.isa<std::int64_t>()) {
+ return lhs.dyn_cast<std::int64_t>() == 1 ? rhs : lhs;
+ } else if (rhs.isa<std::int64_t>()) {
+ return rhs.dyn_cast<std::int64_t>() == 1 ? lhs : rhs;
+ } else {
+ return symbol::Broadcast<symbol::DimExpr>{
+ symbol::List<symbol::DimExpr>{lhs, rhs}};
+ }
+}
+
+} // namespace
+
+bool ShapeBroadcastOp::InferSymbolicShape(
+ pir::ShapeConstraintIRAnalysis *shape_analysis) {
+ pir::Value x = operand_source(0);
+ pir::Value y = operand_source(1);
+ std::string x_id = pir::GetValueId(&x);
+ std::string y_id = pir::GetValueId(&y);
+
+ IR_ENFORCE(shape_analysis->value_id_to_shapeordata_.count(x_id) > 0,
+ "x_id does not exist.");
+ IR_ENFORCE(shape_analysis->value_id_to_shapeordata_.count(y_id) > 0,
+ "y_id does not exist.");
+ const auto &x_data_shape = shape_analysis->value_id_to_shapeordata_.at(x_id);
+ const auto &y_data_shape = shape_analysis->value_id_to_shapeordata_.at(y_id);
+ IR_ENFORCE(x_data_shape.data().has_value(),
+ "Value x comes from ShapeOp, it must have data");
+ IR_ENFORCE(y_data_shape.data().has_value(),
+ "Value y comes from ShapeOp, it must have data");
+ const auto &x_data = x_data_shape.data().value();
+ const auto &y_data = y_data_shape.data().value();
+ IR_ENFORCE(x_data.size() == y_data.size(), "Support same rank temporarily");
+
+ std::vector<symbol::DimExpr> output_data;
+ for (std::size_t i = 0; i < x_data.size(); ++i) {
+ output_data.emplace_back(GetBroadcastDimExpr(x_data.at(i), y_data.at(i)));
+ }
+
+ pir::OpResult res = result(0);
+ std::string res_id = pir::GetValueId(&res);
+ symbol::ShapeOrDataDimExprs output_data_shape =
+ symbol::ShapeOrDataDimExprs::MakeConsistentShapeOrData(output_data);
+ shape_analysis->value_id_to_shapeordata_[res_id] = output_data_shape;
+ return true;
+}
+
 } // namespace pir::shape
 
 IR_DEFINE_EXPLICIT_TYPE_ID(pir::shape::SymbolicDimOp)
@@ -376,3 +457,4 @@ IR_DEFINE_EXPLICIT_TYPE_ID(pir::shape::FromElementsOp)
 IR_DEFINE_EXPLICIT_TYPE_ID(pir::shape::ExtractOp);
 IR_DEFINE_EXPLICIT_TYPE_ID(pir::shape::ConstantIndexOp);
 IR_DEFINE_EXPLICIT_TYPE_ID(pir::shape::IndexCastOp);
+IR_DEFINE_EXPLICIT_TYPE_ID(pir::shape::ShapeBroadcastOp);

paddle/pir/dialect/shape/ir/shape_op.h

@@ -15,6 +15,7 @@
 #pragma once
 
 #include <optional>
+#include "paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape.h"
 #include "paddle/pir/core/builder.h"
 #include "paddle/pir/core/builtin_type_interfaces.h"
 #include "paddle/pir/core/ir_printer.h"
@@ -258,6 +259,27 @@ class IR_API IndexCastOp : public Op<IndexCastOp> {
  void VerifySig() {}
 };
 
+class IR_API ShapeBroadcastOp
+ : public Op<ShapeBroadcastOp,
+ paddle::dialect::InferSymbolicShapeInterface> {
+ public:
+ using Op::Op;
+ static const char *name() { return "shape.shape_broadcast"; }
+ static constexpr const char **attributes_name = nullptr;
+ static constexpr uint32_t attributes_num = 0;
+ static void Build(pir::Builder &builder, // NOLINT
+ pir::OperationArgument &argument, // NOLINT
+ pir::Value x_,
+ pir::Value y_);
+
+ pir::Value x() { return operand_source(0); }
+ pir::Value y() { return operand_source(1); }
+ pir::OpResult out() { return result(0); }
+ void VerifySig() {}
+
+ bool InferSymbolicShape(pir::ShapeConstraintIRAnalysis *shape_analysis);
+};
+
 } // namespace pir::shape
 
 IR_EXPORT_DECLARE_EXPLICIT_TYPE_ID(pir::shape::SymbolicDimOp);
@@ -271,3 +293,4 @@ IR_EXPORT_DECLARE_EXPLICIT_TYPE_ID(pir::shape::FromElementsOp);
 IR_EXPORT_DECLARE_EXPLICIT_TYPE_ID(pir::shape::ExtractOp);
 IR_EXPORT_DECLARE_EXPLICIT_TYPE_ID(pir::shape::ConstantIndexOp);
 IR_EXPORT_DECLARE_EXPLICIT_TYPE_ID(pir::shape::IndexCastOp);
+IR_EXPORT_DECLARE_EXPLICIT_TYPE_ID(pir::shape::ShapeBroadcastOp);