How to use ir::For::Make in Compute？

[zrr1999]

我在开发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)}));
  };

[zrr1999]

https://github.com/zrr1999/CINN/blob/dev_argmax/cinn/hlir/op/contrib/argmin.cc

[zhhsplendid]

嗨，首先回答issue的问题，从原理上，CINN的算子是由CINN形成，所以如你所说，直接用C++语言的for循环是无法形成IR For的，需要使用你问的ir::For

但是对于实现Argmin/max，这个操作和ReduceMin/Max很像，在一个范围内返回一个Min/Max，只是ReduceMin/Max返回的是最小/最大值，而Argmin/max返回的是最小/最大值对应的位置。

因此你可以跟踪ReduceMin/ReduceMax在ir::Expr层别的实现，相应从最小/最大值改为位置，然后在Compute中调用即可。

我所阐述的ReduceMin/Max代码位置在

CINN/cinn/hlir/pe/reduction.h

Line 93 in 0431271

ir::Tensor ReduceMin(const ir::Tensor& A,

[zrr1999]

我在 ReduceMax中发现使用了 lang::ReduceMax，通过进一步查找发现调用的是ir::Reduce::Make，所以我应该直接修改ir::Reduce的实现，我这样理解是对的吗。

例如按我是否应该按如下方式修改，并相应修改其他的部分

struct Reduce : public ExprNode<Reduce> {
  enum ReduceType {
    kSum = 0,
    kSub,
    kMul,
    kDiv,
    kMax,
    kMin,
    kArgMax,
    kArgMin
  };
  ...
};

[zhhsplendid]

我认为可以，LGTM！

[zrr1999]

我找到的最核心的实现部分是在tensor.cc中

CINN/cinn/ir/tensor.cc

Line 326 in 0431271

auto *reduce_node = body().As<ir::Reduce>();

auto *reduce_node = body().As<ir::Reduce>();
  if (reduce_node) {
    final_body = reduce_node->body;
    switch (reduce_node->reduce_type) {
      case ir::Reduce::kSum:
        final_body = Tensor(this)(g_axis) + final_body;
        break;
      case ir::Reduce::kMul:
        final_body = Tensor(this)(g_axis) * final_body;
        break;
      case ir::Reduce::kMax:
        final_body = Max::Make(Tensor(this)(g_axis), final_body);
        break;
      case ir::Reduce::kMin:
        final_body = Min::Make(Tensor(this)(g_axis), final_body);
        break;
      case ir::Reduce::kArgmax:
        auto cur_value = Tensor(this)(g_axis)
        auto max_value = ...;
        auto update = ir::GT::Make(cur_value, max_value);
        final_body    = ir::Select::Make(update, ..., final_body);
        break;
      default:
        CINN_NOT_IMPLEMENTED
    }
  }

我参考其他的实现，直接使用了ir::Max::Make 等方法，但是其更加底层的实现我只找到了cuda的实现，但是对其如何联系起来的并不明确。
因此我计划直接在case中实现，将 final_body当作最大取值的索引值，当前值可以用Tensor(this)(g_axis)，但是我找不到一种好的方法通过final_body得到最大值，我是否应该通过修改其传入的参数，手动增加一个axis参数，或是将init和body扩展为std::vector同时保存最大索引和最大值（TVM似乎采用了这种方法）。