Distributed top
Methods, systems, and apparatus, including computer programs encoded on computer storage media, for distributing a top k computation across multiple computing units of an integrated circuit One of the methods includes computing, by each of the plurality of computing units and for each candidate vector in a respective subset of the candidate vectors assigned to the computing unit, a respective distance between the query vector and the candidate vector; initializing, by the integrated circuit, a cut-off distance value; determining, by the integrated circuit, a final cut-off distance value; and providing, by the integrated circuit and as an output of a top k computation for the query vector and the set of candidate vectors, the candidate vectors that have respective distances that satisfy the final cut-off distance value.
1 . A method performed by an integrated circuit comprising a plurality of computing units that each includes a dedicated memory and a set of multiply accumulate units that allow distances according to distance measures to be computed through multiplication and accumulation operations performed in hardware, the method comprising:
receiving, by the integrated circuit, a query vector;
receiving, by the integrated circuit, a set of candidate vectors;
storing in the dedicated memory, for each of the plurality of computing units and for each candidate vector in a respective subset of the candidate vectors assigned to the computing unit, a respective distance between the query vector and the candidate vector from the respective set of multiply accumulate units;
setting, by the integrated circuit, a cut-off distance value to a predetermined value;
performing, by the integrated circuit, a plurality of iterations comprising, at each of the plurality of iterations:
generating one or more test cut-off distance values for the iteration from the cut-off distance value as of the iteration;
receiving, by each of the plurality of computing units and for each of the test cut-off distance values, a respective count of a number of the respective distances from the computing unit and stored in the designated memory of the computing unit that satisfy the test cut-off distance value, wherein the respective counts are stored in parallel across the plurality of computing units; and
updating, by the integrated circuit, the cut-off distance value using the respective counts from each of the plurality of computing units and for each of the test cut-off distance values; and
after performing the plurality of iterations, providing, by the integrated circuit and as an output of a top k computation for the query vector and the set of candidate vectors, the candidate vectors that have respective distances that satisfy the cut-off distance value.
2 . The method of claim 1 , wherein the setting the cut-off distance value, the generating the one or more test cut-off distance values, and the updating the cut-off distance value are performed by a processing core of the integrated circuit that is separate from the computing units.
3 . The method of claim 2 , wherein performing the plurality of iterations further comprises, at each iteration:
providing, from the processing core and to each of the computing units, data identifying the one or more test cut-off distance values for the iteration; and
providing, from each computing unit and to the processing core, the respective counts for each of the test cut-off distance values.
4 . The method of claim 1 , wherein providing, by the integrated circuit and as an output of a top k computation for the query vector and the set of candidate vectors, the candidate vectors that have respective distances that satisfy the cut-off distance value comprises:
providing, by each of the computing units, the candidate vectors from the respective subset assigned to the computing unit that have respective distances that satisfy the cut-off distance value.
5 . The method of claim 1 , wherein a respective sum of the respective counts from each of the plurality of computing units is associated with each test cut-off distance value.
6 . The method of claim 5 , wherein there is a single test cut-off distance value and wherein the updating further comprises:
in response to the respective sum being less than k, updating the cut-off distance value to be equal to the test cut-off distance value.
7 . The method of claim 5 , wherein there are a plurality of test cut-off distance values and wherein the updating further comprises:
updating the cut-off distance value to be equal to the test cut-off distance value that has a largest respective sum from among the plurality of test cut-off distance values that have respective sums that are less than k.
8 . The method of claim 1 , wherein the cut-off distance value is represented by a sequence of bits, wherein each of the plurality of iterations corresponds to a different bit in the sequence, and wherein generating one or more test cut-off distance values for the iteration from the cut-off distance value as of the iteration comprises generating a test cut-off distance value that has (i) a value of one for the bit corresponding to the iteration and (ii) the same values for the bits corresponding to any preceding iterations as the cut-off distance value as of the iteration.
9 . The method of claim 1 , wherein the cut-off distance value is represented by a sequence of bits, wherein each of the plurality of iterations corresponds to a different plurality of bits in the sequence, and wherein generating one or more test cut-off distance values for the iteration from the cut-off distance value as of the iteration comprises generating a plurality of test cut-off distance values, wherein each of the plurality of test cut-off distance values has (i) a different combination of values from the plurality of bits corresponding to the iteration and (ii) the same values for the bits corresponding to any preceding iterations as the cut-off distance value as of the iteration.
10 . The method of claim 1 , wherein performing, by the integrated circuit, the plurality of iterations, comprises, at a last iteration of the plurality of iterations: setting the cut-off distance value equal to the updated cut-off distance value after the iteration when for the updated cut-off distance value after the iteration, a sum of the respective counts from the plurality of computing units for the updated cut-off distance value after the iteration is equal to k.
11 . One or more non-transitory computer-readable storage media storing instructions that, when executed by an integrated circuit that comprises a plurality of computing units that each includes a dedicated memory and a set of multiply accumulate units that allow distances according to distance measures to be computed through multiplication and accumulation operations performed in hardware, cause the integrated circuit to perform operations comprising:
receiving, by the integrated circuit, a query vector;
receiving, by the integrated circuit, a set of candidate vectors;
storing in the dedicated memory, for each of the plurality of computing units and for each candidate vector in a respective subset of the candidate vectors assigned to the computing unit, a respective distance between the query vector and the candidate vector from the respective set of multiply accumulate units;
setting, by the integrated circuit, a cut-off distance value to a predetermined value;
performing, by the integrated circuit, a plurality of iterations comprising, at each of the plurality of iterations:
generating one or more test cut-off distance values for the iteration from the cut-off distance value as of the iteration;
receiving, by each of the plurality of computing units and for each of the test cut-off distance values, a respective count of a number of the respective distances from the computing unit and stored in the designated memory of the computing unit that satisfy the test cut-off distance value, wherein the respective counts are stored in parallel across the plurality of computing units; and
updating, by the integrated circuit, the cut-off distance value using the respective counts from each of the plurality of computing units and for each of the test cut-off distance values; and
after performing the plurality of iterations, providing, by the integrated circuit and as an output of a top k computation for the query vector and the set of candidate vectors, the candidate vectors that have respective distances that satisfy the cut-off distance value.
12 . An integrated circuit comprising a plurality of computing units that each includes a dedicated memory and a set of multiply accumulate units that allow distances according to distance measures to be computed through multiplication and accumulation operations performed in hardware, wherein the integrated circuit is configured to perform operations comprising:
receiving, by the integrated circuit, a query vector;
receiving, by the integrated circuit, a set of candidate vectors;
storing in the dedicated memory, for each of the plurality of computing units and for each candidate vector in a respective subset of the candidate vectors assigned to the computing unit, a respective distance between the query vector and the candidate vector from the respective set of multiply accumulate units;
setting, by the integrated circuit, a cut-off distance value to a predetermined value;
performing, by the integrated circuit, a plurality of iterations comprising, at each of the plurality of iterations:
generating one or more test cut-off distance values for the iteration from the cut-off distance value as of the iteration;
receiving, by each of the plurality of computing units and for each of the test cut-off distance values, a respective count of a number of the respective distances from the computing unit and stored in the designated memory of the computing unit that satisfy the test cut-off distance value, wherein the respective counts are stored in parallel across the plurality of computing units; and
updating, by the integrated circuit, the cut-off distance value using the respective counts from each of the plurality of computing units and for each of the test cut-off distance values; and
after performing the plurality of iterations, providing, by the integrated circuit and as an output of a top k computation for the query vector and the set of candidate vectors, the candidate vectors that have respective distances that satisfy the cut-off distance value.
13 . The integrated circuit of claim 12 , wherein the setting the cut-off distance value, the generating the one or more test cut-off distance values, and the updating the cut-off distance value are performed by a processing core of the integrated circuit that is separate from the computing units.
14 . The integrated circuit of claim 13 , wherein performing the plurality of iterations further comprises, at each iteration:
providing, from the processing core and to each of the computing units, data identifying the one or more test cut-off distance values for the iteration; and
providing, from each computing unit and to the processing core, the respective counts for each of the test cut-off distance values.
15 . The integrated circuit of claim 12 , wherein providing, by the integrated circuit and as an output of a top k computation for the query vector and the set of candidate vectors, the candidate vectors that have respective distances that satisfy the cut-off distance value comprises:
providing, by each of the computing units, the candidate vectors from the respective subset assigned to the computing unit that have respective distances that satisfy the cut-off distance value.
16 . The integrated circuit of claim 12 , wherein a respective sum of the respective counts from each of the plurality of computing units is associated with each test cut-off distance value.
17 . The integrated circuit of claim 16 , wherein there is a single test cut-off distance value and wherein the updating further comprises:
in response to the respective sum being less than k, updating the cut-off distance value to be equal to the test cut-off distance value.
18 . The integrated circuit of claim 16 , wherein there are a plurality of test cut-off distance values and wherein the updating further comprises:
updating the cut-off distance value to be equal to the test cut-off distance value that has a largest respective sum from among the plurality of test cut-off distance values that have respective sums that are less than k.
19 . The integrated circuit of claim 12 , wherein the cut-off distance value is represented by a sequence of bits, wherein each of the plurality of iterations corresponds to a different bit in the sequence, and wherein generating one or more test cut-off distance values for the iteration from the cut-off distance value as of the iteration comprises generating a test cut-off distance value that has (i) a value of one for the bit corresponding to the iteration and (ii) the same values for the bits corresponding to any preceding iterations as the cut-off distance value as of the iteration.
20 . The integrated circuit of claim 12 , wherein the cut-off distance value is represented by a sequence of bits, wherein each of the plurality of iterations corresponds to a different plurality of bits in the sequence, and wherein generating one or more test cut-off distance values for the iteration from the cut-off distance value as of the iteration comprises generating a plurality of test cut-off distance values, wherein each of the plurality of test cut-off distance values has (i) a different combination of values from the plurality of bits corresponding to the iteration and (ii) the same values for the bits corresponding to any preceding iterations as the cut-off distance value as of the iteration.