[MLIR] tbb fixes (#189)

* fix plier parallel

* use loop to init reduce vars

* transfor to tbb parallel if have parallel attr or outermost loop

* we dont need fix_tls_observer, also do not recreate task arena each time

Author: Ivan Butygin

mlir-compiler/mlir-compiler/src/pipelines/parallel_to_tbb.cpp

@@ -50,7 +50,9 @@ struct ParallelToTbb : public mlir::OpRewritePattern<mlir::scf::ParallelOp>
         {
             return mlir::failure();
         }
-        if (!op->hasAttr(plier::attributes::getParallelName()))
+        bool need_parallel = op->hasAttr(plier::attributes::getParallelName()) ||
+                             !op->getParentOfType<mlir::scf::ParallelOp>();
+        if (!need_parallel)
         {
             return mlir::failure();
         }
@@ -85,31 +87,46 @@ struct ParallelToTbb : public mlir::OpRewritePattern<mlir::scf::ParallelOp>
             auto reduce = rewriter.create<mlir::AllocaOp>(loc, reduce_type);
             auto index = static_cast<unsigned>(it.index());
             reduce_vars[index] = reduce;
-            auto zero = getZeroVal(rewriter, loc, type);
-            mapping.map(op.initVals()[index], zero);
-            for (unsigned i = 0; i < max_concurrency; ++i)
+        }
+
+        auto reduce_init_body_builder = [&](mlir::OpBuilder &builder, mlir::Location loc, mlir::Value index, mlir::ValueRange args)
+        {
+            assert(args.empty());
+            (void)args;
+            for (auto it : llvm::enumerate(reduce_vars))
             {
-                mlir::Value index = rewriter.create<mlir::ConstantIndexOp>(loc, i);
-                rewriter.create<mlir::StoreOp>(loc, zero, reduce, index);
+                auto reduce = it.value();
+                auto type = op.getResultTypes()[it.index()];
+                auto zero = getZeroVal(rewriter, loc, type);
+                builder.create<mlir::StoreOp>(loc, zero, reduce, index);
             }
-        }
+            builder.create<mlir::scf::YieldOp>(loc);
+        };
+
+        auto reduce_lower_bound = rewriter.create<mlir::ConstantIndexOp>(loc, 0);
+        auto reduce_upper_bound = rewriter.create<mlir::ConstantIndexOp>(loc, max_concurrency);
+        auto reduce_step = rewriter.create<mlir::ConstantIndexOp>(loc, 1);
+        rewriter.create<mlir::scf::ForOp>(loc, reduce_lower_bound, reduce_upper_bound, reduce_step, llvm::None, reduce_init_body_builder);
 
         auto& old_body = op.getLoopBody().front();
         auto orig_lower_bound = op.lowerBound().front();
         auto orig_upper_bound = op.upperBound().front();
         auto orig_step = op.step().front();
         auto body_builder = [&](mlir::OpBuilder &builder, ::mlir::Location loc, mlir::Value lower_bound, mlir::Value upper_bound, mlir::Value thread_index)
         {
+            llvm::SmallVector<mlir::Value, 8> initVals(op.initVals().size());
             for (auto it : llvm::enumerate(op.initVals()))
             {
                 auto reduce_var = reduce_vars[it.index()];
                 auto val = builder.create<mlir::LoadOp>(loc, reduce_var, thread_index);
-                mapping.map(it.value(), val);
+                initVals[it.index()] = val;
             }
             auto new_op = mlir::cast<mlir::scf::ParallelOp>(builder.clone(*op, mapping));
+            new_op->removeAttr(plier::attributes::getParallelName());
             assert(new_op->getNumResults() == reduce_vars.size());
             new_op.lowerBoundMutable().assign(lower_bound);
             new_op.upperBoundMutable().assign(upper_bound);
+            new_op.initValsMutable().assign(initVals);
             for (auto it : llvm::enumerate(new_op->getResults()))
             {
                 auto reduce_var = reduce_vars[it.index()];
@@ -119,10 +136,6 @@ struct ParallelToTbb : public mlir::OpRewritePattern<mlir::scf::ParallelOp>
 
         rewriter.create<plier::ParallelOp>(loc, orig_lower_bound, orig_upper_bound, orig_step, body_builder);
 
-        auto reduce_lower_bound = rewriter.create<mlir::ConstantIndexOp>(loc, 0);
-        auto reduce_upper_bound = rewriter.create<mlir::ConstantIndexOp>(loc, max_concurrency);
-        auto reduce_step = rewriter.create<mlir::ConstantIndexOp>(loc, 1);
-
         auto reduce_body_builder = [&](mlir::OpBuilder &builder, mlir::Location loc, mlir::Value index, mlir::ValueRange args)
         {
             assert(args.size() == reduce_vars.size());

numba/np/ufunc/tbbpool.cpp

@@ -42,6 +42,21 @@ Implement parallel vectorize workqueue on top of Intel TBB.
 
 static tbb::task_group *tg = NULL;
 static tbb::task_scheduler_init *tsi = NULL;
+
+namespace
+{
+struct ThreadContext
+{
+    ThreadContext(int n_threads):
+        num_threads(n_threads),
+        arena(n_threads) {}
+
+    int num_threads = 0;
+    tbb::task_arena arena;
+};
+static ThreadContext* thread_context = nullptr;
+}
+
 static int tsi_count = 0;
 
 #ifdef _MSC_VER
@@ -209,15 +224,15 @@ parallel_for(void *fn, char **args, size_t *dimensions, size_t *steps, void *dat
 using parallel_for2_fptr = void(*)(size_t, size_t, size_t, void*);
 static void parallel_for2(size_t lower_bound, size_t upper_bound, size_t step, parallel_for2_fptr func, void* ctx)
 {
-    auto num_threads = get_num_threads();
+    auto context = thread_context;
+    assert(nullptr != context);
+    auto num_threads = context->num_threads;
     if(_DEBUG)
     {
         printf("parallel_for2 %d %d %d %d\n", (int)lower_bound, (int)upper_bound, (int)step, (int)num_threads);
     }
-    tbb::task_arena limited(num_threads);
-    fix_tls_observer observer(limited, num_threads);
 
-    limited.execute([&]
+    context->arena.execute([&]
     {
         size_t count = (upper_bound - lower_bound - 1) / step + 1;
         size_t grain = std::max(size_t(1), std::min(count / num_threads / 2, size_t(64)));
@@ -284,6 +299,8 @@ static void unload_tbb(void)
         tbb::set_assertion_handler(orig);
         delete tsi;
         tsi = NULL;
+        delete thread_context;
+        thread_context = nullptr;
     }
 }
 #endif
@@ -300,6 +317,8 @@ static void launch_threads(int count)
     tg = new tbb::task_group;
     tg->run([] {}); // start creating threads asynchronously
 
+    thread_context = new ThreadContext(count);
+
     _INIT_NUM_THREADS = count;
 
 #ifndef _MSC_VER