[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

why is this line removed?

530

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

138

typo: attribute

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

42

Verifies that ....

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

191

_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

doc: |-
     Numeric casting is performed on the operands to the inner multiply, promoting
     them to the same data type as the accumulator/output.
   implements:
   - LinalgContractionOpInterface
 structured_op: !LinalgStructuredOpConfig
   args:
   - !LinalgOperandDefConfig
     name: A
-    usage: Input
+    kind: input_tensor
     type_var: T1
     shape_map: affine_map<()[s0, s1, s2] -> (s0, s1)>
   - !LinalgOperandDefConfig
     name: B
-    usage: Input
+    kind: input_tensor
     type_var: T2
     shape_map: affine_map<()[s0, s1, s2] -> (s1, s2)>
   - !LinalgOperandDefConfig
     name: C
-    usage: Output
+    kind: output_tensor
     type_var: U
     shape_map: affine_map<()[s0, s1, s2] -> (s0, s2)>
+  - !LinalgOperandDefConfig
+    name: cast
+    kind: type_fn_attr
+    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
     them to the same data type as the accumulator/output.
   implements:
   - LinalgContractionOpInterface
 structured_op: !LinalgStructuredOpConfig
   args:
   - !LinalgOperandDefConfig
     name: A
-    usage: Input
+    kind: input_tensor
     type_var: T1
     shape_map: affine_map<()[s0, s1, s2] -> (s0, s1)>
   - !LinalgOperandDefConfig
     name: B
-    usage: Input
+    kind: input_tensor
     type_var: T2
     shape_map: affine_map<()[s0, s1, s2] -> (s1, s2)>
   - !LinalgOperandDefConfig
     name: C
-    usage: Output
+    kind: output_tensor
     type_var: U
     shape_map: affine_map<()[s0, s1, s2] -> (s0, s2)>
   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:
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
     them to the same data type as the accumulator/output. The quantized variant
     includes zero-point adjustments for the left and right operands of the
     matmul.
 structured_op: !LinalgStructuredOpConfig
   args:
   - !LinalgOperandDefConfig
     name: A
-    usage: Input
+    kind: input_tensor
     type_var: T1
     shape_map: affine_map<()[s0, s1, s2] -> (s0, s1)>
   - !LinalgOperandDefConfig
     name: B
-    usage: Input
+    kind: input_tensor
     type_var: T2
     shape_map: affine_map<()[s0, s1, s2] -> (s1, s2)>
   - !LinalgOperandDefConfig
     name: AZp
-    usage: Input
+    kind: scalar
     type_var: I32
   - !LinalgOperandDefConfig
     name: BZp
-    usage: Input
+    kind: scalar
     type_var: I32
   - !LinalgOperandDefConfig
     name: C
-    usage: Output
+    kind: output_tensor
     type_var: U
     shape_map: affine_map<()[s0, s1, s2] -> (s0, s2)>
   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] -> ()>
     - affine_map<(d0, d1, d2)[s0, s1, s2] -> ()>
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:
     * The right hand side is transposed, whence the 't' in 'mmt'.
     * The input and output tensors have a 4D shape instead of a 2D shape. They
       are interpreted as 2D matrices with one level of 2D tile subdivision,
       whence the 2+2=4 dimensions. The inner tile dimensions are identified with
       '0' suffixes below, for instance the LHS matrix shape (M, K, M0, K0) reads
       as: MxK tiles, each of shape M0xK0.
   implements:
   - LinalgContractionOpInterface
 structured_op: !LinalgStructuredOpConfig
   args:
   - !LinalgOperandDefConfig
     name: lhs
-    usage: Input
+    kind: input_tensor
     type_var: LhsType
     shape_map: affine_map<()[s0, s1, s2, s3, s4, s5] -> (s0, s1, s2, s3)>
   - !LinalgOperandDefConfig
     name: rhs
-    usage: Input
+    kind: input_tensor
     type_var: RhsType
     shape_map: affine_map<()[s0, s1, s2, s3, s4, s5] -> (s4, s1, s5, s3)>
   - !LinalgOperandDefConfig
     name: accum
-    usage: Output
+    kind: output_tensor
     type_var: AccumType
     shape_map: affine_map<()[s0, s1, s2, s3, s4, s5] -> (s0, s4, s2, s5)>
   indexing_maps: !LinalgIndexingMapsConfig
     static_indexing_maps:
     - affine_map<(d0, d1, d2, d3, d4, d5)[s0, s1, s2, s3, s4, s5] -> (d0, d2, d3,
       d5)>
     - affine_map<(d0, d1, d2, d3, d4, d5)[s0, s1, s2, s3, s4, s5] -> (d1, d2, d4,
       d5)>
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
     them to the same data type as the accumulator/output.
   implements:
   - LinalgContractionOpInterface
 structured_op: !LinalgStructuredOpConfig
   args:
   - !LinalgOperandDefConfig
     name: A
-    usage: Input
+    kind: input_tensor
     type_var: T1
     shape_map: affine_map<()[s0, s1, s2, s3] -> (s0, s1, s2)>
   - !LinalgOperandDefConfig
     name: B
-    usage: Input
+    kind: input_tensor
     type_var: T2
     shape_map: affine_map<()[s0, s1, s2, s3] -> (s0, s2, s3)>
   - !LinalgOperandDefConfig
     name: C
-    usage: Output
+    kind: output_tensor
     type_var: U
     shape_map: affine_map<()[s0, s1, s2, s3] -> (s0, s1, s3)>
   indexing_maps: !LinalgIndexingMapsConfig
     static_indexing_maps:
     - affine_map<(d0, d1, d2, d3)[s0, s1, s2, s3] -> (d0, d1, d3)>
     - affine_map<(d0, d1, d2, d3)[s0, s1, s2, s3] -> (d0, d3, d2)>
     - affine_map<(d0, d1, d2, d3)[s0, s1, s2, s3] -> (d0, d1, d2)>
   iterator_types:
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
     them to the same data type as the accumulator/output. The quantized variant
     includes zero-point adjustments for the left and right operands of the
     matmul.
 structured_op: !LinalgStructuredOpConfig
   args:
   - !LinalgOperandDefConfig
     name: A
-    usage: Input
+    kind: input_tensor
     type_var: T1
     shape_map: affine_map<()[s0, s1, s2, s3] -> (s0, s1, s2)>
   - !LinalgOperandDefConfig
     name: B
-    usage: Input
+    kind: input_tensor
     type_var: T2
     shape_map: affine_map<()[s0, s1, s2, s3] -> (s0, s2, s3)>
   - !LinalgOperandDefConfig
     name: AZp
-    usage: Input
+    kind: scalar
     type_var: I32
   - !LinalgOperandDefConfig
     name: BZp
-    usage: Input
+    kind: scalar
     type_var: I32
   - !LinalgOperandDefConfig
     name: C
-    usage: Output
+    kind: output_tensor
     type_var: U
     shape_map: affine_map<()[s0, s1, s2, s3] -> (s0, s1, s3)>
   indexing_maps: !LinalgIndexingMapsConfig
     static_indexing_maps:
     - affine_map<(d0, d1, d2, d3)[s0, s1, s2, s3] -> (d0, d1, d3)>
     - affine_map<(d0, d1, d2, d3)[s0, s1, s2, s3] -> (d0, d3, d2)>
     - affine_map<(d0, d1, d2, d3)[s0, s1, s2, s3] -> ()>
     - affine_map<(d0, d1, d2, d3)[s0, s1, s2, s3] -> ()>
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
     them to the same data type as the accumulator/output.
   implements:
   - LinalgContractionOpInterface
 structured_op: !LinalgStructuredOpConfig
   args:
   - !LinalgOperandDefConfig
     name: A
-    usage: Input
+    kind: input_tensor
     type_var: T1
     shape_map: affine_map<()[s0, s1] -> (s0, s1)>
   - !LinalgOperandDefConfig
     name: y
-    usage: Input
+    kind: input_tensor
     type_var: T2
     shape_map: affine_map<()[s0, s1] -> (s1)>
   - !LinalgOperandDefConfig
     name: x
-    usage: Output
+    kind: output_tensor
     type_var: U
     shape_map: affine_map<()[s0, s1] -> (s0)>
   indexing_maps: !LinalgIndexingMapsConfig
     static_indexing_maps:
     - affine_map<(d0, d1)[s0, s1] -> (d0, d1)>
     - affine_map<(d0, d1)[s0, s1] -> (d1)>
     - affine_map<(d0, d1)[s0, s1] -> (d0)>
   iterator_types:
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
     them to the same data type as the accumulator/output.
   implements:
   - LinalgContractionOpInterface
 structured_op: !LinalgStructuredOpConfig
   args:
   - !LinalgOperandDefConfig
     name: y
-    usage: Input
+    kind: input_tensor
     type_var: T1
     shape_map: affine_map<()[s0, s1] -> (s0)>
   - !LinalgOperandDefConfig
     name: A
-    usage: Input
+    kind: input_tensor
     type_var: T2
     shape_map: affine_map<()[s0, s1] -> (s0, s1)>
   - !LinalgOperandDefConfig
     name: x
-    usage: Output
+    kind: output_tensor
     type_var: U
     shape_map: affine_map<()[s0, s1] -> (s1)>
   indexing_maps: !LinalgIndexingMapsConfig
     static_indexing_maps:
     - affine_map<(d0, d1)[s0, s1] -> (d1)>
     - affine_map<(d0, d1)[s0, s1] -> (d1, d0)>
     - affine_map<(d0, d1)[s0, s1] -> (d0)>
   iterator_types:
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
     them to the same data type as the accumulator/output.
   implements:
   - LinalgContractionOpInterface
 structured_op: !LinalgStructuredOpConfig
   args:
   - !LinalgOperandDefConfig
     name: A
-    usage: Input
+    kind: input_tensor
     type_var: T1
     shape_map: affine_map<()[s0, s1, s2] -> (s0, s1, s2)>
   - !LinalgOperandDefConfig
     name: B
-    usage: Input
+    kind: input_tensor
     type_var: T2
     shape_map: affine_map<()[s0, s1, s2] -> (s0, s2)>
   - !LinalgOperandDefConfig
     name: C
-    usage: Output
+    kind: output_tensor
     type_var: U
     shape_map: affine_map<()[s0, s1, s2] -> (s0, s1)>
   indexing_maps: !LinalgIndexingMapsConfig
     static_indexing_maps:
     - affine_map<(d0, d1, d2)[s0, s1, s2] -> (d0, d1, d2)>
     - affine_map<(d0, d1, d2)[s0, s1, s2] -> (d0, d2)>
     - affine_map<(d0, d1, d2)[s0, s1, s2] -> (d0, d1)>
   iterator_types:
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
     them to the same data type as the accumulator/output.
   implements:
   - LinalgContractionOpInterface
 structured_op: !LinalgStructuredOpConfig
   args:
   - !LinalgOperandDefConfig
     name: A
-    usage: Input
+    kind: input_tensor
     type_var: T1
     shape_map: affine_map<()[s0] -> (s0)>
   - !LinalgOperandDefConfig
     name: B
-    usage: Input
+    kind: input_tensor
     type_var: T2
     shape_map: affine_map<()[s0] -> (s0)>
   - !LinalgOperandDefConfig
     name: C
-    usage: Output
+    kind: output_tensor
     type_var: U
     shape_map: affine_map<()[s0] -> ()>
   indexing_maps: !LinalgIndexingMapsConfig
     static_indexing_maps:
     - affine_map<(d0)[s0] -> (d0)>
     - affine_map<(d0)[s0] -> (d0)>
     - affine_map<(d0)[s0] -> ()>
   iterator_types:
   - 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
+                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
     them to the same data type as the accumulator/output.
   implements:
   - LinalgConvolutionOpInterface
 structured_op: !LinalgStructuredOpConfig
   args:
   - !LinalgOperandDefConfig
     name: I
-    usage: Input
+    kind: input_tensor
     type_var: T1
     shape_map: affine_map<()[s0, s1] -> (s0 + s1)>
   - !LinalgOperandDefConfig
     name: K
-    usage: Input
+    kind: input_tensor
     type_var: T2
     shape_map: affine_map<()[s0, s1] -> (s1)>
   - !LinalgOperandDefConfig
     name: O
-    usage: Output
+    kind: output_tensor
     type_var: U
     shape_map: affine_map<()[s0, s1] -> (s0)>
   indexing_maps: !LinalgIndexingMapsConfig
     static_indexing_maps:
     - affine_map<(d0, d1)[s0, s1] -> (d0 + d1)>
     - affine_map<(d0, d1)[s0, s1] -> (d1)>
     - affine_map<(d0, d1)[s0, s1] -> (d0)>
   iterator_types:
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
     them to the same data type as the accumulator/output.
   implements:
   - LinalgConvolutionOpInterface
 structured_op: !LinalgStructuredOpConfig
   args:
   - !LinalgOperandDefConfig
     name: I
-    usage: Input
+    kind: input_tensor
     type_var: T1
     shape_map: affine_map<()[s0, s1, s2, s3] -> (s0 + s1, s2 + s3)>
   - !LinalgOperandDefConfig
     name: K
-    usage: Input
+    kind: input_tensor
     type_var: T2
     shape_map: affine_map<()[s0, s1, s2, s3] -> (s1, s3)>
   - !LinalgOperandDefConfig
     name: O
-    usage: Output
+    kind: output_tensor
     type_var: U
     shape_map: affine_map<()[s0, s1, s2, s3] -> (s0, s2)>
   indexing_maps: !LinalgIndexingMapsConfig
     static_indexing_maps:
     - affine_map<(d0, d1, d2, d3)[s0, s1, s2, s3] -> (d0 + d2, d1 + d3)>
     - affine_map<(d0, d1, d2, d3)[s0, s1, s2, s3] -> (d2, d3)>
     - affine_map<(d0, d1, d2, d3)[s0, s1, s2, s3] -> (d0, d1)>
   iterator_types:
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
     them to the same data type as the accumulator/output.
   implements:
   - LinalgConvolutionOpInterface
 structured_op: !LinalgStructuredOpConfig
   args:
   - !LinalgOperandDefConfig
     name: I
-    usage: Input
+    kind: input_tensor
     type_var: T1
     shape_map: affine_map<()[s0, s1, s2, s3, s4, s5] -> (s0 + s1, s2 + s3, s4 + s5)>
   - !LinalgOperandDefConfig
     name: K
-    usage: Input
+    kind: input_tensor
     type_var: T2
     shape_map: affine_map<()[s0, s1, s2, s3, s4, s5] -> (s1, s3, s5)>
   - !LinalgOperandDefConfig
     name: O
-    usage: Output
+    kind: output_tensor
     type_var: U
     shape_map: affine_map<()[s0, s1, s2, s3, s4, s5] -> (s0, s2, s4)>
   indexing_maps: !LinalgIndexingMapsConfig
     static_indexing_maps:
     - affine_map<(d0, d1, d2, d3, d4, d5)[s0, s1, s2, s3, s4, s5] -> (d0 + d3, d1
       + d4, d2 + d5)>
     - affine_map<(d0, d1, d2, d3, d4, d5)[s0, s1, s2, s3, s4, s5] -> (d3, d4, d5)>
     - affine_map<(d0, d1, d2, d3, d4, d5)[s0, s1, s2, s3, s4, s5] -> (d0, d1, d2)>
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
     them to the same data type as the accumulator/output.
   implements:
   - LinalgConvolutionOpInterface
 structured_op: !LinalgStructuredOpConfig
   args:
   - !LinalgOperandDefConfig
     name: I
-    usage: Input
+    kind: input_tensor
     type_var: T1
     shape_map: affine_map<()[s0, s1, s2, s3, s4, s5, s6] -> (s0, s1 * s2 + s3 * s4,
       s5)>
   - !LinalgOperandDefConfig
     name: K
-    usage: Input
+    kind: input_tensor
     type_var: T2
     shape_map: affine_map<()[s0, s1, s2, s3, s4, s5, s6] -> (s3, s5, s6)>
   - !LinalgOperandDefConfig
     name: O
-    usage: Output
+    kind: output_tensor
     type_var: U
     shape_map: affine_map<()[s0, s1, s2, s3, s4, s5, s6] -> (s0, s1, s6)>
   - !LinalgOperandDefConfig
     name: strides
-    usage: IndexAttr
+    kind: index_attr
     index_attr_map: affine_map<()[s0, s1, s2, s3, s4, s5, s6] -> (s2)>
     default_vals:
     - 1
   - !LinalgOperandDefConfig
     name: dilations
-    usage: IndexAttr
+    kind: index_attr
     index_attr_map: affine_map<()[s0, s1, s2, s3, s4, s5, s6] -> (s4)>
     default_vals:
     - 1
   indexing_maps: !LinalgIndexingMapsConfig
     static_indexing_maps:
     - affine_map<(d0, d1, d2, d3, d4)[s0, s1, s2, s3, s4, s5, s6] -> (d0, d1 * s2
       + d3 * s4, d4)>
     - affine_map<(d0, d1, d2, d3, d4)[s0, s1, s2, s3, s4, s5, s6] -> (d3, d4, d2)>
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:
       * Input: NHWC.
       * Kernel: HWCF.
 
     Numeric casting is performed on the operands to the inner multiply, promoting
     them to the same data type as the accumulator/output.
   implements:
   - LinalgConvolutionOpInterface
 structured_op: !LinalgStructuredOpConfig
   args:
   - !LinalgOperandDefConfig
     name: I
-    usage: Input
+    kind: input_tensor
     type_var: T1
     shape_map: affine_map<()[s0, s1, s2, s3, s4, s5, s6, s7, s8, s9, s10] -> (s0,
       s1 * s2 + s3 * s4, s5 * s6 + s7 * s8, s9)>
   - !LinalgOperandDefConfig
     name: K
-    usage: Input
+    kind: input_tensor
     type_var: T2
     shape_map: affine_map<()[s0, s1, s2, s3, s4, s5, s6, s7, s8, s9, s10] -> (s3,
       s7, s9, s10)>
   - !LinalgOperandDefConfig
     name: O
-    usage: Output
+    kind: output_tensor
     type_var: U
     shape_map: affine_map<()[s0, s1, s2, s3, s4, s5, s6, s7, s8, s9, s10] -> (s0,
       s1, s5, s10)>
   - !LinalgOperandDefConfig
     name: strides
-    usage: IndexAttr
+    kind: index_attr
     index_attr_map: affine_map<()[s0, s1, s2, s3, s4, s5, s6, s7, s8, s9, s10] ->
       (s2, s6)>
     default_vals:
     - 1
     - 1
   - !LinalgOperandDefConfig
     name: dilations
-    usage: IndexAttr
+    kind: index_attr
     index_attr_map: affine_map<()[s0, s1, s2, s3, s4, s5, s6, s7, s8, s9, s10] ->
       (s4, s8)>
     default_vals:
     - 1
     - 1
   indexing_maps: !LinalgIndexingMapsConfig
     static_indexing_maps:
     - affine_map<(d0, d1, d2, d3, d4, d5, d6)[s0, s1, s2, s3, s4, s5, s6, s7, s8,
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:
       * Input: NHWC.
       * Kernel: HWCF.
 
     Numeric casting is performed on the operands to the inner multiply, promoting
     them to the same data type as the accumulator/output. This includes the zero
     point offsets common to quantized operations.
   implements:
   - LinalgConvolutionOpInterface
 structured_op: !LinalgStructuredOpConfig
   args:
   - !LinalgOperandDefConfig
     name: I
-    usage: Input
+    kind: input_tensor
     type_var: T1
     shape_map: affine_map<()[s0, s1, s2, s3, s4, s5, s6, s7, s8, s9, s10] -> (s0,
       s1 * s2 + s3 * s4, s5 * s6 + s7 * s8, s9)>
   - !LinalgOperandDefConfig
     name: K
-    usage: Input
+    kind: input_tensor
     type_var: T2
     shape_map: affine_map<()[s0, s1, s2, s3, s4, s5, s6, s7, s8, s9, s10] -> (s3,
       s7, s9, s10)>
   - !LinalgOperandDefConfig
     name: IZp
-    usage: Input
+    kind: scalar
     type_var: I32
   - !LinalgOperandDefConfig
     name: KZp
-    usage: Input
+    kind: scalar
     type_var: I32
   - !LinalgOperandDefConfig
     name: O
-    usage: Output
+    kind: output_tensor
     type_var: U
     shape_map: affine_map<()[s0, s1, s2, s3, s4, s5, s6, s7, s8, s9, s10] -> (s0,
       s1, s5, s10)>
   - !LinalgOperandDefConfig
     name: strides
-    usage: IndexAttr
+    kind: index_attr
     index_attr_map: affine_map<()[s0, s1, s2, s3, s4, s5, s6, s7, s8, s9, s10] ->
       (s2, s6)>
     default_vals:
     - 1
     - 1
   - !LinalgOperandDefConfig
     name: dilations
-    usage: IndexAttr
+    kind: index_attr
     index_attr_map: affine_map<()[s0, s1, s2, s3, s4, s5, s6, s7, s8, s9, s10] ->
       (s4, s8)>
     default_vals:
     - 1
     - 1
   indexing_maps: !LinalgIndexingMapsConfig
     static_indexing_maps:
     - affine_map<(d0, d1, d2, d3, d4, d5, d6)[s0, s1, s2, s3, s4, s5, s6, s7, s8,
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:
       * Input: NCHW.
       * Kernel: FCHW.
 
     Numeric casting is performed on the operands to the inner multiply, promoting
     them to the same data type as the accumulator/output.
   implements:
   - LinalgConvolutionOpInterface
 structured_op: !LinalgStructuredOpConfig
   args:
   - !LinalgOperandDefConfig
     name: I
-    usage: Input
+    kind: input_tensor
     type_var: T1
     shape_map: affine_map<()[s0, s1, s2, s3, s4, s5, s6, s7, s8, s9, s10] -> (s0,
       s1, s2 * s3 + s4 * s5, s6 * s7 + s8 * s9)>
   - !LinalgOperandDefConfig
     name: K
-    usage: Input
+    kind: input_tensor
     type_var: T2
     shape_map: affine_map<()[s0, s1, s2, s3, s4, s5, s6, s7, s8, s9, s10] -> (s10,
       s1, s4, s8)>
   - !LinalgOperandDefConfig
     name: O
-    usage: Output
+    kind: output_tensor
     type_var: U
     shape_map: affine_map<()[s0, s1, s2, s3, s4, s5, s6, s7, s8, s9, s10] -> (s0,
       s10, s2, s6)>
   - !LinalgOperandDefConfig
     name: strides
-    usage: IndexAttr
+    kind: index_attr
     index_attr_map: affine_map<()[s0, s1, s2, s3, s4, s5, s6, s7, s8, s9, s10] ->
       (s3, s7)>
     default_vals:
     - 1
     - 1
   - !LinalgOperandDefConfig
     name: dilations
-    usage: IndexAttr
+    kind: index_attr
     index_attr_map: affine_map<()[s0, s1, s2, s3, s4, s5, s6, s7, s8, s9, s10] ->
       (s5, s9)>
     default_vals:
     - 1
     - 1
   indexing_maps: !LinalgIndexingMapsConfig
     static_indexing_maps:
     - affine_map<(d0, d1, d2, d3, d4, d5, d6)[s0, s1, s2, s3, s4, s5, s6, s7, s8,
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
     them to the same data type as the accumulator/output.
   implements:
   - LinalgConvolutionOpInterface
 structured_op: !LinalgStructuredOpConfig
   args:
   - !LinalgOperandDefConfig
     name: I
-    usage: Input
+    kind: input_tensor
     type_var: T1
     shape_map: affine_map<()[s0, s1, s2, s3, s4, s5, s6, s7, s8, s9, s10, s11, s12,
       s13, s14] -> (s0, s1 * s2 + s3 * s4, s5 * s6 + s7 * s8, s9 * s10 + s11 * s12,
       s13)>
   - !LinalgOperandDefConfig
     name: K
-    usage: Input
+    kind: input_tensor
     type_var: T2
     shape_map: affine_map<()[s0, s1, s2, s3, s4, s5, s6, s7, s8, s9, s10, s11, s12,
       s13, s14] -> (s3, s7, s11, s13, s14)>
   - !LinalgOperandDefConfig
     name: O
-    usage: Output
+    kind: output_tensor
     type_var: U
     shape_map: affine_map<()[s0, s1, s2, s3, s4, s5, s6, s7, s8, s9, s10, s11, s12,
       s13, s14] -> (s0, s1, s5, s9, s14)>
   - !LinalgOperandDefConfig
     name: strides
-    usage: IndexAttr
+    kind: index_attr
     index_attr_map: affine_map<()[s0, s1, s2, s3, s4, s5, s6, s7, s8, s9, s10, s11,
       s12, s13, s14] -> (s2, s6, s10)>
     default_vals:
     - 1
     - 1
     - 1
   - !LinalgOperandDefConfig
     name: dilations
-    usage: IndexAttr
+    kind: index_attr
     index_attr_map: affine_map<()[s0, s1, s2, s3, s4, s5, s6, s7, s8, s9, s10, s11,
       s12, s13, s14] -> (s4, s8, s12)>
     default_vals:
     - 1
     - 1
     - 1
   indexing_maps: !LinalgIndexingMapsConfig
     static_indexing_maps:
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
     them to the same data type as the accumulator/output. Multiplier is set to 1
     which is a special case for most depthwise convolutions.
   implements:
   - LinalgConvolutionOpInterface
 structured_op: !LinalgStructuredOpConfig
   args:
   - !LinalgOperandDefConfig
     name: I
-    usage: Input
+    kind: input_tensor
     type_var: T1
     shape_map: affine_map<()[s0, s1, s2, s3, s4, s5] -> (s0, s1 * s2 + s3 * s4, s5)>
   - !LinalgOperandDefConfig
     name: K
-    usage: Input
+    kind: input_tensor
     type_var: T2
     shape_map: affine_map<()[s0, s1, s2, s3, s4, s5] -> (s3, s5)>
   - !LinalgOperandDefConfig
     name: O
-    usage: Output
+    kind: output_tensor
     type_var: U
     shape_map: affine_map<()[s0, s1, s2, s3, s4, s5] -> (s0, s1, s5)>
   - !LinalgOperandDefConfig
     name: strides
-    usage: IndexAttr
+    kind: index_attr
     index_attr_map: affine_map<()[s0, s1, s2, s3, s4, s5] -> (s2)>
     default_vals:
     - 1
   - !LinalgOperandDefConfig
     name: dilations
-    usage: IndexAttr
+    kind: index_attr
     index_attr_map: affine_map<()[s0, s1, s2, s3, s4, s5] -> (s4)>
     default_vals:
     - 1
   indexing_maps: !LinalgIndexingMapsConfig
     static_indexing_maps:
     - affine_map<(d0, d1, d2, d3)[s0, s1, s2, s3, s4, s5] -> (d0, d1 * s2 + d3 * s4,
       d2)>
     - affine_map<(d0, d1, d2, d3)[s0, s1, s2, s3, s4, s5] -> (d3, d2)>
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
     them to the same data type as the accumulator/output. Multiplier is set to 1
     which is a special case for most depthwise convolutions.
   implements:
   - LinalgConvolutionOpInterface
 structured_op: !LinalgStructuredOpConfig
   args:
   - !LinalgOperandDefConfig
     name: I
-    usage: Input
+    kind: input_tensor
     type_var: T1
     shape_map: affine_map<()[s0, s1, s2, s3, s4, s5, s6, s7, s8, s9] -> (s0, s1 *
       s2 + s3 * s4, s5 * s6 + s7 * s8, s9)>
   - !LinalgOperandDefConfig
     name: K
-    usage: Input
+    kind: input_tensor
     type_var: T2
     shape_map: affine_map<()[s0, s1, s2, s3, s4, s5, s6, s7, s8, s9] -> (s3, s7, s9)>
   - !LinalgOperandDefConfig
     name: O
-    usage: Output
+    kind: output_tensor
     type_var: U
     shape_map: affine_map<()[s0, s1, s2, s3, s4, s5, s6, s7, s8, s9] -> (s0, s1, s5,
       s9)>
   - !LinalgOperandDefConfig
     name: strides
-    usage: IndexAttr
+    kind: index_attr
     index_attr_map: affine_map<()[s0, s1, s2, s3, s4, s5, s6, s7, s8, s9] -> (s2,
       s6)>
     default_vals:
     - 1
     - 1
   - !LinalgOperandDefConfig
     name: dilations
-    usage: IndexAttr
+    kind: index_attr
     index_attr_map: affine_map<()[s0, s1, s2, s3, s4, s5, s6, s7, s8, s9] -> (s4,
       s8)>
     default_vals:
     - 1
     - 1
   indexing_maps: !LinalgIndexingMapsConfig
     static_indexing_maps:
     - affine_map<(d0, d1, d2, d3, d4, d5)[s0, s1, s2, s3, s4, s5, s6, s7, s8, s9]
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
     them to the same data type as the accumulator/output.
   implements:
   - LinalgConvolutionOpInterface
 structured_op: !LinalgStructuredOpConfig
   args:
   - !LinalgOperandDefConfig
     name: I
-    usage: Input
+    kind: input_tensor
     type_var: T1
     shape_map: affine_map<()[s0, s1, s2, s3, s4, s5, s6, s7, s8, s9] -> (s0, s1 *
       s2 + s3 * s4, s5 * s6 + s7 * s8, s9)>
   - !LinalgOperandDefConfig
     name: K
-    usage: Input
+    kind: input_tensor
     type_var: T2
     shape_map: affine_map<()[s0, s1, s2, s3, s4, s5, s6, s7, s8, s9] -> (s3, s7, s9)>
   - !LinalgOperandDefConfig
     name: IZp
-    usage: Input
+    kind: scalar
     type_var: I32
   - !LinalgOperandDefConfig
     name: KZp
-    usage: Input
+    kind: scalar
     type_var: I32
   - !LinalgOperandDefConfig
     name: O
-    usage: Output
+    kind: output_tensor
     type_var: U
     shape_map: affine_map<()[s0, s1, s2, s3, s4, s5, s6, s7, s8, s9] -> (s0, s1, s5,
       s9)>
   - !LinalgOperandDefConfig
     name: strides
-    usage: IndexAttr
+    kind: index_attr
     index_attr_map: affine_map<()[s0, s1, s2, s3, s4, s5, s6, s7, s8, s9] -> (s2,
       s6)>
     default_vals:
     - 1
     - 1
   - !LinalgOperandDefConfig
     name: dilations
-    usage: IndexAttr
+    kind: index_attr
     index_attr_map: affine_map<()[s0, s1, s2, s3, s4, s5, s6, s7, s8, s9] -> (s4,
       s8)>
     default_vals:
     - 1
     - 1
   indexing_maps: !LinalgIndexingMapsConfig
     static_indexing_maps:
     - affine_map<(d0, d1, d2, d3, d4, d5)[s0, s1, s2, s3, s4, s5, s6, s7, s8, s9]
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
     them to the same data type as the accumulator/output.
   implements:
   - LinalgConvolutionOpInterface
 structured_op: !LinalgStructuredOpConfig
   args:
   - !LinalgOperandDefConfig
     name: I
-    usage: Input
+    kind: input_tensor
     type_var: T1
     shape_map: affine_map<()[s0, s1, s2, s3, s4, s5, s6, s7, s8, s9, s10] -> (s0,
       s1 * s2 + s3 * s4, s5 * s6 + s7 * s8, s9)>
   - !LinalgOperandDefConfig
     name: K
-    usage: Input
+    kind: input_tensor
     type_var: T2
     shape_map: affine_map<()[s0, s1, s2, s3, s4, s5, s6, s7, s8, s9, s10] -> (s3,
       s7, s9, s10)>
   - !LinalgOperandDefConfig
     name: O
-    usage: Output
+    kind: output_tensor
     type_var: U
     shape_map: affine_map<()[s0, s1, s2, s3, s4, s5, s6, s7, s8, s9, s10] -> (s0,
       s1, s5, s9, s10)>
   - !LinalgOperandDefConfig
     name: strides
-    usage: IndexAttr
+    kind: index_attr
     index_attr_map: affine_map<()[s0, s1, s2, s3, s4, s5, s6, s7, s8, s9, s10] ->
       (s2, s6)>
     default_vals:
     - 1
     - 1
   - !LinalgOperandDefConfig
     name: dilations
-    usage: IndexAttr
+    kind: index_attr
     index_attr_map: affine_map<()[s0, s1, s2, s3, s4, s5, s6, s7, s8, s9, s10] ->
       (s4, s8)>
     default_vals:
     - 1
     - 1
   indexing_maps: !LinalgIndexingMapsConfig
     static_indexing_maps:
     - affine_map<(d0, d1, d2, d3, d4, d5, d6)[s0, s1, s2, s3, s4, s5, s6, s7, s8,
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
     them to the same data type as the accumulator/output.
   implements:
   - LinalgConvolutionOpInterface
 structured_op: !LinalgStructuredOpConfig
   args:
   - !LinalgOperandDefConfig
     name: I
-    usage: Input
+    kind: input_tensor
     type_var: T1
     shape_map: affine_map<()[s0, s1, s2, s3, s4, s5, s6, s7, s8, s9, s10] -> (s0,
       s1 * s2 + s3 * s4, s5 * s6 + s7 * s8, s9)>
   - !LinalgOperandDefConfig
     name: K
-    usage: Input
+    kind: input_tensor
     type_var: T2
     shape_map: affine_map<()[s0, s1, s2, s3, s4, s5, s6, s7, s8, s9, s10] -> (s3,
       s7, s9, s10)>
   - !LinalgOperandDefConfig
     name: IZp
-    usage: Input
+    kind: scalar
     type_var: I32
   - !LinalgOperandDefConfig
     name: KZp
-    usage: Input
+    kind: scalar
     type_var: I32
   - !LinalgOperandDefConfig
     name: O
-    usage: Output
+    kind: output_tensor
     type_var: U
     shape_map: affine_map<()[s0, s1, s2, s3, s4, s5, s6, s7, s8, s9, s10] -> (s0,
       s1, s5, s9, s10)>
   - !LinalgOperandDefConfig
     name: strides
-    usage: IndexAttr
+    kind: index_attr
     index_attr_map: affine_map<()[s0, s1, s2, s3, s4, s5, s6, s7, s8, s9, s10] ->
       (s2, s6)>
     default_vals:
     - 1
     - 1
   - !LinalgOperandDefConfig
     name: dilations
-    usage: IndexAttr
+    kind: index_attr
     index_attr_map: affine_map<()[s0, s1, s2, s3, s4, s5, s6, s7, s8, s9, s10] ->
       (s4, s8)>
     default_vals:
     - 1
     - 1
   indexing_maps: !LinalgIndexingMapsConfig
     static_indexing_maps:
     - affine_map<(d0, d1, d2, d3, d4, d5, d6)[s0, s1, s2, s3, s4, s5, s6, s7, s8,
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
     data type as the accumulator/output.
   implements:
   - LinalgConvolutionOpInterface
 structured_op: !LinalgStructuredOpConfig
   args:
   - !LinalgOperandDefConfig
     name: I
-    usage: Input
+    kind: input_tensor
     type_var: T1
     shape_map: affine_map<()[s0, s1, s2, s3, s4, s5, s6, s7, s8, s9] -> (s0, s1 *
       s2 + s3 * s4, s5 * s6 + s7 * s8, s9)>
   - !LinalgOperandDefConfig
     name: K
-    usage: Input
+    kind: input_tensor
     type_var: T2
     shape_map: affine_map<()[s0, s1, s2, s3, s4, s5, s6, s7, s8, s9] -> (s3, s7)>
   - !LinalgOperandDefConfig
     name: O
-    usage: Output
+    kind: output_tensor
     type_var: U
     shape_map: affine_map<()[s0, s1, s2, s3, s4, s5, s6, s7, s8, s9] -> (s0, s1, s5,
       s9)>
   - !LinalgOperandDefConfig
     name: strides
-    usage: IndexAttr
+    kind: index_attr
     index_attr_map: affine_map<()[s0, s1, s2, s3, s4, s5, s6, s7, s8, s9] -> (s2,
       s6)>
     default_vals:
     - 1
     - 1
   - !LinalgOperandDefConfig
     name: dilations
-    usage: IndexAttr
+    kind: index_attr
     index_attr_map: affine_map<()[s0, s1, s2, s3, s4, s5, s6, s7, s8, s9] -> (s4,
       s8)>
     default_vals:
     - 1
     - 1
   indexing_maps: !LinalgIndexingMapsConfig
     static_indexing_maps:
     - affine_map<(d0, d1, d2, d3, d4, d5)[s0, s1, s2, s3, s4, s5, s6, s7, s8, s9]
- !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
     data type as the accumulator/output.
   implements:
   - LinalgConvolutionOpInterface
 structured_op: !LinalgStructuredOpConfig
   args:
   - !LinalgOperandDefConfig
     name: I
-    usage: Input
+    kind: input_tensor
     type_var: T1
     shape_map: affine_map<()[s0, s1, s2, s3, s4, s5, s6, s7, s8, s9] -> (s0, s1 *
       s2 + s3 * s4, s5 * s6 + s7 * s8, s9)>
   - !LinalgOperandDefConfig
     name: K
-    usage: Input
+    kind: input_tensor
     type_var: T2
     shape_map: affine_map<()[s0, s1, s2, s3, s4, s5, s6, s7, s8, s9] -> (s3, s7)>
   - !LinalgOperandDefConfig
     name: O
-    usage: Output
+    kind: output_tensor
     type_var: U
     shape_map: affine_map<()[s0, s1, s2, s3, s4, s5, s6, s7, s8, s9] -> (s0, s1, s5,
       s9)>
   - !LinalgOperandDefConfig
     name: strides
-    usage: IndexAttr
+    kind: index_attr
     index_attr_map: affine_map<()[s0, s1, s2, s3, s4, s5, s6, s7, s8, s9] -> (s2,
       s6)>
     default_vals:
     - 1
     - 1
   - !LinalgOperandDefConfig
     name: dilations
-    usage: IndexAttr
+    kind: index_attr
     index_attr_map: affine_map<()[s0, s1, s2, s3, s4, s5, s6, s7, s8, s9] -> (s4,
       s8)>
     default_vals:
     - 1
     - 1
   indexing_maps: !LinalgIndexingMapsConfig
     static_indexing_maps:
     - affine_map<(d0, d1, d2, d3, d4, d5)[s0, s1, s2, s3, s4, s5, s6, s7, s8, s9]
- !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
     data type as the accumulator/output.
   implements:
   - LinalgConvolutionOpInterface
 structured_op: !LinalgStructuredOpConfig
   args:
   - !LinalgOperandDefConfig
     name: I
-    usage: Input
+    kind: input_tensor
     type_var: T1
     shape_map: affine_map<()[s0, s1, s2, s3, s4, s5, s6, s7, s8, s9] -> (s0, s1 *
       s2 + s3 * s4, s5 * s6 + s7 * s8, s9)>
   - !LinalgOperandDefConfig
     name: K
-    usage: Input
+    kind: input_tensor
     type_var: T2
     shape_map: affine_map<()[s0, s1, s2, s3, s4, s5, s6, s7, s8, s9] -> (s3, s7)>
   - !LinalgOperandDefConfig
     name: O
-    usage: Output
+    kind: output_tensor
     type_var: U
     shape_map: affine_map<()[s0, s1, s2, s3, s4, s5, s6, s7, s8, s9] -> (s0, s1, s5,
       s9)>
   - !LinalgOperandDefConfig
     name: strides
-    usage: IndexAttr
+    kind: index_attr
     index_attr_map: affine_map<()[s0, s1, s2, s3, s4, s5, s6, s7, s8, s9] -> (s2,
       s6)>
     default_vals:
     - 1
     - 1
   - !LinalgOperandDefConfig
     name: dilations
-    usage: IndexAttr
+    kind: index_attr
     index_attr_map: affine_map<()[s0, s1, s2, s3, s4, s5, s6, s7, s8, s9] -> (s4,
       s8)>
     default_vals:
     - 1
     - 1
   indexing_maps: !LinalgIndexingMapsConfig
     static_indexing_maps:
     - affine_map<(d0, d1, d2, d3, d4, d5)[s0, s1, s2, s3, s4, s5, s6, s7, s8, s9]
- !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
     data type as the accumulator/output.
   implements:
   - LinalgConvolutionOpInterface
 structured_op: !LinalgStructuredOpConfig
   args:
   - !LinalgOperandDefConfig
     name: I
-    usage: Input
+    kind: input_tensor
     type_var: T1
     shape_map: affine_map<()[s0, s1, s2, s3, s4, s5, s6, s7, s8, s9] -> (s0, s1, s2
       * s3 + s4 * s5, s6 * s7 + s8 * s9)>
   - !LinalgOperandDefConfig
     name: K
-    usage: Input
+    kind: input_tensor
     type_var: T2
     shape_map: affine_map<()[s0, s1, s2, s3, s4, s5, s6, s7, s8, s9] -> (s4, s8)>
   - !LinalgOperandDefConfig
     name: O
-    usage: Output
+    kind: output_tensor
     type_var: U
     shape_map: affine_map<()[s0, s1, s2, s3, s4, s5, s6, s7, s8, s9] -> (s0, s1, s2,
       s6)>
   - !LinalgOperandDefConfig
     name: strides
-    usage: IndexAttr
+    kind: index_attr
     index_attr_map: affine_map<()[s0, s1, s2, s3, s4, s5, s6, s7, s8, s9] -> (s3,
       s7)>
     default_vals:
     - 1
     - 1
   - !LinalgOperandDefConfig
     name: dilations
-    usage: IndexAttr
+    kind: index_attr
     index_attr_map: affine_map<()[s0, s1, s2, s3, s4, s5, s6, s7, s8, s9] -> (s5,
       s9)>
     default_vals:
     - 1
     - 1
   indexing_maps: !LinalgIndexingMapsConfig
     static_indexing_maps:
     - affine_map<(d0, d1, d2, d3, d4, d5)[s0, s1, s2, s3, s4, s5, s6, s7, s8, s9]
- !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
     data type as the accumulator/output.
   implements:
   - LinalgConvolutionOpInterface
 structured_op: !LinalgStructuredOpConfig
   args:
   - !LinalgOperandDefConfig
     name: I
-    usage: Input
+    kind: input_tensor
     type_var: T1
     shape_map: affine_map<()[s0, s1, s2, s3, s4, s5, s6, s7, s8, s9] -> (s0, s1 *
       s2 + s3 * s4, s5 * s6 + s7 * s8, s9)>
   - !LinalgOperandDefConfig
     name: K
-    usage: Input
+    kind: input_tensor
     type_var: T2
     shape_map: affine_map<()[s0, s1, s2, s3, s4, s5, s6, s7, s8, s9] -> (s3, s7)>
   - !LinalgOperandDefConfig
     name: O
-    usage: Output
+    kind: output_tensor
     type_var: U
     shape_map: affine_map<()[s0, s1, s2, s3, s4, s5, s6, s7, s8, s9] -> (s0, s1, s5,
       s9)>
   - !LinalgOperandDefConfig
     name: strides
-    usage: IndexAttr
+    kind: index_attr
     index_attr_map: affine_map<()[s0, s1, s2, s3, s4, s5, s6, s7, s8, s9] -> (s2,
       s6)>
     default_vals:
     - 1
     - 1
   - !LinalgOperandDefConfig
     name: dilations
-    usage: IndexAttr
+    kind: index_attr
     index_attr_map: affine_map<()[s0, s1, s2, s3, s4, s5, s6, s7, s8, s9] -> (s4,
       s8)>
     default_vals:
     - 1
     - 1
   indexing_maps: !LinalgIndexingMapsConfig
     static_indexing_maps:
     - affine_map<(d0, d1, d2, d3, d4, d5)[s0, s1, s2, s3, s4, s5, s6, s7, s8, s9]
- !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
     data type as the accumulator/output.
   implements:
   - LinalgConvolutionOpInterface
 structured_op: !LinalgStructuredOpConfig
   args:
   - !LinalgOperandDefConfig
     name: I
-    usage: Input
+    kind: input_tensor
     type_var: T1
     shape_map: affine_map<()[s0, s1, s2, s3, s4, s5, s6, s7, s8, s9] -> (s0, s1 *
       s2 + s3 * s4, s5 * s6 + s7 * s8, s9)>
   - !LinalgOperandDefConfig
     name: K
-    usage: Input
+    kind: input_tensor
     type_var: T2
     shape_map: affine_map<()[s0, s1, s2, s3, s4, s5, s6, s7, s8, s9] -> (s3, s7)>
   - !LinalgOperandDefConfig
     name: O
-    usage: Output
+    kind: output_tensor
     type_var: U
     shape_map: affine_map<()[s0, s1, s2, s3, s4, s5, s6, s7, s8, s9] -> (s0, s1, s5,
       s9)>
   - !LinalgOperandDefConfig
     name: strides
-    usage: IndexAttr
+    kind: index_attr
     index_attr_map: affine_map<()[s0, s1, s2, s3, s4, s5, s6, s7, s8, s9] -> (s2,
       s6)>
     default_vals:
     - 1
     - 1
   - !LinalgOperandDefConfig
     name: dilations
-    usage: IndexAttr
+    kind: index_attr
     index_attr_map: affine_map<()[s0, s1, s2, s3, s4, s5, s6, s7, s8, s9] -> (s4,
       s8)>
     default_vals:
     - 1
     - 1
   indexing_maps: !LinalgIndexingMapsConfig
     static_indexing_maps:
     - affine_map<(d0, d1, d2, d3, d4, d5)[s0, s1, s2, s3, s4, s5, s6, s7, s8, s9]
- !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
     data type as the accumulator/output.
   implements:
   - LinalgConvolutionOpInterface
 structured_op: !LinalgStructuredOpConfig
   args:
   - !LinalgOperandDefConfig
     name: I
-    usage: Input
+    kind: input_tensor
     type_var: T1
     shape_map: affine_map<()[s0, s1, s2, s3, s4, s5, s6, s7, s8, s9, s10, s11, s12,
       s13] -> (s0, s1 * s2 + s3 * s4, s5 * s6 + s7 * s8, s9 * s10 + s11 * s12, s13)>
   - !LinalgOperandDefConfig
     name: K
-    usage: Input
+    kind: input_tensor
     type_var: T2
     shape_map: affine_map<()[s0, s1, s2, s3, s4, s5, s6, s7, s8, s9, s10, s11, s12,
       s13] -> (s3, s7, s11)>
   - !LinalgOperandDefConfig
     name: O
-    usage: Output
+    kind: output_tensor
     type_var: U
     shape_map: affine_map<()[s0, s1, s2, s3, s4, s5, s6, s7, s8, s9, s10, s11, s12,
       s13] -> (s0, s1, s5, s9, s13)>
   - !LinalgOperandDefConfig
     name: strides
-    usage: IndexAttr
+    kind: index_attr
     index_attr_map: affine_map<()[s0, s1, s2, s3, s4, s5, s6, s7, s8, s9, s10, s11,
       s12, s13] -> (s2, s6, s10)>
     default_vals:
     - 1
     - 1
     - 1
   - !LinalgOperandDefConfig
     name: dilations
-    usage: IndexAttr
+    kind: index_attr
     index_attr_map: affine_map<()[s0, s1, s2, s3, s4, s5, s6, s7, s8, s9, s10, s11,
       s12, s13] -> (s4, s8, s12)>
     default_vals:
     - 1
     - 1
     - 1
   indexing_maps: !LinalgIndexingMapsConfig
     static_indexing_maps:
- !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
     data type as the accumulator/output.
   implements:
   - LinalgConvolutionOpInterface
 structured_op: !LinalgStructuredOpConfig
   args:
   - !LinalgOperandDefConfig
     name: I
-    usage: Input
+    kind: input_tensor
     type_var: T1
     shape_map: affine_map<()[s0, s1, s2, s3, s4, s5, s6, s7, s8, s9, s10, s11, s12,
       s13] -> (s0, s1 * s2 + s3 * s4, s5 * s6 + s7 * s8, s9 * s10 + s11 * s12, s13)>
   - !LinalgOperandDefConfig
     name: K
-    usage: Input
+    kind: input_tensor
     type_var: T2
     shape_map: affine_map<()[s0, s1, s2, s3, s4, s5, s6, s7, s8, s9, s10, s11, s12,
       s13] -> (s3, s7, s11)>
   - !LinalgOperandDefConfig
     name: O
-    usage: Output
+    kind: output_tensor
     type_var: U
     shape_map: affine_map<()[s0, s1, s2, s3, s4, s5, s6, s7, s8, s9, s10, s11, s12,
       s13] -> (s0, s1, s5, s9, s13)>
   - !LinalgOperandDefConfig
     name: strides
-    usage: IndexAttr
+    kind: index_attr
     index_attr_map: affine_map<()[s0, s1, s2, s3, s4, s5, s6, s7, s8, s9, s10, s11,
       s12, s13] -> (s2, s6, s10)>
     default_vals:
     - 1
     - 1
     - 1
   - !LinalgOperandDefConfig
     name: dilations
-    usage: IndexAttr
+    kind: index_attr
     index_attr_map: affine_map<()[s0, s1, s2, s3, s4, s5, s6, s7, s8, s9, s10, s11,
       s12, s13] -> (s4, s8, s12)>
     default_vals:
     - 1
     - 1
     - 1
   indexing_maps: !LinalgIndexingMapsConfig
     static_indexing_maps:
- !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
     data type as the accumulator/output.
   implements:
   - LinalgConvolutionOpInterface
 structured_op: !LinalgStructuredOpConfig
   args:
   - !LinalgOperandDefConfig
     name: I
-    usage: Input
+    kind: input_tensor
     type_var: T1
     shape_map: affine_map<()[s0, s1, s2, s3, s4, s5, s6, s7, s8, s9, s10, s11, s12,
       s13] -> (s0, s1 * s2 + s3 * s4, s5 * s6 + s7 * s8, s9 * s10 + s11 * s12, s13)>
   - !LinalgOperandDefConfig
     name: K
-    usage: Input
+    kind: input_tensor
     type_var: T2
     shape_map: affine_map<()[s0, s1, s2, s3, s4, s5, s6, s7, s8, s9, s10, s11, s12,
       s13] -> (s3, s7, s11)>
   - !LinalgOperandDefConfig
     name: O
-    usage: Output
+    kind: output_tensor
     type_var: U
     shape_map: affine_map<()[s0, s1, s2, s3, s4, s5, s6, s7, s8, s9, s10, s11, s12,
       s13] -> (s0, s1, s5, s9, s13)>
   - !LinalgOperandDefConfig
     name: strides
-    usage: IndexAttr
+    kind: index_attr
     index_attr_map: affine_map<()[s0, s1, s2, s3, s4, s5, s6, s7, s8, s9, s10, s11,
       s12, s13] -> (s2, s6, s10)>
     default_vals:
     - 1
     - 1
     - 1
   - !LinalgOperandDefConfig
     name: dilations
-    usage: IndexAttr
+    kind: index_attr
     index_attr_map: affine_map<()[s0, s1, s2, s3, s4, s5, s6, s7, s8, s9, s10, s11,
       s12, s13] -> (s4, s8, s12)>
     default_vals:
     - 1
     - 1
     - 1
   indexing_maps: !LinalgIndexingMapsConfig
     static_indexing_maps:
- !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
     accesses only and is thus rank polymorphic. Numeric casting is performed on
     the value operand, promoting it to the same data type as the output.
 structured_op: !LinalgStructuredOpConfig
   args:
   - !LinalgOperandDefConfig
     name: value
-    usage: Input
+    kind: scalar
     type_var: T1
   - !LinalgOperandDefConfig
     name: O
-    usage: Output
+    kind: output_tensor
     type_var: U
     shape_map: affine_map<() -> ()>
   indexing_maps: !LinalgIndexingMapsConfig
     static_indexing_maps:
     - 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
     generator. It provides no guarantees regarding the distribution of the
     generated random numbers. Instead of generating the random numbers
     sequentially, it instantiates one random number generator per data element
     and runs them in parallel. The seed operand and the indices of the data
     element seed the random number generation. The min and max operands limit
     the range of the generated random numbers.
 structured_op: !LinalgStructuredOpConfig
   args:
   - !LinalgOperandDefConfig
     name: min
-    usage: Input
+    kind: scalar
     type_var: F64
   - !LinalgOperandDefConfig
     name: max
-    usage: Input
+    kind: scalar
     type_var: F64
   - !LinalgOperandDefConfig
     name: seed
-    usage: Input
+    kind: scalar
     type_var: I32
   - !LinalgOperandDefConfig
     name: O
-    usage: Output
+    kind: output_tensor
     type_var: T
     shape_map: affine_map<()[s0, s1] -> (s0, s1)>
   indexing_maps: !LinalgIndexingMapsConfig
     static_indexing_maps:
     - affine_map<(d0, d1)[s0, s1] -> ()>
     - affine_map<(d0, d1)[s0, s1] -> ()>
     - affine_map<(d0, d1)[s0, s1] -> ()>
     - affine_map<(d0, d1)[s0, s1] -> (d0, d1)>
   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
     data type as the accumulator/output.
 structured_op: !LinalgStructuredOpConfig
   args:
   - !LinalgOperandDefConfig
     name: I
-    usage: Input
+    kind: input_tensor
     type_var: T
     shape_map: affine_map<()[s0, s1] -> (s0, s1)>
   - !LinalgOperandDefConfig
     name: O
-    usage: Output
+    kind: output_tensor
     type_var: U
     shape_map: affine_map<()[s0, s1] -> (s0, s1)>
   indexing_maps: !LinalgIndexingMapsConfig
     static_indexing_maps:
     - affine_map<(d0, d1)[s0, s1] -> (d0, d1)>
     - affine_map<(d0, d1)[s0, s1] -> (d0, d1)>
   iterator_types:
   - parallel
   - parallel
   assignments:
   - !ScalarAssign
     arg: O
     value: !ScalarExpression
       arith_fn:
         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.
 

 
 
 ###############################################################################
 # Operand definitions.
 ###############################################################################
 
 
 class OperandKind(Enum):
-  InputTensor = 0
-  Scalar = 1
-  OutputTensor = 2
-  IndexAttr = 3
+  INPUT_TENSOR = 0
+  SCALAR = 1
+  OUTPUT_TENSOR = 2
+  INDEX_ATTR = 3
+  TYPE_FN_ATTR = 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.INPUT_TENSOR)
+
+  def is_tensor(self) -> bool:
+    return (self.kind == OperandKind.INPUT_TENSOR or
+            self.kind == OperandKind.OUTPUT_TENSOR)
+
+  def is_attribute(self) -> bool:
+    return (self.kind == OperandKind.INDEX_ATTR or
+            self.kind == OperandKind.TYPE_FN_ATTR)
+
   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
def __init__(self,
                index_dims: Optional[Sequence[DimDef]] = None,
                output: bool = False):
     if index_dims and len(shape) != len(index_dims):
       raise ValueError(f"Expected the shape rank {len(shape)} to match the "
                        f"number of index_dims {len(index_dims)}")
     if index_dims and any(not isinstance(dim, DimDef) for dim in index_dims):
       raise ValueError(f"TensorDef requires index dims of type DimDef but "
                        f"got {index_dims}")
-    kind = OperandKind.OutputTensor if output else OperandKind.InputTensor
+    kind = OperandKind.OUTPUT_TENSOR if output else OperandKind.INPUT_TENSOR
     self.operand_def = OperandDef(
         kind, type_var=type_var, size_exprs=shape, index_dims=index_dims)
 
   def __getitem__(self, dims: Sequence[AffineExprDef]) -> TensorUse:
     assert self.operand_def.owner, "TensorDef is not attached to an op"
     state = AffineBuildState(
         global_state=self.operand_def.owner._affine_state,
         allow_new_symbols=False)
class ScalarDef(TensorExpression):
   """Scalar operand definition.
 
   Scalar operands are forwarded to the body of the structured op as they are.
   A unique name identifies the scalars and an index determines their position in
   the operation's parameter list.
   """
 
   def __init__(self, type_var: TypeVar):
-    self.operand_def = OperandDef(OperandKind.Scalar, type_var=type_var)
+    self.operand_def = OperandDef(OperandKind.SCALAR, type_var=type_var)
 
   @property
   def scalar_name(self) -> str:
     name = self.operand_def.name
     assert name is not None, "ScalarDef not attached"
     return name
 
   def to_scalar_expression(self) -> ScalarExpression:
if any(not isinstance(size, SymbolDef) for size in sizes):
                        f"but got {sizes}")
     if any(not isinstance(default_val, int) for default_val in default):
       raise ValueError(f"IndexAttrDef requires default values of type int "
                        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)
+        OperandKind.INDEX_ATTR, 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.TYPE_FN_ATTR, default_fn=default.fn_name)
+
+  def __call__(self, type_var: TypeVar, arg: TensorExpression) -> TensorTypeFn:
+    return TensorTypeFn(None, self.operand_def, type_var, arg)
 
 
 ###############################################################################
 # Operation definition.
 ###############################################################################
 
 
 class 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.OUTPUT_TENSOR 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

def __init__(self,
     self.index_attr_map = index_attr_map  # type: Optional[_ir.AffineMap]
     self.indexing_map = None  # type: Optional[_ir.AffineMap]
 
   @property
   def name(self) -> str:
     return self.operand_def.name
 
   @property
-  def type_var(self) -> TypeVar:
-    return self.operand_def.type_var
+  def kind(self) -> OperandKind:
+    return self.operand_def.kind
 
   @property
-  def usage(self) -> str:
-    if self.operand_def.kind == OperandKind.IndexAttr:
-      return "IndexAttr"
-    if self.operand_def.kind == OperandKind.OutputTensor:
-      return "Output"
-    return "Input"
+  def type_var(self) -> TypeVar:
+    return self.operand_def.type_var
 
   def to_yaml_custom_dict(self):
-    self_dict = dict(name=self.name, usage=self.usage)
+    self_dict = dict(name=self.name, kind=self.operand_def.kind.name.lower())
     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 "
for _, read_expr in self.writes:
             self.uses[read_use].indexing_map):
           raise ValueError(
               f"Unexpected multi-read of a tensor with different accesses:"
               f"{read_operand_config} vs {read_use}")
         read_operand_config.indexing_map = self.uses[read_use].indexing_map
 
     # 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:
+      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.TYPE_FN_ATTR):
       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))
       assert local_state.local_dim_count == 0
       affine_map = _ir.AffineMap.get(
           dim_count=0, symbol_count=local_state.symbol_count, exprs=exprs)
-      if operand_def.kind == OperandKind.IndexAttr:
+      if operand_def.kind == OperandKind.INDEX_ATTR:
         self.operands[operand_def] = OperandDefConfig(
             operand_def, index_attr_map=affine_map)
       else:
         self.operands[operand_def] = OperandDefConfig(
             operand_def, shape_map=affine_map)
 
   def add_indexed_operand(self, operand_def: OperandDef):
     with self.context:

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"]
+  in_arg_defs = [
+      d for d in all_arg_defs
+      if d.kind == OperandKind.SCALAR or d.kind == OperandKind.INPUT_TENSOR
+  ]
+  out_arg_defs = [
+      d for d in all_arg_defs if d.kind == OperandKind.OUTPUT_TENSOR
+  ]
+  index_attr_arg_defs = [
+      d for d in all_arg_defs if d.kind == OperandKind.INDEX_ATTR
+  ]
+  type_fn_attr_arg_defs = [
+      d for d in all_arg_defs if d.kind == OperandKind.TYPE_FN_ATTR
+  ]
 
   # 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:
+      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 expr.type_fn.attr_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

 // 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/Integration/Dialect/SparseTensor/taco/tools/mlir_pytaco.py

 #  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
 
 """Experimental MLIR-PyTACO with sparse tensor support.
 
 See http://tensor-compiler.org/ for TACO tensor compiler.
 
 This module implements the Python classes for PyTACO index notation. These
 include classes for data types, tensor dimension formats (aka mode formats),
 tensor dimension orderings (aka mode ordering), tensor storage formats, and
 tensors.
class _Stats:
 
   def get_formats(self, idx: int) -> Tuple[ModeFormat]:
     """Gets the ModeFormats for the ith temporary tensor."""
     return tuple(self._get_element(idx).dst_format.format_pack.formats)
 
 
 class _SparseValueInfo(enum.Enum):
   """Describes how a sparse tensor value is stored.
   _UNPACKED: The sparse tensor value is stored as (coordnates, values) in
     Python.
   _PACKED: The sparse tensor value is stored as a C pointer to a packed MLIR
     sparse tensor.
   """
   _UNPACKED = 0
   _PACKED = 1
 
 
expr_to_input_opnd = {
         for e in expr_inputs
     }
 
     # Emit the expression tree, which produces the value assigned to the
     # destination tensor.
     value = self._emit_expression(expr_to_input_opnd, expr_to_info)
     # Emit the structured op representation for the destination tensor.
     dst_opnd = _emit_operand(op_def, op_info.dst_indices, op_info.dst_name,
-                             lang.OperandKind.OutputTensor)
+                             lang.OperandKind.OUTPUT_TENSOR)
     dst_dim_syms = _mlir_dimensions_from_index_vars(op_info.dst_indices)
     dst_use = lang.TensorUse(dst_opnd, dst_dim_syms)
 
     expr_info = expr_to_info[self]
     # If the structured op reduces some indices, explicitly represent the
     # reduction. This is done by generating a ReduceFn for the dimensions being
     # reduced in the linalg dialect and calling the function with the value
     # being reduced. We only support add reduction currently.
if op_info and not isinstance(expr, Access):
     name = op_info.dst_name
   else:
     # The input is a user provided tensor.
     assert isinstance(expr, Access)
     indices = expr.indices
     name = expr.tensor.name
 
   dim_sym = _mlir_symbols_from_index_vars(indices)
-  opnd = lang.OperandDef(lang.OperandKind.InputTensor, lang.T, dim_sym)
+  opnd = lang.OperandDef(lang.OperandKind.INPUT_TENSOR, lang.T, dim_sym)
   op_def.add_operand(name, opnd)
   return opnd

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
+    kind: output_tensor
     type_var: T
     shape_map: affine_map<()[s0, s1] -> (s0, s1)>
+  - !LinalgOperandDefConfig
+    name: cast
+    kind: type_fn_attr
+    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.
metadata: !LinalgOpMetadata
   doc: |-
     Title.
 
     Detailed description.
 structured_op: !LinalgStructuredOpConfig
   args:
   - !LinalgOperandDefConfig
     name: I
-    usage: Input
+    kind: input_tensor
     type_var: T
     shape_map: affine_map<()[s0, s1, s2, s3] -> (s0, s1)>
   - !LinalgOperandDefConfig
     name: O
-    usage: Output
+    kind: output_tensor
     type_var: T
     shape_map: affine_map<()[s0, s1, s2, s3] -> (s0, s1)>
   - !LinalgOperandDefConfig
     name: strides
-    usage: IndexAttr
+    kind: index_attr
     index_attr_map: affine_map<()[s0, s1, s2, s3] -> (s2, s3)>
     default_vals:
     - 1
     - 2
   indexing_maps: !LinalgIndexingMapsConfig
     static_indexing_maps:
     - affine_map<(d0, d1)[s0, s1, s2, s3] -> (d1 * s2, d0 * s3)>
     - affine_map<(d0, d1)[s0, s1, s2, s3] -> (d0, d1)>
metadata: !LinalgOpMetadata
   doc: |-
     Title.
 
     Detailed description.
 structured_op: !LinalgStructuredOpConfig
   args:
   - !LinalgOperandDefConfig
     name: value
-    usage: Input
+    kind: scalar
     type_var: T1
   - !LinalgOperandDefConfig
     name: O
-    usage: Output
+    kind: output_tensor
     type_var: U
     shape_map: affine_map<() -> ()>
   indexing_maps: !LinalgIndexingMapsConfig
     static_indexing_maps:
     - affine_map<() -> ()>
     - affine_map<() -> ()>
   iterator_types: []
   assignments:

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

 # RUN: %PYTHON -m mlir.dialects.linalg.opdsl.dump_oplib --file %s | FileCheck %s
 
 from mlir.dialects.linalg.opdsl.lang import *
 
 
 # CHECK: ---
 # CHECK-LABEL: matmul
 # CHECK: args:
 # CHECK:     name: A
-# CHECK:     usage: Input
+# CHECK:     kind: input_tensor
 # CHECK:     type_var: T
 # CHECK:     shape_map: affine_map<()[s0, s1, s2] -> (s0, s1)>
 # CHECK:     name: B
-# CHECK:     usage: Input
+# CHECK:     kind: input_tensor
 # CHECK:     type_var: T
 # CHECK:     shape_map: affine_map<()[s0, s1, s2] -> (s1, s2)>
 # CHECK:     name: C
-# CHECK:     usage: Output
+# CHECK:     kind: output_tensor
 # CHECK:     type_var: U
 # CHECK:     shape_map: affine_map<()[s0, s1, s2] -> (s0, s2)>
+# CHECK:     name: cast
+# CHECK:     kind: type_fn_attr
+# CHECK:     default_fn: cast
 @linalg_structured_op
 def matmul(
     A=TensorDef(T, S.M, S.K),
     B=TensorDef(T, S.K, S.N),
-    C=TensorDef(U, S.M, S.N, output=True)):
-  C[D.m, D.n] += TypeFn.cast(U, A[D.m, D.k]) * TypeFn.cast(U, B[D.k, D.n])
+    C=TensorDef(U, S.M, S.N, output=True),
+    cast=TypeFnAttrDef(default=TypeFn.cast)):
+  C[D.m, D.n] += cast(U, A[D.m, D.k]) * cast(U, B[D.k, D.n])
 
 
 # CHECK: ---
 # CHECK-LABEL: fill
 # CHECK: args:
 # CHECK:     name: value
-# CHECK:     usage: Input
+# CHECK:     kind: scalar
 # CHECK-NOT: shape_map:
 # CHECK:     type_var: T
 @linalg_structured_op
 def fill(value=ScalarDef(T), O=TensorDef(T, S.M, S.K, output=True)):
   O[D.m, D.n] = value
 
 
 # CHECK: ---
 # CHECK-LABEL: strided_copy
 # CHECK: args:
 # CHECK:     name: I
-# CHECK:     usage: Input
+# CHECK:     kind: input_tensor
 # CHECK:     type_var: T
 # CHECK:     shape_map: affine_map<()[s0, s1, s2, s3, s4, s5] -> (s0, s1)>
 # CHECK:     name: O
-# CHECK:     usage: Output
+# CHECK:     kind: output_tensor
 # CHECK:     type_var: T
 # CHECK:     shape_map: affine_map<()[s0, s1, s2, s3, s4, s5] -> (s2, s3)>
 # CHECK:     name: strides
-# CHECK:     usage: IndexAttr
+# CHECK:     kind: index_attr
 # CHECK:     index_attr_map: affine_map<()[s0, s1, s2, s3, s4, s5] -> (s4, s5)>
 # CHECK:     default_vals:
 # CHECK:     - 1
 # CHECK:     - 2
 @linalg_structured_op
 def strided_copy(
     I=TensorDef(T, S.IH, S.IW),
     O=TensorDef(T, S.OH, S.OW, output=True),
     strides=IndexAttrDef(S.SH, S.SW, default=[1, 2])):
   O[D.oh, D.ow] = I[D.oh * S.SH, D.ow * S.SW]

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

 # CHECK:        operands:
 # CHECK:          arith_fn:
 # CHECK:            fn_name: mul
 # CHECK:            operands:
 # CHECK:              type_fn:
 # CHECK:                type_var: U
 # CHECK:                operands:
 # CHECK:                  scalar_arg: A
+# CHECK:                attr_name: cast
 # CHECK:              type_fn:
 # CHECK:                type_var: U
 # CHECK:                operands:
 # CHECK:                  scalar_arg: B
+# CHECK:                attr_name: cast
 @linalg_structured_op
 def matmul(
     A=TensorDef(T, S.M, S.K),
     B=TensorDef(T, S.K, S.N),
-    C=TensorDef(U, S.M, S.N, output=True)):
-  C[D.m, D.n] += TypeFn.cast(U, A[D.m, D.k]) * TypeFn.cast(U, B[D.k, D.n])
+    C=TensorDef(U, S.M, S.N, output=True),
+    cast=TypeFnAttrDef(default=TypeFn.cast)):
+  C[D.m, D.n] += cast(U, A[D.m, D.k]) * cast(U, B[D.k, D.n])
 
 
 # CHECK: ---
 # CHECK-LABEL: constants
 # CHECK: assignments:
 # CHECK:  -
 # CHECK:    arg: O
 # CHECK:      arith_fn:

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/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 LinalgOperandDefKind {
+  InputTensor,
+  Scalar,
+  OutputTensor,
+  IndexAttr,
+  TypeFnAttr
+};
 
 struct LinalgOperandDef {
   std::string name;
-  LinalgOperandDefUsage usage;
+  LinalgOperandDefKind kind;
   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.mapRequired("kind", info.kind);
     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);
+struct ScalarEnumerationTraits<LinalgOperandDefKind> {
+  static void enumeration(IO &io, LinalgOperandDefKind &value) {
+    io.enumCase(value, "input_tensor", LinalgOperandDefKind::InputTensor);
+    io.enumCase(value, "scalar", LinalgOperandDefKind::InputTensor);
+    io.enumCase(value, "output_tensor", LinalgOperandDefKind::OutputTensor);
+    io.enumCase(value, "index_attr", LinalgOperandDefKind::IndexAttr);
+    io.enumCase(value, "type_fn_attr", LinalgOperandDefKind::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,
if (typeVar == "I64")
     return std::string("helper.getIntegerType(64)");
   if (typeVar == "F32")
     return std::string("helper.getFloat32Type()");
   if (typeVar == "F64")
     return std::string("helper.getFloat64Type()");
 
   // Search all argument types.
   for (const auto &it : llvm::enumerate(args)) {
-    if (it.value().usage != LinalgOperandDefUsage::Input &&
-        it.value().usage != LinalgOperandDefUsage::Output)
+    if (it.value().kind != LinalgOperandDefKind::InputTensor &&
+        it.value().kind != LinalgOperandDefKind::Scalar &&
+        it.value().kind != LinalgOperandDefKind::OutputTensor)
       continue;
     if (it.value().typeVar.getValue() == typeVar)
       return llvm::formatv("block.getArgument({0}).getType()", it.index())
           .str();
   }
 
   return None;
 }
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.kind == LinalgOperandDefKind::IndexAttr ||
+               arg.kind == LinalgOperandDefKind::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.kind == LinalgOperandDefKind::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.kind == LinalgOperandDefKind::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.kind == LinalgOperandDefKind::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.kind == LinalgOperandDefKind::InputTensor ||
+               arg.kind == LinalgOperandDefKind::Scalar ||
+               arg.kind == LinalgOperandDefKind::OutputTensor;
       });
 
   // 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();
 
)FMT";
         // {1}: Symbol position
         // {2}: Attribute index
         static const char structuredOpAccessAttrFormat[] = R"FMT(
 int64_t cst{1} = self.{0}().getValues<int64_t>()[{2}];
 exprs.push_back(getAffineConstantExpr(cst{1}, context));
 )FMT";
         // Update all symbol bindings mapped to an attribute.
         for (LinalgOperandDef &arg : opConfig.structuredOp->args) {
-          if (arg.usage != LinalgOperandDefUsage::IndexAttr)
+          if (arg.kind != LinalgOperandDefKind::IndexAttr)
             continue;
           assert(arg.indexAttrMap.hasValue());
           for (auto &en :
                llvm::enumerate(arg.indexAttrMap->affineMap().getResults())) {
             if (auto symbol = en.value().dyn_cast<AffineSymbolExpr>()) {
               symbolBindings[symbol.getPosition()] =
                   llvm::formatv(structuredOpAccessAttrFormat, arg.name,
                                 symbol.getPosition(), en.index());
std::string {0}::getLibraryCallName() {{
   return generateLibraryCallName(getOperation());
 }
 )FMT";
     os << llvm::formatv(structuredOpGetLibraryCallFormat, className);
   }
 
   // hasDynamicIndexingMaps() and verifyIndexingMapRequiredAttributes()
   if (llvm::any_of(opConfig.structuredOp->args, [](LinalgOperandDef &arg) {
-        return arg.usage == LinalgOperandDefUsage::IndexAttr;
+        return arg.kind == LinalgOperandDefKind::IndexAttr;
       })) {
     std::vector<std::string> attrVerifications;
     for (LinalgOperandDef &arg : opConfig.structuredOp->args) {
-      if (arg.usage != LinalgOperandDefUsage::IndexAttr)
+      if (arg.kind != LinalgOperandDefKind::IndexAttr)
         continue;
       assert(arg.indexAttrMap.hasValue());
       // Verify index attribute. Paramters:
       // {0}: Attribute name
       // {1}: Attribute size
       static const char attrFmt[] = R"FMT(
 if (auto attr = op->getAttrOfType<DenseElementsAttr>("{0}")) {{
   if (!attr.getType().getElementType().isInteger(64))
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::Output)
+      if (arg.kind != LinalgOperandDefKind::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.kind != LinalgOperandDefKind::OutputTensor)
         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;
       ++generatedAssignmentCount;
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.kind == LinalgOperandDefKind::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",