tensorflow
[spirv] Add a pass to decorate types with layout info
If a composite type is used for global variables in StorageBuffer/Uniform/PushConstant storage class, then it must be explicitly laid out. Right now lowering passes need to make sure the generated SPIR-V global variables follow the above rule by themselves. This is duplicating the work. We should consider to have a single pass to go through all global variables in a spv.module and change their types to have layout decorations.
The layout rules can potentially be a CLI option to the pass. But for Vulkan ML, we only need Vulkan Standard Buffer Layout.
@antiagainst can you please comment on the following, do I get you right?
For each global variable in spv.module, if variable is not a composite type(spv.struct, spv.array, vector), we have to change its type to spv.struct. Something like this:
spv.globalVariable @var1 bind(0, 0) : !spv.ptr<f32, StroageBuffer>
to
spv.globalVariable @var1 bind(0, 0) : !spv.ptr<!spv.struct<f32 [0]>, StorageBuffer>
In this case global variable will be laid out be serializer, decorated with Offset.
I think that's a separate issue. For this issue, we should just handle existing !spv.struct in StorageBuffer/Uniform/PushConstant storage class but without layout decorations. An example is turning
spv.globalVariable @var1 bind(0, 0) : !spv.ptr<!spv.struct<f32>, StorageBuffer>
into
spv.globalVariable @var1 bind(0, 0) : !spv.ptr<!spv.struct<f32 [0]>, StorageBuffer>
Sorry for chiping in late here, but I think a good place to put this would be in the conversion framework into SPIR-V. Essentially make sure all the type conversion are handled appropriately when lowering into the SPIR-V dialect. The next thing I am going to work on is to provide interface in the conversion framework that simplifies this. This would be a nice place to implement the above. Some other things need to fall into place there first before this can be added there.
@denis0x0D : If you havent started on this, I can take this up. If you already have, then no worries. I will let you implement this as a pass and see how I can make use of that in the conversion framework
@MaheshRavishankar I've started, but I'm fine to cancel it, if there a better place for this logic.
@antiagainst @MaheshRavishankar I can start to work on passes which are needed for vulkan-runner; to be able to run vulkan-runner simliar to cuda/cpu-runner. Thanks.
@denis0x0D : It depends on how much effort it is to finish it up if you have already started. If you are almost done then it would be great to have (especially the tests would be very helpful). I can reuse the code for the lowering. Hoping that this information will help you prioritize.
Sorry for chiping in late here, but I think a good place to put this would be in the conversion framework into SPIR-V. Essentially make sure all the type conversion are handled appropriately when lowering into the SPIR-V dialect.
I am not sure I grasp well this "laid out" thing, but it this is a requirement for generating a valid SPIR-V module and you're putting this in the lowering / conversion: then it seems like something that is in the scope of the verifier to enforce (the IR should not be "invalid").
I think they are two connected issues. Yes we should add verification rules to catch missing layouts. During the lowering, we also need to generate properly laid out spv.struct types for global resources. This issue focuses on the later. (Although I think the layout calculation rule can be shared. :)
By the way, I'm using an OpRewritePattern machinery to replace an global op with type !spv.ptr<!spv.struct<...>> without layoutinfo to !spv.ptr<!spv.struct<...>[...]>> with layout info. I hope it's a right way.