[mlir][OpDSL] Rename `AttributeDef` to `IndexAttrDef`.

[gysit]

Summary

Renaming AttributeDef to IndexAttrDef prepares OpDSL to support different kinds of attributes and more closely reflects the purpose of the attribute.

mlir/docs/Dialects/Linalg/OpDSL.md

 actual computation.
 
 The following example demonstrates the use of attributes:
 
 ```python
 @linalg_structured_op
 def strided_copy(I=TensorDef(T, S.IH, S.IW),
                  O=TensorDef(T, S.OH, S.OW, output=True),
-                 strides=AttributeDef(S.SH, S.SW)):
+                 strides=IndexAttrDef(S.SH, S.SW)):
   """Copy a subset of the input tensor elements to the output tensor"""
   O[D.oh, D.ow] = I[D.oh * S.SH, D.ow * S.SW]
 ```
 
 The operation implements a strided copy from the input tensor `I` to the output
-tensor `O`. The `strides` attribute is bound to an `AttributeDef`. It defines
+tensor `O`. The `strides` attribute is bound to an `IndexAttrDef`. It defines
 the symbols `S.SH` and `S.SW`, which are used to index the input tensor `I`.
 When instantiating the operation, the attribute is set using a named argument:
 
 ```python
 strided_copy(in_tensor, outs=[out_tensor], strides=[1,2])
 ```
 
 The `strides` vector elements substitute the symbols `S.SH` and `S.SW` in the

 The following example demonstrates the index dimension annotation:
 
 ```python
 @linalg_structured_op
 def pooling_poly(
     I=TensorDef(T1, S.N, S.H, S.W, S.C),
     K=TensorDef(T2, S.KH, S.KW, index_dims=[D.kh, D.kw]),
     O=TensorDef(U, S.N, S.OH, S.OW, S.C, output=True),
-    strides=AttributeDef(S.SH, S.SW),
-    dilations=AttributeDef(S.DH, S.DW)):
+    strides=IndexAttrDef(S.SH, S.SW),
+    dilations=IndexAttrDef(S.DH, S.DW)):
   O[D.n, D.oh, D.ow, D.c] += \
       cast(U, I[D.n, D.oh * S.SH + D.kh * S.DH, D.ow * S.SW + D.kw * S.DW, D.c])
 ```
 
 The pooling operation does not access the shape-only tensor `K`. Instead, the
 shapes `S.KH` and `S.KW` specify the iteration domain for the reduction
 dimensions `D.kh` and `D.kw`.
 

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

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:
     return ScalarArg(self.scalar_name).expr()
 
 
-class AttributeDef:
+class IndexAttrDef:
   """Index Attribute definition.
 
   Index attributes provide a way to define and set symbols that can be used in
   indexing expressions. Every attribute specifies a tuple of symbols that at
   compile-time are replaced by integer values.
   """
-  yaml_tag = "!LinalgAttributeDef"
 
   def __init__(self, *sizes: SymbolDef):
     if any(not isinstance(size, SymbolDef) for size in sizes):
-      raise ValueError(f"AttributeDef requires sizes of type SymbolDef but got "
+      raise ValueError(f"IndexAttrDef requires sizes of type SymbolDef but got "
                        f"{sizes}")
     self.operand_def = OperandDef(OperandKind.Attribute, I64, size_exprs=sizes)
 
 
 class Comprehension:
   """Represents a single comprehension."""
 
   def __init__(self, *bindings: Tuple[TensorUse, TensorExpression]):
class OpInterfaceDef:
 
   def __init__(self, cpp_name: str):
     self.cpp_name = cpp_name
 
 
 ContractionOpInterface = OpInterfaceDef("LinalgContractionOpInterface")
 ConvolutionOpInterface = OpInterfaceDef("LinalgConvolutionOpInterface")
 
+
 class OpMetadataDef(YAMLObject):
   """Metadata about the op (generally not behavior impacting)."""
   yaml_tag = "!LinalgOpMetadata"
 
   def __init__(self, name: str, cpp_class_name: Optional[str],
                doc: Optional[str]):
     self.name = name
     self.cpp_class_name = cpp_class_name if cpp_class_name is not None else name

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

 from .config import *
 from .emitter import *
 
 _CONTEXT = threading.local()
 
 StructuredOpOuts = Union[ir.Operation, ir.OpView, ir.OpResultList,
                          Sequence[Union[ir.Value, ir.Operation, ir.OpView]]]
 
+
 @contextmanager
 def bind_op_def(model: LinalgOpDef):
   if hasattr(_CONTEXT, "current_op_def"):
     raise ValueError("Cannot recursively define an operation")
   _CONTEXT.current_op_def = model
   try:
     yield model
   finally:
def linalg_structured_op(dsl_func=None,
   tc_model = 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, AttributeDef)):
+    if isinstance(param_default, (TensorDef, ScalarDef, IndexAttrDef)):
       tc_model.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 AttributeDef(...): "
+          f"TensorDef(...), ScalarDef(...), or IndexAttrDef(...): "
           f"Found {param_name}: {param_default}")
     dsl_func_args.append(param_default)
 
   # Invoke the DSL func to finish populating the model.
   with bind_op_def(tc_model):
     dsl_func(*dsl_func_args)
 
   # TODO: The returned callable should be an IR emitter but that is not

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

def matmul(
 
   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] += 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.
 
   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] += cast_unsigned(U, A[D.m, D.k]) * cast_unsigned(U, B[D.k, D.n])
 
+
 @linalg_structured_op
 def quantized_matmul(
     A=TensorDef(T1, S.M, S.K),
     B=TensorDef(T2, S.K, S.N),
     AZp=ScalarDef(I32),
     BZp=ScalarDef(I32),
     C=TensorDef(U, S.M, S.N, output=True)):
   """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.
   """
   domain(D.m, D.n, D.k)
-  C[D.m, D.n] += (cast(U, A[D.m, D.k]) - cast(U, AZp)) * (cast(U, B[D.k, D.n]) - cast(U, BZp))
+  C[D.m, D.n] += (cast(U, A[D.m, D.k]) - cast(U, AZp)) * (
+      cast(U, B[D.k, D.n]) - cast(U, BZp))
+
 
 @linalg_structured_op
-def mmt4d(lhs=TensorDef(TV.LhsType, S.M, S.K, S.M0, S.K0),
-          rhs=TensorDef(TV.RhsType, S.N, S.K, S.N0, S.K0),
-          accum=TensorDef(TV.AccumType, S.M, S.N, S.M0, S.N0,
-                                  output=True)):
+def mmt4d(
+    lhs=TensorDef(TV.LhsType, S.M, S.K, S.M0, S.K0),
+    accum=TensorDef(TV.AccumType, S.M, S.N, S.M0, S.N0, output=True)):
   """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.
   """
   domain(D.m, D.n, D.k, D.m0, D.n0, D.k0)
   implements(ContractionOpInterface)
-  accum[D.m, D.n, D.m0, D.n0] += cast(TV.AccumType, lhs[D.m, D.k, D.m0, D.k0]) * cast(TV.AccumType, rhs[D.n, D.k, D.n0, D.k0])
+  accum[D.m, D.n, D.m0,
+        D.n0] += cast(TV.AccumType, lhs[D.m, D.k, D.m0, D.k0]) * cast(
+            TV.AccumType, rhs[D.n, D.k, D.n0, D.k0])
+
 
 @linalg_structured_op
 def batch_matmul(
     A=TensorDef(T1, Batch, S.M, S.K),
     B=TensorDef(T2, Batch, S.K, S.N),
     C=TensorDef(U, Batch, S.M, S.N, output=True)):
   """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.
   """
   domain(D.b, D.m, D.n, D.k)
   implements(ContractionOpInterface)
   C[D.b, D.m, D.n] += cast(U, A[D.b, D.m, D.k]) * cast(U, B[D.b, D.k, D.n])
 
+
 @linalg_structured_op
 def quantized_batch_matmul(
     A=TensorDef(T1, Batch, S.M, S.K),
     B=TensorDef(T2, Batch, S.K, S.N),
     AZp=ScalarDef(I32),
     BZp=ScalarDef(I32),
     C=TensorDef(U, Batch, S.M, S.N, output=True)):
   """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.
   """
   domain(D.b, D.m, D.n, D.k)
-  C[D.b, D.m, D.n] += (cast(U, A[D.b, D.m, D.k]) - cast(U, AZp)) * (cast(U, B[D.b, D.k, D.n]) - cast(U, BZp))
+  C[D.b, D.m, D.n] += (cast(U, A[D.b, D.m, D.k]) - cast(U, AZp)) * (
+      cast(U, B[D.b, D.k, D.n]) - cast(U, BZp))
 
 
 @linalg_structured_op
 def matvec(
     A=TensorDef(T1, S.M, S.N),
     y=TensorDef(T2, S.N),
     x=TensorDef(U, S.M, output=True)):
   """Performs a matrix-vector multiplication.
def vecmat(
   implements(ContractionOpInterface)
   x[D.n] += cast(U, y[D.m]) * cast(U, A[D.m, D.n])
 
 
 @linalg_structured_op
 def batch_matvec(
     A=TensorDef(T1, Batch, S.M, S.K),
     B=TensorDef(T2, Batch, S.K),
-    C=TensorDef(U,  Batch, S.M, output=True)):
+    C=TensorDef(U, Batch, S.M, output=True)):
   """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.
   """
   domain(D.b, D.m, D.k)
   implements(ContractionOpInterface)
   C[D.b, D.m] += cast(U, A[D.b, D.m, D.k]) * cast(U, B[D.b, D.k])
 
 
 @linalg_structured_op
 def dot(
     A=TensorDef(T1, S.M), B=TensorDef(T2, S.M), C=TensorDef(U, output=True)):
   """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(ContractionOpInterface)
   C[None] += cast(U, A[D.m]) * cast(U, B[D.m])
 
+
 @linalg_structured_op
 def conv_1d(
     I=TensorDef(T1, S.OW + S.KW),
     K=TensorDef(T2, S.KW),
     O=TensorDef(U, S.OW, output=True)):
   """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(ConvolutionOpInterface)
   domain(D.ow, D.kw)
-  O[D.ow] += cast(
-      U, I[D.ow + D.kw]) * cast(U, K[D.kw])
+  O[D.ow] += cast(U, I[D.ow + D.kw]) * cast(U, K[D.kw])
+
 
 @linalg_structured_op
 def conv_2d(
     I=TensorDef(T1, S.OH + S.KH, S.OW + S.KW),
     K=TensorDef(T2, S.KH, S.KW),
     O=TensorDef(U, S.OH, S.OW, output=True)):
   """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(ConvolutionOpInterface)
   domain(D.oh, D.ow, D.kh, D.kw)
-  O[D.oh, D.ow] += cast(
-      U, I[D.oh + D.kh, D.ow + D.kw]) * cast(U, K[D.kh, D.kw])
+  O[D.oh, D.ow] += cast(U, I[D.oh + D.kh, D.ow + D.kw]) * cast(U, K[D.kh, D.kw])
+
 
 @linalg_structured_op
 def conv_3d(
     I=TensorDef(T1, S.OD + S.KD, S.OH + S.KH, S.OW + S.KW),
     K=TensorDef(T2, S.KD, S.KH, S.KW),
     O=TensorDef(U, S.OD, S.OH, S.OW, output=True)):
   """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(ConvolutionOpInterface)
   domain(D.od, D.oh, D.ow, D.kd, D.kh, D.kw)
-  O[D.od, D.oh, D.ow] += cast(
-      U, I[D.od + D.kd, D.oh + D.kh, D.ow + D.kw]) * cast(U, K[D.kd, D.kh, D.kw])
+  O[D.od, D.oh,
+    D.ow] += cast(U, I[D.od + D.kd, D.oh + D.kh, D.ow + D.kw]) * cast(
+        U, K[D.kd, D.kh, D.kw])
+
 
 @linalg_structured_op
 def conv_1d_nwc_wcf(
     I=TensorDef(T1, S.N, S.OW * S.SW + S.KW * S.DW, S.C),
     K=TensorDef(T2, S.KW, S.C, S.F),
     O=TensorDef(U, S.N, S.OW, S.F, output=True),
-    strides=AttributeDef(S.SW),
-    dilations=AttributeDef(S.DW)):
+    strides=IndexAttrDef(S.SW),
+    dilations=IndexAttrDef(S.DW)):
   """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(ConvolutionOpInterface)
   domain(D.n, D.ow, D.f, D.kw, D.c)
-  O[D.n, D.ow, D.f] += cast(
-      U, I[D.n, D.ow * S.SW + D.kw * S.DW, D.c
-           ]) * cast(U, K[D.kw, D.c, D.f])
+  O[D.n, D.ow, D.f] += cast(U, I[D.n, D.ow * S.SW + D.kw * S.DW, D.c]) * cast(
+      U, K[D.kw, D.c, D.f])
+
 
 @linalg_structured_op
 def conv_2d_nhwc_hwcf(
-    I=TensorDef(T1, S.N, S.OH * S.SH + S.KH * S.DH, S.OW * S.SW + S.KW * S.DW, S.C),
+    I=TensorDef(T1, S.N, S.OH * S.SH + S.KH * S.DH, S.OW * S.SW + S.KW * S.DW,
+                S.C),
     K=TensorDef(T2, S.KH, S.KW, S.C, S.F),
     O=TensorDef(U, S.N, S.OH, S.OW, S.F, output=True),
-    strides=AttributeDef(S.SH, S.SW),
-    dilations=AttributeDef(S.DH, S.DW)):
+    strides=IndexAttrDef(S.SH, S.SW),
+    dilations=IndexAttrDef(S.DH, S.DW)):
   """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(ConvolutionOpInterface)
   domain(D.n, D.oh, D.ow, D.f, D.kh, D.kw, D.c)
   O[D.n, D.oh, D.ow, D.f] += cast(
-      U, I[D.n, D.oh * S.SH + D.kh * S.DH, D.ow * S.SW + D.kw * S.DW, D.c
-           ]) * cast(U, K[D.kh, D.kw, D.c, D.f])
+      U, I[D.n, D.oh * S.SH + D.kh * S.DH, D.ow * S.SW + D.kw * S.DW,
+           D.c]) * cast(U, K[D.kh, D.kw, D.c, D.f])
+
 
 @linalg_structured_op
 def conv_2d_nhwc_hwcf_q(
-    I=TensorDef(T1, S.N, S.OH * S.SH + S.KH * S.DH, S.OW * S.SW + S.KW * S.DW, S.C),
+    I=TensorDef(T1, S.N, S.OH * S.SH + S.KH * S.DH, S.OW * S.SW + S.KW * S.DW,
+                S.C),
     K=TensorDef(T2, S.KH, S.KW, S.C, S.F),
     IZp=ScalarDef(I32),
     KZp=ScalarDef(I32),
     O=TensorDef(U, S.N, S.OH, S.OW, S.F, output=True),
-    strides=AttributeDef(S.SH, S.SW),
-    dilations=AttributeDef(S.DH, S.DW)):
+    strides=IndexAttrDef(S.SH, S.SW),
+    dilations=IndexAttrDef(S.DH, S.DW)):
   """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(ConvolutionOpInterface)
   domain(D.n, D.oh, D.ow, D.f, D.kh, D.kw, D.c)
-  O[D.n, D.oh, D.ow, D.f] += (cast(
-      U, I[D.n, D.oh * S.SH + D.kh * S.DH, D.ow * S.SW + D.kw * S.DW, D.c
-           ]) - cast(U, IZp)) * (cast(U, K[D.kh, D.kw, D.c, D.f]) - cast(U, KZp))
+  O[D.n, D.oh, D.ow,
+    D.f] += (cast(
+        U, I[D.n, D.oh * S.SH + D.kh * S.DH, D.ow * S.SW + D.kw * S.DW, D.c]) -
+             cast(U, IZp)) * (
+                 cast(U, K[D.kh, D.kw, D.c, D.f]) - cast(U, KZp))
+
 
 @linalg_structured_op
 def conv_2d_nchw_fchw(
-    I=TensorDef(T1, S.N, S.C, S.OH * S.SH + S.KH * S.DH, S.OW * S.SW + S.KW * S.DW),
+    I=TensorDef(T1, S.N, S.C, S.OH * S.SH + S.KH * S.DH,
+                S.OW * S.SW + S.KW * S.DW),
     K=TensorDef(T2, S.F, S.C, S.KH, S.KW),
     O=TensorDef(U, S.N, S.F, S.OH, S.OW, output=True),
-    strides=AttributeDef(S.SH, S.SW),
-    dilations=AttributeDef(S.DH, S.DW)):
+    strides=IndexAttrDef(S.SH, S.SW),
+    dilations=IndexAttrDef(S.DH, S.DW)):
   """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(ConvolutionOpInterface)
   domain(D.n, D.f, D.oh, D.ow, D.c, D.kh, D.kw)
   O[D.n, D.f, D.oh, D.ow] += cast(
-      U, I[D.n, D.c, D.oh * S.SH + D.kh * S.DH, D.ow * S.SW + D.kw * S.DW
-           ]) * cast(U, K[D.f, D.c, D.kh, D.kw])
+      U, I[D.n, D.c, D.oh * S.SH + D.kh * S.DH,
+           D.ow * S.SW + D.kw * S.DW]) * cast(U, K[D.f, D.c, D.kh, D.kw])
+
 
 @linalg_structured_op
 def conv_3d_ndhwc_dhwcf(
-    I=TensorDef(T1, S.N, S.OD * S.SD + S.KD * S.DD,
-                S.OH * S.SH + S.KH * S.DH, S.OW * S.SW + S.KW * S.DW, S.C),
+    I=TensorDef(T1, S.N, S.OD * S.SD + S.KD * S.DD, S.OH * S.SH + S.KH * S.DH,
+                S.OW * S.SW + S.KW * S.DW, S.C),
     K=TensorDef(T2, S.KD, S.KH, S.KW, S.C, S.F),
     O=TensorDef(U, S.N, S.OD, S.OH, S.OW, S.F, output=True),
-    strides=AttributeDef(S.SD, S.SH, S.SW),
-    dilations=AttributeDef(S.DD, S.DH, S.DW)):
+    strides=IndexAttrDef(S.SD, S.SH, S.SW),
+    dilations=IndexAttrDef(S.DD, S.DH, S.DW)):
   """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(ConvolutionOpInterface)
   domain(D.n, D.od, D.oh, D.ow, D.f, D.kd, D.kh, D.kw, D.c)
   O[D.n, D.od, D.oh, D.ow, D.f] += cast(
-      U, I[D.n, D.od * S.SD + D.kd * S.DD, D.oh * S.SH + D.kh * S.DH, D.ow * S.SW + D.kw * S.DW, D.c
-           ]) * cast(U, K[D.kd, D.kh, D.kw, D.c, D.f])
+      U, I[D.n, D.od * S.SD + D.kd * S.DD, D.oh * S.SH + D.kh * S.DH,
+           D.ow * S.SW + D.kw * S.DW, D.c]) * cast(
+               U, K[D.kd, D.kh, D.kw, D.c, D.f])
+
 
 @linalg_structured_op
 def depthwise_conv_1d_nwc_wc(
     I=TensorDef(T1, S.N, S.OW * S.SW + S.KW * S.DW, S.IC),
     K=TensorDef(T2, S.KW, S.IC),
     O=TensorDef(U, S.N, S.OW, S.IC, output=True),
-    strides=AttributeDef(S.SW),
-    dilations=AttributeDef(S.DW)):
+    strides=IndexAttrDef(S.SW),
+    dilations=IndexAttrDef(S.DW)):
   """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(ConvolutionOpInterface)
   domain(D.n, D.ow, D.ic, D.kw)
   O[D.n, D.ow, D.ic] += \
       cast(U, I[D.n, D.ow * S.SW + D.kw * S.DW, D.ic]) * \
       cast(U, K[D.kw, D.ic])
 
+
 @linalg_structured_op
 def depthwise_conv_2d_nhwc_hwc(
-    I=TensorDef(T1, S.N, S.OH * S.SH + S.KH * S.DH, S.OW * S.SW + S.KW * S.DW, S.IC),
+    I=TensorDef(T1, S.N, S.OH * S.SH + S.KH * S.DH, S.OW * S.SW + S.KW * S.DW,
+                S.IC),
     K=TensorDef(T2, S.KH, S.KW, S.IC),
     O=TensorDef(U, S.N, S.OH, S.OW, S.IC, output=True),
-    strides=AttributeDef(S.SH, S.SW),
-    dilations=AttributeDef(S.DH, S.DW)):
+    strides=IndexAttrDef(S.SH, S.SW),
+    dilations=IndexAttrDef(S.DH, S.DW)):
   """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(ConvolutionOpInterface)
   domain(D.n, D.oh, D.ow, D.ic, D.kh, D.kw)
   O[D.n, D.oh, D.ow, D.ic] += cast(
       U, I[D.n, D.oh * S.SH + D.kh * S.DH, D.ow * S.SW + D.kw * S.DW,
            D.ic]) * cast(U, K[D.kh, D.kw, D.ic])
 
+
 @linalg_structured_op
 def depthwise_conv_2d_nhwc_hwc_q(
-    I=TensorDef(T1, S.N, S.OH * S.SH + S.KH * S.DH, S.OW * S.SW + S.KW * S.DW, S.IC),
+    I=TensorDef(T1, S.N, S.OH * S.SH + S.KH * S.DH, S.OW * S.SW + S.KW * S.DW,
+                S.IC),
     K=TensorDef(T2, S.KH, S.KW, S.IC),
     IZp=ScalarDef(I32),
     KZp=ScalarDef(I32),
     O=TensorDef(U, S.N, S.OH, S.OW, S.IC, output=True),
-    strides=AttributeDef(S.SH, S.SW),
-    dilations=AttributeDef(S.DH, S.DW)):
+    strides=IndexAttrDef(S.SH, S.SW),
+    dilations=IndexAttrDef(S.DH, S.DW)):
   """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(ConvolutionOpInterface)
   domain(D.n, D.oh, D.ow, D.ic, D.kh, D.kw)
   O[D.n, D.oh, D.ow, D.ic] += (
       (cast(U, I[D.n, D.oh * S.SH + D.kh * S.DH, D.ow * S.SW + D.kw * S.DW,
                  D.ic]) - cast(U, IZp)) *
       (cast(U, K[D.kh, D.kw, D.ic]) - cast(U, KZp)))
 
+
 @linalg_structured_op
 def depthwise_conv_2d_nhwc_hwcm(
-    I=TensorDef(T1, S.N, S.OH * S.SH + S.KH * S.DH, S.OW * S.SW + S.KW * S.DW, S.IC),
+    I=TensorDef(T1, S.N, S.OH * S.SH + S.KH * S.DH, S.OW * S.SW + S.KW * S.DW,
+                S.IC),
     K=TensorDef(T2, S.KH, S.KW, S.IC, S.CM),
     O=TensorDef(U, S.N, S.OH, S.OW, S.IC, S.CM, output=True),
-    strides=AttributeDef(S.SH, S.SW),
-    dilations=AttributeDef(S.DH, S.DW)):
+    strides=IndexAttrDef(S.SH, S.SW),
+    dilations=IndexAttrDef(S.DH, S.DW)):
   """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(ConvolutionOpInterface)
   domain(D.n, D.oh, D.ow, D.ic, D.cm, D.kh, D.kw)
   O[D.n, D.oh, D.ow, D.ic, D.cm] += cast(
       U, I[D.n, D.oh * S.SH + D.kh * S.DH, D.ow * S.SW + D.kw * S.DW,
            D.ic]) * cast(U, K[D.kh, D.kw, D.ic, D.cm])
 
+
 @linalg_structured_op
 def depthwise_conv_2d_nhwc_hwcm_q(
-    I=TensorDef(T1, S.N, S.OH * S.SH + S.KH * S.DH, S.OW * S.SW + S.KW * S.DW, S.IC),
+    I=TensorDef(T1, S.N, S.OH * S.SH + S.KH * S.DH, S.OW * S.SW + S.KW * S.DW,
+                S.IC),
     K=TensorDef(T2, S.KH, S.KW, S.IC, S.CM),
     IZp=ScalarDef(I32),
     KZp=ScalarDef(I32),
     O=TensorDef(U, S.N, S.OH, S.OW, S.IC, S.CM, output=True),
-    strides=AttributeDef(S.SH, S.SW),
-    dilations=AttributeDef(S.DH, S.DW)):
+    strides=IndexAttrDef(S.SH, S.SW),
+    dilations=IndexAttrDef(S.DH, S.DW)):
   """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(ConvolutionOpInterface)
   domain(D.n, D.oh, D.ow, D.ic, D.cm, D.kh, D.kw)
   O[D.n, D.oh, D.ow, D.ic, D.cm] += (
       (cast(U, I[D.n, D.oh * S.SH + D.kh * S.DH, D.ow * S.SW + D.kw * S.DW,
                  D.ic]) - cast(U, IZp)) *
       (cast(U, K[D.kh, D.kw, D.ic, D.cm]) - cast(U, KZp)))
 
 
 @linalg_structured_op
 def pooling_nhwc_sum(
-    I=TensorDef(T1, S.N, S.OH * S.SH + S.KH * S.DH, S.OW * S.SW + S.KW * S.DW, S.C),
+    I=TensorDef(T1, S.N, S.OH * S.SH + S.KH * S.DH, S.OW * S.SW + S.KW * S.DW,
+                S.C),
     K=TensorDef(T2, S.KH, S.KW, index_dims=[D.kh, D.kw]),
     O=TensorDef(U, S.N, S.OH, S.OW, S.C, output=True),
-    strides=AttributeDef(S.SH, S.SW),
-    dilations=AttributeDef(S.DH, S.DW)):
+    strides=IndexAttrDef(S.SH, S.SW),
+    dilations=IndexAttrDef(S.DH, S.DW)):
   """Performs sum pooling.
 
   Numeric casting is performed on the input operand, promoting it to the same
   data type as the accumulator/output.
   """
   implements(ConvolutionOpInterface)
   domain(D.n, D.oh, D.ow, D.kh, D.kw, D.c)
   O[D.n, D.oh, D.ow, D.c] += cast(
       U, I[D.n, D.oh * S.SH + D.kh * S.DH, D.ow * S.SW + D.kw * S.DW, D.c])
 
 
 @linalg_structured_op
 def pooling_nhwc_max(
-    I=TensorDef(T1, S.N, S.OH * S.SH + S.KH * S.DH, S.OW * S.SW + S.KW * S.DW, S.C),
+    I=TensorDef(T1, S.N, S.OH * S.SH + S.KH * S.DH, S.OW * S.SW + S.KW * S.DW,
+                S.C),
     K=TensorDef(T2, S.KH, S.KW, index_dims=[D.kh, D.kw]),
     O=TensorDef(U, S.N, S.OH, S.OW, S.C, output=True),
-    strides=AttributeDef(S.SH, S.SW),
-    dilations=AttributeDef(S.DH, S.DW)):
+    strides=IndexAttrDef(S.SH, S.SW),
+    dilations=IndexAttrDef(S.DH, S.DW)):
   """Performs max pooling.
 
   Numeric casting is performed on the input operand, promoting it to the same
   data type as the accumulator/output.
   """
   implements(ConvolutionOpInterface)
   domain(D.n, D.oh, D.ow, D.kh, D.kw, D.c)
   O[D.n, D.oh, D.ow, D.c] = ReduceFn.max(D.kh, D.kw)(
       cast(U, I[D.n, D.oh * S.SH + D.kh * S.DH, D.ow * S.SW + D.kw * S.DW,
                 D.c]))
 
+
 @linalg_structured_op
 def pooling_nhwc_max_unsigned(
-    I=TensorDef(T1, S.N, S.OH * S.SH + S.KH * S.DH, S.OW * S.SW + S.KW * S.DW, S.C),
+    I=TensorDef(T1, S.N, S.OH * S.SH + S.KH * S.DH, S.OW * S.SW + S.KW * S.DW,
+                S.C),
     K=TensorDef(T2, S.KH, S.KW, index_dims=[D.kh, D.kw]),
     O=TensorDef(U, S.N, S.OH, S.OW, S.C, output=True),
-    strides=AttributeDef(S.SH, S.SW),
-    dilations=AttributeDef(S.DH, S.DW)):
+    strides=IndexAttrDef(S.SH, S.SW),
+    dilations=IndexAttrDef(S.DH, S.DW)):
   """Performs unsigned max pooling.
 
   Numeric casting is performed on the input operand, promoting it to the same
   data type as the accumulator/output.
   """
   implements(ConvolutionOpInterface)
   domain(D.n, D.oh, D.ow, D.kh, D.kw, D.c)
   O[D.n, D.oh, D.ow, D.c] = ReduceFn.max_unsigned(D.kh, D.kw)(
       cast_unsigned(
           U, I[D.n, D.oh * S.SH + D.kh * S.DH, D.ow * S.SW + D.kw * S.DW, D.c]))
 
+
 @linalg_structured_op
 def pooling_nchw_max(
-    I=TensorDef(T1, S.N, S.C, S.OH * S.SH + S.KH * S.DH, S.OW * S.SW + S.KW * S.DW),
+    I=TensorDef(T1, S.N, S.C, S.OH * S.SH + S.KH * S.DH,
+                S.OW * S.SW + S.KW * S.DW),
     K=TensorDef(T2, S.KH, S.KW, index_dims=[D.kh, D.kw]),
     O=TensorDef(U, S.N, S.C, S.OH, S.OW, output=True),
-    strides=AttributeDef(S.SH, S.SW),
-    dilations=AttributeDef(S.DH, S.DW)):
+    strides=IndexAttrDef(S.SH, S.SW),
+    dilations=IndexAttrDef(S.DH, S.DW)):
   """Performs max pooling.
 
   Numeric casting is performed on the input operand, promoting it to the same
   data type as the accumulator/output.
   """
   implements(ConvolutionOpInterface)
   domain(D.n, D.c, D.oh, D.ow, D.kh, D.kw)
   O[D.n, D.c, D.oh, D.ow] = ReduceFn.max(D.kh, D.kw)(
-      cast(U, I[D.n, D.c, D.oh * S.SH + D.kh * S.DH, D.ow * S.SW + D.kw * S.DW,
-                ]))
+      cast(U, I[D.n, D.c, D.oh * S.SH + D.kh * S.DH,
+                D.ow * S.SW + D.kw * S.DW,]))
+
 
 @linalg_structured_op
 def pooling_nhwc_min(
-    I=TensorDef(T1, S.N, S.OH * S.SH + S.KH * S.DH, S.OW * S.SW + S.KW * S.DW, S.C),
+    I=TensorDef(T1, S.N, S.OH * S.SH + S.KH * S.DH, S.OW * S.SW + S.KW * S.DW,
+                S.C),
     K=TensorDef(T2, S.KH, S.KW, index_dims=[D.kh, D.kw]),
     O=TensorDef(U, S.N, S.OH, S.OW, S.C, output=True),
-    strides=AttributeDef(S.SH, S.SW),
-    dilations=AttributeDef(S.DH, S.DW)):
+    strides=IndexAttrDef(S.SH, S.SW),
+    dilations=IndexAttrDef(S.DH, S.DW)):
   """Performs min pooling.
 
   Numeric casting is performed on the input operand, promoting it to the same
   data type as the accumulator/output.
   """
   implements(ConvolutionOpInterface)
   domain(D.n, D.oh, D.ow, D.kh, D.kw, D.c)
   O[D.n, D.oh, D.ow, D.c] = ReduceFn.min(D.kh, D.kw)(
       cast(U, I[D.n, D.oh * S.SH + D.kh * S.DH, D.ow * S.SW + D.kw * S.DW,
                 D.c]))
 
+
 @linalg_structured_op
 def pooling_nhwc_min_unsigned(
-    I=TensorDef(T1, S.N, S.OH * S.SH + S.KH * S.DH, S.OW * S.SW + S.KW * S.DW, S.C),
+    I=TensorDef(T1, S.N, S.OH * S.SH + S.KH * S.DH, S.OW * S.SW + S.KW * S.DW,
+                S.C),
     K=TensorDef(T2, S.KH, S.KW, index_dims=[D.kh, D.kw]),
     O=TensorDef(U, S.N, S.OH, S.OW, S.C, output=True),
-    strides=AttributeDef(S.SH, S.SW),
-    dilations=AttributeDef(S.DH, S.DW)):
+    strides=IndexAttrDef(S.SH, S.SW),
+    dilations=IndexAttrDef(S.DH, S.DW)):
   """Performs unsigned min pooling.
 
   Numeric casting is performed on the input operand, promoting it to the same
   data type as the accumulator/output.
   """
   implements(ConvolutionOpInterface)
   domain(D.n, D.oh, D.ow, D.kh, D.kw, D.c)
   O[D.n, D.oh, D.ow, D.c] = ReduceFn.min_unsigned(D.kh, D.kw)(
       cast_unsigned(
           U, I[D.n, D.oh * S.SH + D.kh * S.DH, D.ow * S.SW + D.kw * S.DW, D.c]))
 
+
 @linalg_structured_op
 def pooling_ndhwc_sum(
     I=TensorDef(T1, S.N, S.OD * S.SD + S.KD * S.DD, S.OH * S.SH + S.KH * S.DH,
                 S.OW * S.SW + S.KW * S.DW, S.C),
     K=TensorDef(T2, S.KD, S.KH, S.KW, index_dims=[D.kd, D.kh, D.kw]),
     O=TensorDef(U, S.N, S.OD, S.OH, S.OW, S.C, output=True),
-    strides=AttributeDef(S.SD, S.SH, S.SW),
-    dilations=AttributeDef(S.DD, S.DH, S.DW)):
+    strides=IndexAttrDef(S.SD, S.SH, S.SW),
+    dilations=IndexAttrDef(S.DD, S.DH, S.DW)):
   """Performs 3D sum pooling.
 
   Numeric casting is performed on the input operand, promoting it to the same
   data type as the accumulator/output.
   """
   implements(ConvolutionOpInterface)
   domain(D.n, D.od, D.oh, D.ow, D.kd, D.kh, D.kw, D.c)
   O[D.n, D.od, D.oh, D.ow, D.c] += cast(
       U, I[D.n, D.od * S.SD + D.kd * S.DD, D.oh * S.SH + D.kh * S.DH,
            D.ow * S.SW + D.kw * S.DW, D.c])
 
 
 @linalg_structured_op
 def pooling_ndhwc_max(
     I=TensorDef(T1, S.N, S.OD * S.SD + S.KD * S.DD, S.OH * S.SH + S.KH * S.DH,
                 S.OW * S.SW + S.KW * S.DW, S.C),
     K=TensorDef(T2, S.KD, S.KH, S.KW, index_dims=[D.kd, D.kh, D.kw]),
     O=TensorDef(U, S.N, S.OD, S.OH, S.OW, S.C, output=True),
-    strides=AttributeDef(S.SD, S.SH, S.SW),
-    dilations=AttributeDef(S.DD, S.DH, S.DW)):
+    strides=IndexAttrDef(S.SD, S.SH, S.SW),
+    dilations=IndexAttrDef(S.DD, S.DH, S.DW)):
   """Performs 3D max pooling.
 
   Numeric casting is performed on the input operand, promoting it to the same
   data type as the accumulator/output.
   """
   implements(ConvolutionOpInterface)
   domain(D.n, D.od, D.oh, D.ow, D.kd, D.kh, D.kw, D.c)
   O[D.n, D.od, D.oh, D.ow, D.c] = ReduceFn.max(D.kd, D.kh, D.kw)(
       cast(
           U, I[D.n, D.od * S.SD + D.kd * S.DD, D.oh * S.SH + D.kh * S.DH,
                D.ow * S.SW + D.kw * S.DW, D.c]))
 
 
 @linalg_structured_op
 def pooling_ndhwc_min(
     I=TensorDef(T1, S.N, S.OD * S.SD + S.KD * S.DD, S.OH * S.SH + S.KH * S.DH,
                 S.OW * S.SW + S.KW * S.DW, S.C),
     K=TensorDef(T2, S.KD, S.KH, S.KW, index_dims=[D.kd, D.kh, D.kw]),
     O=TensorDef(U, S.N, S.OD, S.OH, S.OW, S.C, output=True),
-    strides=AttributeDef(S.SD, S.SH, S.SW),
-    dilations=AttributeDef(S.DD, S.DH, S.DW)):
+    strides=IndexAttrDef(S.SD, S.SH, S.SW),
+    dilations=IndexAttrDef(S.DD, S.DH, S.DW)):
   """Performs 3D min pooling.
 
   Numeric casting is performed on the input operand, promoting it to the same
   data type as the accumulator/output.
   """
   implements(ConvolutionOpInterface)
   domain(D.n, D.od, D.oh, D.ow, D.kd, D.kh, D.kw, D.c)
   O[D.n, D.od, D.oh, D.ow, D.c] = ReduceFn.min(D.kd, D.kh, D.kw)(

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

 #   IMPL-DAG:  Value [[VAL2:[a-z0-9]+]] = helper.index(1);
 #   IMPL-DAG:  Value [[VAL3:[a-z0-9]+]] = helper.cast(block.getArgument(0).getType(), [[VAL2]], true);
 #   IMPL-DAG:  Value [[VAL4:[a-z0-9]+]] = helper.applyfn__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=AttributeDef(S.SM, S.SN)):
+#           strides=IndexAttrDef(S.SM, S.SN)):
 #   """Title.
 
 #   Detailed description.
 #   """
 #   O[D.m, D.n] = I[D.n * S.SM, D.m * S.SN]
 
 --- !LinalgOpConfig
 metadata: !LinalgOpMetadata

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

 # CHECK:     name: strides
 # CHECK:     usage: IndexAttribute
 # CHECK:     type_var: I64
 # CHECK:     attribute_map: affine_map<()[s0, s1, s2, s3, s4, s5] -> (s4, s5)>
 @linalg_structured_op
 def strided_copy(
     I=TensorDef(T, S.IH, S.IW),
     O=TensorDef(T, S.OH, S.OW, output=True),
-    strides=AttributeDef(S.SH, S.SW)):
+    strides=IndexAttrDef(S.SH, S.SW)):
   O[D.oh, D.ow] = I[D.oh * S.SH, D.ow * S.SW]

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

 T2 = TV.T2
 
 
 @linalg_structured_op
 def conv_poly(
     I=TensorDef(T1, S.N, S.IH, S.IW, S.C),
     K=TensorDef(T2, S.KH, S.KW, S.C),
     O=TensorDef(U, S.N, S.OH, S.OW, S.C, output=True),
-    strides=AttributeDef(S.SH, S.SW),
-    dilations=AttributeDef(S.DH, S.DW)):
+    strides=IndexAttrDef(S.SH, S.SW),
+    dilations=IndexAttrDef(S.DH, S.DW)):
   domain(D.n, D.oh, D.ow, D.kh, D.kw, D.c)
   O[D.n, D.oh, D.ow, D.c] += cast(
       U, I[D.n, D.oh * S.SH + D.kh * S.DH, D.ow * S.SW + D.kw * S.DW,
            D.c]) * cast(U, K[D.kh, D.kw, D.c])
 
 
 with Context() as ctx, Location.unknown():
   module = Module.create()

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

 T2 = TV.T2
 
 
 @linalg_structured_op
 def pooling_max_poly(
     I=TensorDef(T1, S.N, S.H, S.W, S.C),
     K=TensorDef(T2, S.KH, S.KW, index_dims=[D.kh, D.kw]),
     O=TensorDef(U, S.N, S.OH, S.OW, S.C, output=True),
-    strides=AttributeDef(S.SH, S.SW),
-    dilations=AttributeDef(S.DH, S.DW)):
+    strides=IndexAttrDef(S.SH, S.SW),
+    dilations=IndexAttrDef(S.DH, S.DW)):
   domain(D.n, D.oh, D.ow, D.kh, D.kw, D.c)
   O[D.n, D.oh, D.ow, D.c] = ReduceFn.max(D.kh, D.kw)(
       cast(U, I[D.n, D.oh * S.SH + D.kh * S.DH, D.ow * S.SW + D.kw * S.DW,
                 D.c]))
 
 
 @linalg_structured_op
 def pooling_max_unsigned_poly(
     I=TensorDef(T1, S.N, S.H, S.W, S.C),
     K=TensorDef(T2, S.KH, S.KW, index_dims=[D.kh, D.kw]),
     O=TensorDef(U, S.N, S.OH, S.OW, S.C, output=True),
-    strides=AttributeDef(S.SH, S.SW),
-    dilations=AttributeDef(S.DH, S.DW)):
+    strides=IndexAttrDef(S.SH, S.SW),
+    dilations=IndexAttrDef(S.DH, S.DW)):
   domain(D.n, D.oh, D.ow, D.kh, D.kw, D.c)
   O[D.n, D.oh, D.ow, D.c] = ReduceFn.max_unsigned(D.kh, D.kw)(
       cast_unsigned(
           U, I[D.n, D.oh * S.SH + D.kh * S.DH, D.ow * S.SW + D.kw * S.DW, D.c]))
 
 
 @linalg_structured_op
 def pooling_min_poly(
     I=TensorDef(T1, S.N, S.H, S.W, S.C),
     K=TensorDef(T2, S.KH, S.KW, index_dims=[D.kh, D.kw]),
     O=TensorDef(U, S.N, S.OH, S.OW, S.C, output=True),
-    strides=AttributeDef(S.SH, S.SW),
-    dilations=AttributeDef(S.DH, S.DW)):
+    strides=IndexAttrDef(S.SH, S.SW),
+    dilations=IndexAttrDef(S.DH, S.DW)):
   domain(D.n, D.oh, D.ow, D.kh, D.kw, D.c)
   O[D.n, D.oh, D.ow, D.c] = ReduceFn.min(D.kh, D.kw)(
       cast(U, I[D.n, D.oh * S.SH + D.kh * S.DH, D.ow * S.SW + D.kw * S.DW,
                 D.c]))
 
 
 @linalg_structured_op
 def pooling_min_unsigned_poly(
     I=TensorDef(T1, S.N, S.H, S.W, S.C),
     K=TensorDef(T2, S.KH, S.KW, index_dims=[D.kh, D.kw]),
     O=TensorDef(U, S.N, S.OH, S.OW, S.C, output=True),
-    strides=AttributeDef(S.SH, S.SW),
-    dilations=AttributeDef(S.DH, S.DW)):
+    strides=IndexAttrDef(S.SH, S.SW),
+    dilations=IndexAttrDef(S.DH, S.DW)):
   domain(D.n, D.oh, D.ow, D.kh, D.kw, D.c)
   O[D.n, D.oh, D.ow, D.c] = ReduceFn.min_unsigned(D.kh, D.kw)(
       cast_unsigned(
           U, I[D.n, D.oh * S.SH + D.kh * S.DH, D.ow * S.SW + D.kw * S.DW, D.c]))
 
 
 with Context() as ctx, Location.unknown():
   module = Module.create()