Mutate free variables in CommReducer in cache_write

[vinx13]

VarReplacer in cache_write should mutate combiner of Reduce expression because CommReducer may contain free variables from outer loop.
see https://discuss.tvm.ai/t/cache-write-does-not-replace-vars-in-reducer-identity/1402

cc @tqchen

[tqchen]

Thanks, @vinx13 can you please add a regression test case, as per https://docs.tvm.ai/contribute/code_review.html#ensure-test-coverage

[sgrechanik-h]

@tqchen If variables from the outer loop are allowed in CommReducer, shouldn't we make it a proper Expr and treat combiners of Reduce expressions as proper subexpressions?

[tqchen]

I agree with @sgrechanik-h this is something that needs discussion. The question is whether we should always restrict CommReducer to be pure function or allow it to be a closure, I do agree that maybe having a pure function restriction makese sense. Thanks for pointing this out.

As a followup to this, maybe a proper change would be have checks that restricts CommReducer to be pure and report error properly.

[sgrechanik-h]

I'm not sure, I guess there must be a practical example where a combiner using an outer loop variable is necessary, since @vinx13 started to fix it. On the other hand, in some cases such dependencies may be expressed with tuple inputs if CommReducer's purity is important.

[vinx13]

@sgrechanik-h
My use case is roi pooling. Each output element is maximum of a roi. If roi is empty, the output is 0, instead of the -MAX_FLOAT. In this case, identity value of combiner depends on outer loop index.
Note that I cannot nest the reduction in a if-else statement since reductions can only be in the top level https://github.com/dmlc/tvm/blob/151f550b2ee5050c74f632679a1c1d86fcea2482/src/op/compute_op.cc#L536