Skip to content

Commit 1c899b1

Browse files
scanbergscanberg
committed
gto fixes
1 parent 79670d4 commit 1c899b1

File tree

3 files changed

+54
-22
lines changed

3 files changed

+54
-22
lines changed

‎src/md_gto.c‎

Lines changed: 43 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -123,6 +123,22 @@ static GLuint get_gto_density_program(void) {
123123
return program;
124124
}
125125

126+
static GLuint get_gto_density_grad_program(void) {
127+
static GLuint program = 0;
128+
if (!program) {
129+
GLuint shader = glCreateShader(GL_COMPUTE_SHADER);
130+
if (md_gl_shader_compile(shader, (str_t) { (const char*)eval_gto_density_grad_comp, eval_gto_density_grad_comp_size }, 0, 0)) {
131+
GLuint prog = glCreateProgram();
132+
if (md_gl_program_attach_and_link(prog, &shader, 1)) {
133+
program = prog;
134+
}
135+
}
136+
glDeleteShader(shader);
137+
}
138+
return program;
139+
}
140+
141+
126142
static GLuint get_buffer(size_t size) {
127143
GLuint id = 0;
128144
glCreateBuffers(1, &id);
@@ -347,7 +363,7 @@ void md_gto_grid_evaluate_GPU(uint32_t vol_tex, const md_grid_t* vol_grid, const
347363
md_gto_grid_evaluate_orb_GPU(vol_tex, vol_grid, &orb, mode);
348364
}
349365

350-
void md_gto_grid_evaluate_matrix_GPU(uint32_t vol_tex, const md_grid_t* grid, const md_gto_data_t* gto_data, const float* upper_triangular_matrix_data, size_t matrix_dim) {
366+
void md_gto_grid_evaluate_matrix_GPU(uint32_t vol_tex, const md_grid_t* grid, const md_gto_data_t* gto_data, const float* upper_triangular_matrix_data, size_t matrix_dim, bool output_gradients_and_value) {
351367
ASSERT(grid);
352368
ASSERT(gto_data);
353369
ASSERT(upper_triangular_matrix_data);
@@ -364,17 +380,33 @@ void md_gto_grid_evaluate_matrix_GPU(uint32_t vol_tex, const md_grid_t* grid, co
364380
glGetTexLevelParameteriv(GL_TEXTURE_3D, 0, GL_TEXTURE_INTERNAL_FORMAT, &format);
365381
glBindTexture(GL_TEXTURE_3D, 0);
366382

367-
switch (format) {
368-
case GL_R16F:
369-
case GL_R32F:
370-
break;
371-
default:
372-
// Not good
373-
MD_LOG_ERROR("Unrecognized internal format of supplied volume texture");
374-
goto done;
383+
384+
GLuint program = 0;
385+
if (output_gradients_and_value) {
386+
switch (format) {
387+
case GL_RGBA16F:
388+
case GL_RGBA32F:
389+
break;
390+
default:
391+
// Not good
392+
MD_LOG_ERROR("Unrecognized internal format of supplied volume texture");
393+
goto done;
394+
}
395+
program = get_gto_density_grad_program();
396+
}
397+
else {
398+
switch (format) {
399+
case GL_R16F:
400+
case GL_R32F:
401+
break;
402+
default:
403+
// Not good
404+
MD_LOG_ERROR("Unrecognized internal format of supplied volume texture");
405+
goto done;
406+
}
407+
program = get_gto_density_program();
375408
}
376409

377-
GLuint program = get_gto_density_program();
378410
if (!program) {
379411
MD_LOG_ERROR("Program not found?!");
380412
goto done;
@@ -507,7 +539,7 @@ void md_gto_grid_evaluate_GPU(uint32_t vol_tex, const md_grid_t* vol_grid, const
507539

508540
}
509541

510-
void md_gto_grid_evaluate_matrix_GPU(uint32_t vol_tex, const md_grid_t* grid, const md_gto_data_t* gto_data, const float* matrix_data, size_t matrix_dim) {
542+
void md_gto_grid_evaluate_matrix_GPU(uint32_t vol_tex, const md_grid_t* grid, const md_gto_data_t* gto_data, const float* matrix_data, size_t matrix_dim, bool output_grad_and_value) {
511543

512544
}
513545

‎src/md_gto.h‎

Lines changed: 6 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -111,18 +111,17 @@ void md_gto_grid_evaluate(float* out_grid_values, const md_grid_t* grid, const m
111111
// - out_grid_values: The grid to write the evaluated values to, should have length 'grid->dim[0] * grid->dim[1] * grid->dim[2]'
112112
// - grid: The grid defining the samples to evaluate
113113
// - gto_data: The gto data to evaluate
114-
// - matrix: The matrix coefficients (This is an upper triangular matrix of size num_cgtos x num_cgtos, where num_cgtos is the number of CGTOs in
115-
// gto_data)
116-
// - mode: evaluation mode
117-
void md_gto_grid_evaluate_matrix(float* out_grid_values, const md_grid_t* grid, const md_gto_data_t* gto_data, const float* matrix_data);
114+
// - matrix: The matrix coefficients (Upper triangular format)
115+
// - matrix_dim: The dimension of the square matrix (should match the number of CGTOs in gto_data)
116+
void md_gto_grid_evaluate_matrix(float* out_grid_values, const md_grid_t* grid, const md_gto_data_t* gto_data, const float* matrix_data, size_t matrix_dim);
118117

119118
// Evaluates CGTOs defined by gto_data with a given 'density' matrix
120119
// - vol_tex: The texture handle to the volume
121120
// - grid: The grid defining the location of samples to evaluate
122121
// - gto_data: The gto data to evaluate
123-
// - matrix: The matrix coefficients
124-
// - mode: evaluation mode
125-
void md_gto_grid_evaluate_matrix_GPU(uint32_t vol_tex, const md_grid_t* grid, const md_gto_data_t* gto_data, const float* matrix_data, size_t matrix_dim);
122+
// - matrix_data: The matrix coefficients (Upper triangular format)
123+
// - matrix_dim: The dimension of the square matrix (should match the number of CGTOs in gto_data)
124+
void md_gto_grid_evaluate_matrix_GPU(uint32_t vol_tex, const md_grid_t* grid, const md_gto_data_t* gto_data, const float* matrix_data, size_t matrix_dim, bool output_grad_and_value);
126125

127126
// Evaluate GTOs over subportion of a grid
128127
// - out_grid_values: The grid to write the evaluated values to, should have length 'grid->dim[0] * grid->dim[1] * grid->dim[2]'

‎src/shaders/eval_gto_density_grad.comp‎

Lines changed: 5 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -40,6 +40,7 @@ layout(std140, binding = 0) uniform UniformBlock {
4040
mat4 index_to_world;
4141
vec4 step;
4242
uint D_matrix_dim;
43+
uint operator;
4344
};
4445

4546
layout(binding = 0) writeonly restrict uniform image3D out_vol;
@@ -322,7 +323,7 @@ void main() {
322323
vec4 phi_tile_mu[TILE_SIZE]; // vec4(φ_μ, ∇φ_μ) in registers
323324
vec4 phi_tile_nu[TILE_SIZE]; // vec4(φ_ν, ∇φ_ν) in registers
324325
vec3 coord = (index_to_world * vec4(gl_GlobalInvocationID.xyz, 1.0)).xyz;
325-
vec4 rho_grad = vec4(0.0); // (ρ, ∇ρ_x, ∇ρ_y, ∇ρ_z)
326+
vec4 grad_rho = vec4(0.0); // (∇ρ_x, ∇ρ_y, ∇ρ_z, ρ)
326327

327328
if (num_screened_cgtos == 0u) {
328329
if (all(lessThan(gl_GlobalInvocationID.xyz, imageSize(out_vol)))) {
@@ -347,7 +348,7 @@ void main() {
347348

348349
eval_phis(phi_tile_mu, coord);
349350

350-
rho_grad += symmetric_contraction(phi_tile_mu, D_tile);
351+
grad_rho += symmetric_contraction(phi_tile_mu, D_tile);
351352

352353
for (uint tile_j = tile_i + 1; tile_j < num_tiles; ++tile_j) {
353354
// OFF DIAGONAL TILE
@@ -364,12 +365,12 @@ void main() {
364365

365366
eval_phis(phi_tile_nu, coord);
366367

367-
rho_grad += contraction(phi_tile_mu, phi_tile_nu, D_tile);
368+
grad_rho += contraction(phi_tile_mu, phi_tile_nu, D_tile);
368369
}
369370
}
370371

371372
// Write result: vec4(∇ρ_x, ∇ρ_y, ∇ρ_z, ρ)
372373
if (all(lessThan(gl_GlobalInvocationID.xyz, imageSize(out_vol)))) {
373-
imageStore(out_vol, ivec3(gl_GlobalInvocationID.xyz), rho_grad);
374+
imageStore(out_vol, ivec3(gl_GlobalInvocationID.xyz), grad_rho);
374375
}
375376
}

0 commit comments

Comments
 (0)