Merge pull request #784 from hvdijk/vecz-barrier-id

[vecz] Track barriers by ID.

Author: hvdijk

modules/compiler/compiler_pipeline/include/compiler/utils/barrier_regions.h

@@ -31,6 +31,8 @@
 #include <llvm/Transforms/Utils/ValueMapper.h>
 #include <multi_llvm/llvm_version.h>
 
+#include <map>
+
 #include "pass_functions.h"
 
 namespace llvm {
@@ -82,10 +84,11 @@ struct BarrierRegion {
   BarrierSchedule schedule = BarrierSchedule::Unordered;
 };
 
-using BarrierGraph = llvm::SmallVector<BarrierRegion, 8>;
-
 class Barrier {
  public:
+  /// @brief Type for ids of new kernel functions
+  using kernel_id_map_t = std::map<unsigned, llvm::Function *>;
+
   Barrier(llvm::Module &m, llvm::Function &f, bool IsDebug)
       : live_var_mem_ty_(nullptr),
         size_t_bytes(compiler::utils::getSizeTypeBytes(m)),
@@ -106,6 +109,9 @@ class Barrier {
   /// @brief returns the maximum alignment of the barrier struct
   unsigned getLiveVarMaxAlignment() const { return max_live_var_alignment; }
 
+  /// @brief gets the split subkernels
+  const kernel_id_map_t &getSubkernels() const { return kernel_id_map_; }
+
   /// @brief gets the split subkernel for the given barrier id
   llvm::Function *getSubkernel(unsigned id) const {
     return kernel_id_map_.find(id)->second;
@@ -116,7 +122,7 @@ class Barrier {
 
   llvm::CallInst *getBarrierCall(unsigned id) const {
     return llvm::dyn_cast_or_null<llvm::CallInst>(
-        barrier_graph[id - kBarrier_FirstID].barrier_inst);
+        barrier_region_id_map_.find(id)->second.barrier_inst);
   }
 
   /// @brief gets the size of the fixed sized part of the barrier struct
@@ -135,12 +141,12 @@ class Barrier {
 
   /// @brief gets the barrier IDs of the successors of the given barrier region
   const llvm::SmallVectorImpl<unsigned> &getSuccessorIds(unsigned id) const {
-    return barrier_graph[id - kBarrier_FirstID].successor_ids;
+    return barrier_region_id_map_.find(id)->second.successor_ids;
   }
 
   /// @brief gets the barrier IDs of the successors of the given barrier region
   BarrierSchedule getSchedule(unsigned id) const {
-    return barrier_graph[id - kBarrier_FirstID].schedule;
+    return barrier_region_id_map_.find(id)->second.schedule;
   }
 
   /// @brief replaces a subkernel with a given function
@@ -223,8 +229,8 @@ class Barrier {
   using live_variable_scalables_map_t = live_variable_index_map_t;
   /// @brief Type for ids of barriers
   using barrier_id_map_t = llvm::DenseMap<llvm::BasicBlock *, unsigned>;
-  /// @brief Type for ids of new kernel functions
-  using kernel_id_map_t = llvm::DenseMap<unsigned, llvm::Function *>;
+  /// @brief Type for ids of barrier regions
+  using barrier_region_id_map_t = std::map<unsigned, BarrierRegion>;
   /// @brief Type for map from ids to fence instructions
   using fence_id_map_t = llvm::DenseMap<unsigned, llvm::FenceInst *>;
   /// @brief Type between block and instruction for barrier.
@@ -245,6 +251,8 @@ class Barrier {
   live_variable_scalables_map_t live_variable_scalables_map_;
   /// @brief Keep ids of barriers.
   barrier_id_map_t barrier_id_map_;
+  /// @brief Look up a barrier region by its id.
+  barrier_region_id_map_t barrier_region_id_map_;
   /// @brief Keep ids of barriers.
   kernel_id_map_t kernel_id_map_;
   /// @brief Keep struct types for live variables' memory layout.
@@ -269,8 +277,6 @@ class Barrier {
 
   size_t size_t_bytes;
 
-  BarrierGraph barrier_graph;
-
   llvm::Module &module_;
   llvm::Function &func_;
 

modules/compiler/compiler_pipeline/source/barrier_regions.cpp

@@ -437,11 +437,12 @@ void compiler::utils::Barrier::Run(llvm::ModuleAnalysisManager &mam) {
   bi_ = &mam.getResult<BuiltinInfoAnalysis>(module_);
   FindBarriers();
 
+  kernel_id_map_[kBarrier_EndID] = nullptr;
+
   if (barriers_.empty()) {
     // If there are no barriers, we can use the original function as the
     // single barrier region.
-    barrier_graph.emplace_back();
-    auto &node = barrier_graph.back();
+    auto &node = barrier_region_id_map_[kBarrier_FirstID];
     node.entry = &func_.getEntryBlock();
     node.id = kBarrier_FirstID;
     node.successor_ids.push_back(kBarrier_EndID);
@@ -513,11 +514,9 @@ void compiler::utils::Barrier::FindBarriers() {
         if (callee != nullptr) {
           const auto B = bi_->analyzeBuiltin(*callee);
           if (BuiltinInfo::isMuxBuiltinWithWGBarrierID(B.ID)) {
-            unsigned id = ~0u;
             auto *const id_param = call_inst->getOperand(0);
-            if (auto *const id_param_c = dyn_cast<ConstantInt>(id_param)) {
-              id = id_param_c->getZExtValue();
-            }
+            auto *const id_param_c = cast<ConstantInt>(id_param);
+            const auto id = id_param_c->getZExtValue();
             orderedBarriers.emplace_back(id, call_inst);
           }
         }
@@ -548,13 +547,15 @@ void compiler::utils::Barrier::SplitBlockwithBarrier() {
     exit_stub = MakeStubFunction("__barrier_exit", module_, stub_cc);
   }
 
-  barrier_graph.emplace_back();
-  auto &node = barrier_graph.back();
+  auto &node = barrier_region_id_map_[kBarrier_FirstID];
   node.entry = &func_.getEntryBlock();
   node.id = kBarrier_FirstID;
 
-  unsigned barrier_id = kBarrier_StartNewID;
   for (CallInst *split_point : barriers_) {
+    // ID identifying which barrier invoked stub used as argument to call.
+    auto *id = cast<ConstantInt>(split_point->getOperand(0));
+    const auto barrier_id = kBarrier_StartNewID + id->getZExtValue();
+
     if (is_debug_) {
       assert(entry_stub != nullptr);  // Guaranteed as is_debug_ is const.
       assert(exit_stub != nullptr);   // Guaranteed as is_debug_ is const.
@@ -564,10 +565,6 @@ void compiler::utils::Barrier::SplitBlockwithBarrier() {
       // them at a point where live variables have already been loaded. This
       // info won't be available till later.
 
-      // ID identifying which barrier invoked stub used as argument to call.
-      // This number monotonically increases from 0 for each barrier.
-      auto id = ConstantInt::get(Type::getInt32Ty(module_.getContext()),
-                                 barrier_id - kBarrier_StartNewID);
       // Call invoking entry stub
       auto entry_caller = CallInst::Create(entry_stub, id);
       entry_caller->setDebugLoc(split_point->getDebugLoc());
@@ -583,10 +580,9 @@ void compiler::utils::Barrier::SplitBlockwithBarrier() {
           std::make_pair(entry_caller, exit_caller);
     }
 
-    barrier_graph.emplace_back();
-    auto &node = barrier_graph.back();
+    auto &node = barrier_region_id_map_[barrier_id];
     node.barrier_inst = split_point;
-    node.id = barrier_id++;
+    node.id = barrier_id;
     node.schedule = getBarrierSchedule(*split_point);
 
     // Our scan implementation requires a linear work-item ordering, to loop
@@ -603,7 +599,7 @@ void compiler::utils::Barrier::SplitBlockwithBarrier() {
   // We have to gather the basic block data after splitting, because we
   // might not be processing barriers in program order, and things can get
   // awfully confused.
-  for (auto &node : barrier_graph) {
+  for (auto &[i, node] : barrier_region_id_map_) {
     if (node.barrier_inst) {
       auto *const bb = node.barrier_inst->getParent();
       barrier_id_map_[bb] = node.id;
@@ -770,7 +766,7 @@ void compiler::utils::Barrier::FindLiveVariables() {
     }
   }
 
-  for (auto &region : barrier_graph) {
+  for (auto &[i, region] : barrier_region_id_map_) {
     GatherBarrierRegionBlocks(region);
     GatherBarrierRegionUses(region, func_args);
     whole_live_variables_set_.set_union(region.uses_int);
@@ -1150,9 +1146,9 @@ Function *compiler::utils::Barrier::GenerateNewKernel(BarrierRegion &region) {
     } else if (ReturnInst *ret =
                    dyn_cast<ReturnInst>(cloned_bb->getTerminator())) {
       // Change return instruction with end barrier number.
-      ConstantInt *cst_zero =
+      ConstantInt *cst_endid =
           ConstantInt::get(Type::getInt32Ty(context), kBarrier_EndID);
-      ReturnInst *new_ret = ReturnInst::Create(context, cst_zero);
+      ReturnInst *new_ret = ReturnInst::Create(context, cst_endid);
       new_ret->insertBefore(ret->getIterator());
       ret->replaceAllUsesWith(new_ret);
       ret->eraseFromParent();
@@ -1450,7 +1446,7 @@ BasicBlock *compiler::utils::Barrier::CloneBasicBlock(
 void compiler::utils::Barrier::SeperateKernelWithBarrier() {
   if (barriers_.empty()) return;
 
-  for (auto &region : barrier_graph) {
+  for (auto &[i, region] : barrier_region_id_map_) {
     kernel_id_map_[region.id] = GenerateNewKernel(region);
   }
 
@@ -1467,15 +1463,10 @@ void compiler::utils::Barrier::SeperateKernelWithBarrier() {
 
   LLVM_DEBUG({
     for (const auto &Kid : kernel_id_map_) {
-      dbgs() << "1. kernel_id[" << Kid.first << "] = " << Kid.second->getName()
+      dbgs() << "kernel_id[" << Kid.first << "] = " << Kid.second->getName()
              << "\n";
     }
 
-    for (unsigned I = kBarrier_FirstID;
-         I < kernel_id_map_.size() + kBarrier_FirstID; I++) {
-      dbgs() << "2. kernel_id[" << I << "] = " << kernel_id_map_[I]->getName()
-             << "\n";
-    }
     dbgs() << "\n\n" << module_ << "\n\n";
   });
 }

modules/compiler/compiler_pipeline/source/work_item_loops_pass.cpp

@@ -1552,13 +1552,14 @@ Function *compiler::utils::WorkItemLoopsPass::makeWrapperFunction(
   AllocaInst *nextID =
       entryIR.CreateAlloca(index_type, nullptr, "next_barrier_id");
 
-  SmallVector<BasicBlock *, 8> bbs;
-  const unsigned num_blocks = barrierMain.getNumSubkernels();
-  assert(!emitTail || barrierTail->getNumSubkernels() == num_blocks);
-
-  for (unsigned i = kBarrier_EndID; i <= num_blocks; i++) {
-    BasicBlock *bb = BasicBlock::Create(context, "sw.bb", new_wrapper);
-    bbs.push_back(bb);
+  std::map<unsigned, BasicBlock *> bbs;
+  // The vectorized kernel has been further optimized and may have removed
+  // unreachable barriers that are still present in the scalar kernel. But if
+  // they are unreachable, we know they must also be unreachable in the scalar
+  // kernel even if we have not yet detected that.
+
+  for (auto &[i, subkernel] : barrierMain.getSubkernels()) {
+    bbs[i] = BasicBlock::Create(context, "sw.bb", new_wrapper);
   }
 
   ScheduleGenerator schedule(M, barrierMain, barrierTail, BI);
@@ -1584,7 +1585,9 @@ Function *compiler::utils::WorkItemLoopsPass::makeWrapperFunction(
   // Branch directly into the first basic block.
   entryIR.CreateBr(bbs[kBarrier_FirstID]);
 
-  for (unsigned i = kBarrier_EndID; i <= num_blocks; i++) {
+  for (auto &[i_, subkernel_] : barrierMain.getSubkernels()) {
+    auto i = i_;
+
     // Keep it linear
     BasicBlock *const block = bbs[i];
     block->moveAfter(&new_wrapper->back());
@@ -1663,7 +1666,7 @@ Function *compiler::utils::WorkItemLoopsPass::makeWrapperFunction(
 
       if (num_succ == 1) {
         // If there is only one successor, we can branch directly to it
-        exitIR.CreateBr(bbs[successors.front()]);
+        exitIR.CreateBr(bbs.find(successors.front())->second);
       } else if (num_succ == 2) {
         // If there are exactly two successors, we can use a conditional branch
         auto *const bb_id = ConstantInt::get(index_type, successors[0]);
@@ -1673,8 +1676,8 @@ Function *compiler::utils::WorkItemLoopsPass::makeWrapperFunction(
         auto *const cmp_id =
             CmpInst::Create(Instruction::ICmp, CmpInst::ICMP_EQ, ld_next_id,
                             bb_id, "", br_block);
-        BranchInst::Create(bbs[successors[0]], bbs[successors[1]], cmp_id,
-                           br_block);
+        BranchInst::Create(bbs.find(successors[0])->second,
+                           bbs.find(successors[1])->second, cmp_id, br_block);
 
         exitIR.CreateBr(br_block);
       } else if (num_succ == 0) {
@@ -1700,9 +1703,9 @@ Function *compiler::utils::WorkItemLoopsPass::makeWrapperFunction(
         LoadInst *const ld_next_id =
             new LoadInst(index_type, nextID, "", switch_body);
         SwitchInst *const sw = SwitchInst::Create(
-            ld_next_id, bbs[successors[0]], num_succ, switch_body);
+            ld_next_id, bbs.find(successors[0])->second, num_succ, switch_body);
         for (const auto i : successors) {
-          sw->addCase(ConstantInt::get(index_type, i), bbs[i]);
+          sw->addCase(ConstantInt::get(index_type, i), bbs.find(i)->second);
         }
         exitIR.CreateBr(switch_body);
       }

modules/compiler/test/lit/passes/barriers-mismatch.ll

@@ -0,0 +1,87 @@
+; Copyright (C) Codeplay Software Limited
+;
+; Licensed under the Apache License, Version 2.0 (the "License") with LLVM
+; Exceptions; you may not use this file except in compliance with the License.
+; You may obtain a copy of the License at
+;
+;     https://github.com/codeplaysoftware/oneapi-construction-kit/blob/main/LICENSE.txt
+;
+; Unless required by applicable law or agreed to in writing, software
+; distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+; WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
+; License for the specific language governing permissions and limitations
+; under the License.
+;
+; SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
+; RUN: muxc --passes work-item-loops,verify -S %s  | FileCheck %s
+
+; This test checks the validity of a set of main/tail loops when the
+; barriers between vector and scalar kernels do not match up.
+
+target triple = "spir64-unknown-unknown"
+target datalayout = "e-i64:64-v16:16-v24:32-v32:32-v48:64-v96:128-v192:256-v256:256-v512:512-v1024:1024"
+
+define spir_kernel void @foo(ptr addrspace(1) %a) !codeplay_ca_vecz.base !0 {
+entry:
+  store i1 false, ptr addrspace(1) %a, align 4
+  call void @__mux_work_group_barrier(i32 0, i32 1, i32 272) #8
+  %load = load i1, ptr addrspace(1) %a, align 4
+  call void @__mux_work_group_barrier(i32 1, i32 1, i32 272) #8
+  br i1 %load, label %bb.1, label %bb.2
+
+bb.1:
+  store i1 poison, ptr addrspace(1) poison
+  call void @__mux_work_group_barrier(i32 2, i32 1, i32 272) #8
+  br label %bb.3
+
+bb.2:
+  store i1 true, ptr addrspace(1) %a
+  call void @__mux_work_group_barrier(i32 3, i32 1, i32 272) #8
+  br label %bb.3
+
+bb.3:
+  call void @__mux_work_group_barrier(i32 4, i32 1, i32 272) #8
+  ret void
+}
+
+define spir_kernel void @__vecz_v2_foo(ptr addrspace(1) %a) #0 !codeplay_ca_vecz.derived !2 {
+entry:
+  store i1 false, ptr addrspace(1) %a, align 4
+  call void @__mux_work_group_barrier(i32 0, i32 1, i32 272)
+  %load = load i1, ptr addrspace(1) %a, align 4
+  call void @__mux_work_group_barrier(i32 1, i32 1, i32 272)
+  %0 = xor i1 %load, true
+  call void @llvm.assume(i1 %0)
+  store i1 true, ptr addrspace(1) %a, align 1
+  call void @__mux_work_group_barrier(i32 3, i32 1, i32 272)
+  call void @__mux_work_group_barrier(i32 4, i32 1, i32 272)
+  ret void
+}
+
+declare void @__mux_work_group_barrier(i32, i32, i32)
+
+attributes #0 = { norecurse nounwind "mux-kernel"="entry-point" "mux-base-fn-name"="foo"}
+attributes #1 = { nounwind "mux-barrier-schedule"="linear" }
+
+; The block which has conditional undefined behavior in the scalar kernel.
+;
+; CHECK-LABEL: define internal spir_func i32 @foo.mux-barrier-region.6(ptr addrspace(1) %0, ptr %1)
+; CHECK:       bb.1:
+; CHECK-NEXT:    store i1 poison, ptr addrspace(1) poison, align 1
+;
+; The block which only follows after undefined behavior in the scalar kernel.
+;
+; CHECK-LABEL: define internal spir_func i32 @foo.mux-barrier-region.7(ptr addrspace(1) %0, ptr %1)
+; CHECK-NEXT:  barrier2:
+;
+; Check that we do not call the unreachable region.
+;
+; CHECK-LABEL: define spir_kernel void @foo.mux-barrier-wrapper(ptr addrspace(1) %a)
+; CHECK:         call spir_func i32 @foo.mux-barrier-region.6(
+; CHECK-NOT:     call spir_func i32 @foo.mux-barrier-region.7(
+
+; Vectorized by 2
+!0 = !{!1, ptr @__vecz_v2_foo}
+!1 = !{i32 2, i32 0, i32 0, i32 0}
+!2 = !{!1, ptr @foo}