[Vulkan] Rewrote PointerValueTypeRewrite transform #8528
Conversation
Lunderberg
requested review from
areusch,
comaniac,
jroesch,
junrushao,
kparzysz-quic,
masahi,
merrymercy,
tqchen,
vinx13,
yzhliu and
ZihengJiang
as code owners
Lunderberg
force-pushed
the
vulkan_loadnode
branch
from
1584b8d
to
d7c1376
Compare
Lunderberg
force-pushed
the
vulkan_loadnode
branch
3 times, most recently
from
f4e069d
to
9230a0f
Compare
| * runtimes. Once all runtimes support vectorized buffer elements, these | ||
| * parameters can be removed. | ||
| */ | ||
| class VectorTypeRewriter : public StmtExprMutator { |
There was a problem hiding this comment.
I think I understand what this pass does, but I don't get the following point:
My understanding is that
- For SPIR-V target, this pass should be run
- For C-codegen targets, this should be nop
How these two conflicting requirements are satisfied?
There was a problem hiding this comment.
Ah, good point. Currently, the tir.transform.PointerValueTypeRewrite pass is only used in the SPIR-V codegen step, while the tir.transform.StorageRewrite pass is used as part of tvm.lower. These two passes both share the PointerValueTypeRewrite function. Currently, those two conflicting requirements are handled by having different arguments to PointerValueTypeRewrite for those two cases.
Unfortunately, for the C-codegen, it isn't quite a nop, as that would have been simpler to handle. Instead, the AllocateNode, any function parameters, and references to those variables get rewritten, but the StoreNode and LoadNode do not have their indices rewritten to account for the different variable type. I'm trying to determine the best way to handle that on the C codegens, but for now the previous behavior can be maintained with the boolean options to PointerValueTypeRewrite.
There was a problem hiding this comment.
Oh I didn't know that PointerValueTypeRewrite is called from build_vulkan.cc !
There was a problem hiding this comment.
I was surprised as well, as my mental model was that all changes to the TIR graph occur before being passed to the codegen. This extra call is fairly vulkan-specific, and from talking with @tqchen exists to adjust the pointer-type through which Vulkan accesses an array. For cuda, it's unnecessary because the pointer can be cast to the desired output type, but Vulkan doesn't allow those pointer casts. Instead, we need to choose one specific type for each pointer that is passed in, and stick with it through the entire PrimFunc.
My preference would be to pull it out into a target-aware optimization pass, which could be added into either tvm.lower or tvm.build. I've put together some comments over on discuss, and will tag you there as well.
In C-style codegen, pointer types can be freely cast between scalar and vectorized types (e.g. `float16x4* <-> float16*`). In SPIR-V, these are separate types, and no such casting is allowed. This was previously handled by having a special-case for `Ramp(base, stride=1, lanes)` in the codegen. That method didn't cover all possible cases, including Broadcast nodes used as indices. PointerValueTypeRewrite previously re-wrote the AllocateNode and parameter pointer types, but didn't update the Load/Store node. This change tracks which variables can be updated to a vectorized type, and then updates all references to those. This includes removing the `RampNode`, as the vectorization is then included as part of the variable type.
- Added explicit TODO(Lunderberg) for follow-ups - Pass `checker.info_map_` instead of `checker` to `VectorTypeRewriter`
A single memory allocation may have more than one type of data stored within it. This allows the PointerTypeRewrite pass to recognize if a function only uses the pointer as a particular base type. This wasn't an issue in C-based codegen, but is required for Vulkan. Since Vulkan shaders do not permit type-casting, the cast must be done when passing the pointer argument into the shader.
Lunderberg
force-pushed
the
vulkan_loadnode
branch
from
f5d1ef6
to
a5f23fe
Compare
|
cc @tqchen if you want to review |
|
Thanks @Lunderberg |
* [Vulkan] Rewrote PointerValueTypeRewrite transform In C-style codegen, pointer types can be freely cast between scalar and vectorized types (e.g. `float16x4* <-> float16*`). In SPIR-V, these are separate types, and no such casting is allowed. This was previously handled by having a special-case for `Ramp(base, stride=1, lanes)` in the codegen. That method didn't cover all possible cases, including Broadcast nodes used as indices. PointerValueTypeRewrite previously re-wrote the AllocateNode and parameter pointer types, but didn't update the Load/Store node. This change tracks which variables can be updated to a vectorized type, and then updates all references to those. This includes removing the `RampNode`, as the vectorization is then included as part of the variable type. * [StorageRewrite] Updates as recommended in review. - Added explicit TODO(Lunderberg) for follow-ups - Pass `checker.info_map_` instead of `checker` to `VectorTypeRewriter` * [Vulkan] Allow for pointer rewrites that change base type. A single memory allocation may have more than one type of data stored within it. This allows the PointerTypeRewrite pass to recognize if a function only uses the pointer as a particular base type. This wasn't an issue in C-based codegen, but is required for Vulkan. Since Vulkan shaders do not permit type-casting, the cast must be done when passing the pointer argument into the shader. Co-authored-by: Eric Lunderberg <elunderberg@octoml.ai>
* [Vulkan] Rewrote PointerValueTypeRewrite transform In C-style codegen, pointer types can be freely cast between scalar and vectorized types (e.g. `float16x4* <-> float16*`). In SPIR-V, these are separate types, and no such casting is allowed. This was previously handled by having a special-case for `Ramp(base, stride=1, lanes)` in the codegen. That method didn't cover all possible cases, including Broadcast nodes used as indices. PointerValueTypeRewrite previously re-wrote the AllocateNode and parameter pointer types, but didn't update the Load/Store node. This change tracks which variables can be updated to a vectorized type, and then updates all references to those. This includes removing the `RampNode`, as the vectorization is then included as part of the variable type. * [StorageRewrite] Updates as recommended in review. - Added explicit TODO(Lunderberg) for follow-ups - Pass `checker.info_map_` instead of `checker` to `VectorTypeRewriter` * [Vulkan] Allow for pointer rewrites that change base type. A single memory allocation may have more than one type of data stored within it. This allows the PointerTypeRewrite pass to recognize if a function only uses the pointer as a particular base type. This wasn't an issue in C-based codegen, but is required for Vulkan. Since Vulkan shaders do not permit type-casting, the cast must be done when passing the pointer argument into the shader. Co-authored-by: Eric Lunderberg <elunderberg@octoml.ai>
In C-style codegen, pointer types can be freely cast between scalar and vectorized types (e.g.
float16x4* <-> float16*). In SPIR-V, these are separate types, and no such casting is allowed. This was previously handled by having a special-case forRamp(base, stride=1, lanes)in the codegen. That method didn't cover all possible cases, including Broadcast nodes used as indices.PointerValueTypeRewrite previously re-wrote the AllocateNode and parameter pointer types, but didn't update the Load/Store node. This change tracks which variables can be updated to a vectorized type, and then updates all references to those. This includes removing the
RampNode, as the vectorization is then included as part of the variable type.The existing behavior of the StorageRewrite pass, which uses much of the same code, is maintained, to avoid breaking codegen that relies on it.