在TVM中写in-place计算
刚刚接触TVM的人都觉得TVM总是将输入Tensor变为输出Tensor，由此只能产生新的Tensor，所以in-place计算不能表示。
比如写一个compute：
1
2
3
4
5
6
| import tvm
import numpy as np
A = tvm.placeholder((32, 32), name="A", dtype="float32")
B = tvm.compute((32, 32), lambda i, j: A[i, j] * A[i, j], name="B")
|
我们希望的计算是A自己的原地计算，但是compute会产生一个B，而不是原来的A了。
其实在compute上看来的确是不能进行原地计算，但是要注意compute只是一种符号表达而已，所谓原地计算的关键在于存储位置是同一个。
所以其实原地计算是可以进行的，方法就如下面写的那样：
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
| import tvm
import numpy as np
A = tvm.placeholder((32, 32), name="A", dtype="float32")
B = tvm.compute((32, 32), lambda i, j: A[i, j] * A[i, j], name="B")
s = tvm.create_schedule(B.op)
print(tvm.lower(s, [A, B], simple_mode=True))
a = np.random.uniform(size=(32, 32)).astype("float32")
print(a * a)
ctx = tvm.cpu()
a_tvm = tvm.nd.array(a, ctx)
f = tvm.build(s, [A, B], "llvm")
f(a_tvm, a_tvm)
print(a_tvm)
|
运行一下，输出的两个张量是相同的：
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
| produce B {
for (i, 0, 32) {
for (j, 0, 32) {
B[((i*32) + j)] = (A[((i*32) + j)]*A[((i*32) + j)])
}
}
}
[[0.52790415 0.6495383 0.7994851 ... 0.04518243 0.3339507 0.13561812]
[0.1334706 0.2570305 0.01465752 ... 0.06633044 0.9376439 0.7971087 ]
[0.18213475 0.08067138 0.32445127 ... 0.00524587 0.01037789 0.9083783 ]
...
[0.00298729 0.10308312 0.1025297 ... 0.29211318 0.5990425 0.8084331 ]
[0.11054678 0.2299594 0.05939106 ... 0.05123955 0.07831784 0.8725885 ]
[0.6684032 0.1826053 0.16524051 ... 0.73242 0.08689215 0.21580918]]
[[0.52790415 0.6495383 0.7994851 ... 0.04518243 0.3339507 0.13561812]
[0.1334706 0.2570305 0.01465752 ... 0.06633044 0.9376439 0.7971087 ]
[0.18213475 0.08067138 0.32445127 ... 0.00524587 0.01037789 0.9083783 ]
...
[0.00298729 0.10308312 0.1025297 ... 0.29211318 0.5990425 0.8084331 ]
[0.11054678 0.2299594 0.05939106 ... 0.05123955 0.07831784 0.8725885 ]
[0.6684032 0.1826053 0.16524051 ... 0.73242 0.08689215 0.21580918]]
|
这里的关键点就在于传入的实际张量用同一个，就可以做到原地计算了。
