[CIR] Backport CountOf VLA with Array element type

Author: AmrDeveloper

clang/lib/CIR/CodeGen/CIRGenExprScalar.cpp

@@ -3028,28 +3028,44 @@ mlir::Value ScalarExprEmitter::VisitVAArgExpr(VAArgExpr *VE) {
 mlir::Value ScalarExprEmitter::VisitUnaryExprOrTypeTraitExpr(
     const UnaryExprOrTypeTraitExpr *E) {
   QualType TypeToSize = E->getTypeOfArgument();
-  if (E->getKind() == UETT_SizeOf) {
+  if (auto kind = E->getKind();
+      kind == UETT_SizeOf || kind == UETT_DataSizeOf | kind == UETT_CountOf) {
     if (const VariableArrayType *VAT =
             CGF.getContext().getAsVariableArrayType(TypeToSize)) {
 
-      if (E->isArgumentType()) {
-        // sizeof(type) - make sure to emit the VLA size.
-        CGF.emitVariablyModifiedType(TypeToSize);
-      } else {
-        // C99 6.5.3.4p2: If the argument is an expression of type
-        // VLA, it is evaluated.
-        CGF.emitIgnoredExpr(E->getArgumentExpr());
+      // For _Countof, we only want to evaluate if the extent is actually
+      // variable as opposed to a multi-dimensional array whose extent is
+      // constant but whose element type is variable.
+      bool evaluateExtent = true;
+      if (kind == UETT_CountOf && VAT->getElementType()->isArrayType()) {
+        evaluateExtent =
+            !VAT->getSizeExpr()->isIntegerConstantExpr(CGF.getContext());
       }
 
-      auto VlaSize = CGF.getVLASize(VAT);
-      mlir::Value size = VlaSize.NumElts;
+      if (evaluateExtent) {
+        if (E->isArgumentType()) {
+          // sizeof(type) - make sure to emit the VLA size.
+          CGF.emitVariablyModifiedType(TypeToSize);
+        } else {
+          // C99 6.5.3.4p2: If the argument is an expression of type
+          // VLA, it is evaluated.
+          CGF.emitIgnoredExpr(E->getArgumentExpr());
+        }
+
+        // For _Countof, we just want to return the size of a single dimension.
+        if (kind == UETT_CountOf)
+          return CGF.getVLAElements1D(VAT).NumElts;
+
+        auto VlaSize = CGF.getVLASize(VAT);
+        mlir::Value size = VlaSize.NumElts;
 
-      // Scale the number of non-VLA elements by the non-VLA element size.
-      CharUnits eltSize = CGF.getContext().getTypeSizeInChars(VlaSize.Type);
-      if (!eltSize.isOne())
-        size = Builder.createMul(size, CGF.CGM.getSize(eltSize).getValue());
+        // Scale the number of non-VLA elements by the non-VLA element size.
+        CharUnits eltSize = CGF.getContext().getTypeSizeInChars(VlaSize.Type);
+        if (!eltSize.isOne())
+          size = Builder.createMul(size, CGF.CGM.getSize(eltSize).getValue());
 
-      return size;
+        return size;
+      }
     }
   } else if (E->getKind() == UETT_OpenMPRequiredSimdAlign) {
     llvm_unreachable("NYI");

clang/lib/CIR/CodeGen/CIRGenFunction.cpp

@@ -1808,6 +1808,14 @@ CIRGenFunction::getVLASize(const VariableArrayType *type) {
   return {numElements, elementType};
 }
 
+CIRGenFunction::VlaSizePair
+CIRGenFunction::getVLAElements1D(const VariableArrayType *vla) {
+  mlir::Value vlaSize = VLASizeMap[vla->getSizeExpr()];
+  assert(vlaSize && "no size for VLA!");
+  assert(vlaSize->getType() == sizeTy);
+  return {vlaSize, vla->getElementType()};
+}
+
 // TODO(cir): most part of this function can be shared between CIRGen
 // and traditional LLVM codegen
 void CIRGenFunction::emitVariablyModifiedType(QualType type) {

clang/lib/CIR/CodeGen/CIRGenFunction.h

@@ -683,6 +683,10 @@ class CIRGenFunction : public CIRGenTypeCache {
     VlaSizePair(mlir::Value NE, QualType T) : NumElts(NE), Type(T) {}
   };
 
+  /// Return the number of elements for a single dimension
+  /// for the given array type.
+  VlaSizePair getVLAElements1D(const VariableArrayType *vla);
+
   /// Returns an MLIR value that corresponds to the size,
   /// in non-variably-sized elements, of a variable length array type,
   /// plus that largest non-variably-sized element type.  Assumes that

clang/test/CIR/CodeGen/count-of.c

@@ -0,0 +1,52 @@
+// RUN: %clang_cc1 -std=c2y -triple x86_64-unknown-linux-gnu -Wno-unused-value -fclangir -emit-cir %s -o %t.cir
+// RUN: FileCheck --input-file=%t.cir %s -check-prefix=CIR
+// RUN: %clang_cc1 -std=c2y -triple x86_64-unknown-linux-gnu -Wno-unused-value -fclangir -emit-llvm %s -o %t-cir.ll
+// RUN: FileCheck --input-file=%t-cir.ll %s -check-prefix=LLVM
+// RUN: %clang_cc1 -std=c2y -triple x86_64-unknown-linux-gnu -Wno-unused-value -emit-llvm %s -o %t.ll
+// RUN: FileCheck --input-file=%t.ll %s -check-prefix=OGCG
+
+unsigned long vla_with_array_element_type_with_const_size() {
+  long size;
+  return _Countof(int[5][size]);
+}
+
+// CIR: %[[RET_ADDR:.*]] = cir.alloca !u64i, !cir.ptr<!u64i>, ["__retval"]
+// CIR: %[[SIZE_ADDR:.*]] = cir.alloca !s64i, !cir.ptr<!s64i>, ["size"]
+// CIR: %[[CONST_5:.*]] = cir.const #cir.int<5> : !u64i
+// CIR: cir.store %[[CONST_5]], %[[RET_ADDR]] : !u64i, !cir.ptr<!u64i>
+// CIR: %[[RET_VAL:.*]] = cir.load %[[RET_ADDR]] : !cir.ptr<!u64i>, !u64i
+// CIR: cir.return %[[RET_VAL]] : !u64i
+
+// LLVM: %[[RET_ADDR:.*]] = alloca i64, i64 1, align 8
+// LLVM: %[[SIZE_ADDR:.*]] = alloca i64, i64 1, align 8
+// LLVM: store i64 5, ptr %[[RET_ADDR]], align 8
+// LLVM: %[[RET_VAL:.*]] = load i64, ptr %[[RET_ADDR]], align 8
+// LLVM: ret i64 %[[RET_VAL]]
+
+// OGCG: %[[SIZE_ADDR:.*]] = alloca i64, align 8
+// OGCG: ret i64 5
+
+unsigned long vla_with_array_element_type_non_const_size() {
+  long size;
+  return _Countof(int[size][size]);
+}
+
+// CIR: %[[REET_ADDR:.*]] = cir.alloca !u64i, !cir.ptr<!u64i>, ["__retval"]
+// CIR: %[[SIZE_ADDR:.*]] = cir.alloca !s64i, !cir.ptr<!s64i>, ["size"]
+// CIR: %[[TMP_SIZE:.*]] = cir.load {{.*}} %[[SIZE_ADDR]] : !cir.ptr<!s64i>, !s64i
+// CIR: %[[TMP_SIZE_U64:.*]] = cir.cast integral %[[TMP_SIZE]] : !s64i -> !u64i
+// CIR: cir.store %[[TMP_SIZE_U64]], %[[RET_ADDR]] : !u64i, !cir.ptr<!u64i>
+// CIR: %[[TMP_RET:.*]] = cir.load %[[RET_ADDR]] : !cir.ptr<!u64i>, !u64i
+// CIR: cir.return %[[TMP_RET]] : !u64i
+
+// LLVM: %[[RET_ADDR:.*]] = alloca i64, i64 1, align 8
+// LLVM: %[[SIZE_ADDR:.*]] = alloca i64, i64 1, align 8
+// LLVM: %[[TMP_SIZE:.*]] = load i64, ptr %[[SIZE_ADDR]], align 8
+// LLVM: store i64 %[[TMP_SIZE]], ptr %[[RET_ADDR]], align 8
+// LLVM: %[[TMP_RET:.*]] = load i64, ptr %[[RET_ADDR]], align 8
+// LLVM: ret i64 %[[TMP_RET]]
+
+// OGCG: %[[SIZE_ADDR:.*]] = alloca i64, align 8
+// OGCG: %[[TMP_SIZE:.*]] = load i64, ptr %[[SIZE_ADDR]], align 8
+// OGCG: %[[TMP_SIZE_2:.*]] = load i64, ptr %[[SIZE_ADDR]], align 8
+// OGCG: ret i64 %[[TMP_SIZE]]