[CINN]add CacheGradOpSymbolicShapeInterface (#65343)

Hongqing-work · web-flow · commit a93c2911f9d7 · 2024-06-27T10:37:19.000+08:00
* [CINN]add CacheGradOpSymbolicShapeInterface

* fix bug

* split func

* fix

* fix error msg

* fix name
diff --git a/paddle/cinn/hlir/dialect/operator/transforms/add_cinn_pass.cc b/paddle/cinn/hlir/dialect/operator/transforms/add_cinn_pass.cc
@@ -88,6 +88,18 @@ bool HasDynamicShape(const pir::Program& program) {
 }
 }  // namespace
 
+void ApplyShapeOptimizationPass(
+    ::pir::Program* program,
+    const std::function<std::shared_ptr<::pir::PassManager>()>&
+        CreatePassManager) {
+  std::shared_ptr<pir::PassManager> pass_manager = CreatePassManager();
+  bool has_dynamic_shape = HasDynamicShape(*program);
+  if (has_dynamic_shape) {
+    pass_manager->AddPass(pir::CreateShapeOptimizationPass());
+  }
+  pass_manager->Run(program);
+}
+
 void ApplyPdToCinnPass(
     ::pir::Program* program,
     const std::function<std::shared_ptr<::pir::PassManager>()>&
@@ -115,7 +127,6 @@ void ApplyCinnPreprocessPass(
   bool has_dynamic_shape = HasDynamicShape(*program);
 
   if (has_dynamic_shape) {
-    pass_manager->AddPass(pir::CreateShapeOptimizationPass());
     pass_manager->AddPass(
         cinn::dialect::ir::CreateFuseShapeOpsIntoGenerateShapeOpPass());
     pass_manager->AddPass(pir::CreateDeadCodeEliminationPass());
@@ -235,6 +246,7 @@ void ApplyCinnPass(::pir::Program* program,
       .file_name("original_programs.py")
       .dump_symbolic_shape(FLAGS_logging_pir_py_code_dump_symbolic_dims)
       .SaveIfFlagEnabled();
+  ApplyShapeOptimizationPass(program, CreatePassManager);
   ApplyPdToCinnPass(program, CreatePassManager);
   ApplyCinnPreprocessPass(program, CreatePassManager);
   ApplyBuildGroupOpPass(program, CreatePassManager);
diff --git a/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/cache_grad_op_symbolic_shape.h b/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/cache_grad_op_symbolic_shape.h
@@ -0,0 +1,23 @@
+// Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     https://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#pragma once
+
+#include "paddle/pir/include/dialect/shape/interface/infer_symbolic_shape/cache_grad_op_symbolic_shape.h"
+
+namespace paddle::dialect {
+using CacheGradOpSymbolicShapeInterface =
+    pir::CacheGradOpSymbolicShapeInterface;
+
+}  // namespace paddle::dialect
diff --git a/paddle/pir/include/dialect/shape/interface/infer_symbolic_shape/cache_grad_op_symbolic_shape.h b/paddle/pir/include/dialect/shape/interface/infer_symbolic_shape/cache_grad_op_symbolic_shape.h
@@ -0,0 +1,61 @@
+// Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     https://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#pragma once
+
+#include "paddle/pir/include/dialect/shape/utils/shape_analysis.h"
+
+// The CacheGradOpSymbolicShapeInterface is used by forward operators to cache
+// the symbolic shape information of the corresponding backward operators in the
+// computational graph. With this interface, most backward operators do not need
+// to implement the `InferSymbolicShapeInterface`.
+
+namespace pir {
+
+class CacheGradOpSymbolicShapeInterface
+    : public pir::OpInterfaceBase<CacheGradOpSymbolicShapeInterface> {
+ public:
+  /// Defined these methods with the interface.
+  struct Concept {
+    explicit Concept(bool (*cache_grad_op_symbolic_shape)(
+        pir::Operation *op, pir::InferSymbolicShapeContext *infer_context))
+        : cache_grad_op_symbolic_shape(cache_grad_op_symbolic_shape) {}
+    bool (*cache_grad_op_symbolic_shape)(
+        pir::Operation *op, pir::InferSymbolicShapeContext *infer_context);
+  };
+
+  template <class ConcreteOp>
+  struct Model : public Concept {
+    static inline bool CacheGradOpSymbolicShape(
+        pir::Operation *op, pir::InferSymbolicShapeContext *infer_context) {
+      return op->dyn_cast<ConcreteOp>().CacheGradOpSymbolicShape(infer_context);
+    }
+
+    Model() : Concept(CacheGradOpSymbolicShape) {}
+  };
+
+  /// Constructor
+  CacheGradOpSymbolicShapeInterface(pir::Operation *op, Concept *impl)
+      : pir::OpInterfaceBase<CacheGradOpSymbolicShapeInterface>(op),
+        impl_(impl) {}
+
+  bool CacheGradOpSymbolicShape(pir::InferSymbolicShapeContext *infer_context);
+
+ private:
+  Concept *impl_;
+};
+
+}  // namespace pir
+
+IR_DECLARE_EXPLICIT_TYPE_ID(pir::CacheGradOpSymbolicShapeInterface)
diff --git a/paddle/pir/src/dialect/shape/interface/infer_symbolic_shape/cache_grad_op_symbolic_shape.cc b/paddle/pir/src/dialect/shape/interface/infer_symbolic_shape/cache_grad_op_symbolic_shape.cc
@@ -0,0 +1,32 @@
+// Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     https://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "paddle/pir/include/dialect/shape/interface/infer_symbolic_shape/cache_grad_op_symbolic_shape.h"
+
+// This file implements the cache_grad_op_symbolic_shape interface for both
+// paddle and cinn operators.
+
+// Add `interfaces : pir::CacheGradOpSymbolicShapeInterface` in relative
+// yaml file to corresponding op.
+
+namespace pir {
+
+bool CacheGradOpSymbolicShapeInterface::CacheGradOpSymbolicShape(
+    pir::InferSymbolicShapeContext *infer_context) {
+  return impl_->cache_grad_op_symbolic_shape(operation(), infer_context);
+}
+
+}  // namespace pir
+
+IR_DEFINE_EXPLICIT_TYPE_ID(pir::CacheGradOpSymbolicShapeInterface)
diff --git a/paddle/pir/src/dialect/shape/transforms/shape_optimization_pass.cc b/paddle/pir/src/dialect/shape/transforms/shape_optimization_pass.cc
@@ -14,10 +14,12 @@
 
 #include "paddle/pir/include/dialect/shape/transforms/shape_optimization_pass.h"
 
+#include "paddle/common/errors.h"
 #include "paddle/common/flags.h"
 #include "paddle/pir/include/core/builtin_type.h"
 #include "paddle/pir/include/core/dialect.h"
 #include "paddle/pir/include/core/ir_printer.h"
+#include "paddle/pir/include/dialect/shape/interface/infer_symbolic_shape/cache_grad_op_symbolic_shape.h"
 #include "paddle/pir/include/dialect/shape/interface/infer_symbolic_shape/infer_symbolic_shape.h"
 #include "paddle/pir/include/dialect/shape/ir/shape_attribute.h"
 #include "paddle/pir/include/dialect/shape/ir/shape_dialect.h"
@@ -243,105 +245,131 @@ class ShapeOptimizationPass : public pir::Pass {
 
 }  // namespace
 
-void InferSymExprForBlock(const Block& block,
-                          InferSymbolicShapeContext* infer_context) {
-  for (auto& op : block) {
-    std::vector<symbol::ShapeOrDataDimExprs> input_shape_or_data;
-    for (auto& input : op.operands_source()) {
-      input_shape_or_data.emplace_back(
-          infer_context->GetShapeOrDataForValue(input));
+void InferSymExprForOp(Operation* op,
+                       InferSymbolicShapeContext* infer_context,
+                       const InferSymbolicShapeCacheKey& op_infer_cache_key) {
+  auto infer_symbolic_shape_interface =
+      op->dyn_cast<pir::InferSymbolicShapeInterface>();
+  if (infer_symbolic_shape_interface) {
+    PrintOpInfo(op);
+    PADDLE_ENFORCE_EQ(
+        infer_symbolic_shape_interface.InferSymbolicShape(infer_context),
+        true,
+        common::errors::Fatal("InferSymbolicShape for %s failed.", op->name()));
+
+    if (op->num_results() > 0) {
+      // TODO(lanxianghit): deal with the ops which have more than 1
+      // ACTUAL results
+      pir::shape::SetShapeAttrForOp(
+          op, infer_context->GetShapeOrDataForValue(op->result(0)));
     }
-    InferSymbolicShapeCacheKey op_infer_cache_key(op, input_shape_or_data);
-    auto infer_symbolic_shape_interface =
-        op.dyn_cast<pir::InferSymbolicShapeInterface>();
-    if (infer_symbolic_shape_interface) {
-      PrintOpInfo(&op);
-      PADDLE_ENFORCE_EQ(
-          infer_symbolic_shape_interface.InferSymbolicShape(infer_context),
-          true,
-          "InferSymbolicShape for %s failed.",
-          op.name());
-
-      if (op.num_results() > 0) {
-        // TODO(lanxianghit): deal with the ops which have more than 1
-        // ACTUAL results
-        pir::shape::SetShapeAttrForOp(
-            &op, infer_context->GetShapeOrDataForValue(op.result(0)));
-      }
-    } else {
-      LOG(WARNING) << op.name()
-                   << " DOES NOT have InferSymbolicShapeInterface!";
-      const bool all_outs_static_dims = [&] {
-        bool all_static_dims = true;
-        for (uint32_t i = 0; i < op.num_results(); ++i) {
-          if (IsStaticShape(op.result(i))) {
-            continue;
-          } else {
-            all_static_dims = false;
-            break;
-          }
-        }
-        return all_static_dims;
-      }();
-
-      if (all_outs_static_dims) {
-        for (uint32_t i = 0; i < op.num_results(); ++i) {
-          infer_context->SetSymbolForValueByStaticShape(op.result(i));
-        }
-      } else {
-        if (infer_context->GetOpInferSymbolicShapeCache(op_infer_cache_key)
-                .has_value()) {
-          std::vector<symbol::ShapeOrDataDimExprs> cached_result_shape_or_data =
-              infer_context->GetOpInferSymbolicShapeCache(op_infer_cache_key)
-                  .value();
-          CHECK(cached_result_shape_or_data.size() == op.num_results());
-          for (uint32_t i = 0; i < op.num_results(); ++i) {
-            infer_context->SetShapeOrDataForValue(
-                op.result(i), cached_result_shape_or_data[i]);
-          }
+  } else {
+    LOG(WARNING) << op->name() << " DOES NOT have InferSymbolicShapeInterface!";
+    const bool all_outs_static_dims = [&] {
+      bool all_static_dims = true;
+      for (uint32_t i = 0; i < op->num_results(); ++i) {
+        if (IsStaticShape(op->result(i))) {
+          continue;
         } else {
-          // risk set
-          for (uint32_t i = 0; i < op.num_results(); ++i) {
-            infer_context->SetSymbolForValueByStaticShape(op.result(i));
-          }
+          all_static_dims = false;
+          break;
         }
       }
-    }
+      return all_static_dims;
+    }();
 
-    // set infer_context share cache
-    const auto& SetOpInferSymbolicShapeCache = [&]() {
-      std::vector<symbol::ShapeOrDataDimExprs> result_shape_or_data;
-      for (auto& result : op.results()) {
-        result_shape_or_data.emplace_back(
-            infer_context->GetShapeOrDataForValue(result));
+    if (all_outs_static_dims) {
+      for (uint32_t i = 0; i < op->num_results(); ++i) {
+        infer_context->SetSymbolForValueByStaticShape(op->result(i));
       }
+    } else {
       if (infer_context->GetOpInferSymbolicShapeCache(op_infer_cache_key)
               .has_value()) {
         std::vector<symbol::ShapeOrDataDimExprs> cached_result_shape_or_data =
             infer_context->GetOpInferSymbolicShapeCache(op_infer_cache_key)
                 .value();
-        // check whether the result_shape_or_data is consistent with the cached
-        // TODO(Hongqing-work): delete check and only set cache for op without
-        // InferSymbolicShapeInterface after fixing all warnings.
-        if (cached_result_shape_or_data.size() != result_shape_or_data.size()) {
+        CHECK(cached_result_shape_or_data.size() == op->num_results());
+        for (uint32_t i = 0; i < op->num_results(); ++i) {
+          infer_context->SetShapeOrDataForValue(op->result(i),
+                                                cached_result_shape_or_data[i]);
+        }
+      } else {
+        // risk set
+        for (uint32_t i = 0; i < op->num_results(); ++i) {
+          infer_context->SetSymbolForValueByStaticShape(op->result(i));
+        }
+      }
+    }
+  }
+}
+
+void CacheForwardOpSymbolicShape(
+    Operation* op,
+    InferSymbolicShapeContext* infer_context,
+    const InferSymbolicShapeCacheKey& op_infer_cache_key) {
+  std::vector<symbol::ShapeOrDataDimExprs> result_shape_or_data;
+  const auto& CheckInferSymbolicShapeCacheConsistency =
+      [&](const InferSymbolicShapeCacheValue& infer_result,
+          const InferSymbolicShapeCacheValue& cache_result) {
+        if (infer_result.size() != cache_result.size()) {
           LOG(WARNING) << "cached shape is not consistent with real shape";
         } else {
-          for (uint32_t i = 0; i < op.num_results(); ++i) {
-            if (cached_result_shape_or_data[i] != result_shape_or_data[i]) {
+          for (uint32_t i = 0; i < cache_result.size(); ++i) {
+            if (infer_result[i] != cache_result[i]) {
               LOG(WARNING) << "cached shape is not consistent with real shape";
               VLOG(3) << "InferSymbolicShapeCacheKey is: "
                       << op_infer_cache_key;
-              VLOG(3) << "cached shape is: " << cached_result_shape_or_data[i];
-              VLOG(3) << "real shape is: " << result_shape_or_data[i];
+              VLOG(3) << "cached shape is: " << cache_result[i];
+              VLOG(3) << "real shape is: " << infer_result[i];
             }
           }
         }
-      } else {
-        infer_context->SetOpInferSymbolicShapeCache(op_infer_cache_key,
-                                                    result_shape_or_data);
-      }
-    };
-    SetOpInferSymbolicShapeCache();
+      };
+  for (const auto& result : op->results()) {
+    result_shape_or_data.emplace_back(
+        infer_context->GetShapeOrDataForValue(result));
+  }
+  if (infer_context->GetOpInferSymbolicShapeCache(op_infer_cache_key)
+          .has_value()) {
+    std::vector<symbol::ShapeOrDataDimExprs> cached_result_shape_or_data =
+        infer_context->GetOpInferSymbolicShapeCache(op_infer_cache_key).value();
+    // TODO(Hongqing-work): delete check and only set cache for op without
+    // InferSymbolicShapeInterface after fixing all warnings.
+    CheckInferSymbolicShapeCacheConsistency(result_shape_or_data,
+                                            cached_result_shape_or_data);
+  } else {
+    infer_context->SetOpInferSymbolicShapeCache(op_infer_cache_key,
+                                                result_shape_or_data);
+  }
+}
+
+void CacheBackwardOpSymbolicShape(Operation* op,
+                                  InferSymbolicShapeContext* infer_context) {
+  auto cache_grad_op_symbolic_shape_interface =
+      op->dyn_cast<pir::CacheGradOpSymbolicShapeInterface>();
+  if (cache_grad_op_symbolic_shape_interface) {
+    VLOG(3) << "CacheBackwardOpSymbolicShape for: " << op->name();
+    PADDLE_ENFORCE_EQ(
+        cache_grad_op_symbolic_shape_interface.CacheGradOpSymbolicShape(
+            infer_context),
+        true,
+        common::errors::Fatal("CacheBackwardOpSymbolicShape for %s failed.",
+                              op->name()));
+  }
+}
+
+void InferSymExprForBlock(const Block& block,
+                          InferSymbolicShapeContext* infer_context) {
+  for (auto& op : block) {
+    std::vector<symbol::ShapeOrDataDimExprs> input_shape_or_data;
+    for (auto& input : op.operands_source()) {
+      input_shape_or_data.emplace_back(
+          infer_context->GetShapeOrDataForValue(input));
+    }
+    InferSymbolicShapeCacheKey op_infer_cache_key(op, input_shape_or_data);
+    InferSymExprForOp(&op, infer_context, op_infer_cache_key);
+    CacheForwardOpSymbolicShape(&op, infer_context, op_infer_cache_key);
+    CacheBackwardOpSymbolicShape(&op, infer_context);
     DebugPrintOpInfo(&op, infer_context);
     CheckInferSymWithInferMeta(&op, infer_context);
   }
