From f8903dd35b8af073b7fea50e8a29bd7364e50d74 Mon Sep 17 00:00:00 2001 From: Steven Lyubomirsky Date: Fri, 2 Feb 2024 16:20:18 -0500 Subject: [PATCH 1/3] Indicate in doc comments which passes need dataflow blocks --- include/tvm/relax/transform.h | 16 ++++++++++++---- python/tvm/relax/transform/transform.py | 13 +++++++++---- src/relax/transform/dead_code_elimination.cc | 8 +++++--- 3 files changed, 26 insertions(+), 11 deletions(-) diff --git a/include/tvm/relax/transform.h b/include/tvm/relax/transform.h index efe30e5cbb50..753e2621b785 100644 --- a/include/tvm/relax/transform.h +++ b/include/tvm/relax/transform.h @@ -214,7 +214,7 @@ TVM_DLL Pass BindSymbolicVars(Map binding_map, Optional func_name = NullOpt); /*! - * \brief Fold constant expressions. + * \brief Fold constant expressions within dataflow blocks. * * \return The Pass. */ @@ -477,6 +477,8 @@ TVM_DLL Pass Gradient(String func_name, Optional> require_grads = Nul * This must be True if the created composite functions are intended to be offloaded to * an external backend without using the MergeCompositeFunctions pass. * \return The Pass. + * + * \note Only operates within dataflow blocks. */ TVM_DLL Pass FuseOpsByPattern(const tvm::Array& patterns, bool bind_constants = true, bool annotate_codegen = false); @@ -548,6 +550,7 @@ TVM_DLL Pass AlterOpImpl(const Map& op_impl_map, * \brief Layout conversion pass. * \param desired_layouts The desired layouts for some operators. * \return The Pass. + * \note Operates only on dataflow blocks. */ TVM_DLL Pass ConvertLayout(Map> desired_layouts); @@ -564,10 +567,13 @@ TVM_DLL Pass ConvertToDataflow(int min_size = 2); * \brief Dead code elimination. * \sa RemoveAllUnused * Currently it removes: - * 1. Unused local VarBindings in a DataflowBlock. - * 2. Unused DataflowBlocks in a function. - * 3. Unused Relax functions in the module. + * 1. Unused local VarBindings + * (those where the bound var is unused and no impure operation is used). + * 2. Unused Relax functions in the module. * We detect the call chain from the entry function, and remove all unused functions. + * + * Any binding blocks that are left empty will be removed by the normalizer. + * * \return The Pass. */ TVM_DLL Pass DeadCodeElimination(Array entry_functions); @@ -589,6 +595,8 @@ TVM_DLL Pass DataflowUseInplaceCalls(); * \param fp16_input_names The names of function parameters whose dtype should become fp16. The * function signature would change accordingly. * \return The Pass. + * + * \note Mainly operates within dataflow blocks. */ TVM_DLL Pass ToMixedPrecision(const DataType& out_dtype, Optional> fp16_input_names = NullOpt); diff --git a/python/tvm/relax/transform/transform.py b/python/tvm/relax/transform/transform.py index e360c09392f3..b3283d7eb8ef 100644 --- a/python/tvm/relax/transform/transform.py +++ b/python/tvm/relax/transform/transform.py @@ -584,7 +584,7 @@ def RunCodegen( def FoldConstant() -> tvm.ir.transform.Pass: - """Fold constant expressions. + """Fold constant expressions within dataflow blocks. Returns ------- @@ -764,6 +764,8 @@ def FuseOpsByPattern( The end result is similar to FuseOps, but fusion is driven completely by the provided patterns. + Note: Only operates within dataflow blocks. + Parameters ---------- patterns : List[Union[FusionPattern, Tuple]] @@ -1172,11 +1174,12 @@ def DeadCodeElimination(entry_functions: Optional[List[str]] = None) -> tvm.ir.t """Remove dead code in the IRModule. Currently it removes: - 1. Unused local VarBindings in a DataflowBlock. - 2. Unused DataflowBlocks in a function. - 3. Unused Relax functions in the module. + 1. Unused local VarBindings + (those where the bound var is unused and no impure operation is used). + 2. Unused Relax functions in the module. We detect the call chain from the entry function, and remove all unused functions. + Any binding blocks that are left empty will be removed by the normalizer. Notes ----- @@ -1203,6 +1206,8 @@ def ToMixedPrecision( """Automatic mixed precision pass. Currently the pass assumes the input module to be fp32 only, and will automatically cast fp32 to fp16 for certain ops. + Note: Mainly operates within dataflow blocks. + Parameters ---------- out_dtype : str diff --git a/src/relax/transform/dead_code_elimination.cc b/src/relax/transform/dead_code_elimination.cc index 248e4c1c00b7..73f66d2ef362 100644 --- a/src/relax/transform/dead_code_elimination.cc +++ b/src/relax/transform/dead_code_elimination.cc @@ -24,10 +24,12 @@ * \sa tvm/relax/ir/binding_rewrite.cc * * Currently it removes: - * 1. Unused local VarBindings in a DataflowBlock. - * 2. Unused DataflowBlocks in a function. - * 3. Unused Relax functions in the module. + * 1. Unused local VarBindings + * (those where the bound var is unused and no impure operation is used). + * 2. Unused Relax functions in the module. * We detect the call chain from the entry function, and remove all unused functions. + * + * Any binding blocks that are left empty will be removed by the normalizer. */ #include From 4b6add0a568c4e4eb9d3b6dca3b523668012fc23 Mon Sep 17 00:00:00 2001 From: Steven Lyubomirsky Date: Fri, 2 Feb 2024 16:26:13 -0500 Subject: [PATCH 2/3] Also encourage users to use ConvertToDataflow --- include/tvm/relax/transform.h | 19 ++++++++++++------- python/tvm/relax/transform/transform.py | 16 +++++++++++++--- 2 files changed, 25 insertions(+), 10 deletions(-) diff --git a/include/tvm/relax/transform.h b/include/tvm/relax/transform.h index 753e2621b785..ccc8fe8c07a5 100644 --- a/include/tvm/relax/transform.h +++ b/include/tvm/relax/transform.h @@ -216,6 +216,8 @@ TVM_DLL Pass BindSymbolicVars(Map binding_map, /*! * \brief Fold constant expressions within dataflow blocks. * + * \note ConvertToDataflow may need to be called first to provide dataflow blocks. + * * \return The Pass. */ TVM_DLL Pass FoldConstant(); @@ -458,6 +460,8 @@ class PatternCheckContext : public ObjectRef { * of the return value as the target. If it is not specified, the first return value will be the * target. * \return The Pass. + * + * \note ConvertToDataflow may need to be called first to provide dataflow blocks. */ TVM_DLL Pass Gradient(String func_name, Optional> require_grads = NullOpt, int target_index = 0); @@ -477,8 +481,8 @@ TVM_DLL Pass Gradient(String func_name, Optional> require_grads = Nul * This must be True if the created composite functions are intended to be offloaded to * an external backend without using the MergeCompositeFunctions pass. * \return The Pass. - * - * \note Only operates within dataflow blocks. + * + * \note Only operates within dataflow blocks. ConvertToDataflow may need to be called first. */ TVM_DLL Pass FuseOpsByPattern(const tvm::Array& patterns, bool bind_constants = true, bool annotate_codegen = false); @@ -550,7 +554,7 @@ TVM_DLL Pass AlterOpImpl(const Map& op_impl_map, * \brief Layout conversion pass. * \param desired_layouts The desired layouts for some operators. * \return The Pass. - * \note Operates only on dataflow blocks. + * \note Operates only on dataflow blocks. ConvertToDataflow may need to be called first. */ TVM_DLL Pass ConvertLayout(Map> desired_layouts); @@ -571,9 +575,9 @@ TVM_DLL Pass ConvertToDataflow(int min_size = 2); * (those where the bound var is unused and no impure operation is used). * 2. Unused Relax functions in the module. * We detect the call chain from the entry function, and remove all unused functions. - * + * * Any binding blocks that are left empty will be removed by the normalizer. - * + * * \return The Pass. */ TVM_DLL Pass DeadCodeElimination(Array entry_functions); @@ -584,6 +588,7 @@ TVM_DLL Pass DeadCodeElimination(Array entry_functions); * Supported operators will be replaced by calls to `call_tir_inplace` that invoke in-place * PrimFunc implementations of those operators (which are based on the legalizations of those * operators). + * \note ConvertToDataflow may need to be called first to provide dataflow blocks. * \return The pass. */ TVM_DLL Pass DataflowUseInplaceCalls(); @@ -595,8 +600,8 @@ TVM_DLL Pass DataflowUseInplaceCalls(); * \param fp16_input_names The names of function parameters whose dtype should become fp16. The * function signature would change accordingly. * \return The Pass. - * - * \note Mainly operates within dataflow blocks. + * + * \note Mainly operates within dataflow blocks. ConvertToDataflow may need to be called first. */ TVM_DLL Pass ToMixedPrecision(const DataType& out_dtype, Optional> fp16_input_names = NullOpt); diff --git a/python/tvm/relax/transform/transform.py b/python/tvm/relax/transform/transform.py index b3283d7eb8ef..b2aaa3e331a1 100644 --- a/python/tvm/relax/transform/transform.py +++ b/python/tvm/relax/transform/transform.py @@ -52,7 +52,7 @@ def Gradient( """Reverse-mode automatic differentiation. This pass will differentiate one function in the IRModule. Now the input function must have only - one dataflow block. + one dataflow block (ConvertToDataflow may need to be called first). For a given function specified by `func_name`, it generates a new function with the name `func_name + "_adjoint"`. The new function computes the gradient of the **differentiation @@ -260,6 +260,8 @@ def DataflowUseInplaceCalls() -> tvm.ir.transform.Pass: in-place PrimFunc implementations of those operators (which are based on the legalizations of those operators). + Note: ConvertToDataflow may need to be called first to provide dataflow blocks. + Returns ------- ret: tvm.ir.transform.Pass @@ -282,6 +284,8 @@ def ConvertToDataflow(min_size: int = 2) -> tvm.ir.transform.Pass: """A pass that converts consecutive dataflow operations inside binding blocks into dataflow blocks. + Note: ConvertToDataflow may need to be called first. + Params ------ min_size: int @@ -395,6 +399,8 @@ def RewriteDataflowReshape() -> tvm.ir.transform.Pass: operation at runtime, instead of doing real data copy. Here "reshape-like" includes reshape, expand_dims, flatten, etc. + Note: Operates only in dataflow blocks. ConvertToDataflow may need to be called first. + Returns ------- ret : tvm.ir.transform.Pass @@ -586,6 +592,8 @@ def RunCodegen( def FoldConstant() -> tvm.ir.transform.Pass: """Fold constant expressions within dataflow blocks. + Note: ConvertToDataflow may need to be called first to provide dataflow blocks. + Returns ------- ret: tvm.ir.transform.Pass @@ -651,6 +659,8 @@ def FuseOps(fuse_opt_level=-1) -> tvm.ir.transform.Pass: A follow-up pass named "FuseTIR" will generate a TIR PrimFunc for each grouped function. + Note: ConvertToDataflow may need to be called first to provide dataflow blocks. + Parameters ---------- fuse_opt_level : int @@ -764,7 +774,7 @@ def FuseOpsByPattern( The end result is similar to FuseOps, but fusion is driven completely by the provided patterns. - Note: Only operates within dataflow blocks. + Note: Only operates within dataflow blocks. ConvertToDataflow may need to be called first. Parameters ---------- @@ -1206,7 +1216,7 @@ def ToMixedPrecision( """Automatic mixed precision pass. Currently the pass assumes the input module to be fp32 only, and will automatically cast fp32 to fp16 for certain ops. - Note: Mainly operates within dataflow blocks. + Note: Mainly operates within dataflow blocks. ConvertToDataflow may need to be called first. Parameters ---------- From f3044d5f512dab8059e1259b58844e1393a77066 Mon Sep 17 00:00:00 2001 From: Steven Lyubomirsky Date: Mon, 5 Feb 2024 17:15:45 -0500 Subject: [PATCH 3/3] Whitespace --- include/tvm/relax/transform.h | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/include/tvm/relax/transform.h b/include/tvm/relax/transform.h index ccc8fe8c07a5..027fd6f824db 100644 --- a/include/tvm/relax/transform.h +++ b/include/tvm/relax/transform.h @@ -460,7 +460,7 @@ class PatternCheckContext : public ObjectRef { * of the return value as the target. If it is not specified, the first return value will be the * target. * \return The Pass. - * + * * \note ConvertToDataflow may need to be called first to provide dataflow blocks. */ TVM_DLL Pass Gradient(String func_name, Optional> require_grads = NullOpt,