[MLIR] Add sincos fusion pass #161413

ashermancinelli · 2025-09-30T18:06:32Z

We see performance improvements from using sincos to reuse calculations in hot loops that compute sin() and cos() of the same operand. Add a pass to identify sin() and cos() calls in the same block with the same operand and fast-math flags, and fuse them into a sincos op.

Note: The fusion pass checks for exact equality and not structural equality (op.getOperand() == operand); I believe the latter would do CSE-type substitution, while this pass ought to only fuse sin and cos calls if their operands are exactly the same SSA values. What do you think?

Follow-up to:

We see performance improvements from using sincos to reuse calculations in hot loops that compute sin() and cos() on the same operand. Add a pass to identify sin() and cos() calls in the same block with the same operand and fast-math flags, and fuse them into a sincos op. Follow-up to: * llvm#160561 * llvm#160772

llvmbot · 2025-09-30T18:06:55Z

@llvm/pr-subscribers-mlir-math

@llvm/pr-subscribers-mlir

Author: Asher Mancinelli (ashermancinelli)

Changes

We see performance improvements from using sincos to reuse calculations
in hot loops that compute sin() and cos() of the same operand.
Add a pass to identify sin() and cos() calls in the same block with the
same operand and fast-math flags, and fuse them into a sincos op.

Note: The fusion pass checks for exact equality and not structural equality (op.getOperand() == operand); I believe the latter would do CSE-type substitution, while this pass ought to only fuse sin and cos calls if their operands are exactly the same SSA values. What do you think?

Follow-up to:

#160561
#160772

Full diff: https://github.com/llvm/llvm-project/pull/161413.diff

4 Files Affected:

(modified) mlir/include/mlir/Dialect/Math/Transforms/Passes.td (+8)
(modified) mlir/lib/Dialect/Math/Transforms/CMakeLists.txt (+1)
(added) mlir/lib/Dialect/Math/Transforms/SincosFusion.cpp (+77)
(added) mlir/test/Dialect/Math/sincos-fusion.mlir (+64)

diff --git a/mlir/include/mlir/Dialect/Math/Transforms/Passes.td b/mlir/include/mlir/Dialect/Math/Transforms/Passes.td
index 4d415aeac8f58..48346abd84285 100644
--- a/mlir/include/mlir/Dialect/Math/Transforms/Passes.td
+++ b/mlir/include/mlir/Dialect/Math/Transforms/Passes.td
@@ -64,4 +64,12 @@ def MathExpandOpsPass : Pass<"math-expand-ops"> {
   ];
 }
 
+def MathSincosFusionPass : Pass<"math-sincos-fusion"> {
+  let summary = "Fuse sin and cos operations.";
+  let description = [{
+    Fuse sin and cos operations into a sincos operation.
+  }];
+  let dependentDialects = ["math::MathDialect"];
+}
+
 #endif // MLIR_DIALECT_MATH_TRANSFORMS_PASSES
diff --git a/mlir/lib/Dialect/Math/Transforms/CMakeLists.txt b/mlir/lib/Dialect/Math/Transforms/CMakeLists.txt
index ff62b515533c3..8899c3a1d1a42 100644
--- a/mlir/lib/Dialect/Math/Transforms/CMakeLists.txt
+++ b/mlir/lib/Dialect/Math/Transforms/CMakeLists.txt
@@ -3,6 +3,7 @@ add_mlir_dialect_library(MLIRMathTransforms
   ExpandOps.cpp
   ExtendToSupportedTypes.cpp
   PolynomialApproximation.cpp
+  SincosFusion.cpp
   UpliftToFMA.cpp
 
   ADDITIONAL_HEADER_DIRS
diff --git a/mlir/lib/Dialect/Math/Transforms/SincosFusion.cpp b/mlir/lib/Dialect/Math/Transforms/SincosFusion.cpp
new file mode 100644
index 0000000000000..a373cf70b5541
--- /dev/null
+++ b/mlir/lib/Dialect/Math/Transforms/SincosFusion.cpp
@@ -0,0 +1,77 @@
+//===- SincosFusion.cpp - Fuse sin/cos into sincos -----------------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Dialect/Math/IR/Math.h"
+#include "mlir/Dialect/Math/Transforms/Passes.h"
+#include "mlir/IR/PatternMatch.h"
+#include "mlir/Pass/Pass.h"
+#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
+
+using namespace mlir;
+using namespace mlir::math;
+
+namespace {
+
+/// Fuse a math.sin and math.cos in the same block that use the same operand and
+/// have identical fastmath flags into a single math.sincos.
+struct SincosFusionPattern : OpRewritePattern<math::SinOp> {
+  using OpRewritePattern<math::SinOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(math::SinOp sinOp,
+                                PatternRewriter &rewriter) const override {
+    Value operand = sinOp.getOperand();
+    mlir::arith::FastMathFlags sinFastMathFlags = sinOp.getFastmath();
+
+    math::CosOp cosOp = nullptr;
+    sinOp->getBlock()->walk([&](math::CosOp op) {
+      if (op.getOperand() == operand && op.getFastmath() == sinFastMathFlags) {
+        cosOp = op;
+        return WalkResult::interrupt();
+      }
+      return WalkResult::advance();
+    });
+
+    if (!cosOp)
+      return failure();
+
+    Type elemType = sinOp.getType();
+    auto sincos = rewriter.create<math::SincosOp>(
+        sinOp.getLoc(), TypeRange{elemType, elemType}, operand,
+        sinOp.getFastmathAttr());
+
+    rewriter.replaceOp(sinOp, sincos.getSin());
+    rewriter.replaceOp(cosOp, sincos.getCos());
+    return success();
+  }
+};
+
+} // namespace
+
+namespace mlir::math {
+#define GEN_PASS_DEF_MATHSINCOSFUSIONPASS
+#include "mlir/Dialect/Math/Transforms/Passes.h.inc"
+} // namespace mlir::math
+
+namespace {
+
+struct MathSincosFusionPass final
+    : math::impl::MathSincosFusionPassBase<MathSincosFusionPass> {
+  using MathSincosFusionPassBase::MathSincosFusionPassBase;
+
+  void runOnOperation() override {
+    RewritePatternSet patterns(&getContext());
+    patterns.add<SincosFusionPattern>(&getContext());
+
+    GreedyRewriteConfig config;
+    if (failed(
+            applyPatternsGreedily(getOperation(), std::move(patterns), config)))
+      return signalPassFailure();
+  }
+};
+
+} // namespace
diff --git a/mlir/test/Dialect/Math/sincos-fusion.mlir b/mlir/test/Dialect/Math/sincos-fusion.mlir
new file mode 100644
index 0000000000000..9abf576c858bd
--- /dev/null
+++ b/mlir/test/Dialect/Math/sincos-fusion.mlir
@@ -0,0 +1,64 @@
+// RUN: mlir-opt -math-sincos-fusion %s | FileCheck %s
+
+// CHECK-LABEL:   func.func @sincos_fusion(
+// CHECK-SAME:      %[[ARG0:.*]]: f32,
+// CHECK-SAME:      %[[ARG1:.*]]: f32) -> (f32, f32, f32, f32) {
+// CHECK:           %[[VAL_0:.*]], %[[VAL_1:.*]] = math.sincos %[[ARG0]] : f32
+// CHECK:           %[[VAL_2:.*]], %[[VAL_3:.*]] = math.sincos %[[ARG1]] : f32
+// CHECK:           return %[[VAL_0]], %[[VAL_1]], %[[VAL_3]], %[[VAL_2]] : f32, f32, f32, f32
+// CHECK:         }
+func.func @sincos_fusion(%arg0 : f32, %arg1 : f32) -> (f32, f32, f32, f32) {
+    %0 = math.sin %arg0 : f32
+    %1 = math.cos %arg0 : f32
+
+    %2 = math.cos %arg1 : f32
+    %3 = math.sin %arg1 : f32
+
+    func.return %0, %1, %2, %3 : f32, f32, f32, f32
+}
+
+// CHECK-LABEL:   func.func @sincos_fusion_no_match_fmf(
+// CHECK-SAME:      %[[ARG0:.*]]: f32) -> (f32, f32) {
+// CHECK:           %[[VAL_0:.*]] = math.sin %[[ARG0]] fastmath<contract> : f32
+// CHECK:           %[[VAL_1:.*]] = math.cos %[[ARG0]] : f32
+// CHECK:           return %[[VAL_0]], %[[VAL_1]] : f32, f32
+// CHECK:         }
+func.func @sincos_fusion_no_match_fmf(%arg0 : f32) -> (f32, f32) {
+    %0 = math.sin %arg0 fastmath<contract> : f32
+    %1 = math.cos %arg0 : f32
+    func.return %0, %1 : f32, f32
+}
+
+// CHECK-LABEL:   func.func @sincos_no_fusion_different_block(
+// CHECK-SAME:      %[[ARG0:.*]]: f32,
+// CHECK-SAME:      %[[ARG1:.*]]: i1) -> f32 {
+// CHECK:           %[[VAL_0:.*]] = scf.if %[[ARG1]] -> (f32) {
+// CHECK:             %[[VAL_1:.*]] = math.sin %[[ARG0]] : f32
+// CHECK:             scf.yield %[[VAL_1]] : f32
+// CHECK:           } else {
+// CHECK:             %[[VAL_2:.*]] = math.cos %[[ARG0]] : f32
+// CHECK:             scf.yield %[[VAL_2]] : f32
+// CHECK:           }
+// CHECK:           return %[[VAL_0]] : f32
+// CHECK:         }
+func.func @sincos_no_fusion_different_block(%arg0 : f32, %flag : i1) -> f32 {
+  %0 = scf.if %flag -> f32 {
+    %s = math.sin %arg0 : f32
+    scf.yield %s : f32
+  } else {
+    %c = math.cos %arg0 : f32
+    scf.yield %c : f32
+  }
+  func.return %0 : f32
+}
+
+// CHECK-LABEL:   func.func @sincos_fusion_preserve_fastmath(
+// CHECK-SAME:      %[[ARG0:.*]]: f32) -> (f32, f32) {
+// CHECK:           %[[VAL_0:.*]], %[[VAL_1:.*]] = math.sincos %[[ARG0]] fastmath<contract> : f32
+// CHECK:           return %[[VAL_0]], %[[VAL_1]] : f32, f32
+// CHECK:         }
+func.func @sincos_fusion_preserve_fastmath(%arg0 : f32) -> (f32, f32) {
+    %0 = math.sin %arg0 fastmath<contract> : f32
+    %1 = math.cos %arg0 fastmath<contract> : f32
+    func.return %0, %1 : f32, f32
+}

mlir/lib/Dialect/Math/Transforms/SincosFusion.cpp

clementval

LGTM

vzakhari · 2025-09-30T18:39:05Z

mlir/lib/Dialect/Math/Transforms/SincosFusion.cpp

+        sinOp.getLoc(), TypeRange{elemType, elemType}, operand,
+        sinOp.getFastmathAttr());
+
+    rewriter.replaceOp(sinOp, sincos.getSin());


Please try an example where cos precedes sin, and there is a use of cos before sin.

The insertion point should be the dominating operation. Thanks for pointing that out!

vzakhari

Thank you, Asher!

vzakhari · 2025-09-30T19:50:18Z

while this pass ought to only fuse sin and cos calls if their operands are exactly the same SSA values.

I think this is reasonable for this pass. It is the CSE's job to optimize other cases.

mlir/lib/Dialect/Math/Transforms/SincosFusion.cpp

ashermancinelli · 2025-09-30T20:31:52Z

@Hardcode84 thank you for the review!

…sincos-fusion-pass

We see performance improvements from using sincos to reuse calculations in hot loops that compute sin() and cos() of the same operand. Add a pass to identify sin() and cos() calls in the same block with the same operand and fast-math flags, and fuse them into a sincos op. Follow-up to: * llvm#160561 * llvm#160772

ashermancinelli requested review from clementval and vzakhari September 30, 2025 18:06

ashermancinelli self-assigned this Sep 30, 2025

ashermancinelli added mlir mlir:math labels Sep 30, 2025

clementval reviewed Sep 30, 2025

View reviewed changes

mlir/lib/Dialect/Math/Transforms/SincosFusion.cpp Outdated Show resolved Hide resolved

clementval approved these changes Sep 30, 2025

View reviewed changes

vzakhari reviewed Sep 30, 2025

View reviewed changes

Insertion point should be the dominant op

3e072be

vzakhari approved these changes Sep 30, 2025

View reviewed changes

ashermancinelli requested review from Hardcode84, ftynse and joker-eph September 30, 2025 20:04

Hardcode84 approved these changes Sep 30, 2025

View reviewed changes

mlir/lib/Dialect/Math/Transforms/SincosFusion.cpp Outdated Show resolved Hide resolved

mlir/lib/Dialect/Math/Transforms/SincosFusion.cpp Outdated Show resolved Hide resolved

Apply review suggestions

3a17b7b

Merge branch 'main' of https://github.com/llvm/llvm-project into ajm/…

48dd58b

…sincos-fusion-pass

ashermancinelli merged commit ca1ff80 into llvm:main Oct 1, 2025
9 checks passed

aemerson mentioned this pull request Oct 16, 2025

[clang-reorder-fields] Check for flexible array member (#160262) #163697

Closed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

[MLIR] Add sincos fusion pass #161413

[MLIR] Add sincos fusion pass #161413

Uh oh!

ashermancinelli commented Sep 30, 2025 •

edited

Loading

Uh oh!

llvmbot commented Sep 30, 2025 •

edited

Loading

Uh oh!

Uh oh!

clementval left a comment

Uh oh!

vzakhari Sep 30, 2025

Uh oh!

ashermancinelli Sep 30, 2025

Uh oh!

vzakhari left a comment

Uh oh!

vzakhari commented Sep 30, 2025

Uh oh!

Uh oh!

Uh oh!

ashermancinelli commented Sep 30, 2025

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

[MLIR] Add sincos fusion pass #161413

[MLIR] Add sincos fusion pass #161413

Uh oh!

Conversation

ashermancinelli commented Sep 30, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

llvmbot commented Sep 30, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Uh oh!

clementval left a comment

Choose a reason for hiding this comment

Uh oh!

vzakhari Sep 30, 2025

Choose a reason for hiding this comment

Uh oh!

ashermancinelli Sep 30, 2025

Choose a reason for hiding this comment

Uh oh!

vzakhari left a comment

Choose a reason for hiding this comment

Uh oh!

vzakhari commented Sep 30, 2025

Uh oh!

Uh oh!

Uh oh!

ashermancinelli commented Sep 30, 2025

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

ashermancinelli commented Sep 30, 2025 •

edited

Loading

llvmbot commented Sep 30, 2025 •

edited

Loading