Vulkan: implement instance norm by rewriting it to group norm - #22783
Open
msluszniak wants to merge 1 commit into
Open
Vulkan: implement instance norm by rewriting it to group norm#22783msluszniak wants to merge 1 commit into
msluszniak wants to merge 1 commit into
Conversation
`aten._native_batch_norm_legit.no_stats` has no Vulkan implementation, so every `nn.InstanceNorm2d` is a graph break. Architectures that normalize in each block pay for that repeatedly: the fast neural style TransformerNet at 640x640 splits into 16 delegate subgraphs with 126 nodes left on the CPU, and copies its activations out to the CPU and back 15 times per inference. `F.instance_norm` reshapes its input to [1, N*C, H, W] before lowering to this overload, so with a batch dim of 1 it is the same reduction as group norm with one group per channel. ReplaceInstanceNormPass rewrites it that way and the existing group norm kernels cover it, so no new shader is needed. A batch dim above 1 is a different reduction and is left alone. add_native_group_norm_node() prepacks weight and bias, so the pass only fires when both trace back to a constant. Identity `aten.repeat` (every repeat factor 1, which the existing shape equality check implies) now folds away as well: `F.instance_norm` emits `weight.repeat(b)`, which would otherwise leave the affine params as runtime tensors that group norm cannot prepack. Measured on a Mali-G76 (Galaxy S10+), style transfer at 640x640, fp16, per-exec slope interleaved over 3 rounds on a warm device: 2265.1 ms -> 1469.4 ms median, a 1.54x speedup, with non-overlapping ranges. The partition goes from 16 subgraphs to 1 and from 126 CPU nodes to 0. Output accuracy against the fp32 reference improves slightly as well, cosine 0.999993 -> 0.999994 and max abs diff 0.041 -> 0.031, because the fp16 round trips at the delegate boundaries are gone.
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/22783
Note: Links to docs will display an error until the docs builds have been completed.
|
pytorch-bot
Bot
added
the
module: vulkan
meta-cla
Bot
added
the
CLA Signed
This PR needs a
|
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
aten._native_batch_norm_legit.no_statshas no Vulkan implementation, so everynn.InstanceNorm2dis a graph break. Architectures that normalize in each block pay for that repeatedly: the fast neural style TransformerNet at 640x640 splits into 16 delegate subgraphs with 126 nodes left on the CPU, and copies its activations out to the CPU and back 15 times per inference.F.instance_normreshapes its input to[1, N*C, H, W]before lowering to this overload, so with a batch dim of 1 it is the same reduction as group norm with one group per channel.ReplaceInstanceNormPassrewrites it that way, which the existing group norm kernels already cover, so no new shader is needed. A batch dim above 1 is a different reduction and is left alone.add_native_group_norm_node()prepacks weight and bias, so the pass only fires when both trace back to a constant. Identityaten.repeatnow folds away as well:F.instance_normemitsweight.repeat(b), which would otherwise leave the affine params as runtime tensors that group norm cannot prepack.Measured on a Mali-G76 (Galaxy S10+), style transfer at 640x640, fp16, per-exec slope interleaved over 3 rounds on a warm device:
1.54x, non-overlapping. The partition goes from 16 subgraphs to 1 and from 126 CPU nodes to 0.
Accuracy against the fp32 reference improves slightly, cosine 0.999993 -> 0.999994 and max abs diff 0.041 -> 0.031, because the fp16 round trips at the delegate boundaries are gone.
cc @SS-JIA @manuelcandales @digantdesai @cbilgin