hrxdb 0.3.1

A GPU-powered vector database for unified-memory systems, built on HRX and Loom
Documentation 
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
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162

// SIMT GEMM: a workgroup scores 64 corpus rows against 8–64 queries. Corpus
// components enter LDS once in FP16 and expand to FP32 in registers. Queries
// and accumulation remain FP32; components accumulate in increasing order.
amdgpu.target<gfx1151> @target {subgroup_size = 32}
config.decl @db.batch : %value: index where [range(%value, 8, 64), mul(%value, 8)]
config.decl @db.scan.dimensions : %value: index where [range(%value, 128, 16384), mul(%value, 128)]
kernel.def target(@target) export("batch_scan") @batch_scan(%count_arg: i32) {
  %one = index.constant 1 : index
  %four = index.constant 4 : index
  %rows = index.constant 64 : index
  %extra = index.constant 63 : index
  %batch = config.get @db.batch : index
  %threads = index.mul %batch, %four : index
  %raw = index.cast %count_arg : i32 to index
  %count = index.assume %raw [range(%raw, 1, 262144)] : index
  %rounded = index.add %count, %extra : index
  %groups = index.div %rounded, %rows : index
  kernel.launch.config workgroups(%groups, %one, %one) workgroup_size(%threads, %one, %one) : index
} launch(%count_arg: i32, %masked_arg: i32, %row_base_arg: i32, %total_arg: i32, %data: buffer, %queries: buffer, %norms: buffer, %bitmap: buffer, %scores: buffer) {
  %zero = index.constant 0 : offset
  %i0 = index.constant 0 : index
  %i1 = index.constant 1 : index
  %i4 = index.constant 4 : index
  %i8 = index.constant 8 : index
  %c32 = index.constant 32 : index
  %c64 = index.constant 64 : index
  %c68 = index.constant 68 : index
  %c512 = index.constant 512 : index
  %batch = config.get @db.batch : index
  %d = config.get @db.scan.dimensions : index
  %threads = index.mul %batch, %i4 : index
  %query_groups = index.div %batch, %i4 : index
  %raw = index.cast %count_arg : i32 to index
  %count = index.assume %raw [range(%raw, 1, 262144)] : index
  %tid0 = kernel.workitem.id<x> : index
  %tid = index.assume %tid0 [range(%tid0, 0, 255)] : index
  %group0 = kernel.workgroup.id<x> : index
  %group = index.assume %group0 [range(%group0, 0, 4095)] : index
  %row_base = index.mul %group, %c64 : index
  %row_group = index.div %tid, %query_groups : index
  %query_group = index.rem %tid, %query_groups : index
  %local_row_base = index.mul %row_group, %i4 : index
  %local_query_base = index.mul %query_group, %i4 : index
  %dv = buffer.view %data[%zero] : buffer -> view<[%count]x[%d]xf16>
  %qv = buffer.view %queries[%zero] : buffer -> view<[%d]x[%batch]xf32>
  %nv = buffer.view %norms[%zero] : buffer -> view<[%count]xf32>
  %sv = buffer.view %scores[%zero] : buffer -> view<[%batch]x[%count]xf32>
  %xbytes = index.constant 4352 : offset
  %query_words = index.mul %c32, %batch : index
  %four_offset = index.constant 4 : offset
  %qbytes = index.scale %query_words, %four_offset : index, offset -> offset
  %xscratch = buffer.alloca<workgroup> align(16) %xbytes : buffer
  %qscratch = buffer.alloca<workgroup> align(16) %qbytes : buffer
  // Padding breaks bank conflicts when transposing corpus components into LDS.
  %xs = buffer.view %xscratch[%zero] : buffer -> view<32x68xf16>
  %qs = buffer.view %qscratch[%zero] : buffer -> view<32x[%batch]xf32>
  %f0 = scalar.constant 0.0 : f32
  %initial = vector.splat %f0 : vector<16xf32>
  %query_vectors = index.mul %i8, %batch : index
  %accumulated = scf.for %base = [%i0 to %d step %c32](%acc = %initial : vector<16xf32>) -> (vector<16xf32>) {
    scf.for %v = [%tid to %c512 step %threads] {
      %r = index.div %v, %i8 : index
      %column_group = index.rem %v, %i8 : index
      %column = index.mul %column_group, %i4 : index
      %row = index.add %row_base, %r : index
      %valid = index.cmp ult, %row, %count : index
      %safe = scf.select %valid, %row, %i0 : index
      %at = index.assume %safe [range(%safe, 0, 262143), lt(%safe, %count)] : index
      %dim = index.add %base, %column : index
      %values = vector.load %dv[%at, %dim] : view<[%count]x[%d]xf16> -> vector<4xf16>
      scf.for %j = [%i0 to %i4 step %i1] unroll {
        %col = index.add %column, %j : index
        %value = vector.extract %values[%j] : vector<4xf16> -> f16
        view.store %value, %xs[%col, %r] : f16, view<32x68xf16>
      }
    }
    scf.for %v = [%tid to %query_vectors step %threads] {
      %col = index.div %v, %query_groups : index
      %qgroup = index.rem %v, %query_groups : index
      %q = index.mul %qgroup, %i4 : index
      %dim = index.add %base, %col : index
      %values = vector.load %qv[%dim, %q] : view<[%d]x[%batch]xf32> -> vector<4xf32>
      vector.store %values, %qs[%col, %q] : vector<4xf32>, view<32x[%batch]xf32>
    }
    kernel.barrier<workgroup> scope(workgroup) ordering(acq_rel)
    %next = scf.for %col = [%i0 to %c32 step %i1](%values = %acc : vector<16xf32>) -> (vector<16xf32>) unroll {
      %xh = vector.load %xs[%col, %local_row_base] : view<32x68xf16> -> vector<4xf16>
      %xvalues = vector.extf %xh : vector<4xf16> to vector<4xf32>
      %qvalues = vector.load %qs[%col, %local_query_base] : view<32x[%batch]xf32> -> vector<4xf32>
      %updated = scf.for %r = [%i0 to %i4 step %i1](%rs = %values : vector<16xf32>) -> (vector<16xf32>) unroll {
        %row = index.add %local_row_base, %r : index
        %x = vector.extract %xvalues[%r] : vector<4xf32> -> f32
        %rbase = index.mul %r, %i4 : index
        %new_row = scf.for %q = [%i0 to %i4 step %i1](%sums = %rs : vector<16xf32>) -> (vector<16xf32>) unroll {
          %slot = index.add %rbase, %q : index
          %old = vector.extract %sums[%slot] : vector<16xf32> -> f32
          %y = vector.extract %qvalues[%q] : vector<4xf32> -> f32
          %vx = vector.splat %x : vector<1xf32>
          %vy = vector.splat %y : vector<1xf32>
          %sum = vector.dotf %vx, %vy, %old : vector<1xf32>, vector<1xf32>, f32
          %inserted = vector.insert %sum into %sums[%slot] : f32, vector<16xf32>
          scf.yield %inserted : vector<16xf32>
        }
        scf.yield %new_row : vector<16xf32>
      }
      // Bound operand lifetimes to one column instead of hoisting loads from
      // all 32 unrolled columns. This is a compiler hint, not a GPU barrier.
      scf.schedule.fence
      scf.yield %updated : vector<16xf32>
    }
    kernel.barrier<workgroup> scope(workgroup) ordering(acq_rel)
    scf.yield %next : vector<16xf32>
  }
  %zero_i32 = scalar.constant 0 : i32
  %masked = scalar.cmpi ne, %masked_arg, %zero_i32 : i32
  %c31 = index.constant 31 : index
  %total0 = index.cast %total_arg : i32 to index
  %total = index.assume %total0 [range(%total0, 1, 1073741824)] : index
  %global_base0 = index.cast %row_base_arg : i32 to index
  %global_base = index.assume %global_base0 [range(%global_base0, 0, 1073741823)] : index
  %rounded = index.add %total, %c31 : index
  %words = index.div %rounded, %c32 : index
  %bv = buffer.view %bitmap[%zero] : buffer -> view<[%words]xi32>
  scf.for %r = [%i0 to %i4 step %i1] unroll {
    %local = index.add %local_row_base, %r : index
    %row = index.add %row_base, %local : index
    %valid = index.cmp ult, %row, %count : index
    scf.if %valid {
      %at = index.assume %row [range(%row, 0, 262143), lt(%row, %count)] : index
      %inverse = view.load %nv[%at] : view<[%count]xf32> -> f32
      %excluded = scf.if %masked -> (i1) {
        %global0 = index.add %global_base, %at : index
        %global = index.assume %global0 [range(%global0, 0, 1073741823), lt(%global0, %total)] : index
        %word0 = index.div %global, %c32 : index
        %word = index.assume %word0 [range(%word0, 0, 33554431), lt(%word0, %words)] : index
        %bit = index.rem %global, %c32 : index
        %bits = view.load %bv[%word] : view<[%words]xi32> -> i32
        %shift = index.cast %bit : index to i32
        %shifted = scalar.shrui %bits, %shift : i32
        %one = scalar.constant 1 : i32
        %flag = scalar.andi %shifted, %one : i32
        %test = scalar.cmpi ne, %flag, %zero_i32 : i32
        scf.yield %test : i1
      } else {
        %no = scalar.constant false : i1
        scf.yield %no : i1
      }
      %rbase = index.mul %r, %i4 : index
      scf.for %q = [%i0 to %i4 step %i1] unroll {
        %query = index.add %local_query_base, %q : index
        %slot = index.add %rbase, %q : index
        %sum = vector.extract %accumulated[%slot] : vector<16xf32> -> f32
        %cosine = scalar.mulf %sum, %inverse : f32
        %badbits = scalar.constant -8388608 : i32
        %neg = scalar.bitcast %badbits : i32 to f32
        %score = scf.select %excluded, %neg, %cosine : f32
        view.store %score, %sv[%query, %at] : f32, view<[%batch]x[%count]xf32>
      }
    }
  }
  kernel.return
}