Open spectrometerHBH opened 3 years ago
I agree that clarifying part of the iterator space as affine might be something we want to think a bit more
Some addtional thoughts on correctness.
# Cannot parallel i
for i in tir.grid(8):
with tir.block([8]) as [vi]:
tir.write(A[vi : vi + 2])
A[vi] = vi
A[vi + 1] = vi
# Cannot reorder i, j
for i, j in tir.grid(8, 8)
with tir.block([8, 8]) as [vi, vj]:
tir.bind(vi, i)
tir.bind(vj, j)
A[vi, vj] = B[vi - 1, vj + 1]
B[vi, vj] = 1
We design Block
and expect it to be a level of isolation. But it may not work so well. Here are some cases:
with tir.block([1, 1]):
wmma.sync(A, B, C)
Tensorcore is expected to be a warp-level operation. It requires 32-threadIdx working together. However, by looking at the block signature, we can not know the constraint. Information lacking may influence additional scheduling.
Things are not only happened in TensorCore. Most of the opaque intrinsics have their unique constraint.
# Cannot reorder i, k
for i, k in tir.grid(...):
with tir.block([8, tir.reduce(8)]) as [vi, vk]:
tir.bind(vi, i)
tir.bind(vk, k)
A[i] += B[i, k]
C[i] = A[i]
data_par
and do not use it without carefully thinking.Great point, one thing we need to keep in mind is our usecase. For example, it would be great to be able to have a clear block isolation for wmma because this seems to be the key need for our major usecase.
We can however, tighten up some of the corner cases as long as we have the things we need in the search space.
Recently @Hzfengsy brought up a question regarding affine binding and related schedule primitives. After brief discussions, I put my thoughts here for further discussions.
Intro case
The above code is a simplified pooling operator, and if we do following schedule transformations, we will get
Note that the cache read block's binding is not affine. But we may still want to
Affine binding and parallelization
A clear motivation for affine binding is shown below
The block above is indeed a complete block, but it is incorrect to parallel(io). The current parallel algorithm doesn't reject such cases. If we add C in its read buffers, then the complete block check will work. But I still think if it is worth discussing whether it is OK to do so.
Affine binding and reordering
It's incorrect to reorder k, io, ii.
Affine binding and blockization
Previously, I implemented subspace division for affine bindings to do blockization, since we need to generate reasonable bindings for the result outer and inner blocks. If we face the above tensorization need, affine binding doesn't work anymore.
A somewhat ad-hoc fix I can come up with is to generalized subspace division to work under vi = affine(outer loops) + some_inner _loop.
Would be great to hear your opinions.
cc @tqchen @Hzfengsy @junrushao1994 @MasterJH5574 @jinhongyii @yzh119