[mlir][linalg][python] Adapt the OpDSL to use scalars.

[gysit]

Summary

The patch replaces the existing capture functionality by scalar operands that have been introduced by https://reviews.llvm.org/D104109. Scalar operands behave as tensor operands except for the fact that they are not indexed. As a result ScalarDefs can be accessed directly as no indexing expression is needed.

The patch only updates the OpDSL. The C++ side is updated by a follow up patch.

[gysit]

Rebase.

[stellaraccident]

This is good - thanks. I'm glad we were able to back away from the more complicated capture semantics and just directly represent scalars.

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

211

Type: "... an index determines their position [in] the operation's..."

[gysit]

Address comment.

mlir/python/mlir/dialects/linalg/opdsl/lang/comprehension.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
 """Model classes representing a tensor comprehension.
 
 These classes model the language more at an AST level as evaluated. Reasoning
 about it typically involves processing this form into config objects that
 represent actual op definitions (i.e. YAML).
 """
 
-from typing import Any, Dict, List, Optional, Sequence, Set, Tuple, Union
+from typing import Any, Dict, List, Optional, Sequence, Set, Tuple
 
 from mlir import ir as _ir
 
 from .affine import *
 from .scalar_expr import *
 from .types import *
 from .yaml_helper import *
 
def visit_affine_exprs(expr):
           ind.visit_affine_exprs(visit_dim_def)
       if isinstance(expr, ReduceApply):
         for ind in expr.reduce.reduce_dims:
           ind.visit_affine_exprs(visit_dim_def)
 
     self.visit_tensor_exprs(visit_affine_exprs)
     return results
 
-  def collect_uses(self, uses: Set["TensorUse"]):
+  def collect_tensor_uses(self, uses: Set["TensorUse"]):
     """Collects all TensorUses reachable through this expression."""
 
     def visit_tensor_use(expr):
       if isinstance(expr, TensorUse):
         uses.add(expr)
 
     self.visit_tensor_exprs(visit_tensor_use)
 
   def collect_indices(self, indices: Set["index"]):
     """Collects all index accesses reachable through this expression."""
 
     def visit_index(expr):
       if isinstance(expr, index):
         indices.add(expr)
 
     self.visit_tensor_exprs(visit_index)
 
-  def collect_captures(self, captures: Set["CaptureDef"]):
-    """Collects all CaptureDefs reachable through this expression."""
+  def collect_scalar_uses(self, uses: Set["ScalarDef"]):
+    """Collects all ScalarDefs reachable through this expression."""
 
-    def visit_capture_def(expr):
-      if isinstance(expr, CaptureDef):
-        captures.add(expr)
+    def visit_scalar_def(expr):
+      if isinstance(expr, ScalarDef):
+        uses.add(expr)
 
-    self.visit_tensor_exprs(visit_capture_def)
+    self.visit_tensor_exprs(visit_scalar_def)
 
   def __add__(self, rhs: "TensorExpression") -> "TensorExpression":
     return PrimFn.add(self, rhs)
 
   def __mul__(self, rhs) -> "TensorExpression":
     return PrimFn.mul(self, rhs)
 
   def __sub__(self, rhs) -> "TensorExpression":
class TensorUse(TensorExpression):
   Note that forming a comprehension via direct assignment is performed through
   __setitem__ on the TensorDef level. However, performing a reduction with
   compound ops (+=, *=, etc) is done by doing a:
     TensorDef.__getitem__
     TensorUse.__iadd__
     TensorDef.__setitem__
   """
 
-  def __init__(self, tensor_def: "TensorDef", indices: Sequence[AffineExprDef]):
-    self.tensor_def = tensor_def
+  def __init__(self, operand_def: "OperandDef",
+               indices: Sequence[AffineExprDef]):
+    self.operand_def = operand_def
     self.indices = tuple(indices)
 
   def to_scalar_expression(self) -> ScalarExpression:
-    assert self.tensor_def.tensor_name is not None
-    return ScalarArg(self.tensor_def.tensor_name).expr()
+    return ScalarArg(self.tensor_name).expr()
 
   @property
   def tensor_name(self) -> str:
-    n = self.tensor_def.tensor_name
-    assert n is not None, "TensorDef not attached"
-    return n
+    name = self.operand_def.name
+    assert name is not None, "TensorDef not attached"
+    return name
 
   def __iadd__(self, rhs: TensorExpression) -> TensorExpression:
     return ReduceFn.add(*self._compute_reduce_dims(rhs))(rhs)
 
   def _compute_reduce_dims(self, rhs: TensorExpression) -> Set[DimDef]:
     """For implicit reductions, computes default reduction dims.
 
     Assumes that the rhs is the expression being reduced and self is being
     reduced into. Any indices referenced on the rhs and not in self are
     considered reduction dims and will be ordered as encountered on the rhs.
     """
     rhs_dims = rhs._get_all_dim_defs()
     lhs_dims = self._get_all_dim_defs()
     return rhs_dims - lhs_dims
 
   def __repr__(self):
     return f"{self.tensor_name}[{', '.join([repr(i) for i in self.indices])}]"
 
 
-class TensorDef:
-  """Bookkeeping of a single registered tensor, held in dict by name."""
+class OperandDef:
+  """Definition of a Tensor or Scalar operand passed to an operation."""
 
-  def __init__(self,
-               type_var: TypeVar,
-               *shape: AffineExprDef,
-               indexing_map: Optional[_ir.AffineMap] = None,
-               output: bool = False):
+  def __init__(self, type_var: TypeVar, shape: Sequence[AffineExprDef],
+               scalar: bool, output: bool):
     if not isinstance(type_var, TypeVar):
-      raise ValueError(f"TensorDef requires a TypeVar. Got: {repr(type_var)}")
+      raise ValueError(f"OperandDef requires a TypeVar. Got: {repr(type_var)}")
     self.owner = None  # type: Optional["LinalgOpDef"]
     self.type_var = type_var
     self.shape = shape
-    self.indexing_map = indexing_map
+    self.scalar = scalar
     self.output = output
-    self.tensor_name = None  # type: Optional[str]
+    self.name = None  # type: Optional[str]
     self.registered_index = -1  # type: int
 
-  @property
-  def rank(self) -> int:
-    """The rank of the tensor."""
-    return len(self.shape)
-
-  def attach(self, index: int, tensor_name: str, owner: "LinalgOpDef"):
+  def attach(self, index: int, name: str, owner: "LinalgOpDef"):
     if self.owner:
-      raise ValueError(f"TensorDef already registered with op: {self}")
+      raise ValueError(f"OperandDef already registered with op: {self}")
     self.registered_index = index
-    self.tensor_name = tensor_name
+    self.name = name
     self.owner = owner
 
+  def __hash__(self):
+    return hash(id(self))
+
+  def __repr__(self):
+    output = "OUTPUT " if self.output else ""
+    scalar = "SCALAR " if self.scalar else ""
+    return (f"{self.name}:OperandDef({output}{scalar}"
+            f"{repr(self.type_var)}, shape={self.shape})")
+
+
+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.
+  """
+
+  def __init__(self,
+               type_var: TypeVar,
+               *shape: AffineExprDef,
+               output: bool = False):
+    self.operand_def = OperandDef(type_var, shape, False, output)
+
   def __getitem__(self, dims) -> TensorUse:
-    assert self.owner, "TensorDef is not attached to an op"
+    assert self.operand_def.owner, "TensorDef is not attached to an op"
     state = AffineBuildState(
-        global_state=self.owner._affine_state, allow_new_symbols=False)
+        global_state=self.operand_def.owner._affine_state,
+        allow_new_symbols=False)
     if not isinstance(dims, tuple):
       dims = (dims,)  # Handle single subscript case.
     # Special case: (None) is a 0d-scalar use.
     if dims == (None,):
       dims = ()
 
     exprs = []
     for expr_def in dims:
       if not isinstance(expr_def, AffineExprDef):
         raise KeyError(
             "A TensorDef can only be subscripted by a tuple of affine dims")
       exprs.append(expr_def)
-    return TensorUse(self, exprs)
+    return TensorUse(self.operand_def, exprs)
 
   def __setitem__(self, dims, value):
     """Creates a new 1:1 comprehension by binding this tensor to an expression.
 
     Note that due to the way assignment works in Python, we have to capture
     direct assignment as a setitem on the TensorDef.
     """
     if not isinstance(value, TensorExpression):
       raise ValueError(f"Only TensorExpressions can be assigned to TensorDefs. "
                        f"Got: {repr(value)}")
     use = self[dims]
     comp = Comprehension((use, value))
-    self.owner.comprehensions.append(comp)
+    self.operand_def.owner.comprehensions.append(comp)
 
-  def __hash__(self):
-    return hash(id(self))
 
-  def __repr__(self):
-    output = "OUTPUT " if self.output else ""
-    return (f"{self.tensor_name}:TensorDef({output}{repr(self.type_var)}, "
-            f"shape={self.shape})")
-
-
-class CaptureDef(TensorExpression):
-  """Defines an SSA value captured by the operation.
+class ScalarDef(TensorExpression):
+  """Scalar operand definition.
 
-  The captured SSA values are not indexed by the indexing_maps of the
-  structured op (as opposed to memrefs and tensors). A unique name
-  identifies the captures and an index determines their position the
-  operation's parameter list.
+  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
+  the operation's parameter list.
   """
 
   def __init__(self, type_var: TypeVar):
-    if not isinstance(type_var, TypeVar):
-      raise ValueError(f"CaptureDef requires a TypeVar. Got: {repr(type_var)}")
-    self.owner = None  # type: Optional["LinalgOpDef"]
-    self.type_var = type_var
-    self.capture_name = None  # type: Optional[str]
-    self.registered_index = -1  # type: int
+    self.operand_def = OperandDef(type_var, (), True, False)
 
-  def attach(self, index: int, capture_name: str, owner: "LinalgOpDef"):
-    if self.owner:
-      raise ValueError(f"CaptureDef already registered with op: {self}")
-    self.registered_index = index
-    self.capture_name = capture_name
-    self.owner = owner
+  @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:
-    return ScalarCapture(self.capture_name).expr()
-
-  def __repr__(self):
-    return (f"{self.capture_name}:CaptureDef({repr(self.type_var)})")
+    return ScalarArg(self.scalar_name).expr()
 
 
 class Comprehension:
   """Represents a single comprehension."""
 
   def __init__(self, *bindings: Tuple[TensorUse, TensorExpression]):
     self.definitions = list()  # List[TensorUse]
     self.values = list()  # List[TensorExpression]
class LinalgOpDef:
   """Definition of a linalg op."""
 
   def __init__(self,
                name: str,
                cpp_class_name: Optional[str] = None,
                doc: Optional[str] = None):
     self.metadata = OpMetadataDef(
         name=name, cpp_class_name=cpp_class_name, doc=doc)
-    self.registered_tensors = dict()  # type: Dict[str, TensorDef]
-    self.registered_captures = dict()  # type: Dict[str, CaptureDef]
+    self.registered_operands = dict()  # type: Dict[str, OperandDef]
     self.comprehensions = list()  # type: List[Comprehension]
     self._affine_state = AffineBuildState()
 
   @property
-  def inputs(self) -> Sequence[TensorDef]:
-    return [t for t in self.registered_tensors.values() if not t.output]
-
-  @property
-  def outputs(self) -> Sequence[TensorDef]:
-    return [t for t in self.registered_tensors.values() if t.output]
+  def outputs(self) -> Sequence[OperandDef]:
+    return [
+        operand for operand in self.registered_operands.values()
+        if operand.output
+    ]
 
-  def add_tensor(self, tensor_name: str, tensor: TensorDef):
-    """Registers a tensor."""
-    if tensor_name in self.registered_tensors:
-      raise ValueError(f"Tensor {tensor_name} is already registered "
-                       f"to {self.registered_tensors['tensor_name']}")
-    tensor.attach(len(self.registered_tensors), tensor_name, self)
-    self.registered_tensors[tensor_name] = tensor
-
-  def add_capture(self, capture_name: str, capture: CaptureDef):
-    """Registers a capture."""
-    if capture_name in self.registered_captures:
-      raise ValueError(f"Capture {capture_name} is already registered "
-                       f"to {self.registered_captures['capture_name']}")
-    capture.attach(len(self.registered_captures), capture_name, self)
-    self.registered_captures[capture_name] = capture
+  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']}")
+    if not operand.output and self.outputs:
+      raise ValueError(f"The operand {name} is an input registered after "
+                       f"the output {self.outputs[-1]}")
+    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, tensor in self.registered_tensors.items():
-      lines.append(f"  {tensor}")
-    for name, capture in self.registered_captures.items():
-      lines.append(f"  {capture}")
+    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

 
 from typing import Dict, Optional
 
 from mlir import ir as _ir
 
 from .comprehension import *
 from .yaml_helper import *
 
-__all__ = [
-    "LinalgStructuredOpConfig",
-    "LinalgOpConfig",
-    "TensorDefConfig",
-]
+__all__ = ["LinalgStructuredOpConfig", "LinalgOpConfig", "OperandDefConfig"]
 
 
 def _serialize_affine_map(affine_map: _ir.AffineMap) -> str:
   with affine_map.context:
     # Affine map printing/parsing is via an AffineMap attr.
     attr = _ir.AffineMapAttr.get(affine_map)
     return str(attr)
 
 
 class TensorUseConfig:
   """Wrapper around a TensorUse with additional context-bound state."""
 
   def __init__(self, tensor_use: TensorUse, indexing_map: _ir.AffineMap):
     self.tensor_use = tensor_use
     self.indexing_map = indexing_map
 
   def __repr__(self):
     return f"Use({self.tensor_use}, indexing_map={self.indexing_map})"
 
 
-class TensorDefConfig(YAMLObject):
-  """Wrapper around a TensorDef with additional context-bound state."""
-  yaml_tag = "LinalgTensorDef"
-
-  def __init__(self, tensor_def: TensorDef, shape_map: _ir.AffineMap):
-    self.tensor_def = tensor_def
-    self.shape_map = shape_map
+class OperandDefConfig(YAMLObject):
+  """Wrapper containing an operand definition with additional state."""
+  yaml_tag = "!LinalgOperandDefConfig"
+
+  def __init__(self,
+               operand_def: OperandDef,
+               shape_map: Optional[_ir.AffineMap] = None):
+    self.operand_def = operand_def
+    self.shape_map = shape_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
+
+  @property
   def usage(self) -> str:
-    if self.tensor_def.output:
+    if self.operand_def.output:
       return "output"
-    else:
-      return "input"
 
   def to_yaml_custom_dict(self):
-    return dict(
-        name=self.tensor_def.tensor_name,
-        usage=self.usage,
-        shape=_serialize_affine_map(self.shape_map),
-        element_type_var=self.tensor_def.type_var.name,
-    )
-
-  def __repr__(self):
-    return f"Def({self.tensor_def}, shape_map={self.shape_map}, indexing_map={self.indexing_map})"
-
-
-class CaptureDefConfig(YAMLObject):
-  """Wrapper around a CaptureDef."""
-  yaml_tag = "LinalgCaptureDef"
-
-  def __init__(self, capture_def: CaptureDef):
-    self.capture_def = capture_def
-
-  def to_yaml_custom_dict(self):
-    return dict(
-        name=self.capture_def.capture_name,
-        type_var=self.capture_def.type_var.name,
-    )
+    self_dict = dict(name=self.name)
+    self_dict["usage"] = self.usage
+    if not self.operand_def.scalar:
+      self_dict["shape"] = _serialize_affine_map(self.shape_map)
+    self_dict["type_var"] = self.type_var.name
+    return self_dict
 
   def __repr__(self):
-    return f"Def({self.capture_def})"
+    return (f"OperandDefConfig({self.operand_def}, "
+            f"shape_map={self.shape_map}, indexing_map={self.indexing_map})")
 
 
 class LinalgIndexingMapsConfig(YAMLObject):
   """Abstracts the style of indexing maps that the op exports.
 
   Presently only static (tied to the op name) indexing maps are supported. In
   the future, it is expected that we will have additional variants:
     - Dynamic based on attributes
class LinalgStructuredOpConfig(YAMLObject):
   yaml_tag = "!LinalgStructuredOpConfig"
 
   def __init__(self,
                comprehension: Comprehension,
                context: Optional[_ir.Context] = None):
     self.context = context if context is not None else _ir.Context()
     self.affine_state = AffineBuildState()
     self.writes = list()  # type: List[Tuple[TensorUse, TensorExpression]]
-    self.tensor_args = dict()  # type: Dict[TensorDef, TensorDefConfig]
-    self.capture_args = dict()  # type: Dict[CaptureDef, CaptureDefConfig]
+    self.operands = dict()  # type: Dict[OperandDef, OperandDefConfig]
     self.uses = dict()  # type: Dict[TensorUse, TensorUseConfig]
 
     # Compute the ordered set of writes and collect the tensor, capture, and
     # index uses.
-    collected_uses = set()
-    collected_captures = set()
+    collected_tensor_uses = set()
+    collected_scalar_uses = set()
     collected_indices = set()
     for write_use, read_use in zip(comprehension.definitions,
                                    comprehension.values):
       self.writes.append((write_use, read_use))
 
     for write_use, read_use in self.writes:
-      collected_uses.add(write_use)
-      read_use.collect_uses(collected_uses)
-      read_use.collect_captures(collected_captures)
+      collected_tensor_uses.add(write_use)
+      read_use.collect_tensor_uses(collected_tensor_uses)
+      read_use.collect_scalar_uses(collected_scalar_uses)
       read_use.collect_indices(collected_indices)
 
     # Need to add all definitions before uses, so process twice.
-    for use in collected_uses:
-      self.add_tensor_arg(use.tensor_def)
-    for capture in collected_captures:
-      self.add_capture_arg(capture)
-    for use in collected_uses:
-      self.add_use(use)
+    for use in collected_tensor_uses:
+      self.add_operand(use.operand_def)
+    for use in collected_scalar_uses:
+      self.add_operand(use.operand_def)
+    for use in collected_tensor_uses:
+      self.add_tensor_use(use)
 
     # Now normalize all defs and uses indexing maps now that full count of
     # dims and symbols are known.
     for cuse in self.uses.values():
       cuse.indexing_map = self._normalize_affine_map(cuse.indexing_map)
-    for cdef in self.tensor_args.values():
+    for cdef in self.operands.values():
+      if not cdef.operand_def.scalar:
         cdef.shape_map = self._normalize_affine_map(
             cdef.shape_map, with_dims=False)
 
     # Now for each write use, propagate the indexing maps from the use to the
     # tensor, ensuring that there are not conflicts.
     for write_use, _ in self.writes:
-      write_tensor_def = self.tensor_args[write_use.tensor_def]
-      if write_tensor_def.indexing_map:
+      write_tensor_config = self.operands[write_use.operand_def]
+      if write_tensor_config.indexing_map:
         raise ValueError(
-            f"Unexpected multi-write to a single tensor: {write_tensor_def}")
-      write_tensor_def.indexing_map = self.uses[write_use].indexing_map
+            f"Unexpected multi-write to a single tensor: {write_tensor_config}")
+      write_tensor_config.indexing_map = self.uses[write_use].indexing_map
 
     # For each read use, propagate the indexing maps from the use to the
     # tensor, ensuring that there are not conflicts.
     for _, read_expr in self.writes:
       read_uses = set()  # type: Set[TensorUse]
-      read_expr.collect_uses(read_uses)
+      read_expr.collect_tensor_uses(read_uses)
       for read_use in read_uses:
-        read_tensor_def = self.tensor_args[read_use.tensor_def]
-        if (read_tensor_def.indexing_map and
-            read_tensor_def.indexing_map != self.uses[read_use].indexing_map):
+        read_operand_config = self.operands[read_use.operand_def]
+        if (read_operand_config.indexing_map and
+            read_operand_config.indexing_map !=
+            self.uses[read_use].indexing_map):
           raise ValueError(
               f"Unexpected multi-read of a tensor with different accesses:"
-              f"{read_tensor_def} vs {read_use}")
-        read_tensor_def.indexing_map = self.uses[read_use].indexing_map
+              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.scalar:
+        operand_config.indexing_map = self._create_empty_affine_map()
 
     # Sanity check that all defs have an indexing map.
-    assert all(d.indexing_map for d in self.tensor_args.values()), (
-        f"Missing indexing map on TensorDef: {self.tensor_args}")
+    assert all(d.indexing_map for d in self.operands.values()), (
+        f"Missing indexing map on OperandConfigDef: {self.operands}")
 
     # 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 "
           f"dims. Got: {all_reduction_dims}")
     self.reduction_dims = next(iter(all_reduction_dims))
 
     # Check the index dimension exists and resolve
     for index in collected_indices:
       if index.dim_def.dimname not in self.affine_state.all_dims:
         raise ValueError(
             f"The dimension {index.dim.dimname} is not part of the iteration "
             f"domain {self.affine_state.all_dims}")
       index.resolve_dimension_name(self.affine_state)
 
     # Generate the scalar assignments (used to build a body).
     self.assignments = [
         ScalarAssign(write_use.tensor_name, read_expr.to_scalar_expression())
         for write_use, read_expr in self.writes
     ]
 
   @property
-  def ordered_tensor_args(self) -> Sequence[TensorDefConfig]:
-    return sorted(
-        self.tensor_args.values(),
-        key=lambda tdc: tdc.tensor_def.registered_index)
-
-  @property
-  def ordered_tensor_uses(self) -> Sequence[TensorUseConfig]:
+  def ordered_operands(self) -> Sequence[OperandDefConfig]:
     return sorted(
-        self.uses.values(),
-        key=lambda tuc: tuc.tensor_use.tensor_def.registered_index)
-
-  @property
-  def ordered_capture_args(self) -> Sequence[CaptureDefConfig]:
-    return sorted(
-        self.capture_args.values(),
-        key=lambda cdc: cdc.capture_def.registered_index)
+        self.operands.values(),
+        key=lambda operand: operand.operand_def.registered_index)
 
   @property
   def ordered_dims(self) -> Sequence[Tuple[str, int]]:
     """Gets the ordered list of dim bindings (symbolic name, position).
 
     TODO: The original parser relies on parse ordering to arrive at the
     iterator types, but that ordering is not defined on the Python side, so
     this may be ambiguous.
     """
     return list(self.affine_state.all_dims.items())
 
   @property
   def indexing_maps(self) -> Sequence[_ir.AffineMap]:
-    return [use.indexing_map for use in self.ordered_tensor_uses]
+    return [d.indexing_map for d in self.ordered_operands]
 
   @property
   def iterator_types(self) -> Sequence[str]:
 
     def get_type(symbolic_name, position):
       for reduction_dim_expr in self.reduction_dims:
         if reduction_dim_expr.dimname == symbolic_name:
           return "reduction"
       return "parallel"
 
     return [get_type(*dim) for dim in self.ordered_dims]
 
-  def add_tensor_arg(self, tensor_def: TensorDef):
-    if tensor_def in self.tensor_args:
+  def add_operand(self, operand_def: OperandDef):
+    if operand_def in self.operands:
+      return
+    if operand_def.scalar:
+      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 tensor_def.shape:
+      for expr in operand_def.shape:
         exprs.append(expr.build(state=local_state))
       assert local_state.local_dim_count == 0
-      indexing_map = _ir.AffineMap.get(
+      shape_map = _ir.AffineMap.get(
           dim_count=0, symbol_count=local_state.symbol_count, exprs=exprs)
+      def_config = OperandDefConfig(operand_def, shape_map)
+      self.operands[operand_def] = def_config
 
-      def_config = TensorDefConfig(tensor_def, indexing_map)
-      self.tensor_args[tensor_def] = def_config
-
-  def add_use(self, tensor_use: TensorUse):
+  def add_tensor_use(self, tensor_use: TensorUse):
     if tensor_use in self.uses:
       return
     with self.context:
       local_state = AffineBuildState(
           global_state=self.affine_state, allow_new_symbols=False)
       exprs = []
       for expr in tensor_use.indices:
         exprs.append(expr.build(state=local_state))
       assert local_state.local_symbol_count == 0
       indexing_map = _ir.AffineMap.get(
           dim_count=local_state.dim_count,
           symbol_count=local_state.symbol_count,
           exprs=exprs)
 
       use_config = TensorUseConfig(tensor_use, indexing_map)
       self.uses[tensor_use] = use_config
 
-  def add_capture_arg(self, capture_def: CaptureDef):
-    if capture_def in self.capture_args:
-      return
-    def_config = CaptureDefConfig(capture_def)
-    self.capture_args[capture_def] = def_config
+  def _create_empty_affine_map(self) -> _ir.AffineMap:
+    """Create an affine map with an empty range."""
+    with self.context:
+      return _ir.AffineMap.get(
+          dim_count=self.affine_state.dim_count,
+          symbol_count=self.affine_state.symbol_count,
+          exprs=list())
 
   def _normalize_affine_map(self,
                             affine_map: _ir.AffineMap,
                             with_dims: bool = True) -> _ir.AffineMap:
     """Normalizes an indexing map to have the max known symbols and dims."""
     with self.context:
       return _ir.AffineMap.get(
           dim_count=self.affine_state.dim_count if with_dims else 0,
           symbol_count=self.affine_state.symbol_count,
           exprs=list(affine_map.results))
 
   def to_yaml_custom_dict(self):
-    self_dict = dict(args=self.ordered_tensor_args)
-    if self.ordered_capture_args:
-      self_dict["captures"] = self.ordered_capture_args
+    self_dict = dict(args=self.ordered_operands)
     # TODO: Refactor the hierarchy internally when supporting more
     # than static (preserving this serialized form).
     self_dict["indexing_maps"] = LinalgIndexingMapsConfig(
         static_indexing_maps=self.indexing_maps)
     self_dict["iterator_types"] = self.iterator_types
     self_dict["assignments"] = self.assignments
     return self_dict
 
   def __repr__(self):
     lines = [f"LinalgGenericOpConfig(reduction_dims={self.reduction_dims},"]
-    lines.append("tensor_args=[")
-    for def_config in self.ordered_tensor_args:
-      lines.append(f"  {repr(def_config)}")
-    lines.append("], capture_args=[")
-    for def_config in self.ordered_capture_args:
+    lines.append("operands=[")
+    for def_config in self.ordered_operands:
       lines.append(f"  {repr(def_config)}")
     lines.append("], indexing_maps=[")
     for m in self.indexing_maps:
       lines.append(f"  {repr(m)}")
     lines.append(f"], iterator_types=[")
     for t in self.iterator_types:
       lines.append(f"  {t}")
     lines.append("])")

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

def __call__(self, *args, emit_generic: bool = False, **kwargs):
     """Emits the corresponding op definition as IR.
 
     Most arguments are passed through to the underlying emitter. The following
     are 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).
     """
-    op_configs = LinalgOpConfig.from_linalg_op_def(self.model,
-                                                   context=ir.Context.current)
+    op_configs = LinalgOpConfig.from_linalg_op_def(
+        self.model, 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))
+    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]
     if op_config.structured_op:
       if emit_generic:
         return emit_generic_structured_op(op_config.structured_op, *args,
                                           **kwargs)
       else:
-        return emit_named_structured_op(
-          op_config.structured_op, self.op_name,
-          self.model.metadata.cpp_class_name, *args, **kwargs)
+        return emit_named_structured_op(op_config.structured_op, self.op_name,
+                                        self.model.metadata.cpp_class_name,
+                                        *args, **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(tc_def_op,
-                             op_name=op_name,
-                             op_class_name=op_class_name)
+    return functools.partial(
+        tc_def_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(name=op_name,
-                         cpp_class_name=op_class_name,
-                         doc=inspect.getdoc(dsl_func))
+  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):
-      tc_model.add_tensor(param_name, param_default)
-    elif isinstance(param_default, CaptureDef):
-      tc_model.add_capture(param_name, param_default)
+      tc_model.add_operand(param_name, param_default.operand_def)
+    elif isinstance(param_default, ScalarDef):
+      tc_model.add_operand(param_name, param_default.operand_def)
     else:
       raise ValueError(f"@tc_def_op function parameters must be defaulted as "
-                       f"TensorDef(...) or CaptureDef(...): Found {param_name}"
+                       f"TensorDef(...) or ScalarDef(...): Found {param_name}"
                        f": {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
   # upstreamed yet.
   return DefinedOpCallable(op_name, tc_model)
 
 
 def implements(*interfaces: OpInterfaceDef):
   current_op_def().metadata.implements.extend(interfaces)

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

def isa(cls: Type, ty: Type):
   try:
     cls(ty)
     return True
   except ValueError:
     return False
 
 
 def prepare_common_structured_op(op_config: LinalgStructuredOpConfig,
-                                 *ins: Value, outs: Sequence[Value],
-                                 captures: Sequence[Value]):
-  all_arg_defs = op_config.ordered_tensor_args
+                                 *ins: Value, outs: Sequence[Value]):
+  all_arg_defs = op_config.ordered_operands
   in_arg_defs = [arg for arg in all_arg_defs if arg.usage == "input"]
   out_arg_defs = [arg for arg in all_arg_defs if arg.usage == "output"]
-  capture_arg_defs = op_config.ordered_capture_args
 
   # Verify outs and captures are sequences.
   if not isinstance(outs, Sequence):
     raise ValueError(f"Expected named argument outs to have type Sequence "
                      f"but got {type(outs)}")
-  if not isinstance(captures, Sequence):
-    raise ValueError(f"Expected named argument captures to have type Sequence "
-                     f"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}")
-  if captures and len(captures) != len(capture_arg_defs):
-    raise ValueError(f"Expected {len(capture_arg_defs)} captures but got "
-                     f"{len(captures)} for {op_config}")
 
   outs, out_types = _infer_structured_outs(op_config, in_arg_defs, ins,
                                            out_arg_defs, outs)
 
   result_types = [t for t in out_types if isa(RankedTensorType, t)]
 
   # Initialize the type dictionary with the predefined types.
   type_mapping = dict()  # type: Dict[str, Type]
   type_mapping["F32"] = F32Type.get()
   type_mapping["F64"] = F64Type.get()
   type_mapping["I32"] = IntegerType.get_signless(32)
   type_mapping["I64"] = IntegerType.get_signless(64)
 
   # Extract type vars for input/output based types.
-  for arg_def, arg_element_type in zip(
-      in_arg_defs + out_arg_defs,
-      _get_shaped_element_types_from_values(*ins, *outs)):
-    _add_type_mapping(arg_def.tensor_def.type_var.name, arg_element_type,
-                      type_mapping)
-
-  # Extract type vars for captures and compute capture argument mapping.
-  capture_arg_mapping = dict()  # type: Dict[str, Value]
-  for arg_def, capture_value in zip(capture_arg_defs, captures):
-    _add_type_mapping(arg_def.capture_def.type_var.name, capture_value.type,
-                      type_mapping)
-    capture_arg_mapping[arg_def.capture_def.capture_name] = capture_value
+  block_arg_types = list()  # type: List[Type]
+  for arg_def, arg_element_type in zip(in_arg_defs + out_arg_defs,
+                                       _get_types_from_values(*ins, *outs)):
+    _add_type_mapping(arg_def, arg_element_type, type_mapping, block_arg_types)
 
   # Emit the generic op.
   # TODO: Support emission of pure memref form.
   indexing_maps_attr = ArrayAttr.get([
       AffineMapAttr.get(am)
       # TODO: linalg verification does not currently allow symbols.
       # Compress them for now.
       for am in AffineMap.compress_unused_symbols(op_config.indexing_maps,
                                                   Context.current)
   ])
   iterator_types_attr = ArrayAttr.get(
       [StringAttr.get(s) for s in op_config.iterator_types])
 
   return (all_arg_defs, in_arg_defs, out_arg_defs, outs, result_types,
-          type_mapping, capture_arg_mapping, indexing_maps_attr,
-          iterator_types_attr)
+          type_mapping, indexing_maps_attr, iterator_types_attr,
+          block_arg_types)
 
 
 def emit_generic_structured_op(op_config: LinalgStructuredOpConfig,
                                *ins: Value,
-                               outs: Sequence[Value] = (),
-                               captures: Sequence[Value] = ()):
+                               outs: Sequence[Value] = ()):
   all_arg_defs, in_arg_defs, out_arg_defs, outs, result_types, type_mapping, \
-  capture_arg_mapping, indexing_maps_attr, iterator_types_attr = \
-     prepare_common_structured_op(op_config, *ins, outs = outs,
-                                  captures=captures)
+  indexing_maps_attr, iterator_types_attr, block_arg_types = \
+     prepare_common_structured_op(op_config, *ins, outs = outs)
 
   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_tensor_def_names(*in_arg_defs, *out_arg_defs)
-  block_arg_types = _get_shaped_element_types_from_values(*ins, *outs)
+  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,
-                                capture_arg_mapping)
+    body_builder = _BodyBuilder(type_mapping, block_arg_mapping)
     for assignment in op_config.assignments:
       body_builder.assign(assignment)
-    body_builder.yield_outputs(*_get_tensor_def_names(*out_arg_defs))
+    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: Sequence[Value] = (),
-                             captures: Sequence[Value] = ()):
+                             outs: Sequence[Value] = ()):
   all_arg_defs, in_arg_defs, out_arg_defs, outs, result_types, type_mapping, \
-  capture_arg_mapping, indexing_maps_attr, iterator_types_attr = \
-     prepare_common_structured_op(op_config, *ins, outs = outs,
-                                  captures = captures)
+  indexing_maps_attr, iterator_types_attr, block_arg_types = \
+     prepare_common_structured_op(op_config, *ins, outs = outs)
 
   # 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}")
def emit_named_structured_op(op_config: LinalgStructuredOpConfig,
   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],
-               capture_arg_mapping: Dict[str, Value]):
+               block_arg_mapping: Dict[str, Value]):
     self.type_mapping = type_mapping
     self.block_arg_mapping = block_arg_mapping
-    self.capture_arg_mapping = capture_arg_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)
 
   def expression(self, expr: ScalarExpression) -> Value:
     if expr.scalar_arg:
       try:
         return self.block_arg_mapping[expr.scalar_arg.arg]
       except KeyError:
         raise ValueError(f"Argument {expr.scalar_arg.arg} is not bound for "
                          f"this structured op.")
-    elif expr.scalar_capture:
-      try:
-        return self.capture_arg_mapping[expr.scalar_capture.capture]
-      except KeyError:
-        raise ValueError(
-            f"Capture {expr.scalar_capture.capture} is not bound for "
-            f"this structured op.")
     elif expr.scalar_const:
       value_attr = Attribute.parse(expr.scalar_const.value)
       return std.ConstantOp(value_attr.type, value_attr).result
     elif expr.scalar_index:
       dim_attr = IntegerAttr.get(
           IntegerType.get_signless(64), expr.scalar_index.dim)
       return linalg.IndexOp(IndexType.get(), dim_attr).result
     elif expr.scalar_apply:
elif expr.symbolic_cast:
       return self.cast(expr.symbolic_cast.to_type.name, operand_value)
     raise NotImplementedError(f"Unimplemented scalar body expression: {expr}")
 
   def cast(self, type_var_name: str, operand: Value) -> Value:
     try:
       to_type = self.type_mapping[type_var_name]
     except KeyError:
       raise ValueError(f"Unbound type variable '{type_var_name}' ("
-                       f"expected one of {self.type_mappings.keys()}")
+                       f"expected one of {self.type_mapping.keys()}")
     if operand.type == to_type:
       return operand
     if _is_integer_type(to_type):
       return self._cast_to_integer(to_type, operand)
     elif _is_floating_point_type(to_type):
       return self._cast_to_floating_point(to_type, operand)
 
   def _cast_to_integer(self, to_type: Type, operand: Value) -> Value:
def _eval_mul(self, lhs: Value, rhs: Value) -> Value:
     if _is_floating_point_type(lhs.type):
       return std.MulFOp(lhs.type, lhs, rhs).result
     if _is_integer_type(lhs.type) or _is_index_type(lhs.type):
       return std.MulIOp(lhs.type, lhs, rhs).result
     raise NotImplementedError("Unsupported 'mul' operand: {lhs}")
 
 
 def _infer_structured_outs(op_config: LinalgStructuredOpConfig,
-                           in_arg_defs: Sequence[TensorDefConfig],
+                           in_arg_defs: Sequence[OperandDefConfig],
                            ins: Sequence[Value],
-                           out_arg_defs: Sequence[TensorDefConfig],
+                           out_arg_defs: Sequence[OperandDefConfig],
                            outs: Sequence[Value]):
   """Infers implicit outs and output types.
 
   Respects existing contents of outs if not empty.
 
   Returns:
     normalized outs, output types
   """
   # If outs were explicitly provided, we accept them verbatim.
   if outs:
     return outs, [out.type for out in outs]
 
   raise NotImplementedError(f"Output tensor inference not yet supported for "
                             "structured ops")
 
 
-def _get_shaped_element_types_from_values(*values: Value) -> Sequence[Type]:
+def _get_types_from_values(*values: Value) -> Sequence[Type]:
   types = []
   for v in values:
-    try:
-      t = ShapedType(v.type)
-    except Exception as e:
-      raise ValueError(f"Expected ShapedType but got {v}") from e
-    types.append(t.element_type)
+    types.append(v.type)
   return types
 
 
-def _get_tensor_def_names(
-    *tensor_def_configs: TensorDefConfig) -> Sequence[str]:
-  return [tdc.tensor_def.tensor_name for tdc in tensor_def_configs]
+def _get_operand_def_names(*operand_configs: OperandDefConfig) -> Sequence[str]:
+  return [odc.operand_def.name for odc in operand_configs]
 
 
-def _add_type_mapping(name: str, type: Type, type_mapping: Dict[str, Type]):
+def _add_type_mapping(operand_config: OperandDefConfig, operand_type: Type,
+                      type_mapping: Dict[str, Type],
+                      block_arg_types: Sequence[Type]):
+  element_or_self_type = operand_type
+  # Get the element type for tensor operands and the type itself for scalars.
+  if operand_config.operand_def.shape:
+    try:
+      element_or_self_type = ShapedType(operand_type).element_type
+    except Exception as e:
+      raise ValueError(f"Expected ShapedType but got {operand_type}") from e
+  name = operand_config.type_var.name
   if name in type_mapping:
-    if type_mapping[name] != type:
+    if type_mapping[name] != element_or_self_type:
       raise ValueError(f"Cannot overwrite type mapping {name} = "
-                       f"{type_mapping[name]} by type {type}")
-  type_mapping[name] = type
+                       f"{type_mapping[name]} by type {element_or_self_type}")
+  type_mapping[name] = element_or_self_type
+  block_arg_types.append(element_or_self_type)
 
 
 def _is_floating_point_type(t: Type) -> bool:
   # TODO: Create a FloatType in the Python API and implement the switch
   # there.
   return (F64Type.isinstance(t) or F32Type.isinstance(t) or
           F16Type.isinstance(t) or BF16Type.isinstance(t))
 

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

 
 from .yaml_helper import *
 from .types import *
 
 __all__ = [
     "ScalarAssign",
     "ScalarApplyFn",
     "ScalarArg",
-    "ScalarCapture",
     "ScalarConst",
     "ScalarIndex",
     "ScalarExpression",
     "ScalarSymbolicCast",
 ]
 
 
 class ScalarApplyFn:
class ScalarArg:
 
   def expr(self) -> "ScalarExpression":
     return ScalarExpression(scalar_arg=self)
 
   def __repr__(self):
     return f"(ScalarArg({self.arg})"
 
 
-class ScalarCapture:
-  """A type of ScalarExpression that references a named capture."""
-
-  def __init__(self, capture: str):
-    self.capture = capture
-
-  def expr(self) -> "ScalarExpression":
-    return ScalarExpression(scalar_capture=self)
-
-  def __repr__(self):
-    return f"(ScalarCapture({self.capture})"
-
-
 class ScalarConst:
   """A type of ScalarExpression representing a constant."""
 
   def __init__(self, value: str):
     self.value = value
 
   def expr(self) -> "ScalarExpression":
     return ScalarExpression(scalar_const=self)

 
 
 class ScalarExpression(YAMLObject):
   """An expression on scalar values.
 
   Can be one of:
     - ScalarApplyFn
     - ScalarArg
-    - ScalarCapture
     - ScalarConst
     - ScalarIndex
     - ScalarSymbolicCast
   """
   yaml_tag = "!ScalarExpression"
 
   def __init__(self,
                scalar_apply: Optional[ScalarApplyFn] = None,
                scalar_arg: Optional[ScalarArg] = None,
-               scalar_capture: Optional[ScalarCapture] = None,
                scalar_const: Optional[ScalarConst] = None,
                scalar_index: Optional[ScalarIndex] = None,
                symbolic_cast: Optional[ScalarSymbolicCast] = None):
-    if (bool(scalar_apply) + bool(scalar_arg) + bool(scalar_capture) +
-        bool(scalar_const) + bool(scalar_index) + bool(symbolic_cast)) != 1:
-      raise ValueError(
-          "One of 'scalar_apply', 'scalar_arg', 'scalar_capture', 'scalar_const', "
+    if (bool(scalar_apply) + bool(scalar_arg) + bool(scalar_const) +
+        bool(scalar_index) + bool(symbolic_cast)) != 1:
+      raise ValueError("One of 'scalar_apply', 'scalar_arg', 'scalar_const', "
                        "'scalar_index', 'symbolic_cast' must be specified")
     self.scalar_apply = scalar_apply
     self.scalar_arg = scalar_arg
-    self.scalar_capture = scalar_capture
     self.scalar_const = scalar_const
     self.scalar_index = scalar_index
     self.symbolic_cast = symbolic_cast
 
   def to_yaml_custom_dict(self):
     if self.scalar_apply:
       return dict(
           scalar_apply=dict(
               fn_name=self.scalar_apply.fn_name,
               operands=list(self.scalar_apply.operands),
           ))
     elif self.scalar_arg:
       return dict(scalar_arg=self.scalar_arg.arg)
-    elif self.scalar_capture:
-      return dict(scalar_capture=self.scalar_capture.capture)
     elif self.scalar_const:
       return dict(scalar_const=self.scalar_const.value)
     elif self.scalar_index:
       return dict(scalar_index=self.scalar_index.dim)
     elif self.symbolic_cast:
       # 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.

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

def dot(
   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 fill_rng_2d(O=TensorDef(T, S.M, S.N, output=True)):
+def fill_rng_2d(
+    min=ScalarDef(F64),
+    max=ScalarDef(F64),
+    seed=ScalarDef(I32),
+    O=TensorDef(T, S.M, S.N, output=True)):
   """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.
-
-  Note: The captures are hard-coded till there is capture support on the C++
-  side.
   """
-  min = cast(F64, const(-1000))
-  max = cast(F64, const(+1000))
-  seed = cast(I32, const(42))
   multiplier = cast(I32, const(1103515245))
   increment = cast(I32, const(12345))
   rand1 = (cast(I32, index(D.m)) + seed) * multiplier + increment
   rand2 = (cast(I32, index(D.n)) + rand1) * multiplier + increment
   inv_range = cast(F64, const(2.3283064e-10))
   offset = cast(F64, const(2147483647))
   scaling = (max - min) * inv_range
   O[D.m, D.n] = cast(T, (offset + cast(F64, rand2)) * scaling + min)

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:     shape: affine_map<()[s0, s1, s2] -> (s0, s2)>
-# CHECK:     element_type_var: T
+# CHECK:     type_var: T
 # CHECK:     name: B
 # CHECK:     usage: input
 # CHECK:     shape: affine_map<()[s0, s1, s2] -> (s2, s1)>
-# CHECK:     element_type_var: T
+# CHECK:     type_var: T
 # CHECK:     name: C
 # CHECK:     usage: output
 # CHECK:     shape: affine_map<()[s0, s1, s2] -> (s0, s1)>
-# CHECK:     element_type_var: U
+# CHECK:     type_var: U
 @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] += cast(U, A[D.m, D.k]) * cast(U, B[D.k, D.n])
 
 
 # CHECK: ---
 # CHECK-LABEL: fill
-# CHECK: captures:
-# CHECK: - !<LinalgCaptureDef>
-# CHECK:   name: value
-# CHECK:   type_var: T
+# CHECK: args:
+# CHECK:     name: value
+# CHECK:     usage: input
+# CHECK-NOT: shape:
+# CHECK:     type_var: T
 @linalg_structured_op
-def fill(O=TensorDef(T, S.M, S.K, output=True), value=CaptureDef(T)):
+def fill(value=ScalarDef(T), O=TensorDef(T, S.M, S.K, output=True)):
   O[D.m, D.n] = value

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

def indices(O=TensorDef(T, S.M, S.K, output=True)):
   O[D.m, D.n] = index(D.n) + index(D.m)
 
 
 # CHECK: ---
 # CHECK-LABEL: fill
 # CHECK: assignments:
 # CHECK:  -
 # CHECK:    arg: O
-# CHECK:      scalar_capture: value
+# CHECK:      scalar_arg: value
 @linalg_structured_op
-def fill(O=TensorDef(T, S.M, S.K, output=True), value=CaptureDef(T)):
+def fill(value=ScalarDef(T), O=TensorDef(T, S.M, S.K, output=True)):
   O[D.m, D.n] = value

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

def matmul_poly(
     A=TensorDef(TV.T1, S.M, S.K),
     B=TensorDef(TV.T2, S.K, S.N),
     C=TensorDef(U, S.M, S.N, output=True)):
   C[D.m, D.n] += cast(U, A[D.m, D.k]) * cast(U, B[D.k, D.n])
 
 
 @linalg_structured_op
 def fill_rng(
-    O=TensorDef(T, S.M, S.N, output=True),
-    min=CaptureDef(F64),
-    max=CaptureDef(F64),
-    seed=CaptureDef(I32)):
+    min=ScalarDef(F64),
+    max=ScalarDef(F64),
+    seed=ScalarDef(I32),
+    O=TensorDef(T, S.M, S.N, output=True)):
   multiplier = cast(I32, const(1103515245))
   increment = cast(I32, const(12345))
   rand1 = (cast(I32, index(D.m)) + seed) * multiplier + increment
   rand2 = (cast(I32, index(D.n)) + rand1) * multiplier + increment
   inv_range = cast(F64, const(2.3283064e-10))
   offset = cast(F64, const(2147483647))
   scaling = (max - min) * inv_range
   O[D.m, D.n] = cast(T, (offset + cast(F64, rand2)) * scaling + min)
with InsertionPoint(module.body):
     # CHECK-NEXT: -> tensor<4x8xf32>
     @builtin.FuncOp.from_py_func(
         RankedTensorType.get((4, 16), f64), RankedTensorType.get((16, 8), f64),
         RankedTensorType.get((4, 8), f32))
     def test_f64f64f32_matmul(lhs, rhs, init_result):
       return matmul_poly(lhs, rhs, outs=[init_result])
 
     # CHECK-LABEL: @test_fill_rng
-    # CHECK-SAME:  %{{.*}} tensor<4x16xi32>, %[[MIN:.+]]: f64, %[[MAX:.+]]: f64, %[[SEED:.+]]: i32
+    # CHECK:      ^{{.*}}(%[[MIN:.+]]: f64, %[[MAX:.+]]: f64, %[[SEED:.+]]: i32, %{{.*}}
     # CHECK-DAG:    %[[IDX0:.+]] = linalg.index 0 : index
     # CHECK-DAG:    %[[IDX1:.+]] = linalg.index 1 : index
     # CHECK-DAG:    %[[IDX0_CAST:.+]] = index_cast %[[IDX0]] : index to i32
     # CHECK-DAG:    %[[IDX1_CAST:.+]] = index_cast %[[IDX1]] : index to i32
     # CHECK-DAG:    %[[RND0:.+]] = addi %[[IDX0_CAST]], %[[SEED]] : i32
     # CHECK-DAG:    %[[CST0:.+]] = constant 1103515245 : i64
     # CHECK-DAG:    %[[CST0_CAST:.+]] = trunci %[[CST0]] : i64 to i32
     # CHECK-DAG:    %[[CST1:.+]] = constant 12345 : i64
     # CHECK-DAG:    %[[CST1_CAST:.+]] = trunci %[[CST1]] : i64 to i32
     # CHECK-DAG:    %[[RND1:.+]] = muli %[[RND0]], %[[CST0_CAST]] : i32
     # CHECK-DAG:    %[[RND2:.+]] = addi %[[RND1]], %[[CST1_CAST]] : i32
     # Skip random number computation for the second index.
     # CHECK-DAG:    %[[DIFF:.+]] = subf %[[MAX]], %[[MIN]] : f64
     # CHECK-DAG:    %[[CST3:.+]] = constant 2.3283063999999999E-10 : f64
     # CHECK-DAG:    %[[FACT:.+]] = mulf %[[DIFF]], %[[CST3]] : f64
     # CHECK-DAG:    %[[RND4:.+]] = mulf %{{.+}}, %[[FACT]] : f64
     # CHECK-DAG:    %[[RND5:.+]] = addf %[[RND4]], %[[MIN]] : f64
     # CHECK-DAG:    %{{.*}} = fptosi %[[RND5]] : f64 to i32
-    @builtin.FuncOp.from_py_func(
-        RankedTensorType.get((4, 16), i32), f64, f64, i32)
-    def test_fill_rng(init_result, min, max, seed):
-      return fill_rng(outs=[init_result], captures=[min, max, seed])
+    @builtin.FuncOp.from_py_func(f64, f64, i32,
+                                 RankedTensorType.get((4, 16), i32))
+    def test_fill_rng(min, max, seed, init_result):
+      return fill_rng(min, max, seed, outs=[init_result])
 
 
 print(module)

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

func @main() -> f32 attributes {llvm.emit_c_interface} {
   // 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} {
   %O = memref.alloc() : memref<4x16xi32>
+  %min = constant -1000.0 : f64
+  %max = constant 1000.0 : f64
+  %seed = constant 42 : i32
 
-  call @fill_on_buffers(%O) :
-    (memref<4x16xi32>) -> ()
+  call @fill_on_buffers(%min, %max, %seed, %O) :
+    (f64, f64, i32, memref<4x16xi32>) -> ()
 
   %c0 = constant 0 : index
   %0 = memref.load %O[%c0, %c0] : memref<4x16xi32>
 
   // TODO: FFI-based solution to allow testing and printing with python code.
   return %0 : i32
 }
 """
with Context() as ctx, Location.unknown():
 
     log("RESULT: ", res[0])
     # CHECK: RESULT: 32.0
 
 
 test_matmul_generic()
 
 
-def test_fill_builtin():
-  with Context() as ctx, Location.unknown():
-    module = Module.create()
-    f64 = F64Type.get()
-    i32 = IntegerType.get_signless(32)
-    with InsertionPoint(module.body):
-
-      @builtin.FuncOp.from_py_func(MemRefType.get((4, 16), i32))
-      def fill_on_buffers(out):
-        linalg.fill_rng_2d(outs=[out])
-
-    execution_engine = ExecutionEngine(transform(module, fill_boiler))
-
-    # TODO: FFI-based solution to allow testing and printing with python code.
-    # Prepare arguments: one result i32.
-    # Arguments must be passed as pointers.
-    c_int_p = ctypes.c_int * 1
-    res = c_int_p(-1)
-    execution_engine.invoke("main", res)
-
-    log("RESULT: ", res[0])
-    # CHECK: RESULT: -480
-
-
-test_fill_builtin()
-
-
 def test_fill_generic():
   with Context() as ctx, Location.unknown():
     module = Module.create()
     f64 = F64Type.get()
     i32 = IntegerType.get_signless(32)
     with InsertionPoint(module.body):
 
-      @builtin.FuncOp.from_py_func(MemRefType.get((4, 16), i32))
-      def fill_on_buffers(out):
-        linalg.fill_rng_2d(outs=[out])
+      @builtin.FuncOp.from_py_func(f64, f64, i32, MemRefType.get((4, 16), i32))
+      def fill_on_buffers(min, max, seed, out):
+        linalg.fill_rng_2d(min, max, seed, outs=[out], emit_generic=True)
 
     execution_engine = ExecutionEngine(transform(module, fill_boiler))
 
     # TODO: FFI-based solution to allow testing and printing with python code.
     # Prepare arguments: one result i32.
     # Arguments must be passed as pointers.
     c_int_p = ctypes.c_int * 1
     res = c_int_p(-1)
     execution_engine.invoke("main", res)
 
     log("RESULT: ", res[0])
     # CHECK: RESULT: -480
 
 
 test_fill_generic()