1	# performance/bang_after_fork: a `!` call that runs after a fork of the same evaluation, so Bend runs it
2	# on the CPU pool, not on the GPU: in its definition, a let before it forks (a parallel let, or a value
3	# that calls a definition under which a parallel let runs), or one of its arguments does.
4	# Heuristics: a bang that follows the host's forks in one evaluation runs on the CPU pool ("bend guide
5	# shaders"); a `!` call is not one of those forks: Bend's radix sort benchmark, three `!` calls in a row,
6	# took the same time with IO steps between them. Measured on a discrete GPU, a numeric bang that took
7	# 0.55 s on the GPU after an IO step took 3.9 s, its time on the CPU pool, after a parallel let, after a
8	# let whose value forked in a call, and with an argument that forked; after a let whose value looped
9	# without forking, it stayed on the GPU. A short call loses little on the CPU pool. An IO step (`<-` in a
10	# `do` block) starts a new evaluation, so a fork before one does not count; nor does a fork in a
11	# definition that calls this one, which the rule does not follow. A `!` call of a definition under which
12	# no parallel let runs, or of one that returns Unit, is left to performance/gpu_fit and
13	# performance/drop_in_bang: they ask to remove the `!`.
14	import Base
15	import bend-lint-api@0.1.1.0/linter/lint.bend as Lint
16	import bend-lint-api@0.1.1.0/shared/lists.bend as Lists
17	import bend-lint-api@0.1.1.0/shared/names.bend as Names
18	import ./deep_call_in_fork.bend as Deep
19	import ./gpu_fit.bend as GpuFit
20	
21	# Types
22	# =====
23	
24	# What judging a file's calls reads: the program's definitions, by the name the check gives them, and where the file
25	# writes `!(`.
26	type Program is Data:
27	  Program{defs: Map<&2, Maybe<&2, GpuFit.Def>>, bangs: List<&2, U32>}
28	
29	# Functions
30	# =========
31	def message(callee: String) -> String:
32	  callee
33	    ++ "! runs after a fork in the same evaluation (no IO step between them), so Bend runs it on the CPU pool, not on the GPU. GPU: measured 3.9 s for a numeric fork tree, against 0.55 s after the fix. CPU (`--gpu off`): no change. JS: no GPU, so no change. Confidence: high, as Bend's scheduler decides it. Put an IO step, such as `t : Nat <- IO.now()`, between the forks and this call, as Bend's slash boss demo does. If this call is short, so that the CPU pool runs it as fast, expect this finding."
34	
35	# The definition an application below `path` calls by name on the host: not with `!`, as a `!` call runs on the GPU,
36	# and a `!` call after it stays there.
37	def host.visit(
38	  +bangs: List<&2, U32>,
39	  +path: List<&2, Lint.Above>,
40	  +node: Lint.Node
41	) -> List<&2, String>:
42	  match node:
43	    case Lint.Node{
44	      Lint.Application{},
45	      n,
46	      s,
47	      Con{Lint.Node{Lint.Reference{t}, rn, Some{Lint.Span{p, beg, end}}, rk, rf}, args},
48	      f
49	    }:
50	      Bool.pick(List<&2, String>, List.contains(~U32, ~U32.is_eq, bangs, end), [], [t])
51	    case _:
52	      []
53	
54	# Whether `nodes` fork on the host when they run: a call among them or under them, made without `!`, is of a
55	# definition under which a parallel let runs. Past the steps the walk takes, it answers no.
56	def forks(nodes: List<&2, Lint.Node>, program: Program) -> Bool:
57	  match program:
58	    case Program{defs, bangs}:
59	      GpuFit.reaches(
60	        4096n,
61	        Some{Lint.visit(~List<&2, U32>, ~String, ~host.visit, bangs, nodes)},
62	        [],
63	        defs,
64	        False{}
65	      )
66	
67	# The nodes of a let (`kids`) that run before the rest of its block, its last node: its binders and values.
68	def let.values(+kids: List<&2, Lint.Node>) -> List<&2, Lint.Node>:
69	  List.take(&2, Lint.Node, kids, Nat.sub(Lists.length(&2, Lint.Node, kids), 1n))
70	
71	# Whether a fork runs before the node below `path`, in the same evaluation, or `found`: a let above it, whose rest of
72	# block it is in, forks (a parallel let, or a value that forks), and no lambda (the rest of an IO step) comes between.
73	def after(path: List<&2, Lint.Above>, +program: Program, found: Bool) -> Bool:
74	  match path found:
75	    case _ True{}:
76	      True{}
77	    case Nil{} False{}:
78	      False{}
79	    case Con{Lint.Above{Lint.Node{Lint.Lambda{b, q}, n, s, k, f}, at}, rest} False{}:
80	      False{}
81	    case Con{Lint.Above{Lint.Node{Lint.Body{}, n, s, k, f}, at}, rest} False{}:
82	      False{}
83	    case Con{Lint.Above{Lint.Node{Lint.Destructure{q}, n, s, +kids, f}, at}, rest} False{}:
84	      after(
85	        rest,
86	        program,
87	        Deep.is_last(
88	          kids,
89	          at
90	        )
91	          && (Nat.is_gt(Deep.binders(kids, 0n), 1n) || forks(let.values(kids), program))
92	      )
93	    case Con{Lint.Above{Lint.Node{Lint.Let{}, n, s, +kids, f}, at}, rest} False{}:
94	      after(rest, program, Deep.is_last(kids, at) && forks(let.values(kids), program))
95	    case Con{a, rest} False{}:
96	      after(rest, program, False{})
97	
98	# The finding about an application below `path`, if it is a `!` call of a definition that forks, made after a fork.
99	def visit(
100	  +program: Program,
101	  +path: List<&2, Lint.Above>,
102	  +node: Lint.Node
103	) -> List<&2, Lint.Finding>:
104	  match program node:
105	    case Program{
106	      +defs,
107	      +bangs
108	    } Lint.Node{
109	      Lint.Application{},
110	      n,
111	      span,
112	      Con{Lint.Node{Lint.Reference{+t}, callee, Some{Lint.Span{p, beg, end}}, k, f}, args},
113	      facts
114	    }:
115	      Lint.finding(
116	        Lint.Warning{},
117	        List.contains(
118	          ~U32,
119	          ~U32.is_eq,
120	          bangs,
121	          end
122	        )
123	          && GpuFit.reaches(4096n, Some{[t]}, [], defs, True{})
124	          && Bool.not(GpuFit.is_drop(Names.find(GpuFit.Def, defs, t), t))
125	          && (forks(args, Program{defs, bangs}) || after(path, Program{defs, bangs}, False{})),
126	        span,
127	        message(callee)
128	      )
129	    case _ _:
130	      []
131	
132	def run(+here: Lint.Here) -> Lint.Ask(List<&2, Lint.Finding>):
133	  Lint.done(
134	    Lint.visit(
135	      ~Program,
136	      ~Lint.Finding,
137	      ~visit,
138	      Program{
139	        GpuFit.program.defs(Lint.Here.sources(here), Map.new(&2, Maybe<&2, GpuFit.Def>)),
140	        GpuFit.bangs(Lint.Here.text(here), 0, [])
141	      },
142	      Lint.Here.nodes(here)
143	    )
144	  )
145	
146	def rule() -> Lint.Rule:
147	  Lint.Rule{Lint.Needs{False{}, False{}}, [], Lint.OnFile{run}}