def product_tree

bendtorch-core@0.1.0.0 · src/tensor/matmul.bend:598

import Base
import bendtorch-core@0.1.0.0/src/tensor/parallel.bend as Parallel

598# The cells base .. base + 2^depth on one thread: `Parallel.tree` with the
599# product cell, but each subtree of eight cells in one row is a row block.
600def product_tree(+layout: Layout, +depth: Nat, +count: Nat, +base: Nat, state: Array<F32> & Array<F32>) ->
601  Array<F32> & (Array<F32> & Array<F32>):
602  match depth:
603    case 0n:
604      Parallel.tree(~Layout, ~Array<F32>, ~Array<F32>, ~product_cell_at, layout, 0n, count, base, state)
605    case 1n+d1:
606      match d1:
607        case 0n:
608          Parallel.tree(~Layout, ~Array<F32>, ~Array<F32>, ~product_cell_at, layout, 1n, count, base, state)
609        case 1n+d2:
610          match d2:
611            case 0n:
612              Parallel.tree(~Layout, ~Array<F32>, ~Array<F32>, ~product_cell_at, layout, 2n, count, base, state)
613            case 1n+d3:
614              match d3:
615                case 0n:
616                  block_or_tree(layout, count, base, state, one_row(layout, count, base))
617                case 1n+d4:
618                  Parallel.node(~Array<F32>, ~Array<F32>, l => r => product_tree(layout, d1, count, Parallel.upper(d1, base), (l, r)),
619                    product_tree(layout, d1, count, base, state))