diff --git a/paddle/fluid/jit/CMakeLists.txt b/paddle/fluid/jit/CMakeLists.txt
index 38d2ae54de32d1bc8dc1cd015d0234ce27e82e4b..85148f1387405cf763ac20704b7f42e3cdb68d21 100644
--- a/paddle/fluid/jit/CMakeLists.txt
+++ b/paddle/fluid/jit/CMakeLists.txt
@@ -1,3 +1,4 @@
+add_subdirectory(function)
 proto_library(paddle_jit_property_proto SRCS property.proto)
 
 cc_library(
@@ -33,8 +34,13 @@ cc_library(
 cc_library(
   jit_layer
   SRCS layer.cc
-  DEPS jit_serializer jit_function_utils jit_serializer_utils
-       jit_compilation_unit jit_function_schema)
+  DEPS jit_serializer
+       jit_function_utils
+       jit_serializer_utils
+       jit_compilation_unit
+       jit_function_schema
+       jit_executor_function
+       jit_pe_function)
 
 if(WITH_TESTING AND NOT WIN32)
   add_custom_target(
diff --git a/paddle/fluid/jit/all.h b/paddle/fluid/jit/all.h
index 5a571a72a28243540c35599992170c4b47c9dbbf..6e768b66f881798e1a9461b9943c6bedc46caae0 100644
--- a/paddle/fluid/jit/all.h
+++ b/paddle/fluid/jit/all.h
@@ -14,7 +14,7 @@
 
 #pragma once
 
-#include "base_function.h"
-#include "layer.h"
-#include "serializer.h"
-#include "serializer_utils.h"
+#include "function/base_function.h"  // NOLINT
+#include "layer.h"                   // NOLINT
+#include "serializer.h"              // NOLINT
+#include "serializer_utils.h"        // NOLINT
diff --git a/paddle/fluid/jit/compilation_unit.cc b/paddle/fluid/jit/compilation_unit.cc
index 5a434fba176d3769680b1a4ecc31edfa68afec00..19c9d38034c781e77911583ce04467f417b197f5 100644
--- a/paddle/fluid/jit/compilation_unit.cc
+++ b/paddle/fluid/jit/compilation_unit.cc
@@ -16,7 +16,7 @@
 
 #include "paddle/phi/core/enforce.h"
 
-#include "paddle/fluid/jit/base_function.h"
+#include "paddle/fluid/jit/function/base_function.h"
 
 namespace paddle {
 namespace jit {
diff --git a/paddle/fluid/jit/function/CMakeLists.txt b/paddle/fluid/jit/function/CMakeLists.txt
new file mode 100644
index 0000000000000000000000000000000000000000..7726ea23865d476327a626571082b0cc9e467b1c
--- /dev/null
+++ b/paddle/fluid/jit/function/CMakeLists.txt
@@ -0,0 +1,9 @@
+cc_library(
+  jit_executor_function
+  SRCS executor_function.cc
+  DEPS executor)
+
+cc_library(
+  jit_pe_function
+  SRCS pe_function.cc
+  DEPS parallel_executor)
diff --git a/paddle/fluid/jit/base_function.h b/paddle/fluid/jit/function/base_function.h
similarity index 100%
rename from paddle/fluid/jit/base_function.h
rename to paddle/fluid/jit/function/base_function.h
diff --git a/paddle/fluid/jit/function/executor_function.cc b/paddle/fluid/jit/function/executor_function.cc
new file mode 100644
index 0000000000000000000000000000000000000000..01e58776b305742dbd519bdd9d8f4c4c6e4ff77b
--- /dev/null
+++ b/paddle/fluid/jit/function/executor_function.cc
@@ -0,0 +1,63 @@
+// Copyright (c) 2022 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
+//
+//     http://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/fluid/jit/function/executor_function.h"
+
+#include "paddle/fluid/framework/program_desc.h"
+#include "paddle/fluid/framework/variable.h"
+#include "paddle/phi/core/enforce.h"
+
+namespace paddle {
+namespace jit {
+
+ExecutorFunction::ExecutorFunction(const std::shared_ptr<FunctionInfo> &info,
+                                   const Name2VariableMap &params_dict,
+                                   const phi::Place &place)
+    : info_(info), place_(place), inner_exe_(place_) {
+  info_->RemoveDescFeedFetch();
+  PADDLE_ENFORCE_GT(
+      static_cast<int64_t>(info_->ProgramDesc().Block(0).OpSize()),
+      0,
+      platform::errors::PreconditionNotMet(
+          "There is no operator in ProgramDesc."));
+  utils::ShareParamsIntoScope(info_->ParamNames(), params_dict, &scope_);
+  VLOG(6) << framework::GenScopeTreeDebugInfo(&scope_);
+}
+
+std::vector<Tensor> ExecutorFunction::operator()(
+    const std::vector<Tensor> &inputs) {
+  auto dense_tensors = utils::ToDenseTensors(inputs);
+  return utils::ToTensors(this->operator()(dense_tensors));
+}
+
+std::vector<DenseTensor> ExecutorFunction::operator()(
+    const std::vector<DenseTensor> &inputs) {
+  utils::ShareIntoScope(info_->InputArgNames(), inputs, &scope_);
+  inner_exe_.Run(info_->ProgramDesc(),
+                 &scope_,
+                 /*blockID=*/0,
+                 false,
+                 true,
+                 info_->OutputArgNames());
+  std::vector<DenseTensor> outputs;
+  utils::FetchOuts(info_->OutputArgNames(), scope_, &outputs);
+  return outputs;
+}
+
+const std::shared_ptr<FunctionInfo> &ExecutorFunction::Info() const {
+  return info_;
+}
+
+}  // namespace jit
+}  // namespace paddle
diff --git a/paddle/fluid/jit/executor_function.h b/paddle/fluid/jit/function/executor_function.h
similarity index 51%
rename from paddle/fluid/jit/executor_function.h
rename to paddle/fluid/jit/function/executor_function.h
index 87a31a91949f36756b47fad27e795fc6715e9f59..5136afaf02566c8875b2734bb75ff6c16d6c41b1 100644
--- a/paddle/fluid/jit/executor_function.h
+++ b/paddle/fluid/jit/function/executor_function.h
@@ -14,17 +14,12 @@
 
 #pragma once
 
-#include <iostream>
-#include <string>
 #include <vector>
 
 #include "paddle/fluid/framework/executor.h"
-#include "paddle/fluid/framework/program_desc.h"
 #include "paddle/fluid/framework/scope.h"
-#include "paddle/fluid/framework/variable.h"
-#include "paddle/phi/core/enforce.h"
 
-#include "paddle/fluid/jit/base_function.h"
+#include "paddle/fluid/jit/function/base_function.h"
 #include "paddle/fluid/jit/function_schema.h"
 #include "paddle/fluid/jit/function_utils.h"
 
@@ -35,39 +30,15 @@ class ExecutorFunction : public BaseFunction {
  public:
   ExecutorFunction(const std::shared_ptr<FunctionInfo> &info,
                    const Name2VariableMap &params_dict,
-                   const phi::Place &place)
-      : info_(info), place_(place), inner_exe_(place_) {
-    info_->RemoveDescFeedFetch();
-    PADDLE_ENFORCE_GT(
-        static_cast<int64_t>(info_->ProgramDesc().Block(0).OpSize()),
-        0,
-        platform::errors::PreconditionNotMet(
-            "There is no operator in ProgramDesc."));
-    utils::ShareParamsIntoScope(info_->ParamNames(), params_dict, &scope_);
-    VLOG(6) << framework::GenScopeTreeDebugInfo(&scope_);
-  }
+                   const phi::Place &place);
 
   ~ExecutorFunction() noexcept {}
 
-  std::vector<Tensor> operator()(const std::vector<Tensor> &inputs) {
-    auto dense_tensors = utils::ToDenseTensors(inputs);
-    return utils::ToTensors(this->operator()(dense_tensors));
-  }
+  std::vector<Tensor> operator()(const std::vector<Tensor> &inputs);
 
-  std::vector<DenseTensor> operator()(const std::vector<DenseTensor> &inputs) {
-    utils::ShareIntoScope(info_->InputArgNames(), inputs, &scope_);
-    inner_exe_.Run(info_->ProgramDesc(),
-                   &scope_,
-                   /*blockID=*/0,
-                   false,
-                   true,
-                   info_->OutputArgNames());
-    std::vector<DenseTensor> outputs;
-    utils::FetchOuts(info_->OutputArgNames(), scope_, &outputs);
-    return outputs;
-  }
+  std::vector<DenseTensor> operator()(const std::vector<DenseTensor> &inputs);
 
-  const std::shared_ptr<FunctionInfo> &Info() const { return info_; }
+  const std::shared_ptr<FunctionInfo> &Info() const;
 
  private:
   std::shared_ptr<FunctionInfo> info_;
diff --git a/paddle/fluid/jit/function/pe_function.cc b/paddle/fluid/jit/function/pe_function.cc
new file mode 100644
index 0000000000000000000000000000000000000000..163200184453b34415bba9fc9fbec2167e9a8ab9
--- /dev/null
+++ b/paddle/fluid/jit/function/pe_function.cc
@@ -0,0 +1,115 @@
+// Copyright (c) 2022 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
+//
+//     http://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/fluid/jit/function/pe_function.h"
+
+#include "paddle/fluid/framework/block_desc.h"
+#include "paddle/fluid/framework/details/build_strategy.h"
+#include "paddle/fluid/framework/program_desc.h"
+#include "paddle/phi/core/enforce.h"
+
+namespace paddle {
+namespace jit {
+
+static ExecutionStrategy GetExecutionStrategy(const platform::Place &place) {
+  ExecutionStrategy execution_strategy;
+
+  auto device_type = platform::Place2DeviceType(place);
+  switch (device_type) {
+    case platform::DeviceType::CPU: {
+      execution_strategy.num_threads_ = 2;
+      break;
+    }
+    case platform::DeviceType::CUDA: {
+      // NOTE: According experiments, one thread is faster in
+      // most model training.
+      execution_strategy.num_threads_ = 1;
+      break;
+    }
+    case platform::DeviceType::XPU: {
+      execution_strategy.num_threads_ = 1;
+      break;
+    }
+    case platform::DeviceType::IPU: {
+      execution_strategy.num_threads_ = 1;
+      break;
+    }
+    default:
+      PADDLE_THROW(platform::errors::Unavailable("Unsupported Device type %d.",
+                                                 device_type));
+  }
+  execution_strategy.use_device_ = device_type;
+
+  return execution_strategy;
+}
+
+PEFunction::PEFunction(const std::shared_ptr<FunctionInfo> &info,
+                       const Name2VariableMap &params_dict,
+                       const phi::Place &place)
+    : info_(info), place_(place) {
+  info_->RemoveDescFeedFetch();
+  PADDLE_ENFORCE_GT(
+      static_cast<int64_t>(info_->ProgramDesc().Block(0).OpSize()),
+      0,
+      platform::errors::PreconditionNotMet(
+          "There is no operator in ProgramDesc."));
+  utils::ShareParamsIntoScope(info_->ParamNames(), params_dict, &scope_);
+  VLOG(6) << framework::GenScopeTreeDebugInfo(&scope_);
+  CreateGraphAndPE();
+}
+
+void PEFunction::CreateGraphAndPE() {
+  framework::details::BuildStrategy build_strategy;
+  auto execution_strategy = GetExecutionStrategy(place_);
+
+  auto &program_desc = info_->ProgramDesc();
+  const framework::BlockDesc &global_block = program_desc.Block(0);
+  int64_t start_op_index = 0;
+  int64_t end_op_index = static_cast<int64_t>(global_block.OpSize());
+
+  graph_ = std::make_shared<Graph>(program_desc, start_op_index, end_op_index);
+  inner_pe_ = std::make_shared<ParallelExecutor>(
+      place_, &scope_, execution_strategy, build_strategy, graph_.get());
+  inner_pe_->PrepareVariables(&scope_);
+  inner_pe_->SkipMemoryReuse(/*scope_idx=*/0, info_->InputArgNames());
+}
+
+std::vector<Tensor> PEFunction::operator()(const std::vector<Tensor> &inputs) {
+  auto dense_tensors = utils::ToDenseTensors(inputs);
+  return utils::ToTensors(this->operator()(dense_tensors));
+}
+
+std::vector<DenseTensor> PEFunction::operator()(
+    const std::vector<DenseTensor> &inputs) {
+  utils::ShareIntoScope(info_->InputArgNames(), inputs, &scope_);
+
+  // update op_handle scope_map in pe->executor_->Graph
+  std::unordered_map<framework::Scope *, framework::Scope *> scope_map = {
+      {inner_pe_->GetLocalScopes().front(), &scope_}};
+  inner_pe_->ResetOpHandleScopeMapOfGraphs(scope_map);
+  // need to recreate tmp variables in new scope
+  inner_pe_->PrepareVariables(&scope_);
+
+  inner_pe_->RunWithoutFetch(info_->OutputArgNames());
+
+  std::vector<DenseTensor> outputs;
+  utils::FetchOuts(info_->OutputArgNames(), scope_, &outputs);
+  scope_.DropKids();
+  return outputs;
+}
+
+const std::shared_ptr<FunctionInfo> &PEFunction::Info() const { return info_; }
+
+}  // namespace jit
+}  // namespace paddle
diff --git a/paddle/fluid/jit/function/pe_function.h b/paddle/fluid/jit/function/pe_function.h
new file mode 100644
index 0000000000000000000000000000000000000000..f847b3702d13953df7b2a8a57848e5a247c4ab3f
--- /dev/null
+++ b/paddle/fluid/jit/function/pe_function.h
@@ -0,0 +1,60 @@
+// Copyright (c) 2022 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
+//
+//     http://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 <vector>
+
+#include "paddle/fluid/framework/details/execution_strategy.h"
+#include "paddle/fluid/framework/ir/graph.h"
+#include "paddle/fluid/framework/parallel_executor.h"
+#include "paddle/fluid/framework/scope.h"
+
+#include "paddle/fluid/jit/function/base_function.h"
+#include "paddle/fluid/jit/function_schema.h"
+#include "paddle/fluid/jit/function_utils.h"
+
+namespace paddle {
+namespace jit {
+
+using ExecutionStrategy = framework::details::ExecutionStrategy;
+using ParallelExecutor = framework::ParallelExecutor;
+using Graph = framework::ir::Graph;
+
+class PEFunction : public BaseFunction {
+ public:
+  PEFunction(const std::shared_ptr<FunctionInfo> &info,
+             const Name2VariableMap &params_dict,
+             const phi::Place &place);
+
+  ~PEFunction() noexcept {}
+
+  void CreateGraphAndPE();
+
+  std::vector<Tensor> operator()(const std::vector<Tensor> &inputs);
+
+  std::vector<DenseTensor> operator()(const std::vector<DenseTensor> &inputs);
+
+  const std::shared_ptr<FunctionInfo> &Info() const;
+
+ private:
+  std::shared_ptr<FunctionInfo> info_;
+  framework::Scope scope_;
+  phi::Place place_;
+  std::shared_ptr<ParallelExecutor> inner_pe_;
+  std::shared_ptr<Graph> graph_;
+};
+
+}  // namespace jit
+}  // namespace paddle
diff --git a/paddle/fluid/jit/layer.cc b/paddle/fluid/jit/layer.cc
index a80b05e45cfbd1a770df8a5fd7ae80c3b7ea9df1..e41f78a48777b7a2d8508591e3d7ed568c1fdc36 100644
--- a/paddle/fluid/jit/layer.cc
+++ b/paddle/fluid/jit/layer.cc
@@ -16,8 +16,8 @@
 
 #include "paddle/fluid/framework/variable.h"
 
-#include "paddle/fluid/jit/base_function.h"
 #include "paddle/fluid/jit/compilation_unit.h"
+#include "paddle/fluid/jit/function/base_function.h"
 #include "paddle/fluid/jit/function_schema.h"
 #include "paddle/phi/core/enforce.h"
 #include "paddle/phi/core/errors.h"
diff --git a/paddle/fluid/jit/layer.h b/paddle/fluid/jit/layer.h
index 4c6c714d37b45d80ebc59b58558bad64377ba130..69304c800fd73cf84ce5e1924cd82bb6b06e1ebd 100644
--- a/paddle/fluid/jit/layer.h
+++ b/paddle/fluid/jit/layer.h
@@ -21,7 +21,7 @@
 #include "paddle/phi/api/include/tensor.h"
 #include "paddle/phi/common/place.h"
 
-#include "base_function.h"  //NOLINT
+#include "function/base_function.h"  //NOLINT
 
 namespace paddle {
 
diff --git a/paddle/fluid/jit/pe_function.h b/paddle/fluid/jit/pe_function.h
deleted file mode 100644
index 809ad5ecbe66231f0cde873a79b3a1ee6eca84c8..0000000000000000000000000000000000000000
--- a/paddle/fluid/jit/pe_function.h
+++ /dev/null
@@ -1,144 +0,0 @@
-// Copyright (c) 2022 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
-//
-//     http://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 <iostream>
-#include <string>
-#include <vector>
-
-#include "paddle/fluid/framework/block_desc.h"
-#include "paddle/fluid/framework/details/build_strategy.h"
-#include "paddle/fluid/framework/details/execution_strategy.h"
-#include "paddle/fluid/framework/executor_cache.h"
-#include "paddle/fluid/framework/ir/graph.h"
-#include "paddle/fluid/framework/program_desc.h"
-#include "paddle/fluid/framework/scope.h"
-#include "paddle/fluid/framework/variable.h"
-#include "paddle/phi/core/enforce.h"
-
-#include "paddle/fluid/jit/base_function.h"
-#include "paddle/fluid/jit/function_schema.h"
-#include "paddle/fluid/jit/function_utils.h"
-
-namespace paddle {
-namespace jit {
-
-using ExecutionStrategy = framework::details::ExecutionStrategy;
-using ParallelExecutor = framework::ParallelExecutor;
-using Graph = framework::ir::Graph;
-
-class PEFunction : public BaseFunction {
- public:
-  PEFunction(const std::shared_ptr<FunctionInfo> &info,
-             const Name2VariableMap &params_dict,
-             const phi::Place &place)
-      : info_(info), place_(place) {
-    info_->RemoveDescFeedFetch();
-    PADDLE_ENFORCE_GT(
-        static_cast<int64_t>(info_->ProgramDesc().Block(0).OpSize()),
-        0,
-        platform::errors::PreconditionNotMet(
-            "There is no operator in ProgramDesc."));
-    utils::ShareParamsIntoScope(info_->ParamNames(), params_dict, &scope_);
-    VLOG(6) << framework::GenScopeTreeDebugInfo(&scope_);
-    CreateGraphAndPE();
-  }
-
-  ~PEFunction() noexcept {}
-
-  static ExecutionStrategy GetExecutionStrategy(const platform::Place &place) {
-    ExecutionStrategy execution_strategy;
-
-    auto device_type = platform::Place2DeviceType(place);
-    switch (device_type) {
-      case platform::DeviceType::CPU: {
-        execution_strategy.num_threads_ = 2;
-        break;
-      }
-      case platform::DeviceType::CUDA: {
-        // NOTE: According experiments, one thread is faster in
-        // most model training.
-        execution_strategy.num_threads_ = 1;
-        break;
-      }
-      case platform::DeviceType::XPU: {
-        execution_strategy.num_threads_ = 1;
-        break;
-      }
-      case platform::DeviceType::IPU: {
-        execution_strategy.num_threads_ = 1;
-        break;
-      }
-      default:
-        PADDLE_THROW(platform::errors::Unavailable(
-            "Unsupported Device type %d.", device_type));
-    }
-    execution_strategy.use_device_ = device_type;
-
-    return execution_strategy;
-  }
-
-  void CreateGraphAndPE() {
-    framework::details::BuildStrategy build_strategy;
-    auto execution_strategy = GetExecutionStrategy(place_);
-
-    auto &program_desc = info_->ProgramDesc();
-    const framework::BlockDesc &global_block = program_desc.Block(0);
-    int64_t start_op_index = 0;
-    int64_t end_op_index = static_cast<int64_t>(global_block.OpSize());
-
-    graph_ =
-        std::make_shared<Graph>(program_desc, start_op_index, end_op_index);
-    inner_pe_ = std::make_shared<ParallelExecutor>(
-        place_, &scope_, execution_strategy, build_strategy, graph_.get());
-    inner_pe_->PrepareVariables(&scope_);
-    inner_pe_->SkipMemoryReuse(/*scope_idx=*/0, info_->InputArgNames());
-  }
-
-  std::vector<Tensor> operator()(const std::vector<Tensor> &inputs) {
-    auto dense_tensors = utils::ToDenseTensors(inputs);
-    return utils::ToTensors(this->operator()(dense_tensors));
-  }
-
-  std::vector<DenseTensor> operator()(const std::vector<DenseTensor> &inputs) {
-    utils::ShareIntoScope(info_->InputArgNames(), inputs, &scope_);
-
-    // update op_handle scope_map in pe->executor_->Graph
-    std::unordered_map<framework::Scope *, framework::Scope *> scope_map = {
-        {inner_pe_->GetLocalScopes().front(), &scope_}};
-    inner_pe_->ResetOpHandleScopeMapOfGraphs(scope_map);
-    // need to recreate tmp variables in new scope
-    inner_pe_->PrepareVariables(&scope_);
-
-    inner_pe_->RunWithoutFetch(info_->OutputArgNames());
-
-    std::vector<DenseTensor> outputs;
-    utils::FetchOuts(info_->OutputArgNames(), scope_, &outputs);
-    scope_.DropKids();
-    return outputs;
-  }
-
-  const std::shared_ptr<FunctionInfo> &Info() const { return info_; }
-
- private:
-  std::shared_ptr<FunctionInfo> info_;
-  framework::Scope scope_;
-  phi::Place place_;
-  std::shared_ptr<ParallelExecutor> inner_pe_;
-  std::shared_ptr<Graph> graph_;
-};
-
-}  // namespace jit
-}  // namespace paddle
diff --git a/paddle/fluid/jit/serializer.cc b/paddle/fluid/jit/serializer.cc
index a9bd5676ad26101a655d6d1b053526cb07f5a59e..6ec321168cdd4c835f4132895a6d12544bead1ef 100644
--- a/paddle/fluid/jit/serializer.cc
+++ b/paddle/fluid/jit/serializer.cc
@@ -20,9 +20,9 @@
 #include "paddle/fluid/framework/variable.h"
 #include "paddle/fluid/platform/device_context.h"
 
-#include "paddle/fluid/jit/executor_function.h"
+#include "paddle/fluid/jit/function/executor_function.h"
+#include "paddle/fluid/jit/function/pe_function.h"
 #include "paddle/fluid/jit/layer.h"
-#include "paddle/fluid/jit/pe_function.h"
 #include "paddle/fluid/jit/property.h"
 #include "paddle/fluid/jit/serializer_utils.h"
 
diff --git a/paddle/fluid/pybind/eager_utils.cc b/paddle/fluid/pybind/eager_utils.cc
index 253291256ef86c74801db4cea40f678ed62216c8..3516a6a96389004496b186ee691a6875586faf78 100644
--- a/paddle/fluid/pybind/eager_utils.cc
+++ b/paddle/fluid/pybind/eager_utils.cc
@@ -22,8 +22,8 @@ limitations under the License. */
 #include "paddle/fluid/framework/convert_utils.h"
 #include "paddle/fluid/framework/scope.h"
 #include "paddle/fluid/framework/scope_guard.h"
-#include "paddle/fluid/jit/executor_function.h"
-#include "paddle/fluid/jit/pe_function.h"
+#include "paddle/fluid/jit/function/executor_function.h"
+#include "paddle/fluid/jit/function/pe_function.h"
 #include "paddle/fluid/memory/allocation/allocator.h"
 #include "paddle/fluid/operators/py_func_op.h"
 #include "paddle/fluid/operators/utils.h"
diff --git a/paddle/fluid/pybind/eager_utils.h b/paddle/fluid/pybind/eager_utils.h
index 4fab8534b7b76705a4286d870952c351d29d7f10..9d261c8a95726f1e582e03da2c5cd8b868d82f4f 100644
--- a/paddle/fluid/pybind/eager_utils.h
+++ b/paddle/fluid/pybind/eager_utils.h
@@ -20,7 +20,7 @@ typedef SSIZE_T ssize_t;
 #include "paddle/fluid/eager/hooks.h"
 #include "paddle/fluid/framework/lod_tensor.h"
 #include "paddle/fluid/framework/tensor.h"
-#include "paddle/fluid/jit/base_function.h"
+#include "paddle/fluid/jit/function/base_function.h"
 #include "paddle/fluid/platform/place.h"
 #include "paddle/phi/common/backend.h"
 #include "paddle/phi/common/data_type.h"
diff --git a/paddle/fluid/pybind/jit.cc b/paddle/fluid/pybind/jit.cc
index 79576e6547f9a1791612e59dc1b50949c07c27f5..4ddae883d19759f7dc62b240d336225a0abd2c81 100644
--- a/paddle/fluid/pybind/jit.cc
+++ b/paddle/fluid/pybind/jit.cc
@@ -18,10 +18,10 @@ limitations under the License. */
 #include "paddle/fluid/imperative/layer.h"
 #include "paddle/fluid/platform/place.h"
 
-#include "paddle/fluid/jit/executor_function.h"
+#include "paddle/fluid/jit/function/executor_function.h"
+#include "paddle/fluid/jit/function/pe_function.h"
 #include "paddle/fluid/jit/function_schema.h"
 #include "paddle/fluid/jit/layer.h"
-#include "paddle/fluid/jit/pe_function.h"
 #include "paddle/fluid/jit/serializer.h"
 
 namespace py = pybind11;
diff --git a/python/setup.py.in b/python/setup.py.in
index 864a0a417c0cf4032ab69a3f61648c8a218eecca..287cb2e08e366758745bfb8410bc0946d269469a 100755
--- a/python/setup.py.in
+++ b/python/setup.py.in
@@ -630,7 +630,7 @@ headers = (
 
 jit_layer_headers = ['layer.h', 'serializer.h', 'serializer_utils.h', 'all.h', 'base_function.h']
 for f in jit_layer_headers:
-    headers += list(find_files(f, '@PADDLE_SOURCE_DIR@/paddle/fluid/jit', recursive=False))
+    headers += list(find_files(f, '@PADDLE_SOURCE_DIR@/paddle/fluid/jit', recursive=True))
 
 if '${WITH_MKLDNN}' == 'ON':
     headers += list(find_files('*', '${MKLDNN_INSTALL_DIR}/include')) # mkldnn