[mlir][OpDSL] Refactor function handling.

[gysit]

Summary

Prepare the OpDSL function handling to introduce more function classes. A follow up commit will split ArithFn into UnaryFn and BinaryFn. This revision prepares the split by adding a function kind enum to handle different function types using a single class on the various levels of the stack (for example, there is now one TensorFn and one ScalarFn).

Depends On D119718

[gysit]

Rebase.

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

structured_op: !LinalgStructuredOpConfig
   iterator_types:
   - parallel
   - parallel
   - reduction
   assignments:
   - !ScalarAssign
     arg: C
     value: !ScalarExpression
-      arith_fn:
+      scalar_fn:
+        kind: arith
         fn_name: add
         operands:
         - !ScalarExpression
           scalar_arg: C
         - !ScalarExpression
-          arith_fn:
+          scalar_fn:
+            kind: arith
             fn_name: mul
             operands:
             - !ScalarExpression
-              type_fn:
+              scalar_fn:
+                kind: type
+                attr_name: cast
                 type_var: U
                 operands:
                 - !ScalarExpression
                   scalar_arg: A
-                attr_name: cast
             - !ScalarExpression
-              type_fn:
+              scalar_fn:
+                kind: type
+                attr_name: cast
                 type_var: U
                 operands:
                 - !ScalarExpression
                   scalar_arg: B
-                attr_name: cast
 --- !LinalgOpConfig
 metadata: !LinalgOpMetadata
   name: matmul_unsigned
   cpp_class_name: MatmulUnsignedOp
   doc: |-
     Performs an unsigned matrix multiplication of two 2D inputs.
 
     Numeric casting is performed on the operands to the inner multiply, promoting
structured_op: !LinalgStructuredOpConfig
   iterator_types:
   - parallel
   - parallel
   - reduction
   assignments:
   - !ScalarAssign
     arg: C
     value: !ScalarExpression
-      arith_fn:
+      scalar_fn:
+        kind: arith
         fn_name: add
         operands:
         - !ScalarExpression
           scalar_arg: C
         - !ScalarExpression
-          arith_fn:
+          scalar_fn:
+            kind: arith
             fn_name: mul
             operands:
             - !ScalarExpression
-              type_fn:
+              scalar_fn:
+                kind: type
+                fn_name: cast_unsigned
                 type_var: U
                 operands:
                 - !ScalarExpression
                   scalar_arg: A
-                fn_name: cast_unsigned
             - !ScalarExpression
-              type_fn:
+              scalar_fn:
+                kind: type
+                fn_name: cast_unsigned
                 type_var: U
                 operands:
                 - !ScalarExpression
                   scalar_arg: B
-                fn_name: cast_unsigned
 --- !LinalgOpConfig
 metadata: !LinalgOpMetadata
   name: quantized_matmul
   cpp_class_name: QuantizedMatmulOp
   doc: |-
     Performs a matrix multiplication of two 2D inputs.
 
     Numeric casting is performed on the operands to the inner multiply, promoting
structured_op: !LinalgStructuredOpConfig
   iterator_types:
   - parallel
   - parallel
   - reduction
   assignments:
   - !ScalarAssign
     arg: C
     value: !ScalarExpression
-      arith_fn:
+      scalar_fn:
+        kind: arith
         fn_name: add
         operands:
         - !ScalarExpression
           scalar_arg: C
         - !ScalarExpression
-          arith_fn:
+          scalar_fn:
+            kind: arith
             fn_name: mul
             operands:
             - !ScalarExpression
-              arith_fn:
+              scalar_fn:
+                kind: arith
                 fn_name: sub
                 operands:
                 - !ScalarExpression
-                  type_fn:
+                  scalar_fn:
+                    kind: type
+                    fn_name: cast
                     type_var: U
                     operands:
                     - !ScalarExpression
                       scalar_arg: A
-                    fn_name: cast
                 - !ScalarExpression
-                  type_fn:
+                  scalar_fn:
+                    kind: type
+                    fn_name: cast
                     type_var: U
                     operands:
                     - !ScalarExpression
                       scalar_arg: AZp
-                    fn_name: cast
             - !ScalarExpression
-              arith_fn:
+              scalar_fn:
+                kind: arith
                 fn_name: sub
                 operands:
                 - !ScalarExpression
-                  type_fn:
+                  scalar_fn:
+                    kind: type
+                    fn_name: cast
                     type_var: U
                     operands:
                     - !ScalarExpression
                       scalar_arg: B
-                    fn_name: cast
                 - !ScalarExpression
-                  type_fn:
+                  scalar_fn:
+                    kind: type
+                    fn_name: cast
                     type_var: U
                     operands:
                     - !ScalarExpression
                       scalar_arg: BZp
-                    fn_name: cast
 --- !LinalgOpConfig
 metadata: !LinalgOpMetadata
   name: mmt4d
   cpp_class_name: Mmt4DOp
   doc: |-
     Performs a matrix-matrix-transpose multiplication of two 4D inputs.
 
     Differences from linalg.matmul:
structured_op: !LinalgStructuredOpConfig
   - reduction
   - parallel
   - parallel
   - reduction
   assignments:
   - !ScalarAssign
     arg: accum
     value: !ScalarExpression
-      arith_fn:
+      scalar_fn:
+        kind: arith
         fn_name: add
         operands:
         - !ScalarExpression
           scalar_arg: accum
         - !ScalarExpression
-          arith_fn:
+          scalar_fn:
+            kind: arith
             fn_name: mul
             operands:
             - !ScalarExpression
-              type_fn:
+              scalar_fn:
+                kind: type
+                fn_name: cast
                 type_var: AccumType
                 operands:
                 - !ScalarExpression
                   scalar_arg: lhs
-                fn_name: cast
             - !ScalarExpression
-              type_fn:
+              scalar_fn:
+                kind: type
+                fn_name: cast
                 type_var: AccumType
                 operands:
                 - !ScalarExpression
                   scalar_arg: rhs
-                fn_name: cast
 --- !LinalgOpConfig
 metadata: !LinalgOpMetadata
   name: batch_matmul
   cpp_class_name: BatchMatmulOp
   doc: |-
     Performs a batched matrix multiplication of two 3D inputs.
 
     Numeric casting is performed on the operands to the inner multiply, promoting
structured_op: !LinalgStructuredOpConfig
   - parallel
   - parallel
   - parallel
   - reduction
   assignments:
   - !ScalarAssign
     arg: C
     value: !ScalarExpression
-      arith_fn:
+      scalar_fn:
+        kind: arith
         fn_name: add
         operands:
         - !ScalarExpression
           scalar_arg: C
         - !ScalarExpression
-          arith_fn:
+          scalar_fn:
+            kind: arith
             fn_name: mul
             operands:
             - !ScalarExpression
-              type_fn:
+              scalar_fn:
+                kind: type
+                fn_name: cast
                 type_var: U
                 operands:
                 - !ScalarExpression
                   scalar_arg: A
-                fn_name: cast
             - !ScalarExpression
-              type_fn:
+              scalar_fn:
+                kind: type
+                fn_name: cast
                 type_var: U
                 operands:
                 - !ScalarExpression
                   scalar_arg: B
-                fn_name: cast
 --- !LinalgOpConfig
 metadata: !LinalgOpMetadata
   name: quantized_batch_matmul
   cpp_class_name: QuantizedBatchMatmulOp
   doc: |-
     Performs a batched matrix multiplication of two 3D inputs.
 
     Numeric casting is performed on the operands to the inner multiply, promoting
structured_op: !LinalgStructuredOpConfig
   - parallel
   - parallel
   - parallel
   - reduction
   assignments:
   - !ScalarAssign
     arg: C
     value: !ScalarExpression
-      arith_fn:
+      scalar_fn:
+        kind: arith
         fn_name: add
         operands:
         - !ScalarExpression
           scalar_arg: C
         - !ScalarExpression
-          arith_fn:
+          scalar_fn:
+            kind: arith
             fn_name: mul
             operands:
             - !ScalarExpression
-              arith_fn:
+              scalar_fn:
+                kind: arith
                 fn_name: sub
                 operands:
                 - !ScalarExpression
-                  type_fn:
+                  scalar_fn:
+                    kind: type
+                    fn_name: cast
                     type_var: U
                     operands:
                     - !ScalarExpression
                       scalar_arg: A
-                    fn_name: cast
                 - !ScalarExpression
-                  type_fn:
+                  scalar_fn:
+                    kind: type
+                    fn_name: cast
                     type_var: U
                     operands:
                     - !ScalarExpression
                       scalar_arg: AZp
-                    fn_name: cast
             - !ScalarExpression
-              arith_fn:
+              scalar_fn:
+                kind: arith
                 fn_name: sub
                 operands:
                 - !ScalarExpression
-                  type_fn:
+                  scalar_fn:
+                    kind: type
+                    fn_name: cast
                     type_var: U
                     operands:
                     - !ScalarExpression
                       scalar_arg: B
-                    fn_name: cast
                 - !ScalarExpression
-                  type_fn:
+                  scalar_fn:
+                    kind: type
+                    fn_name: cast
                     type_var: U
                     operands:
                     - !ScalarExpression
                       scalar_arg: BZp
-                    fn_name: cast
 --- !LinalgOpConfig
 metadata: !LinalgOpMetadata
   name: matvec
   cpp_class_name: MatvecOp
   doc: |-
     Performs a matrix-vector multiplication.
 
     Numeric casting is performed on the operands to the inner multiply, promoting
indexing_maps: !LinalgIndexingMapsConfig
     - affine_map<(d0, d1)[s0, s1] -> (d0)>
   iterator_types:
   - parallel
   - reduction
   assignments:
   - !ScalarAssign
     arg: x
     value: !ScalarExpression
-      arith_fn:
+      scalar_fn:
+        kind: arith
         fn_name: add
         operands:
         - !ScalarExpression
           scalar_arg: x
         - !ScalarExpression
-          arith_fn:
+          scalar_fn:
+            kind: arith
             fn_name: mul
             operands:
             - !ScalarExpression
-              type_fn:
+              scalar_fn:
+                kind: type
+                fn_name: cast
                 type_var: U
                 operands:
                 - !ScalarExpression
                   scalar_arg: A
-                fn_name: cast
             - !ScalarExpression
-              type_fn:
+              scalar_fn:
+                kind: type
+                fn_name: cast
                 type_var: U
                 operands:
                 - !ScalarExpression
                   scalar_arg: y
-                fn_name: cast
 --- !LinalgOpConfig
 metadata: !LinalgOpMetadata
   name: vecmat
   cpp_class_name: VecmatOp
   doc: |-
     Performs a vector-matrix multiplication.
 
     Numeric casting is performed on the operands to the inner multiply, promoting
indexing_maps: !LinalgIndexingMapsConfig
     - affine_map<(d0, d1)[s0, s1] -> (d0)>
   iterator_types:
   - parallel
   - reduction
   assignments:
   - !ScalarAssign
     arg: x
     value: !ScalarExpression
-      arith_fn:
+      scalar_fn:
+        kind: arith
         fn_name: add
         operands:
         - !ScalarExpression
           scalar_arg: x
         - !ScalarExpression
-          arith_fn:
+          scalar_fn:
+            kind: arith
             fn_name: mul
             operands:
             - !ScalarExpression
-              type_fn:
+              scalar_fn:
+                kind: type
+                fn_name: cast
                 type_var: U
                 operands:
                 - !ScalarExpression
                   scalar_arg: y
-                fn_name: cast
             - !ScalarExpression
-              type_fn:
+              scalar_fn:
+                kind: type
+                fn_name: cast
                 type_var: U
                 operands:
                 - !ScalarExpression
                   scalar_arg: A
-                fn_name: cast
 --- !LinalgOpConfig
 metadata: !LinalgOpMetadata
   name: batch_matvec
   cpp_class_name: BatchMatvecOp
   doc: |-
     Performs a batched matrix-vector multiplication.
 
     Numeric casting is performed on the operands to the inner multiply, promoting
structured_op: !LinalgStructuredOpConfig
   iterator_types:
   - parallel
   - parallel
   - reduction
   assignments:
   - !ScalarAssign
     arg: C
     value: !ScalarExpression
-      arith_fn:
+      scalar_fn:
+        kind: arith
         fn_name: add
         operands:
         - !ScalarExpression
           scalar_arg: C
         - !ScalarExpression
-          arith_fn:
+          scalar_fn:
+            kind: arith
             fn_name: mul
             operands:
             - !ScalarExpression
-              type_fn:
+              scalar_fn:
+                kind: type
+                fn_name: cast
                 type_var: U
                 operands:
                 - !ScalarExpression
                   scalar_arg: A
-                fn_name: cast
             - !ScalarExpression
-              type_fn:
+              scalar_fn:
+                kind: type
+                fn_name: cast
                 type_var: U
                 operands:
                 - !ScalarExpression
                   scalar_arg: B
-                fn_name: cast
 --- !LinalgOpConfig
 metadata: !LinalgOpMetadata
   name: dot
   cpp_class_name: DotOp
   doc: |-
     Performs a dot product of two vectors to a scalar result.
 
     Numeric casting is performed on the operands to the inner multiply, promoting
indexing_maps: !LinalgIndexingMapsConfig
     - affine_map<(d0)[s0] -> (d0)>
     - affine_map<(d0)[s0] -> ()>
   iterator_types:
   - reduction
   assignments:
   - !ScalarAssign
     arg: C
     value: !ScalarExpression
-      arith_fn:
+      scalar_fn:
+        kind: arith
         fn_name: add
         operands:
         - !ScalarExpression
           scalar_arg: C
         - !ScalarExpression
-          arith_fn:
+          scalar_fn:
+            kind: arith
             fn_name: mul
             operands:
             - !ScalarExpression
-              type_fn:
+              scalar_fn:
+                kind: type
+                fn_name: cast
                 type_var: U
                 operands:
                 - !ScalarExpression
                   scalar_arg: A
-                fn_name: cast
             - !ScalarExpression
-              type_fn:
+              scalar_fn:
+                kind: type
+                fn_name: cast
                 type_var: U
                 operands:
                 - !ScalarExpression
                   scalar_arg: B
-                fn_name: cast
 --- !LinalgOpConfig
 metadata: !LinalgOpMetadata
   name: conv_1d
   cpp_class_name: Conv1DOp
   doc: |-
     Performs 1-D convolution with no channels.
 
     Numeric casting is performed on the operands to the inner multiply, promoting
indexing_maps: !LinalgIndexingMapsConfig
     - affine_map<(d0, d1)[s0, s1] -> (d0)>
   iterator_types:
   - parallel
   - reduction
   assignments:
   - !ScalarAssign
     arg: O
     value: !ScalarExpression
-      arith_fn:
+      scalar_fn:
+        kind: arith
         fn_name: add
         operands:
         - !ScalarExpression
           scalar_arg: O
         - !ScalarExpression
-          arith_fn:
+          scalar_fn:
+            kind: arith
             fn_name: mul
             operands:
             - !ScalarExpression
-              type_fn:
+              scalar_fn:
+                kind: type
+                fn_name: cast
                 type_var: U
                 operands:
                 - !ScalarExpression
                   scalar_arg: I
-                fn_name: cast
             - !ScalarExpression
-              type_fn:
+              scalar_fn:
+                kind: type
+                fn_name: cast
                 type_var: U
                 operands:
                 - !ScalarExpression
                   scalar_arg: K
-                fn_name: cast
 --- !LinalgOpConfig
 metadata: !LinalgOpMetadata
   name: conv_2d
   cpp_class_name: Conv2DOp
   doc: |-
     Performs 2-D convolution with no channels.
 
     Numeric casting is performed on the operands to the inner multiply, promoting
structured_op: !LinalgStructuredOpConfig
   - parallel
   - parallel
   - reduction
   - reduction
   assignments:
   - !ScalarAssign
     arg: O
     value: !ScalarExpression
-      arith_fn:
+      scalar_fn:
+        kind: arith
         fn_name: add
         operands:
         - !ScalarExpression
           scalar_arg: O
         - !ScalarExpression
-          arith_fn:
+          scalar_fn:
+            kind: arith
             fn_name: mul
             operands:
             - !ScalarExpression
-              type_fn:
+              scalar_fn:
+                kind: type
+                fn_name: cast
                 type_var: U
                 operands:
                 - !ScalarExpression
                   scalar_arg: I
-                fn_name: cast
             - !ScalarExpression
-              type_fn:
+              scalar_fn:
+                kind: type
+                fn_name: cast
                 type_var: U
                 operands:
                 - !ScalarExpression
                   scalar_arg: K
-                fn_name: cast
 --- !LinalgOpConfig
 metadata: !LinalgOpMetadata
   name: conv_3d
   cpp_class_name: Conv3DOp
   doc: |-
     Performs 3-D convolution with no channels.
 
     Numeric casting is performed on the operands to the inner multiply, promoting
structured_op: !LinalgStructuredOpConfig
   - parallel
   - reduction
   - reduction
   - reduction
   assignments:
   - !ScalarAssign
     arg: O
     value: !ScalarExpression
-      arith_fn:
+      scalar_fn:
+        kind: arith
         fn_name: add
         operands:
         - !ScalarExpression
           scalar_arg: O
         - !ScalarExpression
-          arith_fn:
+          scalar_fn:
+            kind: arith
             fn_name: mul
             operands:
             - !ScalarExpression
-              type_fn:
+              scalar_fn:
+                kind: type
+                fn_name: cast
                 type_var: U
                 operands:
                 - !ScalarExpression
                   scalar_arg: I
-                fn_name: cast
             - !ScalarExpression
-              type_fn:
+              scalar_fn:
+                kind: type
+                fn_name: cast
                 type_var: U
                 operands:
                 - !ScalarExpression
                   scalar_arg: K
-                fn_name: cast
 --- !LinalgOpConfig
 metadata: !LinalgOpMetadata
   name: conv_1d_nwc_wcf
   cpp_class_name: Conv1DNwcWcfOp
   doc: |-
     Performs 1-D convolution.
 
     Numeric casting is performed on the operands to the inner multiply, promoting
structured_op: !LinalgStructuredOpConfig
   - parallel
   - parallel
   - reduction
   - reduction
   assignments:
   - !ScalarAssign
     arg: O
     value: !ScalarExpression
-      arith_fn:
+      scalar_fn:
+        kind: arith
         fn_name: add
         operands:
         - !ScalarExpression
           scalar_arg: O
         - !ScalarExpression
-          arith_fn:
+          scalar_fn:
+            kind: arith
             fn_name: mul
             operands:
             - !ScalarExpression
-              type_fn:
+              scalar_fn:
+                kind: type
+                fn_name: cast
                 type_var: U
                 operands:
                 - !ScalarExpression
                   scalar_arg: I
-                fn_name: cast
             - !ScalarExpression
-              type_fn:
+              scalar_fn:
+                kind: type
+                fn_name: cast
                 type_var: U
                 operands:
                 - !ScalarExpression
                   scalar_arg: K
-                fn_name: cast
 --- !LinalgOpConfig
 metadata: !LinalgOpMetadata
   name: conv_2d_nhwc_hwcf
   cpp_class_name: Conv2DNhwcHwcfOp
   doc: |-
     Performs 2-D convolution.
 
     Layout:
structured_op: !LinalgStructuredOpConfig
   - parallel
   - reduction
   - reduction
   - reduction
   assignments:
   - !ScalarAssign
     arg: O
     value: !ScalarExpression
-      arith_fn:
+      scalar_fn:
+        kind: arith
         fn_name: add
         operands:
         - !ScalarExpression
           scalar_arg: O
         - !ScalarExpression
-          arith_fn:
+          scalar_fn:
+            kind: arith
             fn_name: mul
             operands:
             - !ScalarExpression
-              type_fn:
+              scalar_fn:
+                kind: type
+                fn_name: cast
                 type_var: U
                 operands:
                 - !ScalarExpression
                   scalar_arg: I
-                fn_name: cast
             - !ScalarExpression
-              type_fn:
+              scalar_fn:
+                kind: type
+                fn_name: cast
                 type_var: U
                 operands:
                 - !ScalarExpression
                   scalar_arg: K
-                fn_name: cast
 --- !LinalgOpConfig
 metadata: !LinalgOpMetadata
   name: conv_2d_nhwc_hwcf_q
   cpp_class_name: Conv2DNhwcHwcfQOp
   doc: |-
     Performs 2-D convolution with zero point offsets.
 
     Layout:
structured_op: !LinalgStructuredOpConfig
   - parallel
   - reduction
   - reduction
   - reduction
   assignments:
   - !ScalarAssign
     arg: O
     value: !ScalarExpression
-      arith_fn:
+      scalar_fn:
+        kind: arith
         fn_name: add
         operands:
         - !ScalarExpression
           scalar_arg: O
         - !ScalarExpression
-          arith_fn:
+          scalar_fn:
+            kind: arith
             fn_name: mul
             operands:
             - !ScalarExpression
-              arith_fn:
+              scalar_fn:
+                kind: arith
                 fn_name: sub
                 operands:
                 - !ScalarExpression
-                  type_fn:
+                  scalar_fn:
+                    kind: type
+                    fn_name: cast
                     type_var: U
                     operands:
                     - !ScalarExpression
                       scalar_arg: I
-                    fn_name: cast
                 - !ScalarExpression
-                  type_fn:
+                  scalar_fn:
+                    kind: type
+                    fn_name: cast
                     type_var: U
                     operands:
                     - !ScalarExpression
                       scalar_arg: IZp
-                    fn_name: cast
             - !ScalarExpression
-              arith_fn:
+              scalar_fn:
+                kind: arith
                 fn_name: sub
                 operands:
                 - !ScalarExpression
-                  type_fn:
+                  scalar_fn:
+                    kind: type
+                    fn_name: cast
                     type_var: U
                     operands:
                     - !ScalarExpression
                       scalar_arg: K
-                    fn_name: cast
                 - !ScalarExpression
-                  type_fn:
+                  scalar_fn:
+                    kind: type
+                    fn_name: cast
                     type_var: U
                     operands:
                     - !ScalarExpression
                       scalar_arg: KZp
-                    fn_name: cast
 --- !LinalgOpConfig
 metadata: !LinalgOpMetadata
   name: conv_2d_nchw_fchw
   cpp_class_name: Conv2DNchwFchwOp
   doc: |-
     Performs 2-D convolution.
 
     Layout:
structured_op: !LinalgStructuredOpConfig
   - parallel
   - reduction
   - reduction
   - reduction
   assignments:
   - !ScalarAssign
     arg: O
     value: !ScalarExpression
-      arith_fn:
+      scalar_fn:
+        kind: arith
         fn_name: add
         operands:
         - !ScalarExpression
           scalar_arg: O
         - !ScalarExpression
-          arith_fn:
+          scalar_fn:
+            kind: arith
             fn_name: mul
             operands:
             - !ScalarExpression
-              type_fn:
+              scalar_fn:
+                kind: type
+                fn_name: cast
                 type_var: U
                 operands:
                 - !ScalarExpression
                   scalar_arg: I
-                fn_name: cast
             - !ScalarExpression
-              type_fn:
+              scalar_fn:
+                kind: type
+                fn_name: cast
                 type_var: U
                 operands:
                 - !ScalarExpression
                   scalar_arg: K
-                fn_name: cast
 --- !LinalgOpConfig
 metadata: !LinalgOpMetadata
   name: conv_3d_ndhwc_dhwcf
   cpp_class_name: Conv3DNdhwcDhwcfOp
   doc: |-
     Performs 3-D convolution.
 
     Numeric casting is performed on the operands to the inner multiply, promoting
structured_op: !LinalgStructuredOpConfig
   - reduction
   - reduction
   - reduction
   - reduction
   assignments:
   - !ScalarAssign
     arg: O
     value: !ScalarExpression
-      arith_fn:
+      scalar_fn:
+        kind: arith
         fn_name: add
         operands:
         - !ScalarExpression
           scalar_arg: O
         - !ScalarExpression
-          arith_fn:
+          scalar_fn:
+            kind: arith
             fn_name: mul
             operands:
             - !ScalarExpression
-              type_fn:
+              scalar_fn:
+                kind: type
+                fn_name: cast
                 type_var: U
                 operands:
                 - !ScalarExpression
                   scalar_arg: I
-                fn_name: cast
             - !ScalarExpression
-              type_fn:
+              scalar_fn:
+                kind: type
+                fn_name: cast
                 type_var: U
                 operands:
                 - !ScalarExpression
                   scalar_arg: K
-                fn_name: cast
 --- !LinalgOpConfig
 metadata: !LinalgOpMetadata
   name: depthwise_conv_1d_nwc_wc
   cpp_class_name: DepthwiseConv1DNwcWcOp
   doc: |-
     Performs depth-wise 1-D convolution.
 
     Numeric casting is performed on the operands to the inner multiply, promoting
structured_op: !LinalgStructuredOpConfig
   - parallel
   - parallel
   - parallel
   - reduction
   assignments:
   - !ScalarAssign
     arg: O
     value: !ScalarExpression
-      arith_fn:
+      scalar_fn:
+        kind: arith
         fn_name: add
         operands:
         - !ScalarExpression
           scalar_arg: O
         - !ScalarExpression
-          arith_fn:
+          scalar_fn:
+            kind: arith
             fn_name: mul
             operands:
             - !ScalarExpression
-              type_fn:
+              scalar_fn:
+                kind: type
+                fn_name: cast
                 type_var: U
                 operands:
                 - !ScalarExpression
                   scalar_arg: I
-                fn_name: cast
             - !ScalarExpression
-              type_fn:
+              scalar_fn:
+                kind: type
+                fn_name: cast
                 type_var: U
                 operands:
                 - !ScalarExpression
                   scalar_arg: K
-                fn_name: cast
 --- !LinalgOpConfig
 metadata: !LinalgOpMetadata
   name: depthwise_conv_2d_nhwc_hwc
   cpp_class_name: DepthwiseConv2DNhwcHwcOp
   doc: |-
     Performs depth-wise 2-D convolution.
 
     Numeric casting is performed on the operands to the inner multiply, promoting
structured_op: !LinalgStructuredOpConfig
   - parallel
   - parallel
   - reduction
   - reduction
   assignments:
   - !ScalarAssign
     arg: O
     value: !ScalarExpression
-      arith_fn:
+      scalar_fn:
+        kind: arith
         fn_name: add
         operands:
         - !ScalarExpression
           scalar_arg: O
         - !ScalarExpression
-          arith_fn:
+          scalar_fn:
+            kind: arith
             fn_name: mul
             operands:
             - !ScalarExpression
-              type_fn:
+              scalar_fn:
+                kind: type
+                fn_name: cast
                 type_var: U
                 operands:
                 - !ScalarExpression
                   scalar_arg: I
-                fn_name: cast
             - !ScalarExpression
-              type_fn:
+              scalar_fn:
+                kind: type
+                fn_name: cast
                 type_var: U
                 operands:
                 - !ScalarExpression
                   scalar_arg: K
-                fn_name: cast
 --- !LinalgOpConfig
 metadata: !LinalgOpMetadata
   name: depthwise_conv_2d_nhwc_hwc_q
   cpp_class_name: DepthwiseConv2DNhwcHwcQOp
   doc: |-
     Performs depth-wise 2-D convolution.
 
     Numeric casting is performed on the operands to the inner multiply, promoting
structured_op: !LinalgStructuredOpConfig
   - parallel
   - parallel
   - reduction
   - reduction
   assignments:
   - !ScalarAssign
     arg: O
     value: !ScalarExpression
-      arith_fn:
+      scalar_fn:
+        kind: arith
         fn_name: add
         operands:
         - !ScalarExpression
           scalar_arg: O
         - !ScalarExpression
-          arith_fn:
+          scalar_fn:
+            kind: arith
             fn_name: mul
             operands:
             - !ScalarExpression
-              arith_fn:
+              scalar_fn:
+                kind: arith
                 fn_name: sub
                 operands:
                 - !ScalarExpression
-                  type_fn:
+                  scalar_fn:
+                    kind: type
+                    fn_name: cast
                     type_var: U
                     operands:
                     - !ScalarExpression
                       scalar_arg: I
-                    fn_name: cast
                 - !ScalarExpression
-                  type_fn:
+                  scalar_fn:
+                    kind: type
+                    fn_name: cast
                     type_var: U
                     operands:
                     - !ScalarExpression
                       scalar_arg: IZp
-                    fn_name: cast
             - !ScalarExpression
-              arith_fn:
+              scalar_fn:
+                kind: arith
                 fn_name: sub
                 operands:
                 - !ScalarExpression
-                  type_fn:
+                  scalar_fn:
+                    kind: type
+                    fn_name: cast
                     type_var: U
                     operands:
                     - !ScalarExpression
                       scalar_arg: K
-                    fn_name: cast
                 - !ScalarExpression
-                  type_fn:
+                  scalar_fn:
+                    kind: type
+                    fn_name: cast
                     type_var: U
                     operands:
                     - !ScalarExpression
                       scalar_arg: KZp
-                    fn_name: cast
 --- !LinalgOpConfig
 metadata: !LinalgOpMetadata
   name: depthwise_conv_2d_nhwc_hwcm
   cpp_class_name: DepthwiseConv2DNhwcHwcmOp
   doc: |-
     Performs depth-wise 2-D convolution.
 
     Numeric casting is performed on the operands to the inner multiply, promoting
structured_op: !LinalgStructuredOpConfig
   - parallel
   - parallel
   - reduction
   - reduction
   assignments:
   - !ScalarAssign
     arg: O
     value: !ScalarExpression
-      arith_fn:
+      scalar_fn:
+        kind: arith
         fn_name: add
         operands:
         - !ScalarExpression
           scalar_arg: O
         - !ScalarExpression
-          arith_fn:
+          scalar_fn:
+            kind: arith
             fn_name: mul
             operands:
             - !ScalarExpression
-              type_fn:
+              scalar_fn:
+                kind: type
+                fn_name: cast
                 type_var: U
                 operands:
                 - !ScalarExpression
                   scalar_arg: I
-                fn_name: cast
             - !ScalarExpression
-              type_fn:
+              scalar_fn:
+                kind: type
+                fn_name: cast
                 type_var: U
                 operands:
                 - !ScalarExpression
                   scalar_arg: K
-                fn_name: cast
 --- !LinalgOpConfig
 metadata: !LinalgOpMetadata
   name: depthwise_conv_2d_nhwc_hwcm_q
   cpp_class_name: DepthwiseConv2DNhwcHwcmQOp
   doc: |-
     Performs depth-wise 2-D convolution.
 
     Numeric casting is performed on the operands to the inner multiply, promoting
structured_op: !LinalgStructuredOpConfig
   - parallel
   - parallel
   - reduction
   - reduction
   assignments:
   - !ScalarAssign
     arg: O
     value: !ScalarExpression
-      arith_fn:
+      scalar_fn:
+        kind: arith
         fn_name: add
         operands:
         - !ScalarExpression
           scalar_arg: O
         - !ScalarExpression
-          arith_fn:
+          scalar_fn:
+            kind: arith
             fn_name: mul
             operands:
             - !ScalarExpression
-              arith_fn:
+              scalar_fn:
+                kind: arith
                 fn_name: sub
                 operands:
                 - !ScalarExpression
-                  type_fn:
+                  scalar_fn:
+                    kind: type
+                    fn_name: cast
                     type_var: U
                     operands:
                     - !ScalarExpression
                       scalar_arg: I
-                    fn_name: cast
                 - !ScalarExpression
-                  type_fn:
+                  scalar_fn:
+                    kind: type
+                    fn_name: cast
                     type_var: U
                     operands:
                     - !ScalarExpression
                       scalar_arg: IZp
-                    fn_name: cast
             - !ScalarExpression
-              arith_fn:
+              scalar_fn:
+                kind: arith
                 fn_name: sub
                 operands:
                 - !ScalarExpression
-                  type_fn:
+                  scalar_fn:
+                    kind: type
+                    fn_name: cast
                     type_var: U
                     operands:
                     - !ScalarExpression
                       scalar_arg: K
-                    fn_name: cast
                 - !ScalarExpression
-                  type_fn:
+                  scalar_fn:
+                    kind: type
+                    fn_name: cast
                     type_var: U
                     operands:
                     - !ScalarExpression
                       scalar_arg: KZp
-                    fn_name: cast
 --- !LinalgOpConfig
 metadata: !LinalgOpMetadata
   name: pooling_nhwc_sum
   cpp_class_name: PoolingNhwcSumOp
   doc: |-
     Performs sum pooling.
 
     Numeric casting is performed on the input operand, promoting it to the same
structured_op: !LinalgStructuredOpConfig
   - parallel
   - parallel
   - reduction
   - reduction
   assignments:
   - !ScalarAssign
     arg: O
     value: !ScalarExpression
-      arith_fn:
+      scalar_fn:
+        kind: arith
         fn_name: add
         operands:
         - !ScalarExpression
           scalar_arg: O
         - !ScalarExpression
-          type_fn:
+          scalar_fn:
+            kind: type
+            fn_name: cast
             type_var: U
             operands:
             - !ScalarExpression
               scalar_arg: I
-            fn_name: cast
 --- !LinalgOpConfig
 metadata: !LinalgOpMetadata
   name: pooling_nhwc_max
   cpp_class_name: PoolingNhwcMaxOp
   doc: |-
     Performs max pooling.
 
     Numeric casting is performed on the input operand, promoting it to the same
structured_op: !LinalgStructuredOpConfig
   - parallel
   - parallel
   - reduction
   - reduction
   assignments:
   - !ScalarAssign
     arg: O
     value: !ScalarExpression
-      arith_fn:
+      scalar_fn:
+        kind: arith
         fn_name: max
         operands:
         - !ScalarExpression
           scalar_arg: O
         - !ScalarExpression
-          type_fn:
+          scalar_fn:
+            kind: type
+            fn_name: cast
             type_var: U
             operands:
             - !ScalarExpression
               scalar_arg: I
-            fn_name: cast
 --- !LinalgOpConfig
 metadata: !LinalgOpMetadata
   name: pooling_nhwc_max_unsigned
   cpp_class_name: PoolingNhwcMaxUnsignedOp
   doc: |-
     Performs unsigned max pooling.
 
     Numeric casting is performed on the input operand, promoting it to the same
structured_op: !LinalgStructuredOpConfig
   - parallel
   - parallel
   - reduction
   - reduction
   assignments:
   - !ScalarAssign
     arg: O
     value: !ScalarExpression
-      arith_fn:
+      scalar_fn:
+        kind: arith
         fn_name: max_unsigned
         operands:
         - !ScalarExpression
           scalar_arg: O
         - !ScalarExpression
-          type_fn:
+          scalar_fn:
+            kind: type
+            fn_name: cast_unsigned
             type_var: U
             operands:
             - !ScalarExpression
               scalar_arg: I
-            fn_name: cast_unsigned
 --- !LinalgOpConfig
 metadata: !LinalgOpMetadata
   name: pooling_nchw_max
   cpp_class_name: PoolingNchwMaxOp
   doc: |-
     Performs max pooling.
 
     Numeric casting is performed on the input operand, promoting it to the same
structured_op: !LinalgStructuredOpConfig
   - parallel
   - parallel
   - reduction
   - reduction
   assignments:
   - !ScalarAssign
     arg: O
     value: !ScalarExpression
-      arith_fn:
+      scalar_fn:
+        kind: arith
         fn_name: max
         operands:
         - !ScalarExpression
           scalar_arg: O
         - !ScalarExpression
-          type_fn:
+          scalar_fn:
+            kind: type
+            fn_name: cast
             type_var: U
             operands:
             - !ScalarExpression
               scalar_arg: I
-            fn_name: cast
 --- !LinalgOpConfig
 metadata: !LinalgOpMetadata
   name: pooling_nhwc_min
   cpp_class_name: PoolingNhwcMinOp
   doc: |-
     Performs min pooling.
 
     Numeric casting is performed on the input operand, promoting it to the same
structured_op: !LinalgStructuredOpConfig
   - parallel
   - parallel
   - reduction
   - reduction
   assignments:
   - !ScalarAssign
     arg: O
     value: !ScalarExpression
-      arith_fn:
+      scalar_fn:
+        kind: arith
         fn_name: min
         operands:
         - !ScalarExpression
           scalar_arg: O
         - !ScalarExpression
-          type_fn:
+          scalar_fn:
+            kind: type
+            fn_name: cast
             type_var: U
             operands:
             - !ScalarExpression
               scalar_arg: I
-            fn_name: cast
 --- !LinalgOpConfig
 metadata: !LinalgOpMetadata
   name: pooling_nhwc_min_unsigned
   cpp_class_name: PoolingNhwcMinUnsignedOp
   doc: |-
     Performs unsigned min pooling.
 
     Numeric casting is performed on the input operand, promoting it to the same
structured_op: !LinalgStructuredOpConfig
   - parallel
   - parallel
   - reduction
   - reduction
   assignments:
   - !ScalarAssign
     arg: O
     value: !ScalarExpression
-      arith_fn:
+      scalar_fn:
+        kind: arith
         fn_name: min_unsigned
         operands:
         - !ScalarExpression
           scalar_arg: O
         - !ScalarExpression
-          type_fn:
+          scalar_fn:
+            kind: type
+            fn_name: cast_unsigned
             type_var: U
             operands:
             - !ScalarExpression
               scalar_arg: I
-            fn_name: cast_unsigned
 --- !LinalgOpConfig
 metadata: !LinalgOpMetadata
   name: pooling_ndhwc_sum
   cpp_class_name: PoolingNdhwcSumOp
   doc: |-
     Performs 3D sum pooling.
 
     Numeric casting is performed on the input operand, promoting it to the same
structured_op: !LinalgStructuredOpConfig
   - parallel
   - reduction
   - reduction
   - reduction
   assignments:
   - !ScalarAssign
     arg: O
     value: !ScalarExpression
-      arith_fn:
+      scalar_fn:
+        kind: arith
         fn_name: add
         operands:
         - !ScalarExpression
           scalar_arg: O
         - !ScalarExpression
-          type_fn:
+          scalar_fn:
+            kind: type
+            fn_name: cast
             type_var: U
             operands:
             - !ScalarExpression
               scalar_arg: I
-            fn_name: cast
 --- !LinalgOpConfig
 metadata: !LinalgOpMetadata
   name: pooling_ndhwc_max
   cpp_class_name: PoolingNdhwcMaxOp
   doc: |-
     Performs 3D max pooling.
 
     Numeric casting is performed on the input operand, promoting it to the same
structured_op: !LinalgStructuredOpConfig
   - parallel
   - reduction
   - reduction
   - reduction
   assignments:
   - !ScalarAssign
     arg: O
     value: !ScalarExpression
-      arith_fn:
+      scalar_fn:
+        kind: arith
         fn_name: max
         operands:
         - !ScalarExpression
           scalar_arg: O
         - !ScalarExpression
-          type_fn:
+          scalar_fn:
+            kind: type
+            fn_name: cast
             type_var: U
             operands:
             - !ScalarExpression
               scalar_arg: I
-            fn_name: cast
 --- !LinalgOpConfig
 metadata: !LinalgOpMetadata
   name: pooling_ndhwc_min
   cpp_class_name: PoolingNdhwcMinOp
   doc: |-
     Performs 3D min pooling.
 
     Numeric casting is performed on the input operand, promoting it to the same
structured_op: !LinalgStructuredOpConfig
   - parallel
   - reduction
   - reduction
   - reduction
   assignments:
   - !ScalarAssign
     arg: O
     value: !ScalarExpression
-      arith_fn:
+      scalar_fn:
+        kind: arith
         fn_name: min
         operands:
         - !ScalarExpression
           scalar_arg: O
         - !ScalarExpression
-          type_fn:
+          scalar_fn:
+            kind: type
+            fn_name: cast
             type_var: U
             operands:
             - !ScalarExpression
               scalar_arg: I
-            fn_name: cast
 --- !LinalgOpConfig
 metadata: !LinalgOpMetadata
   name: fill_tensor
   cpp_class_name: FillTensorOp
   doc: |-
     Fills the output tensor with the given value.
 
     Works for arbitrary ranked output tensors since the operation performs scalar
indexing_maps: !LinalgIndexingMapsConfig
     static_indexing_maps:
     - affine_map<() -> ()>
     - affine_map<() -> ()>
   iterator_types: []
   assignments:
   - !ScalarAssign
     arg: O
     value: !ScalarExpression
-      type_fn:
+      scalar_fn:
+        kind: type
+        fn_name: cast
         type_var: U
         operands:
         - !ScalarExpression
           scalar_arg: value
-        fn_name: cast
 --- !LinalgOpConfig
 metadata: !LinalgOpMetadata
   name: fill_rng_2d
   cpp_class_name: FillRng2DOp
   doc: |-
     Fills the output tensor with pseudo random numbers.
 
     The operation generations pseudo random numbers using a linear congruential
indexing_maps: !LinalgIndexingMapsConfig
     - affine_map<(d0, d1)[s0, s1] -> (d0, d1)>
   iterator_types:
   - parallel
   - parallel
   assignments:
   - !ScalarAssign
     arg: O
     value: !ScalarExpression
-      type_fn:
+      scalar_fn:
+        kind: type
+        fn_name: cast
         type_var: T
         operands:
         - !ScalarExpression
-          arith_fn:
+          scalar_fn:
+            kind: arith
             fn_name: add
             operands:
             - !ScalarExpression
-              arith_fn:
+              scalar_fn:
+                kind: arith
                 fn_name: mul
                 operands:
                 - !ScalarExpression
-                  arith_fn:
+                  scalar_fn:
+                    kind: arith
                     fn_name: add
                     operands:
                     - !ScalarExpression
-                      type_fn:
+                      scalar_fn:
+                        kind: type
+                        fn_name: cast
                         type_var: F64
                         operands:
                         - !ScalarExpression
                           scalar_const: '2147483647 : i64'
-                        fn_name: cast
                     - !ScalarExpression
-                      type_fn:
+                      scalar_fn:
+                        kind: type
+                        fn_name: cast
                         type_var: F64
                         operands:
                         - !ScalarExpression
-                          arith_fn:
+                          scalar_fn:
+                            kind: arith
                             fn_name: add
                             operands:
                             - !ScalarExpression
-                              arith_fn:
+                              scalar_fn:
+                                kind: arith
                                 fn_name: mul
                                 operands:
                                 - !ScalarExpression
-                                  arith_fn:
+                                  scalar_fn:
+                                    kind: arith
                                     fn_name: add
                                     operands:
                                     - !ScalarExpression
-                                      type_fn:
+                                      scalar_fn:
+                                        kind: type
+                                        fn_name: cast
                                         type_var: I32
                                         operands:
                                         - !ScalarExpression
                                           scalar_index: 1
-                                        fn_name: cast
                                     - !ScalarExpression
-                                      arith_fn:
+                                      scalar_fn:
+                                        kind: arith
                                         fn_name: add
                                         operands:
                                         - !ScalarExpression
-                                          arith_fn:
+                                          scalar_fn:
+                                            kind: arith
                                             fn_name: mul
                                             operands:
                                             - !ScalarExpression
-                                              arith_fn:
+                                              scalar_fn:
+                                                kind: arith
                                                 fn_name: add
                                                 operands:
                                                 - !ScalarExpression
-                                                  type_fn:
+                                                  scalar_fn:
+                                                    kind: type
+                                                    fn_name: cast
                                                     type_var: I32
                                                     operands:
                                                     - !ScalarExpression
                                                       scalar_index: 0
-                                                    fn_name: cast
                                                 - !ScalarExpression
                                                   scalar_arg: seed
                                             - !ScalarExpression
-                                              type_fn:
+                                              scalar_fn:
+                                                kind: type
+                                                fn_name: cast
                                                 type_var: I32
                                                 operands:
                                                 - !ScalarExpression
                                                   scalar_const: '1103515245 : i64'
-                                                fn_name: cast
                                         - !ScalarExpression
-                                          type_fn:
+                                          scalar_fn:
+                                            kind: type
+                                            fn_name: cast
                                             type_var: I32
                                             operands:
                                             - !ScalarExpression
                                               scalar_const: '12345 : i64'
-                                            fn_name: cast
                                 - !ScalarExpression
-                                  type_fn:
+                                  scalar_fn:
+                                    kind: type
+                                    fn_name: cast
                                     type_var: I32
                                     operands:
                                     - !ScalarExpression
                                       scalar_const: '1103515245 : i64'
-                                    fn_name: cast
                             - !ScalarExpression
-                              type_fn:
+                              scalar_fn:
+                                kind: type
+                                fn_name: cast
                                 type_var: I32
                                 operands:
                                 - !ScalarExpression
                                   scalar_const: '12345 : i64'
-                                fn_name: cast
-                        fn_name: cast
                 - !ScalarExpression
-                  arith_fn:
+                  scalar_fn:
+                    kind: arith
                     fn_name: mul
                     operands:
                     - !ScalarExpression
-                      arith_fn:
+                      scalar_fn:
+                        kind: arith
                         fn_name: sub
                         operands:
                         - !ScalarExpression
                           scalar_arg: max
                         - !ScalarExpression
                           scalar_arg: min
                     - !ScalarExpression
-                      type_fn:
+                      scalar_fn:
+                        kind: type
+                        fn_name: cast
                         type_var: F64
                         operands:
                         - !ScalarExpression
                           scalar_const: '2.3283063999999999E-10 : f64'
-                        fn_name: cast
             - !ScalarExpression
               scalar_arg: min
-        fn_name: cast
 --- !LinalgOpConfig
 metadata: !LinalgOpMetadata
   name: soft_plus_2d
   cpp_class_name: SoftPlus2DOp
   doc: |-
     Implements the soft plus operator.
 
     Numeric casting is performed on the input operand, promoting it to the same
indexing_maps: !LinalgIndexingMapsConfig
     - affine_map<(d0, d1)[s0, s1] -> (d0, d1)>
   iterator_types:
   - parallel
   - parallel
   assignments:
   - !ScalarAssign
     arg: O
     value: !ScalarExpression
-      arith_fn:
+      scalar_fn:
+        kind: arith
         fn_name: log
         operands:
         - !ScalarExpression
-          arith_fn:
+          scalar_fn:
+            kind: arith
             fn_name: add
             operands:
             - !ScalarExpression
-              type_fn:
+              scalar_fn:
+                kind: type
+                fn_name: cast
                 type_var: U
                 operands:
                 - !ScalarExpression
                   scalar_const: '1.000000e+00 : f64'
-                fn_name: cast
             - !ScalarExpression
-              arith_fn:
+              scalar_fn:
+                kind: arith
                 fn_name: exp
                 operands:
                 - !ScalarExpression
-                  type_fn:
+                  scalar_fn:
+                    kind: type
+                    fn_name: cast
                     type_var: U
                     operands:
                     - !ScalarExpression
                       scalar_arg: I
-                    fn_name: cast

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

class TensorUse(TensorExpression):
   def __iadd__(self, rhs: TensorExpression) -> "TensorReduceFn":
     return ReduceFnUse(ArithFn.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 TensorArithFn(TensorExpression):
-  """Application of an arithmetic function."""
+class TensorFn(TensorExpression):
+  """Application of a tensor function."""
 
-  def __init__(self, arith_fn: "ArithFnType", args: Sequence[TensorExpression]):
-    self.arith_fn = arith_fn
-    self.args = tuple(args)
-
-  def to_scalar_expression(self) -> ScalarExpression:
-    return ScalarArithFn(self.arith_fn.fn_name,
-                         *[arg.to_scalar_expression() for arg in self.args
-                          ]).expr()
-
-  def visit_tensor_exprs(self, callback: Callable[["TensorExpression"], None]):
-    super().visit_tensor_exprs(callback)
-    for arg in self.args:
-      arg.visit_tensor_exprs(callback)
-
-  def __repr__(self):
-    return f"{repr(self.arith_fn)}({', '.join(repr(a) for a in self.args)})"
-
-
-class TensorTypeFn(TensorExpression):
-  """Application of a type conversion function."""
-
-  def __init__(self, type_fn: Optional["TypeFn"],
-               operand_def: Optional["OperandDef"], type_var: TypeVar,
-               arg: TensorExpression):
-    if bool(type_fn) + bool(operand_def) != 1:
-      raise ValueError("Either 'type_fn' or 'operand_def' must be specified")
-    self.type_fn = type_fn
+  def __init__(self, kind: "FunctionKind", name: Optional[str],
+               operand_def: Optional["OperandDef"], type_var: Optional[TypeVar],
+               args: Sequence[TensorExpression]):
+    if bool(name) + bool(operand_def) != 1:
+      raise ValueError("One of 'name', 'operand_def' must be specified")
+    self.name = name
+    self.kind = kind
     self.operand_def = operand_def
     self.type_var = type_var
-    self.arg = arg
+    self.args = args
 
   def to_scalar_expression(self) -> ScalarExpression:
     if self.operand_def:
-      assert self.operand_def.name, "TypeFnAttr not attached"
-    fn_name = self.type_fn.fn_name if self.type_fn else None
+      assert self.operand_def.name, "Function attribute not attached"
     attr_name = self.operand_def.name if self.operand_def else None
-    return ScalarTypeFn(fn_name, attr_name, self.type_var,
-                        self.arg.to_scalar_expression()).expr()
+    args = [arg.to_scalar_expression() for arg in self.args]
+    return ScalarFn(self.kind, self.name, attr_name, self.type_var, args).expr()
 
   def visit_tensor_exprs(self, callback: Callable[["TensorExpression"], None]):
     super().visit_tensor_exprs(callback)
-    self.arg.visit_tensor_exprs(callback)
+    for arg in self.args:
+      arg.visit_tensor_exprs(callback)
 
   def __repr__(self):
-    return (f"{repr(self.type_fn)}[{repr(self.operand_def)}]"
-            f"({self.type_var}, {self.arg})")
+    name = self.operand_def.name if self.operand_def else self.name
+    return (f"{self.kind.name}.{name}(type_var={self.type_var}, "
+            f"args={', '.join(repr(a) for a in self.args)})")
 
 
 class TensorReduceFn(TensorExpression):
   """Application of a reduction function.
 
   This captures the lhs (initial value) separately from the rhs.
   """
 
   def __init__(self, reduce_use: "ReduceFnUse",
                args: Sequence[TensorExpression]):
     self.reduce_use = reduce_use
     self.lhs = None  # type: Optional[TensorUse]
-    self.args = tuple(args)
+    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 ScalarArithFn(self.reduce_use.arith_fn.fn_name, *full_args).expr()
+    return ScalarFn(FunctionKind.ARITH, self.reduce_use.arith_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)})"
 
def __repr__(self):
     return f"index({repr(self.dim)})"
 
 
 ###############################################################################
 # Function types and function definitions.
 ###############################################################################
 
 
+class FunctionKind(Enum):
+  ARITH = 0
+  TYPE = 1
+
+
 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) -> "TypeFnType":
-    return TensorTypeFn(self, None, type_var, arg)
+  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.
 
class ArithFnType:
 
   An arithmetic function takes one ore more tensor expressions and returns the
   function evaluation result.
   """
 
   def __init__(self, fn_name: str):
     self.fn_name = fn_name
 
-  def __call__(self, *args) -> "TensorArithFn":
-    return TensorArithFn(self, args)
+  def __call__(self, *args) -> "TensorFn":
+    return TensorFn(FunctionKind.ARITH, self.fn_name, None, None, args)
 
   def __repr__(self):
     return f"{self.fn_name}"
 
 
 class ArithFn:
   """Arithmetic function namespace.
 
class TypeFnAttrDef:
 
   def __init__(self, default: "TypeFnType"):
     if not isinstance(default, TypeFnType):
       raise ValueError(f"TypeFnAttrDef requires default of type TypeFnType "
                        f"but got {default}")
     self.operand_def = OperandDef(
         OperandKind.TYPE_FN_ATTR, default_fn=default.fn_name)
 
-  def __call__(self, type_var: TypeVar, arg: TensorExpression) -> TensorTypeFn:
-    return TensorTypeFn(None, self.operand_def, type_var, arg)
+  def __call__(self, type_var: TypeVar, arg: TensorExpression) -> TensorFn:
+    return TensorFn(FunctionKind.TYPE, None, self.operand_def, type_var, [arg])
 
 
 ###############################################################################
 # Operation definition.
 ###############################################################################
 
 
 class Comprehension:

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.arith_fn:
-      fn = self._get_function(f"_arithfn_{expr.arith_fn.fn_name}")
+    elif expr.scalar_fn and expr.scalar_fn.kind == FunctionKind.ARITH:
+      fn = self._get_function(f"_arithfn_{expr.scalar_fn.fn_name}")
       operand_values = [
-          self.expression(operand) for operand in expr.arith_fn.operands
+          self.expression(operand) for operand in expr.scalar_fn.operands
       ]
       return fn(*operand_values)
-    elif expr.type_fn:
-      fn_name = expr.type_fn.fn_name
-      if expr.type_fn.attr_name:
-        fn_name = self.type_fn_attr_mapping[expr.type_fn.attr_name]
+    elif expr.scalar_fn and expr.scalar_fn.kind == FunctionKind.TYPE:
+      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}")
-      operand = self.expression(expr.type_fn.operand)
-      return fn(expr.type_fn.type_var.name, operand)
+      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:
         output_values.append(self.yield_mapping[n])
       except KeyError:

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

 These classes are typically not user accessed and are created as a by-product
 of interpreting a comprehension DSL and model the operations to perform in the
 op body. The class hierarchy is laid out to map well to a form of YAML that
 can be easily consumed from the C++ side, not necessarily for ergonomics.
 """
 
 from typing import Optional, Sequence
 
-from .yaml_helper import *
+from .comprehension import *
 from .types import *
+from .yaml_helper import *
 
 __all__ = [
     "ScalarAssign",
-    "ScalarArithFn",
-    "ScalarTypeFn",
+    "ScalarFn",
     "ScalarArg",
     "ScalarConst",
     "ScalarIndex",
     "ScalarExpression",
 ]
 
 
-class ScalarArithFn:
-  """A type of ScalarExpression that applies an arithmetic function."""
-
-  def __init__(self, fn_name: str, *operands: "ScalarExpression"):
-    self.fn_name = fn_name
-    self.operands = operands
-
-  def expr(self) -> "ScalarExpression":
-    return ScalarExpression(arith_fn=self)
-
-  def __repr__(self):
-    return f"ScalarArithFn<{self.fn_name}>({', '.join(self.operands)})"
-
-
-class ScalarTypeFn:
-  """A type of ScalarExpression that applies a type conversion function."""
+class ScalarFn:
+  """A type of ScalarExpression that applies a function."""
 
-  def __init__(self, fn_name: Optional[str], attr_name: Optional[str],
-               type_var: TypeVar, operand: "ScalarExpression"):
+  def __init__(self, kind: "FunctionKind", fn_name: Optional[str],
+               attr_name: Optional[str], type_var: Optional["TypeVar"],
+               operands: Sequence["ScalarExpression"]):
+    if bool(fn_name) + bool(attr_name) != 1:
+      raise ValueError("One of 'fn_name', 'attr_name' must be specified")
+    self.kind = kind
     self.fn_name = fn_name
     self.attr_name = attr_name
     self.type_var = type_var
-    self.operand = operand
+    self.operands = operands
 
   def expr(self) -> "ScalarExpression":
-    return ScalarExpression(type_fn=self)
+    return ScalarExpression(scalar_fn=self)
 
   def __repr__(self):
-    return (f"ScalarTypeFn<{self.fn_name}[{self.attr_name}]>"
-            f"({self.type_var}, {self.operand})")
+    name = self.fn_name if self.fn_name else self.attr_name
+    return (f"ScalarFn<{self.kind.name}.{name}>(type_var={self.type_var}, "
+            f"operands=[{', '.join(self.operands)}])")
 
 
 class ScalarArg:
   """A type of ScalarExpression that references a named argument."""
 
   def __init__(self, arg: str):
     self.arg = arg
 
class ScalarIndex:
   def __repr__(self):
     return f"(ScalarIndex({self.dim})"
 
 
 class ScalarExpression(YAMLObject):
   """An expression on scalar values.
 
   Can be one of:
-    - ScalarArithFn
-    - ScalarTypeFn
+    - ScalarFn
     - ScalarArg
     - ScalarConst
     - ScalarIndex
-    - ScalarSymbolicCast
   """
   yaml_tag = "!ScalarExpression"
 
   def __init__(self,
-               arith_fn: Optional[ScalarArithFn] = None,
-               type_fn: Optional[ScalarTypeFn] = None,
+               scalar_fn: Optional[ScalarFn] = None,
                scalar_arg: Optional[ScalarArg] = None,
                scalar_const: Optional[ScalarConst] = None,
                scalar_index: Optional[ScalarIndex] = None):
-    if (bool(arith_fn) + bool(type_fn) + bool(scalar_arg) + bool(scalar_const) +
+    if (bool(scalar_fn) + bool(scalar_arg) + bool(scalar_const) +
         bool(scalar_index)) != 1:
-      raise ValueError("One of 'arith_fn', 'type_fn', 'scalar_arg', "
-                       "'scalar_const', 'scalar_index', must be specified")
-    self.arith_fn = arith_fn
-    self.type_fn = type_fn
+      raise ValueError("One of 'scalar_fn', 'scalar_arg', 'scalar_const', or "
+                       "'scalar_index' must be specified")
+    self.scalar_fn = scalar_fn
     self.scalar_arg = scalar_arg
     self.scalar_const = scalar_const
     self.scalar_index = scalar_index
 
   def to_yaml_custom_dict(self):
-    if self.arith_fn:
-      return dict(
-          arith_fn=dict(
-              fn_name=self.arith_fn.fn_name,
-              operands=list(self.arith_fn.operands),
-          ))
-    if self.type_fn:
-      # Note that even though operands must be arity 1, we write it the
-      # same way as for apply because it allows handling code to be more
-      # generic vs having a special form.
-      type_fn_dict = dict(
-          type_var=self.type_fn.type_var.name,
-          operands=[self.type_fn.operand],
-      )
-      if self.type_fn.fn_name:
-        type_fn_dict["fn_name"] = self.type_fn.fn_name
-      if self.type_fn.attr_name:
-        type_fn_dict["attr_name"] = self.type_fn.attr_name
-      return dict(type_fn=type_fn_dict)
+    if self.scalar_fn:
+      scalar_fn_dict = dict(kind=self.scalar_fn.kind.name.lower())
+      if self.scalar_fn.fn_name:
+        scalar_fn_dict["fn_name"] = self.scalar_fn.fn_name
+      if self.scalar_fn.attr_name:
+        scalar_fn_dict["attr_name"] = self.scalar_fn.attr_name
+      if self.scalar_fn.type_var:
+        scalar_fn_dict["type_var"] = self.scalar_fn.type_var.name
+      scalar_fn_dict["operands"] = list(self.scalar_fn.operands)
+      return dict(scalar_fn=scalar_fn_dict)
     elif self.scalar_arg:
       return dict(scalar_arg=self.scalar_arg.arg)
     elif self.scalar_const:
       return dict(scalar_const=self.scalar_const.value)
     elif self.scalar_index:
       return dict(scalar_index=self.scalar_index.dim)
     else:
       raise ValueError(f"Unexpected ScalarExpression type: {self}")

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

indexing_maps: !LinalgIndexingMapsConfig
     - affine_map<(d0, d1)[s0, s1] -> (d0, d1)>
   iterator_types:
   - parallel
   - parallel
   assignments:
   - !ScalarAssign
     arg: O
     value: !ScalarExpression
-      arith_fn:
+      scalar_fn:
+        kind: arith
         fn_name: add
         operands:
         - !ScalarExpression
-          type_fn:
+          scalar_fn:
+            kind: type
+            attr_name: cast
             type_var: T
             operands:
             - !ScalarExpression
               scalar_const: '42 : i64'
-            attr_name: cast
         - !ScalarExpression
-          type_fn:
+          scalar_fn:
+            kind: type
+            attr_name: cast
             type_var: T
             operands:
             - !ScalarExpression
               scalar_index: 1
-            attr_name: cast
 
 # ODS-LABEL:  def Test1Op : LinalgStructuredBase_Op<"test1"
 
 #       ODS:  let summary = [{ Title. }];
 #  ODS-NEXT:  let description = [{
 #  ODS-NEXT:    Detailed description.
 #  ODS-NEXT:  }];
 
indexing_maps: !LinalgIndexingMapsConfig
     static_indexing_maps:
     - affine_map<() -> ()>
     - affine_map<() -> ()>
   iterator_types: []
   assignments:
   - !ScalarAssign
     arg: O
     value: !ScalarExpression
-      type_fn:
+      scalar_fn:
+        kind: type
         fn_name: cast
         type_var: U
         operands:
         - !ScalarExpression
           scalar_arg: value
 
 #       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(

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

 # RUN: %PYTHON -m mlir.dialects.linalg.opdsl.dump_oplib --file %s | FileCheck %s
 
 from mlir.dialects.linalg.opdsl.lang import *
 
 
 # CHECK: ---
 # CHECK-LABEL: matmul
 # CHECK: assignments:
 # CHECK:  -
 # CHECK:    arg: C
 # CHECK:    value:
-# CHECK:      arith_fn:
+# CHECK:      scalar_fn:
 # CHECK:        fn_name: add
 # CHECK:        operands:
-# CHECK:          arith_fn:
+# CHECK:          scalar_fn:
 # CHECK:            fn_name: mul
 # CHECK:            operands:
-# CHECK:              type_fn:
+# CHECK:              scalar_fn:
+# CHECK:                kind: type
+# CHECK:                attr_name: cast
 # CHECK:                type_var: U
 # CHECK:                operands:
 # CHECK:                  scalar_arg: A
+# CHECK:              scalar_fn:
+# CHECK:                kind: type
 # CHECK:                attr_name: cast
-# CHECK:              type_fn:
 # CHECK:                type_var: U
 # CHECK:                operands:
 # CHECK:                  scalar_arg: B
-# CHECK:                attr_name: cast
 @linalg_structured_op
 def matmul(
     A=TensorDef(T, S.M, S.K),
     B=TensorDef(T, S.K, S.N),
     C=TensorDef(U, S.M, S.N, output=True),
     cast=TypeFnAttrDef(default=TypeFn.cast)):
   C[D.m, D.n] += cast(U, A[D.m, D.k]) * cast(U, B[D.k, D.n])
 
 
 # CHECK: ---
 # CHECK-LABEL: constants
 # CHECK: assignments:
 # CHECK:  -
 # CHECK:    arg: O
-# CHECK:      arith_fn:
+# CHECK:      scalar_fn:
+# CHECK:        kind: arith
 # CHECK:        fn_name: sub
 # CHECK:        operands:
-# CHECK:          arith_fn:
+# CHECK:          scalar_fn:
+# CHECK:            kind: arith
 # CHECK:            fn_name: add
 # CHECK:            operands:
-# CHECK:              type_fn:
+# CHECK:              scalar_fn:
+# CHECK:                kind: type
 # CHECK:                type_var: T
 # CHECK:                operands:
 # CHECK:                  scalar_const: '3.1415926535897931 : f64'
-# CHECK:              type_fn:
+# CHECK:              scalar_fn:
+# CHECK:                kind: type
+# CHECK:                fn_name: cast
 # CHECK:                type_var: T
 # CHECK:                operands:
 # CHECK:                  scalar_const: '42 : i64'
-# CHECK:          type_fn:
+# 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
 
 
 # CHECK: ---
 # CHECK-LABEL: indices
 # CHECK: assignments:
 # CHECK:  -
 # CHECK:    arg: O
-# CHECK:      arith_fn:
+# CHECK:      scalar_fn:
+# CHECK:        kind: arith
 # 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/tools/mlir-linalg-ods-gen/mlir-linalg-ods-yaml-gen.cpp

 };
 
 struct LinalgIndexingMapsConfig {
   Optional<SmallVector<SerializedAffineMap>> staticIndexingMaps;
 };
 
 struct ScalarExpression;
 
-struct ScalarArithFn {
-  std::string fnName;
-  // NOTE: Must be pure heap allocated container (not SmallVector)
-  // due to recursive data type.
-  std::vector<ScalarExpression> operands;
-};
+enum class ScalarFnKind { Arith, Type };
 
-struct ScalarTypeFn {
-  std::string typeVar;
+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.
   std::vector<ScalarExpression> operands;
-  Optional<std::string> fnName;
-  Optional<std::string> attrName;
 };
 
 struct ScalarExpression {
   Optional<std::string> arg;
   Optional<std::string> constant;
   Optional<int64_t> index;
-  Optional<ScalarArithFn> arithFn;
-  Optional<ScalarTypeFn> typeFn;
+  Optional<ScalarFn> scalarFn;
 };
 
 struct ScalarAssign {
   std::string arg;
   ScalarExpression value;
 };
 
 struct LinalgStructuredOpConfig {
static void mapping(IO &io, ScalarAssign &info) {
     io.mapRequired("value", info.value);
   }
 };
 
 /// A scalar expression (RHS of an assignment). Must be one of:
 ///   - `scalar_arg`: An operation argument.
 ///   - `scalar_const`: A constant definition.
 ///   - `scalar_index`: An iteration index.
-///   - `arith_fn`: A named arithmetic function (see `ScalarArithFn`).
-///   - `type_fn`: A named type conversion function (see `ScalarTypeFn`).
+///   - `scalar_fn`: A named function (see `ScalarFn`).
 template <>
 struct MappingTraits<ScalarExpression> {
   static void mapping(IO &io, ScalarExpression &info) {
     io.mapOptional("scalar_arg", info.arg);
     io.mapOptional("scalar_const", info.constant);
     io.mapOptional("scalar_index", info.index);
-    io.mapOptional("arith_fn", info.arithFn);
-    io.mapOptional("type_fn", info.typeFn);
+    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, "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)`
 ///   - `mul(lhs, rhs)`
 template <>
-struct MappingTraits<ScalarArithFn> {
-  static void mapping(IO &io, ScalarArithFn &info) {
-    io.mapRequired("fn_name", info.fnName);
-    io.mapRequired("operands", info.operands);
-  }
-};
-
-template <>
-struct MappingTraits<ScalarTypeFn> {
-  static void mapping(IO &io, ScalarTypeFn &info) {
-    io.mapRequired("type_var", info.typeVar);
-    io.mapRequired("operands", info.operands);
+struct MappingTraits<ScalarFn> {
+  static void mapping(IO &io, ScalarFn &info) {
+    io.mapRequired("kind", info.kind);
     io.mapOptional("fn_name", info.fnName);
     io.mapOptional("attr_name", info.attrName);
+    io.mapOptional("type_var", info.typeVar);
+    io.mapRequired("operands", info.operands);
   }
 };
 
 /// Helper mapping which accesses an AffineMapAttr as a serialized string of
 /// the same.
 template <>
 struct ScalarTraits<SerializedAffineMap> {
   static void output(const SerializedAffineMap &value, void *rawYamlContext,
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.arithFn) {
+        if (expression.scalarFn &&
+            expression.scalarFn->kind == ScalarFnKind::Arith) {
           // Apply function.
           // Recursively generate operands.
           SmallVector<std::string> operandCppValues;
-          for (ScalarExpression &operand : expression.arithFn->operands) {
+          for (ScalarExpression &operand : expression.scalarFn->operands) {
             auto operandCppValue = generateExpression(operand);
             if (!operandCppValue)
               return None;
             operandCppValues.push_back(*operandCppValue);
           }
           std::string cppIdent = llvm::formatv("value{0}", ++localCounter);
           stmts.push_back(
               llvm::formatv("Value {0} = helper.arithfn__{1}({2});", cppIdent,
-                            expression.arithFn->fnName,
+                            expression.scalarFn->fnName,
                             interleaveToString(operandCppValues, ", ")));
           return cppIdent;
         }
-        if (expression.typeFn) {
+        if (expression.scalarFn &&
+            expression.scalarFn->kind == ScalarFnKind::Type) {
           // Symbolic cast.
           // Operands must be arity 1.
-          if (expression.typeFn->operands.size() != 1) {
+          if (expression.scalarFn->operands.size() != 1) {
             emitError(genContext.getLoc())
                 << "type conversion operand arity must be 1";
             return None;
           }
           Optional<std::string> operandCppValue =
-              generateExpression(expression.typeFn->operands[0]);
+              generateExpression(expression.scalarFn->operands[0]);
           if (!operandCppValue)
             return None;
 
+          assert(expression.scalarFn->typeVar.hasValue());
           Optional<std::string> typeCppValue =
-              findTypeValue(expression.typeFn->typeVar, args);
+              findTypeValue(expression.scalarFn->typeVar.getValue(), args);
           if (!typeCppValue) {
             emitError(genContext.getLoc())
-                << "type variable " << expression.typeFn->typeVar
+                << "type variable " << expression.scalarFn->typeVar.getValue()
                 << ", used in a type conversion, must map to a predefined or "
                 << "an argument type but it does not";
             return None;
           }
 
           // Use the function name or the attribute to build the type function.
           std::string typeFunc = llvm::formatv(
-              "TypeFn::{0}", expression.typeFn->fnName.getValueOr(""));
-          if (expression.typeFn->attrName) {
+              "TypeFn::{0}", expression.scalarFn->fnName.getValueOr(""));
+          if (expression.scalarFn->attrName) {
             if (llvm::none_of(args, [&](LinalgOperandDef &arg) {
                   return arg.kind == LinalgOperandDefKind::TypeFnAttr &&
-                         arg.name == expression.typeFn->attrName.getValue();
+                         arg.name == expression.scalarFn->attrName.getValue();
                 })) {
               emitError(genContext.getLoc())
                   << "missing type function attribute "
-                  << expression.typeFn->attrName.getValue();
+                  << expression.scalarFn->attrName.getValue();
             }
-            typeFunc = llvm::formatv("{0}Val", *expression.typeFn->attrName);
+            typeFunc = llvm::formatv("{0}Val", *expression.scalarFn->attrName);
           }
           std::string cppIdent = llvm::formatv("value{0}", ++localCounter);
           stmts.push_back(llvm::formatv(
               "Value {0} = helper.buildTypeFn({1}, {2}, {3});", cppIdent,
               typeFunc, typeCppValue.getValue(), *operandCppValue));
           return cppIdent;
         }
         emitError(genContext.getLoc()) << "unknown ScalarExpression type";