Move dyn fusion api to jit/api/module/ (#72638)
Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/72638
All of the other user-facing apis are in this header, this should be too.
Test Plan: Imported from OSS
Reviewed By: malfet
Differential Revision: D34159122
Pulled By: eellison
fbshipit-source-id: 71110ad5543246d0fa822c426ad1cf2f65e017f6
(cherry picked from commit f1862b32e15b32c0e9bc18d3c7769e2a668b2678)
diff --git a/torch/csrc/jit/api/module.h b/torch/csrc/jit/api/module.h
index c2506c6..e16f8b6 100644
--- a/torch/csrc/jit/api/module.h
+++ b/torch/csrc/jit/api/module.h
@@ -301,6 +301,59 @@
Module& module,
const std::vector<std::string>& other_methods = {});
+enum class FusionBehavior { STATIC, DYNAMIC };
+
+using FusionStrategy = std::vector<std::pair<FusionBehavior, size_t>>;
+// FusionStrategy is used to control the type and number of specializations that
+// can occur during fusion
+//
+// Usage: provide a list of pairs (type, depth) where type is one of "STATIC" or
+// "DYNAMIC" and depth is an integer.
+//
+// Behavior - static vs dynamic:
+// - in STATIC fusion, fused ops are compiled to have fixed input shapes. The
+// input shapes are determined based on a number of initial profiling runs.
+// The shape is determined based on some initial profiling runs. For example,
+// if on the first run an input of shape [2, 4] is observed, then the compiled
+// op will only work on shapes of size [2, 4].
+// - in DYNAMIC fusion, fused ops are compiled to have variable input shapes, so
+// that multiple shapes are possible. Dynamic fusion uses "symbolic shapes",
+// where any dimensions of the same value that are observed in profiling runs
+// are assumed to have the same value. For example, if inputs of [2,3,4] and
+// [3,4,5] are observed, then it is assumed that future inputs will have
+// shapes [a,b,c] and [b,c,d] for some values of a,b,c,d.
+//
+// In both cases, we also recompile on new striding behavior, device, or
+// dtype.
+//
+// Behavior - fallback functions & depth:
+// When an input doesn't match the format required by the specialized compiled
+// op, it will run a fallback function.
+// Fallback functions can also recursively be compiled and specialized based
+// on the input shape. Since compilation can be slow, the "depth" parameter is
+// provided to limit the number of specializations that can be compiled,
+// before JIT gives up on recompiling and falls back to a completely un-fused,
+// un-specialized implementation.
+//
+// The list of (type, depth) pairs controls the type of specializations and the
+// number of specializations. For example: [("STATIC", 2), ("DYNAMIC", 2)]
+// indicates that the first two specializations will use static fusions, the
+// following two specializations will use dynamic fusion, and any inputs that
+// satisfy none of the 4 options will run an unfused implementation.
+// Below an example of the fallback function structure is shown, if given a
+// strategy of [("STATIC", 2), ("DYNAMIC", 2)] and if consecutive runs had
+// these input shapes:
+// [2, 2], [3, 3], [4, 4], [3, 5], ...
+//
+// + specialized: statically fused, shape [2, 2]
+// \-> + fallback 1; statically fused, shape [3, 3]
+// \-> + fallback 2; dynamically fused, shape [A, A]
+// \-> + fallback 3: dynamically fused, shape [A, B]
+// \-> final fallback: unspecialized, unfused
+TORCH_API FusionStrategy getFusionStrategy();
+// returns previous strategy
+TORCH_API FusionStrategy setFusionStrategy(FusionStrategy& fusion_strategy);
+
namespace detail {
struct TORCH_API SlotCursor {
diff --git a/torch/csrc/jit/runtime/profiling_graph_executor_impl.h b/torch/csrc/jit/runtime/profiling_graph_executor_impl.h
index af9d1a2..5ae3241 100644
--- a/torch/csrc/jit/runtime/profiling_graph_executor_impl.h
+++ b/torch/csrc/jit/runtime/profiling_graph_executor_impl.h
@@ -1,62 +1,10 @@
#pragma once
#include <torch/csrc/jit/runtime/graph_executor_impl.h>
+#include <torch/csrc/jit/api/module.h>
namespace torch {
namespace jit {
-enum class FusionBehavior { STATIC, DYNAMIC };
-
-using FusionStrategy = std::vector<std::pair<FusionBehavior, size_t>>;
-// FusionStrategy is used to control the type and number of specializations that
-// can occur during fusion
-//
-// Usage: provide a list of pairs (type, depth) where type is one of "STATIC" or
-// "DYNAMIC" and depth is an integer.
-//
-// Behavior - static vs dynamic:
-// - in STATIC fusion, fused ops are compiled to have fixed input shapes. The
-// input shapes are determined based on a number of initial profiling runs.
-// The shape is determined based on some initial profiling runs. For example,
-// if on the first run an input of shape [2, 4] is observed, then the compiled
-// op will only work on shapes of size [2, 4].
-// - in DYNAMIC fusion, fused ops are compiled to have variable input shapes, so
-// that multiple shapes are possible. Dynamic fusion uses "symbolic shapes",
-// where any dimensions of the same value that are observed in profiling runs
-// are assumed to have the same value. For example, if inputs of [2,3,4] and
-// [3,4,5] are observed, then it is assumed that future inputs will have
-// shapes [a,b,c] and [b,c,d] for some values of a,b,c,d.
-//
-// In both cases, we also recompile on new striding behavior, device, or
-// dtype.
-//
-// Behavior - fallback functions & depth:
-// When an input doesn't match the format required by the specialized compiled
-// op, it will run a fallback function.
-// Fallback functions can also recursively be compiled and specialized based
-// on the input shape. Since compilation can be slow, the "depth" parameter is
-// provided to limit the number of specializations that can be compiled,
-// before JIT gives up on recompiling and falls back to a completely un-fused,
-// un-specialized implementation.
-//
-// The list of (type, depth) pairs controls the type of specializations and the
-// number of specializations. For example: [("STATIC", 2), ("DYNAMIC", 2)]
-// indicates that the first two specializations will use static fusions, the
-// following two specializations will use dynamic fusion, and any inputs that
-// satisfy none of the 4 options will run an unfused implementation.
-// Below an example of the fallback function structure is shown, if given a
-// strategy of [("STATIC", 2), ("DYNAMIC", 2)] and if consecutive runs had
-// these input shapes:
-// [2, 2], [3, 3], [4, 4], [3, 5], ...
-//
-// + specialized: statically fused, shape [2, 2]
-// \-> + fallback 1; statically fused, shape [3, 3]
-// \-> + fallback 2; dynamically fused, shape [A, A]
-// \-> + fallback 3: dynamically fused, shape [A, B]
-// \-> final fallback: unspecialized, unfused
-TORCH_API FusionStrategy getFusionStrategy();
-// returns previous strategy
-TORCH_API FusionStrategy setFusionStrategy(FusionStrategy& fusion_strategy);
-
struct TORCH_API ProfilingGraphExecutorImpl : public GraphExecutorImplBase {
ProfilingGraphExecutorImpl(
const std::shared_ptr<Graph>& graph,
