[mlir][OpDSL] Add type function attributes.

[gysit]

Summary

Previously, OpDSL operation used hardcoded type conversion operations (cast or cast_unsigned). Supporting signed and unsigned casts thus meant implementing two different operations. Type function attributes allow us to define a single operation that has a cast type function attribute which at operation instantiation time may be set to cast or cast_unsigned. We may for example, defina a matmul operation with a cast argument:

@linalg_structured_op
def matmul(A=TensorDef(T1, S.M, S.K), B=TensorDef(T2, 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])

When instantiating the operation the attribute may be set to the desired cast function:

linalg.matmul(lhs, rhs, outs=[out], cast=TypeFn.cast_unsigned)

The revsion introduces a enum in the Linalg dialect that maps one-by-one to the type functions defined by OpDSL.

[gysit]

Improve tests.

[gysit]

mlir/include/mlir/Dialect/Linalg/IR/LinalgStructuredOps.td

40

This is needed since the operation has a cast member function to access the cast attribute.

[gysit]

Update OperandKind enum to match coding conventions.

[gysit]

Rebase and cleaner test.

[gysit]

Rebase.

[aartbik]

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

4 ↗

(On Diff #409861)

why is this line removed?

529 ↗

(On Diff #409861)

this, too, seems unrelated to the revision at large?

[gysit]

Attempt to fix autoformatting

[gysit]

At least for now the auto formatting noise seems gone...

[gysit]

Use custom enum format.

[aartbik]

LG for the pytaco part

[stellaraccident]

Aart, would you mind doing a review of the entire patch? (I've got a lot of latency this week and would also like to expand the reviewer pool here -- and you are an active user of this)?

[gysit]

Fix enum mapping and add addtional check to avoid attribute name clashes.

[aartbik]

did a bit more thorough pass over the other parts too

mlir/include/mlir/Dialect/Linalg/IR/LinalgBase.td

63

Although I assume you want the common case to be shorter, there is something to say to use cast_signed and cast_unsigned to make the distinction explicit

65

would a cast_nop make sense
(i.e. use the OpsDSL for cases where there is *no* cast in the instantation, even though pattern defines one)

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

158

Thanks for moving that here. Having this at the wrong place can be very confusing
(even though it typically only defines the destructor)

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

174

attached is a bit strange term here, how about specified, like L164

555

typo

670

clashes -> conflicts

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

129

typo: attribute

mlir/test/Dialect/Linalg/generalize-named-polymorphic-ops.mlir

43

Verifies that ....

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

185

_vals reads weird now that we have _fn, can we rename?

[gysit]

mlir/include/mlir/Dialect/Linalg/IR/LinalgBase.td

63

I would suggest to do this in a separate revision together with max -> max_signed. I will change this in OpDSL as well and many places will be affected. I think OpDSL historically was meant to be signed by default and we added unsigned later. I think now that we can use attributes in most places going for the longer names in the default case should be fine.

65

I an a bit unsure about this one. Both cast and cast_unsigned are nops if the types match. A cast_nop/no_cast would result in a verification failure after op construction and we may emit a warning. It may not be easy to understand though since the region of named operations is hidden. I am thus a bit unsure if this would be a good addition.

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

174

Yes, I switched to "not registered with an op" since this used in most other places. The check in fact makes sure a specific operand has been registered with an operation. Otherwise it has no name yet.

[gysit]

Address comments.

[aartbik]

mlir/include/mlir/Dialect/Linalg/IR/LinalgBase.td

65

Fair enough. I was just thinking out loud here.

[gysit]

mlir/include/mlir/Dialect/Linalg/IR/LinalgBase.td

65

Thanks for the review and especially for thinking out loud! A cast_nop is IMO a good idea! It is just not so easy to implement in the current framework.

mlir/include/mlir/Dialect/Linalg/IR/CMakeLists.txt

 add_custom_target(LinalgOdsGen
   DEPENDS
   MLIRLinalgNamedStructuredOpsYamlIncGen
 )
 add_dependencies(mlir-headers LinalgOdsGen)
 
 add_mlir_dialect(LinalgOps linalg)
 
+set(LLVM_TARGET_DEFINITIONS LinalgOps.td)
+mlir_tablegen(LinalgOpsEnums.h.inc -gen-enum-decls)
+mlir_tablegen(LinalgOpsEnums.cpp.inc -gen-enum-defs)
+add_public_tablegen_target(MLIRLinalgOpsEnumsIncGen)
+add_dependencies(mlir-headers MLIRLinalgOpsEnumsIncGen)
+
+set(LLVM_TARGET_DEFINITIONS LinalgOps.td)
+mlir_tablegen(LinalgOpsAttrDefs.h.inc -gen-attrdef-decls)
+mlir_tablegen(LinalgOpsAttrDefs.cpp.inc -gen-attrdef-defs)
+add_public_tablegen_target(MLIRLinalgOpsAttributesIncGen)
+add_dependencies(mlir-headers MLIRLinalgOpsAttributesIncGen)
+
 add_mlir_doc(LinalgDoc LinalgOps Dialects/ -gen-op-doc)
 add_dependencies(LinalgOpsDocGen LinalgOdsGen)
 
 set(LLVM_TARGET_DEFINITIONS LinalgStructuredOps.td)
 mlir_tablegen(LinalgStructuredOps.h.inc -gen-op-decls)
 mlir_tablegen(LinalgStructuredOps.cpp.inc -gen-op-defs)
 add_public_tablegen_target(MLIRLinalgStructuredOpsIncGen)
 add_dependencies(MLIRLinalgStructuredOpsIncGen LinalgOdsGen)
 add_dependencies(mlir-headers MLIRLinalgStructuredOpsIncGen)
 
 set(LLVM_TARGET_DEFINITIONS LinalgInterfaces.td)
 mlir_tablegen(LinalgInterfaces.h.inc -gen-op-interface-decls)
 mlir_tablegen(LinalgInterfaces.cpp.inc -gen-op-interface-defs)
 add_public_tablegen_target(MLIRLinalgInterfacesIncGen)
 add_dependencies(mlir-headers MLIRLinalgInterfacesIncGen)

mlir/include/mlir/Dialect/Linalg/IR/Linalg.h

 
 //===----------------------------------------------------------------------===//
 // Linalg Dialect
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Dialect/Linalg/IR/LinalgOpsDialect.h.inc"
 
 //===----------------------------------------------------------------------===//
+// Linalg Enums
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Dialect/Linalg/IR/LinalgOpsEnums.h.inc"
+
+//===----------------------------------------------------------------------===//
+// Linalg Attributes
+//===----------------------------------------------------------------------===//
+
+#define GET_ATTRDEF_CLASSES
+#include "mlir/Dialect/Linalg/IR/LinalgOpsAttrDefs.h.inc"
+
+//===----------------------------------------------------------------------===//
 // Linalg Interfaces
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Dialect/Linalg/IR/LinalgInterfaces.h"
 
 //===----------------------------------------------------------------------===//
 // Linalg Dialect Operations
 //===----------------------------------------------------------------------===//
 
 #define GET_OP_CLASSES
 #include "mlir/Dialect/Linalg/IR/LinalgOps.h.inc"
 
 #define GET_OP_CLASSES
 #include "mlir/Dialect/Linalg/IR/LinalgStructuredOps.h.inc"
 
 #endif // MLIR_DIALECT_LINALG_IR_LINALG_H

mlir/include/mlir/Dialect/Linalg/IR/LinalgBase.td

 //===- LinalgBase.td - Linalg dialect base support ---------*- tablegen -*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
 // This is the definition file for base linear algebra support.
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef LINALG_BASE
 #define LINALG_BASE
 
+include "mlir/IR/EnumAttr.td"
 include "mlir/IR/OpBase.td"
 
 def Linalg_Dialect : Dialect {
   let name = "linalg";
   let description = [{
     The `linalg` dialect groups together a set of types, operations and
     transformations that are useful to implement a structured abstraction on
     buffers and tensors. These abstractions are useful for transformations and
let extraClassDeclaration = [{
     RegionBuilderFunType getRegionBuilder(StringRef name) {
       return namedStructuredOpRegionBuilders.lookup(name);
     }
     private:
       llvm::StringMap<RegionBuilderFunType> namedStructuredOpRegionBuilders;
   }];
 }
 
+// Define a TypeFn enum matching the OpDSL TypeFn class.
+def TypeFn : I32EnumAttr<"TypeFn", "", [
+  I32EnumAttrCase<"cast", 0>,
+  I32EnumAttrCase<"cast_unsigned", 1>
+]> {
+  let genSpecializedAttr = 0;
+  let cppNamespace = "::mlir::linalg";
+}
+
+def TypeFnAttr : EnumAttr<Linalg_Dialect, TypeFn, "type_fn">;
+
 #endif // LINALG_BASE

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

- !LinalgOperandDefConfig
     usage: Input
     type_var: T2
     shape_map: affine_map<()[s0, s1, s2] -> (s1, s2)>
   - !LinalgOperandDefConfig
     name: C
     usage: Output
     type_var: U
     shape_map: affine_map<()[s0, s1, s2] -> (s0, s2)>
+  - !LinalgOperandDefConfig
+    name: cast
+    usage: TypeFnAttr
+    default_fn: cast
   indexing_maps: !LinalgIndexingMapsConfig
     static_indexing_maps:
     - affine_map<(d0, d1, d2)[s0, s1, s2] -> (d0, d2)>
     - affine_map<(d0, d1, d2)[s0, s1, s2] -> (d2, d1)>
     - affine_map<(d0, d1, d2)[s0, s1, s2] -> (d0, d1)>
   iterator_types:
   - parallel
   - parallel
   - reduction
   assignments:
   - !ScalarAssign
     arg: C
     value: !ScalarExpression
       arith_fn:
         fn_name: add
         operands:
         - !ScalarExpression
           scalar_arg: C
         - !ScalarExpression
           arith_fn:
             fn_name: mul
             operands:
             - !ScalarExpression
               type_fn:
-                fn_name: cast
                 type_var: U
                 operands:
                 - !ScalarExpression
                   scalar_arg: A
+                attr_name: cast
             - !ScalarExpression
               type_fn:
-                fn_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
value: !ScalarExpression
         - !ScalarExpression
           scalar_arg: C
         - !ScalarExpression
           arith_fn:
             fn_name: mul
             operands:
             - !ScalarExpression
               type_fn:
-                fn_name: cast_unsigned
                 type_var: U
                 operands:
                 - !ScalarExpression
                   scalar_arg: A
+                fn_name: cast_unsigned
             - !ScalarExpression
               type_fn:
-                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
value: !ScalarExpression
             fn_name: mul
             operands:
             - !ScalarExpression
               arith_fn:
                 fn_name: sub
                 operands:
                 - !ScalarExpression
                   type_fn:
-                    fn_name: cast
                     type_var: U
                     operands:
                     - !ScalarExpression
                       scalar_arg: A
+                    fn_name: cast
                 - !ScalarExpression
                   type_fn:
-                    fn_name: cast
                     type_var: U
                     operands:
                     - !ScalarExpression
                       scalar_arg: AZp
+                    fn_name: cast
             - !ScalarExpression
               arith_fn:
                 fn_name: sub
                 operands:
                 - !ScalarExpression
                   type_fn:
-                    fn_name: cast
                     type_var: U
                     operands:
                     - !ScalarExpression
                       scalar_arg: B
+                    fn_name: cast
                 - !ScalarExpression
                   type_fn:
-                    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:
value: !ScalarExpression
         - !ScalarExpression
           scalar_arg: accum
         - !ScalarExpression
           arith_fn:
             fn_name: mul
             operands:
             - !ScalarExpression
               type_fn:
-                fn_name: cast
                 type_var: AccumType
                 operands:
                 - !ScalarExpression
                   scalar_arg: lhs
+                fn_name: cast
             - !ScalarExpression
               type_fn:
-                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
value: !ScalarExpression
         - !ScalarExpression
           scalar_arg: C
         - !ScalarExpression
           arith_fn:
             fn_name: mul
             operands:
             - !ScalarExpression
               type_fn:
-                fn_name: cast
                 type_var: U
                 operands:
                 - !ScalarExpression
                   scalar_arg: A
+                fn_name: cast
             - !ScalarExpression
               type_fn:
-                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
value: !ScalarExpression
             fn_name: mul
             operands:
             - !ScalarExpression
               arith_fn:
                 fn_name: sub
                 operands:
                 - !ScalarExpression
                   type_fn:
-                    fn_name: cast
                     type_var: U
                     operands:
                     - !ScalarExpression
                       scalar_arg: A
+                    fn_name: cast
                 - !ScalarExpression
                   type_fn:
-                    fn_name: cast
                     type_var: U
                     operands:
                     - !ScalarExpression
                       scalar_arg: AZp
+                    fn_name: cast
             - !ScalarExpression
               arith_fn:
                 fn_name: sub
                 operands:
                 - !ScalarExpression
                   type_fn:
-                    fn_name: cast
                     type_var: U
                     operands:
                     - !ScalarExpression
                       scalar_arg: B
+                    fn_name: cast
                 - !ScalarExpression
                   type_fn:
-                    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
value: !ScalarExpression
         - !ScalarExpression
           scalar_arg: x
         - !ScalarExpression
           arith_fn:
             fn_name: mul
             operands:
             - !ScalarExpression
               type_fn:
-                fn_name: cast
                 type_var: U
                 operands:
                 - !ScalarExpression
                   scalar_arg: A
+                fn_name: cast
             - !ScalarExpression
               type_fn:
-                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
value: !ScalarExpression
         - !ScalarExpression
           scalar_arg: x
         - !ScalarExpression
           arith_fn:
             fn_name: mul
             operands:
             - !ScalarExpression
               type_fn:
-                fn_name: cast
                 type_var: U
                 operands:
                 - !ScalarExpression
                   scalar_arg: y
+                fn_name: cast
             - !ScalarExpression
               type_fn:
-                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
value: !ScalarExpression
         - !ScalarExpression
           scalar_arg: C
         - !ScalarExpression
           arith_fn:
             fn_name: mul
             operands:
             - !ScalarExpression
               type_fn:
-                fn_name: cast
                 type_var: U
                 operands:
                 - !ScalarExpression
                   scalar_arg: A
+                fn_name: cast
             - !ScalarExpression
               type_fn:
-                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
value: !ScalarExpression
         - !ScalarExpression
           scalar_arg: C
         - !ScalarExpression
           arith_fn:
             fn_name: mul
             operands:
             - !ScalarExpression
               type_fn:
-                fn_name: cast
                 type_var: U
                 operands:
                 - !ScalarExpression
                   scalar_arg: A
+                fn_name: cast
             - !ScalarExpression
               type_fn:
-                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
value: !ScalarExpression
         - !ScalarExpression
           scalar_arg: O
         - !ScalarExpression
           arith_fn:
             fn_name: mul
             operands:
             - !ScalarExpression
               type_fn:
-                fn_name: cast
                 type_var: U
                 operands:
                 - !ScalarExpression
                   scalar_arg: I
+                fn_name: cast
             - !ScalarExpression
               type_fn:
-                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
value: !ScalarExpression
         - !ScalarExpression
           scalar_arg: O
         - !ScalarExpression
           arith_fn:
             fn_name: mul
             operands:
             - !ScalarExpression
               type_fn:
-                fn_name: cast
                 type_var: U
                 operands:
                 - !ScalarExpression
                   scalar_arg: I
+                fn_name: cast
             - !ScalarExpression
               type_fn:
-                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
value: !ScalarExpression
         - !ScalarExpression
           scalar_arg: O
         - !ScalarExpression
           arith_fn:
             fn_name: mul
             operands:
             - !ScalarExpression
               type_fn:
-                fn_name: cast
                 type_var: U
                 operands:
                 - !ScalarExpression
                   scalar_arg: I
+                fn_name: cast
             - !ScalarExpression
               type_fn:
-                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
value: !ScalarExpression
         - !ScalarExpression
           scalar_arg: O
         - !ScalarExpression
           arith_fn:
             fn_name: mul
             operands:
             - !ScalarExpression
               type_fn:
-                fn_name: cast
                 type_var: U
                 operands:
                 - !ScalarExpression
                   scalar_arg: I
+                fn_name: cast
             - !ScalarExpression
               type_fn:
-                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:
value: !ScalarExpression
         - !ScalarExpression
           scalar_arg: O
         - !ScalarExpression
           arith_fn:
             fn_name: mul
             operands:
             - !ScalarExpression
               type_fn:
-                fn_name: cast
                 type_var: U
                 operands:
                 - !ScalarExpression
                   scalar_arg: I
+                fn_name: cast
             - !ScalarExpression
               type_fn:
-                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:
value: !ScalarExpression
             fn_name: mul
             operands:
             - !ScalarExpression
               arith_fn:
                 fn_name: sub
                 operands:
                 - !ScalarExpression
                   type_fn:
-                    fn_name: cast
                     type_var: U
                     operands:
                     - !ScalarExpression
                       scalar_arg: I
+                    fn_name: cast
                 - !ScalarExpression
                   type_fn:
-                    fn_name: cast
                     type_var: U
                     operands:
                     - !ScalarExpression
                       scalar_arg: IZp
+                    fn_name: cast
             - !ScalarExpression
               arith_fn:
                 fn_name: sub
                 operands:
                 - !ScalarExpression
                   type_fn:
-                    fn_name: cast
                     type_var: U
                     operands:
                     - !ScalarExpression
                       scalar_arg: K
+                    fn_name: cast
                 - !ScalarExpression
                   type_fn:
-                    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:
value: !ScalarExpression
         - !ScalarExpression
           scalar_arg: O
         - !ScalarExpression
           arith_fn:
             fn_name: mul
             operands:
             - !ScalarExpression
               type_fn:
-                fn_name: cast
                 type_var: U
                 operands:
                 - !ScalarExpression
                   scalar_arg: I
+                fn_name: cast
             - !ScalarExpression
               type_fn:
-                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
value: !ScalarExpression
         - !ScalarExpression
           scalar_arg: O
         - !ScalarExpression
           arith_fn:
             fn_name: mul
             operands:
             - !ScalarExpression
               type_fn:
-                fn_name: cast
                 type_var: U
                 operands:
                 - !ScalarExpression
                   scalar_arg: I
+                fn_name: cast
             - !ScalarExpression
               type_fn:
-                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
value: !ScalarExpression
         - !ScalarExpression
           scalar_arg: O
         - !ScalarExpression
           arith_fn:
             fn_name: mul
             operands:
             - !ScalarExpression
               type_fn:
-                fn_name: cast
                 type_var: U
                 operands:
                 - !ScalarExpression
                   scalar_arg: I
+                fn_name: cast
             - !ScalarExpression
               type_fn:
-                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
value: !ScalarExpression
         - !ScalarExpression
           scalar_arg: O
         - !ScalarExpression
           arith_fn:
             fn_name: mul
             operands:
             - !ScalarExpression
               type_fn:
-                fn_name: cast
                 type_var: U
                 operands:
                 - !ScalarExpression
                   scalar_arg: I
+                fn_name: cast
             - !ScalarExpression
               type_fn:
-                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
value: !ScalarExpression
             fn_name: mul
             operands:
             - !ScalarExpression
               arith_fn:
                 fn_name: sub
                 operands:
                 - !ScalarExpression
                   type_fn:
-                    fn_name: cast
                     type_var: U
                     operands:
                     - !ScalarExpression
                       scalar_arg: I
+                    fn_name: cast
                 - !ScalarExpression
                   type_fn:
-                    fn_name: cast
                     type_var: U
                     operands:
                     - !ScalarExpression
                       scalar_arg: IZp
+                    fn_name: cast
             - !ScalarExpression
               arith_fn:
                 fn_name: sub
                 operands:
                 - !ScalarExpression
                   type_fn:
-                    fn_name: cast
                     type_var: U
                     operands:
                     - !ScalarExpression
                       scalar_arg: K
+                    fn_name: cast
                 - !ScalarExpression
                   type_fn:
-                    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
value: !ScalarExpression
         - !ScalarExpression
           scalar_arg: O
         - !ScalarExpression
           arith_fn:
             fn_name: mul
             operands:
             - !ScalarExpression
               type_fn:
-                fn_name: cast
                 type_var: U
                 operands:
                 - !ScalarExpression
                   scalar_arg: I
+                fn_name: cast
             - !ScalarExpression
               type_fn:
-                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
value: !ScalarExpression
             fn_name: mul
             operands:
             - !ScalarExpression
               arith_fn:
                 fn_name: sub
                 operands:
                 - !ScalarExpression
                   type_fn:
-                    fn_name: cast
                     type_var: U
                     operands:
                     - !ScalarExpression
                       scalar_arg: I
+                    fn_name: cast
                 - !ScalarExpression
                   type_fn:
-                    fn_name: cast
                     type_var: U
                     operands:
                     - !ScalarExpression
                       scalar_arg: IZp
+                    fn_name: cast
             - !ScalarExpression
               arith_fn:
                 fn_name: sub
                 operands:
                 - !ScalarExpression
                   type_fn:
-                    fn_name: cast
                     type_var: U
                     operands:
                     - !ScalarExpression
                       scalar_arg: K
+                    fn_name: cast
                 - !ScalarExpression
                   type_fn:
-                    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
- !ScalarAssign
     value: !ScalarExpression
       arith_fn:
         fn_name: add
         operands:
         - !ScalarExpression
           scalar_arg: O
         - !ScalarExpression
           type_fn:
-            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
- !ScalarAssign
     value: !ScalarExpression
       arith_fn:
         fn_name: max
         operands:
         - !ScalarExpression
           scalar_arg: O
         - !ScalarExpression
           type_fn:
-            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
- !ScalarAssign
     value: !ScalarExpression
       arith_fn:
         fn_name: max_unsigned
         operands:
         - !ScalarExpression
           scalar_arg: O
         - !ScalarExpression
           type_fn:
-            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
- !ScalarAssign
     value: !ScalarExpression
       arith_fn:
         fn_name: max
         operands:
         - !ScalarExpression
           scalar_arg: O
         - !ScalarExpression
           type_fn:
-            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
- !ScalarAssign
     value: !ScalarExpression
       arith_fn:
         fn_name: min
         operands:
         - !ScalarExpression
           scalar_arg: O
         - !ScalarExpression
           type_fn:
-            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
- !ScalarAssign
     value: !ScalarExpression
       arith_fn:
         fn_name: min_unsigned
         operands:
         - !ScalarExpression
           scalar_arg: O
         - !ScalarExpression
           type_fn:
-            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
- !ScalarAssign
     value: !ScalarExpression
       arith_fn:
         fn_name: add
         operands:
         - !ScalarExpression
           scalar_arg: O
         - !ScalarExpression
           type_fn:
-            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
- !ScalarAssign
     value: !ScalarExpression
       arith_fn:
         fn_name: max
         operands:
         - !ScalarExpression
           scalar_arg: O
         - !ScalarExpression
           type_fn:
-            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
- !ScalarAssign
     value: !ScalarExpression
       arith_fn:
         fn_name: min
         operands:
         - !ScalarExpression
           scalar_arg: O
         - !ScalarExpression
           type_fn:
-            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
     - affine_map<() -> ()>
     - affine_map<() -> ()>
   iterator_types: []
   assignments:
   - !ScalarAssign
     arg: O
     value: !ScalarExpression
       type_fn:
-        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
structured_op: !LinalgStructuredOpConfig
   iterator_types:
   - parallel
   - parallel
   assignments:
   - !ScalarAssign
     arg: O
     value: !ScalarExpression
       type_fn:
-        fn_name: cast
         type_var: T
         operands:
         - !ScalarExpression
           arith_fn:
             fn_name: add
             operands:
             - !ScalarExpression
               arith_fn:
                 fn_name: mul
                 operands:
                 - !ScalarExpression
                   arith_fn:
                     fn_name: add
                     operands:
                     - !ScalarExpression
                       type_fn:
-                        fn_name: cast
                         type_var: F64
                         operands:
                         - !ScalarExpression
                           scalar_const: '2147483647 : i64'
+                        fn_name: cast
                     - !ScalarExpression
                       type_fn:
-                        fn_name: cast
                         type_var: F64
                         operands:
                         - !ScalarExpression
                           arith_fn:
                             fn_name: add
                             operands:
                             - !ScalarExpression
                               arith_fn:
                                 fn_name: mul
                                 operands:
                                 - !ScalarExpression
                                   arith_fn:
                                     fn_name: add
                                     operands:
                                     - !ScalarExpression
                                       type_fn:
-                                        fn_name: cast
                                         type_var: I32
                                         operands:
                                         - !ScalarExpression
                                           scalar_index: 1
+                                        fn_name: cast
                                     - !ScalarExpression
                                       arith_fn:
                                         fn_name: add
                                         operands:
                                         - !ScalarExpression
                                           arith_fn:
                                             fn_name: mul
                                             operands:
                                             - !ScalarExpression
                                               arith_fn:
                                                 fn_name: add
                                                 operands:
                                                 - !ScalarExpression
                                                   type_fn:
-                                                    fn_name: cast
                                                     type_var: I32
                                                     operands:
                                                     - !ScalarExpression
                                                       scalar_index: 0
+                                                    fn_name: cast
                                                 - !ScalarExpression
                                                   scalar_arg: seed
                                             - !ScalarExpression
                                               type_fn:
-                                                fn_name: cast
                                                 type_var: I32
                                                 operands:
                                                 - !ScalarExpression
                                                   scalar_const: '1103515245 : i64'
+                                                fn_name: cast
                                         - !ScalarExpression
                                           type_fn:
-                                            fn_name: cast
                                             type_var: I32
                                             operands:
                                             - !ScalarExpression
                                               scalar_const: '12345 : i64'
+                                            fn_name: cast
                                 - !ScalarExpression
                                   type_fn:
-                                    fn_name: cast
                                     type_var: I32
                                     operands:
                                     - !ScalarExpression
                                       scalar_const: '1103515245 : i64'
+                                    fn_name: cast
                             - !ScalarExpression
                               type_fn:
-                                fn_name: cast
                                 type_var: I32
                                 operands:
                                 - !ScalarExpression
                                   scalar_const: '12345 : i64'
+                                fn_name: cast
+                        fn_name: cast
                 - !ScalarExpression
                   arith_fn:
                     fn_name: mul
                     operands:
                     - !ScalarExpression
                       arith_fn:
                         fn_name: sub
                         operands:
                         - !ScalarExpression
                           scalar_arg: max
                         - !ScalarExpression
                           scalar_arg: min
                     - !ScalarExpression
                       type_fn:
-                        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
value: !ScalarExpression
         fn_name: log
         operands:
         - !ScalarExpression
           arith_fn:
             fn_name: add
             operands:
             - !ScalarExpression
               type_fn:
-                fn_name: cast
                 type_var: U
                 operands:
                 - !ScalarExpression
                   scalar_const: '1.000000e+00 : f64'
+                fn_name: cast
             - !ScalarExpression
               arith_fn:
                 fn_name: exp
                 operands:
                 - !ScalarExpression
                   type_fn:
-                    fn_name: cast
                     type_var: U
                     operands:
                     - !ScalarExpression
                       scalar_arg: I
+                    fn_name: cast

mlir/include/mlir/Dialect/Linalg/IR/LinalgStructuredOps.td

class LinalgStructuredBase_Op<string mnemonic, list<Trait> props>
   code structuredOpsBaseDecls = [{
     // Return whether the op accesses the iteration indices.
     bool hasIndexSemantics() {
       return !this->getBody()->getOps<IndexOp>().empty();
     }
 
     LogicalResult reifyResultShapes(OpBuilder &b,
         ReifiedRankedShapedTypeDims &reifiedReturnShapes) {
-      return cast<LinalgOp>(getOperation()).reifyResultShapes(b,
+      return llvm::cast<LinalgOp>(getOperation()).reifyResultShapes(b,
           reifiedReturnShapes);
     }
   }];
 }
 
 class LinalgStructured_Op<string mnemonic, list<Trait> props>
   : LinalgStructuredBase_Op<mnemonic, !listconcat(props, [])> {
   code structuredOpsDecls = structuredOpsBaseDecls # [{

mlir/lib/Dialect/Linalg/IR/CMakeLists.txt

 add_mlir_dialect_library(MLIRLinalg
   LinalgInterfaces.cpp
   LinalgOps.cpp
   LinalgDialect.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/Linalg
 
   DEPENDS
   MLIRLinalgInterfacesIncGen
+  MLIRLinalgOpsAttributesIncGen
+  MLIRLinalgOpsEnumsIncGen
   MLIRLinalgOpsIncGen
   MLIRLinalgStructuredOpsIncGen
 
   LINK_LIBS PUBLIC
   MLIRAffine
   MLIRArithmetic
   MLIRBufferization
   MLIRDialectUtils

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

 #include "mlir/IR/Dialect.h"
 #include "mlir/IR/DialectImplementation.h"
 #include "mlir/IR/FunctionInterfaces.h"
 #include "mlir/Parser.h"
 #include "mlir/Support/LLVM.h"
 #include "mlir/Transforms/InliningUtils.h"
 
 #include "llvm/ADT/StringExtras.h"
+#include "llvm/ADT/TypeSwitch.h"
 #include "llvm/Support/raw_ostream.h"
 
 using namespace mlir;
 using namespace mlir::linalg;
 
 //===----------------------------------------------------------------------===//
 // LinalgDialect Dialect Interfaces
 //===----------------------------------------------------------------------===//

 template <typename... OpTypes>
 void addNamedOpBuilders(
     llvm::StringMap<LinalgDialect::RegionBuilderFunType> &map) {
   (void)std::initializer_list<int>{0,
                                    (addNamedOpBuilderImpl<OpTypes>(map), 0)...};
 }
 
 void mlir::linalg::LinalgDialect::initialize() {
+  addAttributes<
+#define GET_ATTRDEF_LIST
+#include "mlir/Dialect/Linalg/IR/LinalgOpsAttrDefs.cpp.inc"
+      >();
   addOperations<
 #define GET_OP_LIST
 #include "mlir/Dialect/Linalg/IR/LinalgOps.cpp.inc"
       >();
   addOperations<
 #define GET_OP_LIST
 #include "mlir/Dialect/Linalg/IR/LinalgStructuredOps.cpp.inc"
       >();
if (!isa<FunctionOpInterface>(op))
                              << " to be used on function-like operations";
     return success();
   }
   if (attr.getName() == LinalgDialect::kMemoizedIndexingMapsAttrName)
     return success();
   return op->emitError() << "attribute '" << attr.getName()
                          << "' not supported by the linalg dialect";
 }
+
+#include "mlir/Dialect/Linalg/IR/LinalgOpsEnums.cpp.inc"
+
+#define GET_ATTRDEF_CLASSES
+#include "mlir/Dialect/Linalg/IR/LinalgOpsAttrDefs.cpp.inc"
+
+#include "mlir/Dialect/Linalg/IR/LinalgOpsDialect.cpp.inc"

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

 #include "llvm/ADT/TypeSwitch.h"
 #include "llvm/Support/FormatVariadic.h"
 #include "llvm/Support/MathExtras.h"
 #include "llvm/Support/raw_ostream.h"
 
 using namespace mlir;
 using namespace mlir::linalg;
 
-#include "mlir/Dialect/Linalg/IR/LinalgOpsDialect.cpp.inc"
-
 /// Forward declarations.
 
 /// Generic entry point to create the block for the region of a LinalgOp.
 /// This is used by both named structured ops created by ods-gen and by manually
 /// defined C++ ops.
 /// This is used by both builders and parsers.
 /// This function creates the block in the region with arguments corresponding
 /// to the elemental types of `inputTypes` and `outputTypes`. The latter are
if (auto toIntType = toType.dyn_cast<IntegerType>()) {
       }
     }
 
     emitWarning(operand.getLoc()) << "could not cast operand of type "
                                   << operand.getType() << " to " << toType;
     return operand;
   }
 
-  // NOLINTNEXTLINE(*-identifier-naming): externally called.
-  Value typefn__cast(Type toType, Value operand) {
+  Value buildTypeFn(TypeFn typeFn, Type toType, Value operand) {
+    switch (typeFn) {
+    case TypeFn::cast:
       return cast(toType, operand, false);
-  }
-
-  // NOLINTNEXTLINE(*-identifier-naming): externally called.
-  Value typefn__cast_unsigned(Type toType, Value operand) {
+    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) {
     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);

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

def _compute_reduce_dims(self, rhs: TensorExpression) -> Set[DimDef]:
     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)
 
   def __repr__(self):
-    return f"{self.tensor_name}[{', '.join([repr(i) for i in self.indices])}]"
+    return (f"{self.operand_def.name}"
+            f"[{', '.join([repr(i) for i in self.indices])}]")
 
 
 class TensorArithFn(TensorExpression):
   """Application of an arithmetic function."""
 
   def __init__(self, arith_fn: "ArithFnType", args: Sequence[TensorExpression]):
     self.arith_fn = arith_fn
     self.args = tuple(args)
class TensorArithFn(TensorExpression):
 
   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: "TypeFn", type_var: TypeVar,
+  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
+    self.operand_def = operand_def
     self.type_var = type_var
     self.arg = arg
 
   def to_scalar_expression(self) -> ScalarExpression:
-    return ScalarTypeFn(self.type_fn.fn_name, self.type_var,
+    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
+    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()
 
   def visit_tensor_exprs(self, callback: Callable[["TensorExpression"], None]):
     super().visit_tensor_exprs(callback)
     self.arg.visit_tensor_exprs(callback)
 
   def __repr__(self):
-    return f"{repr(self.type_fn)}({self.type_var}, {self.arg})"
+    return (f"{repr(self.type_fn)}[{repr(self.operand_def)}]"
+            f"({self.type_var}, {self.arg})")
 
 
 class TensorReduceFn(TensorExpression):
   """Application of a reduction function.
 
   This captures the lhs (initial value) separately from the rhs.
   """
 
class TypeFnType:
   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, type_var, arg)
+    return TensorTypeFn(self, None, type_var, arg)
 
   def __repr__(self):
     return f"{self.fn_name}"
 
 
 class TypeFn:
   """Type conversion function namespace.
 

 ###############################################################################
 
 
 class OperandKind(Enum):
   InputTensor = 0
   Scalar = 1
   OutputTensor = 2
   IndexAttr = 3
+  TypeFnAttr = 4
 
 
 class OperandDef:
   """Definition of an operand passed to an operation.
 
   Keep the meta information of Tensor, Scalar, and Attribute operands and
   provide the shared registration functionality.
   """
 
   def __init__(self,
                kind: OperandKind,
                type_var: Optional[TypeVar] = None,
                size_exprs: Optional[Sequence[AffineExprDef]] = None,
                index_dims: Optional[Sequence[DimDef]] = None,
-               default_vals: Optional[Sequence[int]] = None):
+               default_vals: Optional[Sequence[int]] = None,
+               default_fn: Optional[str] = None):
     if type_var and not isinstance(type_var, TypeVar):
       raise ValueError(
           f"OperandDef requires a TypeVar but got {repr(type_var)}")
     self.owner = None  # type: Optional["LinalgOpDef"]
     self.type_var = type_var
     self.size_exprs = size_exprs
     self.index_dims = index_dims
     self.default_vals = default_vals
+    self.default_fn = default_fn
     self.kind = kind
     self.name = None  # type: Optional[str]
     self.registered_index = -1  # type: int
 
   def attach(self, index: int, name: str, owner: "LinalgOpDef"):
     if self.owner:
       raise ValueError(f"OperandDef already registered with op: {self}")
     self.registered_index = index
     self.name = name
     self.owner = owner
 
+  def is_input(self) -> bool:
+    return (self.kind == OperandKind.Scalar or
+            self.kind == OperandKind.InputTensor)
+
+  def is_tensor(self) -> bool:
+    return (self.kind == OperandKind.InputTensor or
+            self.kind == OperandKind.OutputTensor)
+
+  def is_attribute(self) -> bool:
+    return (self.kind == OperandKind.IndexAttr or
+            self.kind == OperandKind.TypeFnAttr)
+
   def __hash__(self):
     return hash(id(self))
 
   def __repr__(self):
     return (f"{self.name}:OperandDef(kind={self.kind.name}, "
-            f"type={repr(self.type_var)}, size_exprs={self.size_exprs}), "
-            f"index_dims={self.index_dims}, default_vals={self.default_vals})")
+            f"type={repr(self.type_var)}, size_exprs={self.size_exprs}, "
+            f"index_dims={self.index_dims}, default_vals={self.default_vals}, "
+            f"default_fn={self.default_fn})")
 
 
 class TensorDef:
   """Tensor operand definition.
 
   Tensor operands are indexed using the associated indexing_map when forwarded
   to the body of the structured op. A unique name identifies the tensor operands
   and an index determines their position in the operation's parameter list. A
if any(not isinstance(default_val, int) for default_val in default):
                        f"but got {default}")
     if len(sizes) != len(default):
       raise ValueError(f"IndexAttrDef expects {len(sizes)} default values "
                        f"but got {len(default)}")
     self.operand_def = OperandDef(
         OperandKind.IndexAttr, size_exprs=sizes, default_vals=default)
 
 
+class TypeFnAttrDef:
+  """Type conversion function attribute definition.
+
+  Type conversion function attributes provide a way to make type conversions
+  parametrizable. Every attribute specifies a default type conversion function
+  that may be overwritten at operation instantiation time.
+  """
+
+  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.TypeFnAttr, default_fn=default.fn_name)
+
+  def __call__(self, type_var: TypeVar, arg: TensorExpression) -> TensorTypeFn:
+    return TensorTypeFn(None, self.operand_def, type_var, arg)
+
+
 ###############################################################################
 # Operation definition.
 ###############################################################################
 
 
 class Comprehension:
   """Represents a single comprehension."""
 
class LinalgOpDef:
 
   def add_operand(self, name: str, operand: OperandDef):
     """Registers an operand."""
     if name in self.registered_operands:
       raise ValueError(f"The operand {name} is already registered "
                        f"to {self.registered_operands['name']}")
     # Ensure output tensors are registered after input tensors and scalars and
     # attributes are registered after all other operand types.
-    registered_kinds = [
-        operand.kind.value for operand in self.registered_operands.values()
-    ]
-    if registered_kinds:
-      maximum = max(registered_kinds)
-      if maximum > operand.kind.value and maximum > OperandKind.Scalar.value:
-        raise ValueError(
-            f"The operand {name} of kind {operand.kind.name} is registered "
-            f"after an operand of kind {OperandKind(maximum).name}")
+    if operand.is_input() and any(
+        not op_def.is_input() for op_def in self.registered_operands.values()):
+      raise ValueError(f"Input {name} registered after an output or attribute")
+    if operand.kind == OperandKind.OutputTensor and any(
+        op_def.is_attribute() for op_def in self.registered_operands.values()):
+      raise ValueError(f"Output {name} registered after an attribute")
     operand.attach(len(self.registered_operands), name, self)
     self.registered_operands[name] = operand
 
   def __repr__(self):
     lines = [
         f"LinalgOpDef({self.metadata.name} -> {self.metadata.cpp_class_name},"
     ]
     for name, operand in self.registered_operands.items():
       lines.append(f"  {operand}")
     if self.comprehensions:
       lines[-1] += " {"
       for comprehension in self.comprehensions:
         lines.append(f"    {comprehension}")
       lines.append("}")
     return "\n".join(lines)

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

class OperandDefConfig(YAMLObject):
   @property
   def type_var(self) -> TypeVar:
     return self.operand_def.type_var
 
   @property
   def usage(self) -> str:
     if self.operand_def.kind == OperandKind.IndexAttr:
       return "IndexAttr"
+    if self.operand_def.kind == OperandKind.TypeFnAttr:
+      return "TypeFnAttr"
     if self.operand_def.kind == OperandKind.OutputTensor:
       return "Output"
     return "Input"
 
   def to_yaml_custom_dict(self):
     self_dict = dict(name=self.name, usage=self.usage)
     if self.type_var:
       self_dict["type_var"] = self.type_var.name
     if self.shape_map:
       self_dict["shape_map"] = _serialize_affine_map(self.shape_map)
     if self.index_attr_map:
       self_dict["index_attr_map"] = _serialize_affine_map(self.index_attr_map)
     if self.operand_def.default_vals:
       self_dict["default_vals"] = self.operand_def.default_vals
+    if self.operand_def.default_fn:
+      self_dict["default_fn"] = self.operand_def.default_fn
     return self_dict
 
   def __repr__(self):
     return (f"OperandDefConfig({self.operand_def}, "
             f"shape_map={self.shape_map}, "
             f"index_attr_map={self.index_attr_map}, "
             f"indexing_map={self.indexing_map})")
 
with self.context:
       local_state = AffineBuildState(
           global_state=self.affine_state, allow_new_symbols=False)
       for dim in domain:
         dim.build(state=local_state)
 
     # Collect all attribute definitions.
     collected_attr_defs = list()
     for operand in registered_operands:
-      if operand.kind == OperandKind.IndexAttr:
+      if operand.is_attribute():
         collected_attr_defs.append(operand)
 
     # Collect all tensors with manual indexing annotation.
     collected_index_defs = list()
     for operand in registered_operands:
       if operand.index_dims:
         if any(dim not in collected_dim_uses for dim in operand.index_dims):
           raise ValueError(f"Expected all index dims {operand.index_dims} of "
def __init__(self,
 
     # Set the indexing map of all scalar uses to the empty map.
     for operand_config in self.operands.values():
       if operand_config.operand_def.kind == OperandKind.Scalar:
         operand_config.indexing_map = self._get_scalar_map()
 
     # Check all registered tensor and scalar operands have an indexing map.
     for operand in registered_operands:
-      if operand.kind == OperandKind.IndexAttr:
+      if operand.is_attribute():
         continue
-      if not (operand in self.operands and self.operands[operand].indexing_map):
+      if not (operand in self.operands and
+              self.operands[operand].indexing_map):
         raise ValueError(f"Failed to compute an indexing map for operand "
                          f"{operand.name}")
 
     # Collect reduction dims and ensure all the same.
     all_reduction_dims = set(comprehension.all_reduction_dims)
     if len(all_reduction_dims) != 1:
       raise ValueError(
           f"All writes within a generic must have the same reduction "
def get_type(symbolic_name, position):
           return "reduction"
       return "parallel"
 
     return [get_type(*dim) for dim in self.ordered_dims]
 
   def add_operand(self, operand_def: OperandDef):
     if operand_def in self.operands:
       return
-    if operand_def.kind == OperandKind.Scalar:
+    if (operand_def.kind == OperandKind.Scalar or
+        operand_def.kind == OperandKind.TypeFnAttr):
       self.operands[operand_def] = OperandDefConfig(operand_def)
       return
     with self.context:
       local_state = AffineBuildState(
           global_state=self.affine_state, allow_new_dims=False)
       exprs = []
       for expr in operand_def.size_exprs:
         exprs.append(expr.build(state=local_state))

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

def linalg_structured_op(dsl_func=None,
   op_def = LinalgOpDef(
       name=op_name, cpp_class_name=op_class_name, doc=inspect.getdoc(dsl_func))
 
   # Extract arguments and TensorDefs from the signature.
   dsl_func_args = list()
   sig = inspect.signature(dsl_func)
   for param_name, param in sig.parameters.items():
     param_default = param.default
-    if isinstance(param_default, (TensorDef, ScalarDef, IndexAttrDef)):
+    if isinstance(param_default,
+                  (TensorDef, ScalarDef, IndexAttrDef, TypeFnAttrDef)):
       op_def.add_operand(param_name, param_default.operand_def)
     else:
       raise ValueError(
           f"@linalg_structured_op function parameters must be defaulted as "
           f"TensorDef(...), ScalarDef(...), or IndexAttrDef(...): "
           f"Found {param_name}: {param_default}")
     dsl_func_args.append(param_default)
 

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

try:
     cls(ty)
     return True
   except ValueError:
     return False
 
 
 def prepare_common_structured_op(op_config: LinalgStructuredOpConfig,
                                  *ins: Value, outs: ValueList,
-                                 **attrs: Sequence[int]):
+                                 **attrs: Union[Sequence[int], TypeFnType]):
   all_arg_defs = op_config.ordered_operands
   in_arg_defs = [d for d in all_arg_defs if d.usage == "Input"]
   out_arg_defs = [d for d in all_arg_defs if d.usage == "Output"]
   index_attr_arg_defs = [d for d in all_arg_defs if d.usage == "IndexAttr"]
+  type_fn_attr_arg_defs = [d for d in all_arg_defs if d.usage == "TypeFnAttr"]
 
   # Verify outs is a sequence or a list of results.
   if not isinstance(outs, (Sequence, OpResultList)):
     raise ValueError(f"Expected named argument outs to have type Sequence or "
                      f"OpResultLis but got {type(outs)}")
 
   # Arity validation.
   if len(ins) != len(in_arg_defs):
     raise ValueError(f"Expected {len(in_arg_defs)} inputs but got "
                      f"{len(ins)} for {op_config}")
   if outs and len(outs) != len(out_arg_defs):
     raise ValueError(f"Expected {len(out_arg_defs)} outputs but got "
                      f"{len(outs)} for {op_config}")
 
-  # Compute a replacement list for all attribute symbols.
+  # Compute a replacement list for all index attribute symbols.
   expressions = []  # type: Sequence[AffineExpr]
   replacements = []  # type: Sequence[AffineExpr]
   for index_attr in index_attr_arg_defs:
     index_attr_vals = index_attr.operand_def.default_vals
     if index_attr.name in attrs:
       index_attr_vals = attrs.get(index_attr.name)
     assert index_attr_vals, "Index attribute has no value"
     if not all(isinstance(value, int) for value in index_attr_vals):
def prepare_common_structured_op(op_config: LinalgStructuredOpConfig,
   index_attrs = {}  # type: Dict[str, DenseElementAttr]
   for index_attr in index_attr_arg_defs:
     index_attr_vals = attrs.get(index_attr.name)
     # Only forward attributes set to a non-default value.
     if index_attr_vals:
       array = np.array(index_attr_vals, dtype=np.int64)
       index_attrs[index_attr.name] = DenseElementsAttr.get(array)
 
+  # Compute the type function attribue mapping.
+  type_fn_attr_mapping = {}
+  for type_fn_attr in type_fn_attr_arg_defs:
+    attr_val = type_fn_attr.operand_def.default_fn
+    if type_fn_attr.name in attrs:
+      type_fn = attrs.get(type_fn_attr.name)
+      if not isinstance(type_fn, TypeFnType):
+        raise ValueError(f"Attribute {type_fn_attr.name} needs to be of type "
+                         f"TypeFnType but got {type(attr_val)}")
+      attr_val = type_fn.fn_name
+    assert attr_val, "Type function attribute has no value"
+    type_fn_attr_mapping[type_fn_attr.name] = attr_val
+
   return (all_arg_defs, in_arg_defs, out_arg_defs, outs, result_types,
-          type_mapping, indexing_maps_attr, iterator_types_attr,
-          index_attrs, block_arg_types)
+          type_mapping, indexing_maps_attr, iterator_types_attr, index_attrs,
+          type_fn_attr_mapping, block_arg_types)
 
 
 def emit_generic_structured_op(op_config: LinalgStructuredOpConfig, *ins: Value,
                                outs: ValueList, **attrs: Sequence[int]):
   all_arg_defs, in_arg_defs, out_arg_defs, outs, result_types, type_mapping, \
-  indexing_maps_attr, iterator_types_attr, index_attrs, block_arg_types = \
+  indexing_maps_attr, iterator_types_attr, index_attrs, type_fn_attr_mapping, \
+  block_arg_types = \
      prepare_common_structured_op(op_config, *ins, outs = outs, **attrs)
 
   # An operation that accesses only scalars and scalar/rank zero tensors is
   # rank polymorhpic. We implement rank polymorphism by generating different
   # indexing maps and iterators that match the rank of the first output tensor.
   # An operation is rank polymorphic if the iteration domain has rank zero.
   if not iterator_types_attr:
     rank = ShapedType(outs[0].type).rank
     iterator_types_attr = ArrayAttr.get([StringAttr.get("parallel")] * rank)
     scalar_map = AffineMap.get(rank, 0, [])
     tensor_map = AffineMap.get_identity(rank)
     indexing_maps = []
     for arg_def in all_arg_defs:
       if arg_def.operand_def.kind == OperandKind.Scalar:
         indexing_maps.append(scalar_map)
-      if (arg_def.operand_def.kind == OperandKind.InputTensor or
-          arg_def.operand_def.kind == OperandKind.OutputTensor):
+      if arg_def.operand_def.is_tensor():
         indexing_maps.append(tensor_map)
     indexing_maps_attr = ArrayAttr.get(
         [AffineMapAttr.get(am) for am in indexing_maps])
 
   generic_op = linalg.GenericOp(
       result_tensors=result_types,
       inputs=ins,
       outputs=outs,
       indexing_maps=indexing_maps_attr,
       iterator_types=iterator_types_attr,
       doc=None,  # TODO: Make optional.
       library_call=None)  # TODO: Make optional.
 
   # Construct the body.
   block_arg_names = _get_operand_def_names(*in_arg_defs, *out_arg_defs)
   block = generic_op.regions[0].blocks.append(*block_arg_types)
   block_arg_mapping = dict(zip(block_arg_names, block.arguments))
   with InsertionPoint(block):
-    body_builder = _BodyBuilder(type_mapping, block_arg_mapping)
+    body_builder = _BodyBuilder(type_mapping, block_arg_mapping,
+                                type_fn_attr_mapping)
     for assignment in op_config.assignments:
       body_builder.assign(assignment)
     body_builder.yield_outputs(*_get_operand_def_names(*out_arg_defs))
 
   if len(result_types) == 1:
     return generic_op.result
   else:
     return generic_op.results
 
 
 def emit_named_structured_op(op_config: LinalgStructuredOpConfig, op_name: str,
                              op_class_name: str, *ins: Value, outs: ValueList,
                              **attrs: Sequence[int]):
   all_arg_defs, in_arg_defs, out_arg_defs, outs, result_types, type_mapping, \
-  indexing_maps_attr, iterator_types_attr, index_attrs, block_arg_types = \
+  indexing_maps_attr, iterator_types_attr, index_attrs, type_fn_attr_mapping, \
+  block_arg_types = \
      prepare_common_structured_op(op_config, *ins, outs = outs, **attrs)
 
   # If we get here, there must exist a builtin class `op_class_name`.
   ctx = Context.current
   fully_qualified_name = "linalg." + op_name
   if (not ctx.is_registered_operation(fully_qualified_name) or
       not op_class_name in linalg.__dict__.keys()):
     raise NotImplementedError(
         f"Unknown named op_name / op_class_name: {op_name} / {op_class_name}")
 
   # Set the index attributes used to compute the indexing maps.
   named_op = getattr(linalg, op_class_name)(ins, outs, result_types)
   for name, value in index_attrs.items():
     named_op.operation.attributes[name] = value
 
+  # Set the type function attributes.
+  for name, value in type_fn_attr_mapping.items():
+    named_op.operation.attributes[name] = Attribute.parse(
+        f"#linalg<\"type_fn {value}\">")
+
   linalg.fill_builtin_region(named_op.operation)
 
   if len(result_types) == 1:
     return named_op.result
   else:
     return named_op.results
 
 
 class _BodyBuilder:
   """Constructs a structured op body by evaluating assignments."""
 
   def __init__(self, type_mapping: Dict[str, Type],
-               block_arg_mapping: Dict[str, Value]):
+               block_arg_mapping: Dict[str, Value],
+               type_fn_attr_mapping: Dict[str, str]):
     self.type_mapping = type_mapping
     self.block_arg_mapping = block_arg_mapping
+    self.type_fn_attr_mapping = type_fn_attr_mapping
     self.yield_mapping = dict()  # type: Dict[str, Value]
 
   def assign(self, assignment: ScalarAssign):
     if assignment.arg in self.yield_mapping:
       raise ValueError(
           f"Multiple assignments to the same argument are forbidden: "
           f"{assignment}")
     self.yield_mapping[assignment.arg] = self.expression(assignment.value)
elif expr.scalar_index:
       return linalg.IndexOp(dim_attr).result
     elif expr.arith_fn:
       fn = self._get_function(f"_arithfn_{expr.arith_fn.fn_name}")
       operand_values = [
           self.expression(operand) for operand in expr.arith_fn.operands
       ]
       return fn(*operand_values)
     elif expr.type_fn:
-      fn = self._get_function(f"_typefn_{expr.type_fn.fn_name}")
+      fn_name = expr.type_fn.fn_name
+      if not fn_name:
+        fn_name = self.type_fn_attr_mapping[expr.type_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)
     raise NotImplementedError(f"Unimplemented scalar body expression: {expr}")
 
   def yield_outputs(self, *output_names: str):
     output_values = []
     for n in output_names:
       try:

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

class ScalarArithFn:
 
   def __repr__(self):
     return f"ScalarArithFn<{self.fn_name}>({', '.join(self.operands)})"
 
 
 class ScalarTypeFn:
   """A type of ScalarExpression that applies a type conversion function."""
 
-  def __init__(self, fn_name: str, type_var: TypeVar,
-               operand: "ScalarExpression"):
+  def __init__(self, fn_name: Optional[str], attr_name: Optional[str],
+               type_var: TypeVar, operand: "ScalarExpression"):
     self.fn_name = fn_name
+    self.attr_name = attr_name
     self.type_var = type_var
     self.operand = operand
 
   def expr(self) -> "ScalarExpression":
     return ScalarExpression(type_fn=self)
 
   def __repr__(self):
-    return f"ScalarTypeFn<{self.fn_name}>({self.type_var}, {self.operand})"
+    return (f"ScalarTypeFn<{self.fn_name}[{self.attr_name}]>"
+            f"({self.type_var}, {self.operand})")
 
 
 class ScalarArg:
   """A type of ScalarExpression that references a named argument."""
 
   def __init__(self, arg: str):
     self.arg = arg
 
if self.arith_fn:
           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.
-      return dict(
-          type_fn=dict(
-              fn_name=self.type_fn.fn_name,
-              type_var=self.type_fn.type_var.name,
-              operands=[self.type_fn.operand],
-          ))
+      type_fn_dict = dict(
+      )
+      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)
     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/python/mlir/dialects/linalg/opdsl/ops/core_named_ops.py

 from ..lang import *
 
 T1 = TV.T1
 T2 = TV.T2
 
 Batch = S.Batch
 
 
 @linalg_structured_op
 def matmul(
     A=TensorDef(T1, S.M, S.K),
     B=TensorDef(T2, S.K, S.N),
-    C=TensorDef(U, S.M, S.N, output=True)):
+    C=TensorDef(U, S.M, S.N, output=True),
+    cast=TypeFnAttrDef(default=TypeFn.cast)):
   """Performs a matrix multiplication of two 2D inputs.
 
   Numeric casting is performed on the operands to the inner multiply, promoting
   them to the same data type as the accumulator/output.
   """
   domain(D.m, D.n, D.k)
   implements(ContractionOpInterface)
-  C[D.m, D.n] += TypeFn.cast(U, A[D.m, D.k]) * TypeFn.cast(U, B[D.k, D.n])
+  C[D.m, D.n] += cast(U, A[D.m, D.k]) * cast(U, B[D.k, D.n])
 
 
 @linalg_structured_op
 def matmul_unsigned(
     A=TensorDef(T1, S.M, S.K),
     B=TensorDef(T2, S.K, S.N),
     C=TensorDef(U, S.M, S.N, output=True)):
   """Performs an unsigned matrix multiplication of two 2D inputs.

mlir/test/Dialect/Linalg/generalize-named-polymorphic-ops.mlir

 // CHECK-NEXT:   %[[ADD:.+]] = arith.addf %[[C_ARG]], %[[MUL]] : f32
 // CHECK-NEXT:   linalg.yield %[[ADD]] : f32
 // CHECK-NEXT: -> tensor<16x32xf32>
 
 // -----
 
 // Verifies that different argument types is legal.
 func @generalize_matmul_tensor_i16i64i32(%A : tensor<16x8xi16>, %B: tensor<8x32xi64>, %C: tensor<16x32xi32>) -> tensor<16x32xi32> {
-  %0 = linalg.matmul ins(%A, %B: tensor<16x8xi16>, tensor<8x32xi64>)
+  %0 = linalg.matmul {cast = #linalg<"type_fn cast">}
+                     ins(%A, %B: tensor<16x8xi16>, tensor<8x32xi64>)
                           outs(%C: tensor<16x32xi32>) -> tensor<16x32xi32>
   return %0: tensor<16x32xi32>
 }
 
 // CHECK-LABEL: @generalize_matmul_tensor_i16i64i32
 // CHECK:      ^{{.*}}(%[[A_ARG:.+]]: i16, %[[B_ARG:.+]]: i64, %[[C_ARG:.+]]: i32)
 // Verify signed integer extension and truncation.
 // CHECK-NEXT:   %[[A_CAST:.+]] = arith.extsi %[[A_ARG]] : i16 to i32
 // CHECK-NEXT:   %[[B_CAST:.+]] = arith.trunci %[[B_ARG]] : i64 to i32
 // CHECK-NEXT:   %[[MUL:.+]] = arith.muli %[[A_CAST]], %[[B_CAST]] : i32
 // CHECK-NEXT:   %[[ADD:.+]] = arith.addi %[[C_ARG]], %[[MUL]] : i32
 // CHECK-NEXT:   linalg.yield %[[ADD]] : i32
 // CHECK-NEXT: -> tensor<16x32xi32>
 
+
+// -----
+
+// Verifies the cast attribute controls the cast operations used.
+func @generalize_matmul_tensor_i16i64i32_unsigned(%A : tensor<16x8xi16>, %B: tensor<8x32xi64>, %C: tensor<16x32xi32>) -> tensor<16x32xi32> {
+  %0 = linalg.matmul {cast = #linalg<"type_fn cast_unsigned">}
+                     ins(%A, %B: tensor<16x8xi16>, tensor<8x32xi64>)
+                          outs(%C: tensor<16x32xi32>) -> tensor<16x32xi32>
+  return %0: tensor<16x32xi32>
+}
+
+// CHECK-LABEL: @generalize_matmul_tensor_i16i64i32_unsigned
+// CHECK:        = arith.extui
+
 // -----
 
 func @generalize_matmul_tensor_i16i64f32(%A : tensor<16x8xi16>, %B: tensor<8x32xi64>, %C: tensor<16x32xf32>) -> tensor<16x32xf32> {
   %0 = linalg.matmul ins(%A, %B: tensor<16x8xi16>, tensor<8x32xi64>)
                      outs(%C: tensor<16x32xf32>) -> tensor<16x32xf32>
   return %0: tensor<16x32xf32>
 }
 

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

 # RUN: mlir-linalg-ods-yaml-gen %s --o-ods-decl=- | FileCheck %s --check-prefix=ODS
 # RUN: mlir-linalg-ods-yaml-gen %s --o-impl=- | FileCheck %s --check-prefix=IMPL
 
 # @linalg_structured_op
-# def test1(O=TensorDef(T, S.M, S.N, output=True)):
+# def test1(O=TensorDef(T, S.M, S.N, output=True),
+#           cast=TypeFnAttrDef(default=TypeFn.cast)):
 #   """Title.
 
 #   Detailed description.
 #   """
 #   O[D.m, D.n] = cast(T, const(42)) + cast(T, index(D.n))
 
 --- !LinalgOpConfig
 metadata: !LinalgOpMetadata
   name: test1
   cpp_class_name: Test1Op
   doc: |-
     Title.
 
     Detailed description.
 structured_op: !LinalgStructuredOpConfig
   args:
   - !LinalgOperandDefConfig
     name: O
     usage: Output
     type_var: T
     shape_map: affine_map<()[s0, s1] -> (s0, s1)>
+  - !LinalgOperandDefConfig
+    name: cast
+    usage: TypeFnAttr
+    default_fn: cast
   indexing_maps: !LinalgIndexingMapsConfig
     static_indexing_maps:
     - affine_map<(d0, d1)[s0, s1] -> (d0, d1)>
   iterator_types:
   - parallel
   - parallel
   assignments:
   - !ScalarAssign
     arg: O
     value: !ScalarExpression
       arith_fn:
         fn_name: add
         operands:
         - !ScalarExpression
           type_fn:
-            fn_name: cast
             type_var: T
             operands:
             - !ScalarExpression
               scalar_const: '42 : i64'
+            attr_name: cast
         - !ScalarExpression
           type_fn:
-            fn_name: cast_unsigned
             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:  }];
 
 #       ODS:  let arguments =
 #  ODS-NEXT:    Variadic<AnyType>:$inputs,
-#  ODS-NEXT:    Variadic<AnyShaped>:$outputs
+#  ODS-NEXT:    Variadic<AnyShaped>:$outputs,
+#  ODS-NEXT:    DefaultValuedAttr<TypeFnAttr, "TypeFn::cast">:$cast
 
 #       ODS:  let builders =
 #       ODS:  (ins "TypeRange":$resultTensorTypes, "ValueRange":$inputs,
 #  ODS-NEXT:       "ValueRange":$outputs,
 #  ODS-NEXT:       CArg<"ArrayRef<NamedAttribute>", "{}">:$attributes),
 
+#       ODS:  (ins "TypeRange":$resultTensorTypes, "ValueRange":$inputs,
+#  ODS-NEXT:       "ValueRange":$outputs, "Attribute":$cast,
+#  ODS-NEXT:       CArg<"ArrayRef<NamedAttribute>", "{}">:$attributes),
+
 #       ODS:    $_state.addOperands(inputs);
 #  ODS-NEXT:    $_state.addOperands(outputs);
 #  ODS-NEXT:    $_state.addTypes(resultTensorTypes);
+#  ODS-NEXT:    $_state.addAttribute("cast", cast)
 #  ODS-NEXT:    $_state.addAttributes(attributes);
 #  ODS-NEXT:    $_state.addAttribute(
 #  ODS-NEXT:      "operand_segment_sizes",
 #  ODS-NEXT:      $_builder.getI32VectorAttr({
 #  ODS-NEXT:        static_cast<int32_t>(inputs.size()),
 #  ODS-NEXT:        static_cast<int32_t>(outputs.size())}));
 #  ODS-NEXT:    createAndFillStructuredOpRegion<Test1Op>(
 #  ODS-NEXT:      $_builder,
 #  ODS-NEXT:      $_state,
 #  ODS-NEXT:      TypeRange(inputs),
 #  ODS-NEXT:      TypeRange(outputs)
 
 # IMPL-LABEL:  void Test1Op::regionBuilder(ImplicitLocOpBuilder &b,
 #  IMPL-NEXT:    Block &block, ArrayRef<NamedAttribute> attrs)
+#       IMPL:  TypeFn castVal = TypeFn::cast;
+#  IMPL-NEXT:  auto castIter = llvm::find_if(attrs, [&](const NamedAttribute &attr) {
+#  IMPL-NEXT:                                return attr.getName() == "cast"; });
+#  IMPL-NEXT:  if (castIter != attrs.end()) {
+#  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.typefn__cast(block.getArgument(0).getType(), [[VAL0]]);
+#   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.typefn__cast_unsigned(block.getArgument(0).getType(), [[VAL2]]);
+#   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]]);
 
 
 # @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.

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

def matmul_mono(
   domain(D.m, D.n, D.k)
   C[D.m, D.n] += A[D.m, D.k] * B[D.k, D.n]
 
 
 @linalg_structured_op
 def matmul_poly(
     A=TensorDef(T1, S.M, S.K),
     B=TensorDef(T2, S.K, S.N),
-    C=TensorDef(U, S.M, S.N, output=True)):
+    C=TensorDef(U, S.M, S.N, output=True),
+    cast=TypeFnAttrDef(default=TypeFn.cast)):
   domain(D.m, D.n, D.k)
-  C[D.m, D.n] += TypeFn.cast(U, A[D.m, D.k]) * TypeFn.cast(U, B[D.k, D.n])
-
-
-@linalg_structured_op
-def matmul_unsigned_poly(
-    A=TensorDef(T1, S.M, S.K),
-    B=TensorDef(T2, S.K, S.N),
-    C=TensorDef(U, S.M, S.N, output=True)):
-  domain(D.m, D.n, D.k)
-  C[D.m, D.n] += TypeFn.cast_unsigned(U, A[D.m, D.k]) * TypeFn.cast_unsigned(
-      U, B[D.k, D.n])
+  C[D.m, D.n] += cast(U, A[D.m, D.k]) * cast(U, B[D.k, D.n])
 
 
 with Context() as ctx, Location.unknown():
   module = Module.create()
   f16 = F16Type.get()
   f32 = F32Type.get()
   f64 = F64Type.get()
   i8 = IntegerType.get_signless(8)
with InsertionPoint(module.body):
 
     # CHECK-LABEL: @test_i8i8i32_matmul_unsigned
     # CHECK:   = arith.extui
     # CHECK:   = arith.extui
     @builtin.FuncOp.from_py_func(
         RankedTensorType.get((4, 16), i8), RankedTensorType.get((16, 8), i8),
         RankedTensorType.get((4, 8), i32))
     def test_i8i8i32_matmul_unsigned(lhs, rhs, init_result):
-      return matmul_unsigned_poly(lhs, rhs, outs=[init_result])
+      return matmul_poly(
+          lhs, rhs, outs=[init_result], cast=TypeFn.cast_unsigned)
 
     # CHECK-LABEL: @test_i8i16i32_matmul
     # CHECK:      ^{{.*}}(%[[A_ARG:.+]]: i8, %[[B_ARG:.+]]: i16, %[[C_ARG:.+]]: i32)
     # CHECK-NEXT:   %[[A_CAST:.+]] = arith.extsi %[[A_ARG]] : i8 to i32
     # CHECK-NEXT:   %[[B_CAST:.+]] = arith.extsi %[[B_ARG]] : i16 to i32
     # CHECK-NEXT:   %[[MUL:.+]] = arith.muli %[[A_CAST]], %[[B_CAST]] : i32
     # CHECK-NEXT:   %[[ADD:.+]] = arith.addi %[[C_ARG]], %[[MUL]] : i32
     # CHECK-NEXT:   linalg.yield %[[ADD]] : i32
with InsertionPoint(module.body):
 
     # CHECK-LABEL: @test_i8i8f32_matmul_unsigned
     # CHECK:   = arith.uitofp
     # CHECK:   = arith.uitofp
     @builtin.FuncOp.from_py_func(
         RankedTensorType.get((4, 16), i8), RankedTensorType.get((16, 8), i8),
         RankedTensorType.get((4, 8), f32))
     def test_i8i8f32_matmul_unsigned(lhs, rhs, init_result):
-      return matmul_unsigned_poly(lhs, rhs, outs=[init_result])
+      return matmul_poly(
+          lhs, rhs, outs=[init_result], cast=TypeFn.cast_unsigned)
 
     # CHECK-LABEL: @test_f16f16f32_matmul
     # CHECK:      ^{{.*}}(%[[A_ARG:.+]]: f16, %[[B_ARG:.+]]: f16, %[[C_ARG:.+]]: f32)
     # CHECK-NEXT:   %[[A_CAST:.+]] = arith.extf %[[A_ARG]] : f16 to f32
     # CHECK-NEXT:   %[[B_CAST:.+]] = arith.extf %[[B_ARG]] : f16 to f32
     # CHECK-NEXT:   %[[MUL:.+]] = arith.mulf %[[A_CAST]], %[[B_CAST]] : f32
     # CHECK-NEXT:   %[[ADD:.+]] = arith.addf %[[C_ARG]], %[[MUL]] : f32
     # CHECK-NEXT:   linalg.yield %[[ADD]] : f32

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

 T2 = TV.T2
 
 
 @linalg_structured_op
 def pooling_max_poly(
     I=TensorDef(T1, S.N, S.H, S.W, S.C),
     K=TensorDef(T2, S.KH, S.KW, index_dims=[D.kh, D.kw]),
     O=TensorDef(U, S.N, S.OH, S.OW, S.C, output=True),
+    cast=TypeFnAttrDef(default=TypeFn.cast),
     strides=IndexAttrDef(S.SH, S.SW, default=[1, 1]),
     dilations=IndexAttrDef(S.DH, S.DW, default=[1, 1])):
   domain(D.n, D.oh, D.ow, D.kh, D.kw, D.c)
   O[D.n, D.oh, D.ow, D.c] = ReduceFn.max[D.kh, D.kw](
-      TypeFn.cast(
-          U, I[D.n, D.oh * S.SH + D.kh * S.DH, D.ow * S.SW + D.kw * S.DW, D.c]))
+      cast(U, I[D.n, D.oh * S.SH + D.kh * S.DH, D.ow * S.SW + D.kw * S.DW,
+                D.c]))
 
 
 @linalg_structured_op
 def pooling_max_unsigned_poly(
     I=TensorDef(T1, S.N, S.H, S.W, S.C),
     K=TensorDef(T2, S.KH, S.KW, index_dims=[D.kh, D.kw]),
     O=TensorDef(U, S.N, S.OH, S.OW, S.C, output=True),
+    cast=TypeFnAttrDef(default=TypeFn.cast_unsigned),
     strides=IndexAttrDef(S.SH, S.SW, default=[1, 1]),
     dilations=IndexAttrDef(S.DH, S.DW, default=[1, 1])):
   domain(D.n, D.oh, D.ow, D.kh, D.kw, D.c)
   O[D.n, D.oh, D.ow, D.c] = ReduceFn.max_unsigned[D.kh, D.kw](
-      TypeFn.cast_unsigned(
-          U, I[D.n, D.oh * S.SH + D.kh * S.DH, D.ow * S.SW + D.kw * S.DW, D.c]))
+      cast(U, I[D.n, D.oh * S.SH + D.kh * S.DH, D.ow * S.SW + D.kw * S.DW,
+                D.c]))
 
 
 @linalg_structured_op
 def pooling_min_poly(
     I=TensorDef(T1, S.N, S.H, S.W, S.C),
     K=TensorDef(T2, S.KH, S.KW, index_dims=[D.kh, D.kw]),
     O=TensorDef(U, S.N, S.OH, S.OW, S.C, output=True),
+    cast=TypeFnAttrDef(default=TypeFn.cast),
     strides=IndexAttrDef(S.SH, S.SW, default=[1, 1]),
     dilations=IndexAttrDef(S.DH, S.DW, default=[1, 1])):
   domain(D.n, D.oh, D.ow, D.kh, D.kw, D.c)
   O[D.n, D.oh, D.ow, D.c] = ReduceFn.min[D.kh, D.kw](
-      TypeFn.cast(
-          U, I[D.n, D.oh * S.SH + D.kh * S.DH, D.ow * S.SW + D.kw * S.DW, D.c]))
+      cast(U, I[D.n, D.oh * S.SH + D.kh * S.DH, D.ow * S.SW + D.kw * S.DW,
+                D.c]))
 
 
 @linalg_structured_op
 def pooling_min_unsigned_poly(
     I=TensorDef(T1, S.N, S.H, S.W, S.C),
     K=TensorDef(T2, S.KH, S.KW, index_dims=[D.kh, D.kw]),
     O=TensorDef(U, S.N, S.OH, S.OW, S.C, output=True),
+    cast=TypeFnAttrDef(default=TypeFn.cast_unsigned),
     strides=IndexAttrDef(S.SH, S.SW, default=[1, 1]),
     dilations=IndexAttrDef(S.DH, S.DW, default=[1, 1])):
   domain(D.n, D.oh, D.ow, D.kh, D.kw, D.c)
   O[D.n, D.oh, D.ow, D.c] = ReduceFn.min_unsigned[D.kh, D.kw](
-      TypeFn.cast_unsigned(
-          U, I[D.n, D.oh * S.SH + D.kh * S.DH, D.ow * S.SW + D.kw * S.DW, D.c]))
+      cast(U, I[D.n, D.oh * S.SH + D.kh * S.DH, D.ow * S.SW + D.kw * S.DW,
+                D.c]))
 
 
 with Context() as ctx, Location.unknown():
   module = Module.create()
   f32 = F32Type.get()
   i32 = IntegerType.get_signless(32)
   with InsertionPoint(module.body):
 
with InsertionPoint(module.body):
     @builtin.FuncOp.from_py_func(
         RankedTensorType.get((1, 4, 16, 1), f32),
         RankedTensorType.get((2, 2), f32),
         RankedTensorType.get((1, 2, 4, 1), f32))
     def test_f32f32_min_pooling(input, shape, init_result):
       return pooling_min_poly(
           input, shape, outs=[init_result], strides=[2, 4], dilations=[1, 2])
 
+    # CHECK-LABEL: @test_f32i32_min_pooling_cast_attr
+    # CHECK:   = arith.fptoui
+    # CHECK:   = arith.minsi
+    @builtin.FuncOp.from_py_func(
+        RankedTensorType.get((1, 4, 16, 1), f32),
+        RankedTensorType.get((2, 2), f32),
+        RankedTensorType.get((1, 2, 4, 1), i32))
+    def test_f32i32_min_pooling_cast_attr(input, shape, init_result):
+      return pooling_min_poly(
+          input,
+          shape,
+          outs=[init_result],
+          cast=TypeFn.cast_unsigned,
+          strides=[2, 4],
+          dilations=[1, 2])
 
 print(module)

mlir/test/python/dialects/linalg/ops.py

 # RUN: %PYTHON %s | FileCheck %s
 
 from mlir.ir import *
 from mlir.dialects import builtin
 from mlir.dialects import linalg
 from mlir.dialects import std
 from mlir.dialects import arith
 
+from mlir.dialects.linalg.opdsl.lang import *
 
 def run(f):
   print("\nTEST:", f.__name__)
   f()
   return f
 
 
 # CHECK-LABEL: TEST: testInitTensor
with InsertionPoint(module.body):
 
       @builtin.FuncOp.from_py_func(
           RankedTensorType.get((4, 16), f32), RankedTensorType.get((16, 8),
                                                                    f32))
       def named_form(lhs, rhs):
         init_result = linalg.InitTensorOp([4, 8], f32)
         # First check the named form with custom format
         #      CHECK: linalg.matmul
+        #      CHECK: cast = #linalg<"type_fn cast_unsigned">
         #  CHECK-NOT: linalg.memoized_indexing_maps
         # CHECK-SAME:    ins(%{{.*}} : tensor<4x16xf32>, tensor<16x8xf32>)
         # CHECK-SAME:   outs(%{{.*}} : tensor<4x8xf32>)
         # CHECK-SAME:   -> tensor<4x8xf32>
         # CHECK-NEXT: return
-        return linalg.matmul(lhs, rhs, outs=[init_result.result])
+        return linalg.matmul(
+            lhs, rhs, outs=[init_result.result], cast=TypeFn.cast_unsigned)
 
   print(module)
 
 
 # CHECK-LABEL: TEST: testNamedStructuredOpGenericForm
 @run
 def testNamedStructuredOpGenericForm():
   with Context() as ctx, Location.unknown():

mlir/test/python/integration/dialects/linalg/opsrun.py

 # RUN: %PYTHON %s 2>&1 | FileCheck %s
 
 import ctypes
 import sys
 from mlir.ir import *
 from mlir.dialects import builtin
 from mlir.dialects import linalg
 from mlir.dialects import std
 from mlir.passmanager import *
 from mlir.execution_engine import *
 
+from mlir.dialects.linalg.opdsl.lang import *
 
 # Log everything to stderr and flush so that we have a unified stream to match
 # errors/info emitted by MLIR to stderr.
 def log(*args):
   print(*args, file=sys.stderr)
   sys.stderr.flush()
 
 
 matmul_boiler = """
 func @main() -> f32 attributes {llvm.emit_c_interface} {
   %v0 = arith.constant 0.0 : f32
-  %v1 = arith.constant 1.0 : f32
+  %v1 = arith.constant -1 : i8
   %v2 = arith.constant 2.0 : f32
 
-  %A = memref.alloc() : memref<4x16xf32>
+  %A = memref.alloc() : memref<4x16xi8>
   %B = memref.alloc() : memref<16x8xf32>
-  %C = memref.alloc() : memref<4x8xf32>
-  linalg.fill(%v1, %A) : f32, memref<4x16xf32>
+  %C0 = memref.alloc() : memref<4x8xf32>
+  %C1 = memref.alloc() : memref<4x8xf32>
+  linalg.fill(%v1, %A) : i8, memref<4x16xi8>
   linalg.fill(%v2, %B) : f32, memref<16x8xf32>
-  linalg.fill(%v0, %C) : f32, memref<4x8xf32>
+  linalg.fill(%v0, %C0) : f32, memref<4x8xf32>
+  linalg.fill(%v0, %C1) : f32, memref<4x8xf32>
 
-  call @matmul_on_buffers(%A, %B, %C) :
-    (memref<4x16xf32>, memref<16x8xf32>, memref<4x8xf32>) -> ()
+  call @matmul_signed_on_buffers(%A, %B, %C0) :
+    (memref<4x16xi8>, memref<16x8xf32>, memref<4x8xf32>) -> ()
+  call @matmul_unsigned_on_buffers(%A, %B, %C1) :
+    (memref<4x16xi8>, memref<16x8xf32>, memref<4x8xf32>) -> ()
 
   %c0 = arith.constant 0 : index
-  %0 = memref.load %C[%c0, %c0] : memref<4x8xf32>
+  %res0 = memref.load %C0[%c0, %c0] : memref<4x8xf32>
+  %res1 = memref.load %C1[%c0, %c0] : memref<4x8xf32>
+
+  %0 = arith.addf %res0, %res1 : f32
 
   // TODO: FFI-based solution to allow testing and printing with python code.
   return %0 : f32
 }
 """
 
 fill_boiler = """
 func @main() -> i32 attributes {llvm.emit_c_interface} {
def transform(module, boilerplate):
   pm.run(mod)
   return mod
 
 
 def test_matmul_builtin():
   with Context() as ctx, Location.unknown():
     module = Module.create()
     f32 = F32Type.get()
+    i8 = IntegerType.get_signless(8)
     with InsertionPoint(module.body):
 
       @builtin.FuncOp.from_py_func(
-          MemRefType.get((4, 16), f32), MemRefType.get((16, 8), f32),
+          MemRefType.get((4, 16), i8), MemRefType.get((16, 8), f32),
           MemRefType.get((4, 8), f32))
-      def matmul_on_buffers(lhs, rhs, out):
+      def matmul_signed_on_buffers(lhs, rhs, out):
         linalg.matmul(lhs, rhs, outs=[out])
 
+      @builtin.FuncOp.from_py_func(
+          MemRefType.get((4, 16), i8), MemRefType.get((16, 8), f32),
+          MemRefType.get((4, 8), f32))
+      def matmul_unsigned_on_buffers(lhs, rhs, out):
+        linalg.matmul(lhs, rhs, outs=[out], cast=TypeFn.cast_unsigned)
+
     execution_engine = ExecutionEngine(transform(module, matmul_boiler))
 
     # TODO: FFI-based solution to allow testing and printing with python code.
     # Prepare arguments: one result f32.
     # Arguments must be passed as pointers.
     c_float_p = ctypes.c_float * 1
     res = c_float_p(-1.)
     execution_engine.invoke("main", res)
 
     log("RESULT: ", res[0])
-    # CHECK: RESULT: 32.0
+    # matmul_signed_on_buffers: -1 * 2.0 * 16 = -32
+    # matmul_unsigned_on_buffers: (2^8-1) * 2.0 * 16 = 8160
+    # CHECK: RESULT: 8128
 
 
 test_matmul_builtin()
 
 
 def test_matmul_generic():
   with Context() as ctx, Location.unknown():
     module = Module.create()
     f32 = F32Type.get()
+    i8 = IntegerType.get_signless(8)
     with InsertionPoint(module.body):
 
       @builtin.FuncOp.from_py_func(
-          MemRefType.get((4, 16), f32), MemRefType.get((16, 8), f32),
+          MemRefType.get((4, 16), i8), MemRefType.get((16, 8), f32),
           MemRefType.get((4, 8), f32))
-      def matmul_on_buffers(lhs, rhs, out):
+      def matmul_signed_on_buffers(lhs, rhs, out):
         linalg.matmul(lhs, rhs, outs=[out], emit_generic=True)
 
+      @builtin.FuncOp.from_py_func(
+          MemRefType.get((4, 16), i8), MemRefType.get((16, 8), f32),
+          MemRefType.get((4, 8), f32))
+      def matmul_unsigned_on_buffers(lhs, rhs, out):
+        linalg.matmul(
+            lhs, rhs, outs=[out], cast=TypeFn.cast_unsigned, emit_generic=True)
+
     execution_engine = ExecutionEngine(transform(module, matmul_boiler))
 
     # TODO: FFI-based solution to allow testing and printing with python code.
     # Prepare arguments: one result f32.
     # Arguments must be passed as pointers.
     c_float_p = ctypes.c_float * 1
     res = c_float_p(-1.)
     execution_engine.invoke("main", res)
 
     log("RESULT: ", res[0])
-    # CHECK: RESULT: 32.0
+    # matmul_signed_on_buffers = -1 * 2.0 * 16 = -32
+    # matmul_unsigned_on_buffers = (2^8-1) * 2.0 * 16 = 8160
+    # CHECK: RESULT: 8128
 
 
 test_matmul_generic()
 
 
 def test_fill_builtin():
   with Context() as ctx, Location.unknown():
     module = Module.create()

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

 };
 
 struct SerializedAffineMap {
   AffineMapAttr affineMapAttr;
 
   AffineMap affineMap() { return affineMapAttr.getValue(); }
 };
 
-enum class LinalgOperandDefUsage { Input, Output, IndexAttr };
+enum class LinalgOperandDefUsage { Input, Output, IndexAttr, TypeFnAttr };
 
 struct LinalgOperandDef {
   std::string name;
   LinalgOperandDefUsage usage;
   Optional<std::string> typeVar;
   Optional<SerializedAffineMap> shapeMap;
   Optional<SerializedAffineMap> indexAttrMap;
   Optional<SmallVector<int64_t>> defaultVals;
+  Optional<std::string> defaultFn;
 };
 
 enum class LinalgIteratorTypeDef {
   parallel,
   reduction,
 };
 
 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;
 };
 
 struct ScalarTypeFn {
-  std::string fnName;
   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;

 ///   - `shape_map`: An optional AffineMap from all op symbols to the shape of
 ///     the argument. Only tensor arguments have a `shape_map`. Each shape must
 ///     be normalized over the same list of symbols and have no dimension
 ///     inputs.
 ///   - `index_attr_map`: An optional AffineMap from all op symbols to the
 ///     index attribute symbols. During op creation these symbols are replaced
 ///     by the corresponding `name` index attribue values. Only index attribute
 ///     arguments have an `index_attr_map`.
 ///   - `default_vals`: An optional default initialization for index attribute
 ///     arguments.
+///   - `default_fn`: An optional default initialization for function attribute
+///     arguments.
 template <>
 struct MappingTraits<LinalgOperandDef> {
   static void mapping(IO &io, LinalgOperandDef &info) {
     io.mapRequired("name", info.name);
     io.mapRequired("usage", info.usage);
     io.mapOptional("type_var", info.typeVar);
     io.mapOptional("shape_map", info.shapeMap);
     io.mapOptional("index_attr_map", info.indexAttrMap);
     io.mapOptional("default_vals", info.defaultVals);
+    io.mapOptional("default_fn", info.defaultFn);
   }
 };
 
 /// Usage enum for a named argument.
 template <>
 struct ScalarEnumerationTraits<LinalgOperandDefUsage> {
   static void enumeration(IO &io, LinalgOperandDefUsage &value) {
     io.enumCase(value, "Input", LinalgOperandDefUsage::Input);
     io.enumCase(value, "Output", LinalgOperandDefUsage::Output);
     io.enumCase(value, "IndexAttr", LinalgOperandDefUsage::IndexAttr);
+    io.enumCase(value, "TypeFnAttr", LinalgOperandDefUsage::TypeFnAttr);
   }
 };
 
 /// Iterator type enum.
 template <>
 struct ScalarEnumerationTraits<LinalgIteratorTypeDef> {
   static void enumeration(IO &io, LinalgIteratorTypeDef &value) {
     io.enumCase(value, "parallel", LinalgIteratorTypeDef::parallel);
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("fn_name", info.fnName);
     io.mapRequired("type_var", info.typeVar);
     io.mapRequired("operands", info.operands);
+    io.mapOptional("fn_name", info.fnName);
+    io.mapOptional("attr_name", info.attrName);
   }
 };
 
 /// 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,
static const char structuredOpBuilderFormat[] = R"FMT(
   , OpBuilder<
   (ins "TypeRange":$resultTensorTypes, "ValueRange":$inputs,
        "ValueRange":$outputs, {1},
        CArg<"ArrayRef<NamedAttribute>", "{{}">:$attributes),
   [{{
     $_state.addOperands(inputs);
     $_state.addOperands(outputs);
     $_state.addTypes(resultTensorTypes);
+    {2}
+    $_state.addAttributes(attributes);
     $_state.addAttribute(
       "operand_segment_sizes",
       $_builder.getI32VectorAttr({{
         static_cast<int32_t>(inputs.size()),
         static_cast<int32_t>(outputs.size())}));
     createAndFillStructuredOpRegion<{0}>(
       $_builder,
       $_state,
       TypeRange(inputs),
       TypeRange(outputs));
-    {2}
-    $_state.addAttributes(attributes);
   }]>
 )FMT";
 
 // The iterator_types() method for structured ops. Parameters:
 // {0}: Class name
 // {1}: Comma interleaved iterator type names.
 static const char structuredOpIteratorTypesFormat[] =
     R"FMT(
)FMT";
     StringRef summary, description;
     std::tie(summary, description) =
         StringRef(*opConfig.metadata->doc).trim().split('\n');
     doc = llvm::formatv(structuredOpDocFmt, summary.trim(), description.trim());
   }
 
   interfaceNameList = interleaveToString(opConfig.metadata->implements, ", ");
 
-  // Assemble the attribute specific logic required for the op definition.
   if (llvm::any_of(opConfig.structuredOp->args, [](LinalgOperandDef &arg) {
-        return arg.usage == LinalgOperandDefUsage::IndexAttr;
+        return arg.usage == LinalgOperandDefUsage::IndexAttr ||
+               arg.usage == LinalgOperandDefUsage::TypeFnAttr;
       })) {
     SmallVector<std::string> attrDefs;
     SmallVector<std::string> attrParams;
     SmallVector<std::string> attrStmts;
     for (LinalgOperandDef &arg : opConfig.structuredOp->args) {
-      if (arg.usage != LinalgOperandDefUsage::IndexAttr)
-        continue;
+      static const char paramFmt[] = "\"Attribute\":${0}";
+      static const char stmtFmt[] = "$_state.addAttribute(\"{0}\", {0});";
+      // Add the type conversion attributes to the op definition and builders.
+      if (arg.usage == LinalgOperandDefUsage::TypeFnAttr) {
+        assert(arg.defaultFn.hasValue());
+        static const char typeFmt[] = "TypeFn::{0}";
+        static const char defFmt[] = "DefaultValuedAttr<{0}, \"{1}\">:${2}";
+        attrDefs.push_back(llvm::formatv(defFmt, "TypeFnAttr",
+                                         llvm::formatv(typeFmt, arg.defaultFn),
+                                         arg.name));
+        attrParams.push_back(llvm::formatv(paramFmt, arg.name));
+        attrStmts.push_back(llvm::formatv(stmtFmt, arg.name));
+      }
+      // Add the index attributes to the op definition and builders.
+      if (arg.usage == LinalgOperandDefUsage::IndexAttr) {
         assert(arg.indexAttrMap.hasValue());
         assert(arg.defaultVals.hasValue());
         size_t size = arg.indexAttrMap->affineMap().getNumResults();
         assert(arg.defaultVals.getValue().size() == size);
         static const char typeFmt[] = "RankedI64ElementsAttr<[{0}]>";
-      static const char defFmt[] = "DefaultValuedAttr<{0}, \"{1}\">:${2}";
-      static const char paramFmt[] = "\"Attribute\":${0}";
-      static const char stmtFmt[] = "$_state.addAttribute(\"{0}\", {0});";
+        static const char defFmt[] = "DefaultValuedAttr<{0}, \"{ {1} }\">:${2}";
         std::string defaultVals;
         llvm::raw_string_ostream ss(defaultVals);
-      ss << "{ ";
         llvm::interleave(
             arg.defaultVals.getValue(), ss,
             [&](int64_t val) { ss << "static_cast<int64_t>(" << val << ")"; },
             ", ");
-      ss << " }";
         attrDefs.push_back(llvm::formatv(defFmt, llvm::formatv(typeFmt, size),
                                          ss.str(), arg.name));
         attrParams.push_back(llvm::formatv(paramFmt, arg.name));
         attrStmts.push_back(llvm::formatv(stmtFmt, arg.name));
       }
-    attrList = ",\n" + llvm::join(attrDefs, ",\n");
+    }
+    if (llvm::any_of(opConfig.structuredOp->args, [](LinalgOperandDef &arg) {
+          return arg.usage == LinalgOperandDefUsage::IndexAttr;
+        })) {
       attrMethods = R"(
         bool hasDynamicIndexingMaps();
         LogicalResult verifyIndexingMapRequiredAttributes();
       )";
+    }
+    attrList = ",\n" + llvm::join(attrDefs, ",\n");
     attrBuilder = llvm::formatv(
         structuredOpBuilderFormat, opConfig.metadata->cppClassName,
         llvm::join(attrParams, ", "), llvm::join(attrStmts, "\n"));
   }
 
   os << llvm::formatv(structuredOpOdsHeaderFormat,
                       opConfig.metadata->cppClassName, opConfig.metadata->name,
                       interfaceNameList, doc, attrList, attrBuilder,
generateNamedGenericOpDefns(LinalgOpConfig &opConfig,
 
   // Implementation banner.
   std::string bannerComment = llvm::formatv("Implementation of {0}", className);
   os << llvm::formatv(bannerFormat, bannerComment);
 
   // Compute the number of scalar and tensor arguments.
   int64_t numOfArgs =
       llvm::count_if(opConfig.structuredOp->args, [](LinalgOperandDef &arg) {
-        return arg.usage != LinalgOperandDefUsage::IndexAttr;
+        return arg.usage == LinalgOperandDefUsage::Input ||
+               arg.usage == LinalgOperandDefUsage::Output;
       });
 
   // An operation that accesses only scalars and scalar/rank zero tensors is
   // rank polymorhpic. We implement rank polymorphism by generating different
   // indexing maps and iterators that match the rank of the first output tensor.
   // An operation is rank polymorphic if the iteration domain has rank zero.
   bool isRankPolymorphic = opConfig.structuredOp->iteratorTypes.empty();
 
os << llvm::formatv(structuredOpVerifyIndexingMapRequiredAttributes,
                         className, llvm::join(attrVerifications, "\n"));
   }
 
   // regionBuilder()
   {
     // Generates a regionBuilder method. Parameters.
     // {0}: Class name
     // {1}: Number of args
-    // {2}: Statements
+    // {2}: Attributes
+    // {3}: Statements
     static const char structuredOpRegionBuilderFormat[] = R"FMT(
 void {0}::regionBuilder(ImplicitLocOpBuilder &b,
                         Block &block, ArrayRef<NamedAttribute> attrs) {{
   assert({1} > 0 && block.getNumArguments() == {1} &&
          "{0} regionBuilder expects {1} (>=0) args");
   RegionBuilderHelper helper(block.getArgument(0).getContext(), block);
   SmallVector<Value> yields;
   {2}
+  {3}
   helper.yieldOutputs(yields);
 }
 )FMT";
     auto &args = opConfig.structuredOp->args;
     auto &assignments = opConfig.structuredOp->assignments;
     size_t generatedAssignmentCount = 0;
     int localCounter = 0;
+    SmallVector<std::string> attrs;
     SmallVector<std::string> stmts;
     for (LinalgOperandDef &arg : args) {
+      if (arg.usage != LinalgOperandDefUsage::TypeFnAttr)
+        continue;
+      // Obtain the type function attribute values. Parameters.
+      // {0}: attribute name
+      // {1}: default type function name
+      static const char attrDef[] = R"FMT(
+TypeFn {0}Val = TypeFn::{1};
+auto {0}Iter = llvm::find_if(attrs, [&](const NamedAttribute &attr) {{
+                              return attr.getName() == "{0}"; });
+if ({0}Iter != attrs.end()) {{
+  if (auto attr = {0}Iter->getValue().dyn_cast<TypeFnAttr>())
+    {0}Val = attr.getValue();
+}
+)FMT";
+      attrs.push_back(llvm::formatv(attrDef, arg.name, arg.defaultFn));
+    }
+    for (LinalgOperandDef &arg : args) {
       if (arg.usage != LinalgOperandDefUsage::Output)
         continue;
 
       // Find the assignment that correlates with the argument.
       ScalarAssign *assignment = findAssignment(arg.name, assignments);
       if (!assignment)
         return emitError(genContext.getLoc())
                << "no assignment found for output argument " << arg.name;
for (LinalgOperandDef &arg : args) {
               findTypeValue(expression.typeFn->typeVar, args);
           if (!typeCppValue) {
             emitError(genContext.getLoc())
                 << "type variable " << expression.typeFn->typeVar
                 << ", 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) {
+            if (llvm::none_of(args, [&](LinalgOperandDef &arg) {
+                  return arg.usage == LinalgOperandDefUsage::TypeFnAttr &&
+                         arg.name == expression.typeFn->attrName.getValue();
+                })) {
+              emitError(genContext.getLoc())
+                  << "missing type function attribute "
+                  << expression.typeFn->attrName.getValue();
+            }
+            typeFunc = llvm::formatv("{0}Val", *expression.typeFn->attrName);
+          }
           std::string cppIdent = llvm::formatv("value{0}", ++localCounter);
-          stmts.push_back(
-              llvm::formatv("Value {0} = helper.typefn__{1}({2}, {3});",
-                            cppIdent, expression.typeFn->fnName,
-                            typeCppValue.getValue(), *operandCppValue));
+          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";
         return None;
       };
       Optional<std::string> cppValue = generateExpression(assignment->value);
       if (!cppValue)
         return failure();
       stmts.push_back(llvm::formatv("yields.push_back({0});", cppValue));
     }
 
     if (generatedAssignmentCount != assignments.size())
       return emitError(genContext.getLoc())
              << "mismatched number of assignments vs output arguments";
 
     os << llvm::formatv(structuredOpRegionBuilderFormat, className, numOfArgs,
+                        interleaveToString(attrs, "\n  "),
                         interleaveToString(stmts, "\n  "));
   }
 
   // Parser and printer.
   os << llvm::formatv(structuredOpParserFormat, className);
 
   // Canonicalizers and folders.
   os << llvm::formatv(structuredOpFoldersFormat, className);

utils/bazel/llvm-project-overlay/mlir/BUILD.bazel

tbl_outs = [
         ),
         (
             [
                 "-gen-dialect-defs",
                 "-dialect=linalg",
             ],
             "include/mlir/Dialect/Linalg/IR/LinalgOpsDialect.cpp.inc",
         ),
+        (
+            ["-gen-enum-decls"],
+            "include/mlir/Dialect/Linalg/IR/LinalgOpsEnums.h.inc",
+        ),
+        (
+            ["-gen-enum-defs"],
+            "include/mlir/Dialect/Linalg/IR/LinalgOpsEnums.cpp.inc",
+        ),
+        (
+            ["-gen-attrdef-decls"],
+            "include/mlir/Dialect/Linalg/IR/LinalgOpsAttrDefs.h.inc",
+        ),
+        (
+            ["-gen-attrdef-defs"],
+            "include/mlir/Dialect/Linalg/IR/LinalgOpsAttrDefs.cpp.inc",
+        ),
     ],
     tblgen = ":mlir-tblgen",
     td_file = "include/mlir/Dialect/Linalg/IR/LinalgOps.td",
     deps = [":LinalgOpsTdFiles"],
 )
 
 genlinalg(
     name = "LinalgNamedStructuredOpsYamlIncGen",