1// Copyright 2019 The Marl Authors.
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7// https://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15// Minimal assembly implementations of fiber context switching for Unix-based
16// platforms.
17//
18// Note: Unlike makecontext, swapcontext or the Windows fiber APIs, these
19// assembly implementations *do not* save or restore signal masks,
20// floating-point control or status registers, FS and GS segment registers,
21// thread-local storage state nor any SIMD registers. This should not be a
22// problem as the marl scheduler requires fibers to be executed on the same
23// thread throughout their lifetime.
24
25#if defined(__x86_64__)
26#include "osfiber_asm_x64.h"
27#elif defined(__i386__)
28#include "osfiber_asm_x86.h"
29#elif defined(__aarch64__)
30#include "osfiber_asm_aarch64.h"
31#elif defined(__arm__)
32#include "osfiber_asm_arm.h"
33#elif defined(__powerpc64__)
34#include "osfiber_asm_ppc64.h"
35#elif defined(__mips__) && _MIPS_SIM == _ABI64
36#include "osfiber_asm_mips64.h"
37#elif defined(__riscv) && __riscv_xlen == 64
38#include "osfiber_asm_rv64.h"
39#elif defined(__loongarch__) && _LOONGARCH_SIM == _ABILP64
40#include "osfiber_asm_loongarch64.h"
41#else
42#error "Unsupported target"
43#endif
44
45#include "marl/export.h"
46#include "marl/memory.h"
47
48#include <functional>
49#include <memory>
50
51extern "C" {
52
53MARL_EXPORT
54extern void marl_fiber_set_target(marl_fiber_context*,
55 void* stack,
56 uint32_t stack_size,
57 void (*target)(void*),
58 void* arg);
59MARL_EXPORT
60extern void marl_fiber_swap(marl_fiber_context* from,
61 const marl_fiber_context* to);
62
63} // extern "C"
64
65namespace marl {
66
67class OSFiber {
68 public:
69 inline OSFiber(Allocator*);
70 inline ~OSFiber();
71
72 // createFiberFromCurrentThread() returns a fiber created from the current
73 // thread.
74 MARL_NO_EXPORT static inline Allocator::unique_ptr<OSFiber>
75 createFiberFromCurrentThread(Allocator* allocator);
76
77 // createFiber() returns a new fiber with the given stack size that will
78 // call func when switched to. func() must end by switching back to another
79 // fiber, and must not return.
80 MARL_NO_EXPORT static inline Allocator::unique_ptr<OSFiber> createFiber(
81 Allocator* allocator,
82 size_t stackSize,
83 const std::function<void()>& func);
84
85 // switchTo() immediately switches execution to the given fiber.
86 // switchTo() must be called on the currently executing fiber.
87 MARL_NO_EXPORT inline void switchTo(OSFiber*);
88
89 private:
90 MARL_NO_EXPORT
91 static inline void run(OSFiber* self);
92
93 Allocator* allocator;
94 marl_fiber_context context;
95 std::function<void()> target;
96 Allocation stack;
97};
98
99OSFiber::OSFiber(Allocator* allocator) : allocator(allocator) {}
100
101OSFiber::~OSFiber() {
102 if (stack.ptr != nullptr) {
103 allocator->free(stack);
104 }
105}
106
107Allocator::unique_ptr<OSFiber> OSFiber::createFiberFromCurrentThread(
108 Allocator* allocator) {
109 auto out = allocator->make_unique<OSFiber>(allocator);
110 out->context = {};
111 return out;
112}
113
114Allocator::unique_ptr<OSFiber> OSFiber::createFiber(
115 Allocator* allocator,
116 size_t stackSize,
117 const std::function<void()>& func) {
118 Allocation::Request request;
119 request.size = stackSize;
120 request.alignment = 16;
121 request.usage = Allocation::Usage::Stack;
122#if MARL_USE_FIBER_STACK_GUARDS
123 request.useGuards = true;
124#endif
125
126 auto out = allocator->make_unique<OSFiber>(allocator);
127 out->context = {};
128 out->target = func;
129 out->stack = allocator->allocate(request);
130 marl_fiber_set_target(
131 &out->context, out->stack.ptr, static_cast<uint32_t>(stackSize),
132 reinterpret_cast<void (*)(void*)>(&OSFiber::run), out.get());
133 return out;
134}
135
136void OSFiber::run(OSFiber* self) {
137 self->target();
138}
139
140void OSFiber::switchTo(OSFiber* fiber) {
141 marl_fiber_swap(&context, &fiber->context);
142}
143
144} // namespace marl
145