[mlir][OpDSL] Refactor function handling.

Prepare the OpDSL function handling to introduce more function classes. A follow up commit will split ArithFn into UnaryFn and BinaryFn. This revision prepares the split by adding a function kind enum to handle different function types using a single class on the various levels of the stack (for example, there is now one TensorFn and one ScalarFn).

Depends On D119718

Reviewed By: aartbik

Differential Revision: https://reviews.llvm.org/D120108

3 files changed, 73 insertions, 108 deletions

diff --git a/mlir/python/mlir/dialects/linalg/opdsl/lang/comprehension.py b/mlir/python/mlir/dialects/linalg/opdsl/lang/comprehension.py
index 68c08809e16a..d26aa077096c 100644
--- a/mlir/python/mlir/dialects/linalg/opdsl/lang/comprehension.py
+++ b/mlir/python/mlir/dialects/linalg/opdsl/lang/comprehension.py
@@ -133,55 +133,36 @@ class TensorUse(TensorExpression):
             f"[{', '.join([repr(i) for i in self.indices])}]")
 
 
-class TensorArithFn(TensorExpression):
-  """Application of an arithmetic function."""
+class TensorFn(TensorExpression):
+  """Application of a tensor function."""
 
-  def __init__(self, arith_fn: "ArithFnType", args: Sequence[TensorExpression]):
-    self.arith_fn = arith_fn
-    self.args = tuple(args)
-
-  def to_scalar_expression(self) -> ScalarExpression:
-    return ScalarArithFn(self.arith_fn.fn_name,
-                         *[arg.to_scalar_expression() for arg in self.args
-                          ]).expr()
-
-  def visit_tensor_exprs(self, callback: Callable[["TensorExpression"], None]):
-    super().visit_tensor_exprs(callback)
-    for arg in self.args:
-      arg.visit_tensor_exprs(callback)
-
-  def __repr__(self):
-    return f"{repr(self.arith_fn)}({', '.join(repr(a) for a in self.args)})"
-
-
-class TensorTypeFn(TensorExpression):
-  """Application of a type conversion function."""
-
-  def __init__(self, type_fn: Optional["TypeFn"],
-               operand_def: Optional["OperandDef"], type_var: TypeVar,
-               arg: TensorExpression):
-    if bool(type_fn) + bool(operand_def) != 1:
-      raise ValueError("Either 'type_fn' or 'operand_def' must be specified")
-    self.type_fn = type_fn
+  def __init__(self, kind: "FunctionKind", name: Optional[str],
+               operand_def: Optional["OperandDef"], type_var: Optional[TypeVar],
+               args: Sequence[TensorExpression]):
+    if bool(name) + bool(operand_def) != 1:
+      raise ValueError("One of 'name', 'operand_def' must be specified")
+    self.name = name
+    self.kind = kind
     self.operand_def = operand_def
     self.type_var = type_var
-    self.arg = arg
+    self.args = args
 
   def to_scalar_expression(self) -> ScalarExpression:
     if self.operand_def:
-      assert self.operand_def.name, "TypeFnAttr not registered with an op"
-    fn_name = self.type_fn.fn_name if self.type_fn else None
+      assert self.operand_def.name, "TensorFn not registered with an op"
     attr_name = self.operand_def.name if self.operand_def else None
-    return ScalarTypeFn(fn_name, attr_name, self.type_var,
-                        self.arg.to_scalar_expression()).expr()
+    args = [arg.to_scalar_expression() for arg in self.args]
+    return ScalarFn(self.kind, self.name, attr_name, self.type_var, args).expr()
 
   def visit_tensor_exprs(self, callback: Callable[["TensorExpression"], None]):
     super().visit_tensor_exprs(callback)
-    self.arg.visit_tensor_exprs(callback)
+    for arg in self.args:
+      arg.visit_tensor_exprs(callback)
 
   def __repr__(self):
-    return (f"{repr(self.type_fn)}[{repr(self.operand_def)}]"
-            f"({self.type_var}, {self.arg})")
+    name = self.operand_def.name if self.operand_def else self.name
+    return (f"{self.kind.name}.{name}(type_var={self.type_var}, "
+            f"args={', '.join(repr(a) for a in self.args)})")
 
 
 class TensorReduceFn(TensorExpression):
@@ -194,7 +175,7 @@ class TensorReduceFn(TensorExpression):
                args: Sequence[TensorExpression]):
     self.reduce_use = reduce_use
     self.lhs = None  # type: Optional[TensorUse]
-    self.args = tuple(args)
+    self.args = args
 
   def to_scalar_expression(self) -> ScalarExpression:
     if self.lhs is None:
@@ -202,7 +183,8 @@ class TensorReduceFn(TensorExpression):
                        f"bound to its lhs: {self}")
     full_args = [self.lhs.to_scalar_expression()
                 ] + [arg.to_scalar_expression() for arg in self.args]
-    return ScalarArithFn(self.reduce_use.arith_fn.fn_name, *full_args).expr()
+    return ScalarFn(FunctionKind.ARITH, self.reduce_use.arith_fn.fn_name, None,
+                    None, full_args).expr()
 
   def visit_tensor_exprs(self, callback: Callable[["TensorExpression"], None]):
     for arg in self.args:
@@ -259,6 +241,11 @@ class index(TensorExpression):
 ###############################################################################
 
 
+class FunctionKind(Enum):
+  ARITH = 0
+  TYPE = 1
+
+
 class TypeFnType:
   """Type conversion function.
 
@@ -269,8 +256,8 @@ class TypeFnType:
   def __init__(self, fn_name: str):
     self.fn_name = fn_name
 
-  def __call__(self, type_var: TypeVar, arg: TensorExpression) -> "TypeFnType":
-    return TensorTypeFn(self, None, type_var, arg)
+  def __call__(self, type_var: TypeVar, arg: TensorExpression) -> "TensorFn":
+    return TensorFn(FunctionKind.TYPE, self.fn_name, None, type_var, [arg])
 
   def __repr__(self):
     return f"{self.fn_name}"
@@ -301,8 +288,8 @@ class ArithFnType:
   def __init__(self, fn_name: str):
     self.fn_name = fn_name
 
-  def __call__(self, *args) -> "TensorArithFn":
-    return TensorArithFn(self, args)
+  def __call__(self, *args) -> "TensorFn":
+    return TensorFn(FunctionKind.ARITH, self.fn_name, None, None, args)
 
   def __repr__(self):
     return f"{self.fn_name}"
@@ -562,8 +549,8 @@ class TypeFnAttrDef:
     self.operand_def = OperandDef(
         OperandKind.TYPE_FN_ATTR, default_fn=default.fn_name)
 
-  def __call__(self, type_var: TypeVar, arg: TensorExpression) -> TensorTypeFn:
-    return TensorTypeFn(None, self.operand_def, type_var, arg)
+  def __call__(self, type_var: TypeVar, arg: TensorExpression) -> TensorFn:
+    return TensorFn(FunctionKind.TYPE, None, self.operand_def, type_var, [arg])
 
 
 ###############################################################################
diff --git a/mlir/python/mlir/dialects/linalg/opdsl/lang/emitter.py b/mlir/python/mlir/dialects/linalg/opdsl/lang/emitter.py
index fc8c13bfe6ec..07050f56fa64 100644
--- a/mlir/python/mlir/dialects/linalg/opdsl/lang/emitter.py
+++ b/mlir/python/mlir/dialects/linalg/opdsl/lang/emitter.py
@@ -270,19 +270,19 @@ class _BodyBuilder:
       dim_attr = IntegerAttr.get(
           IntegerType.get_signless(64), expr.scalar_index.dim)
       return linalg.IndexOp(dim_attr).result
-    elif expr.arith_fn:
-      fn = self._get_function(f"_arithfn_{expr.arith_fn.fn_name}")
+    elif expr.scalar_fn and expr.scalar_fn.kind == FunctionKind.ARITH:
+      fn = self._get_function(f"_arithfn_{expr.scalar_fn.fn_name}")
       operand_values = [
-          self.expression(operand) for operand in expr.arith_fn.operands
+          self.expression(operand) for operand in expr.scalar_fn.operands
       ]
       return fn(*operand_values)
-    elif expr.type_fn:
-      fn_name = expr.type_fn.fn_name
-      if expr.type_fn.attr_name:
-        fn_name = self.type_fn_attr_mapping[expr.type_fn.attr_name]
+    elif expr.scalar_fn and expr.scalar_fn.kind == FunctionKind.TYPE:
+      fn_name = expr.scalar_fn.fn_name
+      if expr.scalar_fn.attr_name:
+        fn_name = self.type_fn_attr_mapping[expr.scalar_fn.attr_name]
       fn = self._get_function(f"_typefn_{fn_name}")
-      operand = self.expression(expr.type_fn.operand)
-      return fn(expr.type_fn.type_var.name, operand)
+      operand_value = self.expression(expr.scalar_fn.operands[0])
+      return fn(expr.scalar_fn.type_var.name, operand_value)
     raise NotImplementedError(f"Unimplemented scalar body expression: {expr}")
 
   def yield_outputs(self, *output_names: str):
diff --git a/mlir/python/mlir/dialects/linalg/opdsl/lang/scalar_expr.py b/mlir/python/mlir/dialects/linalg/opdsl/lang/scalar_expr.py
index af21b40cf27a..aa894dc10954 100644
--- a/mlir/python/mlir/dialects/linalg/opdsl/lang/scalar_expr.py
+++ b/mlir/python/mlir/dialects/linalg/opdsl/lang/scalar_expr.py
@@ -15,13 +15,13 @@ can be easily consumed from the C++ side, not necessarily for ergonomics.
 
 from typing import Optional, Sequence
 
-from .yaml_helper import *
+from .comprehension import *
 from .types import *
+from .yaml_helper import *
 
 __all__ = [
     "ScalarAssign",
-    "ScalarArithFn",
-    "ScalarTypeFn",
+    "ScalarFn",
     "ScalarArg",
     "ScalarConst",
     "ScalarIndex",
@@ -29,36 +29,27 @@ __all__ = [
 ]
 
 
-class ScalarArithFn:
-  """A type of ScalarExpression that applies an arithmetic function."""
-
-  def __init__(self, fn_name: str, *operands: "ScalarExpression"):
-    self.fn_name = fn_name
-    self.operands = operands
-
-  def expr(self) -> "ScalarExpression":
-    return ScalarExpression(arith_fn=self)
-
-  def __repr__(self):
-    return f"ScalarArithFn<{self.fn_name}>({', '.join(self.operands)})"
-
-
-class ScalarTypeFn:
-  """A type of ScalarExpression that applies a type conversion function."""
+class ScalarFn:
+  """A type of ScalarExpression that applies a function."""
 
-  def __init__(self, fn_name: Optional[str], attr_name: Optional[str],
-               type_var: TypeVar, operand: "ScalarExpression"):
+  def __init__(self, kind: "FunctionKind", fn_name: Optional[str],
+               attr_name: Optional[str], type_var: Optional["TypeVar"],
+               operands: Sequence["ScalarExpression"]):
+    if bool(fn_name) + bool(attr_name) != 1:
+      raise ValueError("One of 'fn_name', 'attr_name' must be specified")
+    self.kind = kind
     self.fn_name = fn_name
     self.attr_name = attr_name
     self.type_var = type_var
-    self.operand = operand
+    self.operands = operands
 
   def expr(self) -> "ScalarExpression":
-    return ScalarExpression(type_fn=self)
+    return ScalarExpression(scalar_fn=self)
 
   def __repr__(self):
-    return (f"ScalarTypeFn<{self.fn_name}[{self.attr_name}]>"
-            f"({self.type_var}, {self.operand})")
+    name = self.fn_name if self.fn_name else self.attr_name
+    return (f"ScalarFn<{self.kind.name}.{name}>(type_var={self.type_var}, "
+            f"operands=[{', '.join(self.operands)}])")
 
 
 class ScalarArg:
@@ -104,51 +95,38 @@ class ScalarExpression(YAMLObject):
   """An expression on scalar values.
 
   Can be one of:
-    - ScalarArithFn
-    - ScalarTypeFn
+    - ScalarFn
     - ScalarArg
     - ScalarConst
     - ScalarIndex
-    - ScalarSymbolicCast
   """
   yaml_tag = "!ScalarExpression"
 
   def __init__(self,
-               arith_fn: Optional[ScalarArithFn] = None,
-               type_fn: Optional[ScalarTypeFn] = None,
+               scalar_fn: Optional[ScalarFn] = None,
                scalar_arg: Optional[ScalarArg] = None,
                scalar_const: Optional[ScalarConst] = None,
                scalar_index: Optional[ScalarIndex] = None):
-    if (bool(arith_fn) + bool(type_fn) + bool(scalar_arg) + bool(scalar_const) +
+    if (bool(scalar_fn) + bool(scalar_arg) + bool(scalar_const) +
         bool(scalar_index)) != 1:
-      raise ValueError("One of 'arith_fn', 'type_fn', 'scalar_arg', "
-                       "'scalar_const', 'scalar_index', must be specified")
-    self.arith_fn = arith_fn
-    self.type_fn = type_fn
+      raise ValueError("One of 'scalar_fn', 'scalar_arg', 'scalar_const', or "
+                       "'scalar_index' must be specified")
+    self.scalar_fn = scalar_fn
     self.scalar_arg = scalar_arg
     self.scalar_const = scalar_const
     self.scalar_index = scalar_index
 
   def to_yaml_custom_dict(self):
-    if self.arith_fn:
-      return dict(
-          arith_fn=dict(
-              fn_name=self.arith_fn.fn_name,
-              operands=list(self.arith_fn.operands),
-          ))
-    if self.type_fn:
-      # Note that even though operands must be arity 1, we write it the
-      # same way as for apply because it allows handling code to be more
-      # generic vs having a special form.
-      type_fn_dict = dict(
-          type_var=self.type_fn.type_var.name,
-          operands=[self.type_fn.operand],
-      )
-      if self.type_fn.fn_name:
-        type_fn_dict["fn_name"] = self.type_fn.fn_name
-      if self.type_fn.attr_name:
-        type_fn_dict["attr_name"] = self.type_fn.attr_name
-      return dict(type_fn=type_fn_dict)
+    if self.scalar_fn:
+      scalar_fn_dict = dict(kind=self.scalar_fn.kind.name.lower())
+      if self.scalar_fn.fn_name:
+        scalar_fn_dict["fn_name"] = self.scalar_fn.fn_name
+      if self.scalar_fn.attr_name:
+        scalar_fn_dict["attr_name"] = self.scalar_fn.attr_name
+      if self.scalar_fn.type_var:
+        scalar_fn_dict["type_var"] = self.scalar_fn.type_var.name
+      scalar_fn_dict["operands"] = list(self.scalar_fn.operands)
+      return dict(scalar_fn=scalar_fn_dict)
     elif self.scalar_arg:
       return dict(scalar_arg=self.scalar_arg.arg)
     elif self.scalar_const: