Hi,
for argmax op, before storage flatten, the dump ir is as below, looks wrong to me
// attr [compute(val2, 0x21b48e0)] realize_scope = ""
realize val2([0, m], [0, n]) {
produce val2 {
for (i, 0, m) {
for (j, 0, n) {
val2(i, j) =exp((val(i, j) + float32(1)))
}
}
}
// attr [compute(T, 0x21b8820)] realize_scope = ""
realize T([0, m]).value[0] {
realize T([0, m]).value[1] {
produce T {
for (i, 0, m) {
T(i).value[0] =-1
T(i).value[1] =-340282346638528859811704183484516925440.000000f
for (k, 0, n) {
T(i).value[0] =select((T(i) >= val2(i, k)), T(i), k)
T(i).value[1] =select((T(i) >= val2(i, k)), T(i), val2(i, k))
}
}
}
}
}
}
after I changed some code in IRPrint.cpp, it’s like below, value should be in front of index, otherwise T([0, m]).value[0] T([0, m]).value[1] look like the same allocate, also confuse me in Provide: T(i).value[0] =select((T(i) >= val2(i, k)), T(i), k)
// attr [compute(val2, 0x22032b0)] realize_scope = ""
realize val2([0, m], [0, n]) {
produce val2 {
for (i, 0, m) {
for (j, 0, n) {
val2(i, j) =exp((val(i, j) + float32(1)))
}
}
}
// attr [compute(T, 0x1cee080)] realize_scope = ""
realize T.value[0]([0, m]) {
realize T.value[1]([0, m]) {
produce T {
for (i, 0, m) {
T.value[0](i) =-1
T.value[1](i) =-340282346638528859811704183484516925440.000000f
for (k, 0, n) {
T.value[0](i) =select((T.value[1](i) >= val2(i, k)), T.value[0](i), k)
T.value[1](i) =select((T.value[1](i) >= val2(i, k)), T.value[1](i), val2(i, k))
}
}
}
}
}
}