[TIR][FIX] check args size when creating prim_func by runtime::Registry (#13809)

Ever-Kid · web-flow · commit 558c994d1a70 · 2023-01-31T10:04:07.000+08:00
diff --git a/src/te/operation/create_primfunc.cc b/src/te/operation/create_primfunc.cc
@@ -574,6 +574,7 @@ TVM_REGISTER_GLOBAL("te.CreatePrimFunc").set_body([](TVMArgs args, TVMRetValue*
   Array<te::Tensor> arg_list = args[0];
   std::optional<DataType> index_dtype_override{std::nullopt};
   // Add conversion to make std::optional compatible with FFI.
+  ICHECK_EQ(args.size(), 2);
   if (args[1].type_code() != kTVMNullptr) {
     index_dtype_override = args[1].operator DataType();
   }

Original file line number	Diff line number	Diff line change
`@@ -574,6 +574,7 @@ TVM_REGISTER_GLOBAL("te.CreatePrimFunc").set_body([](TVMArgs args, TVMRetValue*`
`574`	`574`	`Array<te::Tensor> arg_list = args[0];`
`575`	`575`	`std::optional<DataType> index_dtype_override{std::nullopt};`
`576`	`576`	`// Add conversion to make std::optional compatible with FFI.`
	`577`	`+ ICHECK_EQ(args.size(), 2);`
`577`	`578`	`if (args[1].type_code() != kTVMNullptr) {`
`578`	`579`	`index_dtype_override = args[1].operator DataType();`
`579`	`580`	`}`