[CIR] Fix Address element type problems (#1373)

There were problems with the pointer type and element type of the
Address class getting out of sync. In the traditional codegen the
pointer has no type, so it was sufficient for the Address class to
simply track the type it was supposed to be pointing to. Since ClangIR
pointer values are typed, the Address::withType function wasn't really
doing what it was supposed to. It returned an object with the same
pointer that the original object had, but with a mismatched element
type.

This change updates the Address::withType function to perform a bitcast
to get the expected pointer type before creating a new Address object.
It also adds assertions in the Address class to verify that pointer type
and element type are consistent and updates many places that were
causing those assertions to fire.

These code changes cause extra bitcasts to be emitted in a few places.
Regression tests have been updated as needed to reflect the CIR that is
now generated.

Author: andykaylor

clang/lib/CIR/CodeGen/Address.h

@@ -25,6 +25,9 @@
 
 namespace clang::CIRGen {
 
+// Forward declaration to avoid a circular dependency
+class CIRGenBuilderTy;
+
 // Indicates whether a pointer is known not to be null.
 enum KnownNonNull_t { NotKnownNonNull, KnownNonNull };
 
@@ -64,6 +67,9 @@ class Address {
     assert(pointer && "Pointer cannot be null");
     assert(elementType && "Element type cannot be null");
     assert(!alignment.isZero() && "Alignment cannot be zero");
+
+    assert(mlir::cast<cir::PointerType>(pointer.getType()).getPointee() ==
+           ElementType);
   }
 
   Address(mlir::Value basePtr, mlir::Type elementType,
@@ -104,15 +110,9 @@ class Address {
 
   bool hasOffset() const { return bool(offset); }
 
-  /// Return address with different element type, but same pointer and
-  /// alignment.
-  Address withElementType(mlir::Type ElemTy) const {
-    if (!hasOffset())
-      return Address(getBasePointer(), ElemTy, getAlignment(),
-                     getPointerAuthInfo(), /*Offset=*/nullptr,
-                     isKnownNonNull());
-    return Address(getPointer(), ElemTy, getAlignment(), isKnownNonNull());
-  }
+  /// Return address with different element type, a bitcast pointer, and
+  /// the same alignment.
+  Address withElementType(CIRGenBuilderTy &builder, mlir::Type ElemTy) const;
 
   mlir::Value getPointer() const {
     assert(isValid());
@@ -142,11 +142,17 @@ class Address {
 
   /// Return the type of the pointer value.
   cir::PointerType getType() const {
+    assert(mlir::cast<cir::PointerType>(
+               PointerAndKnownNonNull.getPointer().getType())
+               .getPointee() == ElementType);
     return mlir::cast<cir::PointerType>(getPointer().getType());
   }
 
   mlir::Type getElementType() const {
     assert(isValid());
+    assert(mlir::cast<cir::PointerType>(
+               PointerAndKnownNonNull.getPointer().getType())
+               .getPointee() == ElementType);
     return ElementType;
   }
 

clang/lib/CIR/CodeGen/CIRAsm.cpp

@@ -214,8 +214,9 @@ std::pair<mlir::Value, mlir::Type> CIRGenFunction::emitAsmInputLValue(
         getTargetHooks().isScalarizableAsmOperand(*this, Ty)) {
       Ty = cir::IntType::get(&getMLIRContext(), Size, false);
 
-      return {builder.createLoad(getLoc(Loc),
-                                 InputValue.getAddress().withElementType(Ty)),
+      return {builder.createLoad(
+                  getLoc(Loc),
+                  InputValue.getAddress().withElementType(builder, Ty)),
               mlir::Type()};
     }
   }
@@ -320,7 +321,7 @@ static void emitAsmStores(CIRGenFunction &CGF, const AsmStmt &S,
     // ResultTypeRequiresCast.size() elements of RegResults.
     if ((i < ResultTypeRequiresCast.size()) && ResultTypeRequiresCast[i]) {
       unsigned Size = CGF.getContext().getTypeSize(ResultRegQualTys[i]);
-      Address A = Dest.getAddress().withElementType(ResultRegTypes[i]);
+      Address A = Dest.getAddress().withElementType(Builder, ResultRegTypes[i]);
       if (CGF.getTargetHooks().isScalarizableAsmOperand(CGF, TruncTy)) {
         Builder.createStore(CGF.getLoc(S.getAsmLoc()), Tmp, A);
         continue;
@@ -478,7 +479,8 @@ mlir::LogicalResult CIRGenFunction::emitAsmStmt(const AsmStmt &S) {
       // Otherwise there will be a mis-match if the matrix is also an
       // input-argument which is represented as vector.
       if (isa<MatrixType>(OutExpr->getType().getCanonicalType()))
-        DestAddr = DestAddr.withElementType(convertType(OutExpr->getType()));
+        DestAddr =
+            DestAddr.withElementType(builder, convertType(OutExpr->getType()));
 
       ArgTypes.push_back(DestAddr.getType());
       ArgElemTypes.push_back(DestAddr.getElementType());

clang/lib/CIR/CodeGen/CIRGenAtomic.cpp

@@ -305,7 +305,7 @@ Address AtomicInfo::castToAtomicIntPointer(Address addr) const {
   if (intTy && intTy.getWidth() == AtomicSizeInBits)
     return addr;
   auto ty = CGF.getBuilder().getUIntNTy(AtomicSizeInBits);
-  return addr.withElementType(ty);
+  return addr.withElementType(CGF.getBuilder(), ty);
 }
 
 Address AtomicInfo::convertToAtomicIntPointer(Address Addr) const {
@@ -1243,8 +1243,9 @@ RValue CIRGenFunction::emitAtomicExpr(AtomicExpr *E) {
     if (RValTy->isVoidType())
       return RValue::get(nullptr);
 
-    return convertTempToRValue(Dest.withElementType(convertTypeForMem(RValTy)),
-                               RValTy, E->getExprLoc());
+    return convertTempToRValue(
+        Dest.withElementType(builder, convertTypeForMem(RValTy)), RValTy,
+        E->getExprLoc());
   }
 
   // The memory order is not known at compile-time.  The atomic operations
@@ -1321,8 +1322,10 @@ RValue CIRGenFunction::emitAtomicExpr(AtomicExpr *E) {
 
   if (RValTy->isVoidType())
     return RValue::get(nullptr);
-  return convertTempToRValue(Dest.withElementType(convertTypeForMem(RValTy)),
-                             RValTy, E->getExprLoc());
+
+  return convertTempToRValue(
+      Dest.withElementType(builder, convertTypeForMem(RValTy)), RValTy,
+      E->getExprLoc());
 }
 
 void CIRGenFunction::emitAtomicStore(RValue rvalue, LValue lvalue,

clang/lib/CIR/CodeGen/CIRGenBuilder.cpp

@@ -133,4 +133,16 @@ uint64_t CIRGenBuilderTy::computeOffsetFromGlobalViewIndices(
   }
 
   return offset;
-}
+}
+
+// This can't be defined in Address.h because that file is included by
+// CIRGenBuilder.h
+Address Address::withElementType(CIRGenBuilderTy &builder,
+                                 mlir::Type ElemTy) const {
+  if (!hasOffset())
+    return Address(builder.createPtrBitcast(getBasePointer(), ElemTy), ElemTy,
+                   getAlignment(), getPointerAuthInfo(), /*Offset=*/nullptr,
+                   isKnownNonNull());
+  return Address(builder.createPtrBitcast(getPointer(), ElemTy), ElemTy,
+                 getAlignment(), isKnownNonNull());
+}

clang/lib/CIR/CodeGen/CIRGenBuilder.h

@@ -733,7 +733,7 @@ class CIRGenBuilderTy : public cir::CIRBaseBuilderTy {
     auto ptrTy = getPointerTo(destType);
     auto baseAddr = create<cir::BaseClassAddrOp>(
         loc, ptrTy, addr.getPointer(), mlir::APInt(64, offset), assumeNotNull);
-    return Address(baseAddr, ptrTy, addr.getAlignment());
+    return Address(baseAddr, destType, addr.getAlignment());
   }
 
   Address createDerivedClassAddr(mlir::Location loc, Address addr,
@@ -745,7 +745,7 @@ class CIRGenBuilderTy : public cir::CIRBaseBuilderTy {
     auto ptrTy = getPointerTo(destType);
     auto derivedAddr = create<cir::DerivedClassAddrOp>(
         loc, ptrTy, addr.getPointer(), mlir::APInt(64, offset), assumeNotNull);
-    return Address(derivedAddr, ptrTy, addr.getAlignment());
+    return Address(derivedAddr, destType, addr.getAlignment());
   }
 
   mlir::Value createVTTAddrPoint(mlir::Location loc, mlir::Type retTy,

clang/lib/CIR/CodeGen/CIRGenBuiltinAArch64.cpp

@@ -4502,7 +4502,7 @@ CIRGenFunction::emitAArch64BuiltinExpr(unsigned BuiltinID, const CallExpr *E,
   }
   case NEON::BI__builtin_neon_vld1_dup_v:
   case NEON::BI__builtin_neon_vld1q_dup_v: {
-    Address ptrAddr = PtrOp0.withElementType(vTy.getEltType());
+    Address ptrAddr = PtrOp0.withElementType(builder, vTy.getEltType());
     mlir::Value val = builder.createLoad(getLoc(E->getExprLoc()), ptrAddr);
     cir::VecSplatOp vecSplat =
         builder.create<cir::VecSplatOp>(getLoc(E->getExprLoc()), vTy, val);

clang/lib/CIR/CodeGen/CIRGenCXX.cpp

@@ -408,8 +408,7 @@ void CIRGenModule::emitCXXGlobalVarDeclInit(const VarDecl *varDecl,
   builder.setInsertionPointToStart(block);
   auto getGlobal = builder.createGetGlobal(addr);
 
-  Address declAddr(getGlobal, getGlobal.getType(),
-                   getASTContext().getDeclAlign(varDecl));
+  Address declAddr(getGlobal, getASTContext().getDeclAlign(varDecl));
   assert(performInit && "cannot have constant initializer which needs "
                         "destruction for reference");
   RValue rv = cgf.emitReferenceBindingToExpr(init);

clang/lib/CIR/CodeGen/CIRGenClass.cpp

@@ -1672,7 +1672,7 @@ CIRGenFunction::getAddressOfBaseClass(Address Value,
       VBase, BaseValueTy, not NullCheckValue);
 
   // Cast to the destination type.
-  Value = Value.withElementType(BaseValueTy);
+  Value = Value.withElementType(builder, BaseValueTy);
 
   return Value;
 }
@@ -1894,7 +1894,7 @@ void CIRGenFunction::emitCXXAggrConstructorCall(
     builder.create<cir::ArrayCtor>(
         *currSrcLoc, arrayOp, [&](mlir::OpBuilder &b, mlir::Location loc) {
           auto arg = b.getInsertionBlock()->addArgument(ptrToElmType, loc);
-          Address curAddr = Address(arg, ptrToElmType, eltAlignment);
+          Address curAddr = Address(arg, elementType, eltAlignment);
           auto currAVS = AggValueSlot::forAddr(
               curAddr, type.getQualifiers(), AggValueSlot::IsDestructed,
               AggValueSlot::DoesNotNeedGCBarriers, AggValueSlot::IsNotAliased,

clang/lib/CIR/CodeGen/CIRGenDecl.cpp

@@ -244,11 +244,8 @@ static void emitStoresForConstant(CIRGenModule &CGM, const VarDecl &D,
   // FIXME(cir): This is closer to memcpy behavior but less optimal, instead of
   // copy from a global, we just create a cir.const out of it.
 
-  if (addr.getElementType() != Ty) {
-    auto ptr = addr.getPointer();
-    ptr = builder.createBitcast(ptr.getLoc(), ptr, builder.getPointerTo(Ty));
-    addr = addr.withPointer(ptr, addr.isKnownNonNull());
-  }
+  if (addr.getElementType() != Ty)
+    addr = addr.withElementType(builder, Ty);
 
   auto loc = CGM.getLoc(D.getSourceRange());
   builder.createStore(loc, builder.getConstant(loc, constant), addr);
@@ -1108,7 +1105,7 @@ void CIRGenFunction::emitArrayDestroy(mlir::Value begin, mlir::Value end,
   builder.create<cir::ArrayDtor>(
       *currSrcLoc, begin, [&](mlir::OpBuilder &b, mlir::Location loc) {
         auto arg = b.getInsertionBlock()->addArgument(ptrToElmType, loc);
-        Address curAddr = Address(arg, ptrToElmType, elementAlign);
+        Address curAddr = Address(arg, cirElementType, elementAlign);
         if (useEHCleanup) {
           pushRegularPartialArrayCleanup(arg, arg, elementType, elementAlign,
                                          destroyer);

clang/lib/CIR/CodeGen/CIRGenException.cpp

@@ -237,7 +237,7 @@ void CIRGenFunction::emitAnyExprToExn(const Expr *e, Address addr) {
   // __cxa_allocate_exception returns a void*;  we need to cast this
   // to the appropriate type for the object.
   auto ty = convertTypeForMem(e->getType());
-  Address typedAddr = addr.withElementType(ty);
+  Address typedAddr = addr.withElementType(builder, ty);
 
   // From LLVM's codegen:
   // FIXME: this isn't quite right!  If there's a final unelided call