[examples/python/tensorflow] fixed ABI compatibility issue in JIT

destructor

examples/cpp/dot.cpp

@@ -53,8 +53,26 @@ void matmul(restrict read_only fp32 *A, restrict read_only fp32 *B, fp32 *C,
     fp32 b[TN, 1] = checkb ? *pb : 0;
     c = dot(a, trans(b), c);
   }
+  int32 ridx = get_range_id(0);
+  int32 ridy = get_range_id(1);
   fp32* pc[TM, TN] = C + ryc[newaxis, :]*ldc + rxc[:, newaxis];
-  *pc = c;
+  int32 *plock = locks + ridx + ridy*grid0;
+  while(__atomic_cas(plock, 0, 1));
+  int32 *pcount = plock + grid0*grid1;
+  int32 count = *pcount;
+  int32 countp1 = select(count == GZ - 1, 0, count + 1);
+  int1 checkc0[TM] = rxc < M;
+  int1 checkc1[TN] = ryc < N;
+  int1 checkc[TM, TN] = checkc0[:, newaxis] && checkc1[newaxis, :];
+  if(count == 0) {
+    @checkc *pc = c;
+    *pcount = countp1;
+  }
+  else {
+    @checkc *pc = c + *pc;
+    *pcount = countp1;
+  }
+  __atomic_cas(plock, 1, 0);
 }
 )";
 

examples/python/tensorflow/blocksparse.cpp

@@ -66,8 +66,26 @@ void matmul(restrict read_only fp32 *A, restrict read_only fp32 *B, fp32 *C,
     fp32 b[TN, 1] = checkb ? *pb : 0;
     c = dot(a, trans(b), c);
   }
+  int32 ridx = get_range_id(0);
+  int32 ridy = get_range_id(1);
   fp32* pc[TM, TN] = C + ryc[newaxis, :]*ldc + rxc[:, newaxis];
-  *pc = c;
+  int32 *plock = locks + ridx + ridy*grid0;
+  while(__atomic_cas(plock, 0, 1));
+  int32 *pcount = plock + grid0*grid1;
+  int32 count = *pcount;
+  int32 countp1 = select(count == GZ - 1, 0, count + 1);
+  int1 checkc0[TM] = rxc < M;
+  int1 checkc1[TN] = ryc < N;
+  int1 checkc[TM, TN] = checkc0[:, newaxis] && checkc1[newaxis, :];
+  if(count == 0) {
+    @checkc *pc = c;
+    *pcount = countp1;
+  }
+  else {
+    @checkc *pc = c + *pc;
+    *pcount = countp1;
+  }
+  __atomic_cas(plock, 1, 0);
 }
 )";
 
@@ -137,7 +155,6 @@ class BlockSparseGemmOp : public OpKernel {
     kernel->setArg(10, grid[0]);
     kernel->setArg(11, grid[1]);
     stream->enqueue(kernel, grid, {nthreads, 1, 1});
-    stream->synchronize();
   }
 
 private:

examples/python/tensorflow/blocksparse.py

@@ -3,7 +3,7 @@ import tensorflow as tf
 import numpy as np
 
 data_files_path = tf.resource_loader.get_data_files_path()
-library_dir = '/home/philippe/development/triton/build/examples/python/tensorflow'
+library_dir = '/home/philippe/Development/triton/build/examples/python/tensorflow'
 module = tf.load_op_library(os.path.join(library_dir, 'libtf_blocksparse.so'))
 
 M, N, K = 512, 512, 512

include/triton/jit.h

@@ -70,7 +70,6 @@ public:
         shmem_barriers.run(module);
       }
       vectorize.run(module);
-//      ir::print(module, std::cout);
     }
 
     codegen::tune tune;
@@ -93,6 +92,7 @@ private:
 
 public:
   jit(driver::context* context);
+  ~jit();
   void autotune(const char* name, const char* src, benchmark_t benchmark);
   void add_module(ir::module &module, const std::vector<unsigned>& params = {});
   void add_module(const char* name, const char* src, const std::vector<unsigned>& params = {});

lib/jit.cpp

@@ -71,11 +71,12 @@ std::unique_ptr<llvm::Module> jit::make_llvm_module(ir::module &module, passes_w
   llvm::Module* result = new llvm::Module(module.get_name(), llvm_context_);
   passes.selection.run(module, *result);
   // launch information
-  auto &launch_info_map = launch_info_map_[result->getName()];
-  launch_info_map.global_range_size.clear();
+  launch_information info;
+  info.global_range_size.clear();
   for(unsigned i = 0; i < passes.tune.get_num_global_range(); i++)
-    launch_info_map.global_range_size.push_back(passes.tune.get_global_range_size(i));
-  launch_info_map.num_threads = passes.tune.get_num_threads();
+    info.global_range_size.push_back(passes.tune.get_global_range_size(i));
+  info.num_threads = passes.tune.get_num_threads();
+  launch_info_map_.insert({result->getName(), info});
   return std::unique_ptr<llvm::Module>(result);
 }
 
@@ -93,9 +94,9 @@ std::unique_ptr<ir::module> jit::make_triton_module(const char *name, const char
 
 
 jit::jit(driver::context *context): driver_context_(context),
-                                    target_(context->device()->make_target()) {
-}
+                                    target_(context->device()->make_target()) { }
 
+jit::~jit(){ }
 
 void jit::autotune(const char *name, const char *src, benchmark_t benchmark) {
   // find metaparameters