Suppose I have three ops: A is an input of B while B is an input of C. Is it possible to compute A at an axis of C? E.g. doing something equivalent to
s[A].compute_at(s[C], s[C].op.axis[0])
As far as I know the above code won’t work as the compiler would complain that it couldn’t find A among the inputs of C, which is true as C’s only input is B. I’m trying to do something like
s[A].compute_root()
s[A].compute_at(s[C], s[C].op.axis[0])
with which the above error message doesn’t occur, but I haven’t got a working schedule yet. Does this break any connections between A and B and lead to a wrong result (is it safe to do it like this)?
I’m sorry it’s too difficult to share all of my codes, but I think if someone can give me a brief idea of it that would be awesome as well. Thanks!