tensor_select.bend
stelliferous@0.0.2.0 · 106 lines · 15 definitions
1
2import Base
3import stelliferous@0.0.2.0/tensor.bend as Tensors
4import stelliferous@0.0.2.0/tensor_shape.bend as Shapes
5import stelliferous@0.0.2.0/tensor_layout.bend as Layout
6import stelliferous@0.0.2.0/tensor_lanes.bend as Lanes
7import stelliferous@0.0.2.0/tensor_band.bend as Bands
8import stelliferous@0.0.2.0/tensor_band_schedule.bend as Schedule
9import stelliferous@0.0.2.0/storage_buffer.bend as Storage
10import stelliferous@0.0.2.0/traversal_partition.bend as Partition
11
12def dense(+shape: List<&2,U32>,+count: U32,+workers: U32,mask: Tensors.Tensor<F32>,yes: Tensors.Tensor<F32>,no: Tensors.Tensor<F32>) -> Tensors.Tensor<F32>:
13 Tensors.with_dense(~F32,~Tensors.Tensor<F32>,mask,mask => ms => mf => mt => mw =>
14 Tensors.with_dense(~F32,~Tensors.Tensor<F32>,yes,yes => ys => fill => yt => yw =>
15 Tensors.with_dense(~F32,~Tensors.Tensor<F32>,no,no => ns => nf => nt => nw =>
16 Tensors.Tensor{Lanes.select(U32.to_nat(count),mask,yes,no),shape,fill,Tensors.first_failure(mt,Tensors.first_failure(yt,nt)),workers})))
17
18def expanded(same: Bool,shape: List<&2,U32>,tensor: Tensors.Tensor<F32>) -> Tensors.Tensor<F32>:
19 match same:
20 case True{}: tensor
21 case False{}: Layout.broadcast_to(~F32,tensor,shape)
22
23type BandPlan is Data:
24 BandPlan{geometry: Bands.Geometry,layout: Bands.Plan}
25
26def planned(shape: List<&2,U32>,fill: F32,status: Tensors.Status,workers: U32,+geometry: Bands.Geometry,owners: Bands.Bands(~F32) & Bands.Plan) ->
27 Tensors.Tensor<F32> & Maybe<&2,BandPlan>:
28 (bands,+layout) = owners
29 (Tensors.Banded{bands,geometry,shape,fill,status,workers},Some{BandPlan{geometry,layout}})
30
31def plan(tensor: Tensors.Tensor<F32>) -> Tensors.Tensor<F32> & Maybe<&2,BandPlan>:
32 match tensor:
33 case Tensors.Tensor{buffer,shape,fill,status,workers}: (Tensors.Tensor{buffer,shape,fill,status,workers},None{})
34 case Tensors.Banded{bands,geometry,shape,fill,status,workers}: planned(shape,fill,status,workers,geometry,Bands.plan(~F32,bands))
35
36def common(left: Maybe<&2,BandPlan>,right: Maybe<&2,BandPlan>) -> Maybe<&2,BandPlan>:
37 match left right:
38 case Some{BandPlan{+geometry,left}} Some{BandPlan{other,right}}: Some{BandPlan{geometry,Schedule.alignment(geometry,left,right)}}
39 case Some{left} None{}: Some{left}
40 case None{} right: right
41
42def selected_leaf(geometry: Bands.Geometry,first: U32,rows: U32,
43 input: Bands.Paired<Storage.Buffer<F32>,Bands.Paired<Storage.Buffer<F32>,Storage.Buffer<F32>>>,resource: Unit) -> Bands.Bands(~F32):
44 match input:
45 case Bands.Both{yes,Bands.Both{mask,no}}:
46 +rows = rows
47 Partition.Leaf{Lanes.select(U32.to_nat(Bands.local_count(geometry,rows)),mask,yes,no),first,rows,Unit{}}
48 case Bands.Both{yes,Bands.Only{mask,own_first,own_rows}}: Partition.Leaf{yes,first,rows,Unit{}}
49 case Bands.Only{yes,own_first,own_rows}: Partition.Leaf{yes,own_first,own_rows,Unit{}}
50
51def selected_bands(+geometry: Bands.Geometry,+layout: Bands.Plan,+parallel: Bool,mask: Bands.Bands(~F32),yes: Bands.Bands(~F32),no: Bands.Bands(~F32)) -> Bands.Bands(~F32):
52 Bands.each(~Bands.Paired<Storage.Buffer<F32>,Bands.Paired<Storage.Buffer<F32>,Storage.Buffer<F32>>>,~Storage.Buffer<F32>,~Unit,~Bands.Geometry,
53 ~Bands.unshared,~selected_leaf,~Bands.left_of(~Storage.Buffer<F32>,~Bands.Paired<Storage.Buffer<F32>,Storage.Buffer<F32>>),
54 layout,parallel && Bands.same_image(geometry,layout),parallel,geometry,geometry,
55 Bands.zip(~Storage.Buffer<F32>,~Bands.Paired<Storage.Buffer<F32>,Storage.Buffer<F32>>,yes,
56 Bands.zip(~Storage.Buffer<F32>,~Storage.Buffer<F32>,mask,no)),Unit{})
57
58def dispatch(layout: Maybe<&2,BandPlan>,shape: List<&2,U32>,count: U32,fill: F32,workers: U32,
59 mask: Tensors.Tensor<F32>,yes: Tensors.Tensor<F32>,no: Tensors.Tensor<F32>) -> Tensors.Tensor<F32>:
60 match layout:
61 case None{}: dense(shape,count,workers,mask,yes,no)
62 case Some{BandPlan{+geometry,+layout}}:
63 +workers = workers
64 Tensors.Banded{selected_bands(geometry,layout,Schedule.parallel(geometry,workers,16.0,0.875,layout),
65 Tensors.aligned(~F32,geometry,layout,mask),Tensors.aligned(~F32,geometry,layout,yes),Tensors.aligned(~F32,geometry,layout,no)),
66 geometry,shape,fill,Tensors.Usable{},workers}
67
68def matching(+shape: List<&2,U32>,+count: U32,+fill: F32,+workers: U32,mask: Tensors.Tensor<F32>,yes: Tensors.Tensor<F32>,no: Tensors.Tensor<F32>) -> Tensors.Tensor<F32>:
69 Tensors.then(~Tensors.Tensor<F32>,~Maybe<&2,BandPlan>,~Tensors.Tensor<F32>,plan(mask),mask => mask_plan =>
70 Tensors.then(~Tensors.Tensor<F32>,~Maybe<&2,BandPlan>,~Tensors.Tensor<F32>,plan(yes),yes => yes_plan =>
71 Tensors.then(~Tensors.Tensor<F32>,~Maybe<&2,BandPlan>,~Tensors.Tensor<F32>,plan(no),no => no_plan =>
72 dispatch(common(mask_plan,common(yes_plan,no_plan)),shape,count,fill,workers,mask,yes,no))))
73
74def shapes(same: Bool,+shape: List<&2,U32>,count: Maybe<&2,U32>,mask_shape: List<&2,U32>,yes_shape: List<&2,U32>,no_shape: List<&2,U32>,fill: F32,workers: U32,
75 mask: Tensors.Tensor<F32>,yes: Tensors.Tensor<F32>,no: Tensors.Tensor<F32>) -> Tensors.Tensor<F32>:
76 match same count:
77 case True{} Some{count}: matching(shape,count,fill,workers,mask,yes,no)
78 case False{} Some{count}:
79 dense(shape,count,workers,expanded(Shapes.same_extents(mask_shape,shape),shape,mask),
80 expanded(Shapes.same_extents(yes_shape,shape),shape,yes),expanded(Shapes.same_extents(no_shape,shape),shape,no))
81 case same None{}: Tensors.failed(~F32,"select",Tensors.InvalidShape{shape},shape,fill,workers)
82
83def masked_shape(common: Maybe<&2,List<&2,U32>>,+mask_shape: List<&2,U32>,+yes_shape: List<&2,U32>,+no_shape: List<&2,U32>,values_shape: List<&2,U32>,fill: F32,workers: U32,
84 mask: Tensors.Tensor<F32>,yes: Tensors.Tensor<F32>,no: Tensors.Tensor<F32>) -> Tensors.Tensor<F32>:
85 match common:
86 case Some{+shape}:
87 shapes(Shapes.same_extents(mask_shape,yes_shape) && Shapes.same_extents(yes_shape,no_shape),shape,Tensors.count_of(shape),mask_shape,yes_shape,no_shape,fill,workers,mask,yes,no)
88 case None{}: Tensors.failed(~F32,"select",Tensors.ShapeMismatch{mask_shape,values_shape},mask_shape,fill,workers)
89
90def values_shape(common: Maybe<&2,List<&2,U32>>,+mask_shape: List<&2,U32>,+yes_shape: List<&2,U32>,no_shape: List<&2,U32>,fill: F32,workers: U32,
91 mask: Tensors.Tensor<F32>,yes: Tensors.Tensor<F32>,no: Tensors.Tensor<F32>) -> Tensors.Tensor<F32>:
92 match common:
93 case Some{+shape}: masked_shape(Shapes.broadcast(mask_shape,shape),mask_shape,yes_shape,no_shape,shape,fill,workers,mask,yes,no)
94 case None{}: Tensors.failed(~F32,"select",Tensors.ShapeMismatch{yes_shape,no_shape},yes_shape,fill,workers)
95
96def status(status: Tensors.Status,mask_shape: List<&2,U32>,+yes_shape: List<&2,U32>,+no_shape: List<&2,U32>,fill: F32,workers: U32,
97 mask: Tensors.Tensor<F32>,yes: Tensors.Tensor<F32>,no: Tensors.Tensor<F32>) -> Tensors.Tensor<F32>:
98 match status:
99 case Tensors.Failed{problem}: Tensors.with_status(~F32,Tensors.Failed{problem},mask)
100 case Tensors.Usable{}: values_shape(Shapes.broadcast(yes_shape,no_shape),mask_shape,yes_shape,no_shape,fill,workers,mask,yes,no)
101
102def select(mask: Tensors.Tensor<F32>,yes: Tensors.Tensor<F32>,no: Tensors.Tensor<F32>) -> Tensors.Tensor<F32>:
103 Tensors.with_metadata(~F32,~Tensors.Tensor<F32>,mask,mask => mask_shape => mf => mt => mw =>
104 Tensors.with_metadata(~F32,~Tensors.Tensor<F32>,yes,yes => yes_shape => fill => yt => yw =>
105 Tensors.with_metadata(~F32,~Tensors.Tensor<F32>,no,no => no_shape => nf => nt => nw =>
106 status(Tensors.first_failure(mt,Tensors.first_failure(yt,nt)),mask_shape,yes_shape,no_shape,fill,U32.max(mw,U32.max(yw,nw)),mask,yes,no))))