This repository was archived by the owner on Jan 24, 2024. It is now read-only.
How to use ir::For::Make in Compute？ #892
Description
我在开发argmin算子时。参考了min等同类算子的实现方法后，发现这些全部都涉及到了更底层的开发，
需要改动的地方非常多，似乎不是本次任务所希望的实现方法。
开发中遇到的问题核心就是不知道在Compute中如何正确的使用与具体尺寸相关的循环，
若直接使用普通的for循环，对shape取值得到的是Expr(n)，而使用for(Expr i = Expr(0); i< shape[aixs]; i++)不能正确编译程序，
于是考虑使用 ir::For，但是对其的原理不理解，不知道其 Expr类型的返回值的含义，编写如下代码仍然无法获得理想的结果。
我在对已经实现的所有compute中均未发现类似的可以参考的用法，多数都采用Reduce方法，但是Reduce只支持max/min/sum/mul等，如果要扩展需要修改较多底层的实现，似乎不是本次任务所希望的实现方法，
因此希望可以得到一个在compute中使用与具体尺寸相关的循环的例子参考，或者有一些其他方向上的指导。
auto temp_tensor = Compute(
{shape[real_axis] + 1},
[=](const std::vector<Expr> &indices) -> Expr { return lang::Identity(Expr(3.402823e+38f)); },
output_name + "_temp");
auto compute = [=](const std::vector<Expr> &indices) -> Expr {
std::vector<Expr> cur_indices(indices);
if (!keep_dims) {
cur_indices.insert(cur_indices.begin() + real_axis, Expr(0));
}
CHECK_EQ(cur_indices.size(), ndim);
Var loop_var("k0", Int(32));
cur_indices[real_axis] = Expr(loop_var);
auto value = in_tensor(cur_indices);
auto last_value = temp_tensor(Expr(loop_var) - 1);
auto update = ir::GT::Make(value, last_value);
auto c_v = ir::Select::Make(update, value, last_value);
auto c_i = ir::Select::Make(update, ir::Cast::Make(Float(32), Expr(loop_var)), temp_tensor({Expr(0)}));
auto body1 = ir::Store::Make(temp_tensor, c_v, {Expr(loop_var)});
auto body2 = ir::Store::Make(temp_tensor, c_i, {Expr(0)});
auto body = ir::Block::Make({body1, body2});
auto forloop = ir::For::Make(
loop_var, common::make_const(1), shape[real_axis], ir::ForType::Serial, ir::DeviceAPI::Host, body);
return ir::Cast::Make(Int(32), temp_tensor({Expr(0)}));
};