junrushao1994 commented on a change in pull request #9871:
URL: https://github.com/apache/tvm/pull/9871#discussion_r785360038
##########
File path: src/tir/ir/function.cc
##########
@@ -64,6 +64,67 @@ FuncType PrimFuncNode::func_type_annotation() const {
TVM_REGISTER_NODE_TYPE(PrimFuncNode);
+class TensorIntrinManager {
+ public:
+ Map<String, tir::TensorIntrin> reg;
+
+ static TensorIntrinManager* Global() {
+ static TensorIntrinManager* inst = new TensorIntrinManager();
+ return inst;
+ }
+};
+
+TensorIntrin::TensorIntrin(PrimFunc desc_func, PrimFunc intrin_func) {
+ // check the number of func var is equal
+ CHECK_EQ(desc_func->params.size(), intrin_func->params.size());
+ CHECK_EQ(desc_func->buffer_map.size(), intrin_func->buffer_map.size());
+
+ // check both functions' bodies are directly block
+ const auto* desc_realize =
+
Downcast<BlockRealize>(desc_func->body)->block->body.as<BlockRealizeNode>();
+ const auto* intrin_realize =
+
Downcast<BlockRealize>(intrin_func->body)->block->body.as<BlockRealizeNode>();
+ CHECK(desc_realize != nullptr) << "description function's body expect a
directly block";
+ CHECK(intrin_realize != nullptr) << "intrinsic function's body expect a
directly block";
+
+ const Block& desc_block = desc_realize->block;
+ const Block& intrin_block = intrin_realize->block;
+
+ // check block var number and iter type
+ CHECK_EQ(desc_block->iter_vars.size(), intrin_block->iter_vars.size())
+ << "Two blocks should have the same number of block vars";
+ for (size_t i = 0; i < desc_block->iter_vars.size(); i++) {
+ const IterVar& desc_var = desc_block->iter_vars[i];
+ const IterVar& intrin_var = intrin_block->iter_vars[i];
+ CHECK(desc_var->iter_type == intrin_var->iter_type)
+ << "Block iter_type mismatch between " << desc_var->iter_type << " and
"
+ << intrin_var->iter_type;
+ }
+
+ auto n = make_object<TensorIntrinNode>();
Review comment:
`auto` is fine, just nit
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
