[mlir][OpDSL] Split arithmetic functions.

[gysit]

Summary

Split arithmetic function into unary and binary functions. The revision prepares the introduction of unary and binary function attributes that work similar to type function attributes.

Depends On D120108

[aartbik]

LGTM for the sparse part

[aartbik]

I had a broader look at this at well. It is actually surprising that this change does not touch more tests. Do we need a few more just for this change?

mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp

156

you have two examples of binary, so use two example of unary
(in particular, because _exp could be mistaken for expression at first glance ;-)

mlir/python/mlir/dialects/linalg/opdsl/lang/comprehension.py

331

see in the other revision, should we consider cast_signed?

mlir/python/mlir/dialects/linalg/opdsl/lang/emitter.py

382

here and below, dump both lhs and rhs in error?

[gysit]

In D120109#3341892, @aartbik wrote:

I had a broader look at this at well. It is actually surprising that this change does not touch more tests. Do we need a few more just for this change?

Yes, I think one reason is that the revision before (https://reviews.llvm.org/D120108) - I have added you in meantime - refactors the function handling to make sure all function type share the same base logic below the hood. Additionally, the follow up revsion (https://reviews.llvm.org/D120110) finally introduces function attributes and introduces additional tests. In particular, it has a nice integration test for a configurable pooling operation.

I already ported back a test introduced in https://reviews.llvm.org/D120110 and will also add an additional test on the yaml generation side. There is also a number of existing tests that for example check the lowering of unary functions works end to end (generalize-named-polymorphic-ops.mlir lowers softmax to exps and logs).

mlir/python/mlir/dialects/linalg/opdsl/lang/comprehension.py

331

Yes, I will prepare a separate revsion!

[gysit]

Rebase and address comments.

mlir/docs/Dialects/Linalg/OpDSL.md

 The bulk of language consists of assignment expressions of the form above. The
 iteration dimension order is determined lexically based on the order encountered
 in the expression (following operator precedence if math operators are used).
 TODO: Introduce a directive to fix the dimension bindings.
 
 Reduction dimensions are inferred to be any dimensions on the RHS that are not
 on the LHS.
 
-A number of arithmetic functions are supported:
+A number of unary and binary arithmetic functions are supported:
 
-*   `ArithFn.add(a, b)` (also via overloading the binary `+` operator)
-*   `ArithFn.exp(a)`
-*   `ArithFn.log(a)`
-*   `ArithFn.mul(a, b)` (also via overloading the binary `*` operator)
-*   `ArithFn.max(a, b)`
-*   `ArithFn.min(a, b)`
-*   `ArithFn.sub(a, b)` (also via overloading the binary `-` operator)
-*   `ArithFn.max_unsigned(a, b)`
-*   `ArithFn.min_unsigned(a, b)`
+*   `BinaryFn.add(a, b)` (also via overloading the binary `+` operator)
+*   `BinaryFn.mul(a, b)` (also via overloading the binary `*` operator)
+*   `BinaryFn.max(a, b)`
+*   `BinaryFn.min(a, b)`
+*   `BinaryFn.sub(a, b)` (also via overloading the binary `-` operator)
+*   `BinaryFn.max_unsigned(a, b)`
+*   `BinaryFn.min_unsigned(a, b)`
+*   `UnaryFn.exp(a)`
+*   `UnaryFn.log(a)`
 
 As the integer types are signless, signedness is implement by different
 functions that treat integers as signed or unsigned values.
 
 A subset of the arithmetic functions are supported in reductions. These
 reduction functions can appear as the outermost function on the RHS:
 
 *   `ReduceFn.add` (also overloading the inplace `+=` on a LHS)

mlir/include/mlir/Dialect/Linalg/IR/LinalgNamedStructuredOps.yaml

structured_op: !LinalgStructuredOpConfig
   - parallel
   - parallel
   - reduction
   assignments:
   - !ScalarAssign
     arg: C
     value: !ScalarExpression
       scalar_fn:
-        kind: arith
+        kind: binary
         fn_name: add
         operands:
         - !ScalarExpression
           scalar_arg: C
         - !ScalarExpression
           scalar_fn:
-            kind: arith
+            kind: binary
             fn_name: mul
             operands:
             - !ScalarExpression
               scalar_fn:
                 kind: type
                 attr_name: cast
                 type_var: U
                 operands:
structured_op: !LinalgStructuredOpConfig
   - parallel
   - parallel
   - reduction
   assignments:
   - !ScalarAssign
     arg: C
     value: !ScalarExpression
       scalar_fn:
-        kind: arith
+        kind: binary
         fn_name: add
         operands:
         - !ScalarExpression
           scalar_arg: C
         - !ScalarExpression
           scalar_fn:
-            kind: arith
+            kind: binary
             fn_name: mul
             operands:
             - !ScalarExpression
               scalar_fn:
                 kind: type
                 fn_name: cast_unsigned
                 type_var: U
                 operands:
structured_op: !LinalgStructuredOpConfig
   - parallel
   - parallel
   - reduction
   assignments:
   - !ScalarAssign
     arg: C
     value: !ScalarExpression
       scalar_fn:
-        kind: arith
+        kind: binary
         fn_name: add
         operands:
         - !ScalarExpression
           scalar_arg: C
         - !ScalarExpression
           scalar_fn:
-            kind: arith
+            kind: binary
             fn_name: mul
             operands:
             - !ScalarExpression
               scalar_fn:
-                kind: arith
+                kind: binary
                 fn_name: sub
                 operands:
                 - !ScalarExpression
                   scalar_fn:
                     kind: type
                     fn_name: cast
                     type_var: U
                     operands:
                     - !ScalarExpression
                       scalar_arg: A
                 - !ScalarExpression
                   scalar_fn:
                     kind: type
                     fn_name: cast
                     type_var: U
                     operands:
                     - !ScalarExpression
                       scalar_arg: AZp
             - !ScalarExpression
               scalar_fn:
-                kind: arith
+                kind: binary
                 fn_name: sub
                 operands:
                 - !ScalarExpression
                   scalar_fn:
                     kind: type
                     fn_name: cast
                     type_var: U
                     operands:
structured_op: !LinalgStructuredOpConfig
   - parallel
   - parallel
   - reduction
   assignments:
   - !ScalarAssign
     arg: accum
     value: !ScalarExpression
       scalar_fn:
-        kind: arith
+        kind: binary
         fn_name: add
         operands:
         - !ScalarExpression
           scalar_arg: accum
         - !ScalarExpression
           scalar_fn:
-            kind: arith
+            kind: binary
             fn_name: mul
             operands:
             - !ScalarExpression
               scalar_fn:
                 kind: type
                 fn_name: cast
                 type_var: AccumType
                 operands:
structured_op: !LinalgStructuredOpConfig
   - parallel
   - parallel
   - reduction
   assignments:
   - !ScalarAssign
     arg: C
     value: !ScalarExpression
       scalar_fn:
-        kind: arith
+        kind: binary
         fn_name: add
         operands:
         - !ScalarExpression
           scalar_arg: C
         - !ScalarExpression
           scalar_fn:
-            kind: arith
+            kind: binary
             fn_name: mul
             operands:
             - !ScalarExpression
               scalar_fn:
                 kind: type
                 fn_name: cast
                 type_var: U
                 operands:
structured_op: !LinalgStructuredOpConfig
   - parallel
   - parallel
   - reduction
   assignments:
   - !ScalarAssign
     arg: C
     value: !ScalarExpression
       scalar_fn:
-        kind: arith
+        kind: binary
         fn_name: add
         operands:
         - !ScalarExpression
           scalar_arg: C
         - !ScalarExpression
           scalar_fn:
-            kind: arith
+            kind: binary
             fn_name: mul
             operands:
             - !ScalarExpression
               scalar_fn:
-                kind: arith
+                kind: binary
                 fn_name: sub
                 operands:
                 - !ScalarExpression
                   scalar_fn:
                     kind: type
                     fn_name: cast
                     type_var: U
                     operands:
                     - !ScalarExpression
                       scalar_arg: A
                 - !ScalarExpression
                   scalar_fn:
                     kind: type
                     fn_name: cast
                     type_var: U
                     operands:
                     - !ScalarExpression
                       scalar_arg: AZp
             - !ScalarExpression
               scalar_fn:
-                kind: arith
+                kind: binary
                 fn_name: sub
                 operands:
                 - !ScalarExpression
                   scalar_fn:
                     kind: type
                     fn_name: cast
                     type_var: U
                     operands:
structured_op: !LinalgStructuredOpConfig
   iterator_types:
   - parallel
   - reduction
   assignments:
   - !ScalarAssign
     arg: x
     value: !ScalarExpression
       scalar_fn:
-        kind: arith
+        kind: binary
         fn_name: add
         operands:
         - !ScalarExpression
           scalar_arg: x
         - !ScalarExpression
           scalar_fn:
-            kind: arith
+            kind: binary
             fn_name: mul
             operands:
             - !ScalarExpression
               scalar_fn:
                 kind: type
                 fn_name: cast
                 type_var: U
                 operands:
structured_op: !LinalgStructuredOpConfig
   iterator_types:
   - parallel
   - reduction
   assignments:
   - !ScalarAssign
     arg: x
     value: !ScalarExpression
       scalar_fn:
-        kind: arith
+        kind: binary
         fn_name: add
         operands:
         - !ScalarExpression
           scalar_arg: x
         - !ScalarExpression
           scalar_fn:
-            kind: arith
+            kind: binary
             fn_name: mul
             operands:
             - !ScalarExpression
               scalar_fn:
                 kind: type
                 fn_name: cast
                 type_var: U
                 operands:
structured_op: !LinalgStructuredOpConfig
   - parallel
   - parallel
   - reduction
   assignments:
   - !ScalarAssign
     arg: C
     value: !ScalarExpression
       scalar_fn:
-        kind: arith
+        kind: binary
         fn_name: add
         operands:
         - !ScalarExpression
           scalar_arg: C
         - !ScalarExpression
           scalar_fn:
-            kind: arith
+            kind: binary
             fn_name: mul
             operands:
             - !ScalarExpression
               scalar_fn:
                 kind: type
                 fn_name: cast
                 type_var: U
                 operands:
indexing_maps: !LinalgIndexingMapsConfig
     - affine_map<(d0)[s0] -> ()>
   iterator_types:
   - reduction
   assignments:
   - !ScalarAssign
     arg: C
     value: !ScalarExpression
       scalar_fn:
-        kind: arith
+        kind: binary
         fn_name: add
         operands:
         - !ScalarExpression
           scalar_arg: C
         - !ScalarExpression
           scalar_fn:
-            kind: arith
+            kind: binary
             fn_name: mul
             operands:
             - !ScalarExpression
               scalar_fn:
                 kind: type
                 fn_name: cast
                 type_var: U
                 operands:
structured_op: !LinalgStructuredOpConfig
   iterator_types:
   - parallel
   - reduction
   assignments:
   - !ScalarAssign
     arg: O
     value: !ScalarExpression
       scalar_fn:
-        kind: arith
+        kind: binary
         fn_name: add
         operands:
         - !ScalarExpression
           scalar_arg: O
         - !ScalarExpression
           scalar_fn:
-            kind: arith
+            kind: binary
             fn_name: mul
             operands:
             - !ScalarExpression
               scalar_fn:
                 kind: type
                 fn_name: cast
                 type_var: U
                 operands:
structured_op: !LinalgStructuredOpConfig
   - parallel
   - reduction
   - reduction
   assignments:
   - !ScalarAssign
     arg: O
     value: !ScalarExpression
       scalar_fn:
-        kind: arith
+        kind: binary
         fn_name: add
         operands:
         - !ScalarExpression
           scalar_arg: O
         - !ScalarExpression
           scalar_fn:
-            kind: arith
+            kind: binary
             fn_name: mul
             operands:
             - !ScalarExpression
               scalar_fn:
                 kind: type
                 fn_name: cast
                 type_var: U
                 operands:
structured_op: !LinalgStructuredOpConfig
   - reduction
   - reduction
   - reduction
   assignments:
   - !ScalarAssign
     arg: O
     value: !ScalarExpression
       scalar_fn:
-        kind: arith
+        kind: binary
         fn_name: add
         operands:
         - !ScalarExpression
           scalar_arg: O
         - !ScalarExpression
           scalar_fn:
-            kind: arith
+            kind: binary
             fn_name: mul
             operands:
             - !ScalarExpression
               scalar_fn:
                 kind: type
                 fn_name: cast
                 type_var: U
                 operands:
structured_op: !LinalgStructuredOpConfig
   - parallel
   - reduction
   - reduction
   assignments:
   - !ScalarAssign
     arg: O
     value: !ScalarExpression
       scalar_fn:
-        kind: arith
+        kind: binary
         fn_name: add
         operands:
         - !ScalarExpression
           scalar_arg: O
         - !ScalarExpression
           scalar_fn:
-            kind: arith
+            kind: binary
             fn_name: mul
             operands:
             - !ScalarExpression
               scalar_fn:
                 kind: type
                 fn_name: cast
                 type_var: U
                 operands:
structured_op: !LinalgStructuredOpConfig
   - reduction
   - reduction
   - reduction
   assignments:
   - !ScalarAssign
     arg: O
     value: !ScalarExpression
       scalar_fn:
-        kind: arith
+        kind: binary
         fn_name: add
         operands:
         - !ScalarExpression
           scalar_arg: O
         - !ScalarExpression
           scalar_fn:
-            kind: arith
+            kind: binary
             fn_name: mul
             operands:
             - !ScalarExpression
               scalar_fn:
                 kind: type
                 fn_name: cast
                 type_var: U
                 operands:
structured_op: !LinalgStructuredOpConfig
   - reduction
   - reduction
   - reduction
   assignments:
   - !ScalarAssign
     arg: O
     value: !ScalarExpression
       scalar_fn:
-        kind: arith
+        kind: binary
         fn_name: add
         operands:
         - !ScalarExpression
           scalar_arg: O
         - !ScalarExpression
           scalar_fn:
-            kind: arith
+            kind: binary
             fn_name: mul
             operands:
             - !ScalarExpression
               scalar_fn:
-                kind: arith
+                kind: binary
                 fn_name: sub
                 operands:
                 - !ScalarExpression
                   scalar_fn:
                     kind: type
                     fn_name: cast
                     type_var: U
                     operands:
                     - !ScalarExpression
                       scalar_arg: I
                 - !ScalarExpression
                   scalar_fn:
                     kind: type
                     fn_name: cast
                     type_var: U
                     operands:
                     - !ScalarExpression
                       scalar_arg: IZp
             - !ScalarExpression
               scalar_fn:
-                kind: arith
+                kind: binary
                 fn_name: sub
                 operands:
                 - !ScalarExpression
                   scalar_fn:
                     kind: type
                     fn_name: cast
                     type_var: U
                     operands:
structured_op: !LinalgStructuredOpConfig
   - reduction
   - reduction
   - reduction
   assignments:
   - !ScalarAssign
     arg: O
     value: !ScalarExpression
       scalar_fn:
-        kind: arith
+        kind: binary
         fn_name: add
         operands:
         - !ScalarExpression
           scalar_arg: O
         - !ScalarExpression
           scalar_fn:
-            kind: arith
+            kind: binary
             fn_name: mul
             operands:
             - !ScalarExpression
               scalar_fn:
                 kind: type
                 fn_name: cast
                 type_var: U
                 operands:
structured_op: !LinalgStructuredOpConfig
   - reduction
   - reduction
   - reduction
   assignments:
   - !ScalarAssign
     arg: O
     value: !ScalarExpression
       scalar_fn:
-        kind: arith
+        kind: binary
         fn_name: add
         operands:
         - !ScalarExpression
           scalar_arg: O
         - !ScalarExpression
           scalar_fn:
-            kind: arith
+            kind: binary
             fn_name: mul
             operands:
             - !ScalarExpression
               scalar_fn:
                 kind: type
                 fn_name: cast
                 type_var: U
                 operands:
structured_op: !LinalgStructuredOpConfig
   - parallel
   - parallel
   - reduction
   assignments:
   - !ScalarAssign
     arg: O
     value: !ScalarExpression
       scalar_fn:
-        kind: arith
+        kind: binary
         fn_name: add
         operands:
         - !ScalarExpression
           scalar_arg: O
         - !ScalarExpression
           scalar_fn:
-            kind: arith
+            kind: binary
             fn_name: mul
             operands:
             - !ScalarExpression
               scalar_fn:
                 kind: type
                 fn_name: cast
                 type_var: U
                 operands:
structured_op: !LinalgStructuredOpConfig
   - parallel
   - reduction
   - reduction
   assignments:
   - !ScalarAssign
     arg: O
     value: !ScalarExpression
       scalar_fn:
-        kind: arith
+        kind: binary
         fn_name: add
         operands:
         - !ScalarExpression
           scalar_arg: O
         - !ScalarExpression
           scalar_fn:
-            kind: arith
+            kind: binary
             fn_name: mul
             operands:
             - !ScalarExpression
               scalar_fn:
                 kind: type
                 fn_name: cast
                 type_var: U
                 operands:
structured_op: !LinalgStructuredOpConfig
   - parallel
   - reduction
   - reduction
   assignments:
   - !ScalarAssign
     arg: O
     value: !ScalarExpression
       scalar_fn:
-        kind: arith
+        kind: binary
         fn_name: add
         operands:
         - !ScalarExpression
           scalar_arg: O
         - !ScalarExpression
           scalar_fn:
-            kind: arith
+            kind: binary
             fn_name: mul
             operands:
             - !ScalarExpression
               scalar_fn:
-                kind: arith
+                kind: binary
                 fn_name: sub
                 operands:
                 - !ScalarExpression
                   scalar_fn:
                     kind: type
                     fn_name: cast
                     type_var: U
                     operands:
                     - !ScalarExpression
                       scalar_arg: I
                 - !ScalarExpression
                   scalar_fn:
                     kind: type
                     fn_name: cast
                     type_var: U
                     operands:
                     - !ScalarExpression
                       scalar_arg: IZp
             - !ScalarExpression
               scalar_fn:
-                kind: arith
+                kind: binary
                 fn_name: sub
                 operands:
                 - !ScalarExpression
                   scalar_fn:
                     kind: type
                     fn_name: cast
                     type_var: U
                     operands:
structured_op: !LinalgStructuredOpConfig
   - parallel
   - reduction
   - reduction
   assignments:
   - !ScalarAssign
     arg: O
     value: !ScalarExpression
       scalar_fn:
-        kind: arith
+        kind: binary
         fn_name: add
         operands:
         - !ScalarExpression
           scalar_arg: O
         - !ScalarExpression
           scalar_fn:
-            kind: arith
+            kind: binary
             fn_name: mul
             operands:
             - !ScalarExpression
               scalar_fn:
                 kind: type
                 fn_name: cast
                 type_var: U
                 operands:
structured_op: !LinalgStructuredOpConfig
   - parallel
   - reduction
   - reduction
   assignments:
   - !ScalarAssign
     arg: O
     value: !ScalarExpression
       scalar_fn:
-        kind: arith
+        kind: binary
         fn_name: add
         operands:
         - !ScalarExpression
           scalar_arg: O
         - !ScalarExpression
           scalar_fn:
-            kind: arith
+            kind: binary
             fn_name: mul
             operands:
             - !ScalarExpression
               scalar_fn:
-                kind: arith
+                kind: binary
                 fn_name: sub
                 operands:
                 - !ScalarExpression
                   scalar_fn:
                     kind: type
                     fn_name: cast
                     type_var: U
                     operands:
                     - !ScalarExpression
                       scalar_arg: I
                 - !ScalarExpression
                   scalar_fn:
                     kind: type
                     fn_name: cast
                     type_var: U
                     operands:
                     - !ScalarExpression
                       scalar_arg: IZp
             - !ScalarExpression
               scalar_fn:
-                kind: arith
+                kind: binary
                 fn_name: sub
                 operands:
                 - !ScalarExpression
                   scalar_fn:
                     kind: type
                     fn_name: cast
                     type_var: U
                     operands:
structured_op: !LinalgStructuredOpConfig
   - parallel
   - reduction
   - reduction
   assignments:
   - !ScalarAssign
     arg: O
     value: !ScalarExpression
       scalar_fn:
-        kind: arith
+        kind: binary
         fn_name: add
         operands:
         - !ScalarExpression
           scalar_arg: O
         - !ScalarExpression
           scalar_fn:
             kind: type
             fn_name: cast
structured_op: !LinalgStructuredOpConfig
   - parallel
   - reduction
   - reduction
   assignments:
   - !ScalarAssign
     arg: O
     value: !ScalarExpression
       scalar_fn:
-        kind: arith
+        kind: binary
         fn_name: max
         operands:
         - !ScalarExpression
           scalar_arg: O
         - !ScalarExpression
           scalar_fn:
             kind: type
             fn_name: cast
structured_op: !LinalgStructuredOpConfig
   - parallel
   - reduction
   - reduction
   assignments:
   - !ScalarAssign
     arg: O
     value: !ScalarExpression
       scalar_fn:
-        kind: arith
+        kind: binary
         fn_name: max_unsigned
         operands:
         - !ScalarExpression
           scalar_arg: O
         - !ScalarExpression
           scalar_fn:
             kind: type
             fn_name: cast_unsigned
structured_op: !LinalgStructuredOpConfig
   - parallel
   - reduction
   - reduction
   assignments:
   - !ScalarAssign
     arg: O
     value: !ScalarExpression
       scalar_fn:
-        kind: arith
+        kind: binary
         fn_name: max
         operands:
         - !ScalarExpression
           scalar_arg: O
         - !ScalarExpression
           scalar_fn:
             kind: type
             fn_name: cast
structured_op: !LinalgStructuredOpConfig
   - parallel
   - reduction
   - reduction
   assignments:
   - !ScalarAssign
     arg: O
     value: !ScalarExpression
       scalar_fn:
-        kind: arith
+        kind: binary
         fn_name: min
         operands:
         - !ScalarExpression
           scalar_arg: O
         - !ScalarExpression
           scalar_fn:
             kind: type
             fn_name: cast
structured_op: !LinalgStructuredOpConfig
   - parallel
   - reduction
   - reduction
   assignments:
   - !ScalarAssign
     arg: O
     value: !ScalarExpression
       scalar_fn:
-        kind: arith
+        kind: binary
         fn_name: min_unsigned
         operands:
         - !ScalarExpression
           scalar_arg: O
         - !ScalarExpression
           scalar_fn:
             kind: type
             fn_name: cast_unsigned
structured_op: !LinalgStructuredOpConfig
   - reduction
   - reduction
   - reduction
   assignments:
   - !ScalarAssign
     arg: O
     value: !ScalarExpression
       scalar_fn:
-        kind: arith
+        kind: binary
         fn_name: add
         operands:
         - !ScalarExpression
           scalar_arg: O
         - !ScalarExpression
           scalar_fn:
             kind: type
             fn_name: cast
structured_op: !LinalgStructuredOpConfig
   - reduction
   - reduction
   - reduction
   assignments:
   - !ScalarAssign
     arg: O
     value: !ScalarExpression
       scalar_fn:
-        kind: arith
+        kind: binary
         fn_name: max
         operands:
         - !ScalarExpression
           scalar_arg: O
         - !ScalarExpression
           scalar_fn:
             kind: type
             fn_name: cast
structured_op: !LinalgStructuredOpConfig
   - reduction
   - reduction
   - reduction
   assignments:
   - !ScalarAssign
     arg: O
     value: !ScalarExpression
       scalar_fn:
-        kind: arith
+        kind: binary
         fn_name: min
         operands:
         - !ScalarExpression
           scalar_arg: O
         - !ScalarExpression
           scalar_fn:
             kind: type
             fn_name: cast
- !ScalarAssign
     value: !ScalarExpression
       scalar_fn:
         kind: type
         fn_name: cast
         type_var: T
         operands:
         - !ScalarExpression
           scalar_fn:
-            kind: arith
+            kind: binary
             fn_name: add
             operands:
             - !ScalarExpression
               scalar_fn:
-                kind: arith
+                kind: binary
                 fn_name: mul
                 operands:
                 - !ScalarExpression
                   scalar_fn:
-                    kind: arith
+                    kind: binary
                     fn_name: add
                     operands:
                     - !ScalarExpression
                       scalar_fn:
                         kind: type
                         fn_name: cast
                         type_var: F64
                         operands:
                         - !ScalarExpression
                           scalar_const: '2147483647 : i64'
                     - !ScalarExpression
                       scalar_fn:
                         kind: type
                         fn_name: cast
                         type_var: F64
                         operands:
                         - !ScalarExpression
                           scalar_fn:
-                            kind: arith
+                            kind: binary
                             fn_name: add
                             operands:
                             - !ScalarExpression
                               scalar_fn:
-                                kind: arith
+                                kind: binary
                                 fn_name: mul
                                 operands:
                                 - !ScalarExpression
                                   scalar_fn:
-                                    kind: arith
+                                    kind: binary
                                     fn_name: add
                                     operands:
                                     - !ScalarExpression
                                       scalar_fn:
                                         kind: type
                                         fn_name: cast
                                         type_var: I32
                                         operands:
                                         - !ScalarExpression
                                           scalar_index: 1
                                     - !ScalarExpression
                                       scalar_fn:
-                                        kind: arith
+                                        kind: binary
                                         fn_name: add
                                         operands:
                                         - !ScalarExpression
                                           scalar_fn:
-                                            kind: arith
+                                            kind: binary
                                             fn_name: mul
                                             operands:
                                             - !ScalarExpression
                                               scalar_fn:
-                                                kind: arith
+                                                kind: binary
                                                 fn_name: add
                                                 operands:
                                                 - !ScalarExpression
                                                   scalar_fn:
                                                     kind: type
                                                     fn_name: cast
                                                     type_var: I32
                                                     operands:
value: !ScalarExpression
                                 kind: type
                                 fn_name: cast
                                 type_var: I32
                                 operands:
                                 - !ScalarExpression
                                   scalar_const: '12345 : i64'
                 - !ScalarExpression
                   scalar_fn:
-                    kind: arith
+                    kind: binary
                     fn_name: mul
                     operands:
                     - !ScalarExpression
                       scalar_fn:
-                        kind: arith
+                        kind: binary
                         fn_name: sub
                         operands:
                         - !ScalarExpression
                           scalar_arg: max
                         - !ScalarExpression
                           scalar_arg: min
                     - !ScalarExpression
                       scalar_fn:
structured_op: !LinalgStructuredOpConfig
   iterator_types:
   - parallel
   - parallel
   assignments:
   - !ScalarAssign
     arg: O
     value: !ScalarExpression
       scalar_fn:
-        kind: arith
+        kind: unary
         fn_name: log
         operands:
         - !ScalarExpression
           scalar_fn:
-            kind: arith
+            kind: binary
             fn_name: add
             operands:
             - !ScalarExpression
               scalar_fn:
                 kind: type
                 fn_name: cast
                 type_var: U
                 operands:
                 - !ScalarExpression
                   scalar_const: '1.000000e+00 : f64'
             - !ScalarExpression
               scalar_fn:
-                kind: arith
+                kind: unary
                 fn_name: exp
                 operands:
                 - !ScalarExpression
                   scalar_fn:
                     kind: type
                     fn_name: cast
                     type_var: U
                     operands:
                     - !ScalarExpression
                       scalar_arg: I

mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp

static LogicalResult foldMemRefCastInTiledLoopOp(TiledLoopOp op) {
   }
   return success(folded);
 }
 
 //===----------------------------------------------------------------------===//
 // Region builder helper.
 // TODO: Move this to a utility library.
 // The public methods on this class are referenced directly from generated code
-// and bind by name to math and type conversion functions in the DSL as:
-//   `arithfn__{fnName}`
-//   `typefn__{fnName}`
+// and bind by name to math functions in the DSL as:
+//   `unary__{fnName}`
+//   `binary__{fnName}`
 // Examples:
-//   `arithfn__add`
-//   `arithfn__mul`
-//   `typefn__cast`
+//   `binary__add`
+//   `binary__mul`
+//   `unary__exp`
+//   `unary__log`
 // The naming convention is intentional in order to match snake-cased DSL names.
 // See mlir-linalg-ods-yaml-gen.cpp for the code that mates to this class.
 //
 // Implementations of the math functions must be polymorphic over numeric types,
 // internally performing necessary casts. If the function application makes no
 // sense, then the only recourse is to assert and return nullptr. This can be
 // extended later if it becomes possible to fail construction of the region. The
 // invariant should be enforced at a higher level.
case TypeFn::cast:
       return cast(toType, operand, false);
     case TypeFn::cast_unsigned:
       return cast(toType, operand, true);
     }
     llvm_unreachable("unsupported type conversion function");
   }
 
   // NOLINTNEXTLINE(*-identifier-naming): externally called.
-  Value arithfn__add(Value lhs, Value rhs) {
+  Value binary__add(Value lhs, Value rhs) {
     OpBuilder builder = getBuilder();
     if (isFloatingPoint(lhs))
       return builder.create<arith::AddFOp>(lhs.getLoc(), lhs, rhs);
     if (isInteger(lhs))
       return builder.create<arith::AddIOp>(lhs.getLoc(), lhs, rhs);
     llvm_unreachable("unsupported non numeric type");
   }
 
   // NOLINTNEXTLINE(*-identifier-naming): externally called.
-  Value arithfn__exp(Value x) {
+  Value unary__exp(Value x) {
     OpBuilder builder = getBuilder();
     if (isFloatingPoint(x))
       return builder.create<math::ExpOp>(x.getLoc(), x);
     llvm_unreachable("unsupported non numeric type");
   }
 
   // NOLINTNEXTLINE(*-identifier-naming): externally called.
-  Value arithfn__log(Value x) {
+  Value unary__log(Value x) {
     OpBuilder builder = getBuilder();
     if (isFloatingPoint(x))
       return builder.create<math::LogOp>(x.getLoc(), x);
     llvm_unreachable("unsupported non numeric type");
   }
 
   // NOLINTNEXTLINE(*-identifier-naming): externally called.
-  Value arithfn__sub(Value lhs, Value rhs) {
+  Value binary__sub(Value lhs, Value rhs) {
     OpBuilder builder = getBuilder();
     if (isFloatingPoint(lhs))
       return builder.create<arith::SubFOp>(lhs.getLoc(), lhs, rhs);
     if (isInteger(lhs))
       return builder.create<arith::SubIOp>(lhs.getLoc(), lhs, rhs);
     llvm_unreachable("unsupported non numeric type");
   }
 
   // NOLINTNEXTLINE(*-identifier-naming): externally called.
-  Value arithfn__mul(Value lhs, Value rhs) {
+  Value binary__mul(Value lhs, Value rhs) {
     OpBuilder builder = getBuilder();
     if (isFloatingPoint(lhs))
       return builder.create<arith::MulFOp>(lhs.getLoc(), lhs, rhs);
     if (isInteger(lhs))
       return builder.create<arith::MulIOp>(lhs.getLoc(), lhs, rhs);
     llvm_unreachable("unsupported non numeric type");
   }
 
   // NOLINTNEXTLINE(*-identifier-naming): externally called.
-  Value arithfn__max(Value lhs, Value rhs) {
+  Value binary__max(Value lhs, Value rhs) {
     OpBuilder builder = getBuilder();
     if (isFloatingPoint(lhs))
       return builder.create<arith::MaxFOp>(lhs.getLoc(), lhs, rhs);
     if (isInteger(lhs))
       return builder.create<arith::MaxSIOp>(lhs.getLoc(), lhs, rhs);
     llvm_unreachable("unsupported non numeric type");
   }
 
   // NOLINTNEXTLINE(*-identifier-naming): externally called.
-  Value arithfn__max_unsigned(Value lhs, Value rhs) {
+  Value binary__max_unsigned(Value lhs, Value rhs) {
     OpBuilder builder = getBuilder();
     if (isFloatingPoint(lhs))
       return builder.create<arith::MaxFOp>(lhs.getLoc(), lhs, rhs);
     if (isInteger(lhs))
       return builder.create<arith::MaxUIOp>(lhs.getLoc(), lhs, rhs);
     llvm_unreachable("unsupported non numeric type");
   }
 
   // NOLINTNEXTLINE(*-identifier-naming): externally called.
-  Value arithfn__min(Value lhs, Value rhs) {
+  Value binary__min(Value lhs, Value rhs) {
     OpBuilder builder = getBuilder();
     if (isFloatingPoint(lhs))
       return builder.create<arith::MinFOp>(lhs.getLoc(), lhs, rhs);
     if (isInteger(lhs))
       return builder.create<arith::MinSIOp>(lhs.getLoc(), lhs, rhs);
     llvm_unreachable("unsupported non numeric type");
   }
 
   // NOLINTNEXTLINE(*-identifier-naming): externally called.
-  Value arithfn__min_unsigned(Value lhs, Value rhs) {
+  Value binary__min_unsigned(Value lhs, Value rhs) {
     OpBuilder builder = getBuilder();
     if (isFloatingPoint(lhs))
       return builder.create<arith::MinFOp>(lhs.getLoc(), lhs, rhs);
     if (isInteger(lhs))
       return builder.create<arith::MinUIOp>(lhs.getLoc(), lhs, rhs);
     llvm_unreachable("unsupported non numeric type");
   }
 

mlir/python/mlir/dialects/linalg/opdsl/lang/comprehension.py

def collect_scalar_uses(self, uses: Set["ScalarDef"]):
 
     def visit_scalar_def(expr: "TensorExpression"):
       if isinstance(expr, ScalarDef):
         uses.add(expr)
 
     self.visit_tensor_exprs(visit_scalar_def)
 
   def __add__(self, rhs: "TensorExpression") -> "TensorExpression":
-    return ArithFn.add(self, rhs)
+    return BinaryFn.add(self, rhs)
 
   def __mul__(self, rhs) -> "TensorExpression":
-    return ArithFn.mul(self, rhs)
+    return BinaryFn.mul(self, rhs)
 
   def __sub__(self, rhs) -> "TensorExpression":
-    return ArithFn.sub(self, rhs)
+    return BinaryFn.sub(self, rhs)
 
   def __hash__(self):
     return hash(id(self))
 
 
 class TensorUse(TensorExpression):
   """A used tensor represented by its (tensor_name, indices).
 
def _compute_reduce_dims(self, rhs: TensorExpression) -> Set[DimDef]:
     # dims and will be ordered as encountered on the rhs.
     rhs_dims = set()
     lhs_dims = set()
     rhs.collect_dim_uses(rhs_dims)
     self.collect_dim_uses(lhs_dims)
     return rhs_dims - lhs_dims
 
   def __iadd__(self, rhs: TensorExpression) -> "TensorReduceFn":
-    return ReduceFnUse(ArithFn.add, *self._compute_reduce_dims(rhs))(rhs)
+    return ReduceFnUse(BinaryFn.add, *self._compute_reduce_dims(rhs))(rhs)
 
   def __repr__(self):
     return (f"{self.operand_def.name}"
             f"[{', '.join([repr(i) for i in self.indices])}]")
 
 
 class TensorFn(TensorExpression):
   """Application of a tensor function."""
def __init__(self, reduce_use: "ReduceFnUse",
     self.args = args
 
   def to_scalar_expression(self) -> ScalarExpression:
     if self.lhs is None:
       raise ValueError(f"Cannot scalarize a TensorReduceFn that has not been "
                        f"bound to its lhs: {self}")
     full_args = [self.lhs.to_scalar_expression()
                 ] + [arg.to_scalar_expression() for arg in self.args]
-    return ScalarFn(FunctionKind.ARITH, self.reduce_use.arith_fn.fn_name, None,
-                    None, full_args).expr()
+    return ScalarFn(FunctionKind.BINARY, self.reduce_use.binary_fn.fn_name,
+                    None, None, full_args).expr()
 
   def visit_tensor_exprs(self, callback: Callable[["TensorExpression"], None]):
     for arg in self.args:
       arg.visit_tensor_exprs(callback)
 
   def __repr__(self):
     return f"{repr(self.reduce_use)}({', '.join(repr(a) for a in self.args)})"
 

 
 
 ###############################################################################
 # Function types and function definitions.
 ###############################################################################
 
 
 class FunctionKind(Enum):
-  ARITH = 0
-  TYPE = 1
+  UNARY = 0
+  BINARY = 1
+  TYPE = 2
 
 
-class TypeFnType:
-  """Type conversion function.
+class UnaryFnType:
+  """Unary function.
 
-  A type conversion function takes a target type and a tensor expression and
-  returns the casted tensor expression.
+  A unary function takes one tensor expression and returns the
+  function evaluation result.
   """
 
   def __init__(self, fn_name: str):
     self.fn_name = fn_name
 
-  def __call__(self, type_var: TypeVar, arg: TensorExpression) -> "TensorFn":
-    return TensorFn(FunctionKind.TYPE, self.fn_name, None, type_var, [arg])
+  def __call__(self, exp: TensorExpression) -> "TensorFn":
+    return TensorFn(FunctionKind.UNARY, self.fn_name, None, None, [exp])
 
   def __repr__(self):
     return f"{self.fn_name}"
 
 
-class TypeFn:
-  """Type conversion function namespace.
-
-  As the integer types are signless, signedness is implement by different cast
-  functions that treat integers as signed (`cast`) or unsigned
-  (`cast_unsigned`) values.
-
-  Examples:
-  - cast(I32 -> I64) -> `arith.ExtSIOp`
-  - cast_unsigned(I32 -> I64) -> `arith.ExtUIOp`
-  """
-  cast = TypeFnType("cast")
-  cast_unsigned = TypeFnType("cast_unsigned")
+class UnaryFn:
+  """Unary function namespace."""
+  exp = UnaryFnType("exp")
+  log = UnaryFnType("log")
 
 
-class ArithFnType:
-  """Arithmetic function.
+class BinaryFnType:
+  """Binary function.
 
-  An arithmetic function takes one ore more tensor expressions and returns the
+  A binary function takes two tensor expressions and returns the
   function evaluation result.
   """
 
   def __init__(self, fn_name: str):
     self.fn_name = fn_name
 
-  def __call__(self, *args) -> "TensorFn":
-    return TensorFn(FunctionKind.ARITH, self.fn_name, None, None, args)
+  def __call__(self, arg0: TensorExpression,
+               arg1: TensorExpression) -> "TensorFn":
+    return TensorFn(FunctionKind.BINARY, self.fn_name, None, None, [arg0, arg1])
 
   def __repr__(self):
     return f"{self.fn_name}"
 
 
-class ArithFn:
-  """Arithmetic function namespace.
+class BinaryFn:
+  """Binary function namespace.
 
   As the integer types are signless, signedness is implement by different
   functions that treat integers as signed or unsigned values.
 
   Examples:
   - max -> `arith.MaxSIOp`
   - max_unsinged -> `arith.MaxUIOp`
   """
-  add = ArithFnType("add")
-  exp = ArithFnType("exp")
-  log = ArithFnType("log")
-  mul = ArithFnType("mul")
-  max = ArithFnType("max")
-  min = ArithFnType("min")
-  sub = ArithFnType("sub")
-  max_unsigned = ArithFnType("max_unsigned")
-  min_unsigned = ArithFnType("min_unsigned")
+  add = BinaryFnType("add")
+  mul = BinaryFnType("mul")
+  max = BinaryFnType("max")
+  min = BinaryFnType("min")
+  sub = BinaryFnType("sub")
+  max_unsigned = BinaryFnType("max_unsigned")
+  min_unsigned = BinaryFnType("min_unsigned")
+
+
+class TypeFnType:
+  """Type conversion function.
+
+  A type conversion function takes a target type and a tensor expression and
+  returns the casted tensor expression.
+  """
+
+  def __init__(self, fn_name: str):
+    self.fn_name = fn_name
+
+  def __call__(self, type_var: TypeVar, arg: TensorExpression) -> "TensorFn":
+    return TensorFn(FunctionKind.TYPE, self.fn_name, None, type_var, [arg])
+
+  def __repr__(self):
+    return f"{self.fn_name}"
+
+
+class TypeFn:
+  """Type conversion function namespace.
+
+  As the integer types are signless, signedness is implement by different cast
+  functions that treat integers as signed (`cast`) or unsigned
+  (`cast_unsigned`) values.
+
+  Examples:
+  - cast(I32 -> I64) -> `arith.ExtSIOp`
+  - cast_unsigned(I32 -> I64) -> `arith.ExtUIOp`
+  """
+  cast = TypeFnType("cast")
+  cast_unsigned = TypeFnType("cast_unsigned")
 
 
 class ReduceFnUse:
   """Reduction function use.
 
   A reduction use specifies the reduction function and dimensions.
   """
 
-  def __init__(self, arith_fn: ArithFnType, *reduce_dims: DimDef):
-    self.arith_fn = arith_fn
+  def __init__(self, binary_fn: BinaryFnType, *reduce_dims: DimDef):
+    self.binary_fn = binary_fn
     self.reduce_dims = reduce_dims
 
-  def __call__(self, *args: TensorExpression):
+  def __call__(self, *args: TensorExpression) -> "TensorReduceFn":
     return TensorReduceFn(self, args)
 
   def __repr__(self):
-    return (f"reduce_{self.arith_fn.fn_name}"
+    return (f"reduce_{self.binary_fn.fn_name}"
             f"({', '.join(repr(d) for d in self.reduce_dims)})")
 
 
 class ReduceFnType:
   """Reduction function.
 
-  An arithmetic function that reduces its RHS into its LHS.
+  A binary function that reduces its RHS into its LHS.
   """
 
-  def __init__(self, arith_fn: ArithFnType):
-    if not isinstance(arith_fn, ArithFnType):
-      raise ValueError(f"Reduce expected a ArithFnType but got {arith_fn}")
-    self.arith_fn = arith_fn
+  def __init__(self, binary_fn: BinaryFnType):
+    if not isinstance(binary_fn, BinaryFnType):
+      raise ValueError(f"Reduce expected a BinaryFnType but got {binary_fn}")
+    self.binary_fn = binary_fn
 
   def __getitem__(self, reduce_dims: Tuple[DimDef]) -> ReduceFnUse:
-    return ReduceFnUse(self.arith_fn, *reduce_dims)
+    return ReduceFnUse(self.binary_fn, *reduce_dims)
 
   def __repr__(self):
-    return (f"reduce_{self.arith_fn.fn_name}")
+    return (f"reduce_{self.binary_fn.fn_name}")
 
 
 class ReduceFn:
-  add = ReduceFnType(ArithFn.add)
-  mul = ReduceFnType(ArithFn.mul)
-  max = ReduceFnType(ArithFn.max)
-  min = ReduceFnType(ArithFn.min)
-  max_unsigned = ReduceFnType(ArithFn.max_unsigned)
-  min_unsigned = ReduceFnType(ArithFn.min_unsigned)
+  add = ReduceFnType(BinaryFn.add)
+  mul = ReduceFnType(BinaryFn.mul)
+  max = ReduceFnType(BinaryFn.max)
+  min = ReduceFnType(BinaryFn.min)
+  max_unsigned = ReduceFnType(BinaryFn.max_unsigned)
+  min_unsigned = ReduceFnType(BinaryFn.min_unsigned)
 
 
 ###############################################################################
 # Operand definitions.
 ###############################################################################
 
 
 class OperandKind(Enum):

mlir/python/mlir/dialects/linalg/opdsl/lang/emitter.py

if expr.scalar_arg:
                          f"this structured op.")
     elif expr.scalar_const:
       value_attr = Attribute.parse(expr.scalar_const.value)
       return arith.ConstantOp(value_attr.type, value_attr).result
     elif expr.scalar_index:
       dim_attr = IntegerAttr.get(
           IntegerType.get_signless(64), expr.scalar_index.dim)
       return linalg.IndexOp(dim_attr).result
-    elif expr.scalar_fn and expr.scalar_fn.kind == FunctionKind.ARITH:
-      fn = self._get_function(f"_arithfn_{expr.scalar_fn.fn_name}")
+    elif expr.scalar_fn and expr.scalar_fn.kind is not FunctionKind.TYPE:
+      kind = expr.scalar_fn.kind.name.lower()
+      fn = self._get_function(f"_{kind}_{expr.scalar_fn.fn_name}")
       operand_values = [
           self.expression(operand) for operand in expr.scalar_fn.operands
       ]
       return fn(*operand_values)
-    elif expr.scalar_fn and expr.scalar_fn.kind == FunctionKind.TYPE:
+    elif expr.scalar_fn and expr.scalar_fn.kind is FunctionKind.TYPE:
+      kind = expr.scalar_fn.kind.name.lower()
       fn_name = expr.scalar_fn.fn_name
       if expr.scalar_fn.attr_name:
         fn_name = self.type_fn_attr_mapping[expr.scalar_fn.attr_name]
-      fn = self._get_function(f"_typefn_{fn_name}")
+      fn = self._get_function(f"_{kind}_{fn_name}")
       operand_value = self.expression(expr.scalar_fn.operands[0])
       return fn(expr.scalar_fn.type_var.name, operand_value)
     raise NotImplementedError(f"Unimplemented scalar body expression: {expr}")
 
   def yield_outputs(self, *output_names: str):
     output_values = []
     for n in output_names:
       try:
def _cast_to_floating_point(self, to_type: Type, operand: Value,
     from_width = _get_floating_point_width(operand_type)
     if to_width > from_width:
       return arith.ExtFOp(to_type, operand).result
     elif to_width < from_width:
       return arith.TruncFOp(to_type, operand).result
     raise ValueError(f"Unable to cast body expression from {operand_type} to "
                      f"{to_type}")
 
-  def _typefn_cast(self, type_var_name: str, operand: Value) -> Value:
+  def _type_cast(self, type_var_name: str, operand: Value) -> Value:
     return self._cast(type_var_name, operand, False)
 
-  def _typefn_cast_unsigned(self, type_var_name: str, operand: Value) -> Value:
+  def _type_cast_unsigned(self, type_var_name: str, operand: Value) -> Value:
     return self._cast(type_var_name, operand, True)
 
-  def _arithfn_add(self, lhs: Value, rhs: Value) -> Value:
-    if _is_floating_point_type(lhs.type):
-      return arith.AddFOp(lhs, rhs).result
-    if _is_integer_type(lhs.type) or _is_index_type(lhs.type):
-      return arith.AddIOp(lhs, rhs).result
-    raise NotImplementedError("Unsupported 'add' operand: {lhs}")
-
-  def _arithfn_exp(self, x: Value) -> Value:
+  def _unary_exp(self, x: Value) -> Value:
     if _is_floating_point_type(x.type):
       return math.ExpOp(x).result
     raise NotImplementedError("Unsupported 'exp' operand: {x}")
 
-  def _arithfn_log(self, x: Value) -> Value:
+  def _unary_log(self, x: Value) -> Value:
     if _is_floating_point_type(x.type):
       return math.LogOp(x).result
     raise NotImplementedError("Unsupported 'log' operand: {x}")
 
-  def _arithfn_sub(self, lhs: Value, rhs: Value) -> Value:
+  def _binary_add(self, lhs: Value, rhs: Value) -> Value:
+    if _is_floating_point_type(lhs.type):
+      return arith.AddFOp(lhs, rhs).result
+    if _is_integer_type(lhs.type) or _is_index_type(lhs.type):
+      return arith.AddIOp(lhs, rhs).result
+    raise NotImplementedError("Unsupported 'add' operands: {lhs}, {rhs}")
+
+  def _binary_sub(self, lhs: Value, rhs: Value) -> Value:
     if _is_floating_point_type(lhs.type):
       return arith.SubFOp(lhs, rhs).result
     if _is_integer_type(lhs.type) or _is_index_type(lhs.type):
       return arith.SubIOp(lhs, rhs).result
-    raise NotImplementedError("Unsupported 'sub' operand: {lhs}")
+    raise NotImplementedError("Unsupported 'sub' operands: {lhs}, {rhs}")
 
-  def _arithfn_mul(self, lhs: Value, rhs: Value) -> Value:
+  def _binary_mul(self, lhs: Value, rhs: Value) -> Value:
     if _is_floating_point_type(lhs.type):
       return arith.MulFOp(lhs, rhs).result
     if _is_integer_type(lhs.type) or _is_index_type(lhs.type):
       return arith.MulIOp(lhs, rhs).result
-    raise NotImplementedError("Unsupported 'mul' operand: {lhs}")
+    raise NotImplementedError("Unsupported 'mul' operands: {lhs}, {rhs}")
 
-  def _arithfn_max(self, lhs: Value, rhs: Value) -> Value:
+  def _binary_max(self, lhs: Value, rhs: Value) -> Value:
     if _is_floating_point_type(lhs.type):
       return arith.MaxFOp(lhs, rhs).result
     if _is_integer_type(lhs.type) or _is_index_type(lhs.type):
       return arith.MaxSIOp(lhs, rhs).result
-    raise NotImplementedError("Unsupported 'max' operand: {lhs}")
+    raise NotImplementedError("Unsupported 'max' operands: {lhs}, {rhs}")
 
-  def _arithfn_max_unsigned(self, lhs: Value, rhs: Value) -> Value:
+  def _binary_max_unsigned(self, lhs: Value, rhs: Value) -> Value:
     if _is_floating_point_type(lhs.type):
       return arith.MaxFOp(lhs, rhs).result
     if _is_integer_type(lhs.type) or _is_index_type(lhs.type):
       return arith.MaxUIOp(lhs, rhs).result
-    raise NotImplementedError("Unsupported 'max_unsigned' operand: {lhs}")
+    raise NotImplementedError(
+        "Unsupported 'max_unsigned' operands: {lhs}, {rhs}")
 
-  def _arithfn_min(self, lhs: Value, rhs: Value) -> Value:
+  def _binary_min(self, lhs: Value, rhs: Value) -> Value:
     if _is_floating_point_type(lhs.type):
       return arith.MinFOp(lhs, rhs).result
     if _is_integer_type(lhs.type) or _is_index_type(lhs.type):
       return arith.MinSIOp(lhs, rhs).result
-    raise NotImplementedError("Unsupported 'min' operand: {lhs}")
+    raise NotImplementedError("Unsupported 'min' operands: {lhs}, {rhs}")
 
-  def _arithfn_min_unsigned(self, lhs: Value, rhs: Value) -> Value:
+  def _binary_min_unsigned(self, lhs: Value, rhs: Value) -> Value:
     if _is_floating_point_type(lhs.type):
       return arith.MinFOp(lhs, rhs).result
     if _is_integer_type(lhs.type) or _is_index_type(lhs.type):
       return arith.MinUIOp(lhs, rhs).result
-    raise NotImplementedError("Unsupported 'min_unsigned' operand: {lhs}")
+    raise NotImplementedError(
+        "Unsupported 'min_unsigned' operands: {lhs}, {rhs}")
 
 
 def _infer_structured_outs(
     op_config: LinalgStructuredOpConfig,
     in_arg_defs: Sequence[OperandDefConfig], ins: Sequence[Value],
     out_arg_defs: Sequence[OperandDefConfig],
     outs: Union[Sequence[Value], OpResultList]) -> Tuple[ValueList, List[Type]]:
   """Infers implicit outs and output types.

mlir/python/mlir/dialects/linalg/opdsl/ops/core_named_ops.py

def soft_plus_2d(
     I=TensorDef(T, S.M, S.N), O=TensorDef(U, S.M, S.N, output=True)):
   """Implements the soft plus operator.
 
   Numeric casting is performed on the input operand, promoting it to the same
   data type as the accumulator/output.
   """
   domain(D.m, D.n)
   O[D.m, D.n] = \
-      ArithFn.log(TypeFn.cast(U, const(1.0)) + ArithFn.exp(TypeFn.cast(U, I[D.m, D.n])))
+      UnaryFn.log(TypeFn.cast(U, const(1.0)) + UnaryFn.exp(TypeFn.cast(U, I[D.m, D.n])))

mlir/test/Integration/Dialect/SparseTensor/taco/tools/mlir_pytaco.py

def _gather_input_accesses_index_vars(
     expr: IndexExpr,
     input_accesses: List[Access],
 ) -> None:
   """Collects Access nodes."""
   if isinstance(expr, Access) and expr not in input_accesses:
     input_accesses.append(expr)
 
 
-def _op_to_callable(op: _BinaryOp) -> lang.ArithFnType:
+def _op_to_callable(op: _BinaryOp) -> lang.BinaryFnType:
   """Returns the linalg dialect function object for the given operation."""
   op_to_callable = {
-      operator.add: lang.ArithFn.add,
-      operator.sub: lang.ArithFn.sub,
-      operator.mul: lang.ArithFn.mul,
+      operator.add: lang.BinaryFn.add,
+      operator.sub: lang.BinaryFn.sub,
+      operator.mul: lang.BinaryFn.mul,
   }
   return op_to_callable[op]
 
 
 @dataclasses.dataclass(frozen=True)
 class _BinaryExpr(IndexExpr):
   """The representation for a binary operation.
 

mlir/test/mlir-linalg-ods-gen/test-linalg-ods-yaml-gen.yaml

structured_op: !LinalgStructuredOpConfig
   iterator_types:
   - parallel
   - parallel
   assignments:
   - !ScalarAssign
     arg: O
     value: !ScalarExpression
       scalar_fn:
-        kind: arith
+        kind: binary
         fn_name: add
         operands:
         - !ScalarExpression
           scalar_fn:
             kind: type
             attr_name: cast
             type_var: T
             operands:

 #  IMPL-NEXT:    if (auto attr = castIter->getValue().dyn_cast<TypeFnAttr>())
 #  IMPL-NEXT:      castVal = attr.getValue();
 #  IMPL-NEXT:  }
 
 #       IMPL:  Value [[VAL0:[a-z0-9]+]] = helper.constant("42 : i64");
 #   IMPL-DAG:  Value [[VAL1:[a-z0-9]+]] = helper.buildTypeFn(castVal, block.getArgument(0).getType(), [[VAL0]]);
 #   IMPL-DAG:  Value [[VAL2:[a-z0-9]+]] = helper.index(1);
 #   IMPL-DAG:  Value [[VAL3:[a-z0-9]+]] = helper.buildTypeFn(castVal, block.getArgument(0).getType(), [[VAL2]]);
-#   IMPL-DAG:  Value [[VAL4:[a-z0-9]+]] = helper.arithfn__add([[VAL1]], [[VAL3]]);
+#   IMPL-DAG:  Value [[VAL4:[a-z0-9]+]] = helper.binary__add([[VAL1]], [[VAL3]]);
 
 
 # @linalg_structured_op
 # def test2(I=TensorDef(T, S.M, S.N),
 #           O=TensorDef(T, S.M, S.N, output=True),
 #           strides=IndexAttrDef(S.SM, S.SN, default=[1, 2])):
 #   """Title.
 

 
 #       IMPL:  Test3Op::iterator_types() {
 #  IMPL-NEXT:    int64_t rank = getRank(getOutputOperand(0));
 
 #       IMPL:  Test3Op::indexing_maps() {
 #  IMPL-NEXT:    MLIRContext *context = getContext();
 #  IMPL-NEXT:    AffineMap scalarMap = AffineMap::get(getNumParallelLoops(), 0, context);
 #  IMPL-NEXT:    AffineMap tensorMap = AffineMap::getMultiDimIdentityMap(
+
+
+# @linalg_structured_op
+# def test4(O=TensorDef(T, S.M, S.N, output=True)):
+#   """Title.
+
+#   Detailed description.
+#   """
+#   O[D.m, D.n] = BinaryFn.add(UnaryFn.exp(O[D.m, D.n]), O[D.m, D.n])
+
+--- !LinalgOpConfig
+metadata: !LinalgOpMetadata
+  name: test4
+  cpp_class_name: Test4Op
+  doc: |-
+    Title.
+
+    Detailed description.
+structured_op: !LinalgStructuredOpConfig
+  args:
+  - !LinalgOperandDefConfig
+    name: O
+    kind: output_tensor
+    type_var: T
+    shape_map: affine_map<()[s0, s1] -> (s0, s1)>
+  indexing_maps: !LinalgIndexingMapsConfig
+    static_indexing_maps:
+    - affine_map<(d0, d1)[s0, s1] -> (d0, d1)>
+  iterator_types:
+  - parallel
+  - parallel
+  assignments:
+  - !ScalarAssign
+    arg: O
+    value: !ScalarExpression
+      scalar_fn:
+        kind: binary
+        fn_name: add
+        operands:
+        - !ScalarExpression
+          scalar_fn:
+            kind: unary
+            fn_name: exp
+            operands:
+            - !ScalarExpression
+              scalar_arg: O
+        - !ScalarExpression
+          scalar_arg: O
+
+# IMPL-LABEL:  void Test4Op::regionBuilder(ImplicitLocOpBuilder &b,
+#  IMPL-NEXT:    Block &block, ArrayRef<NamedAttribute> attrs)
+
+#       IMPL:  Value [[VAL0:[a-z0-9]+]] = helper.unary__exp(block.getArgument(0))
+#  IMPL-NEXT:  Value [[VAL1:[a-z0-9]+]] = helper.binary__add([[VAL0]], block.getArgument(0))
+#  IMPL-NEXT:  yields.push_back([[VAL1]])

mlir/test/python/dialects/linalg/opdsl/assignments.py

 
 
 # CHECK: ---
 # CHECK-LABEL: constants
 # CHECK: assignments:
 # CHECK:  -
 # CHECK:    arg: O
 # CHECK:      scalar_fn:
-# CHECK:        kind: arith
+# CHECK:        kind: binary
 # CHECK:        fn_name: sub
 # CHECK:        operands:
 # CHECK:          scalar_fn:
-# CHECK:            kind: arith
+# CHECK:            kind: binary
 # CHECK:            fn_name: add
 # CHECK:            operands:
 # CHECK:              scalar_fn:
-# CHECK:                kind: type
-# CHECK:                type_var: T
-# CHECK:                operands:
-# CHECK:                  scalar_const: '3.1415926535897931 : f64'
+# CHECK:                kind: unary
+# CHECK:                fn_name: exp
+# CHECK:                operands:
+# CHECK:                  scalar_fn:
+# CHECK:                    kind: type
+# CHECK:                    type_var: T
+# CHECK:                    operands:
+# CHECK:                      scalar_const: '3.1415926535897931 : f64'
 # CHECK:              scalar_fn:
 # CHECK:                kind: type
 # CHECK:                fn_name: cast
 # CHECK:                type_var: T
 # CHECK:                operands:
 # CHECK:                  scalar_const: '42 : i64'
 # CHECK:          scalar_fn:
 # CHECK:            kind: type
 # CHECK:            fn_name: cast
 # CHECK:            type_var: T
 # CHECK:            operands:
 # CHECK:              scalar_const: '1.{{[0]*}}e+03 : f64'
 @linalg_structured_op
 def constants(O=TensorDef(T, S.M, S.K, output=True)):
   pi = TypeFn.cast(T, const(3.1415926535897931))
   cst42 = TypeFn.cast(T, const(42))
   cst1000 = TypeFn.cast(T, const(1e+3))
-  O[D.m, D.n] = pi + cst42 - cst1000
-
+  O[D.m, D.n] = UnaryFn.exp(pi) + cst42 - cst1000
 
 # CHECK: ---
 # CHECK-LABEL: indices
 # CHECK: assignments:
 # CHECK:  -
 # CHECK:    arg: O
 # CHECK:      scalar_fn:
-# CHECK:        kind: arith
+# CHECK:        kind: binary
 # CHECK:        fn_name: add
 # CHECK:        operands:
 # CHECK:          scalar_index: 1
 # CHECK:          scalar_index: 0
 @linalg_structured_op
 def indices(O=TensorDef(T, S.M, S.K, output=True)):
   O[D.m, D.n] = index(D.n) + index(D.m)
 

mlir/test/python/dialects/linalg/opdsl/emit_misc.py

def fill_rng_poly(
   scaling = (max - min) * inv_range
   O[D.m, D.n] = TypeFn.cast(T,
                             (offset + TypeFn.cast(F64, rand2)) * scaling + min)
 
 
 @linalg_structured_op
 def soft_plus_poly(
     I=TensorDef(T, S.M, S.N), O=TensorDef(U, S.M, S.N, output=True)):
-  O[D.m, D.n] = ArithFn.log(
-      TypeFn.cast(U, const(1.0)) + TypeFn.cast(U, ArithFn.exp(I[D.m, D.n])))
+  O[D.m, D.n] = UnaryFn.log(
+      TypeFn.cast(U, const(1.0)) + TypeFn.cast(U, UnaryFn.exp(I[D.m, D.n])))
 
 
 @linalg_structured_op(op_name="custom_op_name")
 def non_default_op_name(I=TensorDef(T, S.N), O=TensorDef(T, S.N, output=True)):
   O[D.n] = I[D.n]
 
 
 with Context() as ctx, Location.unknown():

mlir/tools/mlir-linalg-ods-gen/mlir-linalg-ods-yaml-gen.cpp

 };
 
 struct LinalgIndexingMapsConfig {
   Optional<SmallVector<SerializedAffineMap>> staticIndexingMaps;
 };
 
 struct ScalarExpression;
 
-enum class ScalarFnKind { Arith, Type };
+enum class ScalarFnKind { Unary, Binary, Type };
 
 struct ScalarFn {
   ScalarFnKind kind;
   Optional<std::string> fnName;
   Optional<std::string> attrName;
   Optional<std::string> typeVar;
   // NOTE: This must be of arity 1, but to break the self-referential cycle,
   // we use a heap allocated vector.
static void mapping(IO &io, ScalarExpression &info) {
     io.mapOptional("scalar_fn", info.scalarFn);
   }
 };
 
 /// Scalar function kind enum.
 template <>
 struct ScalarEnumerationTraits<ScalarFnKind> {
   static void enumeration(IO &io, ScalarFnKind &value) {
-    io.enumCase(value, "arith", ScalarFnKind::Arith);
+    io.enumCase(value, "unary", ScalarFnKind::Unary);
+    io.enumCase(value, "binary", ScalarFnKind::Binary);
     io.enumCase(value, "type", ScalarFnKind::Type);
   }
 };
 
 /// A scalar expression that evaluates a named function.
 /// Functions are generally "math" level and type polymorphic. Builtin
 /// functions include:
 ///   - `add(lhs, rhs)`
for (LinalgOperandDef &arg : args) {
         if (expression.index) {
           // Access an iteration index.
           std::string cppIdent = llvm::formatv("value{0}", ++localCounter);
           stmts.push_back(llvm::formatv("Value {0} = helper.index({1});",
                                         cppIdent, *expression.index));
           return cppIdent;
         }
         if (expression.scalarFn &&
-            expression.scalarFn->kind == ScalarFnKind::Arith) {
+            expression.scalarFn->kind != ScalarFnKind::Type) {
           // Apply function.
           // Recursively generate operands.
           SmallVector<std::string> operandCppValues;
           for (ScalarExpression &operand : expression.scalarFn->operands) {
             auto operandCppValue = generateExpression(operand);
             if (!operandCppValue)
               return None;
             operandCppValues.push_back(*operandCppValue);
           }
+
+          std::string prefix = expression.scalarFn->kind == ScalarFnKind::Unary
+                                   ? "unary"
+                                   : "binary";
           std::string cppIdent = llvm::formatv("value{0}", ++localCounter);
           stmts.push_back(
-              llvm::formatv("Value {0} = helper.arithfn__{1}({2});", cppIdent,
-                            expression.scalarFn->fnName,
+              llvm::formatv("Value {0} = helper.{1}__{2}({3});", cppIdent,
+                            prefix, expression.scalarFn->fnName,
                             interleaveToString(operandCppValues, ", ")));
           return cppIdent;
         }
         if (expression.scalarFn &&
             expression.scalarFn->kind == ScalarFnKind::Type) {
           // Symbolic cast.
           // Operands must be arity 1.
           if (expression.scalarFn->operands.size() != 1) {