Hugging Face's logo

benvanik
/
hrx-loom-serve-example

Text Generation
loom
example
source-defined
Model card Files
xet
 Community
hrx-loom-serve-example
File size: 7,265 Bytes
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
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
func.def @allocation(%records: buffer, %slot: index, %bytes: i64, %clear: i64) {
  %origin = index.constant 0 : offset
  %bounded = index.assume %slot [range(%slot, 0, 10)] : index
  %view = buffer.view %records[%origin] : buffer -> view<11x3xi64>
  %alignment = scalar.constant 256 : i64
  view.store %bytes, %view[%bounded, 0] : i64, view<11x3xi64>
  view.store %alignment, %view[%bounded, 1] : i64, view<11x3xi64>
  view.store %clear, %view[%bounded, 2] : i64, view<11x3xi64>
  func.return
}

// KV is a stable prefix; private views and recurrent slots grow by appending.
func.def public @tiny_storage(%maximum: i64, %span_capacity: i64, %context: i64, %pool: i64, %rows: i64, %pins: i64) -> (buffer, buffer, buffer, buffer, buffer) {
  %origin = index.constant 0 : offset
  %c0 = index.constant 0 : index
  %c1 = index.constant 1 : index
  %c0_i32 = scalar.constant 0 : i32
  %c0_i64 = scalar.constant 0 : i64
  %c1_i64 = scalar.constant 1 : i64
  %c4_i64 = scalar.constant 4 : i64
  %bounded_rows = scalar.assume %rows [range(%rows, 1, 2147483647)] : i64
  %row_count = index.cast %bounded_rows : i64 to index
  %slot_stride = scalar.constant 256 : i64
  %token_bytes = scalar.muli %maximum, %c4_i64 : i64
  %alignment_bias = scalar.constant 255 : i64
  %biased_token_bytes = scalar.addi %token_bytes, %alignment_bias : i64
  %token_slots = scalar.divui %biased_token_bytes, %slot_stride : i64
  %padded_token_bytes = scalar.muli %token_slots, %slot_stride : i64
  %progress_offset = scalar.addi %slot_stride, %padded_token_bytes : i64
  %recurrent_offset = scalar.addi %progress_offset, %slot_stride : i64
  %row_stride = scalar.addi %recurrent_offset, %slot_stride : i64
  %pinned = scalar.cmpi ne, %pins, %c0_i64 : i64
  %twice_rows = scalar.addi %rows, %rows : i64
  %fork_slots = scalar.addi %twice_rows, %pins : i64
  %slots = scf.select %pinned, %fork_slots, %rows : i64
  %private_bytes = scalar.muli %slots, %row_stride : i64
  %pool_bytes = scalar.muli %pool, %c4_i64 : i64
  %biased_pool = scalar.addi %pool_bytes, %alignment_bias : i64
  %pool_slots = scalar.divui %biased_pool, %slot_stride : i64
  %private_origin = scalar.muli %pool_slots, %slot_stride : i64
  %recurrent_origin = scalar.addi %private_origin, %recurrent_offset : i64
  %arena_bytes = scalar.addi %private_origin, %private_bytes : i64
  %c3_i64 = scalar.constant 3 : i64
  %padded = scalar.addi %context, %c3_i64 : i64
  %blocks = scalar.divui %padded, %c4_i64 : i64
  %map_entries = scalar.muli %rows, %blocks : i64
  %map_bytes = scalar.muli %map_entries, %c4_i64 : i64
  %origin_stride = scalar.constant 40 : i64
  %map_origin = scalar.muli %bounded_rows, %origin_stride : i64
  %table_bytes = scalar.addi %map_origin, %map_bytes : i64
  %allocation_bytes = index.constant 264 : offset
  %allocations = buffer.alloca<private> align(8) %allocation_bytes : buffer
  %allocation_count = index.constant 11 : index
  scf.for %i = [%c0 to %allocation_count step %c1] {
    func.call @allocation(%allocations, %i, %c0_i64, %c0_i64) : (buffer, index, i64, i64)
  }
  %residual_bytes = scalar.muli %span_capacity, %c4_i64 : i64
  %header_bytes = scalar.constant 16 : i64
  %record_bytes = scalar.constant 36 : i64
  %plan_bytes = scalar.muli %span_capacity, %record_bytes : i64
  %metadata_bytes = scalar.addi %header_bytes, %plan_bytes : i64
  %c2 = index.constant 2 : index
  %c3 = index.constant 3 : index
  %c4 = index.constant 4 : index
  %c5 = index.constant 5 : index
  func.call @allocation(%allocations, %c0, %residual_bytes, %c0_i64) : (buffer, index, i64, i64)
  func.call @allocation(%allocations, %c1, %arena_bytes, %c0_i64) : (buffer, index, i64, i64)
  func.call @allocation(%allocations, %c2, %metadata_bytes, %c0_i64) : (buffer, index, i64, i64)
  func.call @allocation(%allocations, %c3, %table_bytes, %c0_i64) : (buffer, index, i64, i64)
  func.call @allocation(%allocations, %c4, %token_bytes, %c0_i64) : (buffer, index, i64, i64)
  func.call @allocation(%allocations, %c5, %residual_bytes, %c0_i64) : (buffer, index, i64, i64)

  %view_stride = index.constant 64 : offset
  %view_bytes = index.scale %row_count, %view_stride : index, offset -> offset
  %views = buffer.alloca<private> align(8) %view_bytes : buffer
  %view_records = buffer.view %views[%origin] : buffer -> view<[%row_count]x4x2xi64>
  %origin_bytes = index.cast %map_origin : i64 to offset
  %origins = buffer.alloca<private> align(8) %origin_bytes : buffer
  buffer.fill %c0_i32, %origins[%origin], %origin_bytes : i32
  %origin_records = buffer.view %origins[%origin] : buffer -> view<[%row_count]x5xi64>
  %control_bytes = scalar.constant 12 : i64
  %progress_bytes = scalar.constant 32 : i64
  scf.for %row = [%c0 to %row_count step %c1] {
    %row64 = index.cast %row : index to i64
    %relative = scalar.muli %row64, %row_stride : i64
    %base = scalar.addi %private_origin, %relative : i64
    %tokens = scalar.addi %base, %slot_stride : i64
    %progress = scalar.addi %base, %progress_offset : i64
    view.store %base, %view_records[%row, 0, 0] : i64, view<[%row_count]x4x2xi64>
    view.store %control_bytes, %view_records[%row, 0, 1] : i64, view<[%row_count]x4x2xi64>
    view.store %c0_i64, %view_records[%row, 1, 0] : i64, view<[%row_count]x4x2xi64>
    view.store %c0_i64, %view_records[%row, 1, 1] : i64, view<[%row_count]x4x2xi64>
    view.store %tokens, %view_records[%row, 2, 0] : i64, view<[%row_count]x4x2xi64>
    view.store %token_bytes, %view_records[%row, 2, 1] : i64, view<[%row_count]x4x2xi64>
    view.store %progress, %view_records[%row, 3, 0] : i64, view<[%row_count]x4x2xi64>
    view.store %progress_bytes, %view_records[%row, 3, 1] : i64, view<[%row_count]x4x2xi64>
    view.store %c0_i64, %origin_records[%row, 2] : i64, view<[%row_count]x5xi64>
    %map_stride = scalar.muli %blocks, %c4_i64 : i64
    %map_offset = scalar.muli %row64, %map_stride : i64
    %map_base = scalar.addi %map_origin, %map_offset : i64
    view.store %map_base, %origin_records[%row, 4] : i64, view<[%row_count]x5xi64>
  }
  %geometry_bytes = index.constant 104 : offset
  %geometry = buffer.alloca<private> align(8) %geometry_bytes : buffer
  %geometry_view = buffer.view %geometry[%origin] : buffer -> view<13xi64>
  view.store %c4_i64, %geometry_view[0] : i64, view<13xi64>
  view.store %map_origin, %geometry_view[1] : i64, view<13xi64>
  view.store %c0_i64, %geometry_view[2] : i64, view<13xi64>
  view.store %c4_i64, %geometry_view[3] : i64, view<13xi64>
  view.store %recurrent_origin, %geometry_view[4] : i64, view<13xi64>
  view.store %row_stride, %geometry_view[5] : i64, view<13xi64>
  view.store %c4_i64, %geometry_view[6] : i64, view<13xi64>
  view.store %slots, %geometry_view[7] : i64, view<13xi64>
  view.store %c1_i64, %geometry_view[8] : i64, view<13xi64>
  view.store %c0_i64, %geometry_view[9] : i64, view<13xi64>
  view.store %c1_i64, %geometry_view[10] : i64, view<13xi64>
  view.store %pool_bytes, %geometry_view[11] : i64, view<13xi64>
  %block_bytes = scalar.constant 16 : i64
  view.store %block_bytes, %geometry_view[12] : i64, view<13xi64>
  %empty_bytes = index.constant 0 : offset
  %empty = buffer.alloca<private> align(1) %empty_bytes : buffer
  func.return %allocations, %views, %origins, %empty, %geometry : buffer, buffer, buffer, buffer, buffer
}