cinn_launch_op: skip checking input variables must be used (#37119)

Modify serveral implements on CinnLaunchOp：
1. Skip checking input variables must be used 
2. Move current helper functions to a CinnlaunchContext

Author: CtfGo

paddle/fluid/operators/CMakeLists.txt

@@ -171,6 +171,7 @@ endif()
 if (WITH_CINN)
   op_library(cinn_launch_op SRCS cinn_launch_op.cc cinn_launch_op.cu.cc DEPS transform_desc cinn_compiler cinn ${OP_HEADER_DEPS})
   cc_test(cinn_launch_op_test SRCS cinn_launch_op_test.cc DEPS cinn_compiler cinn_launch_op elementwise_add_op)
+  set_tests_properties(cinn_launch_op_test PROPERTIES ENVIRONMENT OMP_NUM_THREADS=1)
 endif()
 
 # FIXME(typhoonzero): operator deps may not needed.

paddle/fluid/operators/cinn_launch_op.cc

@@ -62,90 +62,102 @@ void DebugCinnCompiledResult(const CinnCompiledObject& result) {
           << "]";
 }
 
-std::vector<std::string> MapPaddleVariablesToCinn(
-    const std::vector<std::string>& paddle_names,
-    const std::unordered_map<std::string, std::string>& paddle2cinn_varmap) {
-  std::vector<std::string> result;
-  result.reserve(result.size());
+void LaunchCinnExecution(const CinnCompiledObject& compiled_obj,
+                         const CinnLaunchContext& context) {
+  compiled_obj.runtime_program->Execute(&context.FinalizeArguments());
+}
+
+CinnLaunchContext::CinnLaunchContext(const CinnCompiledObject& compiled_obj)
+    : paddle2cinn_varmap_(compiled_obj.paddle2cinn_varmap),
+      cinn_scope_(compiled_obj.scope) {
+  auto var_names = cinn_scope_->var_names();
+  cinn_variable_names_.reserve(var_names.size());
   std::transform(
-      paddle_names.begin(), paddle_names.end(), std::back_inserter(result),
-      [&paddle2cinn_varmap](const std::string& pd_name) {
-        PADDLE_ENFORCE_GT(paddle2cinn_varmap.count(pd_name), 0,
-                          platform::errors::NotFound(
-                              "Not found the corresponding cinn variable "
-                              "of paddle variable(%s) in compilation result.",
-                              pd_name));
-        return paddle2cinn_varmap.at(pd_name);
-      });
-  return result;
+      var_names.begin(), var_names.end(),
+      std::inserter(cinn_variable_names_, cinn_variable_names_.end()),
+      [](const auto& name_view) { return std::string(name_view.data()); });
 }
 
-std::vector<CinnTensor> GetCinnTensorsFromCompiledScope(
-    const std::vector<std::string>& cinn_names, const CinnScope& cinn_scope) {
-  std::vector<CinnTensor> result;
-  result.reserve(cinn_names.size());
-  std::transform(cinn_names.begin(), cinn_names.end(),
-                 std::back_inserter(result),
-                 [&cinn_scope](const std::string& var_name) {
-                   PADDLE_ENFORCE_NOT_NULL(
-                       cinn_scope.FindVar(var_name),
-                       platform::errors::NotFound(
-                           "Variable(%s) not found in cinn scope.", var_name));
-                   return cinn_scope.GetTensor(var_name);
-                 });
-  return result;
+bool CinnLaunchContext::IsVariableUsed(const std::string& paddle_name) {
+  return paddle2cinn_varmap_.count(paddle_name) > 0 &&
+         cinn_variable_names_.count(paddle2cinn_varmap_.at(paddle_name)) > 0;
+}
+
+CinnTensor CinnLaunchContext::GetCinnTensor(const std::string& var_name) {
+  PADDLE_ENFORCE_GT(cinn_variable_names_.count(var_name), 0,
+                    platform::errors::NotFound(
+                        "Variable(%s) not found in cinn scope.", var_name));
+  return cinn_scope_->GetTensor(var_name);
+}
+
+std::vector<std::string> CinnLaunchContext::GetInternalVariableNames() {
+  std::unordered_set<std::string> all_parameters(cinn_variable_names_);
+  std::for_each(name2argument_.begin(), name2argument_.end(),
+                [&all_parameters](const auto& name2arg) {
+                  all_parameters.erase(name2arg.first);
+                });
+  return {all_parameters.begin(), all_parameters.end()};
+}
+
+void CinnLaunchContext::MutableTensorData(const std::string& var_name,
+                                          const platform::Place& place,
+                                          LoDTensor* paddle_tensor,
+                                          bool is_internal_var) {
+  auto cinn_name = var_name;
+  if (!is_internal_var) {
+    PADDLE_ENFORCE_EQ(IsVariableUsed(var_name), true,
+                      platform::errors::InvalidArgument(
+                          "Paddle variable(%s) not used by cinn", var_name));
+    cinn_name = paddle2cinn_varmap_.at(var_name);
+  }
+
+  auto cinn_tensor = GetCinnTensor(cinn_name);
+  // TODO(CtfGo): support mutable corresponding c++ type after CINN ready
+  paddle_tensor->mutable_data<float>(
+      framework::make_ddim(cinn_tensor->shape().data()), place);
 }
 
-void CheckTensorEquivalent(const std::string& paddle_name,
-                           const LoDTensor* paddle_tensor,
-                           const CinnTensor& cinn_tensor) {
+void CinnLaunchContext::CheckTensorEquivalent(const std::string& paddle_name,
+                                              const LoDTensor& paddle_tensor,
+                                              const CinnTensor& cinn_tensor) {
   PADDLE_ENFORCE_EQ(
-      paddle_tensor->IsInitialized(), true,
+      paddle_tensor.IsInitialized(), true,
       platform::errors::InvalidArgument(
-          "The tensor in variable(%s) is not initialized.", paddle_name));
+          "Tensor in variable(%s) is not initialized.", paddle_name));
 
   // check dimension
   auto cinn_dims = framework::make_ddim(cinn_tensor->shape().data());
-  PADDLE_ENFORCE_EQ(paddle_tensor->dims(), cinn_dims,
-                    platform::errors::InvalidArgument(
-                        "The tensor dimension in variable(%s) "
-                        "is not equivalent, paddle is [%s] "
-                        "but cinn is [%s].",
-                        paddle_name, paddle_tensor->dims(), cinn_dims));
+  PADDLE_ENFORCE_EQ(paddle_tensor.dims(), cinn_dims,
+                    platform::errors::PreconditionNotMet(
+                        "Tensors' shape in variable(%s) are not equivalent, "
+                        "paddle's shape = [%s], but cinn's shape = [%s].",
+                        paddle_name, paddle_tensor.dims(), cinn_dims));
 
   // TODO(CtfGo): check the underlying data type after CINN ready
 }
 
-void TensorMutableDataWithCinnInfo(const platform::Place& place,
-                                   const CinnTensor& cinn_tensor,
-                                   LoDTensor* paddle_tensor) {
-  // TODO(CtfGo): support mutable corresponding c++ type after CINN ready
-  paddle_tensor->mutable_data<float>(
-      framework::make_ddim(cinn_tensor->shape().data()), place);
-}
-
-std::vector<std::string> SeperateTempVar(
-    const CinnScope& cinn_scope,
-    const std::vector<std::string>& input_cinn_names,
-    const std::vector<std::string>& output_cinn_names) {
-  auto cinn_var_names = cinn_scope.var_names();
-  std::unordered_set<std::string> all_cinn_names;
-  all_cinn_names.reserve(cinn_var_names.size());
-  std::transform(
-      cinn_var_names.begin(), cinn_var_names.end(),
-      std::inserter(all_cinn_names, all_cinn_names.end()),
-      [](const auto& name_view) { return std::string(name_view.data()); });
+void CinnLaunchContext::AssignExternalVariable(const std::string& paddle_name,
+                                               LoDTensor* paddle_tensor) {
+  PADDLE_ENFORCE_EQ(IsVariableUsed(paddle_name), true,
+                    platform::errors::InvalidArgument(
+                        "Paddle variable(%s) not used by cinn", paddle_name));
 
-  auto exclude_fn = [&all_cinn_names](const auto& cinn_name) {
-    all_cinn_names.erase(cinn_name);
-  };
+  const auto& cinn_name = paddle2cinn_varmap_.at(paddle_name);
+  CheckTensorEquivalent(paddle_name, *paddle_tensor, GetCinnTensor(cinn_name));
+  return SetArgument(cinn_name, paddle_tensor);
+}
 
-  std::for_each(input_cinn_names.begin(), input_cinn_names.end(), exclude_fn);
-  std::for_each(output_cinn_names.begin(), output_cinn_names.end(), exclude_fn);
-  return {all_cinn_names.begin(), all_cinn_names.end()};
+void CinnLaunchContext::AssignInternalVariable(const std::string& cinn_name,
+                                               LoDTensor* paddle_tensor) {
+  PADDLE_ENFORCE_GT(cinn_variable_names_.count(cinn_name), 0,
+                    platform::errors::InvalidArgument(
+                        "Variable(%s) not found in cinn socpe.", cinn_name));
+  CheckTensorEquivalent(cinn_name, *paddle_tensor, GetCinnTensor(cinn_name));
+  return SetArgument(cinn_name, paddle_tensor);
 }
 
-std::unique_ptr<cinn_buffer_t> ShareTensorWithCinnBuffer(LoDTensor* tensor) {
+std::unique_ptr<cinn_buffer_t> CinnLaunchContext::ShareTensorWithCinnBuffer(
+    LoDTensor* tensor) {
   // convert paddle dimensions array to cinn format
   std::vector<cinn_dimension_t> cinn_dims(tensor->dims().size());
   for (auto i = 0; i < tensor->dims().size(); ++i) {
@@ -159,17 +171,29 @@ std::unique_ptr<cinn_buffer_t> ShareTensorWithCinnBuffer(LoDTensor* tensor) {
   return cinn_buffer;
 }
 
-void CheckArgumentsNotMissed(
-    const CinnScope& cinn_scope,
-    const std::map<std::string, cinn_pod_value_t>& name2argument) {
-  auto cinn_var_names = cinn_scope.var_names();
-  std::for_each(cinn_var_names.begin(), cinn_var_names.end(),
-                [&name2argument](const auto& name_view) {
-                  PADDLE_ENFORCE_GT(
-                      name2argument.count(name_view.data()), 0,
-                      platform::errors::InvalidArgument(
-                          "Parameter(%s) is not assgined.", name_view.data()));
+void CinnLaunchContext::SetArgument(const std::string& cinn_name,
+                                    LoDTensor* paddle_tensor) {
+  auto buffer = ShareTensorWithCinnBuffer(paddle_tensor);
+  name2argument_.emplace(cinn_name, buffer.get());
+  hold_buffers_.emplace_back(std::move(buffer));
+  VLOG(4) << "SetArgument-" << name2argument_.size() << ": "
+          << "name(" << cinn_name << "), "
+          << "type(" << framework::DataTypeToString(paddle_tensor->type())
+          << "), dims(" << paddle_tensor->dims() << ").";
+}
+
+const std::map<std::string, cinn_pod_value_t>&
+CinnLaunchContext::FinalizeArguments() const {
+  // Check all execution parameters are assigned valued.
+  std::for_each(cinn_variable_names_.begin(), cinn_variable_names_.end(),
+                [this](const auto& var_name) {
+                  PADDLE_ENFORCE_GT(name2argument_.count(var_name), 0,
+                                    platform::errors::InvalidArgument(
+                                        "Variable(%s) is missed for launching "
+                                        "compiled program execution",
+                                        var_name));
                 });
+  return name2argument_;
 }
 
 }  // namespace details