Proposal: Assembly backed opaque types

alichraghi commented 2 months ago

Motivation

Currently, SPIR-V backend has no way distinguish decorated types (in SPIR-V, decoration instructions modifies types/variables/functions/etc attributes). for example, in the below code A and B are the same type therefore the decoration in arrayStride is written twice which produces an invalid module.

fn main() void {
    const A = [3]u32;
    asm volatile (
        \\OpDecorate %T ArrayStride 10
        :
        : [T] "" (A),
    );

    // Somewhere else
    const B = [3]u32;
    asm volatile (
        \\OpDecorate %T ArrayStride 20
        :
        : [T] "" (B),
    );
}

fn arrayStride(comptime T: type, comptime stride: u32) void {
    asm volatile (comptimePrint("OpDecorate %T ArrayStride {}", .{stride})
        :
        : [T] "" (T),
    );
}

           %u32 = OpTypeInt 32 0
             %3 = OpConstant %u32 3
%array_of_3_u32 = OpTypeArray %u32 %3
     OpDecorate %array_of_3_u32 ArrayStride 10
     OpDecorate %array_of_3_u32 ArrayStride 20

Proposal

There are two most straightforward ways to fix this. either extending the zig syntax to allow something like [3]stride(10) u32, or using inline assembly to create opaque types. [3]stride(10) u32 is more pleasing of course, but as there are quite a lot of these decorations, extending the language for every one of them is a bad idea, and they may evolve over time as well. As for inline assembly, after discussions with @Snektron, we reached the conclusion that these types would semantically be an opaque type backed by the assembly written in a typed backend like SPIR-V. so the above example would become:

fn main() void {
    const A = opaque(asm(
      \\        %3 = OpConstant %u32 3
      \\        %A = OpTypeArray %u32 %3
      \\OpDecorate %A ArrayStride 10
      : [array_of_3_u32] "" (->type)
      : [u32] "" (u32),
    )) {
       // Methods goes here
    };

    const B = opaque(asm(
      \\        %3 = OpConstant %u32 3
      \\        %B = OpTypeArray %u32 %3
      \\OpDecorate %B ArrayStride 10
      : [B] "" (->type)
      : [u32] "" (u32),
    )) {
       // Methods goes here
    };
}

You might have also noticed that the types are created inside a function. the reason is global assembly won't allow I/O but that should probably be discussed in another proposal.

Error Messages

The syntax is useless and the behavior is undefined for other backends so perhaps it should emit an error in any non-typed backend (all backends minus SPIR-V currently)

andrewrk commented 1 month ago

You might have also noticed that the types are created inside a function. the reason is global assembly won't allow I/O but that should probably be discussed in another proposal.

Can you explain this a little bit more? What does assembly-backed opaque types have to do with I/O?

Snektron commented 1 month ago

I think it is worded like that too keep this separate from a proposal to allow more I/O for assembly in global contexts, which is currently not allowed I think.

andrewrk commented 1 month ago

Unaccepting for further consideration since @Snektron expressed interest in making a counter proposal.

Snektron commented 1 month ago

Background

SPIR-V is a strongly typed IR format, like LLVM IR especially before LLVM removed pointee types. Most operations need to be ascribed with their result type. Like LLVM IR, it supports the primitve types you'd expect (ints, floats, bools, ...) and also provide the building blocks for composite and custom types (pointers, structs, ...).

The problem is that SPIR-V has additional built-in types for GPU-specific things, for example samplers and images. These are required in practical code, too, for example for textures. When returning a type that is or contains for example an image to there, the generated SPIR-V code needs to refer to that image.

For OpenCL, this is not much of a problem: Images are usually supplied as pointers (for example set as a uniform), which we can cast to an opaque type. In this case, we can deal with most of these types from inline assembly. This is not the case for Vulkan, where we are not allowed to perform such casts. Here we need a way to refer to these types from Zig shader code. While I doubt the use case of actually returning images (for example...), a way to refer to these from builtins or function parameters is vital.

Additionally, SPIR-V supports optional decorations on types that change them in a non-trivial way: For example the ArrayStride decoration can be applied to pointers to set the byte-stride. Currently I'm not completely convincend that these are 100% required from Zig.

Required SPIR-V Types

Concretely, there should be a way to refer to the following types. Some of these are parameterized, see the linked SPIR-V spec for the parameters that they require. Not all of these are immediately of interest.

Note that as far as I can tell, all of these types are opaque. While the OpenCL spec allows implementations to choose a size or sizeof(image2d_t), clang actually throws an error when you try to use the type in a place where its required to be concrete.

Some types are intended to still be used without pointer. For example, to use a cooperative matrix in pseudo code:

cooperative_matrix a = load_matrix(A);
cooperative_matrix b = load_matrix(B);
cooperative_matrix c = cooperative_matmul(a, b);
store_matrix(C, c);

Even though the cooperative matrix is an opaque type, they are still allowed to be used like this. This is similar to bools.

Snektron commented 1 month ago

The counterproposal, or rather nullhypothesis, is to simply add the required types as Zig builtins. The main advantage is that its a little bit more straight forward to implement. The drawbacks are that its a lot of baggage to keep around for all the other backends, I don't like that. Its also a lot less flexibly, in contrast to the original proposal. After some additional thought I still think that the original proposal works better.

Snektron commented 1 month ago

One other alternative I've thought of is a special SPIR-V assembly instruction that allows us to "override" a type during linking. We'd specify a dummy type during codegen and then patch it out. For example:

// This creates `%Image = OpTypeOpaque`
const Image = opaque {
  pub fn whatever(self: *Image) void {
    // self is %Self = OpTypePointer %Image Generic
    asm volatile (
      \\OpDoSomething %Image %self ...
      :: [Image] "" (Image), [self] "" (self)
    );
  }
};

comptime {
  asm (
    \\%ActualImage = OpTypeImage ...
    // Replaces all uses of %Image with %ActualImage, then deletes the %Image declaration from the module.
    \\OpZigReplaceAllUses %Image %ActualImage
    :: [Image] "" (Image)
  );
}

The main advantage of this is that we can also create non-opaque types should we need it. Its also easier to implement, and doesn't require any language extensions. Just a temporary SPIR-V instruction that is removed after linking. It does, however, feel like a much bigger hack.

ziglang / zig