1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
//! Compute pipeline setup
use super::GpuContext;
use anyhow::Result;
use std::sync::Arc;
use tracing::info;
use wgpu::{BindGroupLayout, ComputePipeline};
/// Kangaroo compute pipeline (Clone is cheap - wgpu types are Arc-wrapped)
#[derive(Clone)]
pub struct KangarooPipeline {
pub pipeline: Arc<ComputePipeline>,
pub bind_group_layout: Arc<BindGroupLayout>,
}
impl KangarooPipeline {
/// Create the compute pipeline
pub fn new(ctx: &GpuContext) -> Result<Self> {
info!("Loading shader sources...");
// Use shared shaders from internal gpu_crypto module
let field = crate::gpu_crypto::shaders::FIELD_WGSL;
let curve = crate::gpu_crypto::shaders::CURVE_JACOBIAN_WGSL;
let kangaroo = include_str!("../shaders/kangaroo_jacobian.wgsl");
info!("Creating shader module...");
let shader = ctx.create_shader_module("Kangaroo Shader", &[field, curve, kangaroo]);
info!("Shader module created");
info!("Creating bind group layout...");
// Create bind group layout
let bind_group_layout =
ctx.device
.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
label: Some("Kangaroo Bind Group Layout"),
entries: &[
// Config (uniform)
wgpu::BindGroupLayoutEntry {
binding: 0,
visibility: wgpu::ShaderStages::COMPUTE,
ty: wgpu::BindingType::Buffer {
ty: wgpu::BufferBindingType::Uniform,
has_dynamic_offset: false,
min_binding_size: None,
},
count: None,
},
// Jump points (storage, read_only)
wgpu::BindGroupLayoutEntry {
binding: 1,
visibility: wgpu::ShaderStages::COMPUTE,
ty: wgpu::BindingType::Buffer {
ty: wgpu::BufferBindingType::Storage { read_only: true },
has_dynamic_offset: false,
min_binding_size: None,
},
count: None,
},
// Jump distances (storage, read_only)
wgpu::BindGroupLayoutEntry {
binding: 2,
visibility: wgpu::ShaderStages::COMPUTE,
ty: wgpu::BindingType::Buffer {
ty: wgpu::BufferBindingType::Storage { read_only: true },
has_dynamic_offset: false,
min_binding_size: None,
},
count: None,
},
// Kangaroos (storage, read_write)
wgpu::BindGroupLayoutEntry {
binding: 3,
visibility: wgpu::ShaderStages::COMPUTE,
ty: wgpu::BindingType::Buffer {
ty: wgpu::BufferBindingType::Storage { read_only: false },
has_dynamic_offset: false,
min_binding_size: None,
},
count: None,
},
// DP buffer (storage, read_write)
wgpu::BindGroupLayoutEntry {
binding: 4,
visibility: wgpu::ShaderStages::COMPUTE,
ty: wgpu::BindingType::Buffer {
ty: wgpu::BufferBindingType::Storage { read_only: false },
has_dynamic_offset: false,
min_binding_size: None,
},
count: None,
},
// DP count (storage, read_write atomic)
wgpu::BindGroupLayoutEntry {
binding: 5,
visibility: wgpu::ShaderStages::COMPUTE,
ty: wgpu::BindingType::Buffer {
ty: wgpu::BufferBindingType::Storage { read_only: false },
has_dynamic_offset: false,
min_binding_size: None,
},
count: None,
},
],
});
info!("Bind group layout created");
info!("Creating pipeline layout...");
let pipeline_layout = ctx
.device
.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
label: Some("Kangaroo Pipeline Layout"),
bind_group_layouts: &[&bind_group_layout],
push_constant_ranges: &[],
});
info!("Pipeline layout created");
info!("Creating compute pipeline...");
let pipeline = ctx
.device
.create_compute_pipeline(&wgpu::ComputePipelineDescriptor {
label: Some("Kangaroo Compute Pipeline"),
layout: Some(&pipeline_layout),
module: &shader,
entry_point: Some("main"),
compilation_options: Default::default(),
cache: None,
});
info!("Compute pipeline created");
Ok(Self {
pipeline: Arc::new(pipeline),
bind_group_layout: Arc::new(bind_group_layout),
})
}
}