Incomplete shceduling for "where"

I was trying TVM IR with the following simple computation pattern that creates a set of reference data ref to filter the computation at the next stage. Here is the code I used:

N = tvm.var('N')
V = tvm.var('V')

data = tvm.placeholder((N,V), name='data')
rv = tvm.reduce_axis((0, N), name='rv')

ref  = tvm.compute((N,), lambda n: data[n, 0], name='ref')
masked_data = tvm.compute((N,),
        lambda n: tvm.sum(data[n, rv], rv, ref[rv] == 1),
        name='masked_data')

s = tvm.create_schedule(masked_data.op)
print(tvm.lower(s, [data, ref, masked_data], simple_mode=True))

The output from lower is shown below:

produce masked_data {
  for (n, 0, N) {
    masked_data[n] = 0.000000f
    for (rv, 0, N) {
      if ((ref[rv] == 1.000000f)) {
        masked_data[n] = (masked_data[n] + data[((n*V) + rv)])
      }
    }
  }
}

As can be seen, the computation for ref was missing. Did I misunderstand anything of using TVM IR?

p.s. I know this piece of code looks silly since we can achieve the same functionality without creating ref. On the other hand, this is just the simplified code example to illustrate the problem I’m encountering.

[merrymercy]

I think it is a bug. if you put ref in the body of tvm.sum, tvm can detect the dependency. But it fails when ref is in the condition.

[comaniac]

Sure I’ll take a close look later on, since I have other plans during the Christmas.