Skip to content

Commit

Permalink
[opengl] RW image binding & FP16 support (taichi-dev#7219)
Browse files Browse the repository at this point in the history
Issue: #

### Brief Summary

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
  • Loading branch information
2 people authored and quadpixels committed May 13, 2023
1 parent 0f29526 commit 7124225
Show file tree
Hide file tree
Showing 2 changed files with 102 additions and 50 deletions.
109 changes: 75 additions & 34 deletions taichi/rhi/opengl/opengl_device.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -218,7 +218,8 @@ GLResourceSet &GLResourceSet::image(uint32_t binding,
GLResourceSet &GLResourceSet::rw_image(uint32_t binding,
DeviceAllocation alloc,
int lod) {
TI_NOT_IMPLEMENTED;
rw_image_binding_map_[binding] = alloc.alloc_id;
return *this;
}

GLPipeline::GLPipeline(const PipelineSourceDesc &desc,
Expand All @@ -244,6 +245,16 @@ GLPipeline::GLPipeline(const PipelineSourceDesc &desc,
}
options.vulkan_semantics = false;
glsl.set_common_options(options);
if (GLAD_GL_NV_gpu_shader5) {
glsl.require_extension("GL_NV_gpu_shader5");
glsl.add_header_line(
"uint16_t float16BitsToUint16(float16_t v) { return "
"uint16_t(packFloat2x16(f16vec2(v, 0))); }");
glsl.add_header_line(
"float16_t uint16BitsToFloat16(uint16_t v) { return "
"unpackFloat2x16(uint(v)).x; }");
}
glsl.add_header_line("");
std::string source = glsl.compile();
TI_TRACE("GLSL source: \n{}", source);

Expand Down Expand Up @@ -315,31 +326,37 @@ void GLCommandList::bind_pipeline(Pipeline *p) noexcept {
RhiResult GLCommandList::bind_shader_resources(ShaderResourceSet *res,
int set_index) noexcept {
GLResourceSet *set = static_cast<GLResourceSet *>(res);
auto cmd = std::make_unique<CmdBindResources>();
for (auto &[binding, buffer] : set->ssbo_binding_map()) {
auto cmd = std::make_unique<CmdBindBufferToIndex>();
cmd->buffer = buffer.buffer;
cmd->offset = buffer.offset;
cmd->size = buffer.size;
cmd->index = binding;
recorded_commands_.push_back(std::move(cmd));
auto &bind = cmd->buffers.emplace_back();
bind.buffer = buffer.buffer;
bind.offset = buffer.offset;
bind.size = buffer.size;
bind.index = binding;
}
for (auto &[binding, buffer] : set->ubo_binding_map()) {
auto cmd = std::make_unique<CmdBindBufferToIndex>();
cmd->buffer = buffer.buffer;
cmd->offset = buffer.offset;
cmd->size = buffer.size;
cmd->index = binding;
cmd->target = GL_UNIFORM_BUFFER;
recorded_commands_.push_back(std::move(cmd));
auto &bind = cmd->buffers.emplace_back();
bind.buffer = buffer.buffer;
bind.offset = buffer.offset;
bind.size = buffer.size;
bind.index = binding;
bind.target = GL_UNIFORM_BUFFER;
}
for (auto &[binding, texture] : set->texture_binding_map()) {
auto cmd = std::make_unique<CmdBindTextureToIndex>();
cmd->texture = texture;
cmd->index = binding;
cmd->target = device_->get_gl_image(texture).target;
recorded_commands_.push_back(std::move(cmd));
auto &bind = cmd->textures.emplace_back();
bind.texture = texture;
bind.index = binding;
bind.target = device_->get_gl_image(texture).target;
}
for (auto &[binding, texture] : set->rw_image_binding_map()) {
auto &bind = cmd->textures.emplace_back();
bind.texture = texture;
bind.index = binding;
bind.format = device_->get_gl_image(texture).format;
bind.is_storage = true;
}

recorded_commands_.push_back(std::move(cmd));
return RhiResult::success;
}

Expand Down Expand Up @@ -502,6 +519,20 @@ GLDevice::GLDevice() : stream_(this) {
caps.set(DeviceCapability::spirv_has_int64, true);
caps.set(DeviceCapability::spirv_has_float64, true);
}

if (GLAD_GL_NV_gpu_shader5) {
caps.set(DeviceCapability::spirv_has_int16, true);
caps.set(DeviceCapability::spirv_has_float16, true);
}

if (GLAD_GL_AMD_gpu_shader_int16) {
caps.set(DeviceCapability::spirv_has_int16, true);
}

if (GLAD_GL_AMD_gpu_shader_half_float) {
caps.set(DeviceCapability::spirv_has_float16, true);
}

caps.set(DeviceCapability::spirv_version, 0x10300);
set_caps(std::move(caps));
}
Expand Down Expand Up @@ -671,6 +702,8 @@ DeviceAllocation GLDevice::create_image(const ImageParams &params) {
gl_texture_dims = GL_TEXTURE_1D;
} else if (params.dimension == ImageDimension::d2D) {
gl_texture_dims = GL_TEXTURE_2D;
} else if (params.dimension == ImageDimension::d3D) {
gl_texture_dims = GL_TEXTURE_3D;
}

auto format = format_to_gl_internal_format.at(params.format);
Expand Down Expand Up @@ -787,22 +820,30 @@ void GLCommandList::CmdBindPipeline::execute() {
check_opengl_error("glUseProgram");
}

void GLCommandList::CmdBindBufferToIndex::execute() {
if (size == -1) {
glBindBufferBase(target, index, buffer);
check_opengl_error("glBindBufferBase");
} else {
glBindBufferRange(target, index, buffer, GLintptr(offset),
GLsizeiptr(size));
check_opengl_error("glBindBufferRange");
void GLCommandList::CmdBindResources::execute() {
for (auto &bind : buffers) {
if (bind.size == -1) {
glBindBufferBase(bind.target, bind.index, bind.buffer);
check_opengl_error("glBindBufferBase");
} else {
glBindBufferRange(bind.target, bind.index, bind.buffer,
GLintptr(bind.offset), GLsizeiptr(bind.size));
check_opengl_error("glBindBufferRange");
}
}
for (auto &bind : textures) {
if (bind.is_storage) {
glBindImageTexture(bind.index, bind.texture, /*level=*/0,
/*layered=*/GL_FALSE, /*layer=*/0, GL_READ_WRITE,
bind.format);
check_opengl_error("glBindImageTexture");
} else {
glActiveTexture(GL_TEXTURE0 + bind.index);
check_opengl_error("glActiveTexture");
glBindTexture(bind.target, bind.texture);
check_opengl_error("glBindTexture");
}
}
}

void GLCommandList::CmdBindTextureToIndex::execute() {
glActiveTexture(GL_TEXTURE0 + index);
check_opengl_error("glActiveTexture");
glBindTexture(GL_TEXTURE_2D, texture);
check_opengl_error("glBindTexture");
}

void GLCommandList::CmdBufferBarrier::execute() {
Expand Down
43 changes: 27 additions & 16 deletions taichi/rhi/opengl/opengl_device.h
Original file line number Diff line number Diff line change
Expand Up @@ -39,22 +39,27 @@ class GLResourceSet : public ShaderResourceSet {
size_t size;
};

const std::unordered_map<uint32_t, BufferBinding> &ssbo_binding_map() {
const std::unordered_map<uint32_t, BufferBinding> &ssbo_binding_map() const {
return ssbo_binding_map_;
}

const std::unordered_map<uint32_t, BufferBinding> &ubo_binding_map() {
const std::unordered_map<uint32_t, BufferBinding> &ubo_binding_map() const {
return ubo_binding_map_;
}

const std::unordered_map<uint32_t, GLuint> &texture_binding_map() {
const std::unordered_map<uint32_t, GLuint> &texture_binding_map() const {
return texture_binding_map_;
}

const std::unordered_map<uint32_t, GLuint> &rw_image_binding_map() const {
return rw_image_binding_map_;
}

private:
std::unordered_map<uint32_t, BufferBinding> ssbo_binding_map_;
std::unordered_map<uint32_t, BufferBinding> ubo_binding_map_;
std::unordered_map<uint32_t, GLuint> texture_binding_map_;
std::unordered_map<uint32_t, GLuint> rw_image_binding_map_;
};

class GLPipeline : public Pipeline {
Expand Down Expand Up @@ -132,19 +137,25 @@ class GLCommandList : public CommandList {
void execute() override;
};

struct CmdBindBufferToIndex : public Cmd {
GLuint buffer{0};
GLuint index{0};
GLuint offset{0};
GLuint size{0};
GLenum target{GL_SHADER_STORAGE_BUFFER};
void execute() override;
};

struct CmdBindTextureToIndex : public Cmd {
GLuint texture{0};
GLuint index{0};
GLenum target{GL_TEXTURE_2D};
struct CmdBindResources : public Cmd {
struct BufferBinding {
GLuint buffer{0};
GLuint index{0};
GLuint offset{0};
GLuint size{0};
GLenum target{GL_SHADER_STORAGE_BUFFER};
};

struct TextureBinding {
GLuint texture{0};
GLuint index{0};
GLenum target{GL_TEXTURE_2D};
GLenum format{GL_RGBA32F};
bool is_storage{false};
};

std::vector<BufferBinding> buffers;
std::vector<TextureBinding> textures;
void execute() override;
};

Expand Down

0 comments on commit 7124225

Please sign in to comment.