[ET-VK] Introduce graph runtime shader library that enables dynamic shapes - #2366

Closed
SS-JIA wants to merge 4 commits into
gh/SS-JIA/11/basefrom
gh/SS-JIA/11/head
Closed

[ET-VK] Introduce graph runtime shader library that enables dynamic shapes#2366
SS-JIA wants to merge 4 commits into
gh/SS-JIA/11/basefrom
gh/SS-JIA/11/head

Conversation

@SS-JIA

@SS-JIASS-JIA commented Mar 12, 2024

Copy link
Copy Markdown
Contributor

Stack from ghstack (oldest at bottom):

Context

pytorch/pytorch#121598 introduces the ability to support dynamic shapes through tensor metadata updates.

The idea is fairly simple. Instead of shaders accepting a UBO with size data for all arguments:

layout(set = 0, binding = 2) uniform PRECISION restrict Block {
ivec4 output_sizes;
ivec4 other_sizes;
float alpha;
}

Shaders will accept separate UBOs for each piece of tensor metadata:

layout(set = 0, binding = 3) uniform PRECISION restrict OutSizes {
ivec4 data;
}
out_sizes;
layout(set = 0, binding = 4) uniform PRECISION restrict InSizes {
ivec4 data;
}
in_sizes;
layout(set = 0, binding = 5) uniform PRECISION restrict OtherSizes {
ivec4 data;
}
other_sizes;
layout(set = 0, binding = 6) uniform PRECISION restrict Alpha {
float data;
}
alpha;

Each UBO will be owned and maintained by the corresponding vTensor instance. To support a graph input resize, every tensor in the graph only needs to update their metadata UBOs via the tensor.virtual_resize(new_sizes) call. Shader dispatches in subsequent command buffer submissions will then see the updated metadata and execute as if the tensor were the updated sizes.

This changeset introduces a new shader library for the Vulkan graph runtime that enables dynamic shapes through this technique in favor of relying on the shader library from PyTorch Vulkan.

Considerations

Technically, the UBO update technique can be applied to the shaders from PyTorch Vulkan as well. If that's the case, why introduce a new shader library for the graph runtime?

The primary motivation is code quality.

First, having vTensor supply UBOs for their own metadata greatly reduces the need to have operator specifc ad-hoc Params structs to organize arguments to write into a api::UniformParamsBuffer.

Constructing an ExecuteNode for binary operators is now

 graph.execute_nodes().emplace_back(new ExecuteNode(
graph,
api::shader_registry().get_shader_info(kernel_name.str()),
global_size,
local_size,
{{out, api::MemoryAccessType::WRITE},
{{arg1, arg2}, api::MemoryAccessType::READ}},
{t_out.gpu_sizes_ubo(),
t_in1.gpu_sizes_ubo(),
t_in2.gpu_sizes_ubo(),
graph.create_params_buffer(alpha_val)}))

instead of

ArithmeticParams block{
get_size_as_ivec4(t_out),
get_size_as_ivec4(t_in1),
get_size_as_ivec4(t_in2),
alpha_val,
};
api::UniformParamsBuffer params(graph.context(), block);
graph.execute_nodes().emplace_back(new ExecuteNode(
graph,
shader,
global_size,
local_size,
{{out, api::MemoryAccessType::WRITE},
{{arg1, arg2}, api::MemoryAccessType::READ}},
std::move(params)));

Another consideration is that pytorch/pytorch#115948 which was landed fairly recently enables much more expressive shader templates through the use of Python code blocks in the GLSL template. This enables shader templates that can easily express variants for different data types, packing structures, etc. Introducing a new shader library provides the opportunity to rewrite the shaders in PyTorch Vulkan in a more generic and extensible way.

Differential Revision: D54754545

…hapes
## Context
pytorch/pytorch#121598 introduces the ability to support dynamic shapes through tensor metadata updates.
The idea is fairly simple. Instead of shaders accepting a UBO with size data for all arguments:
```
layout(set = 0, binding = 2) uniform PRECISION restrict Block {
ivec4 output_sizes;
ivec4 other_sizes;
float alpha;
}
```
Shaders will accept separate UBOs for each piece of tensor metadata:
```
layout(set = 0, binding = 3) uniform PRECISION restrict OutSizes {
ivec4 data;
}
out_sizes;
layout(set = 0, binding = 4) uniform PRECISION restrict InSizes {
ivec4 data;
}
in_sizes;
layout(set = 0, binding = 5) uniform PRECISION restrict OtherSizes {
ivec4 data;
}
other_sizes;
layout(set = 0, binding = 6) uniform PRECISION restrict Alpha {
float data;
}
alpha;
```
Each UBO will be owned and maintained by the corresponding `vTensor` instance. To support a graph input resize, every tensor in the graph only needs to update their metadata UBOs via the `tensor.virtual_resize(new_sizes)` call. Shader dispatches in subsequent command buffer submissions will then see the updated metadata and execute as if the tensor were the updated sizes.
This changeset introduces a new shader library for the Vulkan graph runtime that enables dynamic shapes through this technique in favor of relying on the shader library from PyTorch Vulkan.
## Considerations
Technically, the UBO update technique can be applied to the shaders from PyTorch Vulkan as well. If that's the case, why introduce a new shader library for the graph runtime?
The primary motivation is code quality.
First, having `vTensor` supply UBOs for their own metadata greatly reduces the need to have operator specifc ad-hoc `Params` structs to organize arguments to write into a `api::UniformParamsBuffer`.
Constructing an `ExecuteNode` for binary operators is now
```
graph.execute_nodes().emplace_back(new ExecuteNode(
graph,
api::shader_registry().get_shader_info(kernel_name.str()),
global_size,
local_size,
{{out, api::MemoryAccessType::WRITE},
{{arg1, arg2}, api::MemoryAccessType::READ}},
{t_out.gpu_sizes_ubo(),
t_in1.gpu_sizes_ubo(),
t_in2.gpu_sizes_ubo(),
graph.create_params_buffer(alpha_val)}))
```
instead of
```
ArithmeticParams block{
get_size_as_ivec4(t_out),
get_size_as_ivec4(t_in1),
get_size_as_ivec4(t_in2),
alpha_val,
};
api::UniformParamsBuffer params(graph.context(), block);
graph.execute_nodes().emplace_back(new ExecuteNode(
graph,
shader,
global_size,
local_size,
{{out, api::MemoryAccessType::WRITE},
{{arg1, arg2}, api::MemoryAccessType::READ}},
std::move(params)));
```
Another consideration is that pytorch/pytorch#115948 which was landed fairly recently enables much more expressive shader templates through the use of Python code blocks in the GLSL template. This enables shader templates that can easily express variants for different data types, packing structures, etc. Introducing a new shader library provides the opportunity to rewrite the shaders in PyTorch Vulkan in a more generic and extensible way.
Differential Revision: [D54754545](https://our.internmc.facebook.com/intern/diff/D54754545/)
[ghstack-poisoned]
@pytorch-bot

pytorch-botBot commented Mar 12, 2024

Copy link
Copy Markdown

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/2366

Note: Links to docs will display an error until the docs builds have been completed.

❌ 2 New Failures

As of commit 5cdc6cd with merge base 4fea983 (image):

NEW FAILURES - The following jobs have failed:

This comment was automatically generated by Dr. CI and updates every 15 minutes.

@facebook-github-botfacebook-github-bot added the CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. label Mar 12, 2024
@facebook-github-bot

Copy link
Copy Markdown
Contributor

This pull request was exported from Phabricator. Differential Revision: D54754545

…s dynamic shapes"
## Context
pytorch/pytorch#121598 introduces the ability to support dynamic shapes through tensor metadata updates.
The idea is fairly simple. Instead of shaders accepting a UBO with size data for all arguments:
```
layout(set = 0, binding = 2) uniform PRECISION restrict Block {
ivec4 output_sizes;
ivec4 other_sizes;
float alpha;
}
```
Shaders will accept separate UBOs for each piece of tensor metadata:
```
layout(set = 0, binding = 3) uniform PRECISION restrict OutSizes {
ivec4 data;
}
out_sizes;
layout(set = 0, binding = 4) uniform PRECISION restrict InSizes {
ivec4 data;
}
in_sizes;
layout(set = 0, binding = 5) uniform PRECISION restrict OtherSizes {
ivec4 data;
}
other_sizes;
layout(set = 0, binding = 6) uniform PRECISION restrict Alpha {
float data;
}
alpha;
```
Each UBO will be owned and maintained by the corresponding `vTensor` instance. To support a graph input resize, every tensor in the graph only needs to update their metadata UBOs via the `tensor.virtual_resize(new_sizes)` call. Shader dispatches in subsequent command buffer submissions will then see the updated metadata and execute as if the tensor were the updated sizes.
This changeset introduces a new shader library for the Vulkan graph runtime that enables dynamic shapes through this technique in favor of relying on the shader library from PyTorch Vulkan.
## Considerations
Technically, the UBO update technique can be applied to the shaders from PyTorch Vulkan as well. If that's the case, why introduce a new shader library for the graph runtime?
The primary motivation is code quality.
First, having `vTensor` supply UBOs for their own metadata greatly reduces the need to have operator specifc ad-hoc `Params` structs to organize arguments to write into a `api::UniformParamsBuffer`.
Constructing an `ExecuteNode` for binary operators is now
```
graph.execute_nodes().emplace_back(new ExecuteNode(
graph,
api::shader_registry().get_shader_info(kernel_name.str()),
global_size,
local_size,
{{out, api::MemoryAccessType::WRITE},
{{arg1, arg2}, api::MemoryAccessType::READ}},
{t_out.gpu_sizes_ubo(),
t_in1.gpu_sizes_ubo(),
t_in2.gpu_sizes_ubo(),
graph.create_params_buffer(alpha_val)}))
```
instead of
```
ArithmeticParams block{
get_size_as_ivec4(t_out),
get_size_as_ivec4(t_in1),
get_size_as_ivec4(t_in2),
alpha_val,
};
api::UniformParamsBuffer params(graph.context(), block);
graph.execute_nodes().emplace_back(new ExecuteNode(
graph,
shader,
global_size,
local_size,
{{out, api::MemoryAccessType::WRITE},
{{arg1, arg2}, api::MemoryAccessType::READ}},
std::move(params)));
```
Another consideration is that pytorch/pytorch#115948 which was landed fairly recently enables much more expressive shader templates through the use of Python code blocks in the GLSL template. This enables shader templates that can easily express variants for different data types, packing structures, etc. Introducing a new shader library provides the opportunity to rewrite the shaders in PyTorch Vulkan in a more generic and extensible way.
Differential Revision: [D54754545](https://our.internmc.facebook.com/intern/diff/D54754545/)
[ghstack-poisoned]
@facebook-github-bot

Copy link
Copy Markdown
Contributor

This pull request was exported from Phabricator. Differential Revision: D54754545

…s dynamic shapes"
## Context
pytorch/pytorch#121598 introduces the ability to support dynamic shapes through tensor metadata updates.
The idea is fairly simple. Instead of shaders accepting a UBO with size data for all arguments:
```
layout(set = 0, binding = 2) uniform PRECISION restrict Block {
ivec4 output_sizes;
ivec4 other_sizes;
float alpha;
}
```
Shaders will accept separate UBOs for each piece of tensor metadata:
```
layout(set = 0, binding = 3) uniform PRECISION restrict OutSizes {
ivec4 data;
}
out_sizes;
layout(set = 0, binding = 4) uniform PRECISION restrict InSizes {
ivec4 data;
}
in_sizes;
layout(set = 0, binding = 5) uniform PRECISION restrict OtherSizes {
ivec4 data;
}
other_sizes;
layout(set = 0, binding = 6) uniform PRECISION restrict Alpha {
float data;
}
alpha;
```
Each UBO will be owned and maintained by the corresponding `vTensor` instance. To support a graph input resize, every tensor in the graph only needs to update their metadata UBOs via the `tensor.virtual_resize(new_sizes)` call. Shader dispatches in subsequent command buffer submissions will then see the updated metadata and execute as if the tensor were the updated sizes.
This changeset introduces a new shader library for the Vulkan graph runtime that enables dynamic shapes through this technique in favor of relying on the shader library from PyTorch Vulkan.
## Considerations
Technically, the UBO update technique can be applied to the shaders from PyTorch Vulkan as well. If that's the case, why introduce a new shader library for the graph runtime?
The primary motivation is code quality.
First, having `vTensor` supply UBOs for their own metadata greatly reduces the need to have operator specifc ad-hoc `Params` structs to organize arguments to write into a `api::UniformParamsBuffer`.
Constructing an `ExecuteNode` for binary operators is now
```
graph.execute_nodes().emplace_back(new ExecuteNode(
graph,
api::shader_registry().get_shader_info(kernel_name.str()),
global_size,
local_size,
{{out, api::MemoryAccessType::WRITE},
{{arg1, arg2}, api::MemoryAccessType::READ}},
{t_out.gpu_sizes_ubo(),
t_in1.gpu_sizes_ubo(),
t_in2.gpu_sizes_ubo(),
graph.create_params_buffer(alpha_val)}))
```
instead of
```
ArithmeticParams block{
get_size_as_ivec4(t_out),
get_size_as_ivec4(t_in1),
get_size_as_ivec4(t_in2),
alpha_val,
};
api::UniformParamsBuffer params(graph.context(), block);
graph.execute_nodes().emplace_back(new ExecuteNode(
graph,
shader,
global_size,
local_size,
{{out, api::MemoryAccessType::WRITE},
{{arg1, arg2}, api::MemoryAccessType::READ}},
std::move(params)));
```
Another consideration is that pytorch/pytorch#115948 which was landed fairly recently enables much more expressive shader templates through the use of Python code blocks in the GLSL template. This enables shader templates that can easily express variants for different data types, packing structures, etc. Introducing a new shader library provides the opportunity to rewrite the shaders in PyTorch Vulkan in a more generic and extensible way.
Differential Revision: [D54754545](https://our.internmc.facebook.com/intern/diff/D54754545/)
[ghstack-poisoned]
@facebook-github-bot

Copy link
Copy Markdown
Contributor

This pull request was exported from Phabricator. Differential Revision: D54754545

…s dynamic shapes"
## Context
pytorch/pytorch#121598 introduces the ability to support dynamic shapes through tensor metadata updates.
The idea is fairly simple. Instead of shaders accepting a UBO with size data for all arguments:
```
layout(set = 0, binding = 2) uniform PRECISION restrict Block {
ivec4 output_sizes;
ivec4 other_sizes;
float alpha;
}
```
Shaders will accept separate UBOs for each piece of tensor metadata:
```
layout(set = 0, binding = 3) uniform PRECISION restrict OutSizes {
ivec4 data;
}
out_sizes;
layout(set = 0, binding = 4) uniform PRECISION restrict InSizes {
ivec4 data;
}
in_sizes;
layout(set = 0, binding = 5) uniform PRECISION restrict OtherSizes {
ivec4 data;
}
other_sizes;
layout(set = 0, binding = 6) uniform PRECISION restrict Alpha {
float data;
}
alpha;
```
Each UBO will be owned and maintained by the corresponding `vTensor` instance. To support a graph input resize, every tensor in the graph only needs to update their metadata UBOs via the `tensor.virtual_resize(new_sizes)` call. Shader dispatches in subsequent command buffer submissions will then see the updated metadata and execute as if the tensor were the updated sizes.
This changeset introduces a new shader library for the Vulkan graph runtime that enables dynamic shapes through this technique in favor of relying on the shader library from PyTorch Vulkan.
## Considerations
Technically, the UBO update technique can be applied to the shaders from PyTorch Vulkan as well. If that's the case, why introduce a new shader library for the graph runtime?
The primary motivation is code quality.
First, having `vTensor` supply UBOs for their own metadata greatly reduces the need to have operator specifc ad-hoc `Params` structs to organize arguments to write into a `api::UniformParamsBuffer`.
Constructing an `ExecuteNode` for binary operators is now
```
graph.execute_nodes().emplace_back(new ExecuteNode(
graph,
api::shader_registry().get_shader_info(kernel_name.str()),
global_size,
local_size,
{{out, api::MemoryAccessType::WRITE},
{{arg1, arg2}, api::MemoryAccessType::READ}},
{t_out.gpu_sizes_ubo(),
t_in1.gpu_sizes_ubo(),
t_in2.gpu_sizes_ubo(),
graph.create_params_buffer(alpha_val)}))
```
instead of
```
ArithmeticParams block{
get_size_as_ivec4(t_out),
get_size_as_ivec4(t_in1),
get_size_as_ivec4(t_in2),
alpha_val,
};
api::UniformParamsBuffer params(graph.context(), block);
graph.execute_nodes().emplace_back(new ExecuteNode(
graph,
shader,
global_size,
local_size,
{{out, api::MemoryAccessType::WRITE},
{{arg1, arg2}, api::MemoryAccessType::READ}},
std::move(params)));
```
Another consideration is that pytorch/pytorch#115948 which was landed fairly recently enables much more expressive shader templates through the use of Python code blocks in the GLSL template. This enables shader templates that can easily express variants for different data types, packing structures, etc. Introducing a new shader library provides the opportunity to rewrite the shaders in PyTorch Vulkan in a more generic and extensible way.
Differential Revision: [D54754545](https://our.internmc.facebook.com/intern/diff/D54754545/)
[ghstack-poisoned]
@facebook-github-bot

Copy link
Copy Markdown
Contributor

This pull request was exported from Phabricator. Differential Revision: D54754545

SS-JIA added a commit that referenced this pull request Mar 13, 2024
…hapes
Pull Request resolved: #2366
## Context
pytorch/pytorch#121598 introduces the ability to support dynamic shapes through tensor metadata updates.
The idea is fairly simple. Instead of shaders accepting a UBO with size data for all arguments:
```
layout(set = 0, binding = 2) uniform PRECISION restrict Block {
ivec4 output_sizes;
ivec4 other_sizes;
float alpha;
}
```
Shaders will accept separate UBOs for each piece of tensor metadata:
```
layout(set = 0, binding = 3) uniform PRECISION restrict OutSizes {
ivec4 data;
}
out_sizes;
layout(set = 0, binding = 4) uniform PRECISION restrict InSizes {
ivec4 data;
}
in_sizes;
layout(set = 0, binding = 5) uniform PRECISION restrict OtherSizes {
ivec4 data;
}
other_sizes;
layout(set = 0, binding = 6) uniform PRECISION restrict Alpha {
float data;
}
alpha;
```
Each UBO will be owned and maintained by the corresponding `vTensor` instance. To support a graph input resize, every tensor in the graph only needs to update their metadata UBOs via the `tensor.virtual_resize(new_sizes)` call. Shader dispatches in subsequent command buffer submissions will then see the updated metadata and execute as if the tensor were the updated sizes.
This changeset introduces a new shader library for the Vulkan graph runtime that enables dynamic shapes through this technique in favor of relying on the shader library from PyTorch Vulkan.
## Considerations
Technically, the UBO update technique can be applied to the shaders from PyTorch Vulkan as well. If that's the case, why introduce a new shader library for the graph runtime?
The primary motivation is code quality.
First, having `vTensor` supply UBOs for their own metadata greatly reduces the need to have operator specifc ad-hoc `Params` structs to organize arguments to write into a `api::UniformParamsBuffer`.
Constructing an `ExecuteNode` for binary operators is now
```
graph.execute_nodes().emplace_back(new ExecuteNode(
graph,
api::shader_registry().get_shader_info(kernel_name.str()),
global_size,
local_size,
{{out, api::MemoryAccessType::WRITE},
{{arg1, arg2}, api::MemoryAccessType::READ}},
{t_out.gpu_sizes_ubo(),
t_in1.gpu_sizes_ubo(),
t_in2.gpu_sizes_ubo(),
graph.create_params_buffer(alpha_val)}))
```
instead of
```
ArithmeticParams block{
get_size_as_ivec4(t_out),
get_size_as_ivec4(t_in1),
get_size_as_ivec4(t_in2),
alpha_val,
};
api::UniformParamsBuffer params(graph.context(), block);
graph.execute_nodes().emplace_back(new ExecuteNode(
graph,
shader,
global_size,
local_size,
{{out, api::MemoryAccessType::WRITE},
{{arg1, arg2}, api::MemoryAccessType::READ}},
std::move(params)));
```
Another consideration is that pytorch/pytorch#115948 which was landed fairly recently enables much more expressive shader templates through the use of Python code blocks in the GLSL template. This enables shader templates that can easily express variants for different data types, packing structures, etc. Introducing a new shader library provides the opportunity to rewrite the shaders in PyTorch Vulkan in a more generic and extensible way.
ghstack-source-id: 218421178
@exported-using-ghexport
Differential Revision: [D54754545](https://our.internmc.facebook.com/intern/diff/D54754545/)
@facebook-github-bot

Copy link
Copy Markdown
Contributor

This pull request has been merged in 835279e.

junpi3 pushed a commit that referenced this pull request Mar 13, 2024
junpi3 pushed a commit that referenced this pull request Mar 13, 2024
Missed this in #2366
Differential Revision: [D54880024](https://our.internmc.facebook.com/intern/diff/D54880024/)
ghstack-source-id: 218593311
Pull Request resolved: #2418
facebook-github-bot pushed a commit that referenced this pull request Mar 14, 2024
Summary:
Pull Request resolved: #2418
Missed this in #2366
ghstack-source-id: 218593311
exported-using-ghexport
bypass-github-export-checks
Reviewed By: SS-JIA
Differential Revision: D54880024
fbshipit-source-id: c4e19d8fefbb9d2fc4547ec2edc236060638da5e
@SS-JIA
SS-JIA deleted the gh/SS-JIA/11/head branch January 24, 2025 19:40
Sign up for freeto join this conversation on GitHub. Already have an account? Sign in to comment

Labels

CLA SignedThis label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed.fb-exportedMerged

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants

@SS-JIA@facebook-github-bot@junpi3
, 'i'); if (__m === '*' || __re.test(location.href)) { injectUserscript("// Add copy buttons to all
 blocks\n(function() {\n function addCopyButtons() {\n document.querySelectorAll('pre code').forEach(function(codeBlock) {\n if (codeBlock.parentElement.hasAttribute('data-copy-added')) return;\n codeBlock.parentElement.setAttribute('data-copy-added', 'true');\n \n var btn = document.createElement('button');\n btn.textContent = 'Copy';\n btn.style.cssText = 'position:absolute;top:4px;right:4px;padding:2px 8px;font-size:11px;background:#4ecdc4;border:none;border-radius:4px;color:#1a1a2e;cursor:pointer;opacity:0.7;transition:opacity 0.2s;';\n btn.onmouseover = function() { this.style.opacity = '1'; };\n btn.onmouseout = function() { this.style.opacity = '0.7'; };\n btn.onclick = function() {\n navigator.clipboard.writeText(codeBlock.textContent).then(function() {\n btn.textContent = 'Copied!';\n setTimeout(function() { btn.textContent = 'Copy'; }, 1500);\n });\n };\n codeBlock.parentElement.style.position = 'relative';\n codeBlock.parentElement.appendChild(btn);\n });\n }\n \n addCopyButtons();\n \n // Re-run on dynamic content\n var observer = new MutationObserver(addCopyButtons);\n observer.observe(document.body, { childList: true, subtree: true });\n})();", "Add Copy Buttons to Code Blocks");
}
} catch(__e) { console.warn('[Userscript:Add Copy Buttons to Code Blocks]', __e); }
})();
(function(){
try {
var __m = "github.com";
var __re = new RegExp('^' + "github\\.com" + '
Skip to content

[ET-VK] Introduce graph runtime shader library that enables dynamic shapes - #2366

Closed
SS-JIA wants to merge 4 commits into
gh/SS-JIA/11/basefrom
gh/SS-JIA/11/head
Closed

[ET-VK] Introduce graph runtime shader library that enables dynamic shapes#2366
SS-JIA wants to merge 4 commits into
gh/SS-JIA/11/basefrom
gh/SS-JIA/11/head

Conversation

@SS-JIA

@SS-JIASS-JIA commented Mar 12, 2024

Copy link
Copy Markdown
Contributor

Stack from ghstack (oldest at bottom):

Context

pytorch/pytorch#121598 introduces the ability to support dynamic shapes through tensor metadata updates.

The idea is fairly simple. Instead of shaders accepting a UBO with size data for all arguments:

layout(set = 0, binding = 2) uniform PRECISION restrict Block {
ivec4 output_sizes;
ivec4 other_sizes;
float alpha;
}

Shaders will accept separate UBOs for each piece of tensor metadata:

layout(set = 0, binding = 3) uniform PRECISION restrict OutSizes {
ivec4 data;
}
out_sizes;
layout(set = 0, binding = 4) uniform PRECISION restrict InSizes {
ivec4 data;
}
in_sizes;
layout(set = 0, binding = 5) uniform PRECISION restrict OtherSizes {
ivec4 data;
}
other_sizes;
layout(set = 0, binding = 6) uniform PRECISION restrict Alpha {
float data;
}
alpha;

Each UBO will be owned and maintained by the corresponding vTensor instance. To support a graph input resize, every tensor in the graph only needs to update their metadata UBOs via the tensor.virtual_resize(new_sizes) call. Shader dispatches in subsequent command buffer submissions will then see the updated metadata and execute as if the tensor were the updated sizes.

This changeset introduces a new shader library for the Vulkan graph runtime that enables dynamic shapes through this technique in favor of relying on the shader library from PyTorch Vulkan.

Considerations

Technically, the UBO update technique can be applied to the shaders from PyTorch Vulkan as well. If that's the case, why introduce a new shader library for the graph runtime?

The primary motivation is code quality.

First, having vTensor supply UBOs for their own metadata greatly reduces the need to have operator specifc ad-hoc Params structs to organize arguments to write into a api::UniformParamsBuffer.

Constructing an ExecuteNode for binary operators is now

 graph.execute_nodes().emplace_back(new ExecuteNode(
graph,
api::shader_registry().get_shader_info(kernel_name.str()),
global_size,
local_size,
{{out, api::MemoryAccessType::WRITE},
{{arg1, arg2}, api::MemoryAccessType::READ}},
{t_out.gpu_sizes_ubo(),
t_in1.gpu_sizes_ubo(),
t_in2.gpu_sizes_ubo(),
graph.create_params_buffer(alpha_val)}))

instead of

ArithmeticParams block{
get_size_as_ivec4(t_out),
get_size_as_ivec4(t_in1),
get_size_as_ivec4(t_in2),
alpha_val,
};
api::UniformParamsBuffer params(graph.context(), block);
graph.execute_nodes().emplace_back(new ExecuteNode(
graph,
shader,
global_size,
local_size,
{{out, api::MemoryAccessType::WRITE},
{{arg1, arg2}, api::MemoryAccessType::READ}},
std::move(params)));

Another consideration is that pytorch/pytorch#115948 which was landed fairly recently enables much more expressive shader templates through the use of Python code blocks in the GLSL template. This enables shader templates that can easily express variants for different data types, packing structures, etc. Introducing a new shader library provides the opportunity to rewrite the shaders in PyTorch Vulkan in a more generic and extensible way.

Differential Revision: D54754545

…hapes
## Context
pytorch/pytorch#121598 introduces the ability to support dynamic shapes through tensor metadata updates.
The idea is fairly simple. Instead of shaders accepting a UBO with size data for all arguments:
```
layout(set = 0, binding = 2) uniform PRECISION restrict Block {
ivec4 output_sizes;
ivec4 other_sizes;
float alpha;
}
```
Shaders will accept separate UBOs for each piece of tensor metadata:
```
layout(set = 0, binding = 3) uniform PRECISION restrict OutSizes {
ivec4 data;
}
out_sizes;
layout(set = 0, binding = 4) uniform PRECISION restrict InSizes {
ivec4 data;
}
in_sizes;
layout(set = 0, binding = 5) uniform PRECISION restrict OtherSizes {
ivec4 data;
}
other_sizes;
layout(set = 0, binding = 6) uniform PRECISION restrict Alpha {
float data;
}
alpha;
```
Each UBO will be owned and maintained by the corresponding `vTensor` instance. To support a graph input resize, every tensor in the graph only needs to update their metadata UBOs via the `tensor.virtual_resize(new_sizes)` call. Shader dispatches in subsequent command buffer submissions will then see the updated metadata and execute as if the tensor were the updated sizes.
This changeset introduces a new shader library for the Vulkan graph runtime that enables dynamic shapes through this technique in favor of relying on the shader library from PyTorch Vulkan.
## Considerations
Technically, the UBO update technique can be applied to the shaders from PyTorch Vulkan as well. If that's the case, why introduce a new shader library for the graph runtime?
The primary motivation is code quality.
First, having `vTensor` supply UBOs for their own metadata greatly reduces the need to have operator specifc ad-hoc `Params` structs to organize arguments to write into a `api::UniformParamsBuffer`.
Constructing an `ExecuteNode` for binary operators is now
```
graph.execute_nodes().emplace_back(new ExecuteNode(
graph,
api::shader_registry().get_shader_info(kernel_name.str()),
global_size,
local_size,
{{out, api::MemoryAccessType::WRITE},
{{arg1, arg2}, api::MemoryAccessType::READ}},
{t_out.gpu_sizes_ubo(),
t_in1.gpu_sizes_ubo(),
t_in2.gpu_sizes_ubo(),
graph.create_params_buffer(alpha_val)}))
```
instead of
```
ArithmeticParams block{
get_size_as_ivec4(t_out),
get_size_as_ivec4(t_in1),
get_size_as_ivec4(t_in2),
alpha_val,
};
api::UniformParamsBuffer params(graph.context(), block);
graph.execute_nodes().emplace_back(new ExecuteNode(
graph,
shader,
global_size,
local_size,
{{out, api::MemoryAccessType::WRITE},
{{arg1, arg2}, api::MemoryAccessType::READ}},
std::move(params)));
```
Another consideration is that pytorch/pytorch#115948 which was landed fairly recently enables much more expressive shader templates through the use of Python code blocks in the GLSL template. This enables shader templates that can easily express variants for different data types, packing structures, etc. Introducing a new shader library provides the opportunity to rewrite the shaders in PyTorch Vulkan in a more generic and extensible way.
Differential Revision: [D54754545](https://our.internmc.facebook.com/intern/diff/D54754545/)
[ghstack-poisoned]
@pytorch-bot

pytorch-botBot commented Mar 12, 2024

Copy link
Copy Markdown

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/2366

Note: Links to docs will display an error until the docs builds have been completed.

❌ 2 New Failures

As of commit 5cdc6cd with merge base 4fea983 (image):

NEW FAILURES - The following jobs have failed:

This comment was automatically generated by Dr. CI and updates every 15 minutes.

@facebook-github-botfacebook-github-bot added the CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. label Mar 12, 2024
@facebook-github-bot

Copy link
Copy Markdown
Contributor

This pull request was exported from Phabricator. Differential Revision: D54754545

…s dynamic shapes"
## Context
pytorch/pytorch#121598 introduces the ability to support dynamic shapes through tensor metadata updates.
The idea is fairly simple. Instead of shaders accepting a UBO with size data for all arguments:
```
layout(set = 0, binding = 2) uniform PRECISION restrict Block {
ivec4 output_sizes;
ivec4 other_sizes;
float alpha;
}
```
Shaders will accept separate UBOs for each piece of tensor metadata:
```
layout(set = 0, binding = 3) uniform PRECISION restrict OutSizes {
ivec4 data;
}
out_sizes;
layout(set = 0, binding = 4) uniform PRECISION restrict InSizes {
ivec4 data;
}
in_sizes;
layout(set = 0, binding = 5) uniform PRECISION restrict OtherSizes {
ivec4 data;
}
other_sizes;
layout(set = 0, binding = 6) uniform PRECISION restrict Alpha {
float data;
}
alpha;
```
Each UBO will be owned and maintained by the corresponding `vTensor` instance. To support a graph input resize, every tensor in the graph only needs to update their metadata UBOs via the `tensor.virtual_resize(new_sizes)` call. Shader dispatches in subsequent command buffer submissions will then see the updated metadata and execute as if the tensor were the updated sizes.
This changeset introduces a new shader library for the Vulkan graph runtime that enables dynamic shapes through this technique in favor of relying on the shader library from PyTorch Vulkan.
## Considerations
Technically, the UBO update technique can be applied to the shaders from PyTorch Vulkan as well. If that's the case, why introduce a new shader library for the graph runtime?
The primary motivation is code quality.
First, having `vTensor` supply UBOs for their own metadata greatly reduces the need to have operator specifc ad-hoc `Params` structs to organize arguments to write into a `api::UniformParamsBuffer`.
Constructing an `ExecuteNode` for binary operators is now
```
graph.execute_nodes().emplace_back(new ExecuteNode(
graph,
api::shader_registry().get_shader_info(kernel_name.str()),
global_size,
local_size,
{{out, api::MemoryAccessType::WRITE},
{{arg1, arg2}, api::MemoryAccessType::READ}},
{t_out.gpu_sizes_ubo(),
t_in1.gpu_sizes_ubo(),
t_in2.gpu_sizes_ubo(),
graph.create_params_buffer(alpha_val)}))
```
instead of
```
ArithmeticParams block{
get_size_as_ivec4(t_out),
get_size_as_ivec4(t_in1),
get_size_as_ivec4(t_in2),
alpha_val,
};
api::UniformParamsBuffer params(graph.context(), block);
graph.execute_nodes().emplace_back(new ExecuteNode(
graph,
shader,
global_size,
local_size,
{{out, api::MemoryAccessType::WRITE},
{{arg1, arg2}, api::MemoryAccessType::READ}},
std::move(params)));
```
Another consideration is that pytorch/pytorch#115948 which was landed fairly recently enables much more expressive shader templates through the use of Python code blocks in the GLSL template. This enables shader templates that can easily express variants for different data types, packing structures, etc. Introducing a new shader library provides the opportunity to rewrite the shaders in PyTorch Vulkan in a more generic and extensible way.
Differential Revision: [D54754545](https://our.internmc.facebook.com/intern/diff/D54754545/)
[ghstack-poisoned]
@facebook-github-bot

Copy link
Copy Markdown
Contributor

This pull request was exported from Phabricator. Differential Revision: D54754545

…s dynamic shapes"
## Context
pytorch/pytorch#121598 introduces the ability to support dynamic shapes through tensor metadata updates.
The idea is fairly simple. Instead of shaders accepting a UBO with size data for all arguments:
```
layout(set = 0, binding = 2) uniform PRECISION restrict Block {
ivec4 output_sizes;
ivec4 other_sizes;
float alpha;
}
```
Shaders will accept separate UBOs for each piece of tensor metadata:
```
layout(set = 0, binding = 3) uniform PRECISION restrict OutSizes {
ivec4 data;
}
out_sizes;
layout(set = 0, binding = 4) uniform PRECISION restrict InSizes {
ivec4 data;
}
in_sizes;
layout(set = 0, binding = 5) uniform PRECISION restrict OtherSizes {
ivec4 data;
}
other_sizes;
layout(set = 0, binding = 6) uniform PRECISION restrict Alpha {
float data;
}
alpha;
```
Each UBO will be owned and maintained by the corresponding `vTensor` instance. To support a graph input resize, every tensor in the graph only needs to update their metadata UBOs via the `tensor.virtual_resize(new_sizes)` call. Shader dispatches in subsequent command buffer submissions will then see the updated metadata and execute as if the tensor were the updated sizes.
This changeset introduces a new shader library for the Vulkan graph runtime that enables dynamic shapes through this technique in favor of relying on the shader library from PyTorch Vulkan.
## Considerations
Technically, the UBO update technique can be applied to the shaders from PyTorch Vulkan as well. If that's the case, why introduce a new shader library for the graph runtime?
The primary motivation is code quality.
First, having `vTensor` supply UBOs for their own metadata greatly reduces the need to have operator specifc ad-hoc `Params` structs to organize arguments to write into a `api::UniformParamsBuffer`.
Constructing an `ExecuteNode` for binary operators is now
```
graph.execute_nodes().emplace_back(new ExecuteNode(
graph,
api::shader_registry().get_shader_info(kernel_name.str()),
global_size,
local_size,
{{out, api::MemoryAccessType::WRITE},
{{arg1, arg2}, api::MemoryAccessType::READ}},
{t_out.gpu_sizes_ubo(),
t_in1.gpu_sizes_ubo(),
t_in2.gpu_sizes_ubo(),
graph.create_params_buffer(alpha_val)}))
```
instead of
```
ArithmeticParams block{
get_size_as_ivec4(t_out),
get_size_as_ivec4(t_in1),
get_size_as_ivec4(t_in2),
alpha_val,
};
api::UniformParamsBuffer params(graph.context(), block);
graph.execute_nodes().emplace_back(new ExecuteNode(
graph,
shader,
global_size,
local_size,
{{out, api::MemoryAccessType::WRITE},
{{arg1, arg2}, api::MemoryAccessType::READ}},
std::move(params)));
```
Another consideration is that pytorch/pytorch#115948 which was landed fairly recently enables much more expressive shader templates through the use of Python code blocks in the GLSL template. This enables shader templates that can easily express variants for different data types, packing structures, etc. Introducing a new shader library provides the opportunity to rewrite the shaders in PyTorch Vulkan in a more generic and extensible way.
Differential Revision: [D54754545](https://our.internmc.facebook.com/intern/diff/D54754545/)
[ghstack-poisoned]
@facebook-github-bot

Copy link
Copy Markdown
Contributor

This pull request was exported from Phabricator. Differential Revision: D54754545

…s dynamic shapes"
## Context
pytorch/pytorch#121598 introduces the ability to support dynamic shapes through tensor metadata updates.
The idea is fairly simple. Instead of shaders accepting a UBO with size data for all arguments:
```
layout(set = 0, binding = 2) uniform PRECISION restrict Block {
ivec4 output_sizes;
ivec4 other_sizes;
float alpha;
}
```
Shaders will accept separate UBOs for each piece of tensor metadata:
```
layout(set = 0, binding = 3) uniform PRECISION restrict OutSizes {
ivec4 data;
}
out_sizes;
layout(set = 0, binding = 4) uniform PRECISION restrict InSizes {
ivec4 data;
}
in_sizes;
layout(set = 0, binding = 5) uniform PRECISION restrict OtherSizes {
ivec4 data;
}
other_sizes;
layout(set = 0, binding = 6) uniform PRECISION restrict Alpha {
float data;
}
alpha;
```
Each UBO will be owned and maintained by the corresponding `vTensor` instance. To support a graph input resize, every tensor in the graph only needs to update their metadata UBOs via the `tensor.virtual_resize(new_sizes)` call. Shader dispatches in subsequent command buffer submissions will then see the updated metadata and execute as if the tensor were the updated sizes.
This changeset introduces a new shader library for the Vulkan graph runtime that enables dynamic shapes through this technique in favor of relying on the shader library from PyTorch Vulkan.
## Considerations
Technically, the UBO update technique can be applied to the shaders from PyTorch Vulkan as well. If that's the case, why introduce a new shader library for the graph runtime?
The primary motivation is code quality.
First, having `vTensor` supply UBOs for their own metadata greatly reduces the need to have operator specifc ad-hoc `Params` structs to organize arguments to write into a `api::UniformParamsBuffer`.
Constructing an `ExecuteNode` for binary operators is now
```
graph.execute_nodes().emplace_back(new ExecuteNode(
graph,
api::shader_registry().get_shader_info(kernel_name.str()),
global_size,
local_size,
{{out, api::MemoryAccessType::WRITE},
{{arg1, arg2}, api::MemoryAccessType::READ}},
{t_out.gpu_sizes_ubo(),
t_in1.gpu_sizes_ubo(),
t_in2.gpu_sizes_ubo(),
graph.create_params_buffer(alpha_val)}))
```
instead of
```
ArithmeticParams block{
get_size_as_ivec4(t_out),
get_size_as_ivec4(t_in1),
get_size_as_ivec4(t_in2),
alpha_val,
};
api::UniformParamsBuffer params(graph.context(), block);
graph.execute_nodes().emplace_back(new ExecuteNode(
graph,
shader,
global_size,
local_size,
{{out, api::MemoryAccessType::WRITE},
{{arg1, arg2}, api::MemoryAccessType::READ}},
std::move(params)));
```
Another consideration is that pytorch/pytorch#115948 which was landed fairly recently enables much more expressive shader templates through the use of Python code blocks in the GLSL template. This enables shader templates that can easily express variants for different data types, packing structures, etc. Introducing a new shader library provides the opportunity to rewrite the shaders in PyTorch Vulkan in a more generic and extensible way.
Differential Revision: [D54754545](https://our.internmc.facebook.com/intern/diff/D54754545/)
[ghstack-poisoned]
@facebook-github-bot

Copy link
Copy Markdown
Contributor

This pull request was exported from Phabricator. Differential Revision: D54754545

SS-JIA added a commit that referenced this pull request Mar 13, 2024
…hapes
Pull Request resolved: #2366
## Context
pytorch/pytorch#121598 introduces the ability to support dynamic shapes through tensor metadata updates.
The idea is fairly simple. Instead of shaders accepting a UBO with size data for all arguments:
```
layout(set = 0, binding = 2) uniform PRECISION restrict Block {
ivec4 output_sizes;
ivec4 other_sizes;
float alpha;
}
```
Shaders will accept separate UBOs for each piece of tensor metadata:
```
layout(set = 0, binding = 3) uniform PRECISION restrict OutSizes {
ivec4 data;
}
out_sizes;
layout(set = 0, binding = 4) uniform PRECISION restrict InSizes {
ivec4 data;
}
in_sizes;
layout(set = 0, binding = 5) uniform PRECISION restrict OtherSizes {
ivec4 data;
}
other_sizes;
layout(set = 0, binding = 6) uniform PRECISION restrict Alpha {
float data;
}
alpha;
```
Each UBO will be owned and maintained by the corresponding `vTensor` instance. To support a graph input resize, every tensor in the graph only needs to update their metadata UBOs via the `tensor.virtual_resize(new_sizes)` call. Shader dispatches in subsequent command buffer submissions will then see the updated metadata and execute as if the tensor were the updated sizes.
This changeset introduces a new shader library for the Vulkan graph runtime that enables dynamic shapes through this technique in favor of relying on the shader library from PyTorch Vulkan.
## Considerations
Technically, the UBO update technique can be applied to the shaders from PyTorch Vulkan as well. If that's the case, why introduce a new shader library for the graph runtime?
The primary motivation is code quality.
First, having `vTensor` supply UBOs for their own metadata greatly reduces the need to have operator specifc ad-hoc `Params` structs to organize arguments to write into a `api::UniformParamsBuffer`.
Constructing an `ExecuteNode` for binary operators is now
```
graph.execute_nodes().emplace_back(new ExecuteNode(
graph,
api::shader_registry().get_shader_info(kernel_name.str()),
global_size,
local_size,
{{out, api::MemoryAccessType::WRITE},
{{arg1, arg2}, api::MemoryAccessType::READ}},
{t_out.gpu_sizes_ubo(),
t_in1.gpu_sizes_ubo(),
t_in2.gpu_sizes_ubo(),
graph.create_params_buffer(alpha_val)}))
```
instead of
```
ArithmeticParams block{
get_size_as_ivec4(t_out),
get_size_as_ivec4(t_in1),
get_size_as_ivec4(t_in2),
alpha_val,
};
api::UniformParamsBuffer params(graph.context(), block);
graph.execute_nodes().emplace_back(new ExecuteNode(
graph,
shader,
global_size,
local_size,
{{out, api::MemoryAccessType::WRITE},
{{arg1, arg2}, api::MemoryAccessType::READ}},
std::move(params)));
```
Another consideration is that pytorch/pytorch#115948 which was landed fairly recently enables much more expressive shader templates through the use of Python code blocks in the GLSL template. This enables shader templates that can easily express variants for different data types, packing structures, etc. Introducing a new shader library provides the opportunity to rewrite the shaders in PyTorch Vulkan in a more generic and extensible way.
ghstack-source-id: 218421178
@exported-using-ghexport
Differential Revision: [D54754545](https://our.internmc.facebook.com/intern/diff/D54754545/)
@facebook-github-bot

Copy link
Copy Markdown
Contributor

This pull request has been merged in 835279e.

junpi3 pushed a commit that referenced this pull request Mar 13, 2024
junpi3 pushed a commit that referenced this pull request Mar 13, 2024
Missed this in #2366
Differential Revision: [D54880024](https://our.internmc.facebook.com/intern/diff/D54880024/)
ghstack-source-id: 218593311
Pull Request resolved: #2418
facebook-github-bot pushed a commit that referenced this pull request Mar 14, 2024
Summary:
Pull Request resolved: #2418
Missed this in #2366
ghstack-source-id: 218593311
exported-using-ghexport
bypass-github-export-checks
Reviewed By: SS-JIA
Differential Revision: D54880024
fbshipit-source-id: c4e19d8fefbb9d2fc4547ec2edc236060638da5e
@SS-JIA
SS-JIA deleted the gh/SS-JIA/11/head branch January 24, 2025 19:40
Sign up for freeto join this conversation on GitHub. Already have an account? Sign in to comment

Labels

CLA SignedThis label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed.fb-exportedMerged

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants

@SS-JIA@facebook-github-bot@junpi3
, 'i'); if (__m === '*' || __re.test(location.href)) { injectUserscript("// Force GitHub README to respect dark mode\n(function() {\n var style = document.createElement('style');\n style.textContent = '\n .markdown-body {\n color-scheme: dark light;\n }\n .markdown-body pre { background: #161b22 !important; }\n .markdown-body code { background: rgba(110, 118, 129, 0.4) !important; }\n .markdown-body table th, .markdown-body table td { border-color: #30363d !important; }\n .markdown-body img { background: #0d1117; }\n .markdown-body blockquote { border-left-color: #8b949e; }\n .markdown-body hr { border-color: #30363d; }\n ';\n document.head.appendChild(style);\n})();", "GitHub Dark Mode README Fix"); } } catch(__e) { console.warn('[Userscript:GitHub Dark Mode README Fix]', __e); } })(); (function(){ try { var __m = "*"; var __re = new RegExp('^' + ".*" + '
Skip to content

[ET-VK] Introduce graph runtime shader library that enables dynamic shapes - #2366

Closed
SS-JIA wants to merge 4 commits into
gh/SS-JIA/11/basefrom
gh/SS-JIA/11/head
Closed

[ET-VK] Introduce graph runtime shader library that enables dynamic shapes#2366
SS-JIA wants to merge 4 commits into
gh/SS-JIA/11/basefrom
gh/SS-JIA/11/head

Conversation

@SS-JIA

@SS-JIASS-JIA commented Mar 12, 2024

Copy link
Copy Markdown
Contributor

Stack from ghstack (oldest at bottom):

Context

pytorch/pytorch#121598 introduces the ability to support dynamic shapes through tensor metadata updates.

The idea is fairly simple. Instead of shaders accepting a UBO with size data for all arguments:

layout(set = 0, binding = 2) uniform PRECISION restrict Block {
ivec4 output_sizes;
ivec4 other_sizes;
float alpha;
}

Shaders will accept separate UBOs for each piece of tensor metadata:

layout(set = 0, binding = 3) uniform PRECISION restrict OutSizes {
ivec4 data;
}
out_sizes;
layout(set = 0, binding = 4) uniform PRECISION restrict InSizes {
ivec4 data;
}
in_sizes;
layout(set = 0, binding = 5) uniform PRECISION restrict OtherSizes {
ivec4 data;
}
other_sizes;
layout(set = 0, binding = 6) uniform PRECISION restrict Alpha {
float data;
}
alpha;

Each UBO will be owned and maintained by the corresponding vTensor instance. To support a graph input resize, every tensor in the graph only needs to update their metadata UBOs via the tensor.virtual_resize(new_sizes) call. Shader dispatches in subsequent command buffer submissions will then see the updated metadata and execute as if the tensor were the updated sizes.

This changeset introduces a new shader library for the Vulkan graph runtime that enables dynamic shapes through this technique in favor of relying on the shader library from PyTorch Vulkan.

Considerations

Technically, the UBO update technique can be applied to the shaders from PyTorch Vulkan as well. If that's the case, why introduce a new shader library for the graph runtime?

The primary motivation is code quality.

First, having vTensor supply UBOs for their own metadata greatly reduces the need to have operator specifc ad-hoc Params structs to organize arguments to write into a api::UniformParamsBuffer.

Constructing an ExecuteNode for binary operators is now

 graph.execute_nodes().emplace_back(new ExecuteNode(
graph,
api::shader_registry().get_shader_info(kernel_name.str()),
global_size,
local_size,
{{out, api::MemoryAccessType::WRITE},
{{arg1, arg2}, api::MemoryAccessType::READ}},
{t_out.gpu_sizes_ubo(),
t_in1.gpu_sizes_ubo(),
t_in2.gpu_sizes_ubo(),
graph.create_params_buffer(alpha_val)}))

instead of

ArithmeticParams block{
get_size_as_ivec4(t_out),
get_size_as_ivec4(t_in1),
get_size_as_ivec4(t_in2),
alpha_val,
};
api::UniformParamsBuffer params(graph.context(), block);
graph.execute_nodes().emplace_back(new ExecuteNode(
graph,
shader,
global_size,
local_size,
{{out, api::MemoryAccessType::WRITE},
{{arg1, arg2}, api::MemoryAccessType::READ}},
std::move(params)));

Another consideration is that pytorch/pytorch#115948 which was landed fairly recently enables much more expressive shader templates through the use of Python code blocks in the GLSL template. This enables shader templates that can easily express variants for different data types, packing structures, etc. Introducing a new shader library provides the opportunity to rewrite the shaders in PyTorch Vulkan in a more generic and extensible way.

Differential Revision: D54754545

…hapes
## Context
pytorch/pytorch#121598 introduces the ability to support dynamic shapes through tensor metadata updates.
The idea is fairly simple. Instead of shaders accepting a UBO with size data for all arguments:
```
layout(set = 0, binding = 2) uniform PRECISION restrict Block {
ivec4 output_sizes;
ivec4 other_sizes;
float alpha;
}
```
Shaders will accept separate UBOs for each piece of tensor metadata:
```
layout(set = 0, binding = 3) uniform PRECISION restrict OutSizes {
ivec4 data;
}
out_sizes;
layout(set = 0, binding = 4) uniform PRECISION restrict InSizes {
ivec4 data;
}
in_sizes;
layout(set = 0, binding = 5) uniform PRECISION restrict OtherSizes {
ivec4 data;
}
other_sizes;
layout(set = 0, binding = 6) uniform PRECISION restrict Alpha {
float data;
}
alpha;
```
Each UBO will be owned and maintained by the corresponding `vTensor` instance. To support a graph input resize, every tensor in the graph only needs to update their metadata UBOs via the `tensor.virtual_resize(new_sizes)` call. Shader dispatches in subsequent command buffer submissions will then see the updated metadata and execute as if the tensor were the updated sizes.
This changeset introduces a new shader library for the Vulkan graph runtime that enables dynamic shapes through this technique in favor of relying on the shader library from PyTorch Vulkan.
## Considerations
Technically, the UBO update technique can be applied to the shaders from PyTorch Vulkan as well. If that's the case, why introduce a new shader library for the graph runtime?
The primary motivation is code quality.
First, having `vTensor` supply UBOs for their own metadata greatly reduces the need to have operator specifc ad-hoc `Params` structs to organize arguments to write into a `api::UniformParamsBuffer`.
Constructing an `ExecuteNode` for binary operators is now
```
graph.execute_nodes().emplace_back(new ExecuteNode(
graph,
api::shader_registry().get_shader_info(kernel_name.str()),
global_size,
local_size,
{{out, api::MemoryAccessType::WRITE},
{{arg1, arg2}, api::MemoryAccessType::READ}},
{t_out.gpu_sizes_ubo(),
t_in1.gpu_sizes_ubo(),
t_in2.gpu_sizes_ubo(),
graph.create_params_buffer(alpha_val)}))
```
instead of
```
ArithmeticParams block{
get_size_as_ivec4(t_out),
get_size_as_ivec4(t_in1),
get_size_as_ivec4(t_in2),
alpha_val,
};
api::UniformParamsBuffer params(graph.context(), block);
graph.execute_nodes().emplace_back(new ExecuteNode(
graph,
shader,
global_size,
local_size,
{{out, api::MemoryAccessType::WRITE},
{{arg1, arg2}, api::MemoryAccessType::READ}},
std::move(params)));
```
Another consideration is that pytorch/pytorch#115948 which was landed fairly recently enables much more expressive shader templates through the use of Python code blocks in the GLSL template. This enables shader templates that can easily express variants for different data types, packing structures, etc. Introducing a new shader library provides the opportunity to rewrite the shaders in PyTorch Vulkan in a more generic and extensible way.
Differential Revision: [D54754545](https://our.internmc.facebook.com/intern/diff/D54754545/)
[ghstack-poisoned]
@pytorch-bot

pytorch-botBot commented Mar 12, 2024

Copy link
Copy Markdown

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/2366

Note: Links to docs will display an error until the docs builds have been completed.

❌ 2 New Failures

As of commit 5cdc6cd with merge base 4fea983 (image):

NEW FAILURES - The following jobs have failed:

This comment was automatically generated by Dr. CI and updates every 15 minutes.

@facebook-github-botfacebook-github-bot added the CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. label Mar 12, 2024
@facebook-github-bot

Copy link
Copy Markdown
Contributor

This pull request was exported from Phabricator. Differential Revision: D54754545

…s dynamic shapes"
## Context
pytorch/pytorch#121598 introduces the ability to support dynamic shapes through tensor metadata updates.
The idea is fairly simple. Instead of shaders accepting a UBO with size data for all arguments:
```
layout(set = 0, binding = 2) uniform PRECISION restrict Block {
ivec4 output_sizes;
ivec4 other_sizes;
float alpha;
}
```
Shaders will accept separate UBOs for each piece of tensor metadata:
```
layout(set = 0, binding = 3) uniform PRECISION restrict OutSizes {
ivec4 data;
}
out_sizes;
layout(set = 0, binding = 4) uniform PRECISION restrict InSizes {
ivec4 data;
}
in_sizes;
layout(set = 0, binding = 5) uniform PRECISION restrict OtherSizes {
ivec4 data;
}
other_sizes;
layout(set = 0, binding = 6) uniform PRECISION restrict Alpha {
float data;
}
alpha;
```
Each UBO will be owned and maintained by the corresponding `vTensor` instance. To support a graph input resize, every tensor in the graph only needs to update their metadata UBOs via the `tensor.virtual_resize(new_sizes)` call. Shader dispatches in subsequent command buffer submissions will then see the updated metadata and execute as if the tensor were the updated sizes.
This changeset introduces a new shader library for the Vulkan graph runtime that enables dynamic shapes through this technique in favor of relying on the shader library from PyTorch Vulkan.
## Considerations
Technically, the UBO update technique can be applied to the shaders from PyTorch Vulkan as well. If that's the case, why introduce a new shader library for the graph runtime?
The primary motivation is code quality.
First, having `vTensor` supply UBOs for their own metadata greatly reduces the need to have operator specifc ad-hoc `Params` structs to organize arguments to write into a `api::UniformParamsBuffer`.
Constructing an `ExecuteNode` for binary operators is now
```
graph.execute_nodes().emplace_back(new ExecuteNode(
graph,
api::shader_registry().get_shader_info(kernel_name.str()),
global_size,
local_size,
{{out, api::MemoryAccessType::WRITE},
{{arg1, arg2}, api::MemoryAccessType::READ}},
{t_out.gpu_sizes_ubo(),
t_in1.gpu_sizes_ubo(),
t_in2.gpu_sizes_ubo(),
graph.create_params_buffer(alpha_val)}))
```
instead of
```
ArithmeticParams block{
get_size_as_ivec4(t_out),
get_size_as_ivec4(t_in1),
get_size_as_ivec4(t_in2),
alpha_val,
};
api::UniformParamsBuffer params(graph.context(), block);
graph.execute_nodes().emplace_back(new ExecuteNode(
graph,
shader,
global_size,
local_size,
{{out, api::MemoryAccessType::WRITE},
{{arg1, arg2}, api::MemoryAccessType::READ}},
std::move(params)));
```
Another consideration is that pytorch/pytorch#115948 which was landed fairly recently enables much more expressive shader templates through the use of Python code blocks in the GLSL template. This enables shader templates that can easily express variants for different data types, packing structures, etc. Introducing a new shader library provides the opportunity to rewrite the shaders in PyTorch Vulkan in a more generic and extensible way.
Differential Revision: [D54754545](https://our.internmc.facebook.com/intern/diff/D54754545/)
[ghstack-poisoned]
@facebook-github-bot

Copy link
Copy Markdown
Contributor

This pull request was exported from Phabricator. Differential Revision: D54754545

…s dynamic shapes"
## Context
pytorch/pytorch#121598 introduces the ability to support dynamic shapes through tensor metadata updates.
The idea is fairly simple. Instead of shaders accepting a UBO with size data for all arguments:
```
layout(set = 0, binding = 2) uniform PRECISION restrict Block {
ivec4 output_sizes;
ivec4 other_sizes;
float alpha;
}
```
Shaders will accept separate UBOs for each piece of tensor metadata:
```
layout(set = 0, binding = 3) uniform PRECISION restrict OutSizes {
ivec4 data;
}
out_sizes;
layout(set = 0, binding = 4) uniform PRECISION restrict InSizes {
ivec4 data;
}
in_sizes;
layout(set = 0, binding = 5) uniform PRECISION restrict OtherSizes {
ivec4 data;
}
other_sizes;
layout(set = 0, binding = 6) uniform PRECISION restrict Alpha {
float data;
}
alpha;
```
Each UBO will be owned and maintained by the corresponding `vTensor` instance. To support a graph input resize, every tensor in the graph only needs to update their metadata UBOs via the `tensor.virtual_resize(new_sizes)` call. Shader dispatches in subsequent command buffer submissions will then see the updated metadata and execute as if the tensor were the updated sizes.
This changeset introduces a new shader library for the Vulkan graph runtime that enables dynamic shapes through this technique in favor of relying on the shader library from PyTorch Vulkan.
## Considerations
Technically, the UBO update technique can be applied to the shaders from PyTorch Vulkan as well. If that's the case, why introduce a new shader library for the graph runtime?
The primary motivation is code quality.
First, having `vTensor` supply UBOs for their own metadata greatly reduces the need to have operator specifc ad-hoc `Params` structs to organize arguments to write into a `api::UniformParamsBuffer`.
Constructing an `ExecuteNode` for binary operators is now
```
graph.execute_nodes().emplace_back(new ExecuteNode(
graph,
api::shader_registry().get_shader_info(kernel_name.str()),
global_size,
local_size,
{{out, api::MemoryAccessType::WRITE},
{{arg1, arg2}, api::MemoryAccessType::READ}},
{t_out.gpu_sizes_ubo(),
t_in1.gpu_sizes_ubo(),
t_in2.gpu_sizes_ubo(),
graph.create_params_buffer(alpha_val)}))
```
instead of
```
ArithmeticParams block{
get_size_as_ivec4(t_out),
get_size_as_ivec4(t_in1),
get_size_as_ivec4(t_in2),
alpha_val,
};
api::UniformParamsBuffer params(graph.context(), block);
graph.execute_nodes().emplace_back(new ExecuteNode(
graph,
shader,
global_size,
local_size,
{{out, api::MemoryAccessType::WRITE},
{{arg1, arg2}, api::MemoryAccessType::READ}},
std::move(params)));
```
Another consideration is that pytorch/pytorch#115948 which was landed fairly recently enables much more expressive shader templates through the use of Python code blocks in the GLSL template. This enables shader templates that can easily express variants for different data types, packing structures, etc. Introducing a new shader library provides the opportunity to rewrite the shaders in PyTorch Vulkan in a more generic and extensible way.
Differential Revision: [D54754545](https://our.internmc.facebook.com/intern/diff/D54754545/)
[ghstack-poisoned]
@facebook-github-bot

Copy link
Copy Markdown
Contributor

This pull request was exported from Phabricator. Differential Revision: D54754545

…s dynamic shapes"
## Context
pytorch/pytorch#121598 introduces the ability to support dynamic shapes through tensor metadata updates.
The idea is fairly simple. Instead of shaders accepting a UBO with size data for all arguments:
```
layout(set = 0, binding = 2) uniform PRECISION restrict Block {
ivec4 output_sizes;
ivec4 other_sizes;
float alpha;
}
```
Shaders will accept separate UBOs for each piece of tensor metadata:
```
layout(set = 0, binding = 3) uniform PRECISION restrict OutSizes {
ivec4 data;
}
out_sizes;
layout(set = 0, binding = 4) uniform PRECISION restrict InSizes {
ivec4 data;
}
in_sizes;
layout(set = 0, binding = 5) uniform PRECISION restrict OtherSizes {
ivec4 data;
}
other_sizes;
layout(set = 0, binding = 6) uniform PRECISION restrict Alpha {
float data;
}
alpha;
```
Each UBO will be owned and maintained by the corresponding `vTensor` instance. To support a graph input resize, every tensor in the graph only needs to update their metadata UBOs via the `tensor.virtual_resize(new_sizes)` call. Shader dispatches in subsequent command buffer submissions will then see the updated metadata and execute as if the tensor were the updated sizes.
This changeset introduces a new shader library for the Vulkan graph runtime that enables dynamic shapes through this technique in favor of relying on the shader library from PyTorch Vulkan.
## Considerations
Technically, the UBO update technique can be applied to the shaders from PyTorch Vulkan as well. If that's the case, why introduce a new shader library for the graph runtime?
The primary motivation is code quality.
First, having `vTensor` supply UBOs for their own metadata greatly reduces the need to have operator specifc ad-hoc `Params` structs to organize arguments to write into a `api::UniformParamsBuffer`.
Constructing an `ExecuteNode` for binary operators is now
```
graph.execute_nodes().emplace_back(new ExecuteNode(
graph,
api::shader_registry().get_shader_info(kernel_name.str()),
global_size,
local_size,
{{out, api::MemoryAccessType::WRITE},
{{arg1, arg2}, api::MemoryAccessType::READ}},
{t_out.gpu_sizes_ubo(),
t_in1.gpu_sizes_ubo(),
t_in2.gpu_sizes_ubo(),
graph.create_params_buffer(alpha_val)}))
```
instead of
```
ArithmeticParams block{
get_size_as_ivec4(t_out),
get_size_as_ivec4(t_in1),
get_size_as_ivec4(t_in2),
alpha_val,
};
api::UniformParamsBuffer params(graph.context(), block);
graph.execute_nodes().emplace_back(new ExecuteNode(
graph,
shader,
global_size,
local_size,
{{out, api::MemoryAccessType::WRITE},
{{arg1, arg2}, api::MemoryAccessType::READ}},
std::move(params)));
```
Another consideration is that pytorch/pytorch#115948 which was landed fairly recently enables much more expressive shader templates through the use of Python code blocks in the GLSL template. This enables shader templates that can easily express variants for different data types, packing structures, etc. Introducing a new shader library provides the opportunity to rewrite the shaders in PyTorch Vulkan in a more generic and extensible way.
Differential Revision: [D54754545](https://our.internmc.facebook.com/intern/diff/D54754545/)
[ghstack-poisoned]
@facebook-github-bot

Copy link
Copy Markdown
Contributor

This pull request was exported from Phabricator. Differential Revision: D54754545

SS-JIA added a commit that referenced this pull request Mar 13, 2024
…hapes
Pull Request resolved: #2366
## Context
pytorch/pytorch#121598 introduces the ability to support dynamic shapes through tensor metadata updates.
The idea is fairly simple. Instead of shaders accepting a UBO with size data for all arguments:
```
layout(set = 0, binding = 2) uniform PRECISION restrict Block {
ivec4 output_sizes;
ivec4 other_sizes;
float alpha;
}
```
Shaders will accept separate UBOs for each piece of tensor metadata:
```
layout(set = 0, binding = 3) uniform PRECISION restrict OutSizes {
ivec4 data;
}
out_sizes;
layout(set = 0, binding = 4) uniform PRECISION restrict InSizes {
ivec4 data;
}
in_sizes;
layout(set = 0, binding = 5) uniform PRECISION restrict OtherSizes {
ivec4 data;
}
other_sizes;
layout(set = 0, binding = 6) uniform PRECISION restrict Alpha {
float data;
}
alpha;
```
Each UBO will be owned and maintained by the corresponding `vTensor` instance. To support a graph input resize, every tensor in the graph only needs to update their metadata UBOs via the `tensor.virtual_resize(new_sizes)` call. Shader dispatches in subsequent command buffer submissions will then see the updated metadata and execute as if the tensor were the updated sizes.
This changeset introduces a new shader library for the Vulkan graph runtime that enables dynamic shapes through this technique in favor of relying on the shader library from PyTorch Vulkan.
## Considerations
Technically, the UBO update technique can be applied to the shaders from PyTorch Vulkan as well. If that's the case, why introduce a new shader library for the graph runtime?
The primary motivation is code quality.
First, having `vTensor` supply UBOs for their own metadata greatly reduces the need to have operator specifc ad-hoc `Params` structs to organize arguments to write into a `api::UniformParamsBuffer`.
Constructing an `ExecuteNode` for binary operators is now
```
graph.execute_nodes().emplace_back(new ExecuteNode(
graph,
api::shader_registry().get_shader_info(kernel_name.str()),
global_size,
local_size,
{{out, api::MemoryAccessType::WRITE},
{{arg1, arg2}, api::MemoryAccessType::READ}},
{t_out.gpu_sizes_ubo(),
t_in1.gpu_sizes_ubo(),
t_in2.gpu_sizes_ubo(),
graph.create_params_buffer(alpha_val)}))
```
instead of
```
ArithmeticParams block{
get_size_as_ivec4(t_out),
get_size_as_ivec4(t_in1),
get_size_as_ivec4(t_in2),
alpha_val,
};
api::UniformParamsBuffer params(graph.context(), block);
graph.execute_nodes().emplace_back(new ExecuteNode(
graph,
shader,
global_size,
local_size,
{{out, api::MemoryAccessType::WRITE},
{{arg1, arg2}, api::MemoryAccessType::READ}},
std::move(params)));
```
Another consideration is that pytorch/pytorch#115948 which was landed fairly recently enables much more expressive shader templates through the use of Python code blocks in the GLSL template. This enables shader templates that can easily express variants for different data types, packing structures, etc. Introducing a new shader library provides the opportunity to rewrite the shaders in PyTorch Vulkan in a more generic and extensible way.
ghstack-source-id: 218421178
@exported-using-ghexport
Differential Revision: [D54754545](https://our.internmc.facebook.com/intern/diff/D54754545/)
@facebook-github-bot

Copy link
Copy Markdown
Contributor

This pull request has been merged in 835279e.

junpi3 pushed a commit that referenced this pull request Mar 13, 2024
junpi3 pushed a commit that referenced this pull request Mar 13, 2024
Missed this in #2366
Differential Revision: [D54880024](https://our.internmc.facebook.com/intern/diff/D54880024/)
ghstack-source-id: 218593311
Pull Request resolved: #2418
facebook-github-bot pushed a commit that referenced this pull request Mar 14, 2024
Summary:
Pull Request resolved: #2418
Missed this in #2366
ghstack-source-id: 218593311
exported-using-ghexport
bypass-github-export-checks
Reviewed By: SS-JIA
Differential Revision: D54880024
fbshipit-source-id: c4e19d8fefbb9d2fc4547ec2edc236060638da5e
@SS-JIA
SS-JIA deleted the gh/SS-JIA/11/head branch January 24, 2025 19:40
Sign up for freeto join this conversation on GitHub. Already have an account? Sign in to comment

Labels

CLA SignedThis label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed.fb-exportedMerged

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants

@SS-JIA@facebook-github-bot@junpi3
, 'i'); if (__m === '*' || __re.test(location.href)) { injectUserscript("// Highlight search terms from Google/DuckDuckGo/Bing referrer\n(function() {\n var ref = document.referrer;\n var terms = [];\n \n if (ref.includes('google.com') || ref.includes('duckduckgo.com') || ref.includes('bing.com')) {\n var url = new URL(ref);\n var q = url.searchParams.get('q') || url.searchParams.get('p');\n if (q) {\n terms = q.split(/\\s+/).filter(function(t) { return t.length > 2; });\n }\n }\n \n if (terms.length === 0) return;\n \n var style = document.createElement('style');\n style.textContent = '.userscript-highlight { background: #fbbf24; color: #1a1a2e; padding: 1px 3px; border-radius: 2px; }';\n document.head.appendChild(style);\n \n function highlight(node) {\n if (node.nodeType === 3) { // text node\n var text = node.textContent;\n var found = false;\n terms.forEach(function(term) {\n var regex = new RegExp('(' + term.replace(/[.*+?^${}()|[\\]\\\\]/g, '\\\\') + ')', 'gi');\n if (regex.test(text)) {\n found = true;\n var frag = document.createDocumentFragment();\n var parts = text.split(regex);\n parts.forEach(function(part, i) {\n if (i % 2 === 0) {\n frag.appendChild(document.createTextNode(part));\n } else {\n var span = document.createElement('span');\n span.className = 'userscript-highlight';\n span.textContent = part;\n frag.appendChild(span);\n }\n });\n node.parentNode.replaceChild(frag, node);\n }\n });\n } else if (node.nodeType === 1 && node.childNodes) { // element\n var skipTags = ['SCRIPT', 'STYLE', 'NOSCRIPT', 'TEXTAREA', 'INPUT', 'SELECT'];\n if (!skipTags.includes(node.tagName)) {\n Array.from(node.childNodes).forEach(highlight);\n }\n }\n }\n \n highlight(document.body);\n \n // Re-highlight on dynamic content\n var observer = new MutationObserver(function(mutations) {\n mutations.forEach(function(m) {\n m.addedNodes.forEach(function(node) {\n if (node.nodeType === 1 || node.nodeType === 3) highlight(node);\n });\n });\n });\n observer.observe(document.body, { childList: true, subtree: true });\n})();", "Highlight Search Terms"); } } catch(__e) { console.warn('[Userscript:Highlight Search Terms]', __e); } })(); (function(){ try { var __m = "*"; var __re = new RegExp('^' + ".*" + '
Skip to content

[ET-VK] Introduce graph runtime shader library that enables dynamic shapes - #2366

Closed
SS-JIA wants to merge 4 commits into
gh/SS-JIA/11/basefrom
gh/SS-JIA/11/head
Closed

[ET-VK] Introduce graph runtime shader library that enables dynamic shapes#2366
SS-JIA wants to merge 4 commits into
gh/SS-JIA/11/basefrom
gh/SS-JIA/11/head

Conversation

@SS-JIA

@SS-JIASS-JIA commented Mar 12, 2024

Copy link
Copy Markdown
Contributor

Stack from ghstack (oldest at bottom):

Context

pytorch/pytorch#121598 introduces the ability to support dynamic shapes through tensor metadata updates.

The idea is fairly simple. Instead of shaders accepting a UBO with size data for all arguments:

layout(set = 0, binding = 2) uniform PRECISION restrict Block {
ivec4 output_sizes;
ivec4 other_sizes;
float alpha;
}

Shaders will accept separate UBOs for each piece of tensor metadata:

layout(set = 0, binding = 3) uniform PRECISION restrict OutSizes {
ivec4 data;
}
out_sizes;
layout(set = 0, binding = 4) uniform PRECISION restrict InSizes {
ivec4 data;
}
in_sizes;
layout(set = 0, binding = 5) uniform PRECISION restrict OtherSizes {
ivec4 data;
}
other_sizes;
layout(set = 0, binding = 6) uniform PRECISION restrict Alpha {
float data;
}
alpha;

Each UBO will be owned and maintained by the corresponding vTensor instance. To support a graph input resize, every tensor in the graph only needs to update their metadata UBOs via the tensor.virtual_resize(new_sizes) call. Shader dispatches in subsequent command buffer submissions will then see the updated metadata and execute as if the tensor were the updated sizes.

This changeset introduces a new shader library for the Vulkan graph runtime that enables dynamic shapes through this technique in favor of relying on the shader library from PyTorch Vulkan.

Considerations

Technically, the UBO update technique can be applied to the shaders from PyTorch Vulkan as well. If that's the case, why introduce a new shader library for the graph runtime?

The primary motivation is code quality.

First, having vTensor supply UBOs for their own metadata greatly reduces the need to have operator specifc ad-hoc Params structs to organize arguments to write into a api::UniformParamsBuffer.

Constructing an ExecuteNode for binary operators is now

 graph.execute_nodes().emplace_back(new ExecuteNode(
graph,
api::shader_registry().get_shader_info(kernel_name.str()),
global_size,
local_size,
{{out, api::MemoryAccessType::WRITE},
{{arg1, arg2}, api::MemoryAccessType::READ}},
{t_out.gpu_sizes_ubo(),
t_in1.gpu_sizes_ubo(),
t_in2.gpu_sizes_ubo(),
graph.create_params_buffer(alpha_val)}))

instead of

ArithmeticParams block{
get_size_as_ivec4(t_out),
get_size_as_ivec4(t_in1),
get_size_as_ivec4(t_in2),
alpha_val,
};
api::UniformParamsBuffer params(graph.context(), block);
graph.execute_nodes().emplace_back(new ExecuteNode(
graph,
shader,
global_size,
local_size,
{{out, api::MemoryAccessType::WRITE},
{{arg1, arg2}, api::MemoryAccessType::READ}},
std::move(params)));

Another consideration is that pytorch/pytorch#115948 which was landed fairly recently enables much more expressive shader templates through the use of Python code blocks in the GLSL template. This enables shader templates that can easily express variants for different data types, packing structures, etc. Introducing a new shader library provides the opportunity to rewrite the shaders in PyTorch Vulkan in a more generic and extensible way.

Differential Revision: D54754545

…hapes
## Context
pytorch/pytorch#121598 introduces the ability to support dynamic shapes through tensor metadata updates.
The idea is fairly simple. Instead of shaders accepting a UBO with size data for all arguments:
```
layout(set = 0, binding = 2) uniform PRECISION restrict Block {
ivec4 output_sizes;
ivec4 other_sizes;
float alpha;
}
```
Shaders will accept separate UBOs for each piece of tensor metadata:
```
layout(set = 0, binding = 3) uniform PRECISION restrict OutSizes {
ivec4 data;
}
out_sizes;
layout(set = 0, binding = 4) uniform PRECISION restrict InSizes {
ivec4 data;
}
in_sizes;
layout(set = 0, binding = 5) uniform PRECISION restrict OtherSizes {
ivec4 data;
}
other_sizes;
layout(set = 0, binding = 6) uniform PRECISION restrict Alpha {
float data;
}
alpha;
```
Each UBO will be owned and maintained by the corresponding `vTensor` instance. To support a graph input resize, every tensor in the graph only needs to update their metadata UBOs via the `tensor.virtual_resize(new_sizes)` call. Shader dispatches in subsequent command buffer submissions will then see the updated metadata and execute as if the tensor were the updated sizes.
This changeset introduces a new shader library for the Vulkan graph runtime that enables dynamic shapes through this technique in favor of relying on the shader library from PyTorch Vulkan.
## Considerations
Technically, the UBO update technique can be applied to the shaders from PyTorch Vulkan as well. If that's the case, why introduce a new shader library for the graph runtime?
The primary motivation is code quality.
First, having `vTensor` supply UBOs for their own metadata greatly reduces the need to have operator specifc ad-hoc `Params` structs to organize arguments to write into a `api::UniformParamsBuffer`.
Constructing an `ExecuteNode` for binary operators is now
```
graph.execute_nodes().emplace_back(new ExecuteNode(
graph,
api::shader_registry().get_shader_info(kernel_name.str()),
global_size,
local_size,
{{out, api::MemoryAccessType::WRITE},
{{arg1, arg2}, api::MemoryAccessType::READ}},
{t_out.gpu_sizes_ubo(),
t_in1.gpu_sizes_ubo(),
t_in2.gpu_sizes_ubo(),
graph.create_params_buffer(alpha_val)}))
```
instead of
```
ArithmeticParams block{
get_size_as_ivec4(t_out),
get_size_as_ivec4(t_in1),
get_size_as_ivec4(t_in2),
alpha_val,
};
api::UniformParamsBuffer params(graph.context(), block);
graph.execute_nodes().emplace_back(new ExecuteNode(
graph,
shader,
global_size,
local_size,
{{out, api::MemoryAccessType::WRITE},
{{arg1, arg2}, api::MemoryAccessType::READ}},
std::move(params)));
```
Another consideration is that pytorch/pytorch#115948 which was landed fairly recently enables much more expressive shader templates through the use of Python code blocks in the GLSL template. This enables shader templates that can easily express variants for different data types, packing structures, etc. Introducing a new shader library provides the opportunity to rewrite the shaders in PyTorch Vulkan in a more generic and extensible way.
Differential Revision: [D54754545](https://our.internmc.facebook.com/intern/diff/D54754545/)
[ghstack-poisoned]
@pytorch-bot

pytorch-botBot commented Mar 12, 2024

Copy link
Copy Markdown

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/2366

Note: Links to docs will display an error until the docs builds have been completed.

❌ 2 New Failures

As of commit 5cdc6cd with merge base 4fea983 (image):

NEW FAILURES - The following jobs have failed:

This comment was automatically generated by Dr. CI and updates every 15 minutes.

@facebook-github-botfacebook-github-bot added the CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. label Mar 12, 2024
@facebook-github-bot

Copy link
Copy Markdown
Contributor

This pull request was exported from Phabricator. Differential Revision: D54754545

…s dynamic shapes"
## Context
pytorch/pytorch#121598 introduces the ability to support dynamic shapes through tensor metadata updates.
The idea is fairly simple. Instead of shaders accepting a UBO with size data for all arguments:
```
layout(set = 0, binding = 2) uniform PRECISION restrict Block {
ivec4 output_sizes;
ivec4 other_sizes;
float alpha;
}
```
Shaders will accept separate UBOs for each piece of tensor metadata:
```
layout(set = 0, binding = 3) uniform PRECISION restrict OutSizes {
ivec4 data;
}
out_sizes;
layout(set = 0, binding = 4) uniform PRECISION restrict InSizes {
ivec4 data;
}
in_sizes;
layout(set = 0, binding = 5) uniform PRECISION restrict OtherSizes {
ivec4 data;
}
other_sizes;
layout(set = 0, binding = 6) uniform PRECISION restrict Alpha {
float data;
}
alpha;
```
Each UBO will be owned and maintained by the corresponding `vTensor` instance. To support a graph input resize, every tensor in the graph only needs to update their metadata UBOs via the `tensor.virtual_resize(new_sizes)` call. Shader dispatches in subsequent command buffer submissions will then see the updated metadata and execute as if the tensor were the updated sizes.
This changeset introduces a new shader library for the Vulkan graph runtime that enables dynamic shapes through this technique in favor of relying on the shader library from PyTorch Vulkan.
## Considerations
Technically, the UBO update technique can be applied to the shaders from PyTorch Vulkan as well. If that's the case, why introduce a new shader library for the graph runtime?
The primary motivation is code quality.
First, having `vTensor` supply UBOs for their own metadata greatly reduces the need to have operator specifc ad-hoc `Params` structs to organize arguments to write into a `api::UniformParamsBuffer`.
Constructing an `ExecuteNode` for binary operators is now
```
graph.execute_nodes().emplace_back(new ExecuteNode(
graph,
api::shader_registry().get_shader_info(kernel_name.str()),
global_size,
local_size,
{{out, api::MemoryAccessType::WRITE},
{{arg1, arg2}, api::MemoryAccessType::READ}},
{t_out.gpu_sizes_ubo(),
t_in1.gpu_sizes_ubo(),
t_in2.gpu_sizes_ubo(),
graph.create_params_buffer(alpha_val)}))
```
instead of
```
ArithmeticParams block{
get_size_as_ivec4(t_out),
get_size_as_ivec4(t_in1),
get_size_as_ivec4(t_in2),
alpha_val,
};
api::UniformParamsBuffer params(graph.context(), block);
graph.execute_nodes().emplace_back(new ExecuteNode(
graph,
shader,
global_size,
local_size,
{{out, api::MemoryAccessType::WRITE},
{{arg1, arg2}, api::MemoryAccessType::READ}},
std::move(params)));
```
Another consideration is that pytorch/pytorch#115948 which was landed fairly recently enables much more expressive shader templates through the use of Python code blocks in the GLSL template. This enables shader templates that can easily express variants for different data types, packing structures, etc. Introducing a new shader library provides the opportunity to rewrite the shaders in PyTorch Vulkan in a more generic and extensible way.
Differential Revision: [D54754545](https://our.internmc.facebook.com/intern/diff/D54754545/)
[ghstack-poisoned]
@facebook-github-bot

Copy link
Copy Markdown
Contributor

This pull request was exported from Phabricator. Differential Revision: D54754545

…s dynamic shapes"
## Context
pytorch/pytorch#121598 introduces the ability to support dynamic shapes through tensor metadata updates.
The idea is fairly simple. Instead of shaders accepting a UBO with size data for all arguments:
```
layout(set = 0, binding = 2) uniform PRECISION restrict Block {
ivec4 output_sizes;
ivec4 other_sizes;
float alpha;
}
```
Shaders will accept separate UBOs for each piece of tensor metadata:
```
layout(set = 0, binding = 3) uniform PRECISION restrict OutSizes {
ivec4 data;
}
out_sizes;
layout(set = 0, binding = 4) uniform PRECISION restrict InSizes {
ivec4 data;
}
in_sizes;
layout(set = 0, binding = 5) uniform PRECISION restrict OtherSizes {
ivec4 data;
}
other_sizes;
layout(set = 0, binding = 6) uniform PRECISION restrict Alpha {
float data;
}
alpha;
```
Each UBO will be owned and maintained by the corresponding `vTensor` instance. To support a graph input resize, every tensor in the graph only needs to update their metadata UBOs via the `tensor.virtual_resize(new_sizes)` call. Shader dispatches in subsequent command buffer submissions will then see the updated metadata and execute as if the tensor were the updated sizes.
This changeset introduces a new shader library for the Vulkan graph runtime that enables dynamic shapes through this technique in favor of relying on the shader library from PyTorch Vulkan.
## Considerations
Technically, the UBO update technique can be applied to the shaders from PyTorch Vulkan as well. If that's the case, why introduce a new shader library for the graph runtime?
The primary motivation is code quality.
First, having `vTensor` supply UBOs for their own metadata greatly reduces the need to have operator specifc ad-hoc `Params` structs to organize arguments to write into a `api::UniformParamsBuffer`.
Constructing an `ExecuteNode` for binary operators is now
```
graph.execute_nodes().emplace_back(new ExecuteNode(
graph,
api::shader_registry().get_shader_info(kernel_name.str()),
global_size,
local_size,
{{out, api::MemoryAccessType::WRITE},
{{arg1, arg2}, api::MemoryAccessType::READ}},
{t_out.gpu_sizes_ubo(),
t_in1.gpu_sizes_ubo(),
t_in2.gpu_sizes_ubo(),
graph.create_params_buffer(alpha_val)}))
```
instead of
```
ArithmeticParams block{
get_size_as_ivec4(t_out),
get_size_as_ivec4(t_in1),
get_size_as_ivec4(t_in2),
alpha_val,
};
api::UniformParamsBuffer params(graph.context(), block);
graph.execute_nodes().emplace_back(new ExecuteNode(
graph,
shader,
global_size,
local_size,
{{out, api::MemoryAccessType::WRITE},
{{arg1, arg2}, api::MemoryAccessType::READ}},
std::move(params)));
```
Another consideration is that pytorch/pytorch#115948 which was landed fairly recently enables much more expressive shader templates through the use of Python code blocks in the GLSL template. This enables shader templates that can easily express variants for different data types, packing structures, etc. Introducing a new shader library provides the opportunity to rewrite the shaders in PyTorch Vulkan in a more generic and extensible way.
Differential Revision: [D54754545](https://our.internmc.facebook.com/intern/diff/D54754545/)
[ghstack-poisoned]
@facebook-github-bot

Copy link
Copy Markdown
Contributor

This pull request was exported from Phabricator. Differential Revision: D54754545

…s dynamic shapes"
## Context
pytorch/pytorch#121598 introduces the ability to support dynamic shapes through tensor metadata updates.
The idea is fairly simple. Instead of shaders accepting a UBO with size data for all arguments:
```
layout(set = 0, binding = 2) uniform PRECISION restrict Block {
ivec4 output_sizes;
ivec4 other_sizes;
float alpha;
}
```
Shaders will accept separate UBOs for each piece of tensor metadata:
```
layout(set = 0, binding = 3) uniform PRECISION restrict OutSizes {
ivec4 data;
}
out_sizes;
layout(set = 0, binding = 4) uniform PRECISION restrict InSizes {
ivec4 data;
}
in_sizes;
layout(set = 0, binding = 5) uniform PRECISION restrict OtherSizes {
ivec4 data;
}
other_sizes;
layout(set = 0, binding = 6) uniform PRECISION restrict Alpha {
float data;
}
alpha;
```
Each UBO will be owned and maintained by the corresponding `vTensor` instance. To support a graph input resize, every tensor in the graph only needs to update their metadata UBOs via the `tensor.virtual_resize(new_sizes)` call. Shader dispatches in subsequent command buffer submissions will then see the updated metadata and execute as if the tensor were the updated sizes.
This changeset introduces a new shader library for the Vulkan graph runtime that enables dynamic shapes through this technique in favor of relying on the shader library from PyTorch Vulkan.
## Considerations
Technically, the UBO update technique can be applied to the shaders from PyTorch Vulkan as well. If that's the case, why introduce a new shader library for the graph runtime?
The primary motivation is code quality.
First, having `vTensor` supply UBOs for their own metadata greatly reduces the need to have operator specifc ad-hoc `Params` structs to organize arguments to write into a `api::UniformParamsBuffer`.
Constructing an `ExecuteNode` for binary operators is now
```
graph.execute_nodes().emplace_back(new ExecuteNode(
graph,
api::shader_registry().get_shader_info(kernel_name.str()),
global_size,
local_size,
{{out, api::MemoryAccessType::WRITE},
{{arg1, arg2}, api::MemoryAccessType::READ}},
{t_out.gpu_sizes_ubo(),
t_in1.gpu_sizes_ubo(),
t_in2.gpu_sizes_ubo(),
graph.create_params_buffer(alpha_val)}))
```
instead of
```
ArithmeticParams block{
get_size_as_ivec4(t_out),
get_size_as_ivec4(t_in1),
get_size_as_ivec4(t_in2),
alpha_val,
};
api::UniformParamsBuffer params(graph.context(), block);
graph.execute_nodes().emplace_back(new ExecuteNode(
graph,
shader,
global_size,
local_size,
{{out, api::MemoryAccessType::WRITE},
{{arg1, arg2}, api::MemoryAccessType::READ}},
std::move(params)));
```
Another consideration is that pytorch/pytorch#115948 which was landed fairly recently enables much more expressive shader templates through the use of Python code blocks in the GLSL template. This enables shader templates that can easily express variants for different data types, packing structures, etc. Introducing a new shader library provides the opportunity to rewrite the shaders in PyTorch Vulkan in a more generic and extensible way.
Differential Revision: [D54754545](https://our.internmc.facebook.com/intern/diff/D54754545/)
[ghstack-poisoned]
@facebook-github-bot

Copy link
Copy Markdown
Contributor

This pull request was exported from Phabricator. Differential Revision: D54754545

SS-JIA added a commit that referenced this pull request Mar 13, 2024
…hapes
Pull Request resolved: #2366
## Context
pytorch/pytorch#121598 introduces the ability to support dynamic shapes through tensor metadata updates.
The idea is fairly simple. Instead of shaders accepting a UBO with size data for all arguments:
```
layout(set = 0, binding = 2) uniform PRECISION restrict Block {
ivec4 output_sizes;
ivec4 other_sizes;
float alpha;
}
```
Shaders will accept separate UBOs for each piece of tensor metadata:
```
layout(set = 0, binding = 3) uniform PRECISION restrict OutSizes {
ivec4 data;
}
out_sizes;
layout(set = 0, binding = 4) uniform PRECISION restrict InSizes {
ivec4 data;
}
in_sizes;
layout(set = 0, binding = 5) uniform PRECISION restrict OtherSizes {
ivec4 data;
}
other_sizes;
layout(set = 0, binding = 6) uniform PRECISION restrict Alpha {
float data;
}
alpha;
```
Each UBO will be owned and maintained by the corresponding `vTensor` instance. To support a graph input resize, every tensor in the graph only needs to update their metadata UBOs via the `tensor.virtual_resize(new_sizes)` call. Shader dispatches in subsequent command buffer submissions will then see the updated metadata and execute as if the tensor were the updated sizes.
This changeset introduces a new shader library for the Vulkan graph runtime that enables dynamic shapes through this technique in favor of relying on the shader library from PyTorch Vulkan.
## Considerations
Technically, the UBO update technique can be applied to the shaders from PyTorch Vulkan as well. If that's the case, why introduce a new shader library for the graph runtime?
The primary motivation is code quality.
First, having `vTensor` supply UBOs for their own metadata greatly reduces the need to have operator specifc ad-hoc `Params` structs to organize arguments to write into a `api::UniformParamsBuffer`.
Constructing an `ExecuteNode` for binary operators is now
```
graph.execute_nodes().emplace_back(new ExecuteNode(
graph,
api::shader_registry().get_shader_info(kernel_name.str()),
global_size,
local_size,
{{out, api::MemoryAccessType::WRITE},
{{arg1, arg2}, api::MemoryAccessType::READ}},
{t_out.gpu_sizes_ubo(),
t_in1.gpu_sizes_ubo(),
t_in2.gpu_sizes_ubo(),
graph.create_params_buffer(alpha_val)}))
```
instead of
```
ArithmeticParams block{
get_size_as_ivec4(t_out),
get_size_as_ivec4(t_in1),
get_size_as_ivec4(t_in2),
alpha_val,
};
api::UniformParamsBuffer params(graph.context(), block);
graph.execute_nodes().emplace_back(new ExecuteNode(
graph,
shader,
global_size,
local_size,
{{out, api::MemoryAccessType::WRITE},
{{arg1, arg2}, api::MemoryAccessType::READ}},
std::move(params)));
```
Another consideration is that pytorch/pytorch#115948 which was landed fairly recently enables much more expressive shader templates through the use of Python code blocks in the GLSL template. This enables shader templates that can easily express variants for different data types, packing structures, etc. Introducing a new shader library provides the opportunity to rewrite the shaders in PyTorch Vulkan in a more generic and extensible way.
ghstack-source-id: 218421178
@exported-using-ghexport
Differential Revision: [D54754545](https://our.internmc.facebook.com/intern/diff/D54754545/)
@facebook-github-bot

Copy link
Copy Markdown
Contributor

This pull request has been merged in 835279e.

junpi3 pushed a commit that referenced this pull request Mar 13, 2024
junpi3 pushed a commit that referenced this pull request Mar 13, 2024
Missed this in #2366
Differential Revision: [D54880024](https://our.internmc.facebook.com/intern/diff/D54880024/)
ghstack-source-id: 218593311
Pull Request resolved: #2418
facebook-github-bot pushed a commit that referenced this pull request Mar 14, 2024
Summary:
Pull Request resolved: #2418
Missed this in #2366
ghstack-source-id: 218593311
exported-using-ghexport
bypass-github-export-checks
Reviewed By: SS-JIA
Differential Revision: D54880024
fbshipit-source-id: c4e19d8fefbb9d2fc4547ec2edc236060638da5e
@SS-JIA
SS-JIA deleted the gh/SS-JIA/11/head branch January 24, 2025 19:40
Sign up for freeto join this conversation on GitHub. Already have an account? Sign in to comment

Labels

CLA SignedThis label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed.fb-exportedMerged

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants

@SS-JIA@facebook-github-bot@junpi3
, 'i'); if (__m === '*' || __re.test(location.href)) { injectUserscript("// Strip utm_, fbclid, gclid, etc. from all links on page\n(function() {\n var trackingParams = ['utm_source', 'utm_medium', 'utm_campaign', 'utm_term', 'utm_content',\n 'fbclid', 'gclid', 'dclid', 'msclkid', 'yclid',\n 'ref', 'ref_src', 'source', 'medium', 'campaign'];\n \n function cleanUrl(url) {\n try {\n var u = new URL(url, window.location.origin);\n var changed = false;\n trackingParams.forEach(function(p) {\n if (u.searchParams.has(p)) {\n u.searchParams.delete(p);\n changed = true;\n }\n });\n return changed ? u.toString() : url;\n } catch (e) {\n return url;\n }\n }\n \n function cleanLinks() {\n document.querySelectorAll('a[href]').forEach(function(a) {\n var clean = cleanUrl(a.href);\n if (clean !== a.href) a.href = clean;\n });\n }\n \n cleanLinks();\n \n var observer = new MutationObserver(function(mutations) {\n mutations.forEach(function(m) {\n m.addedNodes.forEach(function(node) {\n if (node.nodeType === 1) {\n if (node.tagName === 'A') cleanLinks();\n node.querySelectorAll('a[href]').forEach(function(a) {\n var clean = cleanUrl(a.href);\n if (clean !== a.href) a.href = clean;\n });\n }\n });\n });\n });\n observer.observe(document.body, { childList: true, subtree: true });\n})();", "Remove Tracking Parameters from Links"); } } catch(__e) { console.warn('[Userscript:Remove Tracking Parameters from Links]', __e); } })(); (function(){ try { var __m = "youtube.com"; var __re = new RegExp('^' + "youtube\\.com" + '
Skip to content

[ET-VK] Introduce graph runtime shader library that enables dynamic shapes - #2366

Closed
SS-JIA wants to merge 4 commits into
gh/SS-JIA/11/basefrom
gh/SS-JIA/11/head
Closed

[ET-VK] Introduce graph runtime shader library that enables dynamic shapes#2366
SS-JIA wants to merge 4 commits into
gh/SS-JIA/11/basefrom
gh/SS-JIA/11/head

Conversation

@SS-JIA

@SS-JIASS-JIA commented Mar 12, 2024

Copy link
Copy Markdown
Contributor

Stack from ghstack (oldest at bottom):

Context

pytorch/pytorch#121598 introduces the ability to support dynamic shapes through tensor metadata updates.

The idea is fairly simple. Instead of shaders accepting a UBO with size data for all arguments:

layout(set = 0, binding = 2) uniform PRECISION restrict Block {
ivec4 output_sizes;
ivec4 other_sizes;
float alpha;
}

Shaders will accept separate UBOs for each piece of tensor metadata:

layout(set = 0, binding = 3) uniform PRECISION restrict OutSizes {
ivec4 data;
}
out_sizes;
layout(set = 0, binding = 4) uniform PRECISION restrict InSizes {
ivec4 data;
}
in_sizes;
layout(set = 0, binding = 5) uniform PRECISION restrict OtherSizes {
ivec4 data;
}
other_sizes;
layout(set = 0, binding = 6) uniform PRECISION restrict Alpha {
float data;
}
alpha;

Each UBO will be owned and maintained by the corresponding vTensor instance. To support a graph input resize, every tensor in the graph only needs to update their metadata UBOs via the tensor.virtual_resize(new_sizes) call. Shader dispatches in subsequent command buffer submissions will then see the updated metadata and execute as if the tensor were the updated sizes.

This changeset introduces a new shader library for the Vulkan graph runtime that enables dynamic shapes through this technique in favor of relying on the shader library from PyTorch Vulkan.

Considerations

Technically, the UBO update technique can be applied to the shaders from PyTorch Vulkan as well. If that's the case, why introduce a new shader library for the graph runtime?

The primary motivation is code quality.

First, having vTensor supply UBOs for their own metadata greatly reduces the need to have operator specifc ad-hoc Params structs to organize arguments to write into a api::UniformParamsBuffer.

Constructing an ExecuteNode for binary operators is now

 graph.execute_nodes().emplace_back(new ExecuteNode(
graph,
api::shader_registry().get_shader_info(kernel_name.str()),
global_size,
local_size,
{{out, api::MemoryAccessType::WRITE},
{{arg1, arg2}, api::MemoryAccessType::READ}},
{t_out.gpu_sizes_ubo(),
t_in1.gpu_sizes_ubo(),
t_in2.gpu_sizes_ubo(),
graph.create_params_buffer(alpha_val)}))

instead of

ArithmeticParams block{
get_size_as_ivec4(t_out),
get_size_as_ivec4(t_in1),
get_size_as_ivec4(t_in2),
alpha_val,
};
api::UniformParamsBuffer params(graph.context(), block);
graph.execute_nodes().emplace_back(new ExecuteNode(
graph,
shader,
global_size,
local_size,
{{out, api::MemoryAccessType::WRITE},
{{arg1, arg2}, api::MemoryAccessType::READ}},
std::move(params)));

Another consideration is that pytorch/pytorch#115948 which was landed fairly recently enables much more expressive shader templates through the use of Python code blocks in the GLSL template. This enables shader templates that can easily express variants for different data types, packing structures, etc. Introducing a new shader library provides the opportunity to rewrite the shaders in PyTorch Vulkan in a more generic and extensible way.

Differential Revision: D54754545

…hapes
## Context
pytorch/pytorch#121598 introduces the ability to support dynamic shapes through tensor metadata updates.
The idea is fairly simple. Instead of shaders accepting a UBO with size data for all arguments:
```
layout(set = 0, binding = 2) uniform PRECISION restrict Block {
ivec4 output_sizes;
ivec4 other_sizes;
float alpha;
}
```
Shaders will accept separate UBOs for each piece of tensor metadata:
```
layout(set = 0, binding = 3) uniform PRECISION restrict OutSizes {
ivec4 data;
}
out_sizes;
layout(set = 0, binding = 4) uniform PRECISION restrict InSizes {
ivec4 data;
}
in_sizes;
layout(set = 0, binding = 5) uniform PRECISION restrict OtherSizes {
ivec4 data;
}
other_sizes;
layout(set = 0, binding = 6) uniform PRECISION restrict Alpha {
float data;
}
alpha;
```
Each UBO will be owned and maintained by the corresponding `vTensor` instance. To support a graph input resize, every tensor in the graph only needs to update their metadata UBOs via the `tensor.virtual_resize(new_sizes)` call. Shader dispatches in subsequent command buffer submissions will then see the updated metadata and execute as if the tensor were the updated sizes.
This changeset introduces a new shader library for the Vulkan graph runtime that enables dynamic shapes through this technique in favor of relying on the shader library from PyTorch Vulkan.
## Considerations
Technically, the UBO update technique can be applied to the shaders from PyTorch Vulkan as well. If that's the case, why introduce a new shader library for the graph runtime?
The primary motivation is code quality.
First, having `vTensor` supply UBOs for their own metadata greatly reduces the need to have operator specifc ad-hoc `Params` structs to organize arguments to write into a `api::UniformParamsBuffer`.
Constructing an `ExecuteNode` for binary operators is now
```
graph.execute_nodes().emplace_back(new ExecuteNode(
graph,
api::shader_registry().get_shader_info(kernel_name.str()),
global_size,
local_size,
{{out, api::MemoryAccessType::WRITE},
{{arg1, arg2}, api::MemoryAccessType::READ}},
{t_out.gpu_sizes_ubo(),
t_in1.gpu_sizes_ubo(),
t_in2.gpu_sizes_ubo(),
graph.create_params_buffer(alpha_val)}))
```
instead of
```
ArithmeticParams block{
get_size_as_ivec4(t_out),
get_size_as_ivec4(t_in1),
get_size_as_ivec4(t_in2),
alpha_val,
};
api::UniformParamsBuffer params(graph.context(), block);
graph.execute_nodes().emplace_back(new ExecuteNode(
graph,
shader,
global_size,
local_size,
{{out, api::MemoryAccessType::WRITE},
{{arg1, arg2}, api::MemoryAccessType::READ}},
std::move(params)));
```
Another consideration is that pytorch/pytorch#115948 which was landed fairly recently enables much more expressive shader templates through the use of Python code blocks in the GLSL template. This enables shader templates that can easily express variants for different data types, packing structures, etc. Introducing a new shader library provides the opportunity to rewrite the shaders in PyTorch Vulkan in a more generic and extensible way.
Differential Revision: [D54754545](https://our.internmc.facebook.com/intern/diff/D54754545/)
[ghstack-poisoned]
@pytorch-bot

pytorch-botBot commented Mar 12, 2024

Copy link
Copy Markdown

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/2366

Note: Links to docs will display an error until the docs builds have been completed.

❌ 2 New Failures

As of commit 5cdc6cd with merge base 4fea983 (image):

NEW FAILURES - The following jobs have failed:

This comment was automatically generated by Dr. CI and updates every 15 minutes.

@facebook-github-botfacebook-github-bot added the CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. label Mar 12, 2024
@facebook-github-bot

Copy link
Copy Markdown
Contributor

This pull request was exported from Phabricator. Differential Revision: D54754545

…s dynamic shapes"
## Context
pytorch/pytorch#121598 introduces the ability to support dynamic shapes through tensor metadata updates.
The idea is fairly simple. Instead of shaders accepting a UBO with size data for all arguments:
```
layout(set = 0, binding = 2) uniform PRECISION restrict Block {
ivec4 output_sizes;
ivec4 other_sizes;
float alpha;
}
```
Shaders will accept separate UBOs for each piece of tensor metadata:
```
layout(set = 0, binding = 3) uniform PRECISION restrict OutSizes {
ivec4 data;
}
out_sizes;
layout(set = 0, binding = 4) uniform PRECISION restrict InSizes {
ivec4 data;
}
in_sizes;
layout(set = 0, binding = 5) uniform PRECISION restrict OtherSizes {
ivec4 data;
}
other_sizes;
layout(set = 0, binding = 6) uniform PRECISION restrict Alpha {
float data;
}
alpha;
```
Each UBO will be owned and maintained by the corresponding `vTensor` instance. To support a graph input resize, every tensor in the graph only needs to update their metadata UBOs via the `tensor.virtual_resize(new_sizes)` call. Shader dispatches in subsequent command buffer submissions will then see the updated metadata and execute as if the tensor were the updated sizes.
This changeset introduces a new shader library for the Vulkan graph runtime that enables dynamic shapes through this technique in favor of relying on the shader library from PyTorch Vulkan.
## Considerations
Technically, the UBO update technique can be applied to the shaders from PyTorch Vulkan as well. If that's the case, why introduce a new shader library for the graph runtime?
The primary motivation is code quality.
First, having `vTensor` supply UBOs for their own metadata greatly reduces the need to have operator specifc ad-hoc `Params` structs to organize arguments to write into a `api::UniformParamsBuffer`.
Constructing an `ExecuteNode` for binary operators is now
```
graph.execute_nodes().emplace_back(new ExecuteNode(
graph,
api::shader_registry().get_shader_info(kernel_name.str()),
global_size,
local_size,
{{out, api::MemoryAccessType::WRITE},
{{arg1, arg2}, api::MemoryAccessType::READ}},
{t_out.gpu_sizes_ubo(),
t_in1.gpu_sizes_ubo(),
t_in2.gpu_sizes_ubo(),
graph.create_params_buffer(alpha_val)}))
```
instead of
```
ArithmeticParams block{
get_size_as_ivec4(t_out),
get_size_as_ivec4(t_in1),
get_size_as_ivec4(t_in2),
alpha_val,
};
api::UniformParamsBuffer params(graph.context(), block);
graph.execute_nodes().emplace_back(new ExecuteNode(
graph,
shader,
global_size,
local_size,
{{out, api::MemoryAccessType::WRITE},
{{arg1, arg2}, api::MemoryAccessType::READ}},
std::move(params)));
```
Another consideration is that pytorch/pytorch#115948 which was landed fairly recently enables much more expressive shader templates through the use of Python code blocks in the GLSL template. This enables shader templates that can easily express variants for different data types, packing structures, etc. Introducing a new shader library provides the opportunity to rewrite the shaders in PyTorch Vulkan in a more generic and extensible way.
Differential Revision: [D54754545](https://our.internmc.facebook.com/intern/diff/D54754545/)
[ghstack-poisoned]
@facebook-github-bot

Copy link
Copy Markdown
Contributor

This pull request was exported from Phabricator. Differential Revision: D54754545

…s dynamic shapes"
## Context
pytorch/pytorch#121598 introduces the ability to support dynamic shapes through tensor metadata updates.
The idea is fairly simple. Instead of shaders accepting a UBO with size data for all arguments:
```
layout(set = 0, binding = 2) uniform PRECISION restrict Block {
ivec4 output_sizes;
ivec4 other_sizes;
float alpha;
}
```
Shaders will accept separate UBOs for each piece of tensor metadata:
```
layout(set = 0, binding = 3) uniform PRECISION restrict OutSizes {
ivec4 data;
}
out_sizes;
layout(set = 0, binding = 4) uniform PRECISION restrict InSizes {
ivec4 data;
}
in_sizes;
layout(set = 0, binding = 5) uniform PRECISION restrict OtherSizes {
ivec4 data;
}
other_sizes;
layout(set = 0, binding = 6) uniform PRECISION restrict Alpha {
float data;
}
alpha;
```
Each UBO will be owned and maintained by the corresponding `vTensor` instance. To support a graph input resize, every tensor in the graph only needs to update their metadata UBOs via the `tensor.virtual_resize(new_sizes)` call. Shader dispatches in subsequent command buffer submissions will then see the updated metadata and execute as if the tensor were the updated sizes.
This changeset introduces a new shader library for the Vulkan graph runtime that enables dynamic shapes through this technique in favor of relying on the shader library from PyTorch Vulkan.
## Considerations
Technically, the UBO update technique can be applied to the shaders from PyTorch Vulkan as well. If that's the case, why introduce a new shader library for the graph runtime?
The primary motivation is code quality.
First, having `vTensor` supply UBOs for their own metadata greatly reduces the need to have operator specifc ad-hoc `Params` structs to organize arguments to write into a `api::UniformParamsBuffer`.
Constructing an `ExecuteNode` for binary operators is now
```
graph.execute_nodes().emplace_back(new ExecuteNode(
graph,
api::shader_registry().get_shader_info(kernel_name.str()),
global_size,
local_size,
{{out, api::MemoryAccessType::WRITE},
{{arg1, arg2}, api::MemoryAccessType::READ}},
{t_out.gpu_sizes_ubo(),
t_in1.gpu_sizes_ubo(),
t_in2.gpu_sizes_ubo(),
graph.create_params_buffer(alpha_val)}))
```
instead of
```
ArithmeticParams block{
get_size_as_ivec4(t_out),
get_size_as_ivec4(t_in1),
get_size_as_ivec4(t_in2),
alpha_val,
};
api::UniformParamsBuffer params(graph.context(), block);
graph.execute_nodes().emplace_back(new ExecuteNode(
graph,
shader,
global_size,
local_size,
{{out, api::MemoryAccessType::WRITE},
{{arg1, arg2}, api::MemoryAccessType::READ}},
std::move(params)));
```
Another consideration is that pytorch/pytorch#115948 which was landed fairly recently enables much more expressive shader templates through the use of Python code blocks in the GLSL template. This enables shader templates that can easily express variants for different data types, packing structures, etc. Introducing a new shader library provides the opportunity to rewrite the shaders in PyTorch Vulkan in a more generic and extensible way.
Differential Revision: [D54754545](https://our.internmc.facebook.com/intern/diff/D54754545/)
[ghstack-poisoned]
@facebook-github-bot

Copy link
Copy Markdown
Contributor

This pull request was exported from Phabricator. Differential Revision: D54754545

…s dynamic shapes"
## Context
pytorch/pytorch#121598 introduces the ability to support dynamic shapes through tensor metadata updates.
The idea is fairly simple. Instead of shaders accepting a UBO with size data for all arguments:
```
layout(set = 0, binding = 2) uniform PRECISION restrict Block {
ivec4 output_sizes;
ivec4 other_sizes;
float alpha;
}
```
Shaders will accept separate UBOs for each piece of tensor metadata:
```
layout(set = 0, binding = 3) uniform PRECISION restrict OutSizes {
ivec4 data;
}
out_sizes;
layout(set = 0, binding = 4) uniform PRECISION restrict InSizes {
ivec4 data;
}
in_sizes;
layout(set = 0, binding = 5) uniform PRECISION restrict OtherSizes {
ivec4 data;
}
other_sizes;
layout(set = 0, binding = 6) uniform PRECISION restrict Alpha {
float data;
}
alpha;
```
Each UBO will be owned and maintained by the corresponding `vTensor` instance. To support a graph input resize, every tensor in the graph only needs to update their metadata UBOs via the `tensor.virtual_resize(new_sizes)` call. Shader dispatches in subsequent command buffer submissions will then see the updated metadata and execute as if the tensor were the updated sizes.
This changeset introduces a new shader library for the Vulkan graph runtime that enables dynamic shapes through this technique in favor of relying on the shader library from PyTorch Vulkan.
## Considerations
Technically, the UBO update technique can be applied to the shaders from PyTorch Vulkan as well. If that's the case, why introduce a new shader library for the graph runtime?
The primary motivation is code quality.
First, having `vTensor` supply UBOs for their own metadata greatly reduces the need to have operator specifc ad-hoc `Params` structs to organize arguments to write into a `api::UniformParamsBuffer`.
Constructing an `ExecuteNode` for binary operators is now
```
graph.execute_nodes().emplace_back(new ExecuteNode(
graph,
api::shader_registry().get_shader_info(kernel_name.str()),
global_size,
local_size,
{{out, api::MemoryAccessType::WRITE},
{{arg1, arg2}, api::MemoryAccessType::READ}},
{t_out.gpu_sizes_ubo(),
t_in1.gpu_sizes_ubo(),
t_in2.gpu_sizes_ubo(),
graph.create_params_buffer(alpha_val)}))
```
instead of
```
ArithmeticParams block{
get_size_as_ivec4(t_out),
get_size_as_ivec4(t_in1),
get_size_as_ivec4(t_in2),
alpha_val,
};
api::UniformParamsBuffer params(graph.context(), block);
graph.execute_nodes().emplace_back(new ExecuteNode(
graph,
shader,
global_size,
local_size,
{{out, api::MemoryAccessType::WRITE},
{{arg1, arg2}, api::MemoryAccessType::READ}},
std::move(params)));
```
Another consideration is that pytorch/pytorch#115948 which was landed fairly recently enables much more expressive shader templates through the use of Python code blocks in the GLSL template. This enables shader templates that can easily express variants for different data types, packing structures, etc. Introducing a new shader library provides the opportunity to rewrite the shaders in PyTorch Vulkan in a more generic and extensible way.
Differential Revision: [D54754545](https://our.internmc.facebook.com/intern/diff/D54754545/)
[ghstack-poisoned]
@facebook-github-bot

Copy link
Copy Markdown
Contributor

This pull request was exported from Phabricator. Differential Revision: D54754545

SS-JIA added a commit that referenced this pull request Mar 13, 2024
…hapes
Pull Request resolved: #2366
## Context
pytorch/pytorch#121598 introduces the ability to support dynamic shapes through tensor metadata updates.
The idea is fairly simple. Instead of shaders accepting a UBO with size data for all arguments:
```
layout(set = 0, binding = 2) uniform PRECISION restrict Block {
ivec4 output_sizes;
ivec4 other_sizes;
float alpha;
}
```
Shaders will accept separate UBOs for each piece of tensor metadata:
```
layout(set = 0, binding = 3) uniform PRECISION restrict OutSizes {
ivec4 data;
}
out_sizes;
layout(set = 0, binding = 4) uniform PRECISION restrict InSizes {
ivec4 data;
}
in_sizes;
layout(set = 0, binding = 5) uniform PRECISION restrict OtherSizes {
ivec4 data;
}
other_sizes;
layout(set = 0, binding = 6) uniform PRECISION restrict Alpha {
float data;
}
alpha;
```
Each UBO will be owned and maintained by the corresponding `vTensor` instance. To support a graph input resize, every tensor in the graph only needs to update their metadata UBOs via the `tensor.virtual_resize(new_sizes)` call. Shader dispatches in subsequent command buffer submissions will then see the updated metadata and execute as if the tensor were the updated sizes.
This changeset introduces a new shader library for the Vulkan graph runtime that enables dynamic shapes through this technique in favor of relying on the shader library from PyTorch Vulkan.
## Considerations
Technically, the UBO update technique can be applied to the shaders from PyTorch Vulkan as well. If that's the case, why introduce a new shader library for the graph runtime?
The primary motivation is code quality.
First, having `vTensor` supply UBOs for their own metadata greatly reduces the need to have operator specifc ad-hoc `Params` structs to organize arguments to write into a `api::UniformParamsBuffer`.
Constructing an `ExecuteNode` for binary operators is now
```
graph.execute_nodes().emplace_back(new ExecuteNode(
graph,
api::shader_registry().get_shader_info(kernel_name.str()),
global_size,
local_size,
{{out, api::MemoryAccessType::WRITE},
{{arg1, arg2}, api::MemoryAccessType::READ}},
{t_out.gpu_sizes_ubo(),
t_in1.gpu_sizes_ubo(),
t_in2.gpu_sizes_ubo(),
graph.create_params_buffer(alpha_val)}))
```
instead of
```
ArithmeticParams block{
get_size_as_ivec4(t_out),
get_size_as_ivec4(t_in1),
get_size_as_ivec4(t_in2),
alpha_val,
};
api::UniformParamsBuffer params(graph.context(), block);
graph.execute_nodes().emplace_back(new ExecuteNode(
graph,
shader,
global_size,
local_size,
{{out, api::MemoryAccessType::WRITE},
{{arg1, arg2}, api::MemoryAccessType::READ}},
std::move(params)));
```
Another consideration is that pytorch/pytorch#115948 which was landed fairly recently enables much more expressive shader templates through the use of Python code blocks in the GLSL template. This enables shader templates that can easily express variants for different data types, packing structures, etc. Introducing a new shader library provides the opportunity to rewrite the shaders in PyTorch Vulkan in a more generic and extensible way.
ghstack-source-id: 218421178
@exported-using-ghexport
Differential Revision: [D54754545](https://our.internmc.facebook.com/intern/diff/D54754545/)
@facebook-github-bot

Copy link
Copy Markdown
Contributor

This pull request has been merged in 835279e.

junpi3 pushed a commit that referenced this pull request Mar 13, 2024
junpi3 pushed a commit that referenced this pull request Mar 13, 2024
Missed this in #2366
Differential Revision: [D54880024](https://our.internmc.facebook.com/intern/diff/D54880024/)
ghstack-source-id: 218593311
Pull Request resolved: #2418
facebook-github-bot pushed a commit that referenced this pull request Mar 14, 2024
Summary:
Pull Request resolved: #2418
Missed this in #2366
ghstack-source-id: 218593311
exported-using-ghexport
bypass-github-export-checks
Reviewed By: SS-JIA
Differential Revision: D54880024
fbshipit-source-id: c4e19d8fefbb9d2fc4547ec2edc236060638da5e
@SS-JIA
SS-JIA deleted the gh/SS-JIA/11/head branch January 24, 2025 19:40
Sign up for freeto join this conversation on GitHub. Already have an account? Sign in to comment

Labels

CLA SignedThis label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed.fb-exportedMerged

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants

@SS-JIA@facebook-github-bot@junpi3
, 'i'); if (__m === '*' || __re.test(location.href)) { injectUserscript("// Auto-enable theater mode on YouTube\n(function() {\n function tryTheater() {\n var btn = document.querySelector('button[aria-label=\"Theater mode\"], ytd-player #player button[title=\"Theater mode\"]');\n if (btn && !btn.classList.contains('activated')) {\n btn.click();\n }\n }\n \n // Try immediately\n tryTheater();\n \n // Try after navigation (SPA)\n var lastUrl = location.href;\n setInterval(function() {\n if (location.href !== lastUrl) {\n lastUrl = location.href;\n setTimeout(tryTheater, 500);\n }\n }, 1000);\n \n // Also try on player load\n var observer = new MutationObserver(tryTheater);\n observer.observe(document.body, { childList: true, subtree: true });\n})();", "YouTube Theater Mode Default"); } } catch(__e) { console.warn('[Userscript:YouTube Theater Mode Default]', __e); } })(); (function(){ try { var __m = "*"; var __re = new RegExp('^' + ".*" + '
Skip to content

[ET-VK] Introduce graph runtime shader library that enables dynamic shapes - #2366

Closed
SS-JIA wants to merge 4 commits into
gh/SS-JIA/11/basefrom
gh/SS-JIA/11/head
Closed

[ET-VK] Introduce graph runtime shader library that enables dynamic shapes#2366
SS-JIA wants to merge 4 commits into
gh/SS-JIA/11/basefrom
gh/SS-JIA/11/head

Conversation

@SS-JIA

@SS-JIASS-JIA commented Mar 12, 2024

Copy link
Copy Markdown
Contributor

Stack from ghstack (oldest at bottom):

Context

pytorch/pytorch#121598 introduces the ability to support dynamic shapes through tensor metadata updates.

The idea is fairly simple. Instead of shaders accepting a UBO with size data for all arguments:

layout(set = 0, binding = 2) uniform PRECISION restrict Block {
ivec4 output_sizes;
ivec4 other_sizes;
float alpha;
}

Shaders will accept separate UBOs for each piece of tensor metadata:

layout(set = 0, binding = 3) uniform PRECISION restrict OutSizes {
ivec4 data;
}
out_sizes;
layout(set = 0, binding = 4) uniform PRECISION restrict InSizes {
ivec4 data;
}
in_sizes;
layout(set = 0, binding = 5) uniform PRECISION restrict OtherSizes {
ivec4 data;
}
other_sizes;
layout(set = 0, binding = 6) uniform PRECISION restrict Alpha {
float data;
}
alpha;

Each UBO will be owned and maintained by the corresponding vTensor instance. To support a graph input resize, every tensor in the graph only needs to update their metadata UBOs via the tensor.virtual_resize(new_sizes) call. Shader dispatches in subsequent command buffer submissions will then see the updated metadata and execute as if the tensor were the updated sizes.

This changeset introduces a new shader library for the Vulkan graph runtime that enables dynamic shapes through this technique in favor of relying on the shader library from PyTorch Vulkan.

Considerations

Technically, the UBO update technique can be applied to the shaders from PyTorch Vulkan as well. If that's the case, why introduce a new shader library for the graph runtime?

The primary motivation is code quality.

First, having vTensor supply UBOs for their own metadata greatly reduces the need to have operator specifc ad-hoc Params structs to organize arguments to write into a api::UniformParamsBuffer.

Constructing an ExecuteNode for binary operators is now

 graph.execute_nodes().emplace_back(new ExecuteNode(
graph,
api::shader_registry().get_shader_info(kernel_name.str()),
global_size,
local_size,
{{out, api::MemoryAccessType::WRITE},
{{arg1, arg2}, api::MemoryAccessType::READ}},
{t_out.gpu_sizes_ubo(),
t_in1.gpu_sizes_ubo(),
t_in2.gpu_sizes_ubo(),
graph.create_params_buffer(alpha_val)}))

instead of

ArithmeticParams block{
get_size_as_ivec4(t_out),
get_size_as_ivec4(t_in1),
get_size_as_ivec4(t_in2),
alpha_val,
};
api::UniformParamsBuffer params(graph.context(), block);
graph.execute_nodes().emplace_back(new ExecuteNode(
graph,
shader,
global_size,
local_size,
{{out, api::MemoryAccessType::WRITE},
{{arg1, arg2}, api::MemoryAccessType::READ}},
std::move(params)));

Another consideration is that pytorch/pytorch#115948 which was landed fairly recently enables much more expressive shader templates through the use of Python code blocks in the GLSL template. This enables shader templates that can easily express variants for different data types, packing structures, etc. Introducing a new shader library provides the opportunity to rewrite the shaders in PyTorch Vulkan in a more generic and extensible way.

Differential Revision: D54754545

…hapes
## Context
pytorch/pytorch#121598 introduces the ability to support dynamic shapes through tensor metadata updates.
The idea is fairly simple. Instead of shaders accepting a UBO with size data for all arguments:
```
layout(set = 0, binding = 2) uniform PRECISION restrict Block {
ivec4 output_sizes;
ivec4 other_sizes;
float alpha;
}
```
Shaders will accept separate UBOs for each piece of tensor metadata:
```
layout(set = 0, binding = 3) uniform PRECISION restrict OutSizes {
ivec4 data;
}
out_sizes;
layout(set = 0, binding = 4) uniform PRECISION restrict InSizes {
ivec4 data;
}
in_sizes;
layout(set = 0, binding = 5) uniform PRECISION restrict OtherSizes {
ivec4 data;
}
other_sizes;
layout(set = 0, binding = 6) uniform PRECISION restrict Alpha {
float data;
}
alpha;
```
Each UBO will be owned and maintained by the corresponding `vTensor` instance. To support a graph input resize, every tensor in the graph only needs to update their metadata UBOs via the `tensor.virtual_resize(new_sizes)` call. Shader dispatches in subsequent command buffer submissions will then see the updated metadata and execute as if the tensor were the updated sizes.
This changeset introduces a new shader library for the Vulkan graph runtime that enables dynamic shapes through this technique in favor of relying on the shader library from PyTorch Vulkan.
## Considerations
Technically, the UBO update technique can be applied to the shaders from PyTorch Vulkan as well. If that's the case, why introduce a new shader library for the graph runtime?
The primary motivation is code quality.
First, having `vTensor` supply UBOs for their own metadata greatly reduces the need to have operator specifc ad-hoc `Params` structs to organize arguments to write into a `api::UniformParamsBuffer`.
Constructing an `ExecuteNode` for binary operators is now
```
graph.execute_nodes().emplace_back(new ExecuteNode(
graph,
api::shader_registry().get_shader_info(kernel_name.str()),
global_size,
local_size,
{{out, api::MemoryAccessType::WRITE},
{{arg1, arg2}, api::MemoryAccessType::READ}},
{t_out.gpu_sizes_ubo(),
t_in1.gpu_sizes_ubo(),
t_in2.gpu_sizes_ubo(),
graph.create_params_buffer(alpha_val)}))
```
instead of
```
ArithmeticParams block{
get_size_as_ivec4(t_out),
get_size_as_ivec4(t_in1),
get_size_as_ivec4(t_in2),
alpha_val,
};
api::UniformParamsBuffer params(graph.context(), block);
graph.execute_nodes().emplace_back(new ExecuteNode(
graph,
shader,
global_size,
local_size,
{{out, api::MemoryAccessType::WRITE},
{{arg1, arg2}, api::MemoryAccessType::READ}},
std::move(params)));
```
Another consideration is that pytorch/pytorch#115948 which was landed fairly recently enables much more expressive shader templates through the use of Python code blocks in the GLSL template. This enables shader templates that can easily express variants for different data types, packing structures, etc. Introducing a new shader library provides the opportunity to rewrite the shaders in PyTorch Vulkan in a more generic and extensible way.
Differential Revision: [D54754545](https://our.internmc.facebook.com/intern/diff/D54754545/)
[ghstack-poisoned]
@pytorch-bot

pytorch-botBot commented Mar 12, 2024

Copy link
Copy Markdown

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/2366

Note: Links to docs will display an error until the docs builds have been completed.

❌ 2 New Failures

As of commit 5cdc6cd with merge base 4fea983 (image):

NEW FAILURES - The following jobs have failed:

This comment was automatically generated by Dr. CI and updates every 15 minutes.

@facebook-github-botfacebook-github-bot added the CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. label Mar 12, 2024
@facebook-github-bot

Copy link
Copy Markdown
Contributor

This pull request was exported from Phabricator. Differential Revision: D54754545

…s dynamic shapes"
## Context
pytorch/pytorch#121598 introduces the ability to support dynamic shapes through tensor metadata updates.
The idea is fairly simple. Instead of shaders accepting a UBO with size data for all arguments:
```
layout(set = 0, binding = 2) uniform PRECISION restrict Block {
ivec4 output_sizes;
ivec4 other_sizes;
float alpha;
}
```
Shaders will accept separate UBOs for each piece of tensor metadata:
```
layout(set = 0, binding = 3) uniform PRECISION restrict OutSizes {
ivec4 data;
}
out_sizes;
layout(set = 0, binding = 4) uniform PRECISION restrict InSizes {
ivec4 data;
}
in_sizes;
layout(set = 0, binding = 5) uniform PRECISION restrict OtherSizes {
ivec4 data;
}
other_sizes;
layout(set = 0, binding = 6) uniform PRECISION restrict Alpha {
float data;
}
alpha;
```
Each UBO will be owned and maintained by the corresponding `vTensor` instance. To support a graph input resize, every tensor in the graph only needs to update their metadata UBOs via the `tensor.virtual_resize(new_sizes)` call. Shader dispatches in subsequent command buffer submissions will then see the updated metadata and execute as if the tensor were the updated sizes.
This changeset introduces a new shader library for the Vulkan graph runtime that enables dynamic shapes through this technique in favor of relying on the shader library from PyTorch Vulkan.
## Considerations
Technically, the UBO update technique can be applied to the shaders from PyTorch Vulkan as well. If that's the case, why introduce a new shader library for the graph runtime?
The primary motivation is code quality.
First, having `vTensor` supply UBOs for their own metadata greatly reduces the need to have operator specifc ad-hoc `Params` structs to organize arguments to write into a `api::UniformParamsBuffer`.
Constructing an `ExecuteNode` for binary operators is now
```
graph.execute_nodes().emplace_back(new ExecuteNode(
graph,
api::shader_registry().get_shader_info(kernel_name.str()),
global_size,
local_size,
{{out, api::MemoryAccessType::WRITE},
{{arg1, arg2}, api::MemoryAccessType::READ}},
{t_out.gpu_sizes_ubo(),
t_in1.gpu_sizes_ubo(),
t_in2.gpu_sizes_ubo(),
graph.create_params_buffer(alpha_val)}))
```
instead of
```
ArithmeticParams block{
get_size_as_ivec4(t_out),
get_size_as_ivec4(t_in1),
get_size_as_ivec4(t_in2),
alpha_val,
};
api::UniformParamsBuffer params(graph.context(), block);
graph.execute_nodes().emplace_back(new ExecuteNode(
graph,
shader,
global_size,
local_size,
{{out, api::MemoryAccessType::WRITE},
{{arg1, arg2}, api::MemoryAccessType::READ}},
std::move(params)));
```
Another consideration is that pytorch/pytorch#115948 which was landed fairly recently enables much more expressive shader templates through the use of Python code blocks in the GLSL template. This enables shader templates that can easily express variants for different data types, packing structures, etc. Introducing a new shader library provides the opportunity to rewrite the shaders in PyTorch Vulkan in a more generic and extensible way.
Differential Revision: [D54754545](https://our.internmc.facebook.com/intern/diff/D54754545/)
[ghstack-poisoned]
@facebook-github-bot

Copy link
Copy Markdown
Contributor

This pull request was exported from Phabricator. Differential Revision: D54754545

…s dynamic shapes"
## Context
pytorch/pytorch#121598 introduces the ability to support dynamic shapes through tensor metadata updates.
The idea is fairly simple. Instead of shaders accepting a UBO with size data for all arguments:
```
layout(set = 0, binding = 2) uniform PRECISION restrict Block {
ivec4 output_sizes;
ivec4 other_sizes;
float alpha;
}
```
Shaders will accept separate UBOs for each piece of tensor metadata:
```
layout(set = 0, binding = 3) uniform PRECISION restrict OutSizes {
ivec4 data;
}
out_sizes;
layout(set = 0, binding = 4) uniform PRECISION restrict InSizes {
ivec4 data;
}
in_sizes;
layout(set = 0, binding = 5) uniform PRECISION restrict OtherSizes {
ivec4 data;
}
other_sizes;
layout(set = 0, binding = 6) uniform PRECISION restrict Alpha {
float data;
}
alpha;
```
Each UBO will be owned and maintained by the corresponding `vTensor` instance. To support a graph input resize, every tensor in the graph only needs to update their metadata UBOs via the `tensor.virtual_resize(new_sizes)` call. Shader dispatches in subsequent command buffer submissions will then see the updated metadata and execute as if the tensor were the updated sizes.
This changeset introduces a new shader library for the Vulkan graph runtime that enables dynamic shapes through this technique in favor of relying on the shader library from PyTorch Vulkan.
## Considerations
Technically, the UBO update technique can be applied to the shaders from PyTorch Vulkan as well. If that's the case, why introduce a new shader library for the graph runtime?
The primary motivation is code quality.
First, having `vTensor` supply UBOs for their own metadata greatly reduces the need to have operator specifc ad-hoc `Params` structs to organize arguments to write into a `api::UniformParamsBuffer`.
Constructing an `ExecuteNode` for binary operators is now
```
graph.execute_nodes().emplace_back(new ExecuteNode(
graph,
api::shader_registry().get_shader_info(kernel_name.str()),
global_size,
local_size,
{{out, api::MemoryAccessType::WRITE},
{{arg1, arg2}, api::MemoryAccessType::READ}},
{t_out.gpu_sizes_ubo(),
t_in1.gpu_sizes_ubo(),
t_in2.gpu_sizes_ubo(),
graph.create_params_buffer(alpha_val)}))
```
instead of
```
ArithmeticParams block{
get_size_as_ivec4(t_out),
get_size_as_ivec4(t_in1),
get_size_as_ivec4(t_in2),
alpha_val,
};
api::UniformParamsBuffer params(graph.context(), block);
graph.execute_nodes().emplace_back(new ExecuteNode(
graph,
shader,
global_size,
local_size,
{{out, api::MemoryAccessType::WRITE},
{{arg1, arg2}, api::MemoryAccessType::READ}},
std::move(params)));
```
Another consideration is that pytorch/pytorch#115948 which was landed fairly recently enables much more expressive shader templates through the use of Python code blocks in the GLSL template. This enables shader templates that can easily express variants for different data types, packing structures, etc. Introducing a new shader library provides the opportunity to rewrite the shaders in PyTorch Vulkan in a more generic and extensible way.
Differential Revision: [D54754545](https://our.internmc.facebook.com/intern/diff/D54754545/)
[ghstack-poisoned]
@facebook-github-bot

Copy link
Copy Markdown
Contributor

This pull request was exported from Phabricator. Differential Revision: D54754545

…s dynamic shapes"
## Context
pytorch/pytorch#121598 introduces the ability to support dynamic shapes through tensor metadata updates.
The idea is fairly simple. Instead of shaders accepting a UBO with size data for all arguments:
```
layout(set = 0, binding = 2) uniform PRECISION restrict Block {
ivec4 output_sizes;
ivec4 other_sizes;
float alpha;
}
```
Shaders will accept separate UBOs for each piece of tensor metadata:
```
layout(set = 0, binding = 3) uniform PRECISION restrict OutSizes {
ivec4 data;
}
out_sizes;
layout(set = 0, binding = 4) uniform PRECISION restrict InSizes {
ivec4 data;
}
in_sizes;
layout(set = 0, binding = 5) uniform PRECISION restrict OtherSizes {
ivec4 data;
}
other_sizes;
layout(set = 0, binding = 6) uniform PRECISION restrict Alpha {
float data;
}
alpha;
```
Each UBO will be owned and maintained by the corresponding `vTensor` instance. To support a graph input resize, every tensor in the graph only needs to update their metadata UBOs via the `tensor.virtual_resize(new_sizes)` call. Shader dispatches in subsequent command buffer submissions will then see the updated metadata and execute as if the tensor were the updated sizes.
This changeset introduces a new shader library for the Vulkan graph runtime that enables dynamic shapes through this technique in favor of relying on the shader library from PyTorch Vulkan.
## Considerations
Technically, the UBO update technique can be applied to the shaders from PyTorch Vulkan as well. If that's the case, why introduce a new shader library for the graph runtime?
The primary motivation is code quality.
First, having `vTensor` supply UBOs for their own metadata greatly reduces the need to have operator specifc ad-hoc `Params` structs to organize arguments to write into a `api::UniformParamsBuffer`.
Constructing an `ExecuteNode` for binary operators is now
```
graph.execute_nodes().emplace_back(new ExecuteNode(
graph,
api::shader_registry().get_shader_info(kernel_name.str()),
global_size,
local_size,
{{out, api::MemoryAccessType::WRITE},
{{arg1, arg2}, api::MemoryAccessType::READ}},
{t_out.gpu_sizes_ubo(),
t_in1.gpu_sizes_ubo(),
t_in2.gpu_sizes_ubo(),
graph.create_params_buffer(alpha_val)}))
```
instead of
```
ArithmeticParams block{
get_size_as_ivec4(t_out),
get_size_as_ivec4(t_in1),
get_size_as_ivec4(t_in2),
alpha_val,
};
api::UniformParamsBuffer params(graph.context(), block);
graph.execute_nodes().emplace_back(new ExecuteNode(
graph,
shader,
global_size,
local_size,
{{out, api::MemoryAccessType::WRITE},
{{arg1, arg2}, api::MemoryAccessType::READ}},
std::move(params)));
```
Another consideration is that pytorch/pytorch#115948 which was landed fairly recently enables much more expressive shader templates through the use of Python code blocks in the GLSL template. This enables shader templates that can easily express variants for different data types, packing structures, etc. Introducing a new shader library provides the opportunity to rewrite the shaders in PyTorch Vulkan in a more generic and extensible way.
Differential Revision: [D54754545](https://our.internmc.facebook.com/intern/diff/D54754545/)
[ghstack-poisoned]
@facebook-github-bot

Copy link
Copy Markdown
Contributor

This pull request was exported from Phabricator. Differential Revision: D54754545

SS-JIA added a commit that referenced this pull request Mar 13, 2024
…hapes
Pull Request resolved: #2366
## Context
pytorch/pytorch#121598 introduces the ability to support dynamic shapes through tensor metadata updates.
The idea is fairly simple. Instead of shaders accepting a UBO with size data for all arguments:
```
layout(set = 0, binding = 2) uniform PRECISION restrict Block {
ivec4 output_sizes;
ivec4 other_sizes;
float alpha;
}
```
Shaders will accept separate UBOs for each piece of tensor metadata:
```
layout(set = 0, binding = 3) uniform PRECISION restrict OutSizes {
ivec4 data;
}
out_sizes;
layout(set = 0, binding = 4) uniform PRECISION restrict InSizes {
ivec4 data;
}
in_sizes;
layout(set = 0, binding = 5) uniform PRECISION restrict OtherSizes {
ivec4 data;
}
other_sizes;
layout(set = 0, binding = 6) uniform PRECISION restrict Alpha {
float data;
}
alpha;
```
Each UBO will be owned and maintained by the corresponding `vTensor` instance. To support a graph input resize, every tensor in the graph only needs to update their metadata UBOs via the `tensor.virtual_resize(new_sizes)` call. Shader dispatches in subsequent command buffer submissions will then see the updated metadata and execute as if the tensor were the updated sizes.
This changeset introduces a new shader library for the Vulkan graph runtime that enables dynamic shapes through this technique in favor of relying on the shader library from PyTorch Vulkan.
## Considerations
Technically, the UBO update technique can be applied to the shaders from PyTorch Vulkan as well. If that's the case, why introduce a new shader library for the graph runtime?
The primary motivation is code quality.
First, having `vTensor` supply UBOs for their own metadata greatly reduces the need to have operator specifc ad-hoc `Params` structs to organize arguments to write into a `api::UniformParamsBuffer`.
Constructing an `ExecuteNode` for binary operators is now
```
graph.execute_nodes().emplace_back(new ExecuteNode(
graph,
api::shader_registry().get_shader_info(kernel_name.str()),
global_size,
local_size,
{{out, api::MemoryAccessType::WRITE},
{{arg1, arg2}, api::MemoryAccessType::READ}},
{t_out.gpu_sizes_ubo(),
t_in1.gpu_sizes_ubo(),
t_in2.gpu_sizes_ubo(),
graph.create_params_buffer(alpha_val)}))
```
instead of
```
ArithmeticParams block{
get_size_as_ivec4(t_out),
get_size_as_ivec4(t_in1),
get_size_as_ivec4(t_in2),
alpha_val,
};
api::UniformParamsBuffer params(graph.context(), block);
graph.execute_nodes().emplace_back(new ExecuteNode(
graph,
shader,
global_size,
local_size,
{{out, api::MemoryAccessType::WRITE},
{{arg1, arg2}, api::MemoryAccessType::READ}},
std::move(params)));
```
Another consideration is that pytorch/pytorch#115948 which was landed fairly recently enables much more expressive shader templates through the use of Python code blocks in the GLSL template. This enables shader templates that can easily express variants for different data types, packing structures, etc. Introducing a new shader library provides the opportunity to rewrite the shaders in PyTorch Vulkan in a more generic and extensible way.
ghstack-source-id: 218421178
@exported-using-ghexport
Differential Revision: [D54754545](https://our.internmc.facebook.com/intern/diff/D54754545/)
@facebook-github-bot

Copy link
Copy Markdown
Contributor

This pull request has been merged in 835279e.

junpi3 pushed a commit that referenced this pull request Mar 13, 2024
junpi3 pushed a commit that referenced this pull request Mar 13, 2024
Missed this in #2366
Differential Revision: [D54880024](https://our.internmc.facebook.com/intern/diff/D54880024/)
ghstack-source-id: 218593311
Pull Request resolved: #2418
facebook-github-bot pushed a commit that referenced this pull request Mar 14, 2024
Summary:
Pull Request resolved: #2418
Missed this in #2366
ghstack-source-id: 218593311
exported-using-ghexport
bypass-github-export-checks
Reviewed By: SS-JIA
Differential Revision: D54880024
fbshipit-source-id: c4e19d8fefbb9d2fc4547ec2edc236060638da5e
@SS-JIA
SS-JIA deleted the gh/SS-JIA/11/head branch January 24, 2025 19:40
Sign up for freeto join this conversation on GitHub. Already have an account? Sign in to comment

Labels

CLA SignedThis label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed.fb-exportedMerged

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants

@SS-JIA@facebook-github-bot@junpi3
, 'i'); if (__m === '*' || __re.test(location.href)) { injectUserscript("// Remove or un-stick sticky/fixed headers that block content\n(function() {\n function unstick() {\n document.querySelectorAll('header, nav, [role=\"banner\"], .header, .navbar, .sticky, .fixed-top, [style*=\"position: fixed\"], [style*=\"position:sticky\"]').forEach(function(el) {\n if (el.style.position === 'fixed' || el.style.position === 'sticky' || \n getComputedStyle(el).position === 'fixed' || getComputedStyle(el).position === 'sticky') {\n el.style.position = 'static';\n el.style.top = 'auto';\n el.style.zIndex = 'auto';\n }\n });\n }\n \n unstick();\n \n var observer = new MutationObserver(unstick);\n observer.observe(document.body, { childList: true, subtree: true, attributes: true, attributeFilter: ['style', 'class'] });\n})();", "Kill Sticky Headers"); } } catch(__e) { console.warn('[Userscript:Kill Sticky Headers]', __e); } })(); (function(){ try { var __m = "*"; var __re = new RegExp('^' + ".*" + '
Skip to content

[ET-VK] Introduce graph runtime shader library that enables dynamic shapes - #2366

Closed
SS-JIA wants to merge 4 commits into
gh/SS-JIA/11/basefrom
gh/SS-JIA/11/head
Closed

[ET-VK] Introduce graph runtime shader library that enables dynamic shapes#2366
SS-JIA wants to merge 4 commits into
gh/SS-JIA/11/basefrom
gh/SS-JIA/11/head

Conversation

@SS-JIA

@SS-JIASS-JIA commented Mar 12, 2024

Copy link
Copy Markdown
Contributor

Stack from ghstack (oldest at bottom):

Context

pytorch/pytorch#121598 introduces the ability to support dynamic shapes through tensor metadata updates.

The idea is fairly simple. Instead of shaders accepting a UBO with size data for all arguments:

layout(set = 0, binding = 2) uniform PRECISION restrict Block {
ivec4 output_sizes;
ivec4 other_sizes;
float alpha;
}

Shaders will accept separate UBOs for each piece of tensor metadata:

layout(set = 0, binding = 3) uniform PRECISION restrict OutSizes {
ivec4 data;
}
out_sizes;
layout(set = 0, binding = 4) uniform PRECISION restrict InSizes {
ivec4 data;
}
in_sizes;
layout(set = 0, binding = 5) uniform PRECISION restrict OtherSizes {
ivec4 data;
}
other_sizes;
layout(set = 0, binding = 6) uniform PRECISION restrict Alpha {
float data;
}
alpha;

Each UBO will be owned and maintained by the corresponding vTensor instance. To support a graph input resize, every tensor in the graph only needs to update their metadata UBOs via the tensor.virtual_resize(new_sizes) call. Shader dispatches in subsequent command buffer submissions will then see the updated metadata and execute as if the tensor were the updated sizes.

This changeset introduces a new shader library for the Vulkan graph runtime that enables dynamic shapes through this technique in favor of relying on the shader library from PyTorch Vulkan.

Considerations

Technically, the UBO update technique can be applied to the shaders from PyTorch Vulkan as well. If that's the case, why introduce a new shader library for the graph runtime?

The primary motivation is code quality.

First, having vTensor supply UBOs for their own metadata greatly reduces the need to have operator specifc ad-hoc Params structs to organize arguments to write into a api::UniformParamsBuffer.

Constructing an ExecuteNode for binary operators is now

 graph.execute_nodes().emplace_back(new ExecuteNode(
graph,
api::shader_registry().get_shader_info(kernel_name.str()),
global_size,
local_size,
{{out, api::MemoryAccessType::WRITE},
{{arg1, arg2}, api::MemoryAccessType::READ}},
{t_out.gpu_sizes_ubo(),
t_in1.gpu_sizes_ubo(),
t_in2.gpu_sizes_ubo(),
graph.create_params_buffer(alpha_val)}))

instead of

ArithmeticParams block{
get_size_as_ivec4(t_out),
get_size_as_ivec4(t_in1),
get_size_as_ivec4(t_in2),
alpha_val,
};
api::UniformParamsBuffer params(graph.context(), block);
graph.execute_nodes().emplace_back(new ExecuteNode(
graph,
shader,
global_size,
local_size,
{{out, api::MemoryAccessType::WRITE},
{{arg1, arg2}, api::MemoryAccessType::READ}},
std::move(params)));

Another consideration is that pytorch/pytorch#115948 which was landed fairly recently enables much more expressive shader templates through the use of Python code blocks in the GLSL template. This enables shader templates that can easily express variants for different data types, packing structures, etc. Introducing a new shader library provides the opportunity to rewrite the shaders in PyTorch Vulkan in a more generic and extensible way.

Differential Revision: D54754545

…hapes
## Context
pytorch/pytorch#121598 introduces the ability to support dynamic shapes through tensor metadata updates.
The idea is fairly simple. Instead of shaders accepting a UBO with size data for all arguments:
```
layout(set = 0, binding = 2) uniform PRECISION restrict Block {
ivec4 output_sizes;
ivec4 other_sizes;
float alpha;
}
```
Shaders will accept separate UBOs for each piece of tensor metadata:
```
layout(set = 0, binding = 3) uniform PRECISION restrict OutSizes {
ivec4 data;
}
out_sizes;
layout(set = 0, binding = 4) uniform PRECISION restrict InSizes {
ivec4 data;
}
in_sizes;
layout(set = 0, binding = 5) uniform PRECISION restrict OtherSizes {
ivec4 data;
}
other_sizes;
layout(set = 0, binding = 6) uniform PRECISION restrict Alpha {
float data;
}
alpha;
```
Each UBO will be owned and maintained by the corresponding `vTensor` instance. To support a graph input resize, every tensor in the graph only needs to update their metadata UBOs via the `tensor.virtual_resize(new_sizes)` call. Shader dispatches in subsequent command buffer submissions will then see the updated metadata and execute as if the tensor were the updated sizes.
This changeset introduces a new shader library for the Vulkan graph runtime that enables dynamic shapes through this technique in favor of relying on the shader library from PyTorch Vulkan.
## Considerations
Technically, the UBO update technique can be applied to the shaders from PyTorch Vulkan as well. If that's the case, why introduce a new shader library for the graph runtime?
The primary motivation is code quality.
First, having `vTensor` supply UBOs for their own metadata greatly reduces the need to have operator specifc ad-hoc `Params` structs to organize arguments to write into a `api::UniformParamsBuffer`.
Constructing an `ExecuteNode` for binary operators is now
```
graph.execute_nodes().emplace_back(new ExecuteNode(
graph,
api::shader_registry().get_shader_info(kernel_name.str()),
global_size,
local_size,
{{out, api::MemoryAccessType::WRITE},
{{arg1, arg2}, api::MemoryAccessType::READ}},
{t_out.gpu_sizes_ubo(),
t_in1.gpu_sizes_ubo(),
t_in2.gpu_sizes_ubo(),
graph.create_params_buffer(alpha_val)}))
```
instead of
```
ArithmeticParams block{
get_size_as_ivec4(t_out),
get_size_as_ivec4(t_in1),
get_size_as_ivec4(t_in2),
alpha_val,
};
api::UniformParamsBuffer params(graph.context(), block);
graph.execute_nodes().emplace_back(new ExecuteNode(
graph,
shader,
global_size,
local_size,
{{out, api::MemoryAccessType::WRITE},
{{arg1, arg2}, api::MemoryAccessType::READ}},
std::move(params)));
```
Another consideration is that pytorch/pytorch#115948 which was landed fairly recently enables much more expressive shader templates through the use of Python code blocks in the GLSL template. This enables shader templates that can easily express variants for different data types, packing structures, etc. Introducing a new shader library provides the opportunity to rewrite the shaders in PyTorch Vulkan in a more generic and extensible way.
Differential Revision: [D54754545](https://our.internmc.facebook.com/intern/diff/D54754545/)
[ghstack-poisoned]
@pytorch-bot

pytorch-botBot commented Mar 12, 2024

Copy link
Copy Markdown

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/2366

Note: Links to docs will display an error until the docs builds have been completed.

❌ 2 New Failures

As of commit 5cdc6cd with merge base 4fea983 (image):

NEW FAILURES - The following jobs have failed:

This comment was automatically generated by Dr. CI and updates every 15 minutes.

@facebook-github-botfacebook-github-bot added the CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. label Mar 12, 2024
@facebook-github-bot

Copy link
Copy Markdown
Contributor

This pull request was exported from Phabricator. Differential Revision: D54754545

…s dynamic shapes"
## Context
pytorch/pytorch#121598 introduces the ability to support dynamic shapes through tensor metadata updates.
The idea is fairly simple. Instead of shaders accepting a UBO with size data for all arguments:
```
layout(set = 0, binding = 2) uniform PRECISION restrict Block {
ivec4 output_sizes;
ivec4 other_sizes;
float alpha;
}
```
Shaders will accept separate UBOs for each piece of tensor metadata:
```
layout(set = 0, binding = 3) uniform PRECISION restrict OutSizes {
ivec4 data;
}
out_sizes;
layout(set = 0, binding = 4) uniform PRECISION restrict InSizes {
ivec4 data;
}
in_sizes;
layout(set = 0, binding = 5) uniform PRECISION restrict OtherSizes {
ivec4 data;
}
other_sizes;
layout(set = 0, binding = 6) uniform PRECISION restrict Alpha {
float data;
}
alpha;
```
Each UBO will be owned and maintained by the corresponding `vTensor` instance. To support a graph input resize, every tensor in the graph only needs to update their metadata UBOs via the `tensor.virtual_resize(new_sizes)` call. Shader dispatches in subsequent command buffer submissions will then see the updated metadata and execute as if the tensor were the updated sizes.
This changeset introduces a new shader library for the Vulkan graph runtime that enables dynamic shapes through this technique in favor of relying on the shader library from PyTorch Vulkan.
## Considerations
Technically, the UBO update technique can be applied to the shaders from PyTorch Vulkan as well. If that's the case, why introduce a new shader library for the graph runtime?
The primary motivation is code quality.
First, having `vTensor` supply UBOs for their own metadata greatly reduces the need to have operator specifc ad-hoc `Params` structs to organize arguments to write into a `api::UniformParamsBuffer`.
Constructing an `ExecuteNode` for binary operators is now
```
graph.execute_nodes().emplace_back(new ExecuteNode(
graph,
api::shader_registry().get_shader_info(kernel_name.str()),
global_size,
local_size,
{{out, api::MemoryAccessType::WRITE},
{{arg1, arg2}, api::MemoryAccessType::READ}},
{t_out.gpu_sizes_ubo(),
t_in1.gpu_sizes_ubo(),
t_in2.gpu_sizes_ubo(),
graph.create_params_buffer(alpha_val)}))
```
instead of
```
ArithmeticParams block{
get_size_as_ivec4(t_out),
get_size_as_ivec4(t_in1),
get_size_as_ivec4(t_in2),
alpha_val,
};
api::UniformParamsBuffer params(graph.context(), block);
graph.execute_nodes().emplace_back(new ExecuteNode(
graph,
shader,
global_size,
local_size,
{{out, api::MemoryAccessType::WRITE},
{{arg1, arg2}, api::MemoryAccessType::READ}},
std::move(params)));
```
Another consideration is that pytorch/pytorch#115948 which was landed fairly recently enables much more expressive shader templates through the use of Python code blocks in the GLSL template. This enables shader templates that can easily express variants for different data types, packing structures, etc. Introducing a new shader library provides the opportunity to rewrite the shaders in PyTorch Vulkan in a more generic and extensible way.
Differential Revision: [D54754545](https://our.internmc.facebook.com/intern/diff/D54754545/)
[ghstack-poisoned]
@facebook-github-bot

Copy link
Copy Markdown
Contributor

This pull request was exported from Phabricator. Differential Revision: D54754545

…s dynamic shapes"
## Context
pytorch/pytorch#121598 introduces the ability to support dynamic shapes through tensor metadata updates.
The idea is fairly simple. Instead of shaders accepting a UBO with size data for all arguments:
```
layout(set = 0, binding = 2) uniform PRECISION restrict Block {
ivec4 output_sizes;
ivec4 other_sizes;
float alpha;
}
```
Shaders will accept separate UBOs for each piece of tensor metadata:
```
layout(set = 0, binding = 3) uniform PRECISION restrict OutSizes {
ivec4 data;
}
out_sizes;
layout(set = 0, binding = 4) uniform PRECISION restrict InSizes {
ivec4 data;
}
in_sizes;
layout(set = 0, binding = 5) uniform PRECISION restrict OtherSizes {
ivec4 data;
}
other_sizes;
layout(set = 0, binding = 6) uniform PRECISION restrict Alpha {
float data;
}
alpha;
```
Each UBO will be owned and maintained by the corresponding `vTensor` instance. To support a graph input resize, every tensor in the graph only needs to update their metadata UBOs via the `tensor.virtual_resize(new_sizes)` call. Shader dispatches in subsequent command buffer submissions will then see the updated metadata and execute as if the tensor were the updated sizes.
This changeset introduces a new shader library for the Vulkan graph runtime that enables dynamic shapes through this technique in favor of relying on the shader library from PyTorch Vulkan.
## Considerations
Technically, the UBO update technique can be applied to the shaders from PyTorch Vulkan as well. If that's the case, why introduce a new shader library for the graph runtime?
The primary motivation is code quality.
First, having `vTensor` supply UBOs for their own metadata greatly reduces the need to have operator specifc ad-hoc `Params` structs to organize arguments to write into a `api::UniformParamsBuffer`.
Constructing an `ExecuteNode` for binary operators is now
```
graph.execute_nodes().emplace_back(new ExecuteNode(
graph,
api::shader_registry().get_shader_info(kernel_name.str()),
global_size,
local_size,
{{out, api::MemoryAccessType::WRITE},
{{arg1, arg2}, api::MemoryAccessType::READ}},
{t_out.gpu_sizes_ubo(),
t_in1.gpu_sizes_ubo(),
t_in2.gpu_sizes_ubo(),
graph.create_params_buffer(alpha_val)}))
```
instead of
```
ArithmeticParams block{
get_size_as_ivec4(t_out),
get_size_as_ivec4(t_in1),
get_size_as_ivec4(t_in2),
alpha_val,
};
api::UniformParamsBuffer params(graph.context(), block);
graph.execute_nodes().emplace_back(new ExecuteNode(
graph,
shader,
global_size,
local_size,
{{out, api::MemoryAccessType::WRITE},
{{arg1, arg2}, api::MemoryAccessType::READ}},
std::move(params)));
```
Another consideration is that pytorch/pytorch#115948 which was landed fairly recently enables much more expressive shader templates through the use of Python code blocks in the GLSL template. This enables shader templates that can easily express variants for different data types, packing structures, etc. Introducing a new shader library provides the opportunity to rewrite the shaders in PyTorch Vulkan in a more generic and extensible way.
Differential Revision: [D54754545](https://our.internmc.facebook.com/intern/diff/D54754545/)
[ghstack-poisoned]
@facebook-github-bot

Copy link
Copy Markdown
Contributor

This pull request was exported from Phabricator. Differential Revision: D54754545

…s dynamic shapes"
## Context
pytorch/pytorch#121598 introduces the ability to support dynamic shapes through tensor metadata updates.
The idea is fairly simple. Instead of shaders accepting a UBO with size data for all arguments:
```
layout(set = 0, binding = 2) uniform PRECISION restrict Block {
ivec4 output_sizes;
ivec4 other_sizes;
float alpha;
}
```
Shaders will accept separate UBOs for each piece of tensor metadata:
```
layout(set = 0, binding = 3) uniform PRECISION restrict OutSizes {
ivec4 data;
}
out_sizes;
layout(set = 0, binding = 4) uniform PRECISION restrict InSizes {
ivec4 data;
}
in_sizes;
layout(set = 0, binding = 5) uniform PRECISION restrict OtherSizes {
ivec4 data;
}
other_sizes;
layout(set = 0, binding = 6) uniform PRECISION restrict Alpha {
float data;
}
alpha;
```
Each UBO will be owned and maintained by the corresponding `vTensor` instance. To support a graph input resize, every tensor in the graph only needs to update their metadata UBOs via the `tensor.virtual_resize(new_sizes)` call. Shader dispatches in subsequent command buffer submissions will then see the updated metadata and execute as if the tensor were the updated sizes.
This changeset introduces a new shader library for the Vulkan graph runtime that enables dynamic shapes through this technique in favor of relying on the shader library from PyTorch Vulkan.
## Considerations
Technically, the UBO update technique can be applied to the shaders from PyTorch Vulkan as well. If that's the case, why introduce a new shader library for the graph runtime?
The primary motivation is code quality.
First, having `vTensor` supply UBOs for their own metadata greatly reduces the need to have operator specifc ad-hoc `Params` structs to organize arguments to write into a `api::UniformParamsBuffer`.
Constructing an `ExecuteNode` for binary operators is now
```
graph.execute_nodes().emplace_back(new ExecuteNode(
graph,
api::shader_registry().get_shader_info(kernel_name.str()),
global_size,
local_size,
{{out, api::MemoryAccessType::WRITE},
{{arg1, arg2}, api::MemoryAccessType::READ}},
{t_out.gpu_sizes_ubo(),
t_in1.gpu_sizes_ubo(),
t_in2.gpu_sizes_ubo(),
graph.create_params_buffer(alpha_val)}))
```
instead of
```
ArithmeticParams block{
get_size_as_ivec4(t_out),
get_size_as_ivec4(t_in1),
get_size_as_ivec4(t_in2),
alpha_val,
};
api::UniformParamsBuffer params(graph.context(), block);
graph.execute_nodes().emplace_back(new ExecuteNode(
graph,
shader,
global_size,
local_size,
{{out, api::MemoryAccessType::WRITE},
{{arg1, arg2}, api::MemoryAccessType::READ}},
std::move(params)));
```
Another consideration is that pytorch/pytorch#115948 which was landed fairly recently enables much more expressive shader templates through the use of Python code blocks in the GLSL template. This enables shader templates that can easily express variants for different data types, packing structures, etc. Introducing a new shader library provides the opportunity to rewrite the shaders in PyTorch Vulkan in a more generic and extensible way.
Differential Revision: [D54754545](https://our.internmc.facebook.com/intern/diff/D54754545/)
[ghstack-poisoned]
@facebook-github-bot

Copy link
Copy Markdown
Contributor

This pull request was exported from Phabricator. Differential Revision: D54754545

SS-JIA added a commit that referenced this pull request Mar 13, 2024
…hapes
Pull Request resolved: #2366
## Context
pytorch/pytorch#121598 introduces the ability to support dynamic shapes through tensor metadata updates.
The idea is fairly simple. Instead of shaders accepting a UBO with size data for all arguments:
```
layout(set = 0, binding = 2) uniform PRECISION restrict Block {
ivec4 output_sizes;
ivec4 other_sizes;
float alpha;
}
```
Shaders will accept separate UBOs for each piece of tensor metadata:
```
layout(set = 0, binding = 3) uniform PRECISION restrict OutSizes {
ivec4 data;
}
out_sizes;
layout(set = 0, binding = 4) uniform PRECISION restrict InSizes {
ivec4 data;
}
in_sizes;
layout(set = 0, binding = 5) uniform PRECISION restrict OtherSizes {
ivec4 data;
}
other_sizes;
layout(set = 0, binding = 6) uniform PRECISION restrict Alpha {
float data;
}
alpha;
```
Each UBO will be owned and maintained by the corresponding `vTensor` instance. To support a graph input resize, every tensor in the graph only needs to update their metadata UBOs via the `tensor.virtual_resize(new_sizes)` call. Shader dispatches in subsequent command buffer submissions will then see the updated metadata and execute as if the tensor were the updated sizes.
This changeset introduces a new shader library for the Vulkan graph runtime that enables dynamic shapes through this technique in favor of relying on the shader library from PyTorch Vulkan.
## Considerations
Technically, the UBO update technique can be applied to the shaders from PyTorch Vulkan as well. If that's the case, why introduce a new shader library for the graph runtime?
The primary motivation is code quality.
First, having `vTensor` supply UBOs for their own metadata greatly reduces the need to have operator specifc ad-hoc `Params` structs to organize arguments to write into a `api::UniformParamsBuffer`.
Constructing an `ExecuteNode` for binary operators is now
```
graph.execute_nodes().emplace_back(new ExecuteNode(
graph,
api::shader_registry().get_shader_info(kernel_name.str()),
global_size,
local_size,
{{out, api::MemoryAccessType::WRITE},
{{arg1, arg2}, api::MemoryAccessType::READ}},
{t_out.gpu_sizes_ubo(),
t_in1.gpu_sizes_ubo(),
t_in2.gpu_sizes_ubo(),
graph.create_params_buffer(alpha_val)}))
```
instead of
```
ArithmeticParams block{
get_size_as_ivec4(t_out),
get_size_as_ivec4(t_in1),
get_size_as_ivec4(t_in2),
alpha_val,
};
api::UniformParamsBuffer params(graph.context(), block);
graph.execute_nodes().emplace_back(new ExecuteNode(
graph,
shader,
global_size,
local_size,
{{out, api::MemoryAccessType::WRITE},
{{arg1, arg2}, api::MemoryAccessType::READ}},
std::move(params)));
```
Another consideration is that pytorch/pytorch#115948 which was landed fairly recently enables much more expressive shader templates through the use of Python code blocks in the GLSL template. This enables shader templates that can easily express variants for different data types, packing structures, etc. Introducing a new shader library provides the opportunity to rewrite the shaders in PyTorch Vulkan in a more generic and extensible way.
ghstack-source-id: 218421178
@exported-using-ghexport
Differential Revision: [D54754545](https://our.internmc.facebook.com/intern/diff/D54754545/)
@facebook-github-bot

Copy link
Copy Markdown
Contributor

This pull request has been merged in 835279e.

junpi3 pushed a commit that referenced this pull request Mar 13, 2024
junpi3 pushed a commit that referenced this pull request Mar 13, 2024
Missed this in #2366
Differential Revision: [D54880024](https://our.internmc.facebook.com/intern/diff/D54880024/)
ghstack-source-id: 218593311
Pull Request resolved: #2418
facebook-github-bot pushed a commit that referenced this pull request Mar 14, 2024
Summary:
Pull Request resolved: #2418
Missed this in #2366
ghstack-source-id: 218593311
exported-using-ghexport
bypass-github-export-checks
Reviewed By: SS-JIA
Differential Revision: D54880024
fbshipit-source-id: c4e19d8fefbb9d2fc4547ec2edc236060638da5e
@SS-JIA
SS-JIA deleted the gh/SS-JIA/11/head branch January 24, 2025 19:40
Sign up for freeto join this conversation on GitHub. Already have an account? Sign in to comment

Labels

CLA SignedThis label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed.fb-exportedMerged

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants

@SS-JIA@facebook-github-bot@junpi3
, 'i'); if (__m === '*' || __re.test(location.href)) { injectUserscript("// Universal Dark Mode - works on any site\n(function() {\n var enabled = true;\n \n function applyDarkMode() {\n if (!enabled) return;\n \n // Create style element if it doesn't exist\n var style = document.getElementById('universal-dark-mode-style');\n if (!style) {\n style = document.createElement('style');\n style.id = 'universal-dark-mode-style';\n document.head.appendChild(style);\n }\n \n // Dark mode CSS - inverts colors but preserves images/video\n style.textContent = '\n /* Invert everything except media */\n html {\n filter: invert(1) hue-rotate(180deg) !important;\n background: #1a1a2e !important;\n }\n \n /* Restore images, videos, iframes, canvas */\n img, video, iframe, canvas, svg, picture, [style*=\"background-image\"] {\n filter: invert(1) hue-rotate(180deg) !important;\n }\n \n /* Preserve specific elements that should not be inverted */\n .no-dark-mode, .no-dark-mode *,\n [data-theme=\"light\"], [data-theme=\"light\"],\n .ace_editor, .ace_editor *,\n .CodeMirror, .CodeMirror *,\n .monaco-editor, .monaco-editor *,\n .markdown-body pre, .markdown-body pre *,\n .highlight, .highlight *,\n pre code, pre code * {\n filter: none !important;\n }\n \n /* Fix common UI elements */\n .modal, .popup, .dropdown-menu, .tooltip, .popover {\n filter: invert(1) hue-rotate(180deg) !important;\n background: #2d2d44 !important;\n border-color: #444 !important;\n }\n \n /* Scrollbars */\n ::-webkit-scrollbar { background: #1a1a2e !important; }\n ::-webkit-scrollbar-thumb { background: #444 !important; }\n ::-webkit-scrollbar-thumb:hover { background: #555 !important; }\n \n /* Selection */\n ::selection { background: #4ecdc4 !important; color: #1a1a2e !important; }\n ::-moz-selection { background: #4ecdc4 !important; color: #1a1a2e !important; }\n ';\n }\n \n function removeDarkMode() {\n var style = document.getElementById('universal-dark-mode-style');\n if (style) style.remove();\n }\n \n // Toggle with Alt+Shift+D\n document.addEventListener('keydown', function(e) {\n if (e.altKey && e.shiftKey && e.key === 'D') {\n e.preventDefault();\n enabled = !enabled;\n if (enabled) {\n applyDarkMode();\n console.log('[Universal Dark Mode] Enabled');\n } else {\n removeDarkMode();\n console.log('[Universal Dark Mode] Disabled');\n }\n }\n });\n \n // Apply on load\n applyDarkMode();\n \n // Re-apply on dynamic content\n var observer = new MutationObserver(function(mutations) {\n if (enabled && !document.getElementById('universal-dark-mode-style')) {\n applyDarkMode();\n }\n });\n observer.observe(document.head, { childList: true });\n \n console.log('[Universal Dark Mode] Loaded - Press Alt+Shift+D to toggle');\n})();", "Universal Dark Mode"); } } catch(__e) { console.warn('[Userscript:Universal Dark Mode]', __e); } })(); })();
Skip to content

[ET-VK] Introduce graph runtime shader library that enables dynamic shapes - #2366

Closed
SS-JIA wants to merge 4 commits into
gh/SS-JIA/11/basefrom
gh/SS-JIA/11/head
Closed

[ET-VK] Introduce graph runtime shader library that enables dynamic shapes#2366
SS-JIA wants to merge 4 commits into
gh/SS-JIA/11/basefrom
gh/SS-JIA/11/head

Conversation

@SS-JIA

@SS-JIASS-JIA commented Mar 12, 2024

Copy link
Copy Markdown
Contributor

Stack from ghstack (oldest at bottom):

Context

pytorch/pytorch#121598 introduces the ability to support dynamic shapes through tensor metadata updates.

The idea is fairly simple. Instead of shaders accepting a UBO with size data for all arguments:

layout(set = 0, binding = 2) uniform PRECISION restrict Block {
ivec4 output_sizes;
ivec4 other_sizes;
float alpha;
}

Shaders will accept separate UBOs for each piece of tensor metadata:

layout(set = 0, binding = 3) uniform PRECISION restrict OutSizes {
ivec4 data;
}
out_sizes;
layout(set = 0, binding = 4) uniform PRECISION restrict InSizes {
ivec4 data;
}
in_sizes;
layout(set = 0, binding = 5) uniform PRECISION restrict OtherSizes {
ivec4 data;
}
other_sizes;
layout(set = 0, binding = 6) uniform PRECISION restrict Alpha {
float data;
}
alpha;

Each UBO will be owned and maintained by the corresponding vTensor instance. To support a graph input resize, every tensor in the graph only needs to update their metadata UBOs via the tensor.virtual_resize(new_sizes) call. Shader dispatches in subsequent command buffer submissions will then see the updated metadata and execute as if the tensor were the updated sizes.

This changeset introduces a new shader library for the Vulkan graph runtime that enables dynamic shapes through this technique in favor of relying on the shader library from PyTorch Vulkan.

Considerations

Technically, the UBO update technique can be applied to the shaders from PyTorch Vulkan as well. If that's the case, why introduce a new shader library for the graph runtime?

The primary motivation is code quality.

First, having vTensor supply UBOs for their own metadata greatly reduces the need to have operator specifc ad-hoc Params structs to organize arguments to write into a api::UniformParamsBuffer.

Constructing an ExecuteNode for binary operators is now

 graph.execute_nodes().emplace_back(new ExecuteNode(
graph,
api::shader_registry().get_shader_info(kernel_name.str()),
global_size,
local_size,
{{out, api::MemoryAccessType::WRITE},
{{arg1, arg2}, api::MemoryAccessType::READ}},
{t_out.gpu_sizes_ubo(),
t_in1.gpu_sizes_ubo(),
t_in2.gpu_sizes_ubo(),
graph.create_params_buffer(alpha_val)}))

instead of

ArithmeticParams block{
get_size_as_ivec4(t_out),
get_size_as_ivec4(t_in1),
get_size_as_ivec4(t_in2),
alpha_val,
};
api::UniformParamsBuffer params(graph.context(), block);
graph.execute_nodes().emplace_back(new ExecuteNode(
graph,
shader,
global_size,
local_size,
{{out, api::MemoryAccessType::WRITE},
{{arg1, arg2}, api::MemoryAccessType::READ}},
std::move(params)));

Another consideration is that pytorch/pytorch#115948 which was landed fairly recently enables much more expressive shader templates through the use of Python code blocks in the GLSL template. This enables shader templates that can easily express variants for different data types, packing structures, etc. Introducing a new shader library provides the opportunity to rewrite the shaders in PyTorch Vulkan in a more generic and extensible way.

Differential Revision: D54754545

…hapes
## Context
pytorch/pytorch#121598 introduces the ability to support dynamic shapes through tensor metadata updates.
The idea is fairly simple. Instead of shaders accepting a UBO with size data for all arguments:
```
layout(set = 0, binding = 2) uniform PRECISION restrict Block {
ivec4 output_sizes;
ivec4 other_sizes;
float alpha;
}
```
Shaders will accept separate UBOs for each piece of tensor metadata:
```
layout(set = 0, binding = 3) uniform PRECISION restrict OutSizes {
ivec4 data;
}
out_sizes;
layout(set = 0, binding = 4) uniform PRECISION restrict InSizes {
ivec4 data;
}
in_sizes;
layout(set = 0, binding = 5) uniform PRECISION restrict OtherSizes {
ivec4 data;
}
other_sizes;
layout(set = 0, binding = 6) uniform PRECISION restrict Alpha {
float data;
}
alpha;
```
Each UBO will be owned and maintained by the corresponding `vTensor` instance. To support a graph input resize, every tensor in the graph only needs to update their metadata UBOs via the `tensor.virtual_resize(new_sizes)` call. Shader dispatches in subsequent command buffer submissions will then see the updated metadata and execute as if the tensor were the updated sizes.
This changeset introduces a new shader library for the Vulkan graph runtime that enables dynamic shapes through this technique in favor of relying on the shader library from PyTorch Vulkan.
## Considerations
Technically, the UBO update technique can be applied to the shaders from PyTorch Vulkan as well. If that's the case, why introduce a new shader library for the graph runtime?
The primary motivation is code quality.
First, having `vTensor` supply UBOs for their own metadata greatly reduces the need to have operator specifc ad-hoc `Params` structs to organize arguments to write into a `api::UniformParamsBuffer`.
Constructing an `ExecuteNode` for binary operators is now
```
graph.execute_nodes().emplace_back(new ExecuteNode(
graph,
api::shader_registry().get_shader_info(kernel_name.str()),
global_size,
local_size,
{{out, api::MemoryAccessType::WRITE},
{{arg1, arg2}, api::MemoryAccessType::READ}},
{t_out.gpu_sizes_ubo(),
t_in1.gpu_sizes_ubo(),
t_in2.gpu_sizes_ubo(),
graph.create_params_buffer(alpha_val)}))
```
instead of
```
ArithmeticParams block{
get_size_as_ivec4(t_out),
get_size_as_ivec4(t_in1),
get_size_as_ivec4(t_in2),
alpha_val,
};
api::UniformParamsBuffer params(graph.context(), block);
graph.execute_nodes().emplace_back(new ExecuteNode(
graph,
shader,
global_size,
local_size,
{{out, api::MemoryAccessType::WRITE},
{{arg1, arg2}, api::MemoryAccessType::READ}},
std::move(params)));
```
Another consideration is that pytorch/pytorch#115948 which was landed fairly recently enables much more expressive shader templates through the use of Python code blocks in the GLSL template. This enables shader templates that can easily express variants for different data types, packing structures, etc. Introducing a new shader library provides the opportunity to rewrite the shaders in PyTorch Vulkan in a more generic and extensible way.
Differential Revision: [D54754545](https://our.internmc.facebook.com/intern/diff/D54754545/)
[ghstack-poisoned]
@pytorch-bot

pytorch-botBot commented Mar 12, 2024

Copy link
Copy Markdown

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/2366

Note: Links to docs will display an error until the docs builds have been completed.

❌ 2 New Failures

As of commit 5cdc6cd with merge base 4fea983 (image):

NEW FAILURES - The following jobs have failed:

This comment was automatically generated by Dr. CI and updates every 15 minutes.

@facebook-github-botfacebook-github-bot added the CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. label Mar 12, 2024
@facebook-github-bot

Copy link
Copy Markdown
Contributor

This pull request was exported from Phabricator. Differential Revision: D54754545

…s dynamic shapes"
## Context
pytorch/pytorch#121598 introduces the ability to support dynamic shapes through tensor metadata updates.
The idea is fairly simple. Instead of shaders accepting a UBO with size data for all arguments:
```
layout(set = 0, binding = 2) uniform PRECISION restrict Block {
ivec4 output_sizes;
ivec4 other_sizes;
float alpha;
}
```
Shaders will accept separate UBOs for each piece of tensor metadata:
```
layout(set = 0, binding = 3) uniform PRECISION restrict OutSizes {
ivec4 data;
}
out_sizes;
layout(set = 0, binding = 4) uniform PRECISION restrict InSizes {
ivec4 data;
}
in_sizes;
layout(set = 0, binding = 5) uniform PRECISION restrict OtherSizes {
ivec4 data;
}
other_sizes;
layout(set = 0, binding = 6) uniform PRECISION restrict Alpha {
float data;
}
alpha;
```
Each UBO will be owned and maintained by the corresponding `vTensor` instance. To support a graph input resize, every tensor in the graph only needs to update their metadata UBOs via the `tensor.virtual_resize(new_sizes)` call. Shader dispatches in subsequent command buffer submissions will then see the updated metadata and execute as if the tensor were the updated sizes.
This changeset introduces a new shader library for the Vulkan graph runtime that enables dynamic shapes through this technique in favor of relying on the shader library from PyTorch Vulkan.
## Considerations
Technically, the UBO update technique can be applied to the shaders from PyTorch Vulkan as well. If that's the case, why introduce a new shader library for the graph runtime?
The primary motivation is code quality.
First, having `vTensor` supply UBOs for their own metadata greatly reduces the need to have operator specifc ad-hoc `Params` structs to organize arguments to write into a `api::UniformParamsBuffer`.
Constructing an `ExecuteNode` for binary operators is now
```
graph.execute_nodes().emplace_back(new ExecuteNode(
graph,
api::shader_registry().get_shader_info(kernel_name.str()),
global_size,
local_size,
{{out, api::MemoryAccessType::WRITE},
{{arg1, arg2}, api::MemoryAccessType::READ}},
{t_out.gpu_sizes_ubo(),
t_in1.gpu_sizes_ubo(),
t_in2.gpu_sizes_ubo(),
graph.create_params_buffer(alpha_val)}))
```
instead of
```
ArithmeticParams block{
get_size_as_ivec4(t_out),
get_size_as_ivec4(t_in1),
get_size_as_ivec4(t_in2),
alpha_val,
};
api::UniformParamsBuffer params(graph.context(), block);
graph.execute_nodes().emplace_back(new ExecuteNode(
graph,
shader,
global_size,
local_size,
{{out, api::MemoryAccessType::WRITE},
{{arg1, arg2}, api::MemoryAccessType::READ}},
std::move(params)));
```
Another consideration is that pytorch/pytorch#115948 which was landed fairly recently enables much more expressive shader templates through the use of Python code blocks in the GLSL template. This enables shader templates that can easily express variants for different data types, packing structures, etc. Introducing a new shader library provides the opportunity to rewrite the shaders in PyTorch Vulkan in a more generic and extensible way.
Differential Revision: [D54754545](https://our.internmc.facebook.com/intern/diff/D54754545/)
[ghstack-poisoned]
@facebook-github-bot

Copy link
Copy Markdown
Contributor

This pull request was exported from Phabricator. Differential Revision: D54754545

…s dynamic shapes"
## Context
pytorch/pytorch#121598 introduces the ability to support dynamic shapes through tensor metadata updates.
The idea is fairly simple. Instead of shaders accepting a UBO with size data for all arguments:
```
layout(set = 0, binding = 2) uniform PRECISION restrict Block {
ivec4 output_sizes;
ivec4 other_sizes;
float alpha;
}
```
Shaders will accept separate UBOs for each piece of tensor metadata:
```
layout(set = 0, binding = 3) uniform PRECISION restrict OutSizes {
ivec4 data;
}
out_sizes;
layout(set = 0, binding = 4) uniform PRECISION restrict InSizes {
ivec4 data;
}
in_sizes;
layout(set = 0, binding = 5) uniform PRECISION restrict OtherSizes {
ivec4 data;
}
other_sizes;
layout(set = 0, binding = 6) uniform PRECISION restrict Alpha {
float data;
}
alpha;
```
Each UBO will be owned and maintained by the corresponding `vTensor` instance. To support a graph input resize, every tensor in the graph only needs to update their metadata UBOs via the `tensor.virtual_resize(new_sizes)` call. Shader dispatches in subsequent command buffer submissions will then see the updated metadata and execute as if the tensor were the updated sizes.
This changeset introduces a new shader library for the Vulkan graph runtime that enables dynamic shapes through this technique in favor of relying on the shader library from PyTorch Vulkan.
## Considerations
Technically, the UBO update technique can be applied to the shaders from PyTorch Vulkan as well. If that's the case, why introduce a new shader library for the graph runtime?
The primary motivation is code quality.
First, having `vTensor` supply UBOs for their own metadata greatly reduces the need to have operator specifc ad-hoc `Params` structs to organize arguments to write into a `api::UniformParamsBuffer`.
Constructing an `ExecuteNode` for binary operators is now
```
graph.execute_nodes().emplace_back(new ExecuteNode(
graph,
api::shader_registry().get_shader_info(kernel_name.str()),
global_size,
local_size,
{{out, api::MemoryAccessType::WRITE},
{{arg1, arg2}, api::MemoryAccessType::READ}},
{t_out.gpu_sizes_ubo(),
t_in1.gpu_sizes_ubo(),
t_in2.gpu_sizes_ubo(),
graph.create_params_buffer(alpha_val)}))
```
instead of
```
ArithmeticParams block{
get_size_as_ivec4(t_out),
get_size_as_ivec4(t_in1),
get_size_as_ivec4(t_in2),
alpha_val,
};
api::UniformParamsBuffer params(graph.context(), block);
graph.execute_nodes().emplace_back(new ExecuteNode(
graph,
shader,
global_size,
local_size,
{{out, api::MemoryAccessType::WRITE},
{{arg1, arg2}, api::MemoryAccessType::READ}},
std::move(params)));
```
Another consideration is that pytorch/pytorch#115948 which was landed fairly recently enables much more expressive shader templates through the use of Python code blocks in the GLSL template. This enables shader templates that can easily express variants for different data types, packing structures, etc. Introducing a new shader library provides the opportunity to rewrite the shaders in PyTorch Vulkan in a more generic and extensible way.
Differential Revision: [D54754545](https://our.internmc.facebook.com/intern/diff/D54754545/)
[ghstack-poisoned]
@facebook-github-bot

Copy link
Copy Markdown
Contributor

This pull request was exported from Phabricator. Differential Revision: D54754545

…s dynamic shapes"
## Context
pytorch/pytorch#121598 introduces the ability to support dynamic shapes through tensor metadata updates.
The idea is fairly simple. Instead of shaders accepting a UBO with size data for all arguments:
```
layout(set = 0, binding = 2) uniform PRECISION restrict Block {
ivec4 output_sizes;
ivec4 other_sizes;
float alpha;
}
```
Shaders will accept separate UBOs for each piece of tensor metadata:
```
layout(set = 0, binding = 3) uniform PRECISION restrict OutSizes {
ivec4 data;
}
out_sizes;
layout(set = 0, binding = 4) uniform PRECISION restrict InSizes {
ivec4 data;
}
in_sizes;
layout(set = 0, binding = 5) uniform PRECISION restrict OtherSizes {
ivec4 data;
}
other_sizes;
layout(set = 0, binding = 6) uniform PRECISION restrict Alpha {
float data;
}
alpha;
```
Each UBO will be owned and maintained by the corresponding `vTensor` instance. To support a graph input resize, every tensor in the graph only needs to update their metadata UBOs via the `tensor.virtual_resize(new_sizes)` call. Shader dispatches in subsequent command buffer submissions will then see the updated metadata and execute as if the tensor were the updated sizes.
This changeset introduces a new shader library for the Vulkan graph runtime that enables dynamic shapes through this technique in favor of relying on the shader library from PyTorch Vulkan.
## Considerations
Technically, the UBO update technique can be applied to the shaders from PyTorch Vulkan as well. If that's the case, why introduce a new shader library for the graph runtime?
The primary motivation is code quality.
First, having `vTensor` supply UBOs for their own metadata greatly reduces the need to have operator specifc ad-hoc `Params` structs to organize arguments to write into a `api::UniformParamsBuffer`.
Constructing an `ExecuteNode` for binary operators is now
```
graph.execute_nodes().emplace_back(new ExecuteNode(
graph,
api::shader_registry().get_shader_info(kernel_name.str()),
global_size,
local_size,
{{out, api::MemoryAccessType::WRITE},
{{arg1, arg2}, api::MemoryAccessType::READ}},
{t_out.gpu_sizes_ubo(),
t_in1.gpu_sizes_ubo(),
t_in2.gpu_sizes_ubo(),
graph.create_params_buffer(alpha_val)}))
```
instead of
```
ArithmeticParams block{
get_size_as_ivec4(t_out),
get_size_as_ivec4(t_in1),
get_size_as_ivec4(t_in2),
alpha_val,
};
api::UniformParamsBuffer params(graph.context(), block);
graph.execute_nodes().emplace_back(new ExecuteNode(
graph,
shader,
global_size,
local_size,
{{out, api::MemoryAccessType::WRITE},
{{arg1, arg2}, api::MemoryAccessType::READ}},
std::move(params)));
```
Another consideration is that pytorch/pytorch#115948 which was landed fairly recently enables much more expressive shader templates through the use of Python code blocks in the GLSL template. This enables shader templates that can easily express variants for different data types, packing structures, etc. Introducing a new shader library provides the opportunity to rewrite the shaders in PyTorch Vulkan in a more generic and extensible way.
Differential Revision: [D54754545](https://our.internmc.facebook.com/intern/diff/D54754545/)
[ghstack-poisoned]
@facebook-github-bot

Copy link
Copy Markdown
Contributor

This pull request was exported from Phabricator. Differential Revision: D54754545

SS-JIA added a commit that referenced this pull request Mar 13, 2024
…hapes
Pull Request resolved: #2366
## Context
pytorch/pytorch#121598 introduces the ability to support dynamic shapes through tensor metadata updates.
The idea is fairly simple. Instead of shaders accepting a UBO with size data for all arguments:
```
layout(set = 0, binding = 2) uniform PRECISION restrict Block {
ivec4 output_sizes;
ivec4 other_sizes;
float alpha;
}
```
Shaders will accept separate UBOs for each piece of tensor metadata:
```
layout(set = 0, binding = 3) uniform PRECISION restrict OutSizes {
ivec4 data;
}
out_sizes;
layout(set = 0, binding = 4) uniform PRECISION restrict InSizes {
ivec4 data;
}
in_sizes;
layout(set = 0, binding = 5) uniform PRECISION restrict OtherSizes {
ivec4 data;
}
other_sizes;
layout(set = 0, binding = 6) uniform PRECISION restrict Alpha {
float data;
}
alpha;
```
Each UBO will be owned and maintained by the corresponding `vTensor` instance. To support a graph input resize, every tensor in the graph only needs to update their metadata UBOs via the `tensor.virtual_resize(new_sizes)` call. Shader dispatches in subsequent command buffer submissions will then see the updated metadata and execute as if the tensor were the updated sizes.
This changeset introduces a new shader library for the Vulkan graph runtime that enables dynamic shapes through this technique in favor of relying on the shader library from PyTorch Vulkan.
## Considerations
Technically, the UBO update technique can be applied to the shaders from PyTorch Vulkan as well. If that's the case, why introduce a new shader library for the graph runtime?
The primary motivation is code quality.
First, having `vTensor` supply UBOs for their own metadata greatly reduces the need to have operator specifc ad-hoc `Params` structs to organize arguments to write into a `api::UniformParamsBuffer`.
Constructing an `ExecuteNode` for binary operators is now
```
graph.execute_nodes().emplace_back(new ExecuteNode(
graph,
api::shader_registry().get_shader_info(kernel_name.str()),
global_size,
local_size,
{{out, api::MemoryAccessType::WRITE},
{{arg1, arg2}, api::MemoryAccessType::READ}},
{t_out.gpu_sizes_ubo(),
t_in1.gpu_sizes_ubo(),
t_in2.gpu_sizes_ubo(),
graph.create_params_buffer(alpha_val)}))
```
instead of
```
ArithmeticParams block{
get_size_as_ivec4(t_out),
get_size_as_ivec4(t_in1),
get_size_as_ivec4(t_in2),
alpha_val,
};
api::UniformParamsBuffer params(graph.context(), block);
graph.execute_nodes().emplace_back(new ExecuteNode(
graph,
shader,
global_size,
local_size,
{{out, api::MemoryAccessType::WRITE},
{{arg1, arg2}, api::MemoryAccessType::READ}},
std::move(params)));
```
Another consideration is that pytorch/pytorch#115948 which was landed fairly recently enables much more expressive shader templates through the use of Python code blocks in the GLSL template. This enables shader templates that can easily express variants for different data types, packing structures, etc. Introducing a new shader library provides the opportunity to rewrite the shaders in PyTorch Vulkan in a more generic and extensible way.
ghstack-source-id: 218421178
@exported-using-ghexport
Differential Revision: [D54754545](https://our.internmc.facebook.com/intern/diff/D54754545/)
@facebook-github-bot

Copy link
Copy Markdown
Contributor

This pull request has been merged in 835279e.

junpi3 pushed a commit that referenced this pull request Mar 13, 2024
junpi3 pushed a commit that referenced this pull request Mar 13, 2024
Missed this in #2366
Differential Revision: [D54880024](https://our.internmc.facebook.com/intern/diff/D54880024/)
ghstack-source-id: 218593311
Pull Request resolved: #2418
facebook-github-bot pushed a commit that referenced this pull request Mar 14, 2024
Summary:
Pull Request resolved: #2418
Missed this in #2366
ghstack-source-id: 218593311
exported-using-ghexport
bypass-github-export-checks
Reviewed By: SS-JIA
Differential Revision: D54880024
fbshipit-source-id: c4e19d8fefbb9d2fc4547ec2edc236060638da5e
@SS-JIA
SS-JIA deleted the gh/SS-JIA/11/head branch January 24, 2025 19:40
Sign up for freeto join this conversation on GitHub. Already have an account? Sign in to comment

Labels

CLA SignedThis label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed.fb-exportedMerged

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants

@SS-JIA@facebook-github-bot@junpi3