[mlir] Move the complex support of std.constant to a new complex.constant operation

[rriddle]

Summary

This is part of splitting up the standard dialect.

mlir/include/mlir/Dialect/Complex/IR/Complex.h

 //===- Complex.h - Complex dialect --------------------------------*- C++-*-==//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_DIALECT_COMPLEX_IR_COMPLEX_H_
 #define MLIR_DIALECT_COMPLEX_IR_COMPLEX_H_
 
 #include "mlir/Dialect/Arithmetic/IR/Arithmetic.h"
-#include "mlir/Dialect/StandardOps/IR/Ops.h"
 #include "mlir/IR/BuiltinTypes.h"
 #include "mlir/IR/Dialect.h"
 #include "mlir/IR/OpDefinition.h"
 #include "mlir/IR/OpImplementation.h"
+#include "mlir/Interfaces/InferTypeOpInterface.h"
 #include "mlir/Interfaces/SideEffectInterfaces.h"
-#include "mlir/Interfaces/VectorInterfaces.h"
 
 //===----------------------------------------------------------------------===//
 // Complex Dialect
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Dialect/Complex/IR/ComplexOpsDialect.h.inc"
 
 //===----------------------------------------------------------------------===//
 // Complex Dialect Operations
 //===----------------------------------------------------------------------===//
 
 #define GET_OP_CLASSES
 #include "mlir/Dialect/Complex/IR/ComplexOps.h.inc"
 
 #endif // MLIR_DIALECT_COMPLEX_IR_COMPLEX_H_

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

 def Complex_Dialect : Dialect {
   let name = "complex";
   let cppNamespace = "::mlir::complex";
   let description = [{
     The complex dialect is intended to hold complex numbers creation and
     arithmetic ops.
   }];
 
-  let dependentDialects = ["arith::ArithmeticDialect", "StandardOpsDialect"];
+  let dependentDialects = ["arith::ArithmeticDialect"];
   let hasConstantMaterializer = 1;
   let emitAccessorPrefix = kEmitAccessorPrefix_Prefixed;
 }
 
 #endif // COMPLEX_BASE

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

 //===- ComplexOps.td - Complex op definitions ----------------*- tablegen -*-===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef COMPLEX_OPS
 #define COMPLEX_OPS
 
 include "mlir/Dialect/Complex/IR/ComplexBase.td"
+include "mlir/IR/OpAsmInterface.td"
 include "mlir/Interfaces/InferTypeOpInterface.td"
 include "mlir/Interfaces/SideEffectInterfaces.td"
 
 class Complex_Op<string mnemonic, list<OpTrait> traits = []>
     : Op<Complex_Dialect, mnemonic, traits>;
 
 // Base class for standard arithmetic operations on complex numbers with a
 // floating-point element type. These operations take two operands and return
let description = [{
 
     ```mlir
     %a = complex.add %b, %c : complex<f32>
     ```
   }];
 }
 
 //===----------------------------------------------------------------------===//
+// ConstantOp
+//===----------------------------------------------------------------------===//
+
+def ConstantOp : Complex_Op<"constant", [
+    ConstantLike, NoSideEffect,
+    DeclareOpInterfaceMethods<OpAsmOpInterface, ["getAsmResultNames"]>
+  ]> {
+  let summary = "complex number constant operation";
+  let description = [{
+    The `complex.constant` operation creates a constant complex number from an
+    attribute containing the real and imaginary parts.
+
+    Example:
+
+    ```mlir
+    %a = complex.constant [0.1, -1.0] : complex<f64>
+    ```
+  }];
+
+  let arguments = (ins ArrayAttr:$value);
+  let results = (outs Complex<AnyFloat>:$complex);
+
+  let assemblyFormat = "$value attr-dict `:` type($complex)";
+  let hasFolder = 1;
+  let verifier = [{ return ::verify(*this); }];
+  
+  let extraClassDeclaration = [{
+    /// Returns true if a constant operation can be built with the given value
+    /// and result type.
+    static bool isBuildableWith(Attribute value, Type type);
+  }];
+}
+
+//===----------------------------------------------------------------------===//
 // CreateOp
 //===----------------------------------------------------------------------===//
 
 def CreateOp : Complex_Op<"create",
     [NoSideEffect,
      AllTypesMatch<["real", "imaginary"]>,
      TypesMatchWith<"complex element type matches real operand type",
                     "complex", "real",

mlir/lib/Dialect/Complex/IR/CMakeLists.txt

 add_mlir_dialect_library(MLIRComplex
   ComplexOps.cpp
   ComplexDialect.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/Complex
 
   DEPENDS
   MLIRComplexOpsIncGen
 
   LINK_LIBS PUBLIC
   MLIRArithmetic
   MLIRDialect
+  MLIRInferTypeOpInterface
   MLIRIR
-  MLIRStandard
   )

mlir/lib/Dialect/Complex/IR/ComplexDialect.cpp

 //===- ComplexDialect.cpp - MLIR Complex Dialect --------------------------===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Dialect/Arithmetic/IR/Arithmetic.h"
 #include "mlir/Dialect/Complex/IR/Complex.h"
-#include "mlir/Dialect/StandardOps/IR/Ops.h"
 
 using namespace mlir;
 
 #include "mlir/Dialect/Complex/IR/ComplexOpsDialect.cpp.inc"
 
 void complex::ComplexDialect::initialize() {
   addOperations<
 #define GET_OP_LIST
 #include "mlir/Dialect/Complex/IR/ComplexOps.cpp.inc"
       >();
 }
 
 Operation *complex::ComplexDialect::materializeConstant(OpBuilder &builder,
                                                         Attribute value,
                                                         Type type,
                                                         Location loc) {
-  // TODO complex.constant
-  if (type.isa<ComplexType>())
-    return builder.create<ConstantOp>(loc, type, value);
+  if (complex::ConstantOp::isBuildableWith(value, type)) {
+    return builder.create<complex::ConstantOp>(loc, type,
+                                               value.cast<ArrayAttr>());
+  }
   if (arith::ConstantOp::isBuildableWith(value, type))
     return builder.create<arith::ConstantOp>(loc, type, value);
   return nullptr;
 }

mlir/lib/Dialect/Complex/IR/ComplexOps.cpp

 //===- ComplexOps.cpp - MLIR Complex Operations ---------------------------===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Dialect/Complex/IR/Complex.h"
 #include "mlir/IR/Builders.h"
 
 using namespace mlir;
 using namespace mlir::complex;
 
 //===----------------------------------------------------------------------===//
-// TableGen'd op method definitions
+// ConstantOp
 //===----------------------------------------------------------------------===//
 
-#define GET_OP_CLASSES
-#include "mlir/Dialect/Complex/IR/ComplexOps.cpp.inc"
+OpFoldResult ConstantOp::fold(ArrayRef<Attribute> operands) {
+  assert(operands.empty() && "constant has no operands");
+  return getValue();
+}
+
+void ConstantOp::getAsmResultNames(
+    function_ref<void(Value, StringRef)> setNameFn) {
+  setNameFn(getResult(), "cst");
+}
+
+bool ConstantOp::isBuildableWith(Attribute value, Type type) {
+  if (auto arrAttr = value.dyn_cast<ArrayAttr>()) {
+    auto complexTy = type.dyn_cast<ComplexType>();
+    if (!complexTy)
+      return false;
+    auto complexEltTy = complexTy.getElementType();
+    return arrAttr.size() == 2 && arrAttr[0].getType() == complexEltTy &&
+           arrAttr[1].getType() == complexEltTy;
+  }
+  return false;
+}
+
+static LogicalResult verify(ConstantOp op) {
+  ArrayAttr arrayAttr = op.getValue();
+  if (arrayAttr.size() != 2) {
+    return op.emitOpError(
+        "requires 'value' to be a complex constant, represented as array of "
+        "two values");
+  }
+
+  auto complexEltTy = op.getType().getElementType();
+  if (complexEltTy != arrayAttr[0].getType() ||
+      complexEltTy != arrayAttr[1].getType()) {
+    return op.emitOpError()
+           << "requires attribute's element types (" << arrayAttr[0].getType()
+           << ", " << arrayAttr[1].getType()
+           << ") to match the element type of the op's return type ("
+           << complexEltTy << ")";
+  }
+  return success();
+}
+
+//===----------------------------------------------------------------------===//
+// CreateOp
+//===----------------------------------------------------------------------===//
 
 OpFoldResult CreateOp::fold(ArrayRef<Attribute> operands) {
   assert(operands.size() == 2 && "binary op takes two operands");
   // Fold complex.create(complex.re(op), complex.im(op)).
   if (auto reOp = getOperand(0).getDefiningOp<ReOp>()) {
     if (auto imOp = getOperand(1).getDefiningOp<ImOp>()) {
       if (reOp.getOperand() == imOp.getOperand()) {
         return reOp.getOperand();
       }
     }
   }
   return {};
 }
 
+//===----------------------------------------------------------------------===//
+// ImOp
+//===----------------------------------------------------------------------===//
+
 OpFoldResult ImOp::fold(ArrayRef<Attribute> operands) {
   assert(operands.size() == 1 && "unary op takes 1 operand");
   ArrayAttr arrayAttr = operands[0].dyn_cast_or_null<ArrayAttr>();
   if (arrayAttr && arrayAttr.size() == 2)
     return arrayAttr[1];
   if (auto createOp = getOperand().getDefiningOp<CreateOp>())
     return createOp.getOperand(1);
   return {};
 }
 
+//===----------------------------------------------------------------------===//
+// ReOp
+//===----------------------------------------------------------------------===//
+
 OpFoldResult ReOp::fold(ArrayRef<Attribute> operands) {
   assert(operands.size() == 1 && "unary op takes 1 operand");
   ArrayAttr arrayAttr = operands[0].dyn_cast_or_null<ArrayAttr>();
   if (arrayAttr && arrayAttr.size() == 2)
     return arrayAttr[0];
   if (auto createOp = getOperand().getDefiningOp<CreateOp>())
     return createOp.getOperand(0);
   return {};
 }
+
+//===----------------------------------------------------------------------===//
+// TableGen'd op method definitions
+//===----------------------------------------------------------------------===//
+
+#define GET_OP_CLASSES
+#include "mlir/Dialect/Complex/IR/ComplexOps.cpp.inc"

mlir/lib/Dialect/StandardOps/IR/Ops.cpp

 static void print(OpAsmPrinter &p, ConstantOp &op) {
   p << " ";
   p.printOptionalAttrDict(op->getAttrs(), /*elidedAttrs=*/{"value"});
 
   if (op->getAttrs().size() > 1)
     p << ' ';
   p << op.getValue();
 
-  // If the value is a symbol reference or Array, print a trailing type.
-  if (op.getValue().isa<SymbolRefAttr, ArrayAttr>())
+  // If the value is a symbol reference, print a trailing type.
+  if (op.getValue().isa<SymbolRefAttr>())
     p << " : " << op.getType();
 }
 
 static ParseResult parseConstantOp(OpAsmParser &parser,
                                    OperationState &result) {
   Attribute valueAttr;
   if (parser.parseOptionalAttrDict(result.attributes) ||
       parser.parseAttribute(valueAttr, "value", result.attributes))
     return failure();
 
-  // If the attribute is a symbol reference or array, then we expect a trailing
-  // type.
+  // If the attribute is a symbol reference, then we expect a trailing type.
   Type type;
-  if (!valueAttr.isa<SymbolRefAttr, ArrayAttr>())
+  if (!valueAttr.isa<SymbolRefAttr>())
     type = valueAttr.getType();
   else if (parser.parseColonType(type))
     return failure();
 
   // Add the attribute type to the list.
   return parser.addTypeToList(type, result.types);
 }
 
 /// The constant op requires an attribute, and furthermore requires that it
 /// matches the return type.
 static LogicalResult verify(ConstantOp &op) {
   auto value = op.getValue();
   if (!value)
     return op.emitOpError("requires a 'value' attribute");
 
   Type type = op.getType();
   if (!value.getType().isa<NoneType>() && type != value.getType())
     return op.emitOpError() << "requires attribute's type (" << value.getType()
                             << ") to match op's return type (" << type << ")";
 
-  if (auto complexTy = type.dyn_cast<ComplexType>()) {
-    auto arrayAttr = value.dyn_cast<ArrayAttr>();
-    if (!complexTy || arrayAttr.size() != 2)
-      return op.emitOpError(
-          "requires 'value' to be a complex constant, represented as array of "
-          "two values");
-    auto complexEltTy = complexTy.getElementType();
-    if (complexEltTy != arrayAttr[0].getType() ||
-        complexEltTy != arrayAttr[1].getType()) {
-      return op.emitOpError()
-             << "requires attribute's element types (" << arrayAttr[0].getType()
-             << ", " << arrayAttr[1].getType()
-             << ") to match the element type of the op's return type ("
-             << complexEltTy << ")";
-    }
-    return success();
-  }
-
   if (type.isa<FunctionType>()) {
     auto fnAttr = value.dyn_cast<FlatSymbolRefAttr>();
     if (!fnAttr)
       return op.emitOpError("requires 'value' to be a function reference");
 
     // Try to find the referenced function.
     auto fn =
         op->getParentOfType<ModuleOp>().lookupSymbol<FuncOp>(fnAttr.getValue());

 }
 
 /// Returns true if a constant operation can be built with the given value and
 /// result type.
 bool ConstantOp::isBuildableWith(Attribute value, Type type) {
   // SymbolRefAttr can only be used with a function type.
   if (value.isa<SymbolRefAttr>())
     return type.isa<FunctionType>();
-  // The attribute must have the same type as 'type'.
-  if (!value.getType().isa<NoneType>() && value.getType() != type)
-    return false;
-  // Finally, check that the attribute kind is handled.
-  if (auto arrAttr = value.dyn_cast<ArrayAttr>()) {
-    auto complexTy = type.dyn_cast<ComplexType>();
-    if (!complexTy)
-      return false;
-    auto complexEltTy = complexTy.getElementType();
-    return arrAttr.size() == 2 && arrAttr[0].getType() == complexEltTy &&
-           arrAttr[1].getType() == complexEltTy;
-  }
-  return value.isa<UnitAttr>();
+  // Otherwise, this must be a UnitAttr.
+  return value.isa<UnitAttr>() && type.isa<NoneType>();
 }
 
 //===----------------------------------------------------------------------===//
 // ReturnOp
 //===----------------------------------------------------------------------===//
 
 static LogicalResult verify(ReturnOp op) {
   auto function = cast<FuncOp>(op->getParentOp());

mlir/test/Dialect/Complex/canonicalize.mlir

func @create_of_real_and_imag_different_operand(
   %complex = complex.create %real, %imag : complex<f32>
   return %complex: complex<f32>
 }
 
 // CHECK-LABEL: func @real_of_const(
 func @real_of_const() -> f32 {
   // CHECK: %[[CST:.*]] = arith.constant 1.000000e+00 : f32
   // CHECK-NEXT: return %[[CST]] : f32
-  %complex = constant [1.0 : f32, 0.0 : f32] : complex<f32>
+  %complex = complex.constant [1.0 : f32, 0.0 : f32] : complex<f32>
   %1 = complex.re %complex : complex<f32>
   return %1 : f32
 }
 
 // CHECK-LABEL: func @real_of_create_op(
 func @real_of_create_op() -> f32 {
   // CHECK: %[[CST:.*]] = arith.constant 1.000000e+00 : f32
   // CHECK-NEXT: return %[[CST]] : f32
   %real = arith.constant 1.0 : f32
   %imag = arith.constant 0.0 : f32
   %complex = complex.create %real, %imag : complex<f32>
   %1 = complex.re %complex : complex<f32>
   return %1 : f32
 }
 
 // CHECK-LABEL: func @imag_of_const(
 func @imag_of_const() -> f32 {
   // CHECK: %[[CST:.*]] = arith.constant 0.000000e+00 : f32
   // CHECK-NEXT: return %[[CST]] : f32
-  %complex = constant [1.0 : f32, 0.0 : f32] : complex<f32>
+  %complex = complex.constant [1.0 : f32, 0.0 : f32] : complex<f32>
   %1 = complex.im %complex : complex<f32>
   return %1 : f32
 }
 
 // CHECK-LABEL: func @imag_of_create_op(
 func @imag_of_create_op() -> f32 {
   // CHECK: %[[CST:.*]] = arith.constant 0.000000e+00 : f32
   // CHECK-NEXT: return %[[CST]] : f32
   %real = arith.constant 1.0 : f32
   %imag = arith.constant 0.0 : f32
   %complex = complex.create %real, %imag : complex<f32>
   %1 = complex.im %complex : complex<f32>
   return %1 : f32
 }

mlir/test/Dialect/Complex/invalid.mlir

+// RUN: mlir-opt -split-input-file %s -verify-diagnostics
+
+func @complex_constant_wrong_array_attribute_length() {
+  // expected-error @+1 {{requires 'value' to be a complex constant, represented as array of two values}}
+  %0 = complex.constant [1.0 : f32] : complex<f32>
+  return
+}
+
+// -----
+
+func @complex_constant_wrong_element_types() {
+  // expected-error @+1 {{requires attribute's element types ('f32', 'f32') to match the element type of the op's return type ('f64')}}
+  %0 = complex.constant [1.0 : f32, -1.0 : f32] : complex<f64>
+  return
+}
+
+// -----
+
+func @complex_constant_two_different_element_types() {
+  // expected-error @+1 {{requires attribute's element types ('f32', 'f64') to match the element type of the op's return type ('f64')}}
+  %0 = complex.constant [1.0 : f32, -1.0 : f64] : complex<f64>
+  return
+}

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: complex.constant [1.{{.*}}, -1.{{.*}}] : complex<f64>
+  %cst_f64 = complex.constant [0.1, -1.0] : complex<f64>
+
+  // CHECK: complex.constant [1.{{.*}} : f32, -1.{{.*}} : f32] : complex<f32>
+  %cst_f32 = complex.constant [0.1 : f32, -1.0 : f32] : complex<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>
func @ops(%f: f32) {
 
   // CHECK: complex.sign %[[C]] : complex<f32>
   %sign = complex.sign %complex : complex<f32>
 
   // CHECK: complex.sub %[[C]], %[[C]] : complex<f32>
   %diff = complex.sub %complex, %complex : complex<f32>
   return
 }
-

mlir/test/Dialect/Standard/invalid.mlir

 // RUN: mlir-opt -split-input-file %s -verify-diagnostics
 
 func @unsupported_attribute() {
   // expected-error @+1 {{unsupported 'value' attribute: "" : index}}
   %0 = constant "" : index
   return
 }
 
 // -----
 
-func @complex_constant_wrong_array_attribute_length() {
-  // expected-error @+1 {{requires 'value' to be a complex constant, represented as array of two values}}
-  %0 = constant [1.0 : f32] : complex<f32>
-  return
-}
-
-// -----
-
-func @complex_constant_wrong_element_types() {
-  // expected-error @+1 {{requires attribute's element types ('f32', 'f32') to match the element type of the op's return type ('f64')}}
-  %0 = constant [1.0 : f32, -1.0 : f32] : complex<f64>
-  return
-}
-
-// -----
-
-func @complex_constant_two_different_element_types() {
-  // expected-error @+1 {{requires attribute's element types ('f32', 'f64') to match the element type of the op's return type ('f64')}}
-  %0 = constant [1.0 : f32, -1.0 : f64] : complex<f64>
-  return
-}
-
-// -----
-
 func @return_i32_f32() -> (i32, f32) {
   %0 = arith.constant 1 : i32
   %1 = arith.constant 1. : f32
   return %0, %1 : i32, f32
 }
 
 func @call() {
   // expected-error @+3 {{op result type mismatch at index 0}}
   // expected-note @+2 {{op result types: 'f32', 'i32'}}
   // expected-note @+1 {{function result types: 'i32', 'f32'}}
   %0:2 = call @return_i32_f32() : () -> (f32, i32)
   return
 }

mlir/test/Dialect/Standard/ops.mlir

func @switch_i64(%flag : i64, %caseOperand : i32) {
   ^bb1(%bb1arg : i32):
     return
   ^bb2(%bb2arg : i32):
     return
   ^bb3(%bb3arg : i32):
     return
 }
 
-// CHECK-LABEL: func @constant_complex_f32(
-func @constant_complex_f32() -> complex<f32> {
-  %result = constant [0.1 : f32, -1.0 : f32] : complex<f32>
-  return %result : complex<f32>
-}
-
-// CHECK-LABEL: func @constant_complex_f64(
-func @constant_complex_f64() -> complex<f64> {
-  %result = constant [0.1 : f64, -1.0 : f64] : complex<f64>
-  return %result : complex<f64>
-}
-
 // CHECK-LABEL: func @vector_splat_0d(
 func @vector_splat_0d(%a: f32) -> vector<f32> {
   // CHECK: splat %{{.*}} : vector<f32>
   %0 = splat %a : vector<f32>
   return %0 : vector<f32>
 }