Incomplete shceduling for "where"

[comaniac]

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]

TVM fails to detect the dependence here. You can explicitly add it when creating schedule

s = tvm.create_schedule([ref.op, masked_data.op])

[comaniac]

Thanks, it works!
Would you please confirm if this is a bug or not? I'll take a closer look and try to fix it if so.

[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.

[tqchen]

seems to be a bug that should be fixed

[tqchen]

@comaniac @merrymercy do you mind look into it and submit a fix? Likely we missed something when trace the dependency(in visitor or some recursive pattern)

[comaniac]

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