[mlir][Python][Linalg] Adding const, capture, and index support to the OpDSL.

[gysit]

Summary

The patch extends the OpDSL with support for:

• Constant values
• Capture scalar parameters
• Access the iteration indices using the index operation
• Provide predefined floating point and integer types.

Up to now the patch only supports emitting the new nodes. The C++/yaml path is not fully implemented (there is no implementation at all on the C++ side) and untested. The fill_rng_2d operation defined in emit_structured_generic.py makes use of the new DSL constructs.

[gysit]

The fill_rng_2d operation defined below illustrates the new features implemented in this revision:

@linalg_structured_op
def fill_rng_2d(A=TensorDef(T, S.M, S.N, output=True),
                min=CaptureDef(F64),
                max=CaptureDef(F64),
                seed=CaptureDef(I32)):
  multiplier = const(I32, 1103515245)
  increment = const(I32, 12345)
  temp1 = (cast(I32, index(D.m)) + seed) * multiplier + increment
  temp2 = (cast(I32, index(D.n)) + temp1) * multiplier + increment
  scaling = (max - min) * const(F64, 2.3283064e-10)
  A[D.m, D.n] = cast(T, cast(F64, temp2) * scaling + min)

Running /usr/bin/python3.9 /usr/local/google/home/gysit/Repos/llvm-project/mlir/test/Bindings/Python/dialects/linalg/opdsl/emit_structured_generic.py emits the following ir:

func @test_fill_rng_2d(%arg0: tensor<4x16xi32>, %arg1: f64, %arg2: f64, %arg3: i32) -> tensor<4x16xi32> {
  %0 = linalg.generic {indexing_maps = [#map3], iterator_types = ["parallel", "parallel"]} outs(%arg0 : tensor<4x16xi32>) {
  ^bb0(%arg4: i32):  // no predecessors
    %1 = linalg.index 1 : index
    %2 = index_cast %1 : index to i32
    %3 = linalg.index 0 : index
    %4 = index_cast %3 : index to i32
    %5 = addi %4, %arg3 : i32
    %c1103515245_i32 = constant 1103515245 : i32
    %6 = muli %5, %c1103515245_i32 : i32
    %c12345_i32 = constant 12345 : i32
    %7 = addi %6, %c12345_i32 : i32
    %8 = addi %2, %7 : i32
    %c1103515245_i32_0 = constant 1103515245 : i32
    %9 = muli %8, %c1103515245_i32_0 : i32
    %c12345_i32_1 = constant 12345 : i32
    %10 = addi %9, %c12345_i32_1 : i32
    %11 = sitofp %10 : i32 to f64
    %12 = subf %arg2, %arg1 : f64
    %cst = constant 2.3283063999999999E-10 : f64
    %13 = mulf %12, %cst : f64
    %14 = mulf %11, %13 : f64
    %15 = addf %14, %arg1 : f64
    %16 = fptosi %15 : f64 to i32
    linalg.yield %16 : i32
  } -> tensor<4x16xi32>
  return %0 : tensor<4x16xi32>
}

I expect some iterations will be needed to make the code stable and to add support on the C++ side.

[stellaraccident]

This looks good to me - will leave the rest of the review to Nicolas. We do have docs in the docs/ tree for this language, and it would be good to update them.

mlir/lib/Bindings/Python/mlir/dialects/linalg/opdsl/lang/comprehension.py

200

I'm trying to understand what a CaptureDef can be: I assume from the examples that it must be a primitive/scalar that is supplied as an SSA value. Can it be anything else (i.e. a whole-tensor value that does not participate in indexing/iteration)? Do you foresee this being different from captured attributes (which the old language supports)?

(you don't need to answer me inline if the questions have clear answers: just extend the docs, either here or in the language guide in the docs/ folder)

337

Can you say more about what the allowable type/value are and how this is emitted?

339

I think you want Any here. We likely need to further constrain it later but can let it be anything for now.

349

Somewhat obvious, but can you document more what the dimension index is resolved against?

mlir/lib/Bindings/Python/mlir/dialects/linalg/opdsl/lang/scalar_expr.py

25

Can we just name this ScalarCapture (not abbreviated) for consistency?

mlir/lib/Bindings/Python/mlir/dialects/linalg/opdsl/ops/core_named_ops.py

13 ↗

(On Diff #340827)

Mind submitting these changes NFC separately?

[gysit]

Address comments:

• update the documentation
• improve doc strings
• rename ScalarCap to ScalarCapture
• restrict constants to floats and ints
• use Any instead of object

[gysit]

Thanks for the detailed review comments!

mlir/lib/Bindings/Python/mlir/dialects/linalg/opdsl/lang/comprehension.py

200

I was not aware of the attribute functionality (@nicolasvasilache pointed me to some tc example) and yes there is indeed some overlap. The idea of CaptureDef is to add scalar parameters to the parameter list of the operation as it is done for some Linalg structured operations (e.g. for the fill operation). Capturing values directly would make it harder to isolate the operation / turn it into a library call.

I hope that explains the idea but I am open to other solutions!

[nicolasvasilache]

mlir/docs/Tools/LinalgOpDsl.md

131

do we need the "floating point or integer value" part ?
This should also work with vectors, right?

139

can these be generalized to all scalar and vector types ? (for a subsequent commit).

One of the ideas with such custom ops was to be able to just program with vectors if we wanted.

mlir/lib/Bindings/Python/mlir/dialects/linalg/opdsl/lang/comprehension.py

200

-> "Defines an SSA value captured by the operation" ?

202

-> "The captured SSA values are not indexed by the indexing_maps of the structured op (as opposed to memrefs and tensors)" ?

337

I'd just go for "scalar and vector constants".

mlir/lib/Bindings/Python/mlir/dialects/linalg/opdsl/lang/emitter.py

32

I'd just spell "captures" everywhere, "caps" makes me think of lower/upper.
Since this is python I wouldn't be super surprised to see caps changing code around :)

39

defensive type programming, how fancy ! :)

mlir/lib/Bindings/Python/mlir/dialects/linalg/opdsl/lang/types.py

73

Not for this CL but I am wondering whether we should have these at all (vs a mechanism that just allows building all scalar + vector types) ?

mlir/test/Bindings/Python/dialects/linalg/opdsl/emit_structured_generic.py

35

very nice looking!
Can we spell out the constant into a named lhs ? Is it an epsilon or something else?

[nicolasvasilache]

Nice progress, thanks @gysit !

[gysit]

Address comments:

• Rename caps to captures
• Update documentation
• Cleaner fill_rng definition

[gysit]

mlir/docs/Tools/LinalgOpDsl.md

131

At the moment we are limited to floats and integers since we check this on the python side and since we only have code to emit Integer and Float Attributes. Extending this is possible.

139

Yes this is possible but requires a slightly different design I believe. At the moment I32 etc are simply names that carry no further information. I think we should change that as you suggested to provide a mechanism to generate different vector types etc.

mlir/test/Bindings/Python/dialects/linalg/opdsl/emit_structured_generic.py

35

it is the inverse of the maximal random number / max_int_32

[gysit]

mlir/lib/Bindings/Python/mlir/dialects/linalg/opdsl/lang/comprehension.py

337

I will change this once we support vector constants and types.

[stellaraccident]

mlir/lib/Bindings/Python/mlir/dialects/linalg/opdsl/lang/types.py

73

Yes, I almost commented the same. In some prior versions of this code, I had a more evolved type system but it was incomplete and I stripped it in favor of adding complexity later. I'm fine with this is a stepping stone but it would be good to figure out to where.

mlir/docs/Tools/LinalgOpDsl.md

 expressions).
 * Usage (`output=True`).
 
 The docstring will be transferred to the op definition verbatim.
 
 Special identifying op interfaces can be declared for the op via
 `implements(interface1[, interface2...])`.
 
+## Parameters
+
+Structured operations can take two types of parameters namely input/output
+tensors and captures. Assignment expressions index the tensor parameters to
+access the individual elements, while captures are scalars that can be
+accessed directly.
+
+The following example demonstrates the use of the two parameter types:
+
+```python
+@linalg_structured_op
+def copy_and_scale(I=TensorDef(T, S.M, S.K),
+                   O=TensorDef(T, S.M, S.K, output=True),
+                   val=CaptureDef(T)):
+  """Scale the input by the captured value and store the result"""
+  O[D.m, D.n] = I[D.m, D.n] * val
+```
+
+The operation scales the input tensor `I` scales its elements by the value
+`val` and writes the result to the output tensor `out`. The capture `val` is
+bound to a `CaptureDef`, which specifies the type of the captured value. The
+tensors are bound to a `TensorDef` as demonstrated by the matmul example. All
+parameters appear in the parameter list of the operation:
+
+```python
+fill(in_tensor, outs=[out_tensor], caps=[captured_val])
+```
+
 ## Assignments
 
 The bulk of language consists of assignment expressions of the form above.
 The iteration dimension order is determined lexically based on the order
 encountered in the expression (following operator precedence if math operators
 are used). TODO: Introduce a directive to fix the dimension bindings.
 
 Reduction dimensions are inferred to be any dimensions on the RHS that are not

 Reduction functions can appear as the outer-most function on the RHS:
 
 * `ReduceFn.add` (also overloading the inplace `+=` on a LHS)
 * `ReduceFn.mul`
 * `ReduceFn.max`
 
 There are also special forms:
 
-* `cast(TypeVar, operand)`
+* `cast(TypeVar, operand)` casts the `operand` to the target type `TypeVar`.
+* `const(TypeVar, value)` returns a constant floating point or integer value
+of type `TypeVar`.
+* `index(dim)` returns the iteration index in the given dimension `dim`.
 
 ## Types
 
 All types in assignment expressions are late bound based on actual input
-and output types of constructed ops. Assignment expressions with no `cast`
-calls will generally require uniform types throughout and will fail to
-verify if violated. The presence of a `cast` allows for a limited form of
-numeric type conversion between element types that can be derived from inputs
-and outputs (and in the future, attributes). `cast` calls with a `TypeVar`
-first argument are emitted as `symbolic_cast` primitives in the YAML definition.
-
-Casting will perform `int<->float` type conversions and will perform any
-necessary extension or truncation within type family. Note that presently,
-any integer type is assumed to be signed for the purpose of determing how to
-extend or truncate. Supporting unsigned integer types is left for future work.
+and output types of constructed ops. An exception are the predefined types
+`I32`, `I64`, `F32`, and `F64`. These hardwired types enable intermediate
+computations with a type that is independent of the input and output types.
+For example, parts of floating point computation may require double precision
+arithmetic despite all inputs and outputs being single precision values.
+Assignment expressions with no `cast` calls will generally require uniform
+types throughout and will fail to verify if violated. The presence of a
+`cast` allows for a limited form of numeric type conversion between element
+types that can be derived from inputs and outputs (and in the future,
+attributes). `cast` calls with a `TypeVar` first argument are emitted as
+`symbolic_cast` primitives in the YAML definition.
+
+Casting will perform `int<->float` and `index->int` type conversions and will
+perform any necessary extension or truncation within type family. Note that
+presently, any integer type is assumed to be signed for the purpose of
+determining how to extend or truncate. Supporting unsigned integer types is
+left for future work.
 
 Not all functions are applicable for all numeric types, and on mismatch, op
 verification will fail.

mlir/lib/Bindings/Python/mlir/dialects/linalg/opdsl/lang/affine.py

def __repr__(self):
     return f"{self.ir_ctor.__name__}({repr(self.lhs)}, {repr(self.rhs)})"
 
 
 class DimDef(AffineExprDef):
   """Represents a named dimension.
 
   """
   ALL_DIMS = dict()  # type: Dict[str, "DimDef"]
-  dimname: str
 
   def __new__(cls, dimname: str):
     existing = cls.ALL_DIMS.get(dimname)
     if existing is not None:
       return existing
     new = super().__new__(cls)
     new.dimname = dimname
     cls.ALL_DIMS[dimname] = new
"""Represents a named symbol.
     Symbol(s1)
     >>> s2 = SymbolDef("s2")
     >>> s1 is s2
     False
     >>> s1 is SymbolDef("s1")
     True
   """
   ALL_SYMBOLS = dict()  # type: Dict[str, "SymbolDef"]
-  symname: str
 
   def __new__(cls, symname: str):
     existing = cls.ALL_SYMBOLS.get(symname)
     if existing is not None:
       return existing
     new = super().__new__(cls)
     new.symname = symname
     cls.ALL_SYMBOLS[symname] = new

mlir/lib/Bindings/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 Callable, Dict, List, Optional, Sequence, Set, Tuple, Union
+from typing import Any, Callable, Dict, List, Optional, Sequence, Set, Tuple, Union
 
 from mlir 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_affine_exprs(self, callback):
-    """Visits all affine expressions reachable by the expression."""
-    pass
+  def visit_tensor_exprs(self, callback):
+    """Visits all tensor expression reachable by the expression."""
+    callback(self)
 
   def _get_all_dim_defs(self) -> Set[DimDef]:
     """Recursively gets all DimDef affine expressions that are referenced."""
     results = set()
 
-    def visitor(affine_expr):
-      if isinstance(affine_expr, DimDef):
-        results.add(affine_expr)
+    def visit_dim_def(dim_def):
+        if isinstance(dim_def, DimDef):
+          results.add(dim_def)
+
+    def visit_affine_exprs(expr):
+      if isinstance(expr, TensorUse):
+        for ind in expr.indices:
+          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_affine_exprs(visitor)
+    self.visit_tensor_exprs(visit_affine_exprs)
     return results
 
   def collect_uses(self, uses: Set["TensorUse"]):
     """Collects all TensorUses reachable through this expression."""
-    pass
+    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 visit_capture_def(expr):
+      if isinstance(expr, CaptureDef):
+        captures.add(expr)
+    self.visit_tensor_exprs(visit_capture_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":
def to_scalar_expression(self) -> ScalarExpression:
     return ScalarArg(self.tensor_def.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
 
-  def visit_affine_exprs(self, callback):
-    for ind in self.indices:
-      ind.visit_affine_exprs(callback)
-
-  def collect_uses(self, uses: Set["TensorUse"]):
-    uses.add(self)
-
   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
class TensorDef:
   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 a scalar parameter captured by the operation.
+
+  The captured values are not indexed as opposed to tensors. A unique name
+  identifies the captures 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
+
+  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
+
+  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)})")
 
 class Comprehension:
   """Represents a single comprehension."""
 
   def __init__(self, *bindings: Tuple[TensorUse, TensorExpression]):
     self.definitions = list()  # List[TensorUse]
     self.values = list()  # List[TensorExpression]
 
def __init__(self, prim: PrimFnType, args: Sequence[TensorExpression]):
     self.prim = prim
     self.args = tuple(args)
 
   def to_scalar_expression(self) -> ScalarExpression:
     return ScalarApplyFn(self.prim.prim_name,
                          *[arg.to_scalar_expression() for arg in self.args
                           ]).expr()
 
-  def visit_affine_exprs(self, callback):
+  def visit_tensor_exprs(self, callback):
+    super().visit_tensor_exprs(callback)
     for arg in self.args:
-      arg.visit_affine_exprs(callback)
-
-  def collect_uses(self, uses: Set["TensorUse"]):
-    for arg in self.args:
-      arg.collect_uses(uses)
+      arg.visit_tensor_exprs(callback)
 
   def __repr__(self):
     return f"{repr(self.prim)}({', '.join(repr(a) for a in self.args)})"
 
+class const(TensorExpression):
+  """Returns the given constant floating point or integer value."""
+
+  def __init__(self, type_var: TypeVar, value: Any):
+    if not isinstance(type_var, TypeVar):
+      raise ValueError(f"const requires a TypeVar. Got: {repr(type_var)}")
+    if not (isinstance(value, float) or isinstance(value, int)):
+      raise ValueError(f"const requires int or float. Got: {type(value)}")
+    self.type_var = type_var
+    self.value = value
+
+  def to_scalar_expression(self) -> ScalarExpression:
+    return ScalarConst(self.type_var, self.value).expr()
+
+  def __repr__(self):
+    return f"const({self.type_var}, {self.value})"
+
+class index(TensorExpression):
+  """Returns the iteration index for a given dimension name.
+
+  Resolves the given dimension name to obtain its position in the iteration
+  domain of the operation.
+  """
+
+  def __init__(self, dim : DimDef):
+    self.dim_def = dim
+    self.dim = -1
+
+  def resolve_dimension_name(self, affine_state: AffineBuildState):
+    self.dim = affine_state.get_dim(self.dim_def.dimname)
+
+  def to_scalar_expression(self) -> ScalarExpression:
+    assert self.dim != -1, "Dimension name not resolved"
+    return ScalarIndex(self.dim).expr()
+
+  def __repr__(self):
+    return f"index({repr(self.dim)})"
+
 
 class cast(TensorExpression):
   """Casts the element type to a type (typically symbolic TypeVar)."""
 
   def __init__(self, to_type: TypeVar, operand: TensorExpression):
     self.to_type = to_type
     self.operand = operand
 
   def to_scalar_expression(self) -> ScalarExpression:
     return ScalarSymbolicCast(self.to_type,
                               self.operand.to_scalar_expression()).expr()
 
-  def visit_affine_exprs(self, callback):
-    self.operand.visit_affine_exprs(callback)
-
-  def collect_uses(self, uses: Set["TensorUse"]):
-    self.operand.collect_uses(uses)
+  def visit_tensor_exprs(self, callback):
+    super().visit_tensor_exprs(callback)
+    self.operand.visit_tensor_exprs(callback)
 
   def __repr__(self):
     return f"cast({self.to_type}, {repr(self.operand)})"
 
 
 class ReduceApply(TensorExpression):
   """Application of a reduction.
 
   This captures the lhs separately (initial value) separately from the rhs.
   """
 
   def __init__(self, reduce: ReduceFnType, args: Sequence[TensorExpression]):
     self.reduce = reduce
     self.lhs = None  # type: Optional[TensorUse]
     self.args = tuple(args)
 
   def to_scalar_expression(self) -> ScalarExpression:
     if self.lhs is None:
       raise ValueError(f"Cannot scalarize a ReduceApply that has not been "
                        f"bound to its lhs: {self}")
     full_args = [self.lhs.to_scalar_expression()
                 ] + [arg.to_scalar_expression() for arg in self.args]
     return ScalarApplyFn(self.reduce.operator.prim_name, *full_args).expr()
 
-  def visit_affine_exprs(self, callback):
-    for ind in self.reduce.reduce_dims:
-      ind.visit_affine_exprs(callback)
-    for arg in self.args:
-      arg.visit_affine_exprs(callback)
-
-  def collect_uses(self, uses: Set["TensorUse"]):
+  def visit_tensor_exprs(self, callback):
     for arg in self.args:
-      arg.collect_uses(uses)
+      arg.visit_tensor_exprs(callback)
 
   def __repr__(self):
     return f"{repr(self.reduce)}({', '.join(repr(a) for a in self.args)})"
 
 
 class OpInterfaceDef:
   """An interface that an op implements."""
 
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.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 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 tensor(self, name):
-    """Gets a registered tensor by name."""
-    try:
-      return self.registered_tensors[name]
-    except KeyError:
-      raise KeyError(f"Tensor {name} is not registered")
+  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 __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}")
     if self.comprehensions:
       lines[-1] += " {"
       for comprehension in self.comprehensions:
         lines.append(f"    {comprehension}")
       lines.append("}")
     return "\n".join(lines)

mlir/lib/Bindings/Python/mlir/dialects/linalg/opdsl/lang/config.py

return dict(
         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,
+    )
+
+  def __repr__(self):
+    return f"Def({self.capture_def})"
+
 
 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
     - Dynamic based on operands
class LinalgStructuredOpConfig(YAMLObject):
 
   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.uses = dict()  # type: Dict[TensorUse, TensorUseConfig]
 
-    # Compute the ordered set of writes.
+    # Compute the ordered set of writes and collect the tensor, capture, and
+    # index uses.
     collected_uses = set()
+    collected_captures = 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)
+      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)
 
     # 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():
def __init__(self,
     # 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]:
     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)
+
+  @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())
with self.context:
       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 _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,
+        caps=self.ordered_capture_args,
         # TODO: Refactor the hierarchy internally when supporting more
         # than static (preserving this serialized form).
         indexing_maps=LinalgIndexingMapsConfig(
             static_indexing_maps=self.indexing_maps),
         iterator_types=self.iterator_types,
         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(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("])")
     return "\n".join(lines)

mlir/lib/Bindings/Python/mlir/dialects/linalg/opdsl/lang/dsl.py

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 not isinstance(param_default, TensorDef):
+    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)
+    else:
       raise ValueError(f"@tc_def_op function parameters must be defaulted as "
-                       f"TensorDef(...): Found {param_name}: {param_default}")
+                       f"TensorDef(...) or CaptureDef(...): Found {param_name}"
+                       f": {param_default}")
     dsl_func_args.append(param_default)
-    tc_model.add_tensor(param_name, 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/lib/Bindings/Python/mlir/dialects/linalg/opdsl/lang/emitter.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
 
-from typing import Dict, Sequence
+from typing import Any, Dict, Sequence
 
 from mlir.ir import *
 from mlir.dialects import linalg
 from mlir.dialects import std
 # TODO: resolve name collision for Linalg functionality that is injected inside
 # the _mlir.dialects.linalg directly via pybind.
 from _mlir.dialects.linalg import fill_builtin_region
 
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: Value):
+                                 outs: Sequence[Value],
+                                 caps: Sequence[Value]):
   all_arg_defs = op_config.ordered_tensor_args
   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 caps 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(caps, Sequence):
+    raise ValueError(f"Expected named argument caps 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 caps and len(caps) != len(capture_arg_defs):
+    raise ValueError(f"Expected {len(capture_arg_defs)} captures but got "
+                     f"{len(caps)} 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)]
 
-  # Extract type vars for input/output based types.
+  # 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)):
-    tv_name = arg_def.tensor_def.type_var.name
-    type_mapping[tv_name] = arg_element_type
+    _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, caps):
+    _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
 
   # 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])
-  sparse_attr = ArrayAttr.get(
-      [BoolAttr.get(False) for s in list(ins) + list(outs) if isa(RankedTensorType, s.type)])
-  if len(sparse_attr) == 0:
-    sparse_attr = None
+  # TODO: Add support for sparse operands once there is a stable interface.
 
   return (all_arg_defs, in_arg_defs, out_arg_defs, outs, result_types,
-          type_mapping, indexing_maps_attr, iterator_types_attr, sparse_attr)
+          type_mapping, capture_arg_mapping, indexing_maps_attr,
+          iterator_types_attr, sparse_attr)
 
 
 def emit_generic_structured_op(op_config: LinalgStructuredOpConfig,
                                *ins: Value,
-                               outs: Value = ()):
-  all_arg_defs, in_arg_defs, out_arg_defs, outs, result_types, \
-  type_mapping, indexing_maps_attr, iterator_types_attr, sparse_attr =   \
-     prepare_common_structured_op(op_config, *ins, outs = outs)
+                               outs: Sequence[Value] = (),
+                               caps: 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, sparse_attr = \
+     prepare_common_structured_op(op_config, *ins, outs = outs, caps=caps)
 
   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.
       sparse=sparse_attr)  # 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 = generic_op.regions[0].blocks.append(*block_arg_types)
   block_arg_mapping = dict(zip(block_arg_names, block.arguments))
   with InsertionPoint(block):
-    body_builder = _BodyBuilder(type_mapping, block_arg_mapping)
+    body_builder = _BodyBuilder(type_mapping, block_arg_mapping,
+                                capture_arg_mapping)
     for assignment in op_config.assignments:
       body_builder.assign(assignment)
     body_builder.yield_outputs(*_get_tensor_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: Value = ()):
-  all_arg_defs, in_arg_defs, out_arg_defs, outs, result_types, \
-  type_mapping, indexing_maps_attr, iterator_types_attr, sparse_attr =   \
-     prepare_common_structured_op(op_config, *ins, outs = outs)
+                             outs: Sequence[Value] = (),
+                             caps: 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, sparse_attr = \
+     prepare_common_structured_op(op_config, *ins, outs = outs, caps = caps)
 
   # 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}")
 
   named_op = getattr(linalg, op_class_name)(ins, outs, result_types)
   linalgDialect = ctx.get_dialect_descriptor("linalg")
   fill_builtin_region(linalgDialect, named_op.operation)
   # Note: mlir-linalg-ods-yaml-gen.cpp uses a special linalg.memoized_indexing_maps
-  # attribute that the non-yaml path does not. The non-yaml path hardcodes the 
+  # attribute that the non-yaml path does not. The non-yaml path hardcodes the
   # indexing_maps in C++ directly.
   named_op.operation.attributes["linalg.memoized_indexing_maps"] = indexing_maps_attr
   # iterator_types are hardcoded in C++ both in the yaml and non-yaml path.
 
   if len(result_types) == 1:
     return named_op.result
   else:
     return named_op.results
 
 
 class _BodyBuilder:
   """Constructs a structured op body by evaluating assignments."""
 
-  def __init__(self, type_mapping: Dict[str, Type],
-               block_arg_mapping: Dict[str, Value]):
+  def __init__(self,
+               type_mapping: Dict[str, Type],
+               block_arg_mapping: Dict[str, Value],
+               capture_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:
+      return self.constant(expr.scalar_const.type_var.name, expr.scalar_const.value)
+    elif expr.scalar_index:
+      return self.index(expr.scalar_index.dim)
     elif expr.scalar_apply:
       try:
         fn = getattr(self, f"_eval_{expr.scalar_apply.fn_name}")
       except AttributeError:
         raise ValueError(
             f"Function '{expr.scalar_apply.fn_name}' is not a known "
             "scalar body function")
       operand_values = [
           self.expression(operand) for operand in expr.scalar_apply.operands
       ]
       return fn(*operand_values)
     elif expr.symbolic_cast:
       operand_value = self.expression(expr.symbolic_cast.operand)
       return self.cast(expr.symbolic_cast.to_type.name, operand_value)
     raise NotImplementedError(f"Unimplemented scalar body expression: {expr}")
 
+  def constant(self, type_var_name: str, value: Any) -> Value:
+    try:
+      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()}")
+    try:
+      if(_is_floating_point_type(type)):
+        return std.ConstantOp(type, FloatAttr.get(type, float(value))).result
+      elif(_is_integer_type(type)):
+        return std.ConstantOp(type, IntegerAttr.get(type, int(value))).result
+    except ValueError:
+      raise ValueError(f"Unable to cast value {value} to type {type}")
+    raise NotImplementedError(f"Unimplemented constant type {type}")
+
+  def index(self, dim: int) -> Value:
+    dim_attr = IntegerAttr.get(IntegerType.get_signless(64), dim)
+    return linalg.IndexOp(IndexType.get(), dim_attr).result
+
   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()}")
     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)
-
-    raise ValueError(f"Unable to cast body expression from {operand.type} to "
-                     f"{to_type}")
-
   def _cast_to_integer(self, to_type: Type, operand: Value) -> Value:
     to_width = IntegerType(to_type).width
     operand_type = operand.type
     if _is_floating_point_type(operand_type):
       return std.FPToSIOp(to_type, operand).result
+    if _is_index_type(operand_type):
+      return std.IndexCastOp(to_type, operand).result
     # Assume integer.
     from_width = IntegerType(operand_type).width
     if to_width > from_width:
       return std.SignExtendIOp(to_type, operand).result
     elif to_width < from_width:
       return std.TruncateIOp(to_type, operand).result
     raise ValueError(f"Unable to cast body expression from {operand_type} to "
                      f"{to_type}")
for n in output_names:
       except KeyError:
         raise ValueError(f"Body assignments do not assign all outputs: "
                          f"missing '{n}'")
     linalg.YieldOp(output_values)
 
   def _eval_add(self, lhs: Value, rhs: Value) -> Value:
     if _is_floating_point_type(lhs.type):
       return std.AddFOp(lhs.type, lhs, rhs).result
-    if _is_integer_type(lhs.type):
+    if _is_integer_type(lhs.type) or _is_index_type(lhs.type):
       return std.AddIOp(lhs.type, lhs, rhs).result
     raise NotImplementedError("Unsupported 'add' operand: {lhs}")
 
+  def _eval_sub(self, lhs: Value, rhs: Value) -> Value:
+    if _is_floating_point_type(lhs.type):
+      return std.SubFOp(lhs.type, lhs, rhs).result
+    if _is_integer_type(lhs.type) or _is_index_type(lhs.type):
+      return std.SubIOp(lhs.type, lhs, rhs).result
+    raise NotImplementedError("Unsupported 'sub' operand: {lhs}")
+
   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):
+    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],
                            ins: Sequence[Value],
                            out_arg_defs: Sequence[TensorDefConfig],
for v in values:
     types.append(t.element_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 _add_type_mapping(name: str, type: Type, type_mapping: Dict[str, Type]):
+  if name in type_mapping:
+    if type_mapping[name] != type:
+        raise ValueError(f"Cannot overwrite type mapping {name} = "
+                         f"{type_mapping[name]} by type {type}")
+  type_mapping[name] = 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))
 
-
 def _is_integer_type(t: Type) -> bool:
   return IntegerType.isinstance(t)
 
+def _is_index_type(t: Type) -> bool:
+  return IndexType.isinstance(t)
 
 def _get_floating_point_width(t: Type) -> int:
   # TODO: Create a FloatType in the Python API and implement the switch
   # there.
   if F64Type.isinstance(t):
     return 64
   if F32Type.isinstance(t):
     return 32
   if F16Type.isinstance(t):
     return 16
   if BF16Type.isinstance(t):
     return 16
   raise NotImplementedError(f"Unhandled floating point type switch {t}")

mlir/lib/Bindings/Python/mlir/dialects/linalg/opdsl/lang/scalar_expr.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
 """Models DAGs of scalar math expressions.
 
 Used for generating region bodies at the "math" level where they are still type
 polymorphic. This is modeled to be polymorphic by attribute name for interop
 with serialization schemes that are just plain-old-dicts.
 
 These classes are typically not user accessed and are created as a by-product
 of interpreting a comprehension DSL and model the operations to perform in the
 op body. The class hierarchy is laid out to map well to a form of YAML that
 can be easily consumed from the C++ side, not necessarily for ergonomics.
 """
 
-from typing import Optional, Sequence
+from typing import Any, Optional, Sequence
 
 from .yaml_helper import *
 from .types import *
 
 __all__ = [
     "ScalarAssign",
     "ScalarApplyFn",
     "ScalarArg",
+    "ScalarCapture",
+    "ScalarConst",
+    "ScalarIndex",
     "ScalarExpression",
     "ScalarSymbolicCast",
 ]
 
 
 class ScalarApplyFn:
   """A type of ScalarExpression that applies a named function to operands."""
 
def __init__(self, arg: str):
     self.arg = arg
 
   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, type_var: TypeVar, value: Any):
+    self.type_var = type_var
+    self.value = value
+
+  def expr(self) -> "ScalarExpression":
+    return ScalarExpression(scalar_const=self)
+
+  def __repr__(self):
+    return f"(ScalarConst({self.type_var}, {self.value})"
+
+class ScalarIndex:
+  """A type of ScalarExpression accessing an iteration index."""
+
+  def __init__(self, dim : int):
+    self.dim = dim
+
+  def expr(self) -> "ScalarExpression":
+    return ScalarExpression(scalar_index=self)
+
+  def __repr__(self):
+    return f"(ScalarIndex({self.dim})"
 
 class ScalarSymbolicCast:
   """A type of ScalarExpression that symbolically casts an operand to a TypeVar.
   """
 
   def __init__(self, to_type: TypeVar, operand: "ScalarExpression"):
     self.to_type = to_type
     self.operand = operand
 
   def expr(self) -> "ScalarExpression":
     return ScalarExpression(symbolic_cast=self)
 
   def __repr__(self):
     return f"ScalarSymbolicCast({self.to_type}, {self.operand})"
 
 
 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(symbolic_cast)) != 1:
+    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_block_arg', 'symbolic_cast' must be "
-          "specified")
+          "One of 'scalar_apply', 'scalar_arg', 'scalar_capture', '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=dict(type_var=self.scalar_const.type_var.name,
+                                    attributes=[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.
       return dict(symbolic_cast=dict(type_var=self.symbolic_cast.to_type.name,
                                      operands=[self.symbolic_cast.operand]))
     else:
       raise ValueError(f"Unexpected ScalarExpression type: {self}")

mlir/lib/Bindings/Python/mlir/dialects/linalg/opdsl/lang/types.py

 
 from enum import Enum
 from typing import Dict
 
 __all__ = [
     "TypeVar",
     "TV",
 
+    # Predefined types.
+    "I32",
+    "I64",
+    "F32",
+    "F64",
+
     # TypeVar aliases.
     "T",
     "U",
     "V",
 ]
 
 
 class TypeVar:
class ExpandoTypeVars:
         return cls(n)
 
     return ExpandoTypeVars()
 
 
 # Expando access via TV.foo
 TV = TypeVar.create_expando()
 
+# Predefined types.
+I32 = TV.I32
+I64 = TV.I64
+F32 = TV.F32
+F64 = TV.F64
+
 # Some common type name aliases.
 T = TV.T
 U = TV.U
 V = TV.V

mlir/test/Bindings/Python/dialects/linalg/opdsl/emit_structured_generic.py

 
 
 @linalg_structured_op
 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_2d(A=TensorDef(T, S.M, S.N, output=True),
+                min=CaptureDef(F64),
+                max=CaptureDef(F64),
+                seed=CaptureDef(I32)):
+  multiplier = const(I32, 1103515245)
+  increment = const(I32, 12345)
+  temp1 = (cast(I32, index(D.m)) + seed) * multiplier + increment
+  temp2 = (cast(I32, index(D.n)) + temp1) * multiplier + increment
+  scaling = (max - min) * const(F64, 2.3283064e-10)
+  A[D.m, D.n] = cast(T, cast(F64, temp2) * scaling + min)
 
 with Context() as ctx, Location.unknown():
   module = Module.create()
   f16 = F16Type.get()
   f32 = F32Type.get()
   f64 = F64Type.get()
   i8 = IntegerType.get_signless(8)
   i16 = IntegerType.get_signless(16)
with InsertionPoint(module.body):
     # CHECK-NEXT:   linalg.yield %[[ADD]] : f32
     # 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_2d
+    # CHECK-SAME:  %{{.*}} tensor<4x16xi32>, %[[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 : i32
+    # CHECK-DAG:    %[[CST1:.+]] = constant 12345 : i32
+    # CHECK-DAG:    %[[RND1:.+]] = muli %[[RND0]], %[[CST0]] : i32
+    # CHECK-DAG:    %[[RND2:.+]] = addi %[[RND1]], %[[CST1]] : i32
+    # CHECK:        %[[RND3:.+]] = sitofp %{{.*}} : i32 to f64
+    # CHECK-DAG:    %[[DIFF:.+]] = subf %[[MAX]], %[[MIN]] : f64
+    # CHECK-DAG:    %[[CST2:.+]] = constant 2.3283063999999999E-10 : f64
+    # CHECK-DAG:    %[[FACT:.+]] = mulf %[[DIFF]], %[[CST2]] : f64
+    # CHECK-DAG:    %[[RND4:.+]] = mulf %[[RND3]], %[[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_2d(init_result, min, max, seed):
+      return fill_rng_2d(outs=[init_result], caps=[min, max, seed])
 
 print(module)