tensor_select.bend
stelliferous@0.0.3.0 · 79 lines · 0 definitions
1# Three-input selection: one band kernel over operands brought to one layout.
2import Base
3import stelliferous@0.0.3.0/tensor.bend as Tensors
4import stelliferous@0.0.3.0/tensor_shape.bend as Shapes
5import stelliferous@0.0.3.0/tensor_lanes.bend as Lanes
6import stelliferous@0.0.3.0/tensor_band.bend as Bands
7import stelliferous@0.0.3.0/tensor_band_schedule.bend as Schedule
8import stelliferous@0.0.3.0/storage_buffer.bend as Storage
9import stelliferous@0.0.3.0/traversal_partition.bend as Partition
10
11def expanded(same: Bool,shape: List<&2,U32>,tensor: Tensors.Tensor<F32>) -> Tensors.Tensor<F32>:
12 match same:
13 case True{}: tensor
14 case False{}: Tensors.broadcast_to(~F32,tensor,shape)
15
16def selected_leaf(geometry: Bands.Geometry,first: U32,rows: U32,
17 input: Bands.Paired<Storage.Buffer<F32>,Bands.Paired<Storage.Buffer<F32>,Storage.Buffer<F32>>>,resource: Unit) -> Bands.Bands(~F32):
18 match input:
19 case Bands.Both{yes,Bands.Both{mask,no}}:
20 +rows = rows
21 Partition.Leaf{Lanes.select(U32.to_nat(Bands.local_count(geometry,rows)),mask,yes,no),first,rows,Unit{}}
22 case Bands.Both{yes,Bands.Only{mask,own_first,own_rows}}: Partition.Leaf{yes,first,rows,Unit{}}
23 case Bands.Only{yes,own_first,own_rows}: Partition.Leaf{yes,own_first,own_rows,Unit{}}
24
25def 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):
26 Bands.each(~Bands.Paired<Storage.Buffer<F32>,Bands.Paired<Storage.Buffer<F32>,Storage.Buffer<F32>>>,~Storage.Buffer<F32>,~Unit,~Bands.Geometry,
27 ~Bands.unshared,~selected_leaf,~Bands.left_of(~Storage.Buffer<F32>,~Bands.Paired<Storage.Buffer<F32>,Storage.Buffer<F32>>),
28 layout,parallel,geometry,
29 Bands.zip(~Storage.Buffer<F32>,~Bands.Paired<Storage.Buffer<F32>,Storage.Buffer<F32>>,yes,
30 Bands.zip(~Storage.Buffer<F32>,~Storage.Buffer<F32>,mask,no)),Unit{})
31
32# The three operands meet in one layout, each arrives in it, and one band
33# kernel selects.
34def selected(target: Schedule.Target,shape: List<&2,U32>,fill: F32,+workers: U32,
35 mask: Tensors.Tensor<F32>,yes: Tensors.Tensor<F32>,no: Tensors.Tensor<F32>) -> Tensors.Tensor<F32>:
36 Schedule.Target{+geometry,+layout,split} = target
37 Tensors.Tensor{selected_bands(geometry,layout,Schedule.parallel(geometry,workers,16.0,0.875,layout),
38 Tensors.arrive(~F32,geometry,layout,mask),Tensors.arrive(~F32,geometry,layout,yes),Tensors.arrive(~F32,geometry,layout,no)),
39 geometry,shape,fill,Tensors.Usable{},workers}
40
41def matching(+shape: List<&2,U32>,+fill: F32,+workers: U32,mask: Tensors.Tensor<F32>,yes: Tensors.Tensor<F32>,no: Tensors.Tensor<F32>) -> Tensors.Tensor<F32>:
42 Tensors.then(~Tensors.Tensor<F32>,~Schedule.Target,~Tensors.Tensor<F32>,Tensors.offer(~F32,mask),mask => mask_target =>
43 Tensors.then(~Tensors.Tensor<F32>,~Schedule.Target,~Tensors.Tensor<F32>,Tensors.offer(~F32,yes),yes => yes_target =>
44 Tensors.then(~Tensors.Tensor<F32>,~Schedule.Target,~Tensors.Tensor<F32>,Tensors.offer(~F32,no),no => no_target =>
45 selected(Schedule.meet(mask_target,Schedule.meet(yes_target,no_target)),shape,fill,workers,mask,yes,no))))
46
47def 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,
48 mask: Tensors.Tensor<F32>,yes: Tensors.Tensor<F32>,no: Tensors.Tensor<F32>) -> Tensors.Tensor<F32>:
49 match same count:
50 case True{} Some{count}: matching(shape,fill,workers,mask,yes,no)
51 case False{} Some{count}:
52 matching(shape,fill,workers,expanded(Shapes.same_extents(mask_shape,shape),shape,mask),
53 expanded(Shapes.same_extents(yes_shape,shape),shape,yes),expanded(Shapes.same_extents(no_shape,shape),shape,no))
54 case same None{}: Tensors.failed(~F32,"select",Tensors.InvalidShape{shape},shape,fill,workers)
55
56def 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,
57 mask: Tensors.Tensor<F32>,yes: Tensors.Tensor<F32>,no: Tensors.Tensor<F32>) -> Tensors.Tensor<F32>:
58 match common:
59 case Some{+shape}:
60 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)
61 case None{}: Tensors.failed(~F32,"select",Tensors.ShapeMismatch{mask_shape,values_shape},mask_shape,fill,workers)
62
63def 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,
64 mask: Tensors.Tensor<F32>,yes: Tensors.Tensor<F32>,no: Tensors.Tensor<F32>) -> Tensors.Tensor<F32>:
65 match common:
66 case Some{+shape}: masked_shape(Shapes.broadcast(mask_shape,shape),mask_shape,yes_shape,no_shape,shape,fill,workers,mask,yes,no)
67 case None{}: Tensors.failed(~F32,"select",Tensors.ShapeMismatch{yes_shape,no_shape},yes_shape,fill,workers)
68
69def status(status: Tensors.Status,mask_shape: List<&2,U32>,+yes_shape: List<&2,U32>,+no_shape: List<&2,U32>,fill: F32,workers: U32,
70 mask: Tensors.Tensor<F32>,yes: Tensors.Tensor<F32>,no: Tensors.Tensor<F32>) -> Tensors.Tensor<F32>:
71 match status:
72 case Tensors.Failed{problem}: Tensors.with_status(~F32,Tensors.Failed{problem},mask)
73 case Tensors.Usable{}: values_shape(Shapes.broadcast(yes_shape,no_shape),mask_shape,yes_shape,no_shape,fill,workers,mask,yes,no)
74
75def select(mask: Tensors.Tensor<F32>,yes: Tensors.Tensor<F32>,no: Tensors.Tensor<F32>) -> Tensors.Tensor<F32>:
76 Tensors.with_metadata(~F32,~Tensors.Tensor<F32>,mask,mask => mask_shape => mf => mt => mw =>
77 Tensors.with_metadata(~F32,~Tensors.Tensor<F32>,yes,yes => yes_shape => fill => yt => yw =>
78 Tensors.with_metadata(~F32,~Tensors.Tensor<F32>,no,no => no_shape => nf => nt => nw =>
79 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))))