[clang][WebAssembly] Return aggregate values indirectly in swiftcc

The Swift calling convention on Wasm has historically returned aggregate
values directly at the LLVM IR level due to the use of the generic
`SwiftABIInfo` implementation. The direct return at LLVM IR level will
cause unnecessary stack allocation and memory copies for each aggregate
return value.

Author: kateinoigakukun

clang/lib/CodeGen/ABIInfo.cpp

@@ -247,6 +247,21 @@ ABIInfo::getOptimalVectorMemoryType(llvm::FixedVectorType *T,
 // Pin the vtable to this file.
 SwiftABIInfo::~SwiftABIInfo() = default;
 
+void SwiftABIInfo::countOccupiedRegisters(ArrayRef<llvm::Type *> scalarTypes,
+                                          unsigned &intCount, unsigned &fpCount,
+                                          unsigned maxIntRegisterBitWidth) {
+  for (llvm::Type *type : scalarTypes) {
+    if (type->isPointerTy()) {
+      intCount++;
+    } else if (auto *intTy = dyn_cast<llvm::IntegerType>(type)) {
+      intCount += (intTy->getBitWidth() + maxIntRegisterBitWidth - 1) / maxIntRegisterBitWidth;
+    } else {
+      assert(type->isVectorTy() || type->isFloatingPointTy());
+      fpCount++;
+    }
+  }
+}
+
 /// Does the given lowering require more than the given number of
 /// registers when expanded?
 ///
@@ -262,18 +277,10 @@ SwiftABIInfo::~SwiftABIInfo() = default;
 /// return registers.
 bool SwiftABIInfo::occupiesMoreThan(ArrayRef<llvm::Type *> scalarTypes,
                                     unsigned maxAllRegisters) const {
+  // Use the pointer width as the maximum integer register bit width by default.
+  unsigned ptrWidth = CGT.getTarget().getPointerWidth(LangAS::Default);
   unsigned intCount = 0, fpCount = 0;
-  for (llvm::Type *type : scalarTypes) {
-    if (type->isPointerTy()) {
-      intCount++;
-    } else if (auto intTy = dyn_cast<llvm::IntegerType>(type)) {
-      auto ptrWidth = CGT.getTarget().getPointerWidth(LangAS::Default);
-      intCount += (intTy->getBitWidth() + ptrWidth - 1) / ptrWidth;
-    } else {
-      assert(type->isVectorTy() || type->isFloatingPointTy());
-      fpCount++;
-    }
-  }
+  countOccupiedRegisters(scalarTypes, intCount, fpCount, ptrWidth);
 
   return (intCount + fpCount > maxAllRegisters);
 }

clang/lib/CodeGen/ABIInfo.h

@@ -141,6 +141,9 @@ class SwiftABIInfo {
   CodeGenTypes &CGT;
   bool SwiftErrorInRegister;
 
+  static void countOccupiedRegisters(ArrayRef<llvm::Type *> scalarTypes,
+                                     unsigned &intCount, unsigned &fpCount,
+                                     unsigned maxIntRegisterBitWidth);
   bool occupiesMoreThan(ArrayRef<llvm::Type *> scalarTypes,
                         unsigned maxAllRegisters) const;
 

clang/lib/CodeGen/Targets/WebAssembly.cpp

@@ -45,13 +45,40 @@ class WebAssemblyABIInfo final : public ABIInfo {
                    AggValueSlot Slot) const override;
 };
 
+class WebAssemblySwiftABIInfo final : public SwiftABIInfo {
+  WebAssemblyABIKind Kind;
+public:
+  explicit WebAssemblySwiftABIInfo(CodeGen::CodeGenTypes &CGT, WebAssemblyABIKind K)
+      : SwiftABIInfo(CGT, /*SwiftErrorInRegister=*/false), Kind(K) {}
+
+  bool shouldPassIndirectly(ArrayRef<llvm::Type *> ComponentTys,
+                            bool AsReturnValue) const override {
+    unsigned maxIntRegisterBitWidth = 64;
+    unsigned intCount = 0, fpCount = 0;
+    countOccupiedRegisters(ComponentTys, intCount, fpCount, maxIntRegisterBitWidth);
+
+    if (AsReturnValue) {
+      if (Kind == WebAssemblyABIKind::ExperimentalMV) {
+        // If the experimental multivalue ABI is enabled, try to return up to 2 values for each
+        // of int and fp, which is a very conservative value based on the number of available
+        // physical gp return registers used in the major engines to minimize stack spills
+        // at JIT time.
+        return intCount > 2 || fpCount > 2;
+      }
+      // By default, limit to 1 total register.
+      return (intCount + fpCount > 1);
+    }
+    // For an argument, limit to 4 total registers.
+    return (intCount + fpCount > 4);
+  }
+};
+
 class WebAssemblyTargetCodeGenInfo final : public TargetCodeGenInfo {
 public:
   explicit WebAssemblyTargetCodeGenInfo(CodeGen::CodeGenTypes &CGT,
                                         WebAssemblyABIKind K)
       : TargetCodeGenInfo(std::make_unique<WebAssemblyABIInfo>(CGT, K)) {
-    SwiftInfo =
-        std::make_unique<SwiftABIInfo>(CGT, /*SwiftErrorInRegister=*/false);
+    SwiftInfo = std::make_unique<WebAssemblySwiftABIInfo>(CGT, K);
   }
 
   void setTargetAttributes(const Decl *D, llvm::GlobalValue *GV,

clang/test/CodeGen/WebAssembly/wasm-return-swiftcc.c

@@ -0,0 +1,61 @@
+// RUN: %clang_cc1 -triple wasm32-unknown-unknown %s -emit-llvm -o - | FileCheck %s
+
+typedef struct {
+  int aa;
+  int bb;
+} s1;
+
+// Multiple-element structs should be returned through sret.
+// CHECK: define swiftcc void @return_s1(ptr dead_on_unwind noalias writable sret(%struct.s1) align 4 %agg.result)
+__attribute__((swiftcall))
+s1 return_s1(void) {
+  s1 foo;
+  return foo;
+}
+
+typedef struct {
+  int cc;
+} s2;
+
+// Single-element structs should be returned directly.
+// CHECK: define swiftcc i32 @return_s2()
+__attribute__((swiftcall))
+s2 return_s2(void) {
+  s2 foo;
+  return foo;
+}
+
+typedef struct {
+    char c1[4];
+} s3;
+
+// CHECK: define swiftcc i32 @return_s3()
+__attribute__((swiftcall))
+s3 return_s3(void) {
+  s3 foo;
+  return foo;
+}
+
+typedef struct {
+    int bf1 : 4;
+    int bf2 : 3;
+    int bf3 : 8;
+} s4;
+
+// CHECK: define swiftcc i16 @return_s4()
+__attribute__((swiftcall))
+s4 return_s4(void) {
+  s4 foo;
+  return foo;
+}
+
+typedef struct {
+    long long v;
+} s5;
+
+// CHECK: define swiftcc i64 @return_s5()
+__attribute__((swiftcall))
+s5 return_s5(void) {
+  s5 foo;
+  return foo;
+}

llvm/lib/Target/WebAssembly/WebAssemblyISelLowering.cpp

@@ -20,6 +20,7 @@
 #include "WebAssemblyUtilities.h"
 #include "llvm/CodeGen/CallingConvLower.h"
 #include "llvm/CodeGen/MachineFrameInfo.h"
+#include "llvm/CodeGen/MachineFunction.h"
 #include "llvm/CodeGen/MachineInstrBuilder.h"
 #include "llvm/CodeGen/MachineJumpTableInfo.h"
 #include "llvm/CodeGen/MachineModuleInfo.h"