File size: 5,060 Bytes
dfb774a 999ff90 dfb774a 999ff90 dfb774a 999ff90 dfb774a 999ff90 dfb774a 999ff90 dfb774a 999ff90 dfb774a | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 | func.decl import("input") @require(%condition: i32, %diagnostic: buffer)
func.decl import("json") @members(%json: buffer) -> (buffer, i32, i64)
func.decl import("prepare") @stage(%root: buffer, %tag: i64) -> (i32)
func.decl import("prepare") @config_i64(%stage: i32, %key: buffer, %value: i64)
func.decl import("prepare") @parameter(%stage: i32, %binding: i32, %input: i32, %source: i32, %policy: buffer)
func.decl import("prepare") @source_scope(%scope: buffer) -> (i32)
func.decl @tiny_storage(%maximum: i64, %span_capacity: i64, %context: i64, %pool: i64, %rows: i64, %pins: i64) -> (buffer, buffer, buffer, buffer, buffer)
// One shared command handles every host token/span bucket. The cold program
// owns storage geometry; the host owns physical page IDs and recurrent slots.
func.def public @text_prepare_v1(%prefill: i64, %context: i64, %pool: i64, %rows: i64, %pins: i64, %draft_depth: i32, %shapes: buffer, %shape_count: i64, %weights: buffer, %options: buffer) -> (i64, i32, buffer, buffer, buffer, buffer, buffer, buffer, buffer) {
%c0_source = scalar.constant 0 : i32
%origin = index.constant 0 : offset
%c0 = index.constant 0 : index
%c1 = index.constant 1 : index
%c0_i32 = scalar.constant 0 : i32
%c1_i32 = scalar.constant 1 : i32
%c0_i64 = scalar.constant 0 : i64
%c1_i64 = scalar.constant 1 : i64
%c4_i64 = scalar.constant 4 : i64
%pooled = scalar.cmpi ne, %pool, %c0_i64 : i64
%packed = scalar.cmpi ne, %shape_count, %c0_i64 : i64
%no_draft = scalar.cmpi eq, %draft_depth, %c0_i32 : i32
%has_epochs = scalar.andi %pooled, %packed : i1
%supported_mode = scalar.andi %has_epochs, %no_draft : i1
%fields, %type, %field_count = func.call @members(%options) : (buffer) -> (buffer, i32, i64)
%object_type = scalar.constant 2 : i32
%object = scalar.cmpi eq, %type, %object_type : i32
%empty_options = scalar.cmpi eq, %field_count, %c0_i64 : i64
%valid_options = scalar.andi %object, %empty_options : i1
%supported = scalar.andi %supported_mode, %valid_options : i1
%supported32 = scalar.extui %supported : i1 to i32
%error = buffer.constant utf8("text tiny requires pooled epochs without MTP") : buffer
func.call @require(%supported32, %error) : (i32, buffer)
%root = buffer.constant utf8("advance") : buffer
%stage = func.call @stage(%root, %c0_i64) : (buffer, i64) -> (i32)
%context_key = buffer.constant utf8("tiny.context") : buffer
func.call @config_i64(%stage, %context_key, %context) : (i32, buffer, i64)
%empty = buffer.constant utf8("") : buffer
%policy = buffer.constant utf8("weights.loom") : buffer
%scope = buffer.constant utf8("model") : buffer
%checkpoint = func.call @source_scope(%scope) : (buffer) -> (i32)
func.call @parameter(%stage, %c0_i32, %c0_source, %checkpoint, %policy) : (i32, i32, i32, i32, buffer)
%state_bytes = index.constant 128 : offset
%state = buffer.alloca<private> align(8) %state_bytes : buffer
%state_view = buffer.view %state[%origin] : buffer -> view<1xi32>
view.store %stage, %state_view[0] : i32, view<1xi32>
%bounded_count = scalar.assume %shape_count [range(%shape_count, 1, 2147483647)] : i64
%count = index.cast %bounded_count : i64 to index
%shape_view = buffer.view %shapes[%origin] : buffer -> view<[%count]x2xi64>
%maximum, %span_capacity = scf.for %i = [%c0 to %count step %c1](%previous = %c0_i64 : i64, %previous_spans = %c0_i64 : i64) -> (i64, i64) {
%tokens = view.load %shape_view[%i, 0] : view<[%count]x2xi64> -> i64
%larger = scalar.cmpi ugt, %tokens, %previous : i64
%next = scf.select %larger, %tokens, %previous : i64
%spans = view.load %shape_view[%i, 1] : view<[%count]x2xi64> -> i64
%more_spans = scalar.cmpi ugt, %spans, %previous_spans : i64
%next_spans = scf.select %more_spans, %spans, %previous_spans : i64
scf.yield %next, %next_spans : i64, i64
}
%smaller = scalar.cmpi ult, %maximum, %prefill : i64
%effective = scf.select %smaller, %maximum, %prefill : i64
%tokens_key = buffer.constant utf8("tiny.tokens") : buffer
%spans_key = buffer.constant utf8("tiny.spans") : buffer
func.call @config_i64(%stage, %spans_key, %span_capacity) : (i32, buffer, i64)
func.call @config_i64(%stage, %tokens_key, %maximum) : (i32, buffer, i64)
%layout_origin = index.constant 88 : offset
%layout = buffer.view %state[%layout_origin] : buffer -> view<5xi64>
view.store %maximum, %layout[0] : i64, view<5xi64>
view.store %span_capacity, %layout[1] : i64, view<5xi64>
view.store %context, %layout[2] : i64, view<5xi64>
view.store %pool, %layout[3] : i64, view<5xi64>
view.store %pins, %layout[4] : i64, view<5xi64>
%allocations, %views, %origins, %draft, %geometry = func.call @tiny_storage(%maximum, %span_capacity, %context, %pool, %rows, %pins) : (i64, i64, i64, i64, i64, i64) -> (buffer, buffer, buffer, buffer, buffer)
%terminal = buffer.constant utf8("#") : buffer
func.return %effective, %c1_i32, %state, %terminal, %allocations, %views, %origins, %empty, %geometry : i64, i32, buffer, buffer, buffer, buffer, buffer, buffer, buffer
}
|