436 lines
14 KiB
C++
436 lines
14 KiB
C++
/*
|
|
* Copyright (c) Meta Platforms, Inc. and affiliates.
|
|
* All rights reserved.
|
|
*
|
|
* This source code is licensed under the BSD-style license found in the
|
|
* LICENSE file in the root directory of this source tree.
|
|
*/
|
|
|
|
#if defined(__x86_64__) || defined(__i386__) || \
|
|
(defined(_MSC_VER) && (defined(_M_X64) || defined(_M_IX86)))
|
|
#include <immintrin.h>
|
|
#endif
|
|
#include <algorithm>
|
|
#include <cassert>
|
|
#include <chrono>
|
|
#include <cmath>
|
|
#include <cstdint>
|
|
#include <iomanip>
|
|
#include <iostream>
|
|
#include <map>
|
|
#include <numeric>
|
|
#include <random>
|
|
#include <set>
|
|
#include <vector>
|
|
|
|
#include "./BenchUtils.h"
|
|
#include "fbgemm/Fbgemm.h"
|
|
#include "fbgemm/FbgemmConvert.h"
|
|
#include "src/RefImplementations.h"
|
|
|
|
using namespace std;
|
|
using namespace fbgemm;
|
|
|
|
static vector<vector<int>> GetInputs_() {
|
|
vector<vector<int>> input_dims = {
|
|
// batch size, number of rows of table, emb dim , avg lengthl
|
|
// TODO: Add more inputs
|
|
// Use these -- but they are slow.
|
|
{10, 4000000, 32, 100},
|
|
{10, 4000000, 64, 100},
|
|
{10, 4000000, 128, 100},
|
|
{10, 4000000, 256, 100},
|
|
// Use these for debugging
|
|
// {2, 16, 128, 10},
|
|
// {10, 4000, 128, 100},
|
|
// {10, 4000, 128, 100},
|
|
// {10, 4000, 128, 100},
|
|
};
|
|
return input_dims;
|
|
}
|
|
|
|
void run_benchmark(
|
|
int batch_size,
|
|
int num_rows,
|
|
int embedding_dim,
|
|
int average_len,
|
|
bool normalize_by_lengths,
|
|
bool use_fp16_inputs = false,
|
|
bool use_bf16_inputs = false,
|
|
bool use_32_bit_indices = false,
|
|
bool prefetch = false) {
|
|
// Create embedding table
|
|
default_random_engine generator;
|
|
|
|
vector<float> embedding_table(num_rows * embedding_dim);
|
|
normal_distribution<float> embedding_distribution;
|
|
for (size_t i = 0; i < embedding_table.size(); ++i) {
|
|
embedding_table[i] = embedding_distribution(generator);
|
|
}
|
|
vector<float16> embedding_table_fp16;
|
|
vector<bfloat16> embedding_table_bf16;
|
|
if (use_fp16_inputs) {
|
|
embedding_table_fp16.resize(embedding_table.size());
|
|
FloatToFloat16_simd(
|
|
embedding_table.data(),
|
|
embedding_table_fp16.data(),
|
|
embedding_table.size());
|
|
}
|
|
if (use_bf16_inputs) {
|
|
embedding_table_bf16.resize(embedding_table.size());
|
|
FloatToBfloat16_simd(
|
|
embedding_table.data(),
|
|
embedding_table_bf16.data(),
|
|
embedding_table.size());
|
|
}
|
|
|
|
// Generate lengths
|
|
uniform_int_distribution<int> length_distribution(
|
|
1, std::min(2 * average_len + 1, num_rows));
|
|
vector<int> offsets(batch_size + 1);
|
|
offsets[0] = 0;
|
|
for (int i = 0; i < batch_size; ++i) {
|
|
offsets[i + 1] = offsets[i] + length_distribution(generator);
|
|
}
|
|
|
|
// Compute the number of indices
|
|
int lengths_sum = offsets[batch_size];
|
|
cout << "lengths_sum " << lengths_sum << endl;
|
|
|
|
// Generate indices
|
|
vector<int64_t> indices;
|
|
vector<int32_t> indices_32;
|
|
|
|
vector<int> container(num_rows);
|
|
|
|
// please note we generate unique indices
|
|
for (int i = 0; i < batch_size; ++i) {
|
|
iota(container.begin(), container.end(), 0);
|
|
shuffle(container.begin(), container.end(), generator);
|
|
copy(
|
|
container.begin(),
|
|
container.begin() + (offsets[i + 1] - offsets[i]),
|
|
back_inserter(indices));
|
|
}
|
|
copy(begin(indices), end(indices), back_inserter(indices_32));
|
|
|
|
// Generate weights
|
|
vector<float> weights(lengths_sum);
|
|
for (int i = 0; i < lengths_sum; ++i) {
|
|
weights[i] = embedding_distribution(generator);
|
|
}
|
|
|
|
vector<float> output_sls_ref(batch_size * embedding_dim);
|
|
vector<float> output_slws_ref(output_sls_ref.size()),
|
|
output_sls(output_sls_ref.size()), output_slws(output_sls_ref.size());
|
|
|
|
constexpr int NUM_WARMUP = 4;
|
|
constexpr int NUM_ITER = 10;
|
|
int elem_bytes = use_fp16_inputs ? sizeof(float16) : sizeof(float);
|
|
double bytes = lengths_sum *
|
|
(embedding_dim * elem_bytes + (use_32_bit_indices ? 4 : 8)) +
|
|
batch_size * sizeof(int);
|
|
double bytes_padded = lengths_sum *
|
|
((embedding_dim * elem_bytes + 63) / 64 * 64 +
|
|
(use_32_bit_indices ? 4 : 8)) +
|
|
batch_size * sizeof(int);
|
|
|
|
for (bool has_weight : {false, true}) {
|
|
vector<float>& output_ref = has_weight ? output_slws_ref : output_sls_ref;
|
|
|
|
bool success = false, success_ref = false;
|
|
|
|
if (use_fp16_inputs) {
|
|
if (use_32_bit_indices) {
|
|
success_ref = EmbeddingSpMDM_ref(
|
|
embedding_dim,
|
|
batch_size,
|
|
lengths_sum,
|
|
num_rows,
|
|
embedding_table_fp16.data(),
|
|
indices_32.data(),
|
|
offsets.data(),
|
|
has_weight ? weights.data() : nullptr,
|
|
normalize_by_lengths,
|
|
output_ref.data());
|
|
} else {
|
|
success_ref = EmbeddingSpMDM_ref(
|
|
embedding_dim,
|
|
batch_size,
|
|
lengths_sum,
|
|
num_rows,
|
|
embedding_table_fp16.data(),
|
|
indices.data(),
|
|
offsets.data(),
|
|
has_weight ? weights.data() : nullptr,
|
|
normalize_by_lengths,
|
|
output_ref.data());
|
|
}
|
|
} else if (use_bf16_inputs) {
|
|
if (use_32_bit_indices) {
|
|
success_ref = EmbeddingSpMDM_ref(
|
|
embedding_dim,
|
|
batch_size,
|
|
lengths_sum,
|
|
num_rows,
|
|
embedding_table_bf16.data(),
|
|
indices_32.data(),
|
|
offsets.data(),
|
|
has_weight ? weights.data() : nullptr,
|
|
normalize_by_lengths,
|
|
output_ref.data());
|
|
} else {
|
|
success_ref = EmbeddingSpMDM_ref(
|
|
embedding_dim,
|
|
batch_size,
|
|
lengths_sum,
|
|
num_rows,
|
|
embedding_table_bf16.data(),
|
|
indices.data(),
|
|
offsets.data(),
|
|
has_weight ? weights.data() : nullptr,
|
|
normalize_by_lengths,
|
|
output_ref.data());
|
|
}
|
|
} else {
|
|
if (use_32_bit_indices) {
|
|
success_ref = EmbeddingSpMDM_ref(
|
|
embedding_dim,
|
|
batch_size,
|
|
lengths_sum,
|
|
num_rows,
|
|
embedding_table.data(),
|
|
indices_32.data(),
|
|
offsets.data(),
|
|
has_weight ? weights.data() : nullptr,
|
|
normalize_by_lengths,
|
|
output_ref.data());
|
|
} else {
|
|
success_ref = EmbeddingSpMDM_ref(
|
|
embedding_dim,
|
|
batch_size,
|
|
lengths_sum,
|
|
num_rows,
|
|
embedding_table.data(),
|
|
indices.data(),
|
|
offsets.data(),
|
|
has_weight ? weights.data() : nullptr,
|
|
normalize_by_lengths,
|
|
output_ref.data());
|
|
}
|
|
}
|
|
|
|
auto kernel_fp32_i32 = GenerateEmbeddingSpMDM<float, int32_t>(
|
|
embedding_dim, has_weight, normalize_by_lengths, prefetch ? 16 : 0);
|
|
auto kernel_fp32_i64 = GenerateEmbeddingSpMDM<float, int64_t>(
|
|
embedding_dim, has_weight, normalize_by_lengths, prefetch ? 16 : 0);
|
|
auto kernel_fp16_i32 = GenerateEmbeddingSpMDM<float16, int32_t>(
|
|
embedding_dim, has_weight, normalize_by_lengths, prefetch ? 16 : 0);
|
|
auto kernel_fp16_i64 = GenerateEmbeddingSpMDM<float16, int64_t>(
|
|
embedding_dim, has_weight, normalize_by_lengths, prefetch ? 16 : 0);
|
|
auto kernel_bf16_i32 = GenerateEmbeddingSpMDM<bfloat16, int32_t>(
|
|
embedding_dim,
|
|
has_weight,
|
|
normalize_by_lengths,
|
|
prefetch ? 16 : 0,
|
|
/*is_weight_positional=*/false,
|
|
/*use_offsets=*/true,
|
|
/*isbf16=*/true);
|
|
auto kernel_bf16_i64 = GenerateEmbeddingSpMDM<bfloat16, int64_t>(
|
|
embedding_dim,
|
|
has_weight,
|
|
normalize_by_lengths,
|
|
prefetch ? 16 : 0,
|
|
/*is_weight_positional=*/false,
|
|
/*is_weight_positional=*/true,
|
|
/*isbf16=*/true);
|
|
|
|
vector<float>& output = has_weight ? output_slws : output_sls;
|
|
for (bool flush_cache : {false, true}) {
|
|
double t = measureWithWarmup(
|
|
[&]() {
|
|
if (use_fp16_inputs) {
|
|
if (use_32_bit_indices) {
|
|
success = kernel_fp16_i32(
|
|
batch_size,
|
|
lengths_sum,
|
|
num_rows,
|
|
embedding_table_fp16.data(),
|
|
indices_32.data(),
|
|
offsets.data(),
|
|
has_weight ? weights.data() : nullptr,
|
|
output.data());
|
|
} else {
|
|
success = kernel_fp16_i64(
|
|
batch_size,
|
|
lengths_sum,
|
|
num_rows,
|
|
embedding_table_fp16.data(),
|
|
indices.data(),
|
|
offsets.data(),
|
|
has_weight ? weights.data() : nullptr,
|
|
output.data());
|
|
}
|
|
} else if (use_bf16_inputs) {
|
|
if (use_32_bit_indices) {
|
|
success = kernel_bf16_i32(
|
|
batch_size,
|
|
lengths_sum,
|
|
num_rows,
|
|
embedding_table_bf16.data(),
|
|
indices_32.data(),
|
|
offsets.data(),
|
|
has_weight ? weights.data() : nullptr,
|
|
output.data());
|
|
} else {
|
|
success = kernel_bf16_i64(
|
|
batch_size,
|
|
lengths_sum,
|
|
num_rows,
|
|
embedding_table_bf16.data(),
|
|
indices.data(),
|
|
offsets.data(),
|
|
has_weight ? weights.data() : nullptr,
|
|
output.data());
|
|
}
|
|
} else {
|
|
if (use_32_bit_indices) {
|
|
success = kernel_fp32_i32(
|
|
batch_size,
|
|
lengths_sum,
|
|
num_rows,
|
|
embedding_table.data(),
|
|
indices_32.data(),
|
|
offsets.data(),
|
|
has_weight ? weights.data() : nullptr,
|
|
output.data());
|
|
} else {
|
|
success = kernel_fp32_i64(
|
|
batch_size,
|
|
lengths_sum,
|
|
num_rows,
|
|
embedding_table.data(),
|
|
indices.data(),
|
|
offsets.data(),
|
|
has_weight ? weights.data() : nullptr,
|
|
output.data());
|
|
}
|
|
}
|
|
},
|
|
NUM_WARMUP,
|
|
NUM_ITER,
|
|
[&]() {
|
|
if (flush_cache) {
|
|
cache_evict(embedding_table);
|
|
cache_evict(indices);
|
|
cache_evict(indices_32);
|
|
cache_evict(offsets);
|
|
cache_evict(weights);
|
|
cache_evict(output);
|
|
}
|
|
});
|
|
|
|
// printMatrix(
|
|
// matrix_op_t::NoTranspose,
|
|
// output.data(),
|
|
// batch_size,
|
|
// embedding_dim,
|
|
// embedding_dim,
|
|
// "");
|
|
// cout << "reference data\n";
|
|
// printMatrix(
|
|
// matrix_op_t::NoTranspose,
|
|
// output_ref.data(),
|
|
// batch_size,
|
|
// embedding_dim,
|
|
// embedding_dim,
|
|
// "");
|
|
// Check correctness
|
|
if (!flush_cache) {
|
|
if (success != success_ref) {
|
|
assert(
|
|
false && "ERROR: refernce impl and JIT imp did not both succeed");
|
|
} else if (success) {
|
|
for (size_t i = 0; i < output.size(); ++i) {
|
|
assert(output[i] == output_ref[i]);
|
|
if (output[i] != output_ref[i]) {
|
|
cout << i << " " << output[i] << " " << output_ref[i] << endl;
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
if (has_weight) {
|
|
cout << setw(16) << "SLW(WEIGHTED) ";
|
|
} else {
|
|
cout << setw(16) << "SLS ";
|
|
}
|
|
if (flush_cache) {
|
|
cout << setw(20) << "cache flushed";
|
|
} else {
|
|
cout << setw(20) << "cache not flushed";
|
|
}
|
|
if (prefetch) {
|
|
cout << setw(16) << "prefetch on";
|
|
} else {
|
|
cout << setw(16) << "prefetch off";
|
|
}
|
|
|
|
cout << setw(8) << "b/w" << setw(10) << bytes / 1e9 / t << " GB/s"
|
|
<< setw(20) << "effective b/w: " << setw(16)
|
|
<< bytes_padded / 1e9 / t << "GB/s" << setw(8) << " time "
|
|
<< setw(16) << t << endl;
|
|
} // flush_cache
|
|
} // has_weight
|
|
}
|
|
|
|
int main() {
|
|
vector<vector<int>> inputs(GetInputs_());
|
|
|
|
for (auto& input : inputs) {
|
|
assert(input.size() > 3);
|
|
int batch_size = input[0];
|
|
int num_rows = input[1];
|
|
int embedding_dim = input[2];
|
|
int average_len = input[3];
|
|
|
|
cout << "batch size" << setw(6) << batch_size << setw(10) << "num rows"
|
|
<< setw(16) << num_rows << setw(10) << "emb dim" << setw(6)
|
|
<< embedding_dim << setw(16) << "avg length" << setw(6) << average_len
|
|
<< endl;
|
|
|
|
for (bool normalize_by_lengths : {false, true}) {
|
|
for (std::string input_dtype : {"bf16", "fp16", "fp32"}) {
|
|
for (bool use_32_bit_indices : {false, true}) {
|
|
for (bool prefetch : {false, true}) {
|
|
// args: batch sz, num rows, emb dim, avg len, normalize, use 32b,
|
|
// prefetch
|
|
if (normalize_by_lengths) {
|
|
cout << "Mean ";
|
|
}
|
|
cout << input_dtype << " inputs";
|
|
bool use_fp16_inputs = input_dtype == "fp16" ? true : false;
|
|
bool use_bf16_inputs = input_dtype == "fp16" ? true : false;
|
|
cout << (use_32_bit_indices ? " 32" : " 64") << " bit indices";
|
|
if (prefetch) {
|
|
cout << " with prefetching";
|
|
}
|
|
cout << ", ";
|
|
run_benchmark(
|
|
batch_size,
|
|
num_rows,
|
|
embedding_dim,
|
|
average_len,
|
|
normalize_by_lengths,
|
|
use_fp16_inputs,
|
|
use_bf16_inputs,
|
|
use_32_bit_indices,
|
|
prefetch);
|
|
} // prefetch
|
|
} // use_32_bit_indices
|
|
} // use_fp16_inputs
|
|
} // normalize_by_length
|
|
} // for each input
|
|
return 0;
|
|
}
|