This is an automated email from the ASF dual-hosted git repository.
tqchen pushed a commit to branch unity
in repository https://gitbox.apache.org/repos/asf/tvm.git
The following commit(s) were added to refs/heads/unity by this push:
new d9f3aa32ef [Unity][Analysis] Get symbolic TIR vars from struct info
(#14509)
d9f3aa32ef is described below
commit d9f3aa32ef974d3078f70f75f987e30258b59882
Author: Ruihang Lai <[email protected]>
AuthorDate: Wed Apr 5 13:41:12 2023 -0400
[Unity][Analysis] Get symbolic TIR vars from struct info (#14509)
This PR adds an analysis function which helps to fetch all the TIR
variables that appear in a given struct info object. This is helpful
for passes that need to be aware of symbolic shape dynamism.
---
include/tvm/relax/analysis.h | 8 ++++
python/tvm/relax/analysis/analysis.py | 17 +++++++++
src/relax/analysis/struct_info_analysis.cc | 43 ++++++++++++++++++++++
src/relax/transform/lift_transform_params.cc | 42 +--------------------
.../relax/test_analysis_struct_info_analysis.py | 17 +++++++++
5 files changed, 86 insertions(+), 41 deletions(-)
diff --git a/include/tvm/relax/analysis.h b/include/tvm/relax/analysis.h
index 21dc5ef420..f38a202d49 100644
--- a/include/tvm/relax/analysis.h
+++ b/include/tvm/relax/analysis.h
@@ -260,6 +260,14 @@ TVM_DLL bool IsBaseOf(const StructInfo& base, const
StructInfo& derived,
TVM_DLL StructInfo StructInfoLCA(const StructInfo& lhs, const StructInfo& rhs,
arith::Analyzer* ana = nullptr);
+/*!
+ * \brief Get the TIR variables that appear in the input struct info.
+ * The returned list is deduplicated - each TIR variable will appear at most
once.
+ * \param sinfo The struct info object to be analyzed.
+ * \return The list of TIR variables that appear in the input struct info.
+ */
+TVM_DLL Array<tir::Var> TIRVarsInStructInfo(const StructInfo& sinfo);
+
//-----------------------------------
// General IR analysis
//-----------------------------------
diff --git a/python/tvm/relax/analysis/analysis.py
b/python/tvm/relax/analysis/analysis.py
index 2a2c3d87b8..82d1f7a828 100644
--- a/python/tvm/relax/analysis/analysis.py
+++ b/python/tvm/relax/analysis/analysis.py
@@ -167,6 +167,23 @@ def struct_info_lca(lhs: StructInfo, rhs: StructInfo) ->
StructInfo:
return _ffi_api.StructInfoLCA(lhs, rhs) # type: ignore
+def tir_vars_in_struct_info(sinfo: StructInfo) -> List[tir.Var]:
+ """Get the TIR variables that appear in the input struct info.
+ The returned list is deduplicated - each TIR variable will appear at most
once.
+
+ Parameters
+ ----------
+ sinfo : StructInfo
+ The struct info object to be analyzed.
+
+ Returns
+ -------
+ ret : List[tir.Var]
+ The list of TIR variables that appear in the input struct info.
+ """
+ return _ffi_api.TIRVarsInStructInfo(sinfo) # type: ignore
+
+
def bound_vars(expr: Expr) -> List[Var]:
"""
Return all bound variables from expression expr.
diff --git a/src/relax/analysis/struct_info_analysis.cc
b/src/relax/analysis/struct_info_analysis.cc
index 7dfcd60c95..55f58d8511 100644
--- a/src/relax/analysis/struct_info_analysis.cc
+++ b/src/relax/analysis/struct_info_analysis.cc
@@ -26,6 +26,7 @@
#include <tvm/relax/analysis.h>
#include <tvm/relax/expr_functor.h>
#include <tvm/relax/struct_info_functor.h>
+#include <tvm/tir/analysis.h>
#include <tvm/tir/expr_functor.h>
namespace tvm {
@@ -861,5 +862,47 @@ TVM_REGISTER_GLOBAL("relax.analysis.StructInfoLCA")
return StructInfoLCA(lhs, rhs);
});
+//--------------------------
+// TIRVarsInStructInfo
+//--------------------------
+
+Array<tir::Var> TIRVarsInStructInfo(const StructInfo& sinfo) {
+ struct TIRVarsDetector : public StructInfoVisitor {
+ void VisitShape(Array<PrimExpr> shape) {
+ for (const PrimExpr& value : shape) {
+ Array<tir::Var> vars = tir::UndefinedVars(value);
+ for (const tir::Var& var : vars) {
+ auto insert_res = tir_vars_dedup_set.insert(var.get());
+ if (insert_res.second) {
+ tir_vars.push_back(var);
+ }
+ }
+ }
+ }
+
+ void VisitStructInfo_(const ShapeStructInfoNode* shape_sinfo) final {
+ if (shape_sinfo->values.defined()) {
+ VisitShape(shape_sinfo->values.value());
+ }
+ }
+
+ void VisitStructInfo_(const TensorStructInfoNode* tensor_sinfo) final {
+ if (tensor_sinfo->shape.defined()) {
+ VisitStructInfo(GetStructInfo(tensor_sinfo->shape.value()));
+ }
+ }
+
+ Array<tir::Var> tir_vars;
+ std::unordered_set<const tir::VarNode*> tir_vars_dedup_set;
+ };
+
+ TIRVarsDetector detector;
+ detector(sinfo);
+ return detector.tir_vars;
+}
+
+TVM_REGISTER_GLOBAL("relax.analysis.TIRVarsInStructInfo")
+ .set_body_typed([](const StructInfo& sinfo) { return
TIRVarsInStructInfo(sinfo); });
+
} // namespace relax
} // namespace tvm
diff --git a/src/relax/transform/lift_transform_params.cc
b/src/relax/transform/lift_transform_params.cc
index e0296e6ae5..cbeedb6694 100644
--- a/src/relax/transform/lift_transform_params.cc
+++ b/src/relax/transform/lift_transform_params.cc
@@ -132,46 +132,6 @@ class TransformParamsFuncBuilder : public ExprMutator {
std::unordered_set<Var, ObjectPtrHash, ObjectPtrEqual> outputs_;
};
-bool SInfoContainsSymVar(StructInfo sinfo) {
- struct SymVarDetector : public StructInfoVisitor {
- void VisitStructInfo(const StructInfo& sinfo) final {
- if (contains_sym_var) {
- return;
- }
- StructInfoVisitor::VisitStructInfo(sinfo);
- }
-
- bool CheckShape(Array<PrimExpr> shape) {
- for (const PrimExpr& value : shape) {
- const auto* int_imm = value.as<IntImmNode>();
- if (int_imm == nullptr) {
- contains_sym_var = true;
- return false;
- }
- }
- return true;
- }
-
- void VisitStructInfo_(const ShapeStructInfoNode* shape_sinfo) final {
- if (shape_sinfo->values.defined()) {
- CheckShape(shape_sinfo->values.value());
- }
- }
-
- void VisitStructInfo_(const TensorStructInfoNode* tensor_sinfo) final {
- if (tensor_sinfo->shape.defined()) {
- VisitStructInfo(GetStructInfo(tensor_sinfo->shape.value()));
- }
- }
-
- bool contains_sym_var = false;
- };
-
- SymVarDetector detector;
- detector(sinfo);
- return detector.contains_sym_var;
-}
-
/*!
* \brief Visitor that creates the plan of lifting transform params.
*
@@ -230,7 +190,7 @@ class LiftTransformParamsPlanner : public ExprVisitor {
});
// Cond 4. Do not lift when its struct info contains symbolic variables.
- if (SInfoContainsSymVar(GetStructInfo(binding->var))) {
+ if (!TIRVarsInStructInfo(GetStructInfo(binding->var)).empty()) {
can_lift = false;
}
diff --git a/tests/python/relax/test_analysis_struct_info_analysis.py
b/tests/python/relax/test_analysis_struct_info_analysis.py
index 03b98f8a56..62f44d15dd 100644
--- a/tests/python/relax/test_analysis_struct_info_analysis.py
+++ b/tests/python/relax/test_analysis_struct_info_analysis.py
@@ -557,5 +557,22 @@ def test_struct_info_lca():
_check_lca(fopaque2(), fn_info_shape(1), fopaque2())
+def test_tir_vars_in_struct_info():
+ n, m = tir.Var("n", "int64"), tir.Var("m", "int64")
+ shape0 = rx.ShapeStructInfo([1, n, 3])
+ shape1 = rx.ShapeStructInfo([1, 2 * n, n, m])
+ tensor0 = rx.TensorStructInfo([1, n, 3], "int32")
+ tensor1 = rx.TensorStructInfo([1, 2 * n, n, m], "int32")
+ func = rx.FuncStructInfo(
+ [rx.TensorStructInfo([1, 2 * n, n, m], "int32")],
rx.TensorStructInfo([1, n, 3], "int32")
+ )
+
+
tvm.ir.assert_structural_equal(rx.analysis.tir_vars_in_struct_info(shape0), [n])
+
tvm.ir.assert_structural_equal(rx.analysis.tir_vars_in_struct_info(shape1), [n,
m])
+
tvm.ir.assert_structural_equal(rx.analysis.tir_vars_in_struct_info(tensor0),
[n])
+
tvm.ir.assert_structural_equal(rx.analysis.tir_vars_in_struct_info(tensor1),
[n, m])
+ tvm.ir.assert_structural_equal(rx.analysis.tir_vars_in_struct_info(func),
[n, m])
+
+
if __name__ == "__main__":
tvm.testing.main()
