[Mlir-commits] [mlir] 271c8b7 - [mlir][linalg] Add sin, cos, tan to elementwise operations (#200950)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Tue Jun 9 04:52:52 PDT 2026
Author: Vinit Deodhar
Date: 2026-06-09T12:52:47+01:00
New Revision: 271c8b75cde53d78da5ce971ced8694cca149749
URL: https://github.com/llvm/llvm-project/commit/271c8b75cde53d78da5ce971ced8694cca149749
DIFF: https://github.com/llvm/llvm-project/commit/271c8b75cde53d78da5ce971ced8694cca149749.diff
LOG: [mlir][linalg] Add sin, cos, tan to elementwise operations (#200950)
Add sin, cos, and tan as UnaryFn entries in the linalg dialect, enabling
their use via linalg.elementwise, named ops (linalg.sin, linalg.cos,
linalg.tan), and specialization from linalg.generic.
---------
Co-authored-by: Vinit Deodhar <vinitdeodhar at users.noreply.github.com>
Co-authored-by: Claude Opus 4.6 <noreply at anthropic.com>
Co-authored-by: Vinit Deodhar <vdeodhar at ah-vdeodhar-l.dhcp.mathworks.com>
Added:
Modified:
mlir/include/mlir/Dialect/Linalg/IR/LinalgEnums.td
mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp
mlir/lib/Dialect/Linalg/Transforms/Specialize.cpp
mlir/test/Dialect/Linalg/roundtrip-morphism-linalg-category-ops.mlir
mlir/test/Dialect/Linalg/specialize-generic-ops.mlir
Removed:
################################################################################
diff --git a/mlir/include/mlir/Dialect/Linalg/IR/LinalgEnums.td b/mlir/include/mlir/Dialect/Linalg/IR/LinalgEnums.td
index 1109db973f522..f6b5edd326c85 100644
--- a/mlir/include/mlir/Dialect/Linalg/IR/LinalgEnums.td
+++ b/mlir/include/mlir/Dialect/Linalg/IR/LinalgEnums.td
@@ -29,7 +29,10 @@ def UnaryFn : I32EnumAttr<"UnaryFn", "", [
I32EnumAttrCase<"rsqrt", 9>,
I32EnumAttrCase<"square", 10>,
I32EnumAttrCase<"tanh", 11>,
- I32EnumAttrCase<"erf", 12>
+ I32EnumAttrCase<"erf", 12>,
+ I32EnumAttrCase<"sin", 13>,
+ I32EnumAttrCase<"cos", 14>,
+ I32EnumAttrCase<"tan", 15>
]> {
let genSpecializedAttr = 0;
let cppNamespace = "::mlir::linalg";
diff --git a/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp b/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp
index de7f4d1610bd0..6b61564eb5fcc 100644
--- a/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp
+++ b/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp
@@ -503,6 +503,12 @@ class RegionBuilderHelper {
return math::TanhOp::create(builder, arg.getLoc(), arg);
case UnaryFn::erf:
return math::ErfOp::create(builder, arg.getLoc(), arg);
+ case UnaryFn::sin:
+ return math::SinOp::create(builder, arg.getLoc(), arg);
+ case UnaryFn::cos:
+ return math::CosOp::create(builder, arg.getLoc(), arg);
+ case UnaryFn::tan:
+ return math::TanOp::create(builder, arg.getLoc(), arg);
}
if (emitError) {
emitError() << "unsupported unary function";
diff --git a/mlir/lib/Dialect/Linalg/Transforms/Specialize.cpp b/mlir/lib/Dialect/Linalg/Transforms/Specialize.cpp
index 2ad2774ff1a0c..fd95cfd5d3403 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/Specialize.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/Specialize.cpp
@@ -224,6 +224,17 @@ static FailureOr<LinalgOp> specializeLinalgElementwise(RewriterBase &rewriter,
if (isa<math::ErfOp>(op))
return replaceOp(ErfOp{}, ElementwiseKind::erf);
+ // sin, cos, tan only have the category (elementwise) form, but no
+ // linalg.* named op equivalent.
+ if (emitCategoryOp) {
+ if (isa<math::SinOp>(op))
+ return replaceOp(nullptr, ElementwiseKind::sin);
+ if (isa<math::CosOp>(op))
+ return replaceOp(nullptr, ElementwiseKind::cos);
+ if (isa<math::TanOp>(op))
+ return replaceOp(nullptr, ElementwiseKind::tan);
+ }
+
// At this point, we exhaustively checked the available unary named ops. The
// 1-input generic op might be representable as a `linalg.elementwise` that
// broadcasts a scalar operand. But if we can't emit the category op or
diff --git a/mlir/test/Dialect/Linalg/roundtrip-morphism-linalg-category-ops.mlir b/mlir/test/Dialect/Linalg/roundtrip-morphism-linalg-category-ops.mlir
index f63b795108cdb..5302072777b9f 100644
--- a/mlir/test/Dialect/Linalg/roundtrip-morphism-linalg-category-ops.mlir
+++ b/mlir/test/Dialect/Linalg/roundtrip-morphism-linalg-category-ops.mlir
@@ -32,6 +32,12 @@ func.func @unary_ops(%A: memref<7x14x21xf32>, %Out: memref<7x14x21xf32>) {
ins(%A : memref<7x14x21xf32>) outs(%Out : memref<7x14x21xf32>)
linalg.elementwise kind=#linalg.elementwise_kind<erf>
ins(%A : memref<7x14x21xf32>) outs(%Out : memref<7x14x21xf32>)
+ linalg.elementwise kind=#linalg.elementwise_kind<sin>
+ ins(%A : memref<7x14x21xf32>) outs(%Out : memref<7x14x21xf32>)
+ linalg.elementwise kind=#linalg.elementwise_kind<cos>
+ ins(%A : memref<7x14x21xf32>) outs(%Out : memref<7x14x21xf32>)
+ linalg.elementwise kind=#linalg.elementwise_kind<tan>
+ ins(%A : memref<7x14x21xf32>) outs(%Out : memref<7x14x21xf32>)
return
}
@@ -77,6 +83,15 @@ func.func @unary_ops(%A: memref<7x14x21xf32>, %Out: memref<7x14x21xf32>) {
// CHECK: linalg.elementwise kind=#linalg.elementwise_kind<erf>
// CHECK-SAME: ins(%[[A]] : memref<7x14x21xf32>)
// CHECK-SAME: outs(%[[OUT]] : memref<7x14x21xf32>)
+// CHECK: linalg.elementwise kind=#linalg.elementwise_kind<sin>
+// CHECK-SAME: ins(%[[A]] : memref<7x14x21xf32>)
+// CHECK-SAME: outs(%[[OUT]] : memref<7x14x21xf32>)
+// CHECK: linalg.elementwise kind=#linalg.elementwise_kind<cos>
+// CHECK-SAME: ins(%[[A]] : memref<7x14x21xf32>)
+// CHECK-SAME: outs(%[[OUT]] : memref<7x14x21xf32>)
+// CHECK: linalg.elementwise kind=#linalg.elementwise_kind<tan>
+// CHECK-SAME: ins(%[[A]] : memref<7x14x21xf32>)
+// CHECK-SAME: outs(%[[OUT]] : memref<7x14x21xf32>)
// -----
diff --git a/mlir/test/Dialect/Linalg/specialize-generic-ops.mlir b/mlir/test/Dialect/Linalg/specialize-generic-ops.mlir
index 3d6c2962731c9..46193b4d7e992 100644
--- a/mlir/test/Dialect/Linalg/specialize-generic-ops.mlir
+++ b/mlir/test/Dialect/Linalg/specialize-generic-ops.mlir
@@ -100,7 +100,28 @@ func.func @unary_ops(%A: tensor<?x?x?xf32>, %Out: tensor<?x?x?xf32>) -> tensor<?
%v = math.erf %in : f32
linalg.yield %v : f32
} -> tensor<?x?x?xf32>
- return %12 : tensor<?x?x?xf32>
+ %13 = linalg.generic
+ {indexing_maps = [#umap, #umap], iterator_types = ["parallel", "parallel","parallel"]}
+ ins(%12 : tensor<?x?x?xf32>) outs(%Out : tensor<?x?x?xf32>) {
+ ^bb0(%in: f32, %out: f32):
+ %v = math.sin %in : f32
+ linalg.yield %v : f32
+ } -> tensor<?x?x?xf32>
+ %14 = linalg.generic
+ {indexing_maps = [#umap, #umap], iterator_types = ["parallel", "parallel","parallel"]}
+ ins(%13 : tensor<?x?x?xf32>) outs(%Out : tensor<?x?x?xf32>) {
+ ^bb0(%in: f32, %out: f32):
+ %v = math.cos %in : f32
+ linalg.yield %v : f32
+ } -> tensor<?x?x?xf32>
+ %15 = linalg.generic
+ {indexing_maps = [#umap, #umap], iterator_types = ["parallel", "parallel","parallel"]}
+ ins(%14 : tensor<?x?x?xf32>) outs(%Out : tensor<?x?x?xf32>) {
+ ^bb0(%in: f32, %out: f32):
+ %v = math.tan %in : f32
+ linalg.yield %v : f32
+ } -> tensor<?x?x?xf32>
+ return %15 : tensor<?x?x?xf32>
}
// ALL-LABEL: unary_ops
@@ -186,6 +207,15 @@ func.func @unary_ops(%A: tensor<?x?x?xf32>, %Out: tensor<?x?x?xf32>) -> tensor<?
// CATEGORY: %[[RES12:.+]] = linalg.elementwise kind=#linalg.elementwise_kind<erf>
// CATEGORY-SAME: ins(%[[RES11]] : tensor<?x?x?xf32>)
// CATEGORY-SAME: outs(%[[OUT]] : tensor<?x?x?xf32>) -> tensor<?x?x?xf32>
+// CATEGORY: %[[RES13:.+]] = linalg.elementwise kind=#linalg.elementwise_kind<sin>
+// CATEGORY-SAME: ins(%[[RES12]] : tensor<?x?x?xf32>)
+// CATEGORY-SAME: outs(%[[OUT]] : tensor<?x?x?xf32>) -> tensor<?x?x?xf32>
+// CATEGORY: %[[RES14:.+]] = linalg.elementwise kind=#linalg.elementwise_kind<cos>
+// CATEGORY-SAME: ins(%[[RES13]] : tensor<?x?x?xf32>)
+// CATEGORY-SAME: outs(%[[OUT]] : tensor<?x?x?xf32>) -> tensor<?x?x?xf32>
+// CATEGORY: %[[RES15:.+]] = linalg.elementwise kind=#linalg.elementwise_kind<tan>
+// CATEGORY-SAME: ins(%[[RES14]] : tensor<?x?x?xf32>)
+// CATEGORY-SAME: outs(%[[OUT]] : tensor<?x?x?xf32>) -> tensor<?x?x?xf32>
// -----
More information about the Mlir-commits
mailing list