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

[python][examples] some more cleaning of dot product example

Browse Source
This commit is contained in:
Philippe Tillet
2019-08-30 17:05:03 -07:00
parent 7e0af2118c
commit 5db3a7adfe
2 changed files with 14 additions and 20 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
lib/runtime/function.cc
View File

@@ -217,7 +217,7 @@ std::unique_ptr<driver::module> function::make_bin(ir::module &module, driver::c
dce.run(module);
vectorize.run(module);
dce.run(module);
// ir::print(module, std::cout);
ir::print(module, std::cout);
// generate llvm code
llvm::LLVMContext ctx;
std::unique_ptr<llvm::Module> llvm(new llvm::Module(module.get_name(), ctx));

32
python/examples/dot.py
View File

@@ -35,45 +35,39 @@ src = """
#define SHAPE_B TK, TN
#endif
void dot(TYPE * A,
TYPE * B,
TYPE * C,
void dot(TYPE * A, TYPE * B, TYPE * C,
int M, int N, int K,
int lda __multipleof(8),
int ldb __multipleof(8),
int ldc) {
/* prologue */
// prologue
int ridx = get_program_id(0);
int ridy = get_program_id(1);
int rxa[TM] = ridx * TM + 0 ... TM;
int ryb[TN] = ridy * TN + 0 ... TN;
int rka[TK] = 0 ... TK;
int rkb[TK] = 0 ... TK;
float xc[TM, TN] = 0;
/* pointers for operands */
float c[TM, TN] = 0;
// pointers to operands
TYPE* pa[SHAPE_A] = A + rka[BROADCAST_AK] * STRIDE_AK + rxa[BROADCAST_AM] * STRIDE_AM;
TYPE* pb[SHAPE_B] = B + rkb[BROADCAST_BK] * STRIDE_BK + ryb[BROADCAST_BN] * STRIDE_BN;
/* prefetches operands */
// prefetches operands
TYPE a[SHAPE_A] = *pa;
TYPE b[SHAPE_B] = *pb;
/* reduction loop */
for(int k = K; k > 0; k = k - TK){
xc = USEA @ USEB + xc;
// reduction loop
for(int k = K; k > 0; k-= TK){
c += USEA @ USEB;
pa = pa + TK * STRIDE_AK;
pb = pb + TK * STRIDE_BK;
a = *pa;
b = *pb;
}
/* epilogue */
int rxc[TM] = ridx * TM + (0 ... TM);
int ryc[TN] = ridy * TN + (0 ... TN);
// epilogue
int rxc[TM] = ridx * TM + 0 ... TM;
int ryc[TN] = ridy * TN + 0 ... TN;
TYPE* pc[TM, TN] = C + ryc[newaxis, :] + rxc[:, newaxis] * ldc;
TYPE c[TM, TN] = xc;
bool checkc0[TM] = rxc < M;
bool checkc1[TN] = ryc < N;
bool checkc[TM, TN] = checkc0[:, newaxis] && checkc1[newaxis, :];
*pc = c;
bool checkc[TM, TN] = (rxc < M)[:, newaxis] && (ryc < N)[newaxis, :];
*?(checkc) pc = c;
}
"""