[mlir] Add `complex.abs`, `complex.div` and `complex.mul` to ComplexOps.

[ftynse]

mlir/lib/Conversion/ComplexToLLVM/ComplexToLLVM.cpp

34

camelBack please

mlir/test/Conversion/ComplexToLLVM/convert-to-llvm.mlir

72

Nit: I think we converged to putting % outside the named pattern so %[[NAME]] is still indicative of it being an SSA value.

[herhut]

Thanks for adding these. Beyond coding style this looks great!

mlir/lib/Conversion/ComplexToLLVM/ComplexToLLVM.cpp

265

nit: Sorting is off

mlir/test/Conversion/ComplexToLLVM/convert-to-llvm.mlir

1

Does this work without the extra RUN and the \ at the end?

[pifon2a]

Address the comments.

[pifon2a]

Add complex namespace to examples in ComplexOps.td.

[herhut]

Thanks!

mlir/include/mlir/Dialect/Complex/IR/ComplexOps.td

let description = [{
 
     ```mlir
     %a = add %b, %c : complex<f32>
     ```
   }];
 }
 
 //===----------------------------------------------------------------------===//
+// AbsOp
+//===----------------------------------------------------------------------===//
+
+def AbsOp : Complex_Op<"abs",
+    [NoSideEffect,
+     TypesMatchWith<"complex element type matches result type",
+                    "complex", "result",
+                    "$_self.cast<ComplexType>().getElementType()">]> {
+  let summary = "computes absolute value of a complex number";
+  let description = [{
+    The `abs` op takes a single complex number and computes its absolute value.
+
+    Example:
+
+    ```mlir
+    %a = abs %b : complex<f32>
+    ```
+  }];
+
+  let arguments = (ins Complex<AnyFloat>:$complex);
+  let results = (outs AnyFloat:$result);
+
+  let assemblyFormat = "$complex attr-dict `:` type($complex)";
+}
+
+//===----------------------------------------------------------------------===//
 // CreateOp
 //===----------------------------------------------------------------------===//
 
 def CreateOp : Complex_Op<"create",
     [NoSideEffect,
      AllTypesMatch<["real", "imaginary"]>,
      TypesMatchWith<"complex element type matches real operand type",
                     "complex", "real",
def CreateOp : Complex_Op<"create",
 
   let arguments = (ins AnyFloat:$real, AnyFloat:$imaginary);
   let results = (outs Complex<AnyFloat>:$complex);
 
   let assemblyFormat = "$real `,` $imaginary attr-dict `:` type($complex)";
 }
 
 //===----------------------------------------------------------------------===//
+// DivOp
+//===----------------------------------------------------------------------===//
+
+def DivOp : ComplexArithmeticOp<"div"> {
+  let summary = "complex division";
+  let description = [{
+    The `div` operation takes two complex numbers and returns result of their
+    division:
+
+    ```mlir
+    %a = div %b, %c : complex<f32>
+    ```
+  }];
+}
+
+//===----------------------------------------------------------------------===//
 // ImOp
 //===----------------------------------------------------------------------===//
 
 def ImOp : Complex_Op<"im",
     [NoSideEffect,
      TypesMatchWith<"complex element type matches result type",
                     "complex", "imaginary",
                     "$_self.cast<ComplexType>().getElementType()">]> {
def ImOp : Complex_Op<"im",
 
   let arguments = (ins Complex<AnyFloat>:$complex);
   let results = (outs AnyFloat:$imaginary);
 
   let assemblyFormat = "$complex attr-dict `:` type($complex)";
 }
 
 //===----------------------------------------------------------------------===//
+// MulOp
+//===----------------------------------------------------------------------===//
+
+def MulOp : ComplexArithmeticOp<"mul"> {
+  let summary = "complex multiplication";
+  let description = [{
+    The `mul` operation takes two complex numbers and returns their product:
+
+    ```mlir
+    %a = mul %b, %c : complex<f32>
+    ```
+  }];
+}
+
+//===----------------------------------------------------------------------===//
 // ReOp
 //===----------------------------------------------------------------------===//
 
 def ReOp : Complex_Op<"re",
     [NoSideEffect,
      TypesMatchWith<"complex element type matches result type",
                     "complex", "real",
                     "$_self.cast<ComplexType>().getElementType()">]> {

mlir/lib/Conversion/ComplexToLLVM/ComplexToLLVM.cpp

 #include "mlir/Dialect/Complex/IR/Complex.h"
 #include "mlir/Dialect/LLVMIR/LLVMDialect.h"
 
 using namespace mlir;
 using namespace mlir::LLVM;
 
 namespace {
 
+struct AbsOpConversion : public ConvertOpToLLVMPattern<complex::AbsOp> {
+  using ConvertOpToLLVMPattern<complex::AbsOp>::ConvertOpToLLVMPattern;
+
+  LogicalResult
+  matchAndRewrite(complex::AbsOp op, ArrayRef<Value> operands,
+                  ConversionPatternRewriter &rewriter) const override {
+    complex::AbsOp::Adaptor transformed(operands);
+    auto loc = op.getLoc();
+
+    ComplexStructBuilder complexStruct(transformed.complex());
+    Value real = complexStruct.real(rewriter, op.getLoc());
+    Value imag = complexStruct.imaginary(rewriter, op.getLoc());
+
+    auto fmf = LLVM::FMFAttr::get({}, op.getContext());
+    Value sq_norm = rewriter.create<LLVM::FAddOp>(
+        loc, rewriter.create<LLVM::FMulOp>(loc, real, real, fmf),
+        rewriter.create<LLVM::FMulOp>(loc, imag, imag, fmf), fmf);
+
+    rewriter.replaceOpWithNewOp<LLVM::SqrtOp>(op, sq_norm);
+    return success();
+  }
+};
+
 struct CreateOpConversion : public ConvertOpToLLVMPattern<complex::CreateOp> {
   using ConvertOpToLLVMPattern<complex::CreateOp>::ConvertOpToLLVMPattern;
 
   LogicalResult
   matchAndRewrite(complex::CreateOp complexOp, ArrayRef<Value> operands,
                   ConversionPatternRewriter &rewriter) const override {
     complex::CreateOp::Adaptor transformed(operands);
 
matchAndRewrite(complex::AddOp op, ArrayRef<Value> operands,
     result.setReal(rewriter, loc, real);
     result.setImaginary(rewriter, loc, imag);
 
     rewriter.replaceOp(op, {result});
     return success();
   }
 };
 
+struct DivOpConversion : public ConvertOpToLLVMPattern<complex::DivOp> {
+  using ConvertOpToLLVMPattern<complex::DivOp>::ConvertOpToLLVMPattern;
+
+  LogicalResult
+  matchAndRewrite(complex::DivOp op, ArrayRef<Value> operands,
+                  ConversionPatternRewriter &rewriter) const override {
+    auto loc = op.getLoc();
+    BinaryComplexOperands arg =
+        unpackBinaryComplexOperands<complex::DivOp>(op, operands, rewriter);
+
+    // Initialize complex number struct for result.
+    auto structType = typeConverter->convertType(op.getType());
+    auto result = ComplexStructBuilder::undef(rewriter, loc, structType);
+
+    // Emit IR to add complex numbers.
+    auto fmf = LLVM::FMFAttr::get({}, op.getContext());
+    Value rhs_re = arg.rhs.real();
+    Value rhs_im = arg.rhs.imag();
+    Value lhs_re = arg.lhs.real();
+    Value lhs_im = arg.lhs.imag();
+
+    Value rhs_sq_norm = rewriter.create<LLVM::FAddOp>(
+        loc, rewriter.create<LLVM::FMulOp>(loc, rhs_re, rhs_re, fmf),
+        rewriter.create<LLVM::FMulOp>(loc, rhs_im, rhs_im, fmf), fmf);
+
+    Value result_real = rewriter.create<LLVM::FAddOp>(
+        loc, rewriter.create<LLVM::FMulOp>(loc, lhs_re, rhs_re, fmf),
+        rewriter.create<LLVM::FMulOp>(loc, lhs_im, rhs_im, fmf), fmf);
+
+    Value result_imag = rewriter.create<LLVM::FSubOp>(
+        loc, rewriter.create<LLVM::FMulOp>(loc, lhs_im, rhs_re, fmf),
+        rewriter.create<LLVM::FMulOp>(loc, lhs_re, rhs_im, fmf), fmf);
+
+    result.setReal(
+        rewriter, loc,
+        rewriter.create<LLVM::FDivOp>(loc, result_real, rhs_sq_norm, fmf));
+    result.setImaginary(
+        rewriter, loc,
+        rewriter.create<LLVM::FDivOp>(loc, result_imag, rhs_sq_norm, fmf));
+
+    rewriter.replaceOp(op, {result});
+    return success();
+  }
+};
+
+struct MulOpConversion : public ConvertOpToLLVMPattern<complex::MulOp> {
+  using ConvertOpToLLVMPattern<complex::MulOp>::ConvertOpToLLVMPattern;
+
+  LogicalResult
+  matchAndRewrite(complex::MulOp op, ArrayRef<Value> operands,
+                  ConversionPatternRewriter &rewriter) const override {
+    auto loc = op.getLoc();
+    BinaryComplexOperands arg =
+        unpackBinaryComplexOperands<complex::MulOp>(op, operands, rewriter);
+
+    // Initialize complex number struct for result.
+    auto structType = typeConverter->convertType(op.getType());
+    auto result = ComplexStructBuilder::undef(rewriter, loc, structType);
+
+    // Emit IR to add complex numbers.
+    auto fmf = LLVM::FMFAttr::get({}, op.getContext());
+    Value rhs_re = arg.rhs.real();
+    Value rhs_im = arg.rhs.imag();
+    Value lhs_re = arg.lhs.real();
+    Value lhs_im = arg.lhs.imag();
+
+    Value real = rewriter.create<LLVM::FSubOp>(
+        loc, rewriter.create<LLVM::FMulOp>(loc, rhs_re, lhs_re, fmf),
+        rewriter.create<LLVM::FMulOp>(loc, rhs_im, lhs_im, fmf), fmf);
+
+    Value imag = rewriter.create<LLVM::FAddOp>(
+        loc, rewriter.create<LLVM::FMulOp>(loc, lhs_im, rhs_re, fmf),
+        rewriter.create<LLVM::FMulOp>(loc, lhs_re, rhs_im, fmf), fmf);
+
+    result.setReal(rewriter, loc, real);
+    result.setImaginary(rewriter, loc, imag);
+
+    rewriter.replaceOp(op, {result});
+    return success();
+  }
+};
+
 struct SubOpConversion : public ConvertOpToLLVMPattern<complex::SubOp> {
   using ConvertOpToLLVMPattern<complex::SubOp>::ConvertOpToLLVMPattern;
 
   LogicalResult
   matchAndRewrite(complex::SubOp op, ArrayRef<Value> operands,
                   ConversionPatternRewriter &rewriter) const override {
     auto loc = op.getLoc();
     BinaryComplexOperands arg =

 };
 } // namespace
 
 void mlir::populateComplexToLLVMConversionPatterns(
     LLVMTypeConverter &converter, OwningRewritePatternList &patterns) {
   // clang-format off
   patterns.insert<
       AddOpConversion,
+      AbsOpConversion,
       CreateOpConversion,
+      DivOpConversion,
+      MulOpConversion,
       ImOpConversion,
       ReOpConversion,
       SubOpConversion
     >(converter);
   // clang-format on
 }
 
 namespace {

mlir/test/Conversion/ComplexToLLVM/convert-to-llvm.mlir

-// RUN: mlir-opt %s -split-input-file -convert-complex-to-llvm | FileCheck %s
+// RUN: mlir-opt %s -split-input-file -convert-complex-to-llvm
+// | FileCheck %s
 
 // CHECK-LABEL: llvm.func @complex_numbers()
 // CHECK-NEXT:    %[[REAL0:.*]] = llvm.mlir.constant(1.200000e+00 : f32) : f32
 // CHECK-NEXT:    %[[IMAG0:.*]] = llvm.mlir.constant(3.400000e+00 : f32) : f32
 // CHECK-NEXT:    %[[CPLX0:.*]] = llvm.mlir.undef : !llvm.struct<(f32, f32)>
 // CHECK-NEXT:    %[[CPLX1:.*]] = llvm.insertvalue %[[REAL0]], %[[CPLX0]][0] : !llvm.struct<(f32, f32)>
 // CHECK-NEXT:    %[[CPLX2:.*]] = llvm.insertvalue %[[IMAG0]], %[[CPLX1]][1] : !llvm.struct<(f32, f32)>
 // CHECK-NEXT:    %[[REAL1:.*]] = llvm.extractvalue %[[CPLX2:.*]][0] : !llvm.struct<(f32, f32)>
func @complex_substraction() {
   %a_im = constant 3.4 : f64
   %a = complex.create %a_re, %a_im : complex<f64>
   %b_re = constant 5.6 : f64
   %b_im = constant 7.8 : f64
   %b = complex.create %b_re, %b_im : complex<f64>
   %c = complex.sub %a, %b : complex<f64>
   return
 }
+
+// CHECK-LABEL: llvm.func @complex_div
+func @complex_div(%lhs: complex<f32>, %rhs: complex<f32>) -> complex<f32> {
+  %div = complex.div %lhs, %rhs : complex<f32>
+  return %div : complex<f32>
+}
+// CHECK-LABEL: llvm.func @complex_div(
+// CHECK-SAME:    [[LHS:%.*]]: [[C_TY:!.*]], [[RHS:%.*]]: [[C_TY]]) -> [[C_TY]]
+
+// CHECK: [[LHS_RE:%.*]] = llvm.extractvalue [[LHS]][0] : [[C_TY]]
+// CHECK: [[LHS_IM:%.*]] = llvm.extractvalue [[LHS]][1] : [[C_TY]]
+// CHECK: [[RHS_RE:%.*]] = llvm.extractvalue [[RHS]][0] : [[C_TY]]
+// CHECK: [[RHS_IM:%.*]] = llvm.extractvalue [[RHS]][1] : [[C_TY]]
+
+// CHECK: [[RESULT_0:%.*]] = llvm.mlir.undef : [[C_TY]]
+
+// CHECK: [[RHS_RE_SQ:%.*]] = llvm.fmul [[RHS_RE]], [[RHS_RE]]  : f32
+// CHECK: [[RHS_IM_SQ:%.*]] = llvm.fmul [[RHS_IM]], [[RHS_IM]]  : f32
+// CHECK: [[SQ_NORM:%.*]] = llvm.fadd [[RHS_RE_SQ]], [[RHS_IM_SQ]]  : f32
+
+// CHECK: [[REAL_TMP_0:%.*]] = llvm.fmul [[LHS_RE]], [[RHS_RE]]  : f32
+// CHECK: [[REAL_TMP_1:%.*]] = llvm.fmul [[LHS_IM]], [[RHS_IM]]  : f32
+// CHECK: [[REAL_TMP_2:%.*]] = llvm.fadd [[REAL_TMP_0]], [[REAL_TMP_1]]  : f32
+
+// CHECK: [[IMAG_TMP_0:%.*]] = llvm.fmul [[LHS_IM]], [[RHS_RE]]  : f32
+// CHECK: [[IMAG_TMP_1:%.*]] = llvm.fmul [[LHS_RE]], [[RHS_IM]]  : f32
+// CHECK: [[IMAG_TMP_2:%.*]] = llvm.fsub [[IMAG_TMP_0]], [[IMAG_TMP_1]]  : f32
+
+// CHECK: [[REAL:%.*]] = llvm.fdiv [[REAL_TMP_2]], [[SQ_NORM]]  : f32
+// CHECK: [[RESULT_1:%.*]] = llvm.insertvalue [[REAL]], [[RESULT_0]][0] : [[C_TY]]
+// CHECK: [[IMAG:%.*]] = llvm.fdiv [[IMAG_TMP_2]], [[SQ_NORM]]  : f32
+// CHECK: [[RESULT_2:%.*]] = llvm.insertvalue [[IMAG]], [[RESULT_1]][1] : [[C_TY]]
+// CHECK: llvm.return [[RESULT_2]] : [[C_TY]]
+
+
+// CHECK-LABEL: llvm.func @complex_mul
+func @complex_mul(%lhs: complex<f32>, %rhs: complex<f32>) -> complex<f32> {
+  %mul = complex.mul %lhs, %rhs : complex<f32>
+  return %mul : complex<f32>
+}
+// CHECK-LABEL: llvm.func @complex_mul(
+// CHECK-SAME:    [[LHS:%.*]]: [[C_TY:!.*]], [[RHS:%.*]]: [[C_TY]]) -> [[C_TY]]
+
+// CHECK: [[LHS_RE:%.*]] = llvm.extractvalue [[LHS]][0] : [[C_TY]]
+// CHECK: [[LHS_IM:%.*]] = llvm.extractvalue [[LHS]][1] : [[C_TY]]
+// CHECK: [[RHS_RE:%.*]] = llvm.extractvalue [[RHS]][0] : [[C_TY]]
+// CHECK: [[RHS_IM:%.*]] = llvm.extractvalue [[RHS]][1] : [[C_TY]]
+// CHECK: [[RESULT_0:%.*]] = llvm.mlir.undef : [[C_TY]]
+
+// CHECK: [[REAL_TMP_0:%.*]] = llvm.fmul [[RHS_RE]], [[LHS_RE]]  : f32
+// CHECK: [[REAL_TMP_1:%.*]] = llvm.fmul [[RHS_IM]], [[LHS_IM]]  : f32
+// CHECK: [[REAL:%.*]] = llvm.fsub [[REAL_TMP_0]], [[REAL_TMP_1]]  : f32
+
+// CHECK: [[IMAG_TMP_0:%.*]] = llvm.fmul [[LHS_IM]], [[RHS_RE]]  : f32
+// CHECK: [[IMAG_TMP_1:%.*]] = llvm.fmul [[LHS_RE]], [[RHS_IM]]  : f32
+// CHECK: [[IMAG:%.*]] = llvm.fadd [[IMAG_TMP_0]], [[IMAG_TMP_1]]  : f32
+
+// CHECK: [[RESULT_1:%.*]] = llvm.insertvalue [[REAL]], [[RESULT_0]][0]
+// CHECK: [[RESULT_2:%.*]] = llvm.insertvalue [[IMAG]], [[RESULT_1]][1]
+// CHECK: llvm.return [[RESULT_2]] : [[C_TY]]
+
+// CHECK-LABEL: llvm.func @complex_abs
+func @complex_abs(%arg: complex<f32>) -> f32 {
+  %abs = complex.abs %arg: complex<f32>
+  return %abs : f32
+}
+
+// CHECK-LABEL: llvm.func @complex_abs(
+// CHECK-SAME: [[LHS:%.*]]: [[C_TY:!.*]])
+// CHECK: [[REAL:%.*]] = llvm.extractvalue [[VAL_0]][0] : !llvm.struct<(f32, f32)>
+// CHECK: [[IMAG:%.*]] = llvm.extractvalue [[VAL_0]][1] : !llvm.struct<(f32, f32)>
+// CHECK: [[REAL_SQ:%.*]] = llvm.fmul [[REAL]], [[REAL]]  : f32
+// CHECK: [[IMAG_SQ:%.*]] = llvm.fmul [[IMAG]], [[IMAG]]  : f32
+// CHECK: [[SQ_NORM:%.*]] = llvm.fadd [[REAL_SQ]], [[IMAG_SQ]]  : f32
+// CHECK: [[NORM:%.*]] = "llvm.intr.sqrt"([[SQ_NORM]]) : (f32) -> f32
+// CHECK: llvm.return [[VAL_6]] : f32
+

mlir/test/Dialect/Complex/ops.mlir

 // RUN: mlir-opt %s | mlir-opt | FileCheck %s
 // RUN: mlir-opt %s --mlir-print-op-generic | mlir-opt | FileCheck %s
 
 
 // CHECK-LABEL: func @ops(
 // CHECK-SAME:            [[F:%.*]]: f32) {
 func @ops(%f: f32) {
   // CHECK: [[C:%.*]] = complex.create [[F]], [[F]] : complex<f32>
   %complex = complex.create %f, %f : complex<f32>
 
   // CHECK: complex.re [[C]] : complex<f32>
   %real = complex.re %complex : complex<f32>
 
   // CHECK: complex.im [[C]] : complex<f32>
   %imag = complex.im %complex : complex<f32>
 
+  // CHECK: complex.abs [[C]] : complex<f32>
+  %abs = complex.abs %complex : complex<f32>
+
   // CHECK: complex.add [[C]], [[C]] : complex<f32>
   %sum = complex.add %complex, %complex : complex<f32>
 
+  // CHECK: complex.div [[C]], [[C]] : complex<f32>
+  %div = complex.div %complex, %complex : complex<f32>
+
+  // CHECK: complex.mul [[C]], [[C]] : complex<f32>
+  %prod = complex.mul %complex, %complex : complex<f32>
+
   // CHECK: complex.sub [[C]], [[C]] : complex<f32>
   %diff = complex.sub %complex, %complex : complex<f32>
   return
 }