gaspersic/triton
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

more optimization

Browse Source
This commit is contained in:
Philippe Tillet
2019-06-28 20:22:52 -07:00
parent 83b753512c
commit d8c3d58593
5 changed files with 36 additions and 31 deletions
Download Patch File Download Diff File Expand all files Collapse all files

42
examples/cpp/shift.cpp
View File

@@ -6,17 +6,21 @@
#include "triton/driver/stream.h"
#include "triton/tools/bench.hpp"
#include "triton/dnn/shift.h"
#include "triton/external/half.hpp"
int main() {
typedef half_float::half NumericT;
std::string numeric_t_str = "fp16";
// initialize default compute device
auto context = triton::driver::backend::contexts::get_default();
// initialize just-in-time compiler
triton::jit jit(context);
// initialization
int32_t R = 3, S = 3;
int32_t BS = 4, F = 512;
int32_t BS = 32, F = 1024;
int32_t H = 32, W = 32;
int32_t C = 512;
int32_t C = 1024;
// random shifts
std::vector<int32_t> shift_h(C);
std::vector<int32_t> shift_w(C);
@@ -25,23 +29,23 @@ int main() {
shift_w[c] = rand() % S - S/2;
}
// configuration
triton::dnn::shift shift(BS, C, 1, H, W, 1, R, S, F, shift_h, shift_w);
triton::dnn::shift shift(BS, C, 1, H, W, 1, R, S, F, shift_h, shift_w, numeric_t_str, numeric_t_str);
// host buffers
std::vector<float> hc(shift.c_size());
std::vector<float> rc(shift.c_size());
std::vector<float> ha(shift.a_size());
std::vector<float> hb(shift.b_size());
std::vector<NumericT> ha(shift.a_size());
std::vector<NumericT> hb(shift.b_size());
// device buffers
triton::driver::buffer* dc = triton::driver::buffer::create(context, hc.size()*4);
triton::driver::buffer* da = triton::driver::buffer::create(context, ha.size()*4);
triton::driver::buffer* db = triton::driver::buffer::create(context, hb.size()*4);
triton::driver::buffer* da = triton::driver::buffer::create(context, ha.size()*sizeof(NumericT));
triton::driver::buffer* db = triton::driver::buffer::create(context, hb.size()*sizeof(NumericT));
triton::driver::stream* stream = triton::driver::stream::create(context);
// initialize host
srand(0);
for(size_t i = 0; i < ha.size(); i++)
ha[i] = (float)rand() / RAND_MAX;
ha[i] = (NumericT)rand() / RAND_MAX;
for(size_t i = 0; i < hb.size(); i++)
hb[i] = (float)rand() / RAND_MAX;
hb[i] = (NumericT)rand() / RAND_MAX;
for(size_t i = 0; i < hc.size(); i++)
hc[i] = 0;
// initialize device
@@ -68,23 +72,23 @@ int main() {
// shift
std::vector<unsigned> params = {
32, 2, 128, 16, 2, 128, 16, 8, 2, 2, 4, 2, 8, 8
16, 4, 64, 16, 4, 128, 2, 2, 1, 2, 4, 4, 16, 4
};
std::ostringstream oss;
shift.src(oss);
std::string src = oss.str();
jit.autotune("shift", src.c_str(), benchmark);
// jit.autotune("shift", src.c_str(), benchmark);
jit.add_module("shift", src.c_str(), params);
triton::driver::kernel* kernel = jit.get_function("shift");
triton::jit::launch_information info = jit.get_launch_info("shift");
std::cout << "Performance: " << benchmark(kernel, info) << " TFLOPS " << std::endl;
stream->read(dc, true, 0, hc);
shift.cpu_ref(rc.data(), ha.data(), hb.data());
for(size_t i = 0; i < hc.size(); i++)
if(std::isnan(hc[i]) || std::abs(hc[i] - rc[i])/std::max(hc[i], rc[i]) > 1e-4){
std::cout << i << " " << hc[i] << " " << rc[i] << std::endl;
exit(EXIT_FAILURE);
}
std::cout << "Pass!" << std::endl;
// stream->read(dc, true, 0, hc);
// shift.cpu_ref(rc.data(), ha.data(), hb.data());
// for(size_t i = 0; i < hc.size(); i++)
// if(std::isnan(hc[i]) || std::abs(hc[i] - rc[i])/std::max(hc[i], rc[i]) > 1e-4){
// std::cout << i << " " << hc[i] << " " << rc[i] << std::endl;
// exit(EXIT_FAILURE);
// }
// std::cout << "Pass!" << std::endl;
}

1
include/triton/dnn/shift.h
View File

@@ -107,6 +107,7 @@ public:
private:
int32_t MAX_C_;
int32_t TK_;
// image size
int32_t NB_;
int32_t NC_;

1
lib/codegen/selection.cpp
View File

@@ -1042,6 +1042,7 @@ void selection::lower_tile_instruction(ir::instruction *ins, llvm::IRBuilder<> &
unsigned max_contiguous = axis_info_->get_max_contiguous(ptr);
unsigned alignment = std::min(starting_multiple, max_contiguous);
unsigned vector_size = std::min<unsigned>(result->axis(0).contiguous, alignment);
// vector_size = result->axis(0).contiguous;
std::map<unsigned, Value*> packets;
distributed_tile *TP = (distributed_tile*)tmap_.at(ld->get_pointer_operand());
result->for_each([&](indices_t idx){

6
lib/codegen/tune.cpp
View File

@@ -235,13 +235,13 @@ void tune::run(ir::module &mod) {
continue;
if(dynamic_cast<ir::load_inst*>(i) && i->get_type()->is_tile_ty()){
ir::type *ty = mod.get_builder().get_int32_ty();
std::unique_ptr<ir::metaparameter> tmp(ir::metaparameter::create(ctx, ty, 2, 2));
std::unique_ptr<ir::metaparameter> tmp(ir::metaparameter::create(ctx, ty, 4, 4));
*params_.at(i).at("nts.d0") = *tmp;
}
if(dynamic_cast<ir::dot_inst*>(i) && i->get_type()->is_tile_ty()){
ir::type *ty = mod.get_builder().get_int32_ty();
std::unique_ptr<ir::metaparameter> tmp1(ir::metaparameter::create(ctx, ty, 2, 2));
std::unique_ptr<ir::metaparameter> tmp2(ir::metaparameter::create(ctx, ty, 2, 2));
std::unique_ptr<ir::metaparameter> tmp1(ir::metaparameter::create(ctx, ty, 4, 4));
std::unique_ptr<ir::metaparameter> tmp2(ir::metaparameter::create(ctx, ty, 4, 4));
*params_.at(i).at("nts.d0") = *tmp1;
*params_.at(i).at("nts.d1") = *tmp2;
}

17
lib/dnn/shift.cpp
View File

@@ -29,7 +29,8 @@ shift::shift(int B, int NC,
a_ty_(a_ty), b_ty_(b_ty),
ty_(ty), bias_(bias) {
// max number of channels
MAX_C_ = 1024;
TK_ = 16;
MAX_C_ = 8192 + TK_;
// equivalent matmul
M_ = NB_*AH_*AW_;
N_ = NF_;
@@ -52,14 +53,12 @@ void shift::build_deltas() {
auto offset = [&](unsigned c) {
return c*ld_a_[0] + shift_h_[c]*ld_a_[1] + shift_w_[c]*ld_a_[2];
};
// allocate look-up table
size_t TK = 8;
h_deltas_.resize(MAX_C_);
// populate look-up table
for(unsigned c = 0; c < TK; c++)
for(unsigned c = 0; c < TK_; c++)
h_deltas_[c] = offset(c);
for(unsigned c = 0; c < NC_; c++)
h_deltas_[TK + c] = offset(c + TK) - offset(c);
h_deltas_[TK_ + c] = offset(c + TK_) - offset(c);
}
size_t shift::a_size(){
@@ -116,14 +115,14 @@ void shift::src(std::ostream &os) {
R"(
const tunable int32 TM = {16, 32, 64, 128};
const tunable int32 TN = {16, 32, 64, 128};
const tunable int32 TK = {8};
const tunable int32 TK = {)" << TK_ << R"(};
__constant__ int32* delta = alloc_const int32[)" << MAX_C_ << R"(];
void shift(restrict read_only align(16) )" << a_ty_ << R"( *a,
restrict read_only align(16) )" << b_ty_ << R"( *b,
fp32 *c,
int32 M, int32 N, int32 K,
multiple_of(4) int32 M, multiple_of(4) int32 N, multiple_of(4) int32 K,
int32 lda,
int32 ABS, int32 AH, int32 AW, int32 AR, int32 AS) {
int32 rxa[TM] = get_global_range[TM](0);
@@ -131,8 +130,8 @@ void shift(restrict read_only align(16) )" << a_ty_ << R"( *a,
int32 rka[TK] = 0 ... TK;
int32 rkb[TK] = 0 ... TK;
fp32 C[TM, TN] = 0;
int32 pad_h = AR/2;
int32 pad_w = AS/2;
int32 pad_h = AR / 2;
int32 pad_w = AS / 2;
int32 rawhc[TM] = rxa / ABS;
int32 raw[TM] = rawhc % AW;
int32 rahc[TM] = rawhc / AW;