[mlir][sparse][taco] Add a few unary operations.

[bixia]

Summary

Add operations -, abs, ceil and floor to the index notation.

Add test cases.

[aartbik]

mlir/test/Integration/Dialect/SparseTensor/taco/test_tensor_unary_ops.py

20

very neat to see all these new ops!

mlir/test/Integration/Dialect/SparseTensor/taco/tools/mlir_pytaco.py

1631

per our offline discussion, maybe clarify this for people like me that are a bit rusty in Python ;-)
Something like "# Forward decls." or so

[bixia]

Changed ... to pass and added comments.

mlir/test/Integration/Dialect/SparseTensor/taco/test_tensor_unary_ops.py

20

Thanks!

mlir/test/Integration/Dialect/SparseTensor/taco/tools/mlir_pytaco.py

1631

I replaced ... with pass and added comments.

[bixia]

Rebased.

mlir/test/Integration/Dialect/SparseTensor/taco/test_tensor_unary_ops.py

+# RUN: SUPPORTLIB=%mlir_runner_utils_dir/libmlir_c_runner_utils%shlibext %PYTHON %s | FileCheck %s
+
+import numpy as np
+import os
+import sys
+
+_SCRIPT_PATH = os.path.dirname(os.path.abspath(__file__))
+sys.path.append(_SCRIPT_PATH)
+from tools import mlir_pytaco_api as pt
+
+i, j = pt.get_index_vars(2)
+A = pt.tensor([2, 3])
+B = pt.tensor([2, 3])
+A.insert([0, 1], 10.3)
+A.insert([1, 1], 40.7)
+A.insert([0, 2], -11.3)
+A.insert([1, 2], -41.7)
+
+B[i, j] = abs(A[i, j])
+indices, values = B.get_coordinates_and_values()
+passed = np.array_equal(indices, [[0, 1], [0, 2], [1, 1], [1, 2]])
+passed += np.allclose(values, [10.3, 11.3, 40.7, 41.7])
+
+B[i, j] = pt.ceil(A[i, j])
+indices, values = B.get_coordinates_and_values()
+passed += np.array_equal(indices, [[0, 1], [0, 2], [1, 1], [1, 2]])
+passed += np.allclose(values, [11, -11, 41, -41])
+
+B[i, j] = pt.floor(A[i, j])
+indices, values = B.get_coordinates_and_values()
+passed += np.array_equal(indices, [[0, 1], [0, 2], [1, 1], [1, 2]])
+passed += np.allclose(values, [10, -12, 40, -42])
+
+B[i, j] = -A[i, j]
+indices, values = B.get_coordinates_and_values()
+passed += np.array_equal(indices, [[0, 1], [0, 2], [1, 1], [1, 2]])
+passed += np.allclose(values, [-10.3, 11.3, -40.7, 41.7])
+
+# CHECK: Number of passed: 8
+print("Number of passed:", passed)

mlir/test/Integration/Dialect/SparseTensor/taco/tools/mlir_pytaco.py

 
 # Bitwidths for pointers and indices.
 _POINTER_BIT_WIDTH = 0
 _INDEX_BIT_WIDTH = 0
 # The entry point to the JIT compiled program.
 _ENTRY_NAME = "main"
 
 # Type aliases for type annotation.
+_UnaryOp = Callable[[Any], Any]
 _BinaryOp = Callable[[Any, Any], Any]
 _ExprVisitor = Callable[..., None]
 _ExprInfoDict = Dict["IndexExpr", "_ExprInfo"]
 _LogicalOp = Callable[[bool, bool], bool]
 _ModeFormatOp = Callable[["ModeFormat", "ModeFormat"], "ModeFormat"]
 _SubtreeLeafChecker = Optional[Callable[..., bool]]
 
 
def _verify_operand_and_build_expr(self, rhs, op: _BinaryOp) -> "_BinaryExpr":
 
     Raises:
       ValueError: If rhs is not an IndexExpr.
     """
     if not isinstance(rhs, IndexExpr):
       raise ValueError(f"Expected IndexExpr: {rhs}")
     return _BinaryExpr(op, self, rhs)
 
+  def _build_unary_expr(self, op: _UnaryOp) -> "_UnaryExpr":
+    """Build a unary expression.
+
+    Args:
+      op: A _UnaryOp object representing the unary operation.
+    """
+    return _UnaryExpr(op, self)
+
   def __add__(self, rhs) -> "_BinaryExpr":
     """Defines the operator +.
 
     Args:
       rhs: The value being added, which could be any Python object from user
         inputs.
 
     Returns:
def __mul__(self, rhs) -> "_BinaryExpr":
     Returns:
       A _BinaryExpr object representing the operation.
 
     Raises:
       ValueError: If rhs is not an IndexExpr.
     """
     return self._verify_operand_and_build_expr(rhs, operator.mul)
 
+  def __abs__(self) -> "_UnaryExpr":
+    """Defines the operator abs.
+
+    Returns:
+      A _UnaryExpr object representing the operation.
+    """
+    return self._build_unary_expr(operator.abs)
+
+  def __neg__(self) -> "_UnaryExpr":
+    """Defines the operator neg.
+
+    Returns:
+      A _UnaryExpr object representing the operation.
+    """
+    return self._build_unary_expr(operator.neg)
+
   def __sub__(self, rhs) -> "_BinaryExpr":
     """Defines the operator -.
 
     Args:
       rhs: The value being subtracted, which could be any Python object from
         user inputs.
 
     Returns:
def _gather_input_accesses_index_vars(
     expr: IndexExpr,
     input_accesses: List[Access],
 ) -> None:
   """Collects Access nodes."""
   if isinstance(expr, Access) and expr not in input_accesses:
     input_accesses.append(expr)
 
 
+def _op_ceil(__a: Any) -> Any:
+  """A _UnaryOp object for operation ceil."""
+  pass
+
+
+def _op_floor(__a: Any) -> Any:
+  """A _UnaryOp object for operation floor."""
+  pass
+
+
+def _op_unary_to_callable(op: _UnaryOp) -> lang.UnaryFnType:
+  """Returns the linalg dialect function object for the given operation."""
+  op_to_callable = {
+      operator.abs: lang.UnaryFn.abs,
+      operator.neg: lang.UnaryFn.negf,
+      _op_ceil: lang.UnaryFn.ceil,
+      _op_floor: lang.UnaryFn.floor,
+  }
+  return op_to_callable[op]
+
+
[e-mail](frozen=True)
+class _UnaryExpr(IndexExpr):
+  """The representation for a Unary operation.
+
+  Attributes:
+  op: A _UnaryOp representing the operation.
+  a: An IndexExpr representing the operand for the operation.
+  """
+  op: _BinaryOp
+  a: IndexExpr
+
+  def __post_init__(self) -> None:
+    """Verifies that the operand being added is an IndexExpr."""
+    assert isinstance(self.a, IndexExpr)
+
+  def _emit_expression(
+      self,
+      expr_to_opnd: Dict[IndexExpr, lang.OperandDef],
+      expr_to_info: _ExprInfoDict,
+  ) -> lang.ScalarExpression:
+    """Emits the expression tree and returns the expression."""
+    # The current expression node is an internal node of the structured op.
+    if self not in expr_to_opnd:
+      a = self.a._emit_expression(expr_to_opnd, expr_to_info)
+      return _op_unary_to_callable(self.op)(a)
+
+    # The current expression is a leaf node of the structured op. That is, it is
+    # a temporary tensor generated by its child structured op.
+    op_info = expr_to_info[self].structop_info
+    assert op_info is not None
+    dims = _mlir_dimensions_from_index_vars(op_info.dst_indices)
+    return lang.TensorUse(expr_to_opnd[self], dims)
+
+  def _visit(self,
+             func: _ExprVisitor,
+             args,
+             *,
+             leaf_checker: _SubtreeLeafChecker = None) -> None:
+    """A post-order visitor."""
+    if leaf_checker is None or not leaf_checker(self, *args):
+      self.a._visit(func, args, leaf_checker=leaf_checker)
+    func(self, *args)
+
+  def dtype(self) -> DType:
+    """Returns the data type of the operation."""
+    return self.a.dtype()
+
+
 def _op_to_callable(op: _BinaryOp) -> lang.BinaryFnType:
   """Returns the linalg dialect function object for the given operation."""
   op_to_callable = {
       operator.add: lang.BinaryFn.add,
       operator.sub: lang.BinaryFn.sub,
       operator.mul: lang.BinaryFn.mul,
   }
   return op_to_callable[op]
 
-
 @dataclasses.dataclass(frozen=True)
 class _BinaryExpr(IndexExpr):
   """The representation for a binary operation.
 
   Attributes:
   op: A _BinaryOp representing the binary operation.
   a: An IndexExpr representing the first operand of the operation.
   b: An IndexExpr representing the second operand of the operation.
if isinstance(expr, Access):
     if expr.tensor.format is None:
       # Treat each dimension of a dense tensor as DENSE for the purpose of
       # calculating temporary tensor storage format.
       mode_formats = tuple([ModeFormat.DENSE] * len(src_dims))
     else:
       mode_formats = tuple(expr.tensor.format.format_pack.formats)
     assert len(src_dims) == len(mode_formats)
     dim_infos = tuple([_DimInfo(d, m) for d, m in zip(src_dims, mode_formats)])
+  elif isinstance(expr, _UnaryExpr):
+    a_info = expr_to_info[expr.a]
+    index_to_dim_info = {
+        i: d for i, d in zip(a_info.src_indices, a_info.dim_infos)
+    }
+    # Here we rely on the fact that dictionaries keep the insertion order for
+    # keys and values.
+    src_indices = tuple(index_to_dim_info.keys())
+    dim_infos = tuple(index_to_dim_info.values())
   else:
     assert isinstance(expr, _BinaryExpr)
     a_info = expr_to_info[expr.a]
     index_to_dim_info = {
         i: d for i, d in zip(a_info.src_indices, a_info.dim_infos)
     }
     b_info = expr_to_info[expr.b]
     _validate_and_collect_dim_info(index_to_dim_info, b_info.src_indices,
) -> None:
   expr_info = expr_to_info[expr]
 
   if isinstance(expr, _BinaryExpr):
     a_info = expr_to_info[expr.a]
     b_info = expr_to_info[expr.b]
     expr_info.acc_reduce_indices = (
         a_info.acc_reduce_indices | b_info.acc_reduce_indices
         | expr_info.reduce_indices)
+  elif isinstance(expr, _UnaryExpr):
+    a_info = expr_to_info[expr.a]
+    expr_info.acc_reduce_indices = (
+        a_info.acc_reduce_indices | expr_info.reduce_indices)
   else:
     assert isinstance(expr, Access)
     # Handle simple reduction expression in the format of A[i] = B[i, j].
     expr_info.acc_reduce_indices = expr_info.reduce_indices
 
 
 
 def _gather_structured_op(
else:
     assert isinstance(expr, Access)
     indices = expr.indices
     name = expr.tensor.name
 
   dim_sym = _mlir_symbols_from_index_vars(indices)
   opnd = lang.OperandDef(lang.OperandKind.INPUT_TENSOR, lang.T, dim_sym)
   op_def.add_operand(name, opnd)
   return opnd
+
+
+def _check_and_build_unary(a: Access, op: _UnaryOp) -> "_UnaryExpr":
+  """Build a unary operation ceil.
+
+    Args:
+      a: The operand, which could be any Python object from user inputs.
+      op: An _UnaryOp object representing the operation.
+
+    Returns:
+      A _UnaryExpr object representing the operation.
+
+    Raises:
+      ValueError: If a is not an IndexExpr.
+    """
+  if not isinstance(a, Access):
+    raise ValueError(f"Expected an Access Operand: {a}")
+  return a._build_unary_expr(op)
+
+
+def ceil(a: Access) -> "_UnaryExpr":
+  """Defines the operation ceil.
+
+    Args:
+      a: The operand, which could be any Python object from user inputs.
+
+    Returns:
+      A _UnaryExpr object representing the operation.
+
+    Raises:
+      ValueError: If a is not an IndexExpr.
+    """
+  return _check_and_build_unary(a, _op_ceil)
+
+
+def floor(a: Access) -> "_UnaryExpr":
+  """Defines the operation floor.
+
+    Args:
+      a: The operand, which could be any Python object from user inputs.
+
+    Returns:
+      A _UnaryExpr object representing the operation.
+
+    Raises:
+      ValueError: If a is not an IndexExpr.
+    """
+  return _check_and_build_unary(a, _op_floor)

mlir/test/Integration/Dialect/SparseTensor/taco/tools/mlir_pytaco_api.py

 needed for the PyTACO API to support the execution of PyTACO programs using the
 MLIR-PyTACO implementation.
 """
 
 from . import mlir_pytaco
 from . import mlir_pytaco_io
 
 # Functions defined by PyTACO API.
+ceil = mlir_pytaco.ceil
+floor = mlir_pytaco.floor
 get_index_vars = mlir_pytaco.get_index_vars
 from_array = mlir_pytaco.Tensor.from_array
 read = mlir_pytaco_io.read
 write = mlir_pytaco_io.write
 
 # Classes defined by PyTACO API.
 dtype = mlir_pytaco.DType
 mode_format = mlir_pytaco.ModeFormat