I am generating C code using tensor expression. Currently, a lot of
intermediate tensors are generated along the way, and I want to know if I can
use e.g., scalar temp variables instead.
Specifically, I am calculating Conv, then batch norm, then ReLU. For
performance, I would want to essentially inline BN and Relu to the Conv loop
(or Conv loop and BN loop to the Relu, or whatever to only have single loop).
However, I cannot inline Conv loop into other loops because Conv loop contains
reduce and TVM spits and error saying reductions are only allowed at the top of
the loop.
Thus, I currently inline the BN loop and use compute_at() to merge the Relu
loop and conv loop.
For example, I am doing something like this in my tensor expression code:
A = te.placeholer(...)
W = te.placeholder(...)
Conv = te.compute( ... my conv code from A, W)
BN = te.compute( .. BatchNorm using Conv ...)
Relu = te.compute ( .. Relu over BN)
s = te.create_schedule([Relu.op])
s[BN].compute_inline()
s[B].compute_at(s[Relu], Relu.op.axis[3])
The generated C code looks something like this:
for(...)
for(...)
for(...)
Conv[...] = 0;
for (...)
for (...)
for (...)
Conv[...] += A[...] * W[...] // Conv loop
Relu[...] = ReLU(BN(Conv[...]))
However, I only need Relu[...] array results, so Conv[...] is better replaced
by something like a scalar variable. I.e, I want this code to be:
for(...)
for(...)
for(...)
tmp = 0;
for (...)
for (...)
for (...)
tmp += A[...] * W[...] // Conv loop
Relu[...] = ReLU(BN(tmp))
Is there a simple way to do this with te?
The explanation seems a bit messy, so hope this makes sense.
Thank you!
---
[Visit
Topic](https://discuss.tvm.apache.org/t/removing-intermediate-tensors-replacing-to-temporary-variable/8402/1)
to respond.
You are receiving this because you enabled mailing list mode.
To unsubscribe from these emails, [click
here](https://discuss.tvm.apache.org/email/unsubscribe/5f8fda10a69621ef097f79aeef80d04e220eecc9ed44c726ab9182c90aa4962c).
