[CIR][Transforms] Add folders for complex operations

clang/include/clang/CIR/Dialect/IR/CIRAttrs.td

@@ -274,6 +274,40 @@ def FPAttr : CIR_Attr<"FP", "fp", [TypedAttrInterface]> {
   }];
 }
 
+//===----------------------------------------------------------------------===//
+// ComplexAttr
+//===----------------------------------------------------------------------===//
+
+def ComplexAttr : CIR_Attr<"Complex", "complex", [TypedAttrInterface]> {
+  let summary = "An attribute that contains a constant complex value";
+  let description = [{
+    The `#cir.complex` attribute contains a constant value of complex number
+    type. The `real` parameter gives the real part of the complex number and the
+    `imag` parameter gives the imaginary part of the complex number.
+
+    The `real` and `imag` parameter must be either an IntAttr or an FPAttr that
+    contains values of the same CIR type.
+  }];
+
+  let parameters = (ins 
+    AttributeSelfTypeParameter<"", "mlir::cir::ComplexType">:$type,
+    "mlir::TypedAttr":$real, "mlir::TypedAttr":$imag);
+
+  let builders = [
+    AttrBuilderWithInferredContext<(ins "mlir::cir::ComplexType":$type,
+                                        "mlir::TypedAttr":$real,
+                                        "mlir::TypedAttr":$imag), [{
+      return $_get(type.getContext(), type, real, imag);
+    }]>,
+  ];
+
+  let genVerifyDecl = 1;
+
+  let assemblyFormat = [{
+    `<` qualified($real) `,` qualified($imag) `>`
+  }];
+}
+
 //===----------------------------------------------------------------------===//
 // ConstPointerAttr
 //===----------------------------------------------------------------------===//

clang/include/clang/CIR/Dialect/IR/CIRDialect.td

@@ -27,6 +27,8 @@ def CIR_Dialect : Dialect {
   let useDefaultAttributePrinterParser = 0;
   let useDefaultTypePrinterParser = 0;
 
+  let hasConstantMaterializer = 1;
+
   let extraClassDeclaration = [{
 
     // Names of CIR parameter attributes.

clang/include/clang/CIR/Dialect/IR/CIROps.td

@@ -1217,6 +1217,7 @@ def ComplexCreateOp : CIR_Op<"complex.create", [Pure, SameTypeOperands]> {
   }];
 
   let hasVerifier = 1;
+  let hasFolder = 1;
 }
 
 //===----------------------------------------------------------------------===//
@@ -1245,6 +1246,7 @@ def ComplexRealOp : CIR_Op<"complex.real", [Pure]> {
   }];
 
   let hasVerifier = 1;
+  let hasFolder = 1;
 }
 
 def ComplexImagOp : CIR_Op<"complex.imag", [Pure]> {
@@ -1269,6 +1271,7 @@ def ComplexImagOp : CIR_Op<"complex.imag", [Pure]> {
   }];
 
   let hasVerifier = 1;
+  let hasFolder = 1;
 }
 
 //===----------------------------------------------------------------------===//

clang/lib/CIR/Dialect/IR/CIRAttrs.cpp

@@ -362,6 +362,26 @@ LogicalResult cir::FPAttr::verify(function_ref<InFlightDiagnostic()> emitError,
   return success();
 }
 
+//===----------------------------------------------------------------------===//
+// ComplexAttr definitions
+//===----------------------------------------------------------------------===//
+
+LogicalResult ComplexAttr::verify(function_ref<InFlightDiagnostic()> emitError,
+                                  mlir::cir::ComplexType type,
+                                  mlir::TypedAttr real, mlir::TypedAttr imag) {
+  auto elemTy = type.getElementTy();
+  if (real.getType() != elemTy) {
+    emitError() << "type of the real part does not match the complex type";
+    return failure();
+  }
+  if (imag.getType() != elemTy) {
+    emitError() << "type of the imaginary part does not match the complex type";
+    return failure();
+  }
+
+  return success();
+}
+
 //===----------------------------------------------------------------------===//
 // CmpThreeWayInfoAttr definitions
 //===----------------------------------------------------------------------===//

clang/lib/CIR/Dialect/IR/CIRDialect.cpp

@@ -124,6 +124,14 @@ void cir::CIRDialect::initialize() {
   addInterfaces<CIROpAsmDialectInterface>();
 }
 
+Operation *cir::CIRDialect::materializeConstant(mlir::OpBuilder &builder,
+                                                mlir::Attribute value,
+                                                mlir::Type type,
+                                                mlir::Location loc) {
+  return builder.create<mlir::cir::ConstantOp>(
+      loc, type, mlir::cast<mlir::TypedAttr>(value));
+}
+
 //===----------------------------------------------------------------------===//
 // Helpers
 //===----------------------------------------------------------------------===//
@@ -344,7 +352,8 @@ static LogicalResult checkConstantTypes(mlir::Operation *op, mlir::Type opType,
     return success();
   }
 
-  if (mlir::isa<mlir::cir::IntAttr, mlir::cir::FPAttr>(attrType)) {
+  if (mlir::isa<mlir::cir::IntAttr, mlir::cir::FPAttr, mlir::cir::ComplexAttr>(
+          attrType)) {
     auto at = cast<TypedAttr>(attrType);
     if (at.getType() != opType) {
       return op->emitOpError("result type (")
@@ -638,6 +647,26 @@ LogicalResult ComplexCreateOp::verify() {
   return success();
 }
 
+OpFoldResult ComplexCreateOp::fold(FoldAdaptor adaptor) {
+  auto real = adaptor.getReal();
+  auto imag = adaptor.getImag();
+
+  if (!real || !imag)
+    return nullptr;
+
+  // When both of real and imag are constants, we can fold the operation into an
+  // `cir.const #cir.complex` operation.
+
+  auto realAttr = mlir::cast<mlir::TypedAttr>(real);
+  auto imagAttr = mlir::cast<mlir::TypedAttr>(imag);
+  assert(realAttr.getType() == imagAttr.getType() &&
+         "real part and imag part should be of the same type");
+
+  auto complexTy =
+      mlir::cir::ComplexType::get(getContext(), realAttr.getType());
+  return mlir::cir::ComplexAttr::get(complexTy, realAttr, imagAttr);
+}
+
 //===----------------------------------------------------------------------===//
 // ComplexRealOp and ComplexImagOp
 //===----------------------------------------------------------------------===//
@@ -650,6 +679,14 @@ LogicalResult ComplexRealOp::verify() {
   return success();
 }
 
+OpFoldResult ComplexRealOp::fold(FoldAdaptor adaptor) {
+  auto input =
+      mlir::cast_if_present<mlir::cir::ComplexAttr>(adaptor.getOperand());
+  if (input)
+    return input.getReal();
+  return nullptr;
+}
+
 LogicalResult ComplexImagOp::verify() {
   if (getType() != getOperand().getType().getElementTy()) {
     emitOpError() << "cir.complex.imag result type does not match operand type";
@@ -658,6 +695,14 @@ LogicalResult ComplexImagOp::verify() {
   return success();
 }
 
+OpFoldResult ComplexImagOp::fold(FoldAdaptor adaptor) {
+  auto input =
+      mlir::cast_if_present<mlir::cir::ComplexAttr>(adaptor.getOperand());
+  if (input)
+    return input.getImag();
+  return nullptr;
+}
+
 //===----------------------------------------------------------------------===//
 // ComplexRealPtrOp and ComplexImagPtrOp
 //===----------------------------------------------------------------------===//

clang/lib/CIR/Dialect/Transforms/CIRSimplify.cpp

@@ -146,7 +146,8 @@ void CIRSimplifyPass::runOnOperation() {
   getOperation()->walk([&](Operation *op) {
     // CastOp here is to perform a manual `fold` in
     // applyOpPatternsAndFold
-    if (isa<BrOp, BrCondOp, ScopeOp, SwitchOp, CastOp, TryOp, UnaryOp>(op))
+    if (isa<BrOp, BrCondOp, ScopeOp, SwitchOp, CastOp, TryOp, UnaryOp,
+            ComplexCreateOp, ComplexRealOp, ComplexImagOp>(op))
       ops.push_back(op);
   });
 

clang/lib/CIR/Lowering/DirectToLLVM/LowerToLLVM.cpp

@@ -1179,6 +1179,30 @@ class CIRConstantLowering
       attr = rewriter.getFloatAttr(
           typeConverter->convertType(op.getType()),
           mlir::cast<mlir::cir::FPAttr>(op.getValue()).getValue());
+    } else if (auto complexTy =
+                   mlir::dyn_cast<mlir::cir::ComplexType>(op.getType())) {
+      auto complexAttr = mlir::cast<mlir::cir::ComplexAttr>(op.getValue());
+      auto complexElemTy = complexTy.getElementTy();
+      auto complexElemLLVMTy = typeConverter->convertType(complexElemTy);
+
+      mlir::Attribute components[2];
+      if (mlir::isa<mlir::cir::IntType>(complexElemTy)) {
+        components[0] = rewriter.getIntegerAttr(
+            complexElemLLVMTy,
+            mlir::cast<mlir::cir::IntAttr>(complexAttr.getReal()).getValue());
+        components[1] = rewriter.getIntegerAttr(
+            complexElemLLVMTy,
+            mlir::cast<mlir::cir::IntAttr>(complexAttr.getImag()).getValue());
+      } else {
+        components[0] = rewriter.getFloatAttr(
+            complexElemLLVMTy,
+            mlir::cast<mlir::cir::FPAttr>(complexAttr.getReal()).getValue());
+        components[1] = rewriter.getFloatAttr(
+            complexElemLLVMTy,
+            mlir::cast<mlir::cir::FPAttr>(complexAttr.getImag()).getValue());
+      }
+
+      attr = rewriter.getArrayAttr(components);
     } else if (mlir::isa<mlir::cir::PointerType>(op.getType())) {
       // Optimize with dedicated LLVM op for null pointers.
       if (mlir::isa<mlir::cir::ConstPtrAttr>(op.getValue())) {