[HLSL][SPIRV] Propagate address space casts to their uses

[s-perron]

In SPIR-V shaders, you cannot cast a pointer from one storage class to another (address space casts in llvm-ir). So we need to try to remove the address space casts when possible, and fails nicely when not possible.

This happens in HLSL for the this pointer. See https://github.com/llvm/wg-hlsl/blob/main/proposals/0021-this-address-space.md for a discussion of the issues.

[efriedma-quic]

This sounds scary: I'd expect that a cast is either illegal in the frontend, or otherwise there's some possible lowering. We aren't allowed to produce an error just because we find the cast might lead to undefined behavior at runtime. And we shouldn't try to "propagate" address-spaces for correctness.

[Keenuts]

This sounds scary

Well, that's the essence of HLSL 🙃

We aren't allowed to produce an error just because we find the cast might lead to undefined behavior at runtime.

This is not about preventing undefined behaviors. This is about being able to even generate the code.

Consider this example:

class A {
    int member;

    void set(int value) {
        member = value;
    }

    int get() {
         return member;
    }
}

cbuffer CB {
    A constant;
};
StructuredBuffer<A> ro_buffer;
static A global;

In DXC, depending on the usage, we either allowed this, or failed, but very late (during validation), with a very unfriendly message:

local.set(local.get());    // Fine
global.set(global.get());  // Fine

ro_buffer[0].set(ro_buffer[0].get());
// DXIL: repro.hlsl:7:16: error: store should be on uav resource.
//       quite clean, tells us the `member = value` is not legal.
// SPIRV: passes compilation because of inline+legalization+addrspace forwarding.

constant.set(constant.get());
// DXIL:  error: cast<X>() argument of incompatible type!
//        no line info, no context, just this cast<X>() error.
// SPIRV: generated SPIR-V is invalid: [VUID-StandaloneSpirv-Uniform-06925] In the Vulkan environment, cannot store to Uniform Blocks
//        No line info, no context, just this validation error.

Those errors exists becaue the this pointer address space remains default, and thus the constraints the storage class implies are not known when generating the code.

What we added in Clang already is explicit address spaces (hlsl_device, hlsl_constant, etc).
Those will allow us to correctly carry the address space constraints.

But there are still issues: Clang is not ready to handle function overloading depending on the this pointer address space, and neither is the HLSL spec.
Meaning the codegen for the snippet above will generate a single version for get and set, with default address spaces. This means Clang will still happily generate a member = value, even if the address space of this is read-only.

Address space propagation is a mean to fix what is completely broken today, it's not perfect, but it's a step forward.

In this example, the constant example would yield something like:

%ptr = getelementptr %A, ptr addrspace(hlsl_constant) %constant, %indices
%val = call i32 @A_get(ptr addrspacecast(ptr addrspace(hlsl_constant) %ptr to ptr));
call void @A_set(ptr addrspacecast(ptr addrspace(hlsl_constant) %ptr to ptr), i32 %val);

Now, if we inline the functions, we get something like:

%ptr = getelementptr %A, ptr addrspace(hlsl_constant) %constant, %indices
%getptr = ptr addrspacecast(ptr addrspace(hlsl_constant) %ptr to ptr
%val = load i32, ptr %getptr
%setptr = ptr addrspacecast(ptr addrspace(hlsl_constant) %ptr to ptr
store i32 %val, ptr %setptr

Propagating the address spaces will yield:

%ptr = getelementptr %A, ptr addrspace(hlsl_constant) %constant, %indices
%val = load i32, ptr addrspace(hlsl_constant) %ptr
store i32 %val, ptr addrspace(hlsl_constant) %ptr

And this will allow us to correctly identify the store as bad, since we store to an hlsl_constant address space, which is illegal. And because we'd do that at the IR level, we still have the context around the instruction, and will be able to surface a easy-to-understand error linking the store to the source code.
(Which is way better than DXC relying on spirv-val, which has no context)

[s-perron]

This sounds scary: I'd expect that a cast is either illegal in the frontend, or otherwise there's some possible lowering. We aren't allowed to produce an error just because we find the cast might lead to undefined behavior at runtime. And we shouldn't try to "propagate" address-spaces for correctness.

I added a link to the design doc which explains the problem with generating correct code in the the frontend. If you have other solutions, we are open to suggestions.

[efriedma-quic]

These address space casts cause issues for SPIR-V because SPIR-V’s Generic storage class is not allowed by the Vulkan environment for SPIR-V,

Oh. The AST contains pointer types which are impossible to represent. That's bad.

I don't think it's a great idea to try to do this in LLVM IR. If everything goes right, it works out, I think? But basically every LLVM IR pass is your enemy: if you ever end up with an address-space cast where the operand isn't a variable, or a GEP of a variable, it's very hard to recover... and LLVM passes love eliminating "duplicate" code. If you are going this route, I'd suggest eliminating the casts as soon as possible, before you start running the full pass pipeline, to ensure consistency and reduce the risk of complications from optimizations.

In terms of doing something in the AST, we do have "deducing this". So using a template to represent the type of "this" isn't completely new... although getting overload resolution to handle address-spaces the way you want could get messy.

[Keenuts]

If you are going this route, I'd suggest eliminating the casts as soon as possible, before you start running the full pass pipeline

Interesting, from what I've tried, I thought it would be simpler to solve this after optimizations: once everything is inlined & DCE cleaned up the noise we have full visibility on the chain between a pointer load and the usage.

Also, we already must deduce types in the backend since logical SPIR-V requires typed pointers, but LLVM-IR now has opaque pointers, so we have some backend mechanism to deduce the pointer types from the IR.

See Keenuts@f2780fe for the current draft of this propagation pass, with the examples we mostly worry about.