tensor_view.bend source

tensor_view.bend on the hub · documented module

# Rank-independent forward-strided metadata. The Array owner is passed separately.# A vector, matrix and tensor use the same descriptor and lookup operation.import Baseimport ./storage_buffer.bend as Storageimport ./tensor_shape.bend as Shapesimport ./machine_limits.bend as Machinetype Axis is Data:  Axis{extent: U32,stride: U32}type View is Data:  View{axes: List<&2,Axis>,offset: U32}# Successful addresses are below logical_length. The checked product and the# remaining-room comparison prevent multiplication and addition from wrapping.def address_when(valid: Bool,offset: U32,increment: U32) -> Maybe<&2,U32>:  match valid:    case False{}: None{}    case True{}: Some{(offset + increment : U32)}def finish_address(valid: Bool,offset: U32) -> Maybe<&2,U32>:  match valid:    case False{}: None{}    case True{}: Some{offset}def advance_product(product: Maybe<&2,U32>,offset: U32,room: U32) -> Maybe<&2,U32>:  match product:    case None{}: None{}    case Some{+increment}: address_when(U32.is_le(increment,room),offset,increment)def advance_inside(inside: Bool,+offset: U32,coordinate: U32,+stride: U32,logical_length: U32) -> Maybe<&2,U32>:  match inside:    case False{}: None{}    case True{}: advance_product(Shapes.checked_multiply(coordinate,stride),offset,(logical_length - 1 - offset : U32))def advance(+offset: U32,coordinate: U32,stride: U32,+logical_length: U32) -> Maybe<&2,U32>:  advance_inside(U32.is_lt(offset,logical_length),offset,coordinate,stride,logical_length)# The recursive step consumes one axis and one coordinate; no flattening/copy.def locate(axes: List<&2,Axis>,coordinates: List<&2,U32>,address: Maybe<&2,U32>,+logical_length: U32) -> Maybe<&2,U32>:  match axes coordinates address:    case Nil{} Nil{} Some{+offset}:      finish_address(U32.is_lt(offset,logical_length),offset)    case Con{Axis{+extent,stride},tail} Con{+coordinate,rest} Some{+offset}:      locate(tail,rest,advance_inside(U32.is_lt(coordinate,extent) && U32.is_lt(offset,logical_length),offset,coordinate,stride,logical_length),logical_length)    case axes coordinates address: None{}def address(view: View,coordinates: List<&2,U32>,logical_length: U32) -> Maybe<&2,U32>:  View{axes,offset} = view  locate(axes,coordinates,Some{offset},logical_length)def axis_extents(axes: List<&2,Axis>) -> List<&2,U32>:  match axes:    case Nil{}: Nil{}    case Con{Axis{extent,stride},tail}: extent <> axis_extents(tail)def shape(view: View) -> Shapes.Shape:  View{axes,offset} = view  Shapes.Shape{axis_extents(axes)}def numel(view: View) -> Maybe<&2,U32>:  Shapes.numel(shape(view))def rank(view: View) -> Nat:  Shapes.rank(shape(view))def extent(view: View,axis: Nat) -> Maybe<&2,U32>:  Shapes.extent(shape(view),axis)def build_axes(reversed: List<&2,U32>,stride: Maybe<&2,U32>,axes: List<&2,Axis>,offset: U32) -> Maybe<&2,View>:  match reversed stride:    case Nil{} Some{value}: Some{View{axes,offset}}    case Con{+size,tail} Some{+value}:      build_axes(tail,Shapes.checked_multiply(value,size),Axis{size,value} <> axes,offset)    case reversed None{}: None{}def contiguous_count(count: Maybe<&2,U32>,extents: List<&2,U32>,offset: U32) -> Maybe<&2,View>:  match count:    case None{}: None{}    case Some{+value}:      build_axes(List.reverse(&2,U32,extents),Some{Bool.pick(U32,U32.is_zero(value),0,1)},Nil{},offset)def contiguous(+dimensions: Shapes.Shape,offset: U32) -> Maybe<&2,View>:  Shapes.Shape{extents} = dimensions  contiguous_count(Shapes.numel(dimensions),extents,offset)def contiguous_product(valid: Bool,extent: U32,inner_count: U32) -> Maybe<&2,U32>:  match valid:    case False{}: None{}    case True{}: Shapes.checked_multiply(extent,inner_count)def contiguous_axis(+extent: U32,stride: U32,inner: Maybe<&2,U32>) -> Maybe<&2,U32>:  match inner:    case None{}: None{}    case Some{+count}: contiguous_product(U32.is_le(extent,1) || U32.is_eq(stride,count),extent,count)# Validate strides from the innermost axis outward. This avoids allocating a# second canonical View just to compare its metadata. Empty views bypass it.def contiguous_size(axes: List<&2,Axis>) -> Maybe<&2,U32>:  match axes:    case []: Some{1}    case Axis{extent,stride} <> tail: contiguous_axis(extent,stride,contiguous_size(tail))def has_count(count: Maybe<&2,U32>) -> Bool:  match count:    case None{}: False{}    case Some{value}: True{}def contiguous_nonempty(empty: Bool,+axes: List<&2,Axis>) -> Bool:  match empty:    case True{}: True{}    case False{}: has_count(contiguous_size(axes))def is_contiguous(view: View) -> Bool:  View{+axes,offset} = view  contiguous_nonempty(Shapes.has_zero(axis_extents(axes)),axes)def last_coordinates(axes: List<&2,Axis>) -> List<&2,U32>:  match axes:    case Nil{}: Nil{}    case Con{Axis{size,stride},tail}: (size - 1 : U32) <> last_coordinates(tail)def found(address: Maybe<&2,U32>) -> Bool:  match address:    case None{}: False{}    case Some{value}: True{}def fits_nonempty(empty: Bool,+view: View,+logical_length: U32) -> Bool:  match empty:    case True{}:      View{axes,offset} = view      U32.is_le(offset,logical_length)    case False{}:      View{+axes,offset} = view      found(address(View{axes,offset},last_coordinates(axes),logical_length))# Forward strides are nonnegative: the last coordinate is the maximum address.# Empty views may point exactly one past the logical end, and perform no reads.def axis_extents_of(view: View) -> List<&2,U32>:  View{axes,offset} = view  axis_extents(axes)def fits(view: View,logical_length: U32) -> Bool:  +view = view  fits_nonempty(Shapes.has_zero(axis_extents_of(view)),view,logical_length)def contains_index(indices: List<&2,Nat>,+axis: Nat) -> Bool:  match indices:    case Nil{}: False{}    case Con{index,tail}: Nat.is_eq(index,axis) || contains_index(tail,axis)def unique_axes(order: List<&2,Nat>) -> Bool:  match order:    case Nil{}: True{}    case Con{axis,+tail}: Bool.not(contains_index(tail,axis)) && unique_axes(tail)def prepend_axis(axis: Maybe<&2,Axis>,view: Maybe<&2,View>) -> Maybe<&2,View>:  match axis view:    case Some{item} Some{View{axes,offset}}: Some{View{item <> axes,offset}}    case axis view: None{}def reorder(order: List<&2,Nat>,+axes: List<&2,Axis>,offset: U32) -> Maybe<&2,View>:  match order:    case Nil{}: Some{View{Nil{},offset}}    case Con{index,tail}:      prepend_axis(List.get(&2,Axis,axes,index),reorder(tail,axes,offset))def permute_when(valid: Bool,view: View,order: List<&2,Nat>) -> Maybe<&2,View>:  match valid:    case False{}: None{}    case True{}:      View{axes,offset} = view      reorder(order,axes,offset)# Order must contain each axis exactly once, including singleton axes.def permute(+view: View,+order: List<&2,Nat>) -> Maybe<&2,View>:  permute_when(Nat.is_eq(rank(view),List.length(&2,Nat,order)) && unique_axes(order),view,order)def slice_step(zero: Bool,+start: U32,+count: U32,+step: U32,+size: U32) -> Bool:  match zero:    case True{}: False{}    case False{}:      U32.is_le(start,size) && (U32.is_zero(count) ||        (U32.is_lt(start,size) && U32.is_le((count - 1 : U32),((size - 1 - start) / step : U32))))def sliced_axis(valid: Bool,count: U32,stride: Maybe<&2,U32>,offset: Maybe<&2,U32>,tail: List<&2,Axis>) -> Maybe<&2,View>:  match valid stride offset:    case True{} Some{stride} Some{offset}: Some{View{Axis{count,stride} <> tail,offset}}    case valid stride offset: None{}def add_offset(shift: Maybe<&2,U32>,+offset: U32) -> Maybe<&2,U32>:  match shift:    case None{}: None{}    case Some{+value}: address_when(U32.is_le(value,(Machine.word_max() - offset : U32)),offset,value)def slice_axes(axes: List<&2,Axis>,axis: Nat,+start: U32,+count: U32,+step: U32,+offset: U32) -> Maybe<&2,View>:  match axes axis:    case Con{Axis{size,+stride},tail} 0n:      sliced_axis(slice_step(U32.is_zero(step),start,count,step,size),count,Shapes.checked_multiply(stride,step),        add_offset(Shapes.checked_multiply(start,stride),offset),tail)    case Con{head,tail} 1n+rest:      prepend_axis(Some{head},slice_axes(tail,rest,start,count,step,offset))    case Nil{} axis: None{}# Explicit start/count/positive step. The resulting view is validated against# the actual Buffer at execution time; arithmetic overflow is rejected here.def slice(view: View,axis: Nat,start: U32,count: U32,step: U32) -> Maybe<&2,View>:  View{axes,offset} = view  slice_axes(axes,axis,start,count,step,offset)# One axis read as extent elements: unchanged when the extents agree, the# single element repeated (stride 0) when the axis has extent 1.def stretched_axis(same: Bool,single: Bool,size: U32,stride: U32,extent: U32) -> Maybe<&2,Axis>:  match same single:    case True{} single: Some{Axis{size,stride}}    case False{} True{}: Some{Axis{extent,0}}    case False{} False{}: None{}def stretched_prepend(axis: Maybe<&2,Axis>,rest: Maybe<&2,List<&2,Axis>>) -> Maybe<&2,List<&2,Axis>>:  match axis rest:    case Some{axis} Some{rest}: Some{axis <> rest}    case axis rest: None{}def stretch_axes(axes: List<&2,Axis>,target: List<&2,U32>) -> Maybe<&2,List<&2,Axis>>:  match axes target:    case Nil{} Nil{}: Some{[]}    case Con{Axis{+size,stride},tail} Con{+extent,rest}:      stretched_prepend(stretched_axis(U32.is_eq(size,extent),U32.is_eq(size,1),size,stride,extent),stretch_axes(tail,rest))    case axes target: None{}# The axes with leading axes of extent 1 added up to rank; the same elements.def with_rank(+axes: List<&2,Axis>,rank: Nat) -> List<&2,Axis>:  List.append(&2,Axis,List.replicate(Axis,Nat.sub(rank,List.length(&2,Axis,axes)),Axis{1,0}),axes)def stretched_view(axes: Maybe<&2,List<&2,Axis>>,offset: U32) -> Maybe<&2,View>:  match axes:    case Some{axes}: Some{View{axes,offset}}    case None{}: None{}# The view read with the extents target under NumPy broadcasting. No element# moves: an axis of extent 1 is repeated with stride 0, and so is an axis the# view lacks.def stretch(view: View,+target: List<&2,U32>) -> Maybe<&2,View>:  View{axes,offset} = view  stretched_view(stretch_axes(with_rank(axes,List.length(&2,U32,target)),target),offset)