Skip to content

Commit 39a6b61

Browse files
TheJeranTheJeran
authored
Merge pull request #551 from EarthyScience/jp/custom-shaders
Custom Shaders
2 parents 7c20841 + 3d25bbb commit 39a6b61

12 files changed

Lines changed: 560 additions & 131 deletions

File tree

‎package.json‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -50,6 +50,7 @@
5050
"@earthyscience/netcdf4-wasm": "^0.1.3",
5151
"@ffmpeg/ffmpeg": "^0.12.15",
5252
"@ffmpeg/util": "^0.12.2",
53+
"@monaco-editor/react": "^4.7.0",
5354
"@radix-ui/react-accordion": "^1.2.12",
5455
"@radix-ui/react-collapsible": "^1.1.12",
5556
"@radix-ui/react-dialog": "^1.1.15",

‎src/GlobalStates/AnalysisStore.ts‎

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@ const ESDC = 'https://s3.bgc-jena.mpg.de:9000/esdl-esdc-v3.0.2/esdc-16d-2.5deg-4
55
type AnalysisState = {
66
analysisMode: boolean;
77
axis: number;
8+
reduceOnAxis: boolean;
89
operation: string;
910
execute: boolean;
1011
useTwo: boolean;
@@ -17,9 +18,14 @@ type AnalysisState = {
1718
reverseDirection: number;
1819
analysisStore: string;
1920
analysisDim: number | null;
21+
customShader: string | undefined;
22+
useEditor: boolean;
23+
executeCustom: boolean;
24+
outputShape: number[];
2025

2126
setAnalysisMode: (analysisMode: boolean) => void;
2227
setAxis: (axis: number) => void;
28+
setReduceOnAxis: (reduceOnAxis: boolean) => void;
2329
setOperation: (operation: string) => void;
2430
setExecute: (execute: boolean) => void;
2531
setUseTwo: (useTwo: boolean) => void;
@@ -32,11 +38,14 @@ type AnalysisState = {
3238
setReverseDirection: (reverseDirection: number) => void;
3339
setAnalysisStore: (analysisStore: string) => void;
3440
setAnalysisDim: (analysisDim: number | null) => void;
41+
setCustomShader: (customShader: string | undefined) => void;
42+
setOutPutShape: (outputShape: number[]) => void;
3543
}
3644

3745
export const useAnalysisStore = create<AnalysisState>((set) => ({
3846
analysisMode: false,
3947
axis: 0,
48+
reduceOnAxis: false,
4049
operation: "Default",
4150
execute: false,
4251
useTwo: false,
@@ -49,9 +58,14 @@ export const useAnalysisStore = create<AnalysisState>((set) => ({
4958
reverseDirection: 0,
5059
analysisStore: ESDC,
5160
analysisDim: null,
61+
customShader: undefined,
62+
useEditor: false,
63+
executeCustom:false,
64+
outputShape: [],
5265

5366
setAnalysisMode: (analysisMode) => set({ analysisMode }),
5467
setAxis: (axis) => set({ axis }),
68+
setReduceOnAxis: (reduceOnAxis) => set({ reduceOnAxis }),
5569
setOperation: (operation) => set({ operation }),
5670
setExecute: (execute) => set({ execute }),
5771
setUseTwo: (useTwo) => set({ useTwo}),
@@ -64,4 +78,6 @@ export const useAnalysisStore = create<AnalysisState>((set) => ({
6478
setReverseDirection: (reverseDirection) => set( { reverseDirection} ),
6579
setAnalysisStore: (analysisStore) => set({ analysisStore }),
6680
setAnalysisDim: (analysisDim) => set({ analysisDim }),
81+
setCustomShader: (customShader) => set({ customShader }),
82+
setOutPutShape: (outputShape) => set({ outputShape }),
6783
}));

‎src/components/computation/WGSLShaders.ts‎

Lines changed: 115 additions & 111 deletions
Original file line numberDiff line numberDiff line change
@@ -7,132 +7,136 @@ export const createShaders = (precision: Precision) => {
77
const enableF16 = precision === 'f16' ? 'enable f16;' : '';
88

99
// #region BOILERPLATES
10-
10+
//They are untabbed for formatting reasons in the editor
1111
const ReductionBoilerPlate = /* WGSL */`
12-
${enableF16}
13-
struct Params {
14-
zStride: u32,
15-
yStride: u32,
16-
xStride: u32,
17-
xSize: u32,
18-
ySize: u32,
19-
reduceDim: u32,
20-
dimLength: u32,
21-
};
22-
@group(0) @binding(0) var<storage, read> inputData: array<${precision}>;
23-
@group(0) @binding(1) var<storage, read_write> outputData: array<${precision}>;
24-
@group(0) @binding(2) var<uniform> params: Params;
25-
26-
@compute @workgroup_size(16, 16, 1)
27-
fn main(@builtin(global_invocation_id) global_id: vec3<u32>) {
28-
let zStride = params.zStride;
29-
let yStride = params.yStride;
30-
let xStride = params.xStride;
31-
let xSize = params.xSize;
32-
let ySize = params.ySize;
33-
let reduceDim = params.reduceDim;
34-
let dimLength = params.dimLength;
35-
36-
let outX = global_id.y;
37-
let outY = global_id.x;
38-
39-
if (outX >= xSize || outY >= ySize) {
40-
return;
41-
}
12+
${enableF16}
13+
struct Params {
14+
zStride: u32,
15+
yStride: u32,
16+
xStride: u32,
17+
xSize: u32,
18+
ySize: u32,
19+
reduceDim: u32,
20+
dimLength: u32,
21+
};
22+
@group(0) @binding(0) var<storage, read> inputData: array<${precision}>;
23+
@group(0) @binding(1) var<storage, read_write> outputData: array<${precision}>;
24+
@group(0) @binding(2) var<uniform> params: Params;
25+
26+
@compute @workgroup_size(16, 16, 1)
27+
fn main(@builtin(global_invocation_id) global_id: vec3<u32>) {
28+
let zStride = params.zStride;
29+
let yStride = params.yStride;
30+
let xStride = params.xStride;
31+
let xSize = params.xSize;
32+
let ySize = params.ySize;
33+
let reduceDim = params.reduceDim;
34+
let dimLength = params.dimLength;
35+
36+
let outX = global_id.y;
37+
let outY = global_id.x;
38+
39+
if (outX >= xSize || outY >= ySize) {
40+
return;
41+
}
4242
`
4343
const ConvolutionBoilerPlate = /* WGSL */`
44-
${enableF16}
45-
struct Params {
46-
xStride: u32,
47-
yStride: u32,
48-
zStride: u32,
49-
xSize: u32,
50-
ySize: u32,
51-
zSize: u32,
52-
workGroups: vec3<u32>,
53-
kernelSize: u32,
54-
kernelDepth: u32
55-
};
56-
@group(0) @binding(0) var<storage, read> inputData: array<${precision}>;
57-
@group(0) @binding(1) var<storage, read_write> outputData: array<${precision}>;
58-
@group(0) @binding(2) var<uniform> params: Params;
59-
60-
@compute @workgroup_size(4, 4, 4)
61-
fn main(@builtin(global_invocation_id) global_id: vec3<u32>) {
62-
let zStride = params.zStride;
63-
let yStride = params.yStride;
64-
let xStride = params.xStride;
65-
let xSize = params.xSize;
66-
let ySize = params.ySize;
67-
let zSize = params.zSize;
68-
let workGroups = params.workGroups;
69-
let kernelSize = params.kernelSize;
70-
let kernelDepth = params.kernelDepth;
71-
72-
let outX = global_id.x;
73-
let outY = global_id.y;
74-
let outZ = global_id.z;
75-
76-
if (outX >= xSize || outY >= ySize || outZ >= zSize) {
77-
return;
78-
}
44+
${enableF16}
45+
struct Params {
46+
xStride: u32,
47+
yStride: u32,
48+
zStride: u32,
49+
xSize: u32,
50+
ySize: u32,
51+
zSize: u32,
52+
workGroups: vec3<u32>,
53+
kernelSize: u32,
54+
kernelDepth: u32
55+
};
56+
@group(0) @binding(0) var<storage, read> inputData: array<${precision}>;
57+
@group(0) @binding(1) var<storage, read_write> outputData: array<${precision}>;
58+
@group(0) @binding(2) var<uniform> params: Params;
59+
60+
@compute @workgroup_size(4, 4, 4)
61+
fn main(@builtin(global_invocation_id) global_id: vec3<u32>) {
62+
let zStride = params.zStride;
63+
let yStride = params.yStride;
64+
let xStride = params.xStride;
65+
let xSize = params.xSize;
66+
let ySize = params.ySize;
67+
let zSize = params.zSize;
68+
let workGroups = params.workGroups;
69+
let kernelSize = params.kernelSize;
70+
let kernelDepth = params.kernelDepth;
71+
72+
let outX = global_id.x;
73+
let outY = global_id.y;
74+
let outZ = global_id.z;
75+
76+
if (outX >= xSize || outY >= ySize || outZ >= zSize) {
77+
return;
78+
}
7979
80-
let total_threads_per_slice = workGroups.x * workGroups.y * 16;
81-
let globalIdx = global_id.z * total_threads_per_slice +
82-
global_id.y * (workGroups.x * 4) +
83-
global_id.x;
80+
let total_threads_per_slice = workGroups.x * workGroups.y * 16;
81+
let globalIdx = global_id.z * total_threads_per_slice +
82+
global_id.y * (workGroups.x * 4) +
83+
global_id.x;
8484
85-
let xy_radius: i32 = i32(kernelSize/2u);
86-
let z_radius: i32 = i32(kernelDepth/2u);
85+
let xy_radius: i32 = i32(kernelSize/2u);
86+
let z_radius: i32 = i32(kernelDepth/2u);
8787
88-
let xy_start: i32 = select(-xy_radius, 0, kernelSize == 1u);
89-
let xy_end: i32 = select(xy_radius + 1, 1, kernelSize == 1u);
90-
let z_start: i32 = select(-z_radius, 0, kernelDepth == 1u);
91-
let z_end: i32 = select(z_radius + 1, 1, kernelDepth == 1u);
88+
let xy_start: i32 = select(-xy_radius, 0, kernelSize == 1u);
89+
let xy_end: i32 = select(xy_radius + 1, 1, kernelSize == 1u);
90+
let z_start: i32 = select(-z_radius, 0, kernelDepth == 1u);
91+
let z_end: i32 = select(z_radius + 1, 1, kernelDepth == 1u);
9292
`
9393
const ConvolutionBoilerPlate2D = /* WGSL */`
94-
${enableF16}
95-
struct Params {
96-
xStride: u32,
97-
yStride: u32,
98-
xSize: u32,
99-
ySize: u32,
100-
kernelSize: u32,
101-
kernelDepth: u32
102-
};
103-
@group(0) @binding(0) var<storage, read> inputData: array<${precision}>;
104-
@group(0) @binding(1) var<storage, read_write> outputData: array<${precision}>;
105-
@group(0) @binding(2) var<uniform> params: Params;
106-
107-
@compute @workgroup_size(16, 16, 1)
108-
fn main(@builtin(global_invocation_id) global_id: vec3<u32>,) {
109-
let xStride = params.xStride;
110-
let yStride = params.yStride;
111-
let xSize = params.xSize;
112-
let ySize = params.ySize;
113-
let kernelSize = params.kernelSize;
114-
115-
let outX = global_id.x;
116-
let outY = global_id.y;
117-
118-
if (outX >= xSize|| outY >= ySize) {
119-
return;
120-
}
94+
${enableF16}
95+
struct Params {
96+
xStride: u32,
97+
yStride: u32,
98+
xSize: u32,
99+
ySize: u32,
100+
kernelSize: u32,
101+
};
102+
@group(0) @binding(0) var<storage, read> inputData: array<${precision}>;
103+
@group(0) @binding(1) var<storage, read_write> outputData: array<${precision}>;
104+
@group(0) @binding(2) var<uniform> params: Params;
105+
106+
@compute @workgroup_size(16, 16, 1)
107+
fn main(@builtin(global_invocation_id) global_id: vec3<u32>,) {
108+
let xStride = params.xStride;
109+
let yStride = params.yStride;
110+
let xSize = params.xSize;
111+
let ySize = params.ySize;
112+
let kernelSize = params.kernelSize;
113+
114+
let outX = global_id.x;
115+
let outY = global_id.y;
116+
117+
if (outX >= xSize|| outY >= ySize) {
118+
return;
119+
}
121120
122-
let globalIdx = outY * xSize + outX;
123-
let thisVal = inputData[globalIdx];
124-
let isNaN: bool = thisVal != thisVal;
125-
if (isNaN){
126-
outputData[globalIdx] = thisVal;
127-
return;
128-
}
121+
let globalIdx = outY * xSize + outX;
122+
let thisVal = inputData[globalIdx];
123+
let isNaN: bool = thisVal != thisVal;
124+
if (isNaN){
125+
outputData[globalIdx] = thisVal;
126+
return;
127+
}
129128
130-
let xy_radius: i32 = i32(kernelSize/2u);
129+
let xy_radius: i32 = i32(kernelSize/2u);
131130
132131
`
133132
// #endregion
134133

135134
const allShaders = {
135+
boilerPlates:{
136+
ReductionBoilerPlate,
137+
ConvolutionBoilerPlate,
138+
ConvolutionBoilerPlate2D
139+
},
136140
// #region REDUCTION SHADERS
137141
MeanReduction: /* wgsl */`
138142
${ReductionBoilerPlate}

0 commit comments

Comments
 (0)