Skip to content

Export aten::relu6 as a single Clip op - #2987

Merged
justinchuby merged 2 commits into
mainfrom
copilot/improved-onnx-export-relu6
Aug 4, 2026
Merged

Export aten::relu6 as a single Clip op#2987
justinchuby merged 2 commits into
mainfrom
copilot/improved-onnx-export-relu6

Conversation

Copilot AI commented Aug 3, 2026

Copy link
Copy Markdown
Contributor

aten::relu6 was exported as a Relu + Min pair, which relies on downstream fusion to become a single clamp. ONNX has Clip, which expresses this directly.

Changes

  • onnxscript/function_libs/torch_lib/ops/nn.py: aten_relu6 now emits Clip(self, 0, 6) with bounds cast to the input dtype.
@torch_op("aten::relu6", trace_only=True)
def aten_relu6(self: TReal) -> TReal:
    """relu6(Tensor self) -> Tensor"""

    zero = op.CastLike(op.Constant(value_int=0), self)
    six = op.CastLike(op.Constant(value_int=6), self)
    return op.Clip(self, zero, six)

Benchmark

Measured on onnxruntime CPU EP (float32, 500 iterations, mean per-run latency) per the issue's request to confirm which form is faster:

shape Clip Relu + Min
(1, 64, 56, 56) 15.5 µs 67.3 µs
(64, 1024) 8.3 µs 18.8 µs

Clip is ~2–4x faster unfused, so the single-op form is preferable even before any graph optimization.

Copilot AI review requested due to automatic review settings August 3, 2026 19:35
Copilot AI removed the request for review from Copilot August 3, 2026 19:35
Copilot AI linked an issue Aug 3, 2026 that may be closed by this pull request
Co-authored-by: justinchuby <11205048+justinchuby@users.noreply.github.com>
Copilot AI requested review from Copilot and removed request for Copilot August 3, 2026 19:39
Copilot AI changed the title [WIP] Improve ONNX export for relu6 layer to use Clip Export aten::relu6 as a single Clip op Aug 3, 2026
Copilot AI requested a review from justinchuby August 3, 2026 19:39
@justinchuby
justinchuby marked this pull request as ready for review August 4, 2026 17:17
Copilot AI lite review requested due to automatic review settings August 4, 2026 17:17
@justinchuby
justinchuby enabled auto-merge (squash) August 4, 2026 17:17
@justinchuby justinchuby added the module: torchlib Related to the torch/aten function lib in development label Aug 4, 2026

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Pull request overview

This PR improves the ONNX export of aten::relu6 by emitting a single Clip op (with min=0, max=6) instead of the previous Relu + Min pair, aligning the export with ONNX’s native clamp representation and avoiding reliance on downstream fusion.

Changes:

  • Update aten_relu6 to return op.Clip(self, zero, six) rather than op.Min(op.Relu(self), six).
  • Ensure Clip bounds are cast to the input dtype via CastLike.

@github-project-automation github-project-automation Bot moved this from Todo to Done in ONNX Script Review Board Aug 4, 2026
@codecov

codecov Bot commented Aug 4, 2026

Copy link
Copy Markdown

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 72.63%. Comparing base (4cd022b) to head (33bda6b).
⚠️ Report is 1 commits behind head on main.
✅ All tests successful. No failed tests found.

Additional details and impacted files
@@           Coverage Diff           @@
##             main    #2987   +/-   ##
=======================================
  Coverage   72.63%   72.63%           
=======================================
  Files         265      265           
  Lines       32204    32205    +1     
  Branches     3041     3041           
=======================================
+ Hits        23391    23392    +1     
  Misses       7779     7779           
  Partials     1034     1034           

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

@justinchuby
justinchuby merged commit 457a260 into main Aug 4, 2026
31 of 35 checks passed
@justinchuby
justinchuby deleted the copilot/improved-onnx-export-relu6 branch August 4, 2026 17:37
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

module: torchlib Related to the torch/aten function lib in development

Projects

Development

Successfully merging this pull request may close these issues.

Improved ONNX export for relu6 layer

4 participants