[CIR][Transforms] Add folders for complex operations (#775)

Lancern · lanza · commit 4ebe6f1a9f2d · 2025-03-13T17:53:47.000-07:00
This PR adds folders for `cir.complex.create`, `cir.complex.real`, and `cir.complex.imag`. This PR adds a new attribute `#cir.complex` that represents a constant complex value. Besides, the CIR dialect does not have a constant materializer yet; this PR adds it. Address #726 .
diff --git a/clang/include/clang/CIR/Dialect/IR/CIRAttrs.td b/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
 //===----------------------------------------------------------------------===//
diff --git a/clang/include/clang/CIR/Dialect/IR/CIRDialect.td b/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.
diff --git a/clang/include/clang/CIR/Dialect/IR/CIROps.td b/clang/include/clang/CIR/Dialect/IR/CIROps.td
@@ -1243,6 +1243,7 @@ def ComplexCreateOp : CIR_Op<"complex.create", [Pure, SameTypeOperands]> {
   }];
 
   let hasVerifier = 1;
+  let hasFolder = 1;
 }
 
 //===----------------------------------------------------------------------===//
@@ -1271,6 +1272,7 @@ def ComplexRealOp : CIR_Op<"complex.real", [Pure]> {
   }];
 
   let hasVerifier = 1;
+  let hasFolder = 1;
 }
 
 def ComplexImagOp : CIR_Op<"complex.imag", [Pure]> {
@@ -1295,6 +1297,7 @@ def ComplexImagOp : CIR_Op<"complex.imag", [Pure]> {
   }];
 
   let hasVerifier = 1;
+  let hasFolder = 1;
 }
 
 //===----------------------------------------------------------------------===//
diff --git a/clang/lib/CIR/Dialect/IR/CIRAttrs.cpp b/clang/lib/CIR/Dialect/IR/CIRAttrs.cpp
@@ -365,6 +365,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
 //===----------------------------------------------------------------------===//
diff --git a/clang/lib/CIR/Dialect/IR/CIRDialect.cpp b/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 (")
@@ -748,6 +757,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
 //===----------------------------------------------------------------------===//
@@ -760,6 +789,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";
@@ -768,6 +805,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
 //===----------------------------------------------------------------------===//
diff --git a/clang/lib/CIR/Dialect/Transforms/CIRSimplify.cpp b/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);
   });
 
diff --git a/clang/lib/CIR/Lowering/DirectToLLVM/LowerToLLVM.cpp b/clang/lib/CIR/Lowering/DirectToLLVM/LowerToLLVM.cpp
@@ -1180,6 +1180,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())) {
diff --git a/clang/test/CIR/CodeGen/complex.c b/clang/test/CIR/CodeGen/complex.c
diff --git a/clang/test/CIR/Lowering/complex.cir b/clang/test/CIR/Lowering/complex.cir
diff --git a/clang/test/CIR/Transforms/complex-fold.cir b/clang/test/CIR/Transforms/complex-fold.cir

Original file line number	Diff line number	Diff line change
`@@ -1243,6 +1243,7 @@ def ComplexCreateOp : CIR_Op<"complex.create", [Pure, SameTypeOperands]> {`
`1243`	`1243`	`}];`
`1244`	`1244`
`1245`	`1245`	`let hasVerifier = 1;`
	`1246`	`+ let hasFolder = 1;`
`1246`	`1247`	`}`
`1247`	`1248`
`1248`	`1249`	`//===----------------------------------------------------------------------===//`
`@@ -1271,6 +1272,7 @@ def ComplexRealOp : CIR_Op<"complex.real", [Pure]> {`
`1271`	`1272`	`}];`
`1272`	`1273`
`1273`	`1274`	`let hasVerifier = 1;`
	`1275`	`+ let hasFolder = 1;`
`1274`	`1276`	`}`
`1275`	`1277`
`1276`	`1278`	`def ComplexImagOp : CIR_Op<"complex.imag", [Pure]> {`
`@@ -1295,6 +1297,7 @@ def ComplexImagOp : CIR_Op<"complex.imag", [Pure]> {`
`1295`	`1297`	`}];`
`1296`	`1298`
`1297`	`1299`	`let hasVerifier = 1;`
	`1300`	`+ let hasFolder = 1;`
`1298`	`1301`	`}`
`1299`	`1302`
`1300`	`1303`	`//===----------------------------------------------------------------------===//`