[mlir][OpDSL] Consistently use the term op_def (NFC).

[gysit]

Summary

... and remove unused type aliases.

Depends On D119003

mlir/python/mlir/dialects/linalg/opdsl/dump_oplib.py

modules.append(
                                 package="mlir.dialects.linalg.opdsl"))
   for i, file_path in enumerate(args.file or []):
     modules.append(load_module_from_file(f"_mlir_eval_oplib{i}", file_path))
   for m in modules:
     for attr_name, value in m.__dict__.items():
       # TODO: This class layering is awkward.
       if isinstance(value, DefinedOpCallable):
         try:
-          linalg_config = LinalgOpConfig.from_linalg_op_def(value.model)
+          linalg_config = LinalgOpConfig.from_linalg_op_def(value.op_def)
         except Exception as e:
           raise ValueError(
-              f"Could not create LinalgOpConfig from {value.model}") from e
+              f"Could not create LinalgOpConfig from {value.op_def}") from e
         configs.extend(linalg_config)
 
   # Print.
   if args.format == "yaml":
     print(yaml_dump_all(configs))
   elif args.format == "repr":
     for config in configs:
       print(repr(config))
 
 
 if __name__ == "__main__":
   main(create_arg_parser().parse_args())

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

__all__ = [
     "AffineBuildState",
     "AffineExprDef",
     "D",
     "DimDef",
     "S",
     "SymbolDef",
 ]
 
-# Type aliases.
-SymbolPosMap = Dict[str, int]
-
 
 class AffineBuildState:
   """Internal state for the AffineExprDef._create impls.
 
   Note that a "local" AffineBuildState can be created relative to a "global"
   AffineBuildState. In that case, any affine expressions built will inherit
   symbol and dim bindings from the global state and will update both as new
   ones are discovered. This allows for building expressions across contexts

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

 from enum import Enum
 
 from ..... import ir as _ir
 from .affine import *
 from .scalar_expr import *
 from .types import *
 from .yaml_helper import *
 
-# Type aliases.
-AffineDimList = Dict[str, _ir.AffineExpr]
-
 
 class TensorExpression:
   """An expression that can appear on the RHS of a comprehension."""
 
   def to_scalar_expression(self) -> ScalarExpression:
     raise NotImplementedError()
 
   def visit_tensor_exprs(self, callback):

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

class LinalgOpConfig(YAMLObject):
   def to_yaml_custom_dict(self):
     self_dict = dict(metadata=self.metadata,)
     if self.structured_op:
       self_dict["structured_op"] = self.structured_op
     return self_dict
 
   @staticmethod
   def from_linalg_op_def(
-      tc_op_def: LinalgOpDef,
+      op_def: LinalgOpDef,
       context: Optional[_ir.Context] = None) -> Sequence["LinalgOpConfig"]:
     """Expands a LinalgOpDef into corresponding Linalg configured ops."""
     # TODO: Many LinalgOpDef patterns need to expand to multiple generics.
     assert len(
-        tc_op_def.comprehensions) == 1, "Only one comprehension supported"
+        op_def.comprehensions) == 1, "Only one comprehension supported"
     return [
         LinalgOpConfig(
-            tc_op_def.metadata,
+            op_def.metadata,
             structured_op=LinalgStructuredOpConfig(
-                tc_op_def.comprehensions[0], tc_op_def.domain,
-                tc_op_def.registered_operands.values(), context)),
+                op_def.comprehensions[0], op_def.domain,
+                op_def.registered_operands.values(), context)),
     ]
 
   def __repr__(self):
     return (f"LinalgOpConfig(metadata={self.metadata},\n"
             f"structured_op={self.structured_op})")

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

 
 _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):
+def bind_op_def(op_def: LinalgOpDef):
   if hasattr(_CONTEXT, "current_op_def"):
     raise ValueError("Cannot recursively define an operation")
-  _CONTEXT.current_op_def = model
+  _CONTEXT.current_op_def = op_def
   try:
-    yield model
+    yield op_def
   finally:
     del _CONTEXT.current_op_def
 
 
 def current_op_def() -> LinalgOpDef:
   try:
     return _CONTEXT.current_op_def
   except AttributeError:
elif isinstance(outs, ir.OpResultList):
     return outs
 
   return [_get_op_result_or_value(o) for o in outs]
 
 
 class DefinedOpCallable:
   """Callable that wraps any defined op function."""
 
-  def __init__(self, op_name: str, model: LinalgOpDef):
+  def __init__(self, op_name: str, op_def: LinalgOpDef):
     self.op_name = op_name
-    self.model = model
+    self.op_def = op_def
 
   def __call__(self, *ins: Union[ir.Operation, ir.OpView, ir.Value],
                outs: StructuredOpOuts, **kwargs):
     """Emits the corresponding op definition as IR.
 
     Most arguments are passed through to the underlying emitter. The following
     keyword argument is interpreted here:
       emit_generic: Emits a generic form as appropriate (default True). If
         False, a named form is emitted (which must have been built in to the
         compiler).
     """
     emit_generic = kwargs.pop("emit_generic", False)
     if not isinstance(emit_generic, bool):
       raise ValueError(f"The named argument 'emit_generic' needs to be "
                        f" of type bool but got {type(emit_generic)}")
 
     op_configs = LinalgOpConfig.from_linalg_op_def(
-        self.model, context=ir.Context.current)
+        self.op_def, context=ir.Context.current)
 
     if len(op_configs) != 1:
       # TODO: Support composite ops.
       raise NotImplementedError(
           f"Emission of composite linalg ops not supported: {op_configs}")
 
     ctx = ir.Context.current
     linalgDialect = ctx.get_dialect_descriptor("linalg")
     fully_qualified_name = "linalg." + self.op_name
     emit_generic = (
         emit_generic or not ctx.is_registered_operation(fully_qualified_name))
 
     op_config = op_configs[0]
     out_values = _prepare_structured_op_outs(outs)
     in_values = [_get_op_result_or_value(i) for i in ins]
     if op_config.structured_op:
       if emit_generic:
         return emit_generic_structured_op(
             op_config.structured_op, *in_values, outs=out_values, **kwargs)
       else:
         return emit_named_structured_op(
             op_config.structured_op,
             self.op_name,
-            self.model.metadata.cpp_class_name,
+            self.op_def.metadata.cpp_class_name,
             *in_values,
             outs=out_values,
             **kwargs)
 
     raise NotImplementedError(
         f"Emission of linalg op type not supported: {op_config}")
 
 
 def linalg_structured_op(dsl_func=None,
                          *,
                          op_name=None,
                          op_class_name=None) -> DefinedOpCallable:
   if dsl_func is None:
     # Curry the keyword args in for delayed application.
     return functools.partial(
         linalg_structured_op, op_name=op_name, op_class_name=op_class_name)
   # Determine default names by introspecting the function.
   if op_name is None:
     op_name = dsl_func.__name__
   if op_class_name is None:
     # Camel case it.
     op_class_name = f"{''.join(x.title() for x in op_name.split('_'))}Op"
 
-  tc_model = LinalgOpDef(
+  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)):
-      tc_model.add_operand(param_name, param_default.operand_def)
+      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)
 
-  # Invoke the DSL func to finish populating the model.
-  with bind_op_def(tc_model):
+  # Invoke the DSL func to finish populating the op definition.
+  with bind_op_def(op_def):
     dsl_func(*dsl_func_args)
 
   # TODO: The returned callable should be an IR emitter but that is not
   # upstreamed yet.
-  return DefinedOpCallable(op_name, tc_model)
+  return DefinedOpCallable(op_name, op_def)
 
 
 def implements(*interfaces: OpInterfaceDef):
   current_op_def().metadata.implements.extend(interfaces)
 
 
 def domain(*dimensions: DimDef):
   if current_op_def().domain:
     raise ValueError(f"Expected only one set of domain dimensions per operator")
   if any(not isinstance(dim, DimDef) for dim in dimensions):
     raise ValueError(f"Expected dimensions of type DimDef but got {dimensions}")
   current_op_def().domain.extend(dimensions)

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

 import numpy as np
 
 __all__ = [
     "emit_generic_structured_op",
     "emit_named_structured_op",
     "ValueList",
 ]
 
+# Type aliases.
 ValueList = Union[Sequence[Value], OpResultList]
 
 
 def isa(cls: Type, ty: Type):
   try:
     cls(ty)
     return True
   except ValueError: