powi is slow

[voidentente]

Premise

I was wondering whether I could replace expressions like x * x with x.powi(2) and went ahead and compared the generated SPIR-V in both cases. Turns out, Rust-GPU fails to optimize even the most obvious opportunities.

Method

I inspected the SPIR-V using https://www.khronos.org/spirv/visualizer/ and the artifact was generated in release mode using rust-gpu 08f98a7 (a commit from Mon Jan 5 12:45:27 2026 +0100), which uses Rust nightly-2025-06-30.

Case: Manual Multiply

This Rust code:

#[spirv(fragment)]
pub fn fragment(out: &mut Vec4, #[spirv(frag_coord)] frag_coord: Vec4) {
    *out = Vec4::splat(frag_coord.x * frag_coord.x);
}

Becomes the following SPIR-V:

OpCapability Shader
OpCapability VulkanMemoryModel
OpMemoryModel Logical Vulkan
OpEntryPoint Fragment %1 "fragment" %2 %3
OpExecutionMode %1 OriginUpperLeft
OpDecorate %2 BuiltIn FragCoord
OpDecorate %3 Location 0
%6 = OpTypeFloat 32
%7 = OpTypeVector %6 4
%8 = OpTypePointer Output %7
%9 = OpTypePointer Input %7
%10 = OpTypeVoid
%11 = OpTypeFunction %10
%2 = OpVariable %9 Input
%3 = OpVariable %8 Output
%1 = OpFunction %10 None %11
%12 = OpLabel
%13 = OpLoad %7 %2
%14 = OpCompositeExtract %6 %13 0
%16 = OpFMul %6 %14 %14
%17 = OpCompositeConstruct %7 %16 %16 %16 %16
OpStore %3 %17
OpReturn
OpFunctionEnd

Case: Powi

The following Rust code:

#[spirv(fragment)]
pub fn fragment(out: &mut Vec4, #[spirv(frag_coord)] frag_coord: Vec4) {
    *out = Vec4::splat(frag_coord.x.powi(2));
}

Becomes the following SPIR-V:

OpCapability Shader
OpCapability VulkanMemoryModel
OpMemoryModel Logical Vulkan
OpEntryPoint Fragment %1 "fragment" %2 %3
OpExecutionMode %1 OriginUpperLeft
OpDecorate %2 BuiltIn FragCoord
OpDecorate %3 Location 0
%14 = OpTypeFloat 32
%15 = OpTypeVector %14 4
%16 = OpTypePointer Output %15
%17 = OpTypePointer Input %15
%18 = OpTypeVoid
%19 = OpTypeFunction %18
%2 = OpVariable %17 Input
%20 = OpTypeBool
%21 = OpTypeInt 32 1
%25 = OpTypeInt 32 0
%30 = OpConstant %25 0
%31 = OpConstant %25 1
%32 = OpConstant %21 1
%33 = OpUndef %14
%34 = OpUndef %25
%3 = OpVariable %16 Output
%127 = OpConstant %25 2
%1 = OpFunction %18 None %19
%35 = OpLabel
%36 = OpLoad %15 %2
%37 = OpCompositeExtract %14 %36 0
OpBranch %75
%75 = OpLabel
%76 = OpPhi %14 (%37 : %35) (%77 : %78)
%79 = OpPhi %25 (%127 : %35) (%80 : %78)
OpLoopMerge %81 %78 None
OpBranch %82
%82 = OpLabel
%83 = OpBitwiseAnd %25 %79 %31
%84 = OpIEqual %20 %83 %30
OpSelectionMerge %85 None
OpBranchConditional %84 %86 %87
%87 = OpLabel
OpBranch %85
%86 = OpLabel
%88 = OpFMul %14 %76 %76
%89 = OpShiftRightLogical %25 %79 %32
OpBranch %85
%85 = OpLabel
%77 = OpPhi %14 (%88 : %86) (%33 : %87)
%80 = OpPhi %25 (%89 : %86) (%34 : %87)
OpBranch %78
%78 = OpLabel
OpBranchConditional %84 %75 %81
%81 = OpLabel
%90 = OpIEqual %20 %79 %31
OpSelectionMerge %91 None
OpBranchConditional %90 %92 %93
%93 = OpLabel
OpBranch %94
%94 = OpLabel
%95 = OpPhi %14 (%76 : %93) (%96 : %97)
%98 = OpPhi %14 (%76 : %93) (%99 : %97)
%100 = OpPhi %25 (%79 : %93) (%101 : %97)
OpLoopMerge %102 %97 None
OpBranch %103
%103 = OpLabel
%104 = OpUGreaterThan %20 %100 %31
OpSelectionMerge %105 None
OpBranchConditional %104 %106 %107
%107 = OpLabel
OpBranch %105
%106 = OpLabel
%108 = OpShiftRightLogical %25 %100 %32
%109 = OpFMul %14 %95 %95
%110 = OpBitwiseAnd %25 %108 %31
%111 = OpIEqual %20 %110 %31
OpSelectionMerge %112 None
OpBranchConditional %111 %113 %114
%114 = OpLabel
OpBranch %112
%113 = OpLabel
%115 = OpFMul %14 %98 %109
OpBranch %112
%112 = OpLabel
%116 = OpPhi %14 (%115 : %113) (%98 : %114)
OpBranch %105
%105 = OpLabel
%96 = OpPhi %14 (%109 : %112) (%33 : %107)
%99 = OpPhi %14 (%116 : %112) (%33 : %107)
%101 = OpPhi %25 (%108 : %112) (%34 : %107)
OpBranch %97
%97 = OpLabel
OpBranchConditional %104 %94 %102
%102 = OpLabel
OpBranch %91
%92 = OpLabel
OpBranch %91
%91 = OpLabel
%117 = OpPhi %14 (%76 : %92) (%98 : %102)
%119 = OpCompositeConstruct %15 %117 %117 %117 %117
OpStore %3 %119
OpReturn
OpFunctionEnd

Proposal

Support "unrolling" x.powi(2) etcetera into multiplications, i.e. x * x, which are significantly faster.

[nazar-pc]

From my experience there are countless examples of trivial optimizations that are not happening, including seemingly completely non-functional DCE: #308

I'd love to see someone doing something about it, I bet my shaders would become a lot faster as the result.

[voidentente]

Concrete proposal:

• x * x is not completely semantically equivalent to x.powi(2) and may yield different results. However, since core::f32::math::powi(x, n) is documented to have unspecified precision, this isn't a breaking change. Side effects may vary. For users that want to forbid altering behavior, provide a flag to disable this optimization.
• Scan calls of core::f32::math::powi(x, n) where the n is known at compile time and is either 0, 1, 2, 3, or a power of two like 4, 8, 16, 32, 64, etc, and substitute with multiplications.
  • n = 4 may be substituted with x *= x; x *= x;
  • n = 8 may be substituted with x *= x; x *= x; x *= x;
  • n = 16 may be substituted with x *= x; x *= x; x *= x; x *= x;
  • etc
• The optimization should be enabled by default, as it is unlikely the user depends on the specific behavior of core::f32::math::powi(x, n).

[voidentente]

Actually, I just realized that only f32::powi generates this long SPIR-V.

• When using core::f32::math::powi, it imports Pow from GLSL.std.450.

OpCapability Shader
OpCapability VulkanMemoryModel
%1 = OpExtInstImport "GLSL.std.450"
OpMemoryModel Logical Vulkan
OpEntryPoint Fragment %2 "fragment" %3 %4
OpExecutionMode %2 OriginUpperLeft
OpDecorate %3 BuiltIn FragCoord
OpDecorate %4 Location 0
%8 = OpTypeFloat 32
%9 = OpTypeVector %8 4
%10 = OpTypePointer Input %9
%11 = OpTypePointer Output %9
%12 = OpTypeVoid
%13 = OpTypeFunction %12
%3 = OpVariable %10 Input
%4 = OpVariable %11 Output
%22 = OpConstant %8 2.0
%2 = OpFunction %12 None %13
%16 = OpLabel
%17 = OpLoad %9 %3
%18 = OpCompositeExtract %8 %17 0
%20 = OpExtInst %8 %1 Pow %18 %22
%21 = OpCompositeConstruct %9 %20 %20 %20 %20
OpStore %4 %21
OpReturn
OpFunctionEnd

• When using libm::powf, it imports Pow from GLSL.std.450.

OpCapability Shader
OpCapability VulkanMemoryModel
%1 = OpExtInstImport "GLSL.std.450"
OpMemoryModel Logical Vulkan
OpEntryPoint Fragment %2 "fragment" %3 %4
OpExecutionMode %2 OriginUpperLeft
OpDecorate %3 BuiltIn FragCoord
OpDecorate %4 Location 0
%7 = OpTypeFloat 32
%8 = OpTypeVector %7 4
%9 = OpTypePointer Input %8
%10 = OpTypePointer Output %8
%11 = OpTypeVoid
%12 = OpTypeFunction %11
%3 = OpVariable %9 Input
%13 = OpConstant %7 2.0
%4 = OpVariable %10 Output
%2 = OpFunction %11 None %12
%14 = OpLabel
%15 = OpLoad %8 %3
%16 = OpCompositeExtract %7 %15 0
%17 = OpExtInst %7 %1 Pow %16 %13
%18 = OpCompositeConstruct %8 %17 %17 %17 %17
OpStore %4 %18
OpReturn
OpFunctionEnd

• When using core::intrinsics::powif32, it imports Pow from GLSL.std.450:

OpCapability Shader
OpCapability VulkanMemoryModel
%1 = OpExtInstImport "GLSL.std.450"
OpMemoryModel Logical Vulkan
OpEntryPoint Fragment %2 "fragment" %3 %4
OpExecutionMode %2 OriginUpperLeft
OpDecorate %3 BuiltIn FragCoord
OpDecorate %4 Location 0
%7 = OpTypeFloat 32
%8 = OpTypeVector %7 4
%9 = OpTypePointer Input %8
%10 = OpTypePointer Output %8
%11 = OpTypeVoid
%12 = OpTypeFunction %11
%3 = OpVariable %9 Input
%4 = OpVariable %10 Output
%21 = OpConstant %7 2.0
%2 = OpFunction %11 None %12
%15 = OpLabel
%16 = OpLoad %8 %3
%17 = OpCompositeExtract %7 %16 0
%19 = OpExtInst %7 %1 Pow %17 %21
%20 = OpCompositeConstruct %8 %19 %19 %19 %19
OpStore %4 %20
OpReturn
OpFunctionEnd

ONLY the inherent method f32::powi, which should internally call core::f32::math::powi, which in turn internally calls core::intrinsics::powif32, generates the "Case: Powi" I sent above.

[voidentente]

Technically, this is two separate issues now.

• One for the optimization
• One for why f32::powi behaves aberrantly

[LegNeato]

I think f32::powi comes from num_traits, right? It doesn't look like we are overriding that impl.