yum-mirror/slang

Making it easier to work with shaders

git clone https://git.yummers.dev/yum-mirror/slang

Yong HeSupport visibility control and default to `internal`. (#3380)11111e573

master
branchesmasterall refs →
1.5 KiB53 linesraw
1//TEST:SIMPLE(filecheck=CUDA): -target cuda -line-directive-mode none
2//TEST:SIMPLE(filecheck=TORCH): -target torch -line-directive-mode none
3
4// Verify that we can output a cuda device function with [CudaKernel].
5
6struct MySubType
7{
8    TorchTensor<float> array[2];
9}
10
11struct MyType
12{
13    float2 v;
14    MySubType sub[2];
15}
16
17struct MyInput
18{
19    TorchTensor<float> inValues;
20    float normalVal;
21}
22
23// CUDA: __global__ void myKernel(TensorView inValues_[[#]], TensorView outValues_[[#]])
24[CudaKernel]
25void myKernel(TensorView<float> inValues, TensorView<float> outValues)
26{
27    if (cudaThreadIdx().x > 0)
28        return;
29    outValues.store(cudaThreadIdx().x, sin(inValues.load(cudaThreadIdx().x)));
30}
31
32// TORCH:      {{^SLANG_PRELUDE_EXPORT$}}
33// TORCH-NEXT: void myKernel(TensorView {{[[:alnum:]_]+}}, TensorView {{[[:alnum:]_]+}});
34//
35// TORCH:      {{^SLANG_PRELUDE_EXPORT$}}
36// TORCH-NEXT: std::tuple<std::tuple<float, float>, std::tuple<std::tuple<std::tuple<torch::Tensor, torch::Tensor>>, std::tuple<std::tuple<torch::Tensor, torch::Tensor>>>> runCompute(std::tuple<torch::Tensor, float> input_[[#]])
37[TorchEntryPoint]
38export __extern_cpp MyType runCompute(MyInput input)
39{
40    MyType rs;
41    var outValues = TorchTensor<float>.alloc(1);
42    let inValues = input.inValues;
43    
44    __dispatch_kernel(myKernel, uint3(1, 1, 1), uint3(32, 1, 1))(inValues, outValues);
45
46    rs.v = float2(1.0, 2.0);
47    rs.sub[0].array[0] = outValues;
48    rs.sub[0].array[1] = inValues;
49
50    rs.sub[1].array[0] = inValues;
51    rs.sub[1].array[1] = outValues;
52    return rs;
53}