From 26cfb6e562f12f8c0d8952120d9df81151dc9c19 Mon Sep 17 00:00:00 2001
From: "Liu, Chen3" <chen3.liu@intel.com>
Date: Fri, 9 Oct 2020 11:20:29 +0800
Subject: [PATCH] [X86] Passing union type through register

For example:

  union M256 {
    double d;
    __m256 m;
  };
  extern void foo1(union M256 A);
  union M256 m1;
  void test() {
    foo1(m1);
  }

clang will pass m1 through stack which does not follow the ABI.

Differential Revision: https://reviews.llvm.org/D78699
---
 clang/lib/CodeGen/TargetInfo.cpp   | 16 ++++++++++------
 clang/test/CodeGen/X86/avx-union.c | 29 +++++++++++++++++++++++++++++
 2 files changed, 39 insertions(+), 6 deletions(-)
 create mode 100644 clang/test/CodeGen/X86/avx-union.c

diff --git a/clang/lib/CodeGen/TargetInfo.cpp b/clang/lib/CodeGen/TargetInfo.cpp
index f39ded3dc31c..d7c279495ace 100644
--- a/clang/lib/CodeGen/TargetInfo.cpp
+++ b/clang/lib/CodeGen/TargetInfo.cpp
@@ -3061,6 +3061,7 @@ void X86_64ABIInfo::classify(QualType Ty, uint64_t OffsetBase,
 
     // Classify the fields one at a time, merging the results.
     unsigned idx = 0;
+    bool IsUnion = RT->isUnionType();
     for (RecordDecl::field_iterator i = RD->field_begin(), e = RD->field_end();
            i != e; ++i, ++idx) {
       uint64_t Offset = OffsetBase + Layout.getFieldOffset(idx);
@@ -3071,14 +3072,17 @@ void X86_64ABIInfo::classify(QualType Ty, uint64_t OffsetBase,
         continue;
 
       // AMD64-ABI 3.2.3p2: Rule 1. If the size of an object is larger than
-      // four eightbytes, or it contains unaligned fields, it has class MEMORY.
+      // eight eightbytes, or it contains unaligned fields, it has class MEMORY.
       //
-      // The only case a 256-bit wide vector could be used is when the struct
-      // contains a single 256-bit element. Since Lo and Hi logic isn't extended
-      // to work for sizes wider than 128, early check and fallback to memory.
+      // The only case a 256-bit or a 512-bit wide vector could be used is when
+      // the struct contains a single 256-bit or 512-bit element. Early check
+      // and fallback to memory.
       //
-      if (Size > 128 && (Size != getContext().getTypeSize(i->getType()) ||
-                         Size > getNativeVectorSizeForAVXABI(AVXLevel))) {
+      // FIXME: Extended the Lo and Hi logic properly to work for size wider
+      // than 128.
+      if (Size > 128 &&
+          ((!IsUnion && Size != getContext().getTypeSize(i->getType())) ||
+           Size > getNativeVectorSizeForAVXABI(AVXLevel))) {
         Lo = Memory;
         postMerge(Size, Lo, Hi);
         return;
diff --git a/clang/test/CodeGen/X86/avx-union.c b/clang/test/CodeGen/X86/avx-union.c
new file mode 100644
index 000000000000..78c366fe1a18
--- /dev/null
+++ b/clang/test/CodeGen/X86/avx-union.c
@@ -0,0 +1,29 @@
+// RUN: %clang_cc1 -w -ffreestanding -triple x86_64-linux-gnu -target-feature +avx -emit-llvm -o %t %s || FileCheck < %t %s --check-prefix=CHECK, AVX
+// RUN: %clang_cc1 -w -ffreestanding -triple x86_64-linux-gnu -target-feature +avx512f -emit-llvm -o %t %s || FileCheck < %t %s --check-prefix=CHECK, AVX512
+// This tests verifies that a union parameter should pass by a vector regitster whose first eightbyte is SSE and the other eightbytes are SSEUP.
+
+typedef int __m256 __attribute__ ((__vector_size__ (32)));
+typedef int __m512 __attribute__ ((__vector_size__ (64)));
+
+union M256 {
+  double d;
+  __m256 m;
+};
+
+union M512 {
+  double d;
+  __m512 m;
+};
+
+extern void foo1(union M256 A);
+extern void foo2(union M512 A);
+union M256 m1;
+union M512 m2;
+// CHECK-LABEL: define dso_local void @test()
+// CHECK:       void @foo1(<4 x double>
+// AVX:         call void @foo2(%union.M512* byval(%union.M512) align 64
+// AVX512:      call void @foo2(<8 x double>
+void test() {
+  foo1(m1);
+  foo2(m2);
+}