Add support for returning structures that contain opaque types (shader-slang#1835)

Introduction
============

Several of our target platforms share a concept of "opaque" types, including resources (`Texture2D`) and samplers (`SamplerState`), which are restricted in how they can be used. GLSL and SPIR-V place very severe restrictions, in that opaque types cannot be used for the type of:

* (mutable) local variables
* (mutable) global variables
* structure fields
* Function result/return
* `out` or `inout` parameters

The HLSL language allows all of these cases, but with the practical caveat that the compiler front-end must be able to statically analyze how opaque types have been used and "optimize away" all of the above cases. For example, it is legal to have a local variable of an opaque type, but at any point where the variable gets used it must be statically known which top-level shader parameter the variable refers to.

Existing Work
=============

In the Slang compiler we need to implement our own passes to detect these "illegal" uses of opaque types and legalize them. The work is basically broken into two distinct steps:

* The existing `legalizeResourceTypes()` pass detects illegal types (e.g., a `struct` that has a field of type `Texture2D`) and replaces them with legal types, sometimes by splitting apart declarations (e.g., a parameter using such a `struct` type gets split into multiple parameters). At a high level, we can think of this as "exposing" opaque types so that they are not hidden inside of nested structures.

* Next, the `specializeResourceOutputs()` pass detects calls to functions that output opaque types (whether by the function return value of `out` / `inout` parameters). The pass analyzes the body of such functions, and tries to isolate the logic that determines their resource-type outputs and hoise that logic into call sites (so that the opaque-type outputs can then be eliminated).

This Change
===========

One important missing case was that the type legalization step was incapable of legalizing types that appear in the result/return type of functions. The existing logic would simply diagnose an internal/unimplemented error if it ecountered a non-simple type in the return position.

At a high-level, supporting this case seems simple enough. Given a function signature like:

```
struct Things { int a; Texture2D b; }

Things myFunc(int x) { ... }
```

we want to split the result type into an "ordinary" result type and then `out` parameters for any opaque-type fields:

```
struct Things_Legal { int a; }

Things_Legal myFunc(int x, out Texture2D result_b) { ... };
```

Similarly, at a call site to a function like this:

```
Things t = myFunc(99);
```

we split the function result into ordinary and opaque-type parts, and pass the latter as `out` parameters:

```
Texture2D t_b;
Things_Legal t = myFunc(99, /*out*/ t_b);
```

The main place where things get tricky is when dealing with `return` sites within the body of a function that needs legalization:

```
Things myFunc(int x) {
    ...
    Things things = ...;
    ...
    return things;
}
```

In theory the answer is simple: a `return` translates into writes to the `out` parameters for any opaque-type data, followed by a return of the ordinary-type part:

```
Things_Legal myFunc(int x, out Texture2D result_b) {
    ...
    Things_Legal things = ...;
    Texture2D things_b = ...;
    ...
    result_b = things_b;
    return things;
}
```

The sticking point here is that this step requires tracking data between the legalization of the parameter list for `myFunc` and legalization of the `return`s in its body, so that we can identify the `result_b` parameter to be able to write to it. The existing type legalization pass was not built with the idea that such communication is commonly needed; it assumes that each instruction can be legalized in isolation, so long as dependencies are respected.

This change adds logic such that the `legalizeFunc()` step sets up a data structure that it used to represent information about how a function (and its parameter list) got legalized, so that the logic for a `return` can make use of that legalized information. Right now the information we track consists of just the list of parameters that were introduced to represent a return/result type.

Testing
=======

In order to confirm what features do/don't work, I added a set of tests that cover a cross-product of opaque type use cases:

* The opaque type can be used in the function result type, an `out` parameter, or an `inout` parameter
* The opaque type can be used "directly" or nested inside a `struct`.

These tests are helpful to make sure we handle the most important cases, but it is worth noting that the coverage is still lacking in that we do not sufficiently test all the options for what the function body might do. An opaque-type function result could be derived from many different sources:

* It could be a global shader parameter
* It could be an `in` or `inout` parameter of the function itself
* It could be wrapped up in one or more structure types
* It could be wrapped up in one or more array types (such that the output of specialization needs to pass around array indices)
* It could involve use of the type as a local variable (including passing it into other functions with result/`out`/`inout` outputs of opaque types)

This change makes it so that we can handle the simplest cases involving result/return types with a wrapper `struct`, and adds test cases that confirm we handle several other cases for `out` and `inout` parameters. Gaining confidence that we cover all the cases that arise in practical shaders will require more work over following changes.

Author: Tim Foley

source/slang/slang-ir-legalize-types.cpp

@@ -171,7 +171,7 @@ struct ResourceOutputSpecializationPass
             oldFunc,
             newFunc);
 
-        // At first `newFunc` is a directclone of `oldFunc`, and thus doesn't
+        // At first `newFunc` is a direct clone of `oldFunc`, and thus doesn't
         // solve any of our problems. We will traverse `oldFunc` and specialize
         // it as needed, while also collecting information that will allow
         // us to rewrite call sites.
@@ -468,8 +468,29 @@ struct ResourceOutputSpecializationPass
         //
         // Any failures along the way cause the whole process to fail.
 
-        for( auto param : func->getParams() )
+        // Note: We are introducing new parameters at the same time as we
+        // iterate over the parameter list, so we cannot just use the
+        // `func->getParams()` convenience accessor. Instead, we manually
+        // iterate over the parameters in a way that avoids invalidation
+        // if we remove the `param` we are working on.
+        //
+        // Note: it might seem odd that we are modifying `func` but will
+        // still bail out on any errors. You might ask: isn't there a chance
+        // that we will end up with the function in a partially-modified state?
+        //
+        // The important thing to remember is that `func` is  *copy* of the
+        // original function, so any modifications we make to it do not
+        // affect the original, so that if we *do* have to bail out we can
+        // leave any call sites intact as calls to the original. The result
+        // is that bailing out here may leave the new/copied function in
+        // a state where it isn't useful, but it also won't have any uses,
+        // and can be eliminated later.
+        //
+        IRParam* nextParam = nullptr;
+        for( IRParam* param = func->getFirstParam(); param; param = nextParam )
         {
+            nextParam = param->getNextParam();
+
             ParamInfo paramInfo;
             SLANG_RETURN_ON_FAIL(maybeSpecializeParam(param, paramInfo, outFuncInfo));
             outFuncInfo.oldParams.add(paramInfo);

source/slang/slang-ir-specialize-resources.cpp

@@ -509,7 +509,7 @@ struct LegalVal
     }
 
     static LegalVal implicitDeref(LegalVal const& val);
-    LegalVal getImplicitDeref();
+    LegalVal getImplicitDeref() const;
 
     static LegalVal pair(RefPtr<PairPseudoVal> pairInfo);
     static LegalVal pair(
@@ -566,6 +566,30 @@ struct WrappedBufferPseudoVal : LegalValImpl
     LegalElementWrapping    elementInfo;
 };
 
+//
+
+    /// Information about a function that has been legalized
+    ///
+    /// This type is used to track any information about the function
+    /// and its signature that might be relevant to the legalization
+    /// of instructions inside the function body.
+    ///
+struct LegalFuncInfo : RefObject
+{
+        /// Any parameters that were added to the function signature
+        /// to represent the function result after legalization.
+        ///
+        /// It is possible that the result type of a function needed
+        /// to be split into multiple types, and as a result a single
+        /// function result couldn't return all of them.
+        ///
+        /// This array is a list of `out` parameters created to represent
+        /// additional function results. Because they are `out` parameters,
+        /// each is a *pointer* to a value of the relevant type.
+        ///
+    List<IRInst*> resultParamVals;
+};
+
 //
 
     /// Context that drives type legalization
@@ -601,6 +625,14 @@ struct IRTypeLegalizationContext
 
     Dictionary<IRType*, LegalType> mapTypeToLegalType;
 
+        /// Map a function to information about how it was legalized.
+        ///
+        /// Note that entries are only created if there is somehting for them
+        /// to represent, so many functions may lack entries in this map even
+        /// after legalization.
+        ///
+    Dictionary<IRFunc*, RefPtr<LegalFuncInfo>> mapFuncToInfo;
+
     IRBuilder* getBuilder() { return builder; }
 
         /// Customization point to decide what types are "special."

source/slang/slang-legalize-types.h

@@ -0,0 +1,55 @@
+// inout-param-opaque-type-in-struct.slang
+
+// Test that a function/method can have an `out` parameter of
+// aggregate type that includes an opaque type
+
+//TEST(compute):COMPARE_COMPUTE:
+
+struct Things
+{
+    int first;
+    RWStructuredBuffer<int> rest;
+}
+
+//TEST_INPUT:set C = new { {1, ubuffer(data=[2 3 4 5], stride=4)}, {6, ubuffer(data=[7 8 9 10], stride=4)} }
+cbuffer C
+{
+    Things gX;
+    Things gY;
+}
+
+void swap(
+    inout Things a,
+    inout Things b)
+{
+    Things t = a;
+    a = b;
+    b = t;
+}
+
+int eval(Things t, int val)
+{
+    return t.first*256 + t.rest[val];
+}
+
+int test(int val)
+{
+    Things f = gX;
+    Things g = gY;
+
+    swap(f, g);
+
+    return (eval(f,val) << 16) + eval(g,val);
+}
+
+//TEST_INPUT:set gOutput = out ubuffer(data=[0 0 0 0], stride=4)
+RWStructuredBuffer<int> gOutput;
+
+[numthreads(4, 1, 1)]
+void computeMain(uint3 dispatchThreadID : SV_DispatchThreadID)
+{
+    uint tid = dispatchThreadID.x;
+    int inVal = tid;
+    int outVal = test(inVal);
+    gOutput[tid] = outVal;
+}

tests/language-feature/types/opaque/inout-param-opaque-type-in-struct.slang

@@ -0,0 +1,4 @@
+6070102
+6080103
+6090104
+60A0105

tests/language-feature/types/opaque/inout-param-opaque-type-in-struct.slang.expected.txt

@@ -0,0 +1,42 @@
+// inout-param-opaque-type.slang
+
+// Test that a function/method can have an `out` parameter of opaque type
+
+//TEST(compute):COMPARE_COMPUTE:
+
+//TEST_INPUT:set gX = ubuffer(data=[16 17 18 19], stride=4)
+RWStructuredBuffer<int> gX;
+
+//TEST_INPUT:set gY = ubuffer(data=[3 6 9 12], stride=4)
+RWStructuredBuffer<int> gY;
+
+void swap(
+    inout RWStructuredBuffer<int> a,
+    inout RWStructuredBuffer<int> b)
+{
+    RWStructuredBuffer<int> t = a;
+    a = b;
+    b = t;
+}
+
+int test(int val)
+{
+    RWStructuredBuffer<int> f = gX;
+    RWStructuredBuffer<int> g = gY;
+
+    swap(f, g);
+
+    return f[val] * 256 + g[val];
+}
+
+//TEST_INPUT:set gOutput = out ubuffer(data=[0 0 0 0], stride=4)
+RWStructuredBuffer<int> gOutput;
+
+[numthreads(4, 1, 1)]
+void computeMain(uint3 dispatchThreadID : SV_DispatchThreadID)
+{
+    uint tid = dispatchThreadID.x;
+    int inVal = tid;
+    int outVal = test(inVal);
+    gOutput[tid] = outVal;
+}

tests/language-feature/types/opaque/inout-param-opaque-type.slang

@@ -0,0 +1,4 @@
+310
+611
+912
+C13

tests/language-feature/types/opaque/inout-param-opaque-type.slang.expected.txt

@@ -0,0 +1,39 @@
+// out-opaque-type-in-struct.slang
+
+// Test that a function/method can have an `out` parameter of
+// aggregate type that includes an opaque type
+
+//TEST(compute):COMPARE_COMPUTE:
+
+struct Things
+{
+    int first;
+    RWStructuredBuffer<int> rest;
+}
+
+//TEST_INPUT:set gThings = new Things { 1, ubuffer(data=[2 3 4 5], stride=4) }
+ConstantBuffer<Things> gThings;
+
+void getThings(out Things outThings)
+{
+    outThings = gThings;
+}
+
+int test(int val)
+{
+    Things things;
+    getThings(things);
+    return things.first * (16 << val) + things.rest[val];
+}
+
+//TEST_INPUT:set gOutput = out ubuffer(data=[0 0 0 0], stride=4)
+RWStructuredBuffer<int> gOutput;
+
+[numthreads(4, 1, 1)]
+void computeMain(uint3 dispatchThreadID : SV_DispatchThreadID)
+{
+    uint tid = dispatchThreadID.x;
+    int inVal = tid;
+    int outVal = test(inVal);
+    gOutput[tid] = outVal;
+}

tests/language-feature/types/opaque/out-param-opaque-type-in-struct.slang

@@ -0,0 +1,4 @@
+12
+23
+44
+85

tests/language-feature/types/opaque/out-param-opaque-type-in-struct.slang.expected.txt

@@ -0,0 +1,33 @@
+// out-opaque-type.slang
+
+// Test that a function/method can have an `out` parameter of opaque type
+
+//TEST(compute):COMPARE_COMPUTE:
+
+//TEST_INPUT:set gThings = ubuffer(data=[16 17 18 19], stride=4)
+RWStructuredBuffer<int> gThings;
+
+
+void getThings(out RWStructuredBuffer<int> things)
+{
+    things = gThings;
+}
+
+int test(int val)
+{
+    RWStructuredBuffer<int> t;
+    getThings(t);
+    return t[val];
+}
+
+//TEST_INPUT:set gOutput = out ubuffer(data=[0 0 0 0], stride=4)
+RWStructuredBuffer<int> gOutput;
+
+[numthreads(4, 1, 1)]
+void computeMain(uint3 dispatchThreadID : SV_DispatchThreadID)
+{
+    uint tid = dispatchThreadID.x;
+    int inVal = tid;
+    int outVal = test(inVal);
+    gOutput[tid] = outVal;
+}

tests/language-feature/types/opaque/out-param-opaque-type.slang

@@ -0,0 +1,4 @@
+10
+11
+12
+13

tests/language-feature/types/opaque/out-param-opaque-type.slang.expected.txt

@@ -0,0 +1,38 @@
+// return-opaque-type-in-struct.slang
+
+// Test that a function/method can return a value of
+// aggregate type that includes an opaque type
+
+//TEST(compute):COMPARE_COMPUTE:
+
+struct Things
+{
+    int first;
+    RWStructuredBuffer<int> rest;
+}
+
+//TEST_INPUT:set gThings = new Things { 1, ubuffer(data=[2 3 4 5], stride=4) }
+ConstantBuffer<Things> gThings;
+
+Things getThings()
+{
+    return gThings;
+}
+
+int test(int val)
+{
+    let things = getThings();
+    return things.first * (16 << val) + things.rest[val];
+}
+
+//TEST_INPUT:set gOutput = out ubuffer(data=[0 0 0 0], stride=4)
+RWStructuredBuffer<int> gOutput;
+
+[numthreads(4, 1, 1)]
+void computeMain(uint3 dispatchThreadID : SV_DispatchThreadID)
+{
+    uint tid = dispatchThreadID.x;
+    int inVal = tid;
+    int outVal = test(inVal);
+    gOutput[tid] = outVal;
+}

tests/language-feature/types/opaque/return-opaque-type-in-struct.slang

@@ -0,0 +1,4 @@
+12
+23
+44
+85

tests/language-feature/types/opaque/return-opaque-type-in-struct.slang.expected.txt

@@ -0,0 +1,32 @@
+// return-opaque-type.slang
+
+// Test that a function/method can return a value of an opaque type.
+
+//TEST(compute):COMPARE_COMPUTE:
+
+struct Stuff
+{
+    RWStructuredBuffer<int> things;
+
+    RWStructuredBuffer<int> getThings() { return things; }
+}
+
+//TEST_INPUT:set gStuff = new Stuff { ubuffer(data=[16 17 18 19], stride=4) }
+ConstantBuffer<Stuff> gStuff;
+
+int test(int val)
+{
+    return gStuff.getThings()[val];
+}
+
+//TEST_INPUT:set gOutput = out ubuffer(data=[0 0 0 0], stride=4)
+RWStructuredBuffer<int> gOutput;
+
+[numthreads(4, 1, 1)]
+void computeMain(uint3 dispatchThreadID : SV_DispatchThreadID)
+{
+    uint tid = dispatchThreadID.x;
+    int inVal = tid;
+    int outVal = test(inVal);
+    gOutput[tid] = outVal;
+}

tests/language-feature/types/opaque/return-opaque-type.slang

@@ -0,0 +1,4 @@
+10
+11
+12
+13