[TIR] Allow IndexMap applied to arguments with different dtypes #13085
Conversation
|
Thanks for contributing to TVM! Please refer to the contributing guidelines https://tvm.apache.org/docs/contribute/ for useful information and tips. Please request code reviews from Reviewers by @-ing them in a comment.
Generated by tvm-bot |
vinx13
force-pushed
the
feat/index_map_generic_dtype
branch
2 times, most recently
from
a31ac49
to
b585fb5
Compare
src/tir/ir/index_map.cc
Outdated
| for (int i = 0; i < static_cast<int>(arguments.size()); ++i) { | ||
| var_map_.Set(index_map_->initial_indices[i], arguments[i]); | ||
| } | ||
| Array<PrimExpr> result = index_map_->final_indices; |
There was a problem hiding this comment.
Can use Array::Map instead of making a copy that is then mutated.
return index_map_->final_indices.Map(...);
src/tir/ir/index_map.cc
Outdated
| /*! | ||
| * \brief Evaluator to compute the mapped indices of a given index map. | ||
| */ | ||
| class IndexMapEvaluator : public DataTypeLegalizer { |
There was a problem hiding this comment.
Should this be part of a more general utility overall? It looks like this is effectively a variation of tvm::tir::Substitute that allows the substituted values to have a different datatype than the replaced var's type. Rather than having a utility that is specific to IndexMap, should this functionality be an option when calling Substitute?
There was a problem hiding this comment.
It inherits from DataTypeLegalizer that contains extra logic than StmtMutator, I think it's better to be separated from IRSubstitute. Perhaps we can introduce another utility function SubstituteWithDataTypePromotion
There was a problem hiding this comment.
Hmm, good point. I had been thinking that the DataTypeLegalizer would have the necessary follow-up steps that should be performed after every substitution in order to have valid TIR, rather having extra logic as well. Since there is extra logic, that would make sense to me to have the separate SubstituteWithDataTypeLegalization.
vinx13
force-pushed
the
feat/index_map_generic_dtype
branch
from
b585fb5
to
da3dcdf
Compare
vinx13
force-pushed
the
feat/index_map_generic_dtype
branch
from
da3dcdf
to
5af29a6
Compare
src/tir/ir/index_map.cc
Outdated
| /*! | ||
| * \brief Evaluator to compute the mapped indices of a given index map. | ||
| */ | ||
| class IndexMapEvaluator : public DataTypeLegalizer { |
There was a problem hiding this comment.
Hmm, good point. I had been thinking that the DataTypeLegalizer would have the necessary follow-up steps that should be performed after every substitution in order to have valid TIR, rather having extra logic as well. Since there is extra logic, that would make sense to me to have the separate SubstituteWithDataTypeLegalization.
areusch
added
needs-triage
…he#13085) * [TIR] Allow IndexMap applied to arguments with different dtypes * address comments * Add SubstituteWithDataTypeLegalization
…he#13085) * [TIR] Allow IndexMap applied to arguments with different dtypes * address comments * Add SubstituteWithDataTypeLegalization
cc @Lunderberg @Hzfengsy