[mlir][linalg] Extend opdsl to support operations on complex types.

[bixia]

Summary

Linalg opdsl now supports negf/add/sub/mul on complex types.

Add a test.

[aartbik]

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

47

do we want a c64 test too?

mlir/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 Callable, Dict, List, Sequence, Tuple, Union
 
 from .....ir import *
 
 from .... import func
 from .... import linalg
 from .... import math
 from .... import arith
+from .... import complex
 from ...._ods_common import get_op_result_or_value as _get_op_result_or_value, get_op_results_or_values as _get_op_results_or_values
 
 from .scalar_expr import *
 from .config import *
 from .comprehension import *
 import numpy as np
 
 __all__ = [
class _BodyBuilder:
   def _unary_floor(self, x: Value) -> Value:
     if _is_floating_point_type(x.type):
       return math.FloorOp(x).result
     raise NotImplementedError("Unsupported 'floor' operand: {x}")
 
   def _unary_negf(self, x: Value) -> Value:
     if _is_floating_point_type(x.type):
       return arith.NegFOp(x).result
+    if _is_complex_type(x.type):
+      return complex.NegOp(x).result
     raise NotImplementedError("Unsupported 'negf' operand: {x}")
 
   def _binary_add(self, lhs: Value, rhs: Value) -> Value:
     if _is_floating_point_type(lhs.type):
       return arith.AddFOp(lhs, rhs).result
     if _is_integer_type(lhs.type) or _is_index_type(lhs.type):
       return arith.AddIOp(lhs, rhs).result
+    if _is_complex_type(lhs.type):
+      return complex.AddOp(lhs, rhs).result
     raise NotImplementedError("Unsupported 'add' operands: {lhs}, {rhs}")
 
   def _binary_sub(self, lhs: Value, rhs: Value) -> Value:
     if _is_floating_point_type(lhs.type):
       return arith.SubFOp(lhs, rhs).result
     if _is_integer_type(lhs.type) or _is_index_type(lhs.type):
       return arith.SubIOp(lhs, rhs).result
+    if _is_complex_type(lhs.type):
+      return complex.SubOp(lhs, rhs).result
     raise NotImplementedError("Unsupported 'sub' operands: {lhs}, {rhs}")
 
   def _binary_mul(self, lhs: Value, rhs: Value) -> Value:
     if _is_floating_point_type(lhs.type):
       return arith.MulFOp(lhs, rhs).result
     if _is_integer_type(lhs.type) or _is_index_type(lhs.type):
       return arith.MulIOp(lhs, rhs).result
+    if _is_complex_type(lhs.type):
+      return complex.MulOp(lhs, rhs).result
     raise NotImplementedError("Unsupported 'mul' operands: {lhs}, {rhs}")
 
   def _binary_max_signed(self, lhs: Value, rhs: Value) -> Value:
     if _is_floating_point_type(lhs.type):
       return arith.MaxFOp(lhs, rhs).result
     if _is_integer_type(lhs.type) or _is_index_type(lhs.type):
       return arith.MaxSIOp(lhs, rhs).result
     raise NotImplementedError("Unsupported 'max' operands: {lhs}, {rhs}")
def _add_type_mapping(operand_config: OperandDefConfig, operand_type: Type,
   if name in type_mapping:
     if type_mapping[name] != element_or_self_type:
       raise ValueError(f"Cannot overwrite type mapping {name} = "
                        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_complex_type(t: Type) -> bool:
+  return ComplexType.isinstance(t)
+
+
 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:

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

 @linalg_structured_op(op_name="custom_op_name")
 def non_default_op_name(I=TensorDef(T, S.N), O=TensorDef(T, S.N, output=True)):
   O[D.n] = I[D.n]
 
 
 with Context() as ctx, Location.unknown():
   module = Module.create()
   f32 = F32Type.get()
+  c32 = ComplexType.get(f32)
   i32 = IntegerType.get_signless(32)
   with InsertionPoint(module.body):
 
     # CHECK-LABEL: @test_f32_const
     # CHECK-DAG:    %[[CST0:.+]] = arith.constant 42 : i64
     # CHECK-DAG:    %[[CST0_CAST:.+]] = arith.uitofp %[[CST0]] : i64 to f32
     # CHECK-DAG:    %[[CST1:.+]] = arith.constant 2.3283063999999999E-10 : f64
     # CHECK-DAG:    %[[CST1_CAST:.+]] = arith.truncf %[[CST1]] : f64 to f32
with InsertionPoint(module.body):
     # CHECK-NEXT:   %[[EXP:.+]] = arith.negf %[[IN]] : f32
     # CHECK-NEXT:   linalg.yield %[[EXP]] : f32
     # CHECK-NEXT: -> tensor<4x16xf32>
     @func.FuncOp.from_py_func(
         RankedTensorType.get((4, 16), f32), RankedTensorType.get((4, 16), f32))
     def test_f32_elemwise_neg(input, init_result):
       return elemwise_unary_poly(input, outs=[init_result], fun=UnaryFn.negf)
 
+    # CHECK-LABEL: @test_c32_elemwise_neg
+    # CHECK:      ^{{.*}}(%[[IN:.+]]: complex<f32>, %[[OUT:.+]]: complex<f32>)
+    # CHECK-NEXT:   %[[EXP:.+]] = complex.neg %[[IN]] : complex<f32>
+    # CHECK-NEXT:   linalg.yield %[[EXP]] : complex<f32>
+    # CHECK-NEXT: -> tensor<4x16xcomplex<f32>>
+    @func.FuncOp.from_py_func(
+        RankedTensorType.get((4, 16), c32), RankedTensorType.get((4, 16), c32))
+    def test_c32_elemwise_neg(input, init_result):
+      return elemwise_unary_poly(input, outs=[init_result], fun=UnaryFn.negf)
+
     # Just check that we don't assert out on name mismatch.
     # CHECK-LABEL: @test_non_default_op_name
     @func.FuncOp.from_py_func(
         RankedTensorType.get((42,), f32), RankedTensorType.get((42,), f32))
     def test_non_default_op_name(input, init_result):
       return non_default_op_name(input, outs=[init_result])
 
 
 print(module)