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.2.2.0/linter/lint.bend as Lint
16	import bend-lint-api@0.2.2.0/shared/lists.bend as Lists
17	import bend-lint-api@0.2.2.0/shared/names.bend as Names
18	import bend-lint-api@0.2.2.0/shared/sets.bend as Sets
19	import ./deep_call_in_fork.bend as Deep
20	import ./gpu_fit.bend as GpuFit
21	
22	# Types
23	# =====
24	
25	# What judging a file's calls reads: the program's definitions, by the name the check gives them; those under which a
26	# parallel let runs (see GpuFit.program.under); and where the file writes `!(`.
27	type Program is Data:
28	  Program{defs: Map<&2, Maybe<&2, Lint.Def>>, under: Set(), bangs: List<&2, U32>}
29	
30	# Functions
31	# =========
32	def message(callee: String) -> String:
33	  callee
34	    ++ "! 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."
35	
36	# The definition an application below `path` calls by name on the host: not with `!`, as a `!` call runs on the GPU,
37	# and a `!` call after it stays there.
38	def host.visit(
39	  +bangs: List<&2, U32>,
40	  +path: List<&2, Lint.Above>,
41	  +node: Lint.Node
42	) -> List<&2, String>:
43	  match node:
44	    case Lint.Node{
45	      Lint.Application{},
46	      n,
47	      s,
48	      Con{Lint.Node{Lint.Reference{t}, rn, Some{Lint.Span{p, beg, end}}, rk, rf}, args},
49	      f
50	    }:
51	      Bool.pick(List<&2, String>, List.contains(~U32, ~U32.is_eq, bangs, end), [], [t])
52	    case _:
53	      []
54	
55	# Whether `nodes` fork on the host when they run: a call among them or under them, made without `!`, is of a
56	# definition under which a parallel let runs.
57	def forks(nodes: List<&2, Lint.Node>, program: Program) -> Bool:
58	  match program:
59	    case Program{defs, under, bangs}:
60	      Sets.has_any(Lint.visit(~List<&2, U32>, ~String, ~host.visit, bangs, nodes), under, False{})
61	
62	# The nodes of a let (`kids`) that run before the rest of its block, its last node: its binders and values.
63	def let.values(+kids: List<&2, Lint.Node>) -> List<&2, Lint.Node>:
64	  List.take(&2, Lint.Node, kids, Nat.sub(Lists.length(&2, Lint.Node, kids), 1n))
65	
66	# Whether the let of `kids` forks before the rest of its block, if the node is in it (`last`): it is a parallel let
67	# (`parallel`), or a value of it forks.
68	def let.forks(last: Bool, parallel: Bool, kids: List<&2, Lint.Node>, program: Program) -> Bool:
69	  match last parallel:
70	    case False{} _:
71	      False{}
72	    case True{} True{}:
73	      True{}
74	    case True{} False{}:
75	      forks(let.values(kids), program)
76	
77	# Whether a fork runs before the node below `path`, in the same evaluation, or `found`: a let above it, whose rest of
78	# block it is in, forks (a parallel let, or a value that forks), and no lambda (the rest of an IO step) comes between.
79	def after(path: List<&2, Lint.Above>, +program: Program, found: Bool) -> Bool:
80	  match path found:
81	    case _ True{}:
82	      True{}
83	    case Nil{} False{}:
84	      False{}
85	    case Con{Lint.Above{Lint.Node{Lint.Lambda{b, q}, n, s, k, f}, at}, rest} False{}:
86	      False{}
87	    case Con{Lint.Above{Lint.Node{Lint.Body{}, n, s, k, f}, at}, rest} False{}:
88	      False{}
89	    case Con{Lint.Above{Lint.Node{Lint.Destructure{q}, n, s, +kids, f}, at}, rest} False{}:
90	      after(
91	        rest,
92	        program,
93	        let.forks(Deep.is_last(kids, at), Nat.is_gt(Deep.binders(kids, 0n), 1n), kids, program)
94	      )
95	    case Con{Lint.Above{Lint.Node{Lint.Let{}, n, s, +kids, f}, at}, rest} False{}:
96	      after(rest, program, let.forks(Deep.is_last(kids, at), False{}, kids, program))
97	    case Con{a, rest} False{}:
98	      after(rest, program, False{})
99	
100	# Whether a fork runs before a call with `args` below `path`, in the same evaluation, if it is `asked`: an argument
101	# forks, or a let above it does.
102	def after.if(
103	  asked: Bool,
104	  args: List<&2, Lint.Node>,
105	  path: List<&2, Lint.Above>,
106	  +program: Program
107	) -> Bool:
108	  match asked:
109	    case True{}:
110	      after(path, program, forks(args, program))
111	    case False{}:
112	      False{}
113	
114	# The finding about an application below `path`, if it is a `!` call of a definition that forks, made after a fork.
115	def visit(
116	  +program: Program,
117	  +path: List<&2, Lint.Above>,
118	  +node: Lint.Node
119	) -> List<&2, Lint.Finding>:
120	  match program node:
121	    case Program{
122	      +defs,
123	      +under,
124	      +bangs
125	    } Lint.Node{
126	      Lint.Application{},
127	      n,
128	      span,
129	      Con{Lint.Node{Lint.Reference{+t}, callee, Some{Lint.Span{p, beg, end}}, k, f}, args},
130	      facts
131	    }:
132	      Lint.finding(
133	        Lint.Warning{},
134	        after.if(
135	          List.contains(
136	            ~U32,
137	            ~U32.is_eq,
138	            bangs,
139	            end
140	          )
141	            && Sets.has(under, t)
142	            && Bool.not(GpuFit.is_drop(Names.find(Lint.Def, defs, t), t)),
143	          args,
144	          path,
145	          Program{defs, under, bangs}
146	        ),
147	        span,
148	        message(callee)
149	      )
150	    case _ _:
151	      []
152	
153	# The findings in the file of `here`, which writes `!(` where `bangs` say: none, if it writes none.
154	def scanned(+here: Lint.Here, bangs: List<&2, U32>) -> Lint.Ask(List<&2, Lint.Finding>):
155	  match bangs:
156	    case Nil{}:
157	      Lint.done([])
158	    case Con{b, rest}:
159	      +sources = Lint.Here.sources(here)
160	      +defs = Lint.Program.defs(sources, Map.new(&2, Maybe<&2, Lint.Def>))
161	      Lint.done(
162	        Lint.visit(
163	          ~Program,
164	          ~Lint.Finding,
165	          ~visit,
166	          Program{defs, GpuFit.program.under(sources, defs), Con{b, rest}},
167	          Lint.Here.nodes(here)
168	        )
169	      )
170	
171	def run(+here: Lint.Here) -> Lint.Ask(List<&2, Lint.Finding>):
172	  scanned(here, GpuFit.bangs(Lint.Here.text(here), 0, []))
173	
174	def rule() -> Lint.Rule:
175	  Lint.Rule{Lint.Needs{False{}, False{}}, [], Lint.OnFile{run}}