| 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 |
| // 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 |
| } |
|
|
