1 # Three-input selection: one band kernel over operands brought to one layout. 2 import Base 3 import ./tensor.bend as Tensors 4 import ./tensor_shape.bend as Shapes 5 import ./tensor_lanes.bend as Lanes 6 import ./tensor_band.bend as Bands 7 import ./tensor_band_schedule.bend as Schedule 8 import ./storage_buffer.bend as Storage 9 import ./traversal_partition.bend as Partition 10 11 def expanded(same: Bool,shape: List<&2,U32>,tensor: Tensors.Tensor) -> Tensors.Tensor: 12 match same: 13 case True{}: tensor 14 case False{}: Tensors.broadcast_to(~F32,tensor,shape) 15 16 def selected_leaf(geometry: Bands.Geometry,first: U32,rows: U32, 17 input: Bands.Paired,Bands.Paired,Storage.Buffer>>,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 25 def 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,Bands.Paired,Storage.Buffer>>,~Storage.Buffer,~Unit,~Bands.Geometry, 27 ~Bands.unshared,~selected_leaf,~Bands.left_of(~Storage.Buffer,~Bands.Paired,Storage.Buffer>), 28 layout,parallel,geometry, 29 Bands.zip(~Storage.Buffer,~Bands.Paired,Storage.Buffer>,yes, 30 Bands.zip(~Storage.Buffer,~Storage.Buffer,mask,no)),Unit{}) 31 32 # The three operands meet in one layout, each arrives in it, and one band 33 # kernel selects. 34 def selected(target: Schedule.Target,shape: List<&2,U32>,fill: F32,+workers: U32, 35 mask: Tensors.Tensor,yes: Tensors.Tensor,no: Tensors.Tensor) -> Tensors.Tensor: 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 41 def matching(+shape: List<&2,U32>,+fill: F32,+workers: U32,mask: Tensors.Tensor,yes: Tensors.Tensor,no: Tensors.Tensor) -> Tensors.Tensor: 42 Tensors.then(~Tensors.Tensor,~Schedule.Target,~Tensors.Tensor,Tensors.offer(~F32,mask),mask => mask_target => 43 Tensors.then(~Tensors.Tensor,~Schedule.Target,~Tensors.Tensor,Tensors.offer(~F32,yes),yes => yes_target => 44 Tensors.then(~Tensors.Tensor,~Schedule.Target,~Tensors.Tensor,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 47 def 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,yes: Tensors.Tensor,no: Tensors.Tensor) -> Tensors.Tensor: 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 56 def 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,yes: Tensors.Tensor,no: Tensors.Tensor) -> Tensors.Tensor: 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 63 def 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,yes: Tensors.Tensor,no: Tensors.Tensor) -> Tensors.Tensor: 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 69 def 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,yes: Tensors.Tensor,no: Tensors.Tensor) -> Tensors.Tensor: 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 75 def select(mask: Tensors.Tensor,yes: Tensors.Tensor,no: Tensors.Tensor) -> Tensors.Tensor: 76 Tensors.with_metadata(~F32,~Tensors.Tensor,mask,mask => mask_shape => mf => mt => mw => 77 Tensors.with_metadata(~F32,~Tensors.Tensor,yes,yes => yes_shape => fill => yt => yw => 78 Tensors.with_metadata(~F32,~Tensors.Tensor,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))))