[Relax][Analysis] Validate FuncStructInfo correctness

src/relax/analysis/well_formed.cc

@@ -90,11 +90,14 @@ class WellFormedChecker : public relax::ExprVisitor,
         WellFormedChecker(obj.as<IRModule>(), check_struct_info);
 
     if (const auto* mod = obj.as<IRModuleNode>()) {
-      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<FunctionNode>()) {
-          Function func = GetRef<Function>(n);
-          well_formed_checker.CheckGlobalVarAndGsymbolConsistency(it.first, func);
+        if (auto opt_func = base_func.as<Function>()) {
+          Function func = opt_func.value();
+          well_formed_checker.CheckGlobalVarAndGsymbolConsistency(gvar, func);
           well_formed_checker.VisitExpr(func);
         }
       }
@@ -146,6 +149,59 @@ class WellFormedChecker : public relax::ExprVisitor,
     }
   }
 
+  void CheckGlobalVarStructInfoConsistency(GlobalVar gvar, BaseFunc func) {
+    if (!gvar->struct_info_.defined() || !func->struct_info_.defined()) {
+      return;
+    }
+
+    if (!IsBaseOf(GetStructInfo(gvar), GetStructInfo(func))) {
+      Malformed(Diagnostic::Error(func->span)
+                << "ValueError: "
+                << "StructInfo of GlobalVar " << gvar
+                << " 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_);
+    }
+  }
+
+  void CheckParamStructInfoConsistency(GlobalVar gvar, BaseFunc base_func) {
+    if (!base_func->struct_info_.defined()) {
+      return;
+    }
+
+    auto func_sinfo = base_func->struct_info_.as<FuncStructInfo>();
+    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<StructInfo> {
+      if (auto func = base_func.as<FunctionNode>()) {
+        return FuncStructInfo(func->params.Map(GetStructInfo), func->ret_struct_info,
+                              func->is_pure);
+      } else if (auto func = base_func.as<tir::PrimFuncNode>()) {
+        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<OpNode>() && !expr->checked_type_.defined()) {
       Malformed(Diagnostic::Error(expr) << "The checked_type_ of Expr " << expr << " is nullptr.");

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 <tvm/relax/expr_functor.h>
+#include <tvm/relax/transform.h>
+
+namespace tvm {
+namespace relax {
+namespace testing {
+
+Function IncorrectImplementationOfBindParamForRelaxFunc(Function func) {
+  auto body = Downcast<relax::SeqExpr>(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

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()