[REQ] Allow wp.zeros backed by shared memory in kernels/funcs

[nvtw]

Description

Allow local arrays to be backed by shared memory. This allows to reduce the amount of registers consumed by local arrays. The proposed API would include an additional parameter to wp.zeros that is only available in kernels. Tiles have a similar mechanic that allows to select their memory type.

Proposed API: left = wp.zeros(shape=(6,), dtype=wp.vec3, mem_type="shared")

Context

Allows to reduce the register pressure on bigger kernels. Especially for collision detection kernels that need local arrays while processing the contacts, it's very hard to get below the register limit of 255. Shared memory backed local arrays could be a tool to help with that. A similar strategy is used under the hood for the stack memory of bvh queries.

[shi-eric]

Not sure what this was marked as "in review" (maybe I accidentally clicked it). Moving to the 1.12 milestone since I don't think this will be ready in time, @nvtw?

[nvtw]

Maybe an accident. I had the feature in the same MR that brought the launch_bounds support but kicked it out again since my implementation failed on some platforms