From the idea to the upstream implementation
The animation shows the idea. This guide explains the compiler concepts, the decisions inside my patch, and how to study a similar open-source problem yourself.
1. Why the same math can produce different programs
PyTorch lets developers describe calculations. Its Inductor compiler decides how to turn those calculations into executable work. A graph is the compiler's representation of the operations and the values flowing between them. Changing that representation can reveal an optimization without changing the intended calculation.
- Batch matrix multiply
- Multiply several pairs of matrices together. Think of processing a stack of small spreadsheets, one pair at a time.
- Broadcast bias
- Add a smaller array of adjustments across a larger result. A row of adjustments can be reused across every row of a matrix.
- Pointwise operation
- Apply a calculation independently to each element. ReLU replaces negative values with zero; GELU is another activation used in neural networks.
- Fusion
- Combine compatible work into one generated kernel, which is a unit of GPU execution. This can avoid intermediate memory traffic and separate launches.
The operation baddbmm combines a batch matrix multiplication with an addition. The original issue showed why separating out a broadcast bias could help when an activation followed it. Optimizing an operation in isolation could miss the better arrangement across its boundary.
2. What I changed in the compiler
I extended an existing optimization for addmm to eligible baddbmm operations. The contribution reused Inductor's graph-pattern machinery and normal pointwise fusion, rather than introducing a separate hand-written GPU kernel.
Before: activation(baddbmm(bias, A, B, alpha=α, beta=β))
After: activation(β * bias + α * bmm(A, B))This is a simplified view of the rewrite. The multiplication stays a matrix multiplication, but the bias addition becomes a visible graph node. The compiler can then consider combining it with the activation. More nodes in the intermediate graph can therefore lead to less separate work in the generated program.
A small worked example. Suppose one matrix product produces [[1, 4], [3, 8]], the broadcast bias is [-10, 2], alpha is 2, and beta is 0.5.
Scaling the product gives [[2, 8], [6, 16]]. Adding the scaled bias [-5, 1] to each row gives [[-3, 9], [1, 17]]. ReLU then gives [[0, 9], [1, 17]].
In the merged implementation, start with should_prefer_unfused_baddbmm, then read unfuse_bias_baddbmm_to_pointwise. The first decides eligibility; the second constructs the replacement graph. Separating those responsibilities makes the optimization easier to reason about.
3. Why this is more than an algebraic substitution
Production compiler work has to respect the surrounding graph and hardware behavior. My patch includes several boundaries on when the rewrite applies:
- GPU and bias shape: the eligibility check requires GPU execution and an input recognized as bias-like or broadcast.
- Consumers: the result's users must be pointwise operations. The standalone case is explicitly covered by a test that keeps
baddbmm. - Scaling: the replacement retains both alpha and beta. Assuming they are always one would change valid programs.
- Precision: the rewrite follows the existing half-dtype configuration guard, including its XPU-specific exception. A useful transformation for one precision or device cannot simply be applied everywhere.
These decisions are visible in the guard and replacement functions. The engineering challenge is selecting a useful transformation while preserving the cases where existing behavior should remain.
4. How to tell whether a compiler change actually worked
A correct output alone cannot show that an optimization fired. The new pattern-matcher tests check numerical results and inspect the generated code.
Look for test_unfuse_broadcast_bias_baddbmm and test_unfuse_broadcast_bias_baddbmm_alpha_beta. Together, they show how a small regression test can protect both the intended transformation and behavior that should stay unchanged.
| Case | Evidence |
|---|---|
| Standalone operation | The output matches PyTorch and the generated code retains the external baddbmm call. |
| GELU consumer | The output matches the reference; the code uses bmm in place of the external baddbmm call. |
| Explicitly expanded bias | A stride-zero expanded input is checked, so broadcasting behavior is covered. |
| Non-default scaling | A ReLU case exercises alpha 0.8 and beta 0.2, with output and code checks. |
The PR discussion also records environment setup problems, CI iteration, and benchmark results. Later measurements include near-parity cases and both faster and slower decomposed examples. The defensible outcome is a merged, guarded optimization opportunity; a universal speedup would require evidence the record does not provide.
For a newcomer, keep three questions separate: Does the answer match? Did the intended code change happen? Does that change help this workload? Each needs its own check.
5. A route into this kind of open-source work
- Start with a reproducible issue. Read the input shapes, device, precision, expected behavior, and measurements in issue #187093. Reduce the problem to one operation and one consumer.
- Find the nearest existing solution. Compare the addmm pattern with my baddbmm extension in the same file. List what can be reused and what changes for batched inputs.
- Read a focused test before the whole compiler. Trace one test from its inputs to its reference output and generated-code assertions.
- Set up the project's development environment. Use PyTorch's contributing guide. Confirm your tests import the intended build; an unrelated installed wheel can make a source change appear ineffective.
- Add a negative case. Ask when your proposed optimization should not run. Then add the smallest regression test that captures that boundary.
Try it yourself: use the worked example above to calculate the answer before looking at the result. Then read the alpha/beta test and identify the assertion that checks correctness and the assertion that checks the rewrite.
6. What this contribution demonstrates
I contributed at the boundary between mathematical equivalence, graph representation, and generated execution. The work combines pattern matching, broadcast semantics, device and precision guards, and regression tests inside PyTorch's compiler.
The upstream commit records the resolved PR and approval by Jason Ansel. For someone learning open source, the useful model is a focused change that fits the existing architecture and carries evidence for its behavior.