Ordering in tuple affects output of tuple_select

[dragly]

I am having trouble with understanding why the order of values in a tuple can affect the output of a function.

The following is a (possibly very suboptimal) implementation of reverse argsort in Halide, where the order of (value, index) or (index, value) in the output causes the result to change.

Starting with some values:

import halide as hl
import numpy as np

col = hl.Var("col")
row = hl.Var("row")
img = hl.Var("img")

values = hl.Func()
values[img] = hl.Float(32).min()
values[0] = 42.0
values[1] = 12.0
values[2] = 15.0
values[3] = 30.0

Argsort can be implemented as follows:

rdom = hl.RDom([(0, 4), (0, 4)])
sorted_value_and_index = hl.Func()
sorted_value_and_index[img] = (hl.Float(32).min(), -1)
sorted_value_and_index[-1] = (hl.Float(32).max(), -1)
old_value_previous = sorted_value_and_index[rdom.y - 1][0]
old_value = sorted_value_and_index[rdom.y][0]
old_index = sorted_value_and_index[rdom.y][1]
condition = (values[rdom.x] > old_value) & (values[rdom.x] < old_value_previous)
sorted_value_and_index[rdom.y] = hl.tuple_select(condition, (values[rdom.x], rdom.x), sorted_value_and_index[rdom.y])

result = np.asanyarray(sorted_value_and_index.realize([4]))
print("Indices:", result[1])
print("Values:", result[0])

However, the resulting output shows -1 for the indices:

Indices: [-1. -1. -1. -1.]
Values: [42. 30. 15. 12.]

On the other hand, if the order in the tuple is swapped:

sorted_value_and_index = hl.Func()
sorted_value_and_index[img] = (-1, hl.Float(32).min())
sorted_value_and_index[-1] = (-1, hl.Float(32).max())
old_value_previous = sorted_value_and_index[rdom.y - 1][1]
old_value = sorted_value_and_index[rdom.y][1]
old_index = sorted_value_and_index[rdom.y][0]
condition = (values[rdom.x] > old_value) & (values[rdom.x] < old_value_previous)
sorted_value_and_index[rdom.y] = hl.tuple_select(condition, (rdom.x, values[rdom.x]), sorted_value_and_index[rdom.y])

result = np.asanyarray(sorted_value_and_index.realize([4]))
print("Indices:", result[0])
print("Values:", result[1])

The indices and values become correct:

Indices: [0. 3. 2. 1.]
Values: [42. 30. 15. 12.]

Note that by looking at the generated code, it makes sense that the order matters, because the tuple_select is effectively evaluated as two subsequent select statements:

 produce sorted_value_and_index {
  let t23 = -1 - max(sorted_value_and_index.0.min.0, -1)
  for (sorted_value_and_index.s0.img.rebased, 0, sorted_value_and_index.s0.img.loop_extent) {
   sorted_value_and_index.0[sorted_value_and_index.s0.img.rebased + t23] = -inff
   sorted_value_and_index.1[sorted_value_and_index.s0.img.rebased + t23] = -1
  }
  sorted_value_and_index.0[-1 - sorted_value_and_index.0.min.0] = inff
  sorted_value_and_index.1[-1 - sorted_value_and_index.0.min.0] = -1
  for (sorted_value_and_index.s2.r5$y, 0, 4) {
   let t24 = sorted_value_and_index.s2.r5$y - sorted_value_and_index.0.min.0
   for (sorted_value_and_index.s2.r5$x, 0, 4) {
    allocate values[float32 * 4]
    produce values {
     values[sorted_value_and_index.s2.r5$x] = -inff
     values[0] = 42.000000f
     values[1] = 12.000000f
     values[2] = 15.000000f
     values[3] = 30.000000f
    }
    consume values {
     let sorted_value_and_index.s2.t0 = values[sorted_value_and_index.s2.r5$x]
     let t19 = sorted_value_and_index.0[t24]
     sorted_value_and_index.0[t24] = select((uint1)bitwise_and(t19 < sorted_value_and_index.s2.t0, sorted_value_and_index.s2.t0 < sorted_value_and_index.0[t24 + -1]), sorted_value_and_index.s2.t0, t19)
     let sorted_value_and_index.s2.t2 = values[sorted_value_and_index.s2.r5$x]
     sorted_value_and_index.1[t24] = select((uint1)bitwise_and(sorted_value_and_index.0[t24] < sorted_value_and_index.s2.t2, sorted_value_and_index.s2.t2 < sorted_value_and_index.0[t24 + -1]), sorted_value_and_index.s2.r5$x, sorted_value_and_index.1[t24])
     free values
    }
   }
  }
 }
}

Is this expected? Or is there maybe some way of requiring some kind of "atomic" update of the tuple that I am missing? I tried calling both sorted_value_and_index.update(1).atomic(False) and sorted_value_and_index.update(1).atomic(True), but that does not affect the output.