[CIR] Upstream initial support for union type (llvm#137501)

Closes llvm#136059

Author: el-ev

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

@@ -531,6 +531,7 @@ def CIR_RecordType : CIR_Type<"Record", "record",
     bool isComplete() const { return !isIncomplete(); };
     bool isIncomplete() const;
 
+    mlir::Type getLargestMember(const mlir::DataLayout &dataLayout) const;
     size_t getNumElements() const { return getMembers().size(); };
     std::string getKindAsStr() {
       switch (getKind()) {

clang/lib/CIR/CodeGen/CIRGenExpr.cpp

@@ -317,20 +317,21 @@ LValue CIRGenFunction::emitLValueForField(LValue base, const FieldDecl *field) {
   }
 
   unsigned recordCVR = base.getVRQualifiers();
-  if (rec->isUnion()) {
-    cgm.errorNYI(field->getSourceRange(), "emitLValueForField: union");
-    return LValue();
-  }
 
-  assert(!cir::MissingFeatures::preservedAccessIndexRegion());
   llvm::StringRef fieldName = field->getName();
-  const CIRGenRecordLayout &layout =
-      cgm.getTypes().getCIRGenRecordLayout(field->getParent());
-  unsigned fieldIndex = layout.getCIRFieldNo(field);
-
+  unsigned fieldIndex;
   assert(!cir::MissingFeatures::lambdaFieldToName());
 
+  if (rec->isUnion())
+    fieldIndex = field->getFieldIndex();
+  else {
+    const CIRGenRecordLayout &layout =
+        cgm.getTypes().getCIRGenRecordLayout(field->getParent());
+    fieldIndex = layout.getCIRFieldNo(field);
+  }
+
   addr = emitAddrOfFieldStorage(addr, field, fieldName, fieldIndex);
+  assert(!cir::MissingFeatures::preservedAccessIndexRegion());
 
   // If this is a reference field, load the reference right now.
   if (fieldType->isReferenceType()) {

clang/lib/CIR/CodeGen/CIRGenExprAggregate.cpp

@@ -1,4 +1,4 @@
-//===--- CIRGenExprAgg.cpp - Emit CIR Code from Aggregate Expressions -----===//
+//===- CIRGenExprAggregrate.cpp - Emit CIR Code from Aggregate Expressions ===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.

clang/lib/CIR/CodeGen/CIRGenRecordLayout.h

@@ -31,7 +31,7 @@ class CIRGenRecordLayout {
 
   /// Map from (non-bit-field) record field to the corresponding cir record type
   /// field no. This info is populated by the record builder.
-  llvm::DenseMap<const clang::FieldDecl *, unsigned> fieldInfo;
+  llvm::DenseMap<const clang::FieldDecl *, unsigned> fieldIdxMap;
 
 public:
   CIRGenRecordLayout(cir::RecordType completeObjectType)
@@ -44,8 +44,8 @@ class CIRGenRecordLayout {
   /// Return cir::RecordType element number that corresponds to the field FD.
   unsigned getCIRFieldNo(const clang::FieldDecl *fd) const {
     fd = fd->getCanonicalDecl();
-    assert(fieldInfo.count(fd) && "Invalid field for record!");
-    return fieldInfo.lookup(fd);
+    assert(fieldIdxMap.count(fd) && "Invalid field for record!");
+    return fieldIdxMap.lookup(fd);
   }
 };
 

clang/lib/CIR/CodeGen/CIRGenRecordLayoutBuilder.cpp

@@ -56,14 +56,15 @@ struct CIRRecordLowering final {
   };
   // The constructor.
   CIRRecordLowering(CIRGenTypes &cirGenTypes, const RecordDecl *recordDecl,
-                    bool isPacked);
+                    bool packed);
 
   /// Constructs a MemberInfo instance from an offset and mlir::Type.
   MemberInfo makeStorageInfo(CharUnits offset, mlir::Type data) {
     return MemberInfo(offset, MemberInfo::InfoKind::Field, data);
   }
 
   void lower();
+  void lowerUnion();
 
   /// Determines if we need a packed llvm struct.
   void determinePacked();
@@ -83,6 +84,10 @@ struct CIRRecordLowering final {
     return CharUnits::fromQuantity(dataLayout.layout.getTypeABIAlignment(Ty));
   }
 
+  bool isZeroInitializable(const FieldDecl *fd) {
+    return cirGenTypes.isZeroInitializable(fd->getType());
+  }
+
   /// Wraps cir::IntType with some implicit arguments.
   mlir::Type getUIntNType(uint64_t numBits) {
     unsigned alignedBits = llvm::PowerOf2Ceil(numBits);
@@ -121,6 +126,13 @@ struct CIRRecordLowering final {
   /// Fills out the structures that are ultimately consumed.
   void fillOutputFields();
 
+  void appendPaddingBytes(CharUnits size) {
+    if (!size.isZero()) {
+      fieldTypes.push_back(getByteArrayType(size));
+      padded = true;
+    }
+  }
+
   CIRGenTypes &cirGenTypes;
   CIRGenBuilderTy &builder;
   const ASTContext &astContext;
@@ -130,12 +142,14 @@ struct CIRRecordLowering final {
   std::vector<MemberInfo> members;
   // Output fields, consumed by CIRGenTypes::computeRecordLayout
   llvm::SmallVector<mlir::Type, 16> fieldTypes;
-  llvm::DenseMap<const FieldDecl *, unsigned> fields;
+  llvm::DenseMap<const FieldDecl *, unsigned> fieldIdxMap;
   cir::CIRDataLayout dataLayout;
 
   LLVM_PREFERRED_TYPE(bool)
   unsigned zeroInitializable : 1;
   LLVM_PREFERRED_TYPE(bool)
+  unsigned zeroInitializableAsBase : 1;
+  LLVM_PREFERRED_TYPE(bool)
   unsigned packed : 1;
   LLVM_PREFERRED_TYPE(bool)
   unsigned padded : 1;
@@ -147,19 +161,19 @@ struct CIRRecordLowering final {
 } // namespace
 
 CIRRecordLowering::CIRRecordLowering(CIRGenTypes &cirGenTypes,
-                                     const RecordDecl *recordDecl,
-                                     bool isPacked)
+                                     const RecordDecl *recordDecl, bool packed)
     : cirGenTypes(cirGenTypes), builder(cirGenTypes.getBuilder()),
       astContext(cirGenTypes.getASTContext()), recordDecl(recordDecl),
       astRecordLayout(
           cirGenTypes.getASTContext().getASTRecordLayout(recordDecl)),
       dataLayout(cirGenTypes.getCGModule().getModule()),
-      zeroInitializable(true), packed(isPacked), padded(false) {}
+      zeroInitializable(true), zeroInitializableAsBase(true), packed(packed),
+      padded(false) {}
 
 void CIRRecordLowering::lower() {
   if (recordDecl->isUnion()) {
-    cirGenTypes.getCGModule().errorNYI(recordDecl->getSourceRange(),
-                                       "lower: union");
+    lowerUnion();
+    assert(!cir::MissingFeatures::bitfields());
     return;
   }
 
@@ -194,7 +208,8 @@ void CIRRecordLowering::fillOutputFields() {
       fieldTypes.push_back(member.data);
     if (member.kind == MemberInfo::InfoKind::Field) {
       if (member.fieldDecl)
-        fields[member.fieldDecl->getCanonicalDecl()] = fieldTypes.size() - 1;
+        fieldIdxMap[member.fieldDecl->getCanonicalDecl()] =
+            fieldTypes.size() - 1;
       // A field without storage must be a bitfield.
       assert(!cir::MissingFeatures::bitfields());
     }
@@ -296,7 +311,7 @@ CIRGenTypes::computeRecordLayout(const RecordDecl *rd, cir::RecordType *ty) {
   assert(!cir::MissingFeatures::bitfields());
 
   // Add all the field numbers.
-  rl->fieldInfo.swap(lowering.fields);
+  rl->fieldIdxMap.swap(lowering.fieldIdxMap);
 
   // Dump the layout, if requested.
   if (getASTContext().getLangOpts().DumpRecordLayouts) {
@@ -306,3 +321,68 @@ CIRGenTypes::computeRecordLayout(const RecordDecl *rd, cir::RecordType *ty) {
   // TODO: implement verification
   return rl;
 }
+
+void CIRRecordLowering::lowerUnion() {
+  CharUnits layoutSize = astRecordLayout.getSize();
+  mlir::Type storageType = nullptr;
+  bool seenNamedMember = false;
+
+  // Iterate through the fields setting bitFieldInfo and the Fields array. Also
+  // locate the "most appropriate" storage type.
+  for (const FieldDecl *field : recordDecl->fields()) {
+    mlir::Type fieldType;
+    if (field->isBitField())
+      cirGenTypes.getCGModule().errorNYI(recordDecl->getSourceRange(),
+                                         "bitfields in lowerUnion");
+    else
+      fieldType = getStorageType(field);
+
+    // This maps a field to its index. For unions, the index is always 0.
+    fieldIdxMap[field->getCanonicalDecl()] = 0;
+
+    // Compute zero-initializable status.
+    // This union might not be zero initialized: it may contain a pointer to
+    // data member which might have some exotic initialization sequence.
+    // If this is the case, then we ought not to try and come up with a "better"
+    // type, it might not be very easy to come up with a Constant which
+    // correctly initializes it.
+    if (!seenNamedMember) {
+      seenNamedMember = field->getIdentifier();
+      if (!seenNamedMember)
+        if (const RecordDecl *fieldRD = field->getType()->getAsRecordDecl())
+          seenNamedMember = fieldRD->findFirstNamedDataMember();
+      if (seenNamedMember && !isZeroInitializable(field)) {
+        zeroInitializable = zeroInitializableAsBase = false;
+        storageType = fieldType;
+      }
+    }
+
+    // Because our union isn't zero initializable, we won't be getting a better
+    // storage type.
+    if (!zeroInitializable)
+      continue;
+
+    // Conditionally update our storage type if we've got a new "better" one.
+    if (!storageType || getAlignment(fieldType) > getAlignment(storageType) ||
+        (getAlignment(fieldType) == getAlignment(storageType) &&
+         getSize(fieldType) > getSize(storageType)))
+      storageType = fieldType;
+
+    // NOTE(cir): Track all union member's types, not just the largest one. It
+    // allows for proper type-checking and retain more info for analisys.
+    fieldTypes.push_back(fieldType);
+  }
+
+  if (!storageType)
+    cirGenTypes.getCGModule().errorNYI(recordDecl->getSourceRange(),
+                                       "No-storage Union NYI");
+
+  if (layoutSize < getSize(storageType))
+    storageType = getByteArrayType(layoutSize);
+  else
+    appendPaddingBytes(layoutSize - getSize(storageType));
+
+  // Set packed if we need it.
+  if (layoutSize % getAlignment(storageType))
+    packed = true;
+}

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

@@ -230,17 +230,34 @@ void RecordType::complete(ArrayRef<Type> members, bool packed, bool padded) {
     llvm_unreachable("failed to complete record");
 }
 
+/// Return the largest member of in the type.
+///
+/// Recurses into union members never returning a union as the largest member.
+Type RecordType::getLargestMember(const ::mlir::DataLayout &dataLayout) const {
+  assert(isUnion() && "Only call getLargestMember on unions");
+  llvm::ArrayRef<Type> members = getMembers();
+  // If the union is padded, we need to ignore the last member,
+  // which is the padding.
+  return *std::max_element(
+      members.begin(), getPadded() ? members.end() - 1 : members.end(),
+      [&](Type lhs, Type rhs) {
+        return dataLayout.getTypeABIAlignment(lhs) <
+                   dataLayout.getTypeABIAlignment(rhs) ||
+               (dataLayout.getTypeABIAlignment(lhs) ==
+                    dataLayout.getTypeABIAlignment(rhs) &&
+                dataLayout.getTypeSize(lhs) < dataLayout.getTypeSize(rhs));
+      });
+}
+
 //===----------------------------------------------------------------------===//
 // Data Layout information for types
 //===----------------------------------------------------------------------===//
 
 llvm::TypeSize
 RecordType::getTypeSizeInBits(const mlir::DataLayout &dataLayout,
                               mlir::DataLayoutEntryListRef params) const {
-  if (isUnion()) {
-    // TODO(CIR): Implement union layout.
-    return llvm::TypeSize::getFixed(8);
-  }
+  if (isUnion())
+    return dataLayout.getTypeSize(getLargestMember(dataLayout));
 
   unsigned recordSize = computeStructSize(dataLayout);
   return llvm::TypeSize::getFixed(recordSize * 8);
@@ -249,10 +266,8 @@ RecordType::getTypeSizeInBits(const mlir::DataLayout &dataLayout,
 uint64_t
 RecordType::getABIAlignment(const ::mlir::DataLayout &dataLayout,
                             ::mlir::DataLayoutEntryListRef params) const {
-  if (isUnion()) {
-    // TODO(CIR): Implement union layout.
-    return 8;
-  }
+  if (isUnion())
+    return dataLayout.getTypeABIAlignment(getLargestMember(dataLayout));
 
   // Packed structures always have an ABI alignment of 1.
   if (getPacked())
@@ -268,8 +283,6 @@ RecordType::computeStructSize(const mlir::DataLayout &dataLayout) const {
   unsigned recordSize = 0;
   uint64_t recordAlignment = 1;
 
-  // We can't use a range-based for loop here because we might be ignoring the
-  // last element.
   for (mlir::Type ty : getMembers()) {
     // This assumes that we're calculating size based on the ABI alignment, not
     // the preferred alignment for each type.

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

@@ -1436,7 +1436,14 @@ static void prepareTypeConverter(mlir::LLVMTypeConverter &converter,
       break;
     // Unions are lowered as only the largest member.
     case cir::RecordType::Union:
-      llvm_unreachable("Lowering of unions is NYI");
+      if (auto largestMember = type.getLargestMember(dataLayout))
+        llvmMembers.push_back(
+            convertTypeForMemory(converter, dataLayout, largestMember));
+      if (type.getPadded()) {
+        auto last = *type.getMembers().rbegin();
+        llvmMembers.push_back(
+            convertTypeForMemory(converter, dataLayout, last));
+      }
       break;
     }
 
@@ -1609,7 +1616,11 @@ mlir::LogicalResult CIRToLLVMGetMemberOpLowering::matchAndRewrite(
     return mlir::success();
   }
   case cir::RecordType::Union:
-    return op.emitError() << "NYI: union get_member lowering";
+    // Union members share the address space, so we just need a bitcast to
+    // conform to type-checking.
+    rewriter.replaceOpWithNewOp<mlir::LLVM::BitcastOp>(op, llResTy,
+                                                       adaptor.getAddr());
+    return mlir::success();
   }
 }