Allow PTLS to be changed in a function (preparation for task migration) #39168

tkf · 2021-01-10T03:23:10Z

This PR tweaks how PTLS getters are generated and lets us insert multiple PTLS getters at arbitrary points in a function. There are two goals motivating this change:

Task migration across threads. To migrate the tasks across threads, we need to refetch PTLS after every point that a function can yield (= most of the function calls).
For Julia-Tapir integration (or, more in general, LLVM-level compiler support for the task system) we are working on, we need to be able to split out chunks of LLVM IR as functions. I think relaxing the current requirement that the PTLS has to be fetched only once at the entry block helps this part of the LLVM pass.

Note that this PR does not cover these goals and so I'm not 100% sure if it is the right approach yet. Also, there are some failures in tests due to this change and so it's not merge-ready. But it'd be great if I can get some feedbacks in the design before I spend more time fixing the bugs.

Demo

I build julia with MIGRATE_TASKS enabled:

diff --git a/src/options.h b/src/options.h
index 3ffbf05b22..5ea220900b 100644
--- a/src/options.h
+++ b/src/options.h
@@ -114,7 +114,7 @@
 #endif

 // allow a suspended Task to restart on a different thread
-//#define MIGRATE_TASKS
+#define MIGRATE_TASKS

 // threading options ----------------------------------------------------------

and then

julia> function f()
           a = Threads.threadid()
           yield()
           b = Threads.threadid()
           (a, b)
       end
f (generic function with 1 method)

julia> @code_llvm debuginfo=:none f()

prints

; Function Attrs: sspstrong
define void @julia_f_175([2 x i64]* noalias nocapture sret([2 x i64]) %0) #0 {
top:
  %thread_ptr4 = call i8* asm "movq %fs:0, $0", "=r"() #5
  %1 = getelementptr i8, i8* %thread_ptr4, i64 -32752
  %2 = bitcast i8* %1 to i16*
  %3 = load i16, i16* %2, align 2
  %4 = sext i16 %3 to i64
  %5 = add nsw i64 %4, 1
  %6 = call nonnull {}* @j_yield_177()
  %thread_ptr = call i8* asm "movq %fs:0, $0", "=r"() #5
  %7 = getelementptr i8, i8* %thread_ptr, i64 -32752
  %8 = bitcast i8* %7 to i16*
  %9 = load i16, i16* %8, align 2
  %10 = sext i16 %9 to i64
  %11 = add nsw i64 %10, 1
  %.sroa.0.0..sroa_idx = getelementptr inbounds [2 x i64], [2 x i64]* %0, i64 0, i64 0
  store i64 %5, i64* %.sroa.0.0..sroa_idx, align 8
  %.sroa.2.0..sroa_idx1 = getelementptr inbounds [2 x i64], [2 x i64]* %0, i64 0, i64 1
  store i64 %11, i64* %.sroa.2.0..sroa_idx1, align 8
  ret void
}

Note that %thread_ptr is loaded twice. Once at the beginning (as usual) and also just after @j_yield_177().

(Using multiple threads in julia with MIGRATE_TASKS crashes julia. I don't know if it's from the scheduler or this PR or both.)

How it works

This patch introduces new LLVM placeholder functions julia.reuse_ptls_states and julia.refetch_ptls_states whose signature is identical to julia.ptls_states. During the codegen (emit_function etc.), it inserts julia.reuse_ptls_states just before each instruction that needs PTLS. It inserts julia.refetch_ptls_states if we need to refetch the PTLS (but it's not used in the normal build that doesn't enable MIGRATE_TASKS). So, in the above example, emit_function produces this LLVM IR (@code_llvm debuginfo=:none optimize=false f()):

; Function Attrs: sspstrong
define void @julia_f_235([2 x i64]* noalias nocapture sret([2 x i64]) %0) #0 {
top:
  %1 = alloca [2 x i64], align 8
  %2 = call {}*** @julia.reuse_ptls_states()
  %3 = bitcast {}*** %2 to i16*
  %4 = getelementptr inbounds i16, i16* %3, i64 8
  %5 = load i16, i16* %4, align 2
  %6 = sext i16 %5 to i64
  %7 = add i64 %6, 1
  %8 = call nonnull {}* @j_yield_237()
  %9 = call {}*** @julia.refetch_ptls_states()
  %10 = call {}*** @julia.reuse_ptls_states()
  %11 = bitcast {}*** %10 to i16*
  %12 = getelementptr inbounds i16, i16* %11, i64 8
  %13 = load i16, i16* %12, align 2
  %14 = sext i16 %13 to i64
  %15 = add i64 %14, 1
  %16 = getelementptr inbounds [2 x i64], [2 x i64]* %1, i32 0, i32 0
  store i64 %7, i64* %16, align 8
  %17 = getelementptr inbounds [2 x i64], [2 x i64]* %1, i32 0, i32 1
  store i64 %15, i64* %17, align 8
  %18 = bitcast [2 x i64]* %0 to i8*
  %19 = bitcast [2 x i64]* %1 to i8*
  call void @llvm.memcpy.p0i8.p0i8.i64(i8* align 8 %18, i8* %19, i64 16, i1 false)
  ret void
}

In this patch, julia.ptls_states is not emitted by the initial codegen anymore. Instead, I created a pass LowerPTLSReuse that inserts julia.ptls_states if required and merges all @julia.refetch_ptls_states() using phi nodes. I think it means that we are manually denoting a particular "task-pure" semantics of the function @julia.reuse_ptls_states. (Ideally, LLVM has the notion of pure/const appropriate for this use case. But, looking at isLoadFromConstGV and its comment, I'm guessing LLVM doesn't have this and implementing the pass ourselves is a reasonable approach.)

The LowerPTLSReuse pass is implemented in such a way that it is (supposed to be) safe to run multiple times. This is because LateLowerGCFrame can introduce new instructions that require PTLS (i.e., new @julia.reuse_ptls_states() calls) so that we need to re-run LowerPTLSReuse again. LowerPTLSReuse is called at the beginning of the LLVM pass to avoid interfering with the optimization passes that may give up when there are calls to opaque function like @julia.reuse_ptls_states.

Next steps

Ideally:

Get OK for the design from compiler devs
Fix the bugs and make the CI happy
Get OK for the implementation and then merge it

tkf · 2021-01-10T03:24:11Z

src/llvm-late-gc-lowering.cpp

+        // Fetch or insert the call to `julia.ptls_states` in the entry block.
+        //
+        // Note: It is OK to use the PTLS of the entry block since the runtime
+        // (`ctx_switch`) is responsible for maintaining the association of PTLS
+        // and GC fame. What the lowering cares is the uses of the PTLS after
+        // re-fetches *other than* the GC frame.
+        auto ptlsStates = ensureEntryBlockPtls(*F);
+        // TODO: Ask for a review.


Is my comment here correct? By ctx_switch, I was referring to this part:

julia/src/task.c

Lines 402 to 403 in 86639e1

// set up global state for new task

ptls->pgcstack = t->gcstack;

tkf · 2021-01-10T03:25:17Z

src/llvm-ptls-reuse.cpp

+    if (isa<SelectInst>(Ptls)) {
+        // TODO: Ask if this branch is required. Since we don't introduce select
+        // for PTLS variables explicitly, this won't be necessary if we can be
+        // sure that other LLVM passes won't introduces select instructions.
+        // TODO: Test this branch if we keep this branch.
+        if (!InsertResult.second) {
+            auto OldPtls = InsertResult.first->second;
+            if (OldPtls->comesBefore(Ptls)) {
+                BbToPtls[BB] = Ptls;


Do we need this isa<SelectInst>(Ptls) branch?

tkf · 2021-01-10T03:29:24Z

src/llvm-final-gc-lowering.cpp

-    if (!ptls_getter)
-        return true;
-
-    // Look for a call to 'julia.ptls_states'.
-    ptlsStates = getPtls(F);
-    if (!ptlsStates)
-        return true;


Were these early returns some kind of optimization? If so, we can add something like

if (!usePtls(F)) return true;

tkf · 2021-01-10T03:52:04Z

src/llvm-ptls-reuse.cpp

+    SmallVector<std::pair<Instruction *, std::unique_ptr<SmallVector<Instruction *>>>>
+        SourceAndReuses;


I have no idea how to do SmallVector-of-SmallVectors in a "modern" C++ (or maybe rather in LLVM infrastructure). Is this an OK approach?

tkf · 2021-01-10T03:54:17Z

src/aotcompile.cpp

+    PM->add(createLowerPTLSReusePass());
+#ifdef JL_DEBUG_BUILD
+    PM->add(createVerifierPass(false));
+#endif


I'll remove PM->add(createVerifierPass(false));s before merge.

yuyichao · 2021-01-10T04:00:04Z

This pushes the cost onto code that does not use any related features. Instead, the cost should be on the code that want to move task between threads by changing setting the correct content of the local storage.

tkf · 2021-01-10T04:08:19Z

By "the cost should be on the code", do you mean the code in the compiler or the user code generated by the compiler?

yuyichao · 2021-01-10T04:44:38Z

do you mean the code in the compiler or the user code generated by the compiler

Doesn't matter. It's "the code that want to move task between threads". If you want to do this in generated code then it's the generated code, if you want to do that in the runtime then it's runtime code.

tkf · 2021-01-10T05:17:27Z

Could you elaborate your reasoning? If you are talking about the cost in the compiler, the cost would be paid only once during the compilation for a given user method. If you are talking about the cost in the generated user code (or also the scheduler), the cost would be paid every time the function is called. I am not following why these distinction "doesn't matter." I think how often we'd need to pay the cost matters when we are assessing the efficiency of the approach.

From your reply, I'm guessing that you are not talking about the cost of compilation but rather the run-time cost of the compiled (generated) code and also the runtime/scheduler. Is this correct?

yuyichao · 2021-01-10T05:28:06Z

Could you elaborate your reasoning? If you are talking about the cost in the compiler, the cost would be paid only once during the compilation for a given user method. If you are talking about the cost in the generated user code (or also the scheduler), the cost would be paid every time the function is called. I am not following why these distinction "doesn't matter." I think how often we'd need to pay the cost matters when we are assessing the efficiency of the approach.

From your reply, I'm guessing that you are not talking about the cost of compilation but rather the run-time cost of the compiled (generated) code and also the runtime/scheduler. Is this correct?

I'm not sure what explanation you want. All of these are talking about the cost in this PR whereas the previous reply was entirely about what I believe it should be. The two replies are talking about completely different things so I'm confused about what you are asking about.

In any case, the only distinction that matters, as I said at the very beginning, is whether you are adding cost to code that uses the feature vs code that doesn't. I also never said that the frequency doesn't matter. All what I said is that what kind of code it is "that want to move task between threads" doesn't matter at all. It's what it does, i.e. move task between thread, that matters. And it matters exactly because this must happens way less frequent than function calls.

tkf · 2021-01-10T06:29:19Z

Thanks for the explanation. I think your reply gives me an extra hint that my guess in the previous comment might be correct. I'll reply assuming that your concern is not the cost of compilation time (i.e., the cost paid one for each method; a penalty for complexifying the compiler) and rather the cost at run-time (i.e., the cost paid for each method call). More specifically, I'm guessing you are referring in particular to the insertion of julia.refetch_ptls_states after every function invocation. I have a few comments on that:

First of all, refetch is not enabled unless you switch on the C macro MIGRATE_TASKS. So, I'd expect normal build of Julia to be as efficient as the one without this PR. Of course, the IR won't be identical since it inserts extra phi nodes. I'm hoping something like CFGSimplifyPass can eliminate the possible costs associated with this kind of IR. I think landing this to the main branch without MIGRATE_TASKS enabled is a good way to assess the cost (hopefully 0) of the extra phi nodes in real applications.

Second, since the optimizer on the Julia IR can inline function calls, let me emphasize that a function call in Julia code at syntax level does not imply a PTLS refetch. Having said that, the current implementation does not support elimination of the refetch after LLVM's inliner. It's also conceivable that we'd have some analysis pass that can prove certain function calls do not yield. In that case, we can safely remove julia.refetch_ptls_states. So, I think these points rather farther backup the design decision to postpone the insertion of julia.ptls_states to the LLVM pass rather than in emit_function. Since we emit a flexible IR, we can do some optimizations down the line.

Third, even if (hypothetically) it turned out supporting the migration of arbitrary task is not possible without a significant sacrifice of the single-thread performance, we can still offer opt-in migration by a special macro/function (say) Threads.@migratable akin to GC.safepoint(). I think the current approach would support this scenario as well since we just need to insert a yield (with a hint to the scheduler it supports migration) and a julia.refetch_ptls_states.

yuyichao · 2021-01-10T06:46:04Z

First of all, refetch is not enabled unless you switch on the C macro MIGRATE_TASKS.

this is not a valid argument unless you promise that this feature will never be on by default.

let me emphasize that a function call in Julia code at syntax level does not imply a PTLS refetch

yes, of course, I did not assume that was the case. Still, you are forcing the cost on code that does not even use task switches because the compiler must make conservative assumptions about function calls to c code and non inlined functions.

we can still offer opt-in migration by a special macro/function (say) Threads.@migratable akin to GC.safepoint()

no you can’t. Safe point is a local property whereas task migration is not unless you guarantee to migrate things back.

tkf · 2021-01-10T09:05:33Z

we can still offer opt-in migration by a special macro/function (say) Threads.@migratable akin to GC.safepoint()

no you can’t. Safe point is a local property whereas task migration is not unless you guarantee to migrate things back.

I'm interested in why you think it's not possible. If you think task migration is possible to implement (putting aside efficiency consideration), then "migration point" support would need to do exactly how a function containing yield would be supported in such compiler. Since all functions called from the function containing the migration point already have returned at the migration point [1], I'm not sure why it's not possible to process the migration point with just the information of the IR of the function, independent of the properties of the callees.

[1] The only way I can imagine that would be a problem is when someone develops a custom sub-scheduler by directly using yieldto to manually suspend a function call (then a migration point in any of the user would migrate the entire sub-scheduler). AFAICT, such a framework is not widely used yet. So, I'm leaning towards to suppose this is an non-issue and add it as a warning in yieldto docstring.

yuyichao · 2021-01-10T12:11:32Z

then "migration point" support would need to do exactly how a function containing yield would be supported in such compiler.

Well, the fact that this PR exists at all proves this point wrong.

Since all functions called from the function containing the migration point already have returned at the migration point,

Errrr, no by definition? If a function contains a point, then by definition it hasn't returned yet when reaching that point.

vchuravy · 2021-01-10T18:06:11Z

So the underlying issue is that we are caching data and that after a task migration that cached data is invalid.

The three approaches I see are:

Conservatively insert refetches after potential yield points
- This PR
- Fairly straightforward to implement by repurposing the current ptls infrastructure
- As pointed out by Yichao increases the runtime cost regardless of whether a task switch occurred
During compilation derive a stack-table that points to where the ptls pointer lives (or even values derived from the ptls but I fear that's a rabbit hole), and during task migration use that table to fix up the ptls pointer

Moves the cost to compilation and actual migration
I think this is similar to what Go uses to support preemption
No clear implementation strategy -- maybe we can use debug-info but that will be unreliable

Use the ptls pointer once to obtain the current task, and lookup information through the task struct instead of ptls

Provides a place for task-migration to change things as needed
increases the size of the task struct

@yuyichao do you see an alternative path towards task migration (and perhaps even task preemption?)

vtjnash · 2021-01-10T18:58:45Z

I've been wanting to do 3 on that list for some time. Just waiting for the release to finish to start working on that. The jl_task_t struct already links back to ptls, so it won't necessarily even change the size. In fact, with a little bit of planning, I think it will work out to be equivalent in every place it matters.

yuyichao · 2021-01-10T19:09:39Z

Yes, 3 is what I mentioned in the other issue a while ago. It puts the cost at the right place.

tkf · 2021-01-10T22:05:24Z

A "disadvantage" of approach 3 is that we can't support LICM involving rand() or array[threadid()] in general (which was the reason why I thought we were OK to accept the cost of approach 1, based on the quick discussion in Julia slack last month). Having said that, it is questionable if we should encourage threadid()-based user-space thread-local storage. For thread-safe RNG, there's already #34852 and we can even generalize it by implementing splitter "hyperobject" as suggested in Cilk. It'd be great if we can use approach 3 as an additional incentive to move away from threadid().

vtjnash · 2021-04-03T15:49:21Z

Superseded by #39220

tkf added 4 commits January 8, 2021 09:37

Separate PTLS getter insertion in two phases; add LowerPTLSReuse pass

bed7ba6

Handle PTLS refetch

bd0e909

Make LowerPTLSReuse idempotent

5e55cfa

Fixing MIGRATE_TASKS

7a38cf7

tkf added compiler:codegen Generation of LLVM IR and native code multithreading Base.Threads and related functionality labels Jan 10, 2021

tkf requested review from vchuravy, vtjnash and JeffBezanson January 10, 2021 03:23

tkf commented Jan 10, 2021

View reviewed changes

tkf mentioned this pull request Jan 13, 2021

Refetch PTLS via Task struct (preparation for task migration) #39220

Closed

vtjnash closed this Apr 3, 2021

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Allow PTLS to be changed in a function (preparation for task migration) #39168

Allow PTLS to be changed in a function (preparation for task migration) #39168

tkf commented Jan 10, 2021 •

edited

Loading

tkf Jan 10, 2021 •

edited

Loading

tkf Jan 10, 2021

tkf Jan 10, 2021

tkf Jan 10, 2021

tkf Jan 10, 2021

yuyichao commented Jan 10, 2021

tkf commented Jan 10, 2021

yuyichao commented Jan 10, 2021

tkf commented Jan 10, 2021

yuyichao commented Jan 10, 2021

tkf commented Jan 10, 2021

yuyichao commented Jan 10, 2021

tkf commented Jan 10, 2021

yuyichao commented Jan 10, 2021

vchuravy commented Jan 10, 2021

vtjnash commented Jan 10, 2021 •

edited

Loading

yuyichao commented Jan 10, 2021

tkf commented Jan 10, 2021

vtjnash commented Apr 3, 2021

	// set up global state for new task
	ptls->pgcstack = t->gcstack;

		SmallVector<std::pair<Instruction , std::unique_ptr<SmallVector<Instruction >>>>
		SourceAndReuses;

Allow PTLS to be changed in a function (preparation for task migration) #39168

Allow PTLS to be changed in a function (preparation for task migration) #39168

Conversation

tkf commented Jan 10, 2021 • edited Loading

Demo

How it works

Next steps

tkf Jan 10, 2021 • edited Loading

Choose a reason for hiding this comment

tkf Jan 10, 2021

Choose a reason for hiding this comment

tkf Jan 10, 2021

Choose a reason for hiding this comment

tkf Jan 10, 2021

Choose a reason for hiding this comment

tkf Jan 10, 2021

Choose a reason for hiding this comment

yuyichao commented Jan 10, 2021

tkf commented Jan 10, 2021

yuyichao commented Jan 10, 2021

tkf commented Jan 10, 2021

yuyichao commented Jan 10, 2021

tkf commented Jan 10, 2021

yuyichao commented Jan 10, 2021

tkf commented Jan 10, 2021

yuyichao commented Jan 10, 2021

vchuravy commented Jan 10, 2021

vtjnash commented Jan 10, 2021 • edited Loading

yuyichao commented Jan 10, 2021

tkf commented Jan 10, 2021

vtjnash commented Apr 3, 2021

tkf commented Jan 10, 2021 •

edited

Loading

tkf Jan 10, 2021 •

edited

Loading

vtjnash commented Jan 10, 2021 •

edited

Loading