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