[Mlir-commits] [mlir] [mlir][linalg] Migrate elementwise named ops from OpDSL to tablegen multiclass (PR #216976)
Javed Absar
llvmlistbot at llvm.org
Wed Aug 19 06:14:33 PDT 2026
https://github.com/javedabsar1 updated https://github.com/llvm/llvm-project/pull/216976
>From bc5e6a62767c25dc1135562e3063e6eec36b4024 Mon Sep 17 00:00:00 2001
From: mabsar <mabsar at qti.qualcommm.com>
Date: Mon, 17 Aug 2026 03:48:20 -0700
Subject: [PATCH 1/2] [mlir][linalg] Add tablegen based linalg.named ops.
Signed-off-by: mabsar <mabsar at qti.qualcommm.com>
---
mlir/include/mlir/Dialect/Linalg/IR/Linalg.h | 22 +
.../mlir/Dialect/Linalg/IR/LinalgInterfaces.h | 3 +
.../Dialect/Linalg/IR/LinalgInterfaces.td | 22 +
.../Linalg/IR/LinalgNamedStructuredOps.yaml | 907 ------------------
.../Dialect/Linalg/IR/LinalgStructuredOps.td | 179 +++-
mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp | 37 +
.../linalg/opdsl/ops/core_named_ops.py | 331 -------
7 files changed, 262 insertions(+), 1239 deletions(-)
diff --git a/mlir/include/mlir/Dialect/Linalg/IR/Linalg.h b/mlir/include/mlir/Dialect/Linalg/IR/Linalg.h
index 9de6d8fd50983..2d2566c4ac072 100644
--- a/mlir/include/mlir/Dialect/Linalg/IR/Linalg.h
+++ b/mlir/include/mlir/Dialect/Linalg/IR/Linalg.h
@@ -131,6 +131,28 @@ std::pair<int64_t, int64_t> getFmrFromWinogradConv2DFmr(WinogradConv2DFmr fmr);
#include "mlir/Dialect/Linalg/IR/LinalgInterfaces.h"
+//===----------------------------------------------------------------------===//
+// Shared utilities for named elementwise ops
+//===----------------------------------------------------------------------===//
+
+namespace mlir::linalg {
+
+/// Builds the body region for a named elementwise op based on the given kind.
+/// Dispatches actual building to one of build[UnaryFn,BinaryFn,TernaryFn].
+void buildElementwiseRegion(ImplicitLocOpBuilder &b, Block &block,
+ ElementwiseKind kind,
+ function_ref<InFlightDiagnostic()> emitError);
+
+/// RegionBuilderFn for all named elementwise ops, parameterized by kind.
+template <ElementwiseKind Kind>
+void elementwiseNamedOpRegionBuilder(
+ ImplicitLocOpBuilder &b, Block &block, ArrayRef<NamedAttribute> attrs,
+ function_ref<InFlightDiagnostic()> emitError) {
+ buildElementwiseRegion(b, block, Kind, emitError);
+}
+
+} // namespace mlir::linalg
+
//===----------------------------------------------------------------------===//
// Linalg Dialect Operations
//===----------------------------------------------------------------------===//
diff --git a/mlir/include/mlir/Dialect/Linalg/IR/LinalgInterfaces.h b/mlir/include/mlir/Dialect/Linalg/IR/LinalgInterfaces.h
index 9ff32216ce042..e879452fb8a29 100644
--- a/mlir/include/mlir/Dialect/Linalg/IR/LinalgInterfaces.h
+++ b/mlir/include/mlir/Dialect/Linalg/IR/LinalgInterfaces.h
@@ -31,6 +31,9 @@ class IteratorTypeAttr;
class LinalgOp;
class GenericOp;
+// Forward declaration needed by ElementwiseOpInterface.
+enum class ElementwiseKind : uint32_t;
+
namespace detail {
/// Implementation of the method that check if given operands
/// can be dropped, i.e. the remaining operands can compute the loop
diff --git a/mlir/include/mlir/Dialect/Linalg/IR/LinalgInterfaces.td b/mlir/include/mlir/Dialect/Linalg/IR/LinalgInterfaces.td
index 9f1e88a040f5f..c56b83863fde8 100644
--- a/mlir/include/mlir/Dialect/Linalg/IR/LinalgInterfaces.td
+++ b/mlir/include/mlir/Dialect/Linalg/IR/LinalgInterfaces.td
@@ -798,4 +798,26 @@ def AggregatedOpInterface : OpInterface<"AggregatedOpInterface"> {
];
}
+def ElementwiseOpInterface : OpInterface<"ElementwiseOpInterface"> {
+ let description = [{
+ Interface for operations that represent elementwise computations. This
+ includes both the generic `linalg.elementwise` op and its named
+ specializations a.k.a. linalg names ops (e.g. `linalg.add`, `linalg.exp`).
+
+ The interface exposes the kind of elementwise operation being performed,
+ allowing transforms to handle all elementwise ops uniformly.
+ }];
+ let cppNamespace = "::mlir::linalg";
+ let methods = [
+ InterfaceMethod<
+ /*desc=*/[{
+ Returns the kind of elementwise operation (e.g. add, exp, mul).
+ }],
+ /*retType=*/"::mlir::linalg::ElementwiseKind",
+ /*methodName=*/"getElementwiseKind",
+ /*args=*/(ins)
+ >
+ ];
+}
+
#endif // LINALG_IR_LINALGINTERFACES
diff --git a/mlir/include/mlir/Dialect/Linalg/IR/LinalgNamedStructuredOps.yaml b/mlir/include/mlir/Dialect/Linalg/IR/LinalgNamedStructuredOps.yaml
index 521afc991063f..828981fe17a3f 100644
--- a/mlir/include/mlir/Dialect/Linalg/IR/LinalgNamedStructuredOps.yaml
+++ b/mlir/include/mlir/Dialect/Linalg/IR/LinalgNamedStructuredOps.yaml
@@ -44,913 +44,6 @@ structured_op: !LinalgStructuredOpConfig
- !ScalarExpression
scalar_arg: I
--- !LinalgOpConfig
-metadata: !LinalgOpMetadata
- name: exp
- cpp_class_name: ExpOp
- doc: |-
- Applies exp(x) elementwise.
-
- No numeric casting is performed on the input operand.
-structured_op: !LinalgStructuredOpConfig
- args:
- - !LinalgOperandDefConfig
- name: I
- kind: input_tensor
- type_var: T1
- shape_map: affine_map<() -> ()>
- - !LinalgOperandDefConfig
- name: O
- kind: output_tensor
- type_var: T1
- shape_map: affine_map<() -> ()>
- indexing_maps: !LinalgIndexingMapsConfig
- static_indexing_maps:
- - affine_map<() -> ()>
- - affine_map<() -> ()>
- iterator_types: []
- assignments:
- - !ScalarAssign
- arg: O
- value: !ScalarExpression
- scalar_fn:
- kind: unary
- fn_name: exp
- operands:
- - !ScalarExpression
- scalar_arg: I
---- !LinalgOpConfig
-metadata: !LinalgOpMetadata
- name: log
- cpp_class_name: LogOp
- doc: |-
- Applies log(x) elementwise.
-
- No numeric casting is performed on the input operand.
-structured_op: !LinalgStructuredOpConfig
- args:
- - !LinalgOperandDefConfig
- name: I
- kind: input_tensor
- type_var: T1
- shape_map: affine_map<() -> ()>
- - !LinalgOperandDefConfig
- name: O
- kind: output_tensor
- type_var: T1
- shape_map: affine_map<() -> ()>
- indexing_maps: !LinalgIndexingMapsConfig
- static_indexing_maps:
- - affine_map<() -> ()>
- - affine_map<() -> ()>
- iterator_types: []
- assignments:
- - !ScalarAssign
- arg: O
- value: !ScalarExpression
- scalar_fn:
- kind: unary
- fn_name: log
- operands:
- - !ScalarExpression
- scalar_arg: I
---- !LinalgOpConfig
-metadata: !LinalgOpMetadata
- name: abs
- cpp_class_name: AbsOp
- doc: |-
- Applies abs(x) elementwise.
-
- No numeric casting is performed on the input operand.
-structured_op: !LinalgStructuredOpConfig
- args:
- - !LinalgOperandDefConfig
- name: I
- kind: input_tensor
- type_var: T1
- shape_map: affine_map<() -> ()>
- - !LinalgOperandDefConfig
- name: O
- kind: output_tensor
- type_var: T1
- shape_map: affine_map<() -> ()>
- indexing_maps: !LinalgIndexingMapsConfig
- static_indexing_maps:
- - affine_map<() -> ()>
- - affine_map<() -> ()>
- iterator_types: []
- assignments:
- - !ScalarAssign
- arg: O
- value: !ScalarExpression
- scalar_fn:
- kind: unary
- fn_name: abs
- operands:
- - !ScalarExpression
- scalar_arg: I
---- !LinalgOpConfig
-metadata: !LinalgOpMetadata
- name: ceil
- cpp_class_name: CeilOp
- doc: |-
- Applies ceil(x) elementwise.
-
- No numeric casting is performed on the input operand.
-structured_op: !LinalgStructuredOpConfig
- args:
- - !LinalgOperandDefConfig
- name: I
- kind: input_tensor
- type_var: T1
- shape_map: affine_map<() -> ()>
- - !LinalgOperandDefConfig
- name: O
- kind: output_tensor
- type_var: T1
- shape_map: affine_map<() -> ()>
- indexing_maps: !LinalgIndexingMapsConfig
- static_indexing_maps:
- - affine_map<() -> ()>
- - affine_map<() -> ()>
- iterator_types: []
- assignments:
- - !ScalarAssign
- arg: O
- value: !ScalarExpression
- scalar_fn:
- kind: unary
- fn_name: ceil
- operands:
- - !ScalarExpression
- scalar_arg: I
---- !LinalgOpConfig
-metadata: !LinalgOpMetadata
- name: floor
- cpp_class_name: FloorOp
- doc: |-
- Applies floor(x) elementwise.
-
- No numeric casting is performed on the input operand.
-structured_op: !LinalgStructuredOpConfig
- args:
- - !LinalgOperandDefConfig
- name: I
- kind: input_tensor
- type_var: T1
- shape_map: affine_map<() -> ()>
- - !LinalgOperandDefConfig
- name: O
- kind: output_tensor
- type_var: T1
- shape_map: affine_map<() -> ()>
- indexing_maps: !LinalgIndexingMapsConfig
- static_indexing_maps:
- - affine_map<() -> ()>
- - affine_map<() -> ()>
- iterator_types: []
- assignments:
- - !ScalarAssign
- arg: O
- value: !ScalarExpression
- scalar_fn:
- kind: unary
- fn_name: floor
- operands:
- - !ScalarExpression
- scalar_arg: I
---- !LinalgOpConfig
-metadata: !LinalgOpMetadata
- name: negf
- cpp_class_name: NegFOp
- doc: |-
- Applies negf(x) elementwise.
-
- No numeric casting is performed on the input operand.
-structured_op: !LinalgStructuredOpConfig
- args:
- - !LinalgOperandDefConfig
- name: I
- kind: input_tensor
- type_var: T1
- shape_map: affine_map<() -> ()>
- - !LinalgOperandDefConfig
- name: O
- kind: output_tensor
- type_var: T1
- shape_map: affine_map<() -> ()>
- indexing_maps: !LinalgIndexingMapsConfig
- static_indexing_maps:
- - affine_map<() -> ()>
- - affine_map<() -> ()>
- iterator_types: []
- assignments:
- - !ScalarAssign
- arg: O
- value: !ScalarExpression
- scalar_fn:
- kind: unary
- fn_name: negf
- operands:
- - !ScalarExpression
- scalar_arg: I
---- !LinalgOpConfig
-metadata: !LinalgOpMetadata
- name: reciprocal
- cpp_class_name: ReciprocalOp
- doc: |-
- Applies reciprocal(x) elementwise.
-
- No numeric casting is performed on the input operand.
-structured_op: !LinalgStructuredOpConfig
- args:
- - !LinalgOperandDefConfig
- name: I
- kind: input_tensor
- type_var: T1
- shape_map: affine_map<() -> ()>
- - !LinalgOperandDefConfig
- name: O
- kind: output_tensor
- type_var: T1
- shape_map: affine_map<() -> ()>
- indexing_maps: !LinalgIndexingMapsConfig
- static_indexing_maps:
- - affine_map<() -> ()>
- - affine_map<() -> ()>
- iterator_types: []
- assignments:
- - !ScalarAssign
- arg: O
- value: !ScalarExpression
- scalar_fn:
- kind: unary
- fn_name: reciprocal
- operands:
- - !ScalarExpression
- scalar_arg: I
---- !LinalgOpConfig
-metadata: !LinalgOpMetadata
- name: round
- cpp_class_name: RoundOp
- doc: |-
- Applies round(x) elementwise.
-
- No numeric casting is performed on the input operand.
-structured_op: !LinalgStructuredOpConfig
- args:
- - !LinalgOperandDefConfig
- name: I
- kind: input_tensor
- type_var: T1
- shape_map: affine_map<() -> ()>
- - !LinalgOperandDefConfig
- name: O
- kind: output_tensor
- type_var: T1
- shape_map: affine_map<() -> ()>
- indexing_maps: !LinalgIndexingMapsConfig
- static_indexing_maps:
- - affine_map<() -> ()>
- - affine_map<() -> ()>
- iterator_types: []
- assignments:
- - !ScalarAssign
- arg: O
- value: !ScalarExpression
- scalar_fn:
- kind: unary
- fn_name: round
- operands:
- - !ScalarExpression
- scalar_arg: I
---- !LinalgOpConfig
-metadata: !LinalgOpMetadata
- name: sqrt
- cpp_class_name: SqrtOp
- doc: |-
- Applies sqrt(x) elementwise.
-
- No numeric casting is performed on the input operand.
-structured_op: !LinalgStructuredOpConfig
- args:
- - !LinalgOperandDefConfig
- name: I
- kind: input_tensor
- type_var: T1
- shape_map: affine_map<() -> ()>
- - !LinalgOperandDefConfig
- name: O
- kind: output_tensor
- type_var: T1
- shape_map: affine_map<() -> ()>
- indexing_maps: !LinalgIndexingMapsConfig
- static_indexing_maps:
- - affine_map<() -> ()>
- - affine_map<() -> ()>
- iterator_types: []
- assignments:
- - !ScalarAssign
- arg: O
- value: !ScalarExpression
- scalar_fn:
- kind: unary
- fn_name: sqrt
- operands:
- - !ScalarExpression
- scalar_arg: I
---- !LinalgOpConfig
-metadata: !LinalgOpMetadata
- name: rsqrt
- cpp_class_name: RsqrtOp
- doc: |-
- Applies rsqrt(x) elementwise.
-
- No numeric casting is performed on the input operand.
-structured_op: !LinalgStructuredOpConfig
- args:
- - !LinalgOperandDefConfig
- name: I
- kind: input_tensor
- type_var: T1
- shape_map: affine_map<() -> ()>
- - !LinalgOperandDefConfig
- name: O
- kind: output_tensor
- type_var: T1
- shape_map: affine_map<() -> ()>
- indexing_maps: !LinalgIndexingMapsConfig
- static_indexing_maps:
- - affine_map<() -> ()>
- - affine_map<() -> ()>
- iterator_types: []
- assignments:
- - !ScalarAssign
- arg: O
- value: !ScalarExpression
- scalar_fn:
- kind: unary
- fn_name: rsqrt
- operands:
- - !ScalarExpression
- scalar_arg: I
---- !LinalgOpConfig
-metadata: !LinalgOpMetadata
- name: square
- cpp_class_name: SquareOp
- doc: |-
- Applies square(x) elementwise.
-
- No numeric casting is performed on the input operand.
-structured_op: !LinalgStructuredOpConfig
- args:
- - !LinalgOperandDefConfig
- name: I
- kind: input_tensor
- type_var: T1
- shape_map: affine_map<() -> ()>
- - !LinalgOperandDefConfig
- name: O
- kind: output_tensor
- type_var: T1
- shape_map: affine_map<() -> ()>
- indexing_maps: !LinalgIndexingMapsConfig
- static_indexing_maps:
- - affine_map<() -> ()>
- - affine_map<() -> ()>
- iterator_types: []
- assignments:
- - !ScalarAssign
- arg: O
- value: !ScalarExpression
- scalar_fn:
- kind: unary
- fn_name: square
- operands:
- - !ScalarExpression
- scalar_arg: I
---- !LinalgOpConfig
-metadata: !LinalgOpMetadata
- name: tanh
- cpp_class_name: TanhOp
- doc: |-
- Applies tanh(x) elementwise.
-
- No numeric casting is performed on the input operand.
-structured_op: !LinalgStructuredOpConfig
- args:
- - !LinalgOperandDefConfig
- name: I
- kind: input_tensor
- type_var: T1
- shape_map: affine_map<() -> ()>
- - !LinalgOperandDefConfig
- name: O
- kind: output_tensor
- type_var: T1
- shape_map: affine_map<() -> ()>
- indexing_maps: !LinalgIndexingMapsConfig
- static_indexing_maps:
- - affine_map<() -> ()>
- - affine_map<() -> ()>
- iterator_types: []
- assignments:
- - !ScalarAssign
- arg: O
- value: !ScalarExpression
- scalar_fn:
- kind: unary
- fn_name: tanh
- operands:
- - !ScalarExpression
- scalar_arg: I
---- !LinalgOpConfig
-metadata: !LinalgOpMetadata
- name: erf
- cpp_class_name: ErfOp
- doc: |-
- Applies erf(x) elementwise.
-
- No numeric casting is performed on the input operand.
-structured_op: !LinalgStructuredOpConfig
- args:
- - !LinalgOperandDefConfig
- name: I
- kind: input_tensor
- type_var: T1
- shape_map: affine_map<() -> ()>
- - !LinalgOperandDefConfig
- name: O
- kind: output_tensor
- type_var: T1
- shape_map: affine_map<() -> ()>
- indexing_maps: !LinalgIndexingMapsConfig
- static_indexing_maps:
- - affine_map<() -> ()>
- - affine_map<() -> ()>
- iterator_types: []
- assignments:
- - !ScalarAssign
- arg: O
- value: !ScalarExpression
- scalar_fn:
- kind: unary
- fn_name: erf
- operands:
- - !ScalarExpression
- scalar_arg: I
---- !LinalgOpConfig
-metadata: !LinalgOpMetadata
- name: add
- cpp_class_name: AddOp
- doc: |-
- Adds two tensors elementwise.
-
- The shapes and element types must be identical. The appropriate casts,
- broadcasts and reductions should be done previously to calling this op.
-
- This means reduction/broadcast/element cast semantics is explicit. Further
- passes can take that into account when lowering this code. For example,
- a `linalg.broadcast` + `linalg.add` sequence can be lowered to a
- `linalg.generic` with different affine maps for the two operands.
-structured_op: !LinalgStructuredOpConfig
- args:
- - !LinalgOperandDefConfig
- name: lhs
- kind: input_tensor
- type_var: T1
- shape_map: affine_map<() -> ()>
- - !LinalgOperandDefConfig
- name: rhs
- kind: input_tensor
- type_var: T1
- shape_map: affine_map<() -> ()>
- - !LinalgOperandDefConfig
- name: O
- kind: output_tensor
- type_var: T1
- shape_map: affine_map<() -> ()>
- indexing_maps: !LinalgIndexingMapsConfig
- static_indexing_maps:
- - affine_map<() -> ()>
- - affine_map<() -> ()>
- - affine_map<() -> ()>
- iterator_types: []
- assignments:
- - !ScalarAssign
- arg: O
- value: !ScalarExpression
- scalar_fn:
- kind: binary
- fn_name: add
- operands:
- - !ScalarExpression
- scalar_arg: lhs
- - !ScalarExpression
- scalar_arg: rhs
---- !LinalgOpConfig
-metadata: !LinalgOpMetadata
- name: sub
- cpp_class_name: SubOp
- doc: |-
- Subtracts two tensors elementwise.
-
- The shapes and element types must be identical. The appropriate casts,
- broadcasts and reductions should be done previously to calling this op.
-
- This means reduction/broadcast/element cast semantics is explicit. Further
- passes can take that into account when lowering this code. For example,
- a `linalg.broadcast` + `linalg.sub` sequence can be lowered to a
- `linalg.generic` with different affine maps for the two operands.
-structured_op: !LinalgStructuredOpConfig
- args:
- - !LinalgOperandDefConfig
- name: lhs
- kind: input_tensor
- type_var: T1
- shape_map: affine_map<() -> ()>
- - !LinalgOperandDefConfig
- name: rhs
- kind: input_tensor
- type_var: T1
- shape_map: affine_map<() -> ()>
- - !LinalgOperandDefConfig
- name: O
- kind: output_tensor
- type_var: T1
- shape_map: affine_map<() -> ()>
- indexing_maps: !LinalgIndexingMapsConfig
- static_indexing_maps:
- - affine_map<() -> ()>
- - affine_map<() -> ()>
- - affine_map<() -> ()>
- iterator_types: []
- assignments:
- - !ScalarAssign
- arg: O
- value: !ScalarExpression
- scalar_fn:
- kind: binary
- fn_name: sub
- operands:
- - !ScalarExpression
- scalar_arg: lhs
- - !ScalarExpression
- scalar_arg: rhs
---- !LinalgOpConfig
-metadata: !LinalgOpMetadata
- name: mul
- cpp_class_name: MulOp
- doc: |-
- Multiplies two tensors elementwise.
-
- The shapes and element types must be identical. The appropriate casts,
- broadcasts and reductions should be done previously to calling this op.
-
- This means reduction/broadcast/element cast semantics is explicit. Further
- passes can take that into account when lowering this code. For example,
- a `linalg.broadcast` + `linalg.mul` sequence can be lowered to a
- `linalg.generic` with different affine maps for the two operands.
-structured_op: !LinalgStructuredOpConfig
- args:
- - !LinalgOperandDefConfig
- name: lhs
- kind: input_tensor
- type_var: T1
- shape_map: affine_map<() -> ()>
- - !LinalgOperandDefConfig
- name: rhs
- kind: input_tensor
- type_var: T1
- shape_map: affine_map<() -> ()>
- - !LinalgOperandDefConfig
- name: O
- kind: output_tensor
- type_var: T1
- shape_map: affine_map<() -> ()>
- indexing_maps: !LinalgIndexingMapsConfig
- static_indexing_maps:
- - affine_map<() -> ()>
- - affine_map<() -> ()>
- - affine_map<() -> ()>
- iterator_types: []
- assignments:
- - !ScalarAssign
- arg: O
- value: !ScalarExpression
- scalar_fn:
- kind: binary
- fn_name: mul
- operands:
- - !ScalarExpression
- scalar_arg: lhs
- - !ScalarExpression
- scalar_arg: rhs
---- !LinalgOpConfig
-metadata: !LinalgOpMetadata
- name: div
- cpp_class_name: DivOp
- doc: |-
- Divides the first tensor by the second tensor, elementwise.
-
- The shapes and element types must be identical. The appropriate casts,
- broadcasts and reductions should be done previously to calling this op.
-
- This means reduction/broadcast/element cast semantics is explicit. Further
- passes can take that into account when lowering this code. For example,
- a `linalg.broadcast` + `linalg.div` sequence can be lowered to a
- `linalg.generic` with different affine maps for the two operands.
-structured_op: !LinalgStructuredOpConfig
- args:
- - !LinalgOperandDefConfig
- name: lhs
- kind: input_tensor
- type_var: T1
- shape_map: affine_map<() -> ()>
- - !LinalgOperandDefConfig
- name: rhs
- kind: input_tensor
- type_var: T1
- shape_map: affine_map<() -> ()>
- - !LinalgOperandDefConfig
- name: O
- kind: output_tensor
- type_var: T1
- shape_map: affine_map<() -> ()>
- indexing_maps: !LinalgIndexingMapsConfig
- static_indexing_maps:
- - affine_map<() -> ()>
- - affine_map<() -> ()>
- - affine_map<() -> ()>
- iterator_types: []
- assignments:
- - !ScalarAssign
- arg: O
- value: !ScalarExpression
- scalar_fn:
- kind: binary
- fn_name: div
- operands:
- - !ScalarExpression
- scalar_arg: lhs
- - !ScalarExpression
- scalar_arg: rhs
---- !LinalgOpConfig
-metadata: !LinalgOpMetadata
- name: div_unsigned
- cpp_class_name: DivUnsignedOp
- doc: |-
- Divides the first tensor by the second tensor, elementwise. For integer
- types, performs an unsigned division.
-
- The shapes and element types must be identical. The appropriate casts,
- broadcasts and reductions should be done previously to calling this op.
-
- This means reduction/broadcast/element cast semantics is explicit. Further
- passes can take that into account when lowering this code. For example,
- a `linalg.broadcast` + `linalg.div` sequence can be lowered to a
- `linalg.generic` with different affine maps for the two operands.
-structured_op: !LinalgStructuredOpConfig
- args:
- - !LinalgOperandDefConfig
- name: lhs
- kind: input_tensor
- type_var: T1
- shape_map: affine_map<() -> ()>
- - !LinalgOperandDefConfig
- name: rhs
- kind: input_tensor
- type_var: T1
- shape_map: affine_map<() -> ()>
- - !LinalgOperandDefConfig
- name: O
- kind: output_tensor
- type_var: T1
- shape_map: affine_map<() -> ()>
- indexing_maps: !LinalgIndexingMapsConfig
- static_indexing_maps:
- - affine_map<() -> ()>
- - affine_map<() -> ()>
- - affine_map<() -> ()>
- iterator_types: []
- assignments:
- - !ScalarAssign
- arg: O
- value: !ScalarExpression
- scalar_fn:
- kind: binary
- fn_name: div_unsigned
- operands:
- - !ScalarExpression
- scalar_arg: lhs
- - !ScalarExpression
- scalar_arg: rhs
---- !LinalgOpConfig
-metadata: !LinalgOpMetadata
- name: max
- cpp_class_name: MaxOp
- doc: |-
- Takes the max (signed) between two inputs, elementwise.
-
- The shapes and element types must be identical. The appropriate casts,
- broadcasts and reductions should be done previously to calling this op.
-
- This means reduction/broadcast/element cast semantics is explicit. Further
- passes can take that into account when lowering this code. For example,
- a `linalg.broadcast` + `linalg.max` sequence can be lowered to a
- `linalg.generic` with different affine maps for the two operands.
-structured_op: !LinalgStructuredOpConfig
- args:
- - !LinalgOperandDefConfig
- name: lhs
- kind: input_tensor
- type_var: T1
- shape_map: affine_map<() -> ()>
- - !LinalgOperandDefConfig
- name: rhs
- kind: input_tensor
- type_var: T1
- shape_map: affine_map<() -> ()>
- - !LinalgOperandDefConfig
- name: O
- kind: output_tensor
- type_var: T1
- shape_map: affine_map<() -> ()>
- indexing_maps: !LinalgIndexingMapsConfig
- static_indexing_maps:
- - affine_map<() -> ()>
- - affine_map<() -> ()>
- - affine_map<() -> ()>
- iterator_types: []
- assignments:
- - !ScalarAssign
- arg: O
- value: !ScalarExpression
- scalar_fn:
- kind: binary
- fn_name: max_signed
- operands:
- - !ScalarExpression
- scalar_arg: lhs
- - !ScalarExpression
- scalar_arg: rhs
---- !LinalgOpConfig
-metadata: !LinalgOpMetadata
- name: min
- cpp_class_name: MinOp
- doc: |-
- Takes the min (signed) between two inputs, elementwise.
-
- The shapes and element types must be identical. The appropriate casts,
- broadcasts and reductions should be done previously to calling this op.
-
- This means reduction/broadcast/element cast semantics is explicit. Further
- passes can take that into account when lowering this code. For example,
- a `linalg.broadcast` + `linalg.min` sequence can be lowered to a
- `linalg.generic` with different affine maps for the two operands.
-structured_op: !LinalgStructuredOpConfig
- args:
- - !LinalgOperandDefConfig
- name: lhs
- kind: input_tensor
- type_var: T1
- shape_map: affine_map<() -> ()>
- - !LinalgOperandDefConfig
- name: rhs
- kind: input_tensor
- type_var: T1
- shape_map: affine_map<() -> ()>
- - !LinalgOperandDefConfig
- name: O
- kind: output_tensor
- type_var: T1
- shape_map: affine_map<() -> ()>
- indexing_maps: !LinalgIndexingMapsConfig
- static_indexing_maps:
- - affine_map<() -> ()>
- - affine_map<() -> ()>
- - affine_map<() -> ()>
- iterator_types: []
- assignments:
- - !ScalarAssign
- arg: O
- value: !ScalarExpression
- scalar_fn:
- kind: binary
- fn_name: min_signed
- operands:
- - !ScalarExpression
- scalar_arg: lhs
- - !ScalarExpression
- scalar_arg: rhs
---- !LinalgOpConfig
-metadata: !LinalgOpMetadata
- name: powf
- cpp_class_name: PowFOp
- doc: |-
- Takes the powf(lhs, rhs) between two inputs, elementwise. For powf(arg, 2) use `linalg.square`.
-
- Only applies to floating point values.
-
- The shapes and element types must be identical. The appropriate casts,
- broadcasts and reductions should be done previously to calling this op.
-
- This means reduction/broadcast/element cast semantics is explicit. Further
- passes can take that into account when lowering this code. For example,
- a `linalg.broadcast` + `linalg.powf` sequence can be lowered to a
- `linalg.generic` with different affine maps for the two operands.
-structured_op: !LinalgStructuredOpConfig
- args:
- - !LinalgOperandDefConfig
- name: lhs
- kind: input_tensor
- type_var: T1
- shape_map: affine_map<() -> ()>
- - !LinalgOperandDefConfig
- name: rhs
- kind: input_tensor
- type_var: T1
- shape_map: affine_map<() -> ()>
- - !LinalgOperandDefConfig
- name: O
- kind: output_tensor
- type_var: T1
- shape_map: affine_map<() -> ()>
- indexing_maps: !LinalgIndexingMapsConfig
- static_indexing_maps:
- - affine_map<() -> ()>
- - affine_map<() -> ()>
- - affine_map<() -> ()>
- iterator_types: []
- assignments:
- - !ScalarAssign
- arg: O
- value: !ScalarExpression
- scalar_fn:
- kind: binary
- fn_name: powf
- operands:
- - !ScalarExpression
- scalar_arg: lhs
- - !ScalarExpression
- scalar_arg: rhs
---- !LinalgOpConfig
-metadata: !LinalgOpMetadata
- name: select
- cpp_class_name: SelectOp
- doc: |-
- Chooses one value based on a binary condition supplied as its first operand.
-
- The shapes and element types must be identical. The appropriate casts,
- broadcasts and reductions should be done previously to calling this op.
-
- This means reduction/broadcast/element cast semantics is explicit. Further
- passes can take that into account when lowering this code. For example,
- a `linalg.broadcast` + `linalg.select` sequence can be lowered to a
- `linalg.generic` with different affine maps for the two operands.
-structured_op: !LinalgStructuredOpConfig
- args:
- - !LinalgOperandDefConfig
- name: cond
- kind: input_tensor
- type_var: U
- shape_map: affine_map<() -> ()>
- - !LinalgOperandDefConfig
- name: lhs
- kind: input_tensor
- type_var: T1
- shape_map: affine_map<() -> ()>
- - !LinalgOperandDefConfig
- name: rhs
- kind: input_tensor
- type_var: T1
- shape_map: affine_map<() -> ()>
- - !LinalgOperandDefConfig
- name: O
- kind: output_tensor
- type_var: T1
- shape_map: affine_map<() -> ()>
- indexing_maps: !LinalgIndexingMapsConfig
- static_indexing_maps:
- - affine_map<() -> ()>
- - affine_map<() -> ()>
- - affine_map<() -> ()>
- - affine_map<() -> ()>
- iterator_types: []
- assignments:
- - !ScalarAssign
- arg: O
- value: !ScalarExpression
- scalar_fn:
- kind: ternary
- fn_name: select
- operands:
- - !ScalarExpression
- scalar_arg: cond
- - !ScalarExpression
- scalar_arg: lhs
- - !ScalarExpression
- scalar_arg: rhs
---- !LinalgOpConfig
metadata: !LinalgOpMetadata
name: quantized_matmul
cpp_class_name: QuantizedMatmulOp
diff --git a/mlir/include/mlir/Dialect/Linalg/IR/LinalgStructuredOps.td b/mlir/include/mlir/Dialect/Linalg/IR/LinalgStructuredOps.td
index 51bfcbec5d1a4..12accf3e82d85 100644
--- a/mlir/include/mlir/Dialect/Linalg/IR/LinalgStructuredOps.td
+++ b/mlir/include/mlir/Dialect/Linalg/IR/LinalgStructuredOps.td
@@ -547,7 +547,8 @@ def BroadcastOp : LinalgStructuredBase_Op<"broadcast", [
//===----------------------------------------------------------------------===//
def ElementwiseOp : LinalgStructuredBase_Op<"elementwise", [
- AttrSizedOperandSegments]> {
+ AttrSizedOperandSegments,
+ DeclareOpInterfaceMethods<ElementwiseOpInterface>]> {
let summary = [{ Performs element-wise operation }];
let description = [{
The attribute `kind` describes arithmetic operation to perform. The
@@ -1216,6 +1217,182 @@ def BatchReduceMatmulOp : LinalgStructuredBase_Op<"batch_reduce_matmul", [
}];
}
+//===----------------------------------------------------------------------===//
+// Elementwise specializations
+//
+// These are registered ops that are specializations of `linalg.elementwise`,
+// also known as linalg (elementwise) named ops.
+// They always use identity indexing maps and fix the operation kind.
+// Their C++ implementations reuse the same helpers
+// (buildStructuredOp, RegionBuilderHelper) as ElementwiseOp.
+//===----------------------------------------------------------------------===//
+
+multiclass ElementwiseNamedOp<string mnemonic, string kind, int numRegionArgs,
+ string opSummary> {
+ def Op : LinalgStructuredBase_Op<mnemonic, [AttrSizedOperandSegments,
+ DeclareOpInterfaceMethods<ElementwiseOpInterface>]> {
+ let summary = opSummary;
+ let description = [{
+ }] # opSummary # [{
+
+
+ The shapes and element types must be identical. The appropriate casts,
+ broadcasts and reductions should be done previously to calling this op.
+ Ideally, named ops should be lowered to linalg.elementwise and then
+ broadcast, transpose can be folded in using indexing maps.
+ }];
+
+ let arguments = (ins
+ Variadic<AnyType>:$inputs,
+ Variadic<AnyShaped>:$outputs
+ );
+ let results = (outs Variadic<AnyRankedTensor>:$result_tensors);
+ let regions = (region AnyRegion:$region);
+
+ let skipDefaultBuilders = 1;
+ let builders = [
+ OpBuilder<
+ (ins "ValueRange":$inputs, "ValueRange":$outputs,
+ CArg<"ArrayRef<NamedAttribute>", "{}">:$attributes),
+ [{
+ buildStructuredOp($_builder, $_state, std::nullopt, inputs, outputs,
+ attributes, }] # NAME # [{Op::getRegionBuilder());
+ }]>,
+ OpBuilder<
+ (ins "TypeRange":$resultTensorTypes, "ValueRange":$inputs,
+ "ValueRange":$outputs,
+ CArg<"ArrayRef<NamedAttribute>", "{}">:$attributes),
+ [{
+ buildStructuredOp($_builder, $_state, resultTensorTypes,
+ inputs, outputs, attributes, }] # NAME # [{Op::getRegionBuilder());
+ }]>,
+ OpBuilder<
+ (ins "TypeRange":$resultTensorTypes, "ValueRange":$operands,
+ CArg<"ArrayRef<NamedAttribute>", "{}">:$attributes),
+ [{
+ $_state.addOperands(operands);
+ $_state.addAttributes(attributes);
+ $_state.addTypes(resultTensorTypes);
+ (void)$_state.addRegion();
+ }]>
+ ];
+
+ let hasCustomAssemblyFormat = 1;
+ let hasFolder = 1;
+
+ let extraClassDeclaration = structuredOpsBaseDecls # [{
+ SmallVector<utils::IteratorType> getIteratorTypesArray() {
+ int64_t rank = getRank(getDpsInitOperand(0));
+ return SmallVector<utils::IteratorType>(rank,
+ utils::IteratorType::parallel);
+ }
+
+ ArrayAttr getIndexingMaps() {
+ unsigned numDims = getRank(getDpsInitOperand(0));
+ MLIRContext *context = getContext();
+ AffineMap scalarMap = AffineMap::get(numDims, 0, context);
+ AffineMap tensorMap = numDims == 0
+ ? scalarMap
+ : AffineMap::getMultiDimIdentityMap(numDims, context);
+ SmallVector<AffineMap> maps;
+ for (OpOperand &opOperand : getOperation()->getOpOperands())
+ maps.push_back(getRank(&opOperand) == 0 ? scalarMap : tensorMap);
+ return Builder(context).getAffineMapArrayAttr(maps);
+ }
+
+ ::mlir::MutableOperandRange getDpsInitsMutable() {
+ return getOutputsMutable();
+ }
+
+ std::string getLibraryCallName() {
+ return generateLibraryCallName(getOperation());
+ }
+
+ static unsigned getNumRegionArgs() { return }] # !cast<string>(numRegionArgs) # [{; }
+
+ static std::function<void(ImplicitLocOpBuilder &,
+ Block &, ArrayRef<NamedAttribute>,
+ function_ref<InFlightDiagnostic()>)>
+ getRegionBuilder() {
+ return elementwiseNamedOpRegionBuilder<ElementwiseKind::}] # kind # [{>;
+ }
+ }];
+
+ let extraClassDefinition = [{
+ ElementwiseKind $cppClass::getElementwiseKind() {
+ return ElementwiseKind::}] # kind # [{;
+ }
+ ParseResult $cppClass::parse(OpAsmParser &parser, OperationState &result) {
+ return ::parseNamedStructuredOp(parser, result,
+ $cppClass::getNumRegionArgs(),
+ $cppClass::getRegionBuilder());
+ }
+ void $cppClass::print(OpAsmPrinter &p) {
+ ::printNamedStructuredOp(p, getOperation(), getInputs(), getOutputs(),
+ {"operandSegmentSizes",
+ "linalg.memoized_indexing_maps"});
+ }
+ LogicalResult $cppClass::fold(FoldAdaptor,
+ SmallVectorImpl<OpFoldResult> &) {
+ return memref::foldMemRefCast(*this);
+ }
+ void $cppClass::getEffects(
+ SmallVectorImpl<SideEffects::EffectInstance<MemoryEffects::Effect>>
+ &effects) {
+ if (hasPureTensorSemantics())
+ return;
+ getGenericEffectsImpl(effects, cast<LinalgOp>(getOperation()));
+ }
+ Speculation::Speculatability $cppClass::getSpeculatability() {
+ return getGenericSpeculatabilityImpl(cast<LinalgOp>(getOperation()));
+ }
+ }];
+ }
+}
+
+// Shorthand wrappers for common arities where kind == mnemonic.
+multiclass UnaryElementwiseOp<string mnemonic>
+ : ElementwiseNamedOp<mnemonic, mnemonic, 2,
+ "Applies " # mnemonic # "(x) elementwise.">;
+
+multiclass BinaryElementwiseOp<string mnemonic, string summary>
+ : ElementwiseNamedOp<mnemonic, mnemonic, 3, summary>;
+
+// --- Unary ops ---
+defm Exp : UnaryElementwiseOp<"exp">;
+defm Log : UnaryElementwiseOp<"log">;
+defm Abs : UnaryElementwiseOp<"abs">;
+defm Ceil : UnaryElementwiseOp<"ceil">;
+defm Floor : UnaryElementwiseOp<"floor">;
+defm NegF : UnaryElementwiseOp<"negf">;
+defm Reciprocal: UnaryElementwiseOp<"reciprocal">;
+defm Round : UnaryElementwiseOp<"round">;
+defm Sqrt : UnaryElementwiseOp<"sqrt">;
+defm Rsqrt : UnaryElementwiseOp<"rsqrt">;
+defm Square : UnaryElementwiseOp<"square">;
+defm Tanh : UnaryElementwiseOp<"tanh">;
+defm Erf : UnaryElementwiseOp<"erf">;
+
+// --- Binary ops ---
+defm Add : BinaryElementwiseOp<"add", "Adds two tensors elementwise.">;
+defm Sub : BinaryElementwiseOp<"sub", "Subtracts two tensors elementwise.">;
+defm Mul : BinaryElementwiseOp<"mul", "Multiplies two tensors elementwise.">;
+defm Div : BinaryElementwiseOp<"div", "Divides two tensors elementwise.">;
+defm DivUnsigned : BinaryElementwiseOp<"div_unsigned",
+ "Unsigned-divides two tensors elementwise.">;
+defm PowF : BinaryElementwiseOp<"powf",
+ "Takes powf(lhs, rhs) elementwise.">;
+
+// Binary ops where kind != mnemonic (signed variants).
+defm Max : ElementwiseNamedOp<"max", "max_signed", 3,
+ "Takes the signed max between two tensors, elementwise.">;
+defm Min : ElementwiseNamedOp<"min", "min_signed", 3,
+ "Takes the signed min between two tensors, elementwise.">;
+
+// --- Ternary ops ---
+defm Select : ElementwiseNamedOp<"select", "select", 4,
+ "Chooses one value based on a binary condition.">;
+
//===----------------------------------------------------------------------===//
// Named Linalg ops, implemented as a declarative configurations of generic ops.
//===----------------------------------------------------------------------===//
diff --git a/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp b/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp
index 170e1edf8a55d..ab55e57dedf7a 100644
--- a/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp
+++ b/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp
@@ -5114,6 +5114,43 @@ Speculation::Speculatability ElementwiseOp::getSpeculatability() {
return getGenericSpeculatabilityImpl(cast<LinalgOp>(getOperation()));
}
+ElementwiseKind ElementwiseOp::getElementwiseKind() { return getKind(); }
+
+//===----------------------------------------------------------------------===//
+// Shared utilities for named elementwise ops (AddOp, SubOp, ExpOp, etc.)
+//===----------------------------------------------------------------------===//
+
+void buildElementwiseRegion(ImplicitLocOpBuilder &b, Block &block,
+ ElementwiseKind kind,
+ function_ref<InFlightDiagnostic()> emitError) {
+ ArityGroupAndKind groupAndKind = getArityGroupAndKind(kind);
+ auto arityGroup = groupAndKind.arityGroup;
+ auto fnKind = groupAndKind.kind;
+
+ unsigned expectedArgs = getArityGroupAsUInt(arityGroup) + 1;
+ assert(block.getNumArguments() == expectedArgs &&
+ "elementwise regionBuilder arg count mismatch");
+
+ RegionBuilderHelper helper(b, block);
+ Value result;
+
+ if (arityGroup == ElementwiseArityGroup::Unary) {
+ result = helper.buildUnaryFn(fnKind.unaryFn, block.getArgument(0));
+ } else if (arityGroup == ElementwiseArityGroup::Binary) {
+ result = helper.buildBinaryFn(fnKind.binaryFn, block.getArgument(0),
+ block.getArgument(1), emitError);
+ } else if (arityGroup == ElementwiseArityGroup::Ternary) {
+ result = helper.buildTernaryFn(fnKind.ternaryFn, block.getArgument(0),
+ block.getArgument(1), block.getArgument(2));
+ } else {
+ assert(false && "unhandled arity group");
+ }
+
+ if (!result)
+ return;
+ helper.yieldOutputs({result});
+}
+
//===----------------------------------------------------------------------===//
// PackOp/UnPackOp Common
//===----------------------------------------------------------------------===//
diff --git a/mlir/python/mlir/dialects/linalg/opdsl/ops/core_named_ops.py b/mlir/python/mlir/dialects/linalg/opdsl/ops/core_named_ops.py
index 9c24f94fcf612..16d76bb07dd88 100644
--- a/mlir/python/mlir/dialects/linalg/opdsl/ops/core_named_ops.py
+++ b/mlir/python/mlir/dialects/linalg/opdsl/ops/core_named_ops.py
@@ -21,337 +21,6 @@ def copy(
O[None] = cast(U, I[None])
- at linalg_structured_op
-def exp(
- I=TensorDef(T1),
- O=TensorDef(T1, output=True),
-):
- """Applies exp(x) elementwise.
-
- No numeric casting is performed on the input operand.
- """
- O[None] = UnaryFn.exp(I[None])
-
-
- at linalg_structured_op
-def log(
- I=TensorDef(T1),
- O=TensorDef(T1, output=True),
-):
- """Applies log(x) elementwise.
-
- No numeric casting is performed on the input operand.
- """
- O[None] = UnaryFn.log(I[None])
-
-
- at linalg_structured_op
-def abs(
- I=TensorDef(T1),
- O=TensorDef(T1, output=True),
-):
- """Applies abs(x) elementwise.
-
- No numeric casting is performed on the input operand.
- """
- O[None] = UnaryFn.abs(I[None])
-
-
- at linalg_structured_op
-def ceil(
- I=TensorDef(T1),
- O=TensorDef(T1, output=True),
-):
- """Applies ceil(x) elementwise.
-
- No numeric casting is performed on the input operand.
- """
- O[None] = UnaryFn.ceil(I[None])
-
-
- at linalg_structured_op
-def floor(
- I=TensorDef(T1),
- O=TensorDef(T1, output=True),
-):
- """Applies floor(x) elementwise.
-
- No numeric casting is performed on the input operand.
- """
- O[None] = UnaryFn.floor(I[None])
-
-
- at linalg_structured_op(op_class_name="NegFOp")
-def negf(
- I=TensorDef(T1),
- O=TensorDef(T1, output=True),
-):
- """Applies negf(x) elementwise.
-
- No numeric casting is performed on the input operand.
- """
- O[None] = UnaryFn.negf(I[None])
-
-
- at linalg_structured_op(op_class_name="ReciprocalOp")
-def reciprocal(
- I=TensorDef(T1),
- O=TensorDef(T1, output=True),
-):
- """Applies reciprocal(x) elementwise.
-
- No numeric casting is performed on the input operand.
- """
- O[None] = UnaryFn.reciprocal(I[None])
-
-
- at linalg_structured_op
-def round(
- I=TensorDef(T1),
- O=TensorDef(T1, output=True),
-):
- """Applies round(x) elementwise.
-
- No numeric casting is performed on the input operand.
- """
- O[None] = UnaryFn.round(I[None])
-
-
- at linalg_structured_op
-def sqrt(
- I=TensorDef(T1),
- O=TensorDef(T1, output=True),
-):
- """Applies sqrt(x) elementwise.
-
- No numeric casting is performed on the input operand.
- """
- O[None] = UnaryFn.sqrt(I[None])
-
-
- at linalg_structured_op
-def rsqrt(
- I=TensorDef(T1),
- O=TensorDef(T1, output=True),
-):
- """Applies rsqrt(x) elementwise.
-
- No numeric casting is performed on the input operand.
- """
- O[None] = UnaryFn.rsqrt(I[None])
-
-
- at linalg_structured_op
-def square(
- I=TensorDef(T1),
- O=TensorDef(T1, output=True),
-):
- """Applies square(x) elementwise.
-
- No numeric casting is performed on the input operand.
- """
- O[None] = UnaryFn.square(I[None])
-
-
- at linalg_structured_op
-def tanh(
- I=TensorDef(T1),
- O=TensorDef(T1, output=True),
-):
- """Applies tanh(x) elementwise.
-
- No numeric casting is performed on the input operand.
- """
- O[None] = UnaryFn.tanh(I[None])
-
-
- at linalg_structured_op
-def erf(
- I=TensorDef(T1),
- O=TensorDef(T1, output=True),
-):
- """Applies erf(x) elementwise.
-
- No numeric casting is performed on the input operand.
- """
- O[None] = UnaryFn.erf(I[None])
-
-
- at linalg_structured_op
-def add(
- lhs=TensorDef(T1),
- rhs=TensorDef(T1),
- O=TensorDef(T1, output=True),
-):
- """Adds two tensors elementwise.
-
- The shapes and element types must be identical. The appropriate casts,
- broadcasts and reductions should be done previously to calling this op.
-
- This means reduction/broadcast/element cast semantics is explicit. Further
- passes can take that into account when lowering this code. For example,
- a `linalg.broadcast` + `linalg.add` sequence can be lowered to a
- `linalg.generic` with different affine maps for the two operands.
- """
- O[None] = BinaryFn.add(lhs[None], rhs[None])
-
-
- at linalg_structured_op
-def sub(
- lhs=TensorDef(T1),
- rhs=TensorDef(T1),
- O=TensorDef(T1, output=True),
-):
- """Subtracts two tensors elementwise.
-
- The shapes and element types must be identical. The appropriate casts,
- broadcasts and reductions should be done previously to calling this op.
-
- This means reduction/broadcast/element cast semantics is explicit. Further
- passes can take that into account when lowering this code. For example,
- a `linalg.broadcast` + `linalg.sub` sequence can be lowered to a
- `linalg.generic` with different affine maps for the two operands.
- """
- O[None] = BinaryFn.sub(lhs[None], rhs[None])
-
-
- at linalg_structured_op
-def mul(
- lhs=TensorDef(T1),
- rhs=TensorDef(T1),
- O=TensorDef(T1, output=True),
-):
- """Multiplies two tensors elementwise.
-
- The shapes and element types must be identical. The appropriate casts,
- broadcasts and reductions should be done previously to calling this op.
-
- This means reduction/broadcast/element cast semantics is explicit. Further
- passes can take that into account when lowering this code. For example,
- a `linalg.broadcast` + `linalg.mul` sequence can be lowered to a
- `linalg.generic` with different affine maps for the two operands.
- """
- O[None] = BinaryFn.mul(lhs[None], rhs[None])
-
-
- at linalg_structured_op
-def div(
- lhs=TensorDef(T1),
- rhs=TensorDef(T1),
- O=TensorDef(T1, output=True),
-):
- """Divides the first tensor by the second tensor, elementwise.
-
- The shapes and element types must be identical. The appropriate casts,
- broadcasts and reductions should be done previously to calling this op.
-
- This means reduction/broadcast/element cast semantics is explicit. Further
- passes can take that into account when lowering this code. For example,
- a `linalg.broadcast` + `linalg.div` sequence can be lowered to a
- `linalg.generic` with different affine maps for the two operands.
- """
- O[None] = BinaryFn.div(lhs[None], rhs[None])
-
-
- at linalg_structured_op
-def div_unsigned(
- lhs=TensorDef(T1),
- rhs=TensorDef(T1),
- O=TensorDef(T1, output=True),
-):
- """Divides the first tensor by the second tensor, elementwise. For integer
- types, performs an unsigned division.
-
- The shapes and element types must be identical. The appropriate casts,
- broadcasts and reductions should be done previously to calling this op.
-
- This means reduction/broadcast/element cast semantics is explicit. Further
- passes can take that into account when lowering this code. For example,
- a `linalg.broadcast` + `linalg.div` sequence can be lowered to a
- `linalg.generic` with different affine maps for the two operands.
- """
- O[None] = BinaryFn.div_unsigned(lhs[None], rhs[None])
-
-
- at linalg_structured_op
-def max(
- lhs=TensorDef(T1),
- rhs=TensorDef(T1),
- O=TensorDef(T1, output=True),
-):
- """Takes the max (signed) between two inputs, elementwise.
-
- The shapes and element types must be identical. The appropriate casts,
- broadcasts and reductions should be done previously to calling this op.
-
- This means reduction/broadcast/element cast semantics is explicit. Further
- passes can take that into account when lowering this code. For example,
- a `linalg.broadcast` + `linalg.max` sequence can be lowered to a
- `linalg.generic` with different affine maps for the two operands.
- """
- O[None] = BinaryFn.max_signed(lhs[None], rhs[None])
-
-
- at linalg_structured_op
-def min(
- lhs=TensorDef(T1),
- rhs=TensorDef(T1),
- O=TensorDef(T1, output=True),
-):
- """Takes the min (signed) between two inputs, elementwise.
-
- The shapes and element types must be identical. The appropriate casts,
- broadcasts and reductions should be done previously to calling this op.
-
- This means reduction/broadcast/element cast semantics is explicit. Further
- passes can take that into account when lowering this code. For example,
- a `linalg.broadcast` + `linalg.min` sequence can be lowered to a
- `linalg.generic` with different affine maps for the two operands.
- """
- O[None] = BinaryFn.min_signed(lhs[None], rhs[None])
-
-
- at linalg_structured_op(op_class_name="PowFOp")
-def powf(
- lhs=TensorDef(T1),
- rhs=TensorDef(T1),
- O=TensorDef(T1, output=True),
-):
- """Takes the powf(lhs, rhs) between two inputs, elementwise. For powf(arg, 2) use `linalg.square`.
-
- Only applies to floating point values.
-
- The shapes and element types must be identical. The appropriate casts,
- broadcasts and reductions should be done previously to calling this op.
-
- This means reduction/broadcast/element cast semantics is explicit. Further
- passes can take that into account when lowering this code. For example,
- a `linalg.broadcast` + `linalg.powf` sequence can be lowered to a
- `linalg.generic` with different affine maps for the two operands.
- """
- O[None] = BinaryFn.powf(lhs[None], rhs[None])
-
-
- at linalg_structured_op
-def select(
- cond=TensorDef(U),
- lhs=TensorDef(T1),
- rhs=TensorDef(T1),
- O=TensorDef(T1, output=True),
-):
- """Chooses one value based on a binary condition supplied as its first operand.
-
- The shapes and element types must be identical. The appropriate casts,
- broadcasts and reductions should be done previously to calling this op.
-
- This means reduction/broadcast/element cast semantics is explicit. Further
- passes can take that into account when lowering this code. For example,
- a `linalg.broadcast` + `linalg.select` sequence can be lowered to a
- `linalg.generic` with different affine maps for the two operands.
- """
- O[None] = TernaryFn.select(cond[None], lhs[None], rhs[None])
-
-
@linalg_structured_op
def quantized_matmul(
A=TensorDef(T1, S.M, S.K),
>From 531149906419b326b737bfc2b6a75df15269be44 Mon Sep 17 00:00:00 2001
From: mabsar <mabsar at qti.qualcommm.com>
Date: Wed, 19 Aug 2026 06:07:43 -0700
Subject: [PATCH 2/2] address review comment.
---
.../mlir/Dialect/Linalg/IR/LinalgInterfaces.h | 1 +
.../Dialect/Linalg/IR/LinalgInterfaces.td | 9 +++++++++
.../Dialect/Linalg/IR/LinalgStructuredOps.td | 9 +++++----
mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp | 19 ++++++++++++-------
4 files changed, 27 insertions(+), 11 deletions(-)
diff --git a/mlir/include/mlir/Dialect/Linalg/IR/LinalgInterfaces.h b/mlir/include/mlir/Dialect/Linalg/IR/LinalgInterfaces.h
index e879452fb8a29..f343e1ceb4ac0 100644
--- a/mlir/include/mlir/Dialect/Linalg/IR/LinalgInterfaces.h
+++ b/mlir/include/mlir/Dialect/Linalg/IR/LinalgInterfaces.h
@@ -33,6 +33,7 @@ class GenericOp;
// Forward declaration needed by ElementwiseOpInterface.
enum class ElementwiseKind : uint32_t;
+enum class ElementwiseArityGroup : uint32_t;
namespace detail {
/// Implementation of the method that check if given operands
diff --git a/mlir/include/mlir/Dialect/Linalg/IR/LinalgInterfaces.td b/mlir/include/mlir/Dialect/Linalg/IR/LinalgInterfaces.td
index c56b83863fde8..859c345b18de4 100644
--- a/mlir/include/mlir/Dialect/Linalg/IR/LinalgInterfaces.td
+++ b/mlir/include/mlir/Dialect/Linalg/IR/LinalgInterfaces.td
@@ -816,6 +816,15 @@ def ElementwiseOpInterface : OpInterface<"ElementwiseOpInterface"> {
/*retType=*/"::mlir::linalg::ElementwiseKind",
/*methodName=*/"getElementwiseKind",
/*args=*/(ins)
+ >,
+ InterfaceMethod<
+ /*desc=*/[{
+ Returns the arity group of this elementwise operation
+ (unary, binary, or ternary).
+ }],
+ /*retType=*/"::mlir::linalg::ElementwiseArityGroup",
+ /*methodName=*/"getArityGroup",
+ /*args=*/(ins)
>
];
}
diff --git a/mlir/include/mlir/Dialect/Linalg/IR/LinalgStructuredOps.td b/mlir/include/mlir/Dialect/Linalg/IR/LinalgStructuredOps.td
index 12accf3e82d85..7219747d4d396 100644
--- a/mlir/include/mlir/Dialect/Linalg/IR/LinalgStructuredOps.td
+++ b/mlir/include/mlir/Dialect/Linalg/IR/LinalgStructuredOps.td
@@ -625,10 +625,6 @@ def ElementwiseOp : LinalgStructuredBase_Op<"elementwise", [
let hasFolder = 1;
let extraClassDeclaration = structuredOpsBaseDecls # [{
- /// Get the arity enum corresponding to the kind of op, e.g. if arg is
- /// `ElementwiseKind::add`, return `ElementwiseArityGroup::Binary`.
- static ElementwiseArityGroup getArityGroup(ElementwiseKind n);
-
/// Both user-specified and default indexing map will always depend on
/// the current Op instance.
static bool hasDynamicIndexingMaps() { return true; }
@@ -1322,6 +1318,11 @@ multiclass ElementwiseNamedOp<string mnemonic, string kind, int numRegionArgs,
ElementwiseKind $cppClass::getElementwiseKind() {
return ElementwiseKind::}] # kind # [{;
}
+ ElementwiseArityGroup $cppClass::getArityGroup() {
+ return ElementwiseArityGroup::}] #
+ !if(!eq(numRegionArgs, 2), "Unary",
+ !if(!eq(numRegionArgs, 3), "Binary", "Ternary")) # [{;
+ }
ParseResult $cppClass::parse(OpAsmParser &parser, OperationState &result) {
return ::parseNamedStructuredOp(parser, result,
$cppClass::getNumRegionArgs(),
diff --git a/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp b/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp
index ab55e57dedf7a..88ed9b48252af 100644
--- a/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp
+++ b/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp
@@ -5029,8 +5029,7 @@ void ElementwiseOp::print(OpAsmPrinter &p) {
p.printAttribute(getKindAttr());
SmallVector<StringRef, 3> elidedAttrs = {"operandSegmentSizes", "kind",
"indexing_maps"};
- unsigned arity =
- getArityGroupAsUInt(getArityGroupAndKind(getKind()).arityGroup);
+ unsigned arity = static_cast<unsigned>(getArityGroup());
unsigned numDims = getResultRank();
SmallVector<Attribute, 3> indexingMaps = llvm::map_to_vector<3>(
@@ -5116,6 +5115,10 @@ Speculation::Speculatability ElementwiseOp::getSpeculatability() {
ElementwiseKind ElementwiseOp::getElementwiseKind() { return getKind(); }
+ElementwiseArityGroup ElementwiseOp::getArityGroup() {
+ return getArityGroupAndKind(getKind()).arityGroup;
+}
+
//===----------------------------------------------------------------------===//
// Shared utilities for named elementwise ops (AddOp, SubOp, ExpOp, etc.)
//===----------------------------------------------------------------------===//
@@ -5134,16 +5137,18 @@ void buildElementwiseRegion(ImplicitLocOpBuilder &b, Block &block,
RegionBuilderHelper helper(b, block);
Value result;
- if (arityGroup == ElementwiseArityGroup::Unary) {
+ switch (arityGroup) {
+ case ElementwiseArityGroup::Unary:
result = helper.buildUnaryFn(fnKind.unaryFn, block.getArgument(0));
- } else if (arityGroup == ElementwiseArityGroup::Binary) {
+ break;
+ case ElementwiseArityGroup::Binary:
result = helper.buildBinaryFn(fnKind.binaryFn, block.getArgument(0),
block.getArgument(1), emitError);
- } else if (arityGroup == ElementwiseArityGroup::Ternary) {
+ break;
+ case ElementwiseArityGroup::Ternary:
result = helper.buildTernaryFn(fnKind.ternaryFn, block.getArgument(0),
block.getArgument(1), block.getArgument(2));
- } else {
- assert(false && "unhandled arity group");
+ break;
}
if (!result)
More information about the Mlir-commits
mailing list