[codegen] bugfix in builder insert point for predicated instructions

examples/cpp/dot.cpp

@@ -53,26 +53,11 @@ 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];
-  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);
+  @checkc *pc = c;
 }
 )";
 

include/triton/jit.h

@@ -58,6 +58,7 @@ public:
                       target_(target) { }
 
     void target_independent(ir::module &module) {
+//        ir::print(module, std::cout);
         optimize_dot.run(module);
         optimize_trans.run(module);
     }

lib/codegen/selection.cpp

@@ -632,11 +632,13 @@ void selection::lower_tile_instruction(ir::instruction *ins, llvm::IRBuilder<> &
   LLVMContext &ctx = builder.getContext();
   Function *fn = block->getParent();
   ir::value *mask = ins->get_mask_pred();
+  BasicBlock *last_block = nullptr;
   auto set_mask_insert_pt = [&](indices_t idx){
     if(mask){
       distributed_tile *mask_tile = (distributed_tile*)tmap_.at(ins->get_mask_pred());
       BasicBlock *block = pmap_.at({mask_tile, idx});
       builder.SetInsertPoint(block->getTerminator());
+      last_block = last_block_.at({mask_tile, idx});
     }
   };
   // store
@@ -646,7 +648,6 @@ void selection::lower_tile_instruction(ir::instruction *ins, llvm::IRBuilder<> &
     ptr->for_each([&](indices_t idx){
       set_mask_insert_pt(idx);
       StoreInst *store = new StoreInst(value->get_value(idx), ptr->get_value(idx));
-//      store->setAlignment(16);
       builder.Insert(store);
     });
   }
@@ -847,8 +848,11 @@ void selection::lower_tile_instruction(ir::instruction *ins, llvm::IRBuilder<> &
       });
     }
   }
-  if(mask)
+  if(mask){
     builder.SetInsertPoint(block);
+    if(last_block)
+      builder.SetInsertPoint(last_block);
+  }
 }
 
 void selection::lower_instruction(ir::instruction *src, IRBuilder<> &builder) {