From 185925697a57902add65fa1f802395f940a0aba6 Mon Sep 17 00:00:00 2001 From: Eric Lunderberg Date: Tue, 2 Apr 2024 18:39:06 -0500 Subject: [PATCH 1/2] [Relax][Analysis] Validate FuncStructInfo correctness Any `FuncStructInfo` annotations should be correct. --- src/relax/analysis/well_formed.cc | 63 +++++++++++++++++-- src/relax/testing/well_formed.cc | 55 ++++++++++++++++ .../python/relax/test_analysis_well_formed.py | 30 +++++++++ 3 files changed, 144 insertions(+), 4 deletions(-) create mode 100644 src/relax/testing/well_formed.cc diff --git a/src/relax/analysis/well_formed.cc b/src/relax/analysis/well_formed.cc index b4a0fc4b9883..a459ba3f61c6 100644 --- a/src/relax/analysis/well_formed.cc +++ b/src/relax/analysis/well_formed.cc @@ -90,11 +90,14 @@ class WellFormedChecker : public relax::ExprVisitor, WellFormedChecker(obj.as(), check_struct_info); if (const auto* mod = obj.as()) { - for (const auto& it : mod->functions) { + for (const auto& [gvar, base_func] : mod->functions) { + well_formed_checker.CheckGlobalVarStructInfoConsistency(gvar, base_func); + well_formed_checker.CheckParamStructInfoConsistency(gvar, base_func); + // visit relax.Function - if (auto* n = it.second.as()) { - Function func = GetRef(n); - well_formed_checker.CheckGlobalVarAndGsymbolConsistency(it.first, func); + if (auto opt_func = base_func.as()) { + Function func = opt_func.value(); + well_formed_checker.CheckGlobalVarAndGsymbolConsistency(gvar, func); well_formed_checker.VisitExpr(func); } } @@ -146,6 +149,58 @@ class WellFormedChecker : public relax::ExprVisitor, } } + void CheckGlobalVarStructInfoConsistency(GlobalVar gvar, BaseFunc func) { + if (!gvar->struct_info_.defined() || !func->struct_info_.defined()) { + return; + } + + if (!StructuralEqual()(gvar->struct_info_, func->struct_info_)) { + Malformed(Diagnostic::Error(func->span) + << "ValueError: " + << "StructInfo of GlobalVar " << gvar + << " does not match the StructInfo of the function. " + << "The GlobalVar has StructInfo " << gvar->struct_info_ + << ", while the function itself has StructInfo " << func->struct_info_); + } + } + + void CheckParamStructInfoConsistency(GlobalVar gvar, BaseFunc base_func) { + if (!base_func->struct_info_.defined()) { + return; + } + + auto func_sinfo = base_func->struct_info_.as(); + if (!func_sinfo) { + Malformed(Diagnostic::Error(base_func->span) + << "TypeError: " + << "Functions must be annotated with FuncStructInfo. " + << "However, function " << gvar << " has struct info " << base_func->struct_info_); + return; + } + + auto expected_sinfo = [&]() -> Optional { + if (auto func = base_func.as()) { + return FuncStructInfo(func->params.Map(GetStructInfo), func->ret_struct_info, + func->is_pure); + } else if (auto func = base_func.as()) { + return GetStructInfo( + tir::PrimFunc(func->params, func->body, func->ret_type, func->buffer_map, func->attrs)); + } else { + return NullOpt; + } + }(); + + if (expected_sinfo.defined() && !StructuralEqual()(func_sinfo, expected_sinfo)) { + Malformed(Diagnostic::Error(base_func->span) + << "ValueError: " + << "If annotated with any StructInfo, " + << "a PrimFunc must have FuncStructInfo that matches the parameters. " + << "However, function " << gvar << " is annotated with struct info " << func_sinfo + << ", but has parameters indicated that it should have struct info " + << expected_sinfo); + } + } + void VisitExpr(const Expr& expr) final { if (!expr.as() && !expr->checked_type_.defined()) { Malformed(Diagnostic::Error(expr) << "The checked_type_ of Expr " << expr << " is nullptr."); diff --git a/src/relax/testing/well_formed.cc b/src/relax/testing/well_formed.cc new file mode 100644 index 000000000000..4f9d6f2072d7 --- /dev/null +++ b/src/relax/testing/well_formed.cc @@ -0,0 +1,55 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you 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 +#include + +namespace tvm { +namespace relax { +namespace testing { + +Function IncorrectImplementationOfBindParamForRelaxFunc(Function func) { + auto body = Downcast(func->body); + body.CopyOnWrite()->blocks = { + BindingBlock({VarBinding(func->params[0], relax::PrimValue::Int64(0))})}; + + auto write_ptr = func.CopyOnWrite(); + write_ptr->params = {}; + write_ptr->body = body; + return func; +} + +TVM_REGISTER_GLOBAL("relax.testing.IncorrectImplementationOfBindParamForRelaxFunc") + .set_body_typed(IncorrectImplementationOfBindParamForRelaxFunc); + +tir::PrimFunc IncorrectImplementationOfBindParamForPrimFunc(tir::PrimFunc func) { + auto new_body = tir::LetStmt(func->params[0], IntImm(func->params[0]->dtype, 0), func->body); + + auto write_ptr = func.CopyOnWrite(); + write_ptr->params = {}; + write_ptr->body = new_body; + return func; +} + +TVM_REGISTER_GLOBAL("relax.testing.IncorrectImplementationOfBindParamForPrimFunc") + .set_body_typed(IncorrectImplementationOfBindParamForPrimFunc); + +} // namespace testing +} // namespace relax +} // namespace tvm diff --git a/tests/python/relax/test_analysis_well_formed.py b/tests/python/relax/test_analysis_well_formed.py index 7deddfd28eb9..4e11334031f9 100644 --- a/tests/python/relax/test_analysis_well_formed.py +++ b/tests/python/relax/test_analysis_well_formed.py @@ -702,5 +702,35 @@ def is_bfloat16_dtype(tensor: T.handle) -> T.bool: assert rx.analysis.well_formed(Module) +def test_incorrect_relax_func_struct_info(): + """relax::Function annotations must be correct""" + + @I.ir_module + class Module: + @R.function + def func(x: R.Prim("int64")) -> R.Prim("int64"): + return x + + arg_binder = tvm.get_global_func("relax.testing.IncorrectImplementationOfBindParamForRelaxFunc") + Module["func"] = arg_binder(Module["func"]) + + assert not rx.analysis.well_formed(Module) + + +def test_incorrect_tir_func_struct_info(): + """If present, PrimFunc annotations must be correct""" + + @I.ir_module + class Module: + @T.prim_func + def func(x: T.int64) -> T.int64: + return x + + arg_binder = tvm.get_global_func("relax.testing.IncorrectImplementationOfBindParamForPrimFunc") + Module["func"] = arg_binder(Module["func"]) + + assert not rx.analysis.well_formed(Module) + + if __name__ == "__main__": tvm.testing.main() From 130ab799b7ca10e47f0600618203b400e5d9102c Mon Sep 17 00:00:00 2001 From: Eric Lunderberg Date: Thu, 4 Apr 2024 17:00:51 -0500 Subject: [PATCH 2/2] Relax the check from StructuralEqual to IsBaseOf --- src/relax/analysis/well_formed.cc | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/src/relax/analysis/well_formed.cc b/src/relax/analysis/well_formed.cc index a459ba3f61c6..6b5704f826e8 100644 --- a/src/relax/analysis/well_formed.cc +++ b/src/relax/analysis/well_formed.cc @@ -154,13 +154,14 @@ class WellFormedChecker : public relax::ExprVisitor, return; } - if (!StructuralEqual()(gvar->struct_info_, func->struct_info_)) { + if (!IsBaseOf(GetStructInfo(gvar), GetStructInfo(func))) { Malformed(Diagnostic::Error(func->span) << "ValueError: " << "StructInfo of GlobalVar " << gvar - << " does not match the StructInfo of the function. " - << "The GlobalVar has StructInfo " << gvar->struct_info_ - << ", while the function itself has StructInfo " << func->struct_info_); + << " must be an accurate description of the function it references. " + << "While it does not need to be an exact match, " + << "the GlobalVar has StructInfo " << gvar->struct_info_ + << ", which is not a superset of the function's StructInfo " << func->struct_info_); } }