config.decl @tiny.context : %value: index where [range(%value, 1, 262144)]

config.decl @tiny.spans : %value: index where [range(%value, 1, 262144)]

config.decl @tiny.tokens : %value: index where [range(%value, 1, 524288)]

config.def @tiny.bias_delta = 0 : i32

// Logical positions resolve through the native row-major page map. The opaque
// row record supplies the byte origin of this model's pooled history plane.
// Admission bounds positions by context and pages by blocks; these facts also
// hold when the helper remains a callable function instead of being inlined.
func.def @history_offset(%origins: buffer, %row: index, %position: index, %blocks: index) -> (offset) {
  %c0 = index.constant 0 : offset
  %c4 = index.constant 4 : index
  %word = index.constant 4 : offset
  %record_stride = index.constant 40 : offset
  %bounded_row = index.assume %row [range(%row, 0, 2147483646)] : index
  %record_offset = index.scale %bounded_row, %record_stride : index, offset -> offset
  %bounded_position = index.assume %position [range(%position, 0, 262143)] : index
  %records = buffer.view %origins[%record_offset] : buffer -> view<5xi64>
  %base64 = view.load %records[2] : view<5xi64> -> i64
  %map64 = view.load %records[4] : view<5xi64> -> i64
  %map_bound = scalar.assume %map64 [range(%map64, 0, 9223372036854775807)] : i64
  %map_origin = index.cast %map_bound : i64 to offset
  %base_positive = scalar.assume %base64 [range(%base64, 0, 9223372036854775807)] : i64
  %base = index.cast %base_positive : i64 to offset
  %page_value = index.div %bounded_position, %c4 : index
  %c1 = index.constant 1 : index
  %last = index.sub %blocks, %c1 : index
  %page = index.assume %page_value [range(%page_value, 0, %last), lt(%page_value, %blocks)] : index
  %map = buffer.view %origins[%map_origin] : buffer -> view<[%blocks]xi32>
  %physical32 = view.load %map[%page] : view<[%blocks]xi32> -> i32
  %bounded_physical = scalar.assume %physical32 [range(%physical32, 0, 1048575)] : i32
  %physical = index.cast %bounded_physical : i32 to index
  %block_position = index.mul %physical, %c4 : index
  %within_page = index.rem %bounded_position, %c4 : index
  %slot = index.add %block_position, %within_page : index
  %bytes = index.scale %slot, %word : index, offset -> offset
  %address = index.add %base, %bytes : offset
  func.return %address : offset
}

// One scalar worker per live span is intentional: this is an ownership/ABI
// witness, not a transformer performance proxy. Chunking cannot affect math.
kernel.def @consume() {
  %spans = config.get @tiny.spans : index
  %c1 = index.constant 1 : index
  kernel.launch.config workgroups(%spans, %c1, %c1) workgroup_size(%c1, %c1, %c1) : index
} launch(%bias: view<1xi32>, %residual: buffer, %metadata: buffer, %origins: buffer, %arena: buffer, %tokens: buffer) {
  %spans = config.get @tiny.spans : index
  %span = kernel.workgroup.id<x> : index
  %origin = index.constant 0 : offset
  %c16 = index.constant 16 : offset
  %header = buffer.view %metadata[%origin] : buffer -> view<4xi32>
  %count32 = view.load %header[0] : view<4xi32> -> i32
  %count = index.cast %count32 : i32 to index
  %live = index.cmp ult, %span, %count : index
  scf.if %live {
    %packets = buffer.view %metadata[%c16] : buffer -> view<[%spans]x9xi32>
    %length32 = view.load %packets[%span, 0] : view<[%spans]x9xi32> -> i32
    %position32 = view.load %packets[%span, 1] : view<[%spans]x9xi32> -> i32
    %source32 = view.load %packets[%span, 2] : view<[%spans]x9xi32> -> i32
    %row32 = view.load %packets[%span, 3] : view<[%spans]x9xi32> -> i32
    %input32 = view.load %packets[%span, 4] : view<[%spans]x9xi32> -> i32
    %length_bound = scalar.assume %length32 [range(%length32, 1, 262144)] : i32
    %position_bound = scalar.assume %position32 [range(%position32, 0, 262143)] : i32
    %source_bound = scalar.assume %source32 [range(%source32, 0, 524287)] : i32
    %row_bound = scalar.assume %row32 [range(%row32, 0, 2147483646)] : i32
    %row_value = index.cast %row_bound : i32 to index
    %input_bound = scalar.assume %input32 [range(%input32, 0, 262143)] : i32
    %length = index.cast %length_bound : i32 to index
    %position = index.cast %position_bound : i32 to index
    %source = index.cast %source_bound : i32 to index
    %row = index.assume %row_value [range(%row_value, 0, 2147483646)] : index
    %input = index.cast %input_bound : i32 to index
    %record_stride = index.constant 40 : offset
    %record_offset = index.scale %row, %record_stride : index, offset -> offset
    %records = buffer.view %origins[%record_offset] : buffer -> view<5xi64>
    %read64 = view.load %records[0] : view<5xi64> -> i64
    %write64 = view.load %records[1] : view<5xi64> -> i64
    %fork64 = view.load %records[3] : view<5xi64> -> i64
    %position64 = scalar.extui %position32 : i32 to i64
    %fork = scalar.cmpi eq, %position64, %fork64 : i64
    %reader64 = scf.select %fork, %read64, %write64 : i64
    %reader_bound = scalar.assume %reader64 [range(%reader64, 0, 9223372036854775807)] : i64
    %writer_bound = scalar.assume %write64 [range(%write64, 0, 9223372036854775807)] : i64
    %reader = index.cast %reader_bound : i64 to offset
    %writer = index.cast %writer_bound : i64 to offset
    %read_state = buffer.view %arena[%reader] : buffer -> view<1xi32>
    %write_state = buffer.view %arena[%writer] : buffer -> view<1xi32>
    %context = config.get @tiny.context : index
    %c0 = index.constant 0 : index
    %c1 = index.constant 1 : index
    %c2 = index.constant 2 : index
    %c3 = index.constant 3 : index
    %c4 = index.constant 4 : index
    %padded = index.add %context, %c3 : index
    %blocks = index.div %padded, %c4 : index
    %c0_i32 = scalar.constant 0 : i32
    %fresh = index.cmp eq, %position, %c0 : index
    %initial_h, %initial_k = scf.if %fresh -> (i32, i32) {
      scf.yield %c0_i32, %c0_i32 : i32, i32
    } else {
      %h = view.load %read_state[0] : view<1xi32> -> i32
      %previous = index.sub %position, %c1 : index
      %offset = func.call @history_offset(%origins, %row, %previous, %blocks) : (buffer, index, index, index) -> (offset)
      %last = buffer.view %arena[%offset] : buffer -> view<1xi32>
      %k = view.load %last[0] : view<1xi32> -> i32
      scf.yield %h, %k : i32, i32
    }
    %checkpoint_bias = view.load %bias[0] : view<1xi32> -> i32
    %bias_delta = config.get @tiny.bias_delta : i32
    %bias_value = scalar.addi %checkpoint_bias, %bias_delta : i32
    %capacity = config.get @tiny.tokens : index
    %c5_i32 = scalar.constant 5 : i32
    %c7_i32 = scalar.constant 7 : i32
    %c1_i32 = scalar.constant 1 : i32
    %c97 = scalar.constant 97 : i32
    %c101 = scalar.constant 101 : i32
    %final_h, %final_k = scf.for %i = [%c0 to %length step %c1](%h = %initial_h : i32, %k = %initial_k : i32) -> (i32, i32) {
      %input_position = index.add %input, %i : index
      %bounded_input = index.assume %input_position [lt(%input_position, %capacity)] : index
      %ids = buffer.view %tokens[%origin] : buffer -> view<[%capacity]xi32>
      %token = view.load %ids[%bounded_input] : view<[%capacity]xi32> -> i32
      %scaled_h = scalar.muli %h, %c5_i32 : i32
      %with_token = scalar.addi %scaled_h, %token : i32
      %with_bias = scalar.addi %with_token, %bias_value : i32
      %next_h = scalar.remui %with_bias, %c97 : i32
      %absolute = index.add %position, %i : index
      %source_delta = index.mul %i, %c2 : index
      %absolute_source = index.add %source, %source_delta : index
      %absolute_source32 = index.cast %absolute_source : index to i32
      %absolute_position32 = scalar.shrui %absolute_source32, %c1_i32 : i32
      %scaled_k = scalar.muli %k, %c7_i32 : i32
      %with_h = scalar.addi %scaled_k, %next_h : i32
      %with_position = scalar.addi %with_h, %absolute_position32 : i32
      %next_k = scalar.remui %with_position, %c101 : i32
      %offset = func.call @history_offset(%origins, %row, %absolute, %blocks) : (buffer, index, index, index) -> (offset)
      %cache = buffer.view %arena[%offset] : buffer -> view<1xi32>
      view.store %next_k, %cache[0] : i32, view<1xi32>
      scf.yield %next_h, %next_k : i32, i32
    }
    view.store %final_h, %write_state[0] : i32, view<1xi32>
    %hidden = buffer.view %residual[%origin] : buffer -> view<[%spans]xi32>
    view.store %final_h, %hidden[%span] : i32, view<[%spans]xi32>
  }
  kernel.return
}

kernel.def @select() {
  %spans = config.get @tiny.spans : index
  %c1 = index.constant 1 : index
  kernel.launch.config workgroups(%spans, %c1, %c1) workgroup_size(%c1, %c1, %c1) : index
} launch(%residual: buffer, %metadata: buffer, %origins: buffer, %arena: buffer, %outputs: buffer) {
  %spans = config.get @tiny.spans : index
  %span = kernel.workgroup.id<x> : index
  %origin = index.constant 0 : offset
  %c16 = index.constant 16 : offset
  %header = buffer.view %metadata[%origin] : buffer -> view<4xi32>
  %count32 = view.load %header[0] : view<4xi32> -> i32
  %count = index.cast %count32 : i32 to index
  %live = index.cmp ult, %span, %count : index
  scf.if %live {
    %packets = buffer.view %metadata[%c16] : buffer -> view<[%spans]x9xi32>
    %flags = view.load %packets[%span, 6] : view<[%spans]x9xi32> -> i32
    %c0_i32 = scalar.constant 0 : i32
    %c1_i32 = scalar.constant 1 : i32
    %bits = scalar.andi %flags, %c1_i32 : i32
    %selected = scalar.cmpi ne, %bits, %c0_i32 : i32
    scf.if %selected {
      %length32 = view.load %packets[%span, 0] : view<[%spans]x9xi32> -> i32
      %position32 = view.load %packets[%span, 1] : view<[%spans]x9xi32> -> i32
      %source32 = view.load %packets[%span, 2] : view<[%spans]x9xi32> -> i32
      %row32 = view.load %packets[%span, 3] : view<[%spans]x9xi32> -> i32
      %end32 = scalar.addi %position32, %length32 : i32
      %c2_i32 = scalar.constant 2 : i32
      %source_length32 = scalar.muli %length32, %c2_i32 : i32
      %source_end32 = scalar.addi %source32, %source_length32 : i32
      %end_bound = scalar.assume %end32 [range(%end32, 1, 262144)] : i32
      %source_end_bound = scalar.assume %source_end32 [range(%source_end32, 1, 524288)] : i32
      %row_bound = scalar.assume %row32 [range(%row32, 0, 2147483646)] : i32
      %row_value = index.cast %row_bound : i32 to index
      %end = index.cast %end_bound : i32 to index
      %row = index.assume %row_value [range(%row_value, 0, 2147483646)] : index
      %c0 = index.constant 0 : index
      %c1 = index.constant 1 : index
      %c3 = index.constant 3 : index
      %c4 = index.constant 4 : index
      %context = config.get @tiny.context : index
      %padded = index.add %context, %c3 : index
      %blocks = index.div %padded, %c4 : index
      %sum = scf.for %position = [%c0 to %end step %c1](%previous = %c0_i32 : i32) -> (i32) {
        %offset = func.call @history_offset(%origins, %row, %position, %blocks) : (buffer, index, index, index) -> (offset)
        %cache = buffer.view %arena[%offset] : buffer -> view<1xi32>
        %value = view.load %cache[0] : view<1xi32> -> i32
        %next = scalar.addi %previous, %value : i32
        scf.yield %next : i32
      }
      %hidden = buffer.view %residual[%origin] : buffer -> view<[%spans]xi32>
      %h = view.load %hidden[%span] : view<[%spans]xi32> -> i32
      %total = scalar.addi %h, %sum : i32
      %c3_i32 = scalar.constant 3 : i32
      %c11_i32 = scalar.constant 11 : i32
      %token = scalar.remui %total, %c3_i32 : i32
      %source_phase = scalar.shrui %source_end_bound, %c1_i32 : i32
      %phase = scalar.remui %source_phase, %c11_i32 : i32
      %terminal = scalar.cmpi eq, %phase, %c0_i32 : i32
      %eos = view.load %header[1] : view<4xi32> -> i32
      %prediction = scf.select %terminal, %eos, %token : i32
      %result = buffer.view %outputs[%origin] : buffer -> view<[%spans]xi32>
      view.store %prediction, %result[%span] : i32, view<[%spans]xi32>
    }
  }
  kernel.return
}

command.program.def public @advance() launch(%weights: buffer, %residual: buffer, %metadata: buffer, %origins: buffer, %arena: buffer, %tokens: buffer, %outputs: buffer) {
  %bias = command.parameter %weights, "bias" : view<1xi32>
  kernel.launch @consume(%bias, %residual, %metadata, %origins, %arena, %tokens) : (view<1xi32>, buffer, buffer, buffer, buffer, buffer)
  kernel.launch @select(%residual, %metadata, %origins, %arena, %outputs) : (buffer, buffer, buffer, buffer, buffer)
  command.return
}