Completing refactor of VLA CodeGen

mikaoP · mikaoP · commit 64176640a16f · 2019-10-30T08:35:18.000Z
Closes 3eb4f59d0eb

Things done:

- Emit QUAL.OSS.CAPTURED bundle in clang
- Gather QUAL.OSS.CAPTURED info in analysis
- assert if a value inside task body is not in QUAL.OSS.CAPTURED (instead of asserting with VLA dims)
- assert if QUAL.OSS.VLA.DIMS are not in QUAL.OSS.CAPTURED
- Add a flag to analysis to output all those QUAL.OSS.VLA.DIMS that are not in QUAL.OSS.CAPTURED
diff --git a/llvm/include/llvm/Analysis/OmpSsRegionAnalysis.h b/llvm/include/llvm/Analysis/OmpSsRegionAnalysis.h
@@ -26,6 +26,7 @@ struct TaskDSAInfo {
 
 // <VLA, VLA_dims>
 using TaskVLADimsInfo = MapVector<Value *, SetVector<Value *>>;
+using TaskCapturedInfo = SetVector<Value *>;
 
 struct DependInfo {
   int SymbolIndex;
@@ -56,6 +57,7 @@ struct TaskInfo {
   TaskDependsInfo DependsInfo;
   Value *Final = nullptr;
   Value *If = nullptr;
+  TaskCapturedInfo CapturedInfo;
   Instruction *Entry;
   Instruction *Exit;
 };
diff --git a/llvm/include/llvm/IR/LLVMContext.h b/llvm/include/llvm/IR/LLVMContext.h
@@ -98,6 +98,7 @@ class LLVMContext {
     OB_oss_dep_weakinout = 12,   // "OB_oss_dep_weakinout"
     OB_oss_final = 13,           // "OB_oss_final"
     OB_oss_if = 14,              // "OB_oss_if"
+    OB_oss_captured = 15,        // "OB_oss_captured"
   };
 
   /// getMDKindID - Return a unique non-zero ID for the specified metadata kind.
diff --git a/llvm/lib/Analysis/OmpSsRegionAnalysis.cpp b/llvm/lib/Analysis/OmpSsRegionAnalysis.cpp
@@ -34,7 +34,8 @@ enum PrintVerbosity {
   PV_Uses,
   PV_UnpackAndConst,
   PV_DsaMissing,
-  PV_DsaVLADimsMissing
+  PV_DsaVLADimsMissing,
+  PV_VLADimsCaptureMissing
 };
 
 static cl::opt<PrintVerbosity>
@@ -46,7 +47,8 @@ PrintVerboseLevel("print-verbosity",
   clEnumValN(PV_Uses, "uses", "Print task layout with uses"),
   clEnumValN(PV_UnpackAndConst, "unpack", "Print task layout with unpack instructions/constexprs needed in dependencies"),
   clEnumValN(PV_DsaMissing, "dsa_missing", "Print task layout with uses without DSA"),
-  clEnumValN(PV_DsaVLADimsMissing, "dsa_vla_dims_missing", "Print task layout with DSAs without VLA info or VLA info without DSAs")));
+  clEnumValN(PV_DsaVLADimsMissing, "dsa_vla_dims_missing", "Print task layout with DSAs without VLA info or VLA info without DSAs"),
+  clEnumValN(PV_VLADimsCaptureMissing, "vla_dims_capture_missing", "Print task layout with VLA dimensions without capture")));
 
 char OmpSsRegionAnalysisPass::ID = 0;
 
@@ -77,6 +79,11 @@ static bool valueInVLADimsBundles(const TaskVLADimsInfo& VLADimsInfo,
   return false;
 }
 
+static bool valueInCapturedBundle(const TaskCapturedInfo& CapturedInfo,
+                                  Value *const V) {
+  return CapturedInfo.count(V);
+}
+
 void OmpSsRegionAnalysisPass::print(raw_ostream &OS, const Module *M) const {
   for (auto it = TaskProgramOrder.begin(); it != TaskProgramOrder.end(); ++it) {
     Instruction *I = it->first;
@@ -140,14 +147,25 @@ void OmpSsRegionAnalysisPass::print(raw_ostream &OS, const Module *M) const {
         DSAVLADimsFreqMap[VLAWithDimsMap.first]++;
       }
       for (const auto &Pair : DSAVLADimsFreqMap) {
-        // It's expected to have only two VLA bundles, the DSA and de dimensions
+        // It's expected to have only two VLA bundles, the DSA and dimensions
         if (Pair.second != 2) {
           dbgs() << "\n";
           dbgs() << std::string((Depth + 1) * PrintSpaceMultiplier, ' ');
           Pair.first->printAsOperand(dbgs(), false);
         }
       }
     }
+    else if (PrintVerboseLevel == PV_VLADimsCaptureMissing) {
+      for (auto &VLAWithDimsMap : Info.VLADimsInfo) {
+        for (Value *const &V : VLAWithDimsMap.second) {
+          if (!valueInCapturedBundle(Info.CapturedInfo, V)) {
+            dbgs() << "\n";
+            dbgs() << std::string((Depth + 1) * PrintSpaceMultiplier, ' ');
+            V->printAsOperand(dbgs(), false);
+          }
+        }
+      }
+    }
 
     dbgs() << "\n";
   }
@@ -245,7 +263,7 @@ static bool insertUniqInstInProgramOrder(SmallVectorImpl<Instruction *> &InstLis
 }
 
 static void gatherUnpackInstructions(const TaskDSAInfo &DSAInfo,
-                                     const TaskVLADimsInfo &VLADimsInfo,
+                                     const TaskCapturedInfo &CapturedInfo,
                                      const OrderedInstructions &OI,
                                      DependInfo &DI,
                                      TaskAnalysisInfo &TAI,
@@ -265,13 +283,12 @@ static void gatherUnpackInstructions(const TaskDSAInfo &DSAInfo,
     Value *Dep = It->second;
     WorkList.erase(It);
     bool IsDSA = valueInDSABundles(DSAInfo, Cur);
-    // TODO: this will be get from captured info
-    bool IsVLADim = valueInVLADimsBundles(VLADimsInfo, Cur);
+    bool IsCaptured = valueInCapturedBundle(CapturedInfo, Cur);
     // Go over all uses until:
     //  1. We get a DSA so assign a symbol index
-    //  2. We get a VLA dimension, so we're done. We don't want to move
-    //  instructions that generate the vla dimension
-    if (!IsDSA && !IsVLADim) {
+    //  2. We get a Capture value, so we're done. We don't want to move
+    //  instructions that generate this
+    if (!IsDSA && !IsCaptured) {
       if (ConstantExpr *CE = dyn_cast<ConstantExpr>(Cur)) {
         for (Use &U : CE->operands()) {
           WorkList.emplace_back(U.get(), Dep);
@@ -300,7 +317,7 @@ static void gatherUnpackInstructions(const TaskDSAInfo &DSAInfo,
 static void gatherDependsInfoFromBundles(const SmallVectorImpl<OperandBundleDef> &OpBundles,
                                     const OrderedInstructions &OI,
                                     const TaskDSAInfo &DSAInfo,
-                                    const TaskVLADimsInfo &VLADimsInfo,
+                                    const TaskCapturedInfo &CapturedInfo,
                                     TaskAnalysisInfo &TAI,
                                     SmallVectorImpl<DependInfo> &DependsList,
                                     SmallVectorImpl<Instruction *> &UnpackInsts,
@@ -317,7 +334,7 @@ static void gatherDependsInfoFromBundles(const SmallVectorImpl<OperandBundleDef>
       DI.Dims.push_back(OBArgs[i]);
     }
 
-    gatherUnpackInstructions(DSAInfo, VLADimsInfo, OI, DI, TAI, UnpackInsts, UnpackConsts);
+    gatherUnpackInstructions(DSAInfo, CapturedInfo, OI, DI, TAI, UnpackInsts, UnpackConsts);
 
     DependsList.push_back(DI);
   }
@@ -327,50 +344,50 @@ static void gatherDependsInfoFromBundles(const SmallVectorImpl<OperandBundleDef>
 static void gatherDependsInfoWithID(const IntrinsicInst *I,
                                     const OrderedInstructions &OI,
                                     const TaskDSAInfo &DSAInfo,
-                                    const TaskVLADimsInfo &VLADimsInfo,
+                                    const TaskCapturedInfo &CapturedInfo,
                                     TaskAnalysisInfo &TAI,
                                     SmallVectorImpl<DependInfo> &DependsList,
                                     SmallVectorImpl<Instruction *> &UnpackInsts,
                                     SetVector<ConstantExpr *> &UnpackConsts,
                                     uint64_t Id) {
   SmallVector<OperandBundleDef, 4> OpBundles;
   getOperandBundlesAsDefsWithID(I, OpBundles, Id);
-  gatherDependsInfoFromBundles(OpBundles, OI, DSAInfo, VLADimsInfo, TAI, DependsList, UnpackInsts, UnpackConsts);
+  gatherDependsInfoFromBundles(OpBundles, OI, DSAInfo, CapturedInfo, TAI, DependsList, UnpackInsts, UnpackConsts);
 }
 
 // Gathers all dependencies needed information
 static void gatherDependsInfo(const IntrinsicInst *I, TaskInfo &TI,
                               TaskAnalysisInfo &TAI,
                               const OrderedInstructions &OI) {
-  gatherDependsInfoWithID(I, OI, TI.DSAInfo, TI.VLADimsInfo, TAI,
+  gatherDependsInfoWithID(I, OI, TI.DSAInfo, TI.CapturedInfo, TAI,
                           TI.DependsInfo.Ins,
                           TI.DependsInfo.UnpackInstructions,
                           TI.DependsInfo.UnpackConstants,
                           LLVMContext::OB_oss_dep_in);
-  gatherDependsInfoWithID(I, OI, TI.DSAInfo, TI.VLADimsInfo, TAI,
+  gatherDependsInfoWithID(I, OI, TI.DSAInfo, TI.CapturedInfo, TAI,
                           TI.DependsInfo.Outs,
                           TI.DependsInfo.UnpackInstructions,
                           TI.DependsInfo.UnpackConstants,
                           LLVMContext::OB_oss_dep_out);
-  gatherDependsInfoWithID(I, OI, TI.DSAInfo, TI.VLADimsInfo, TAI,
+  gatherDependsInfoWithID(I, OI, TI.DSAInfo, TI.CapturedInfo, TAI,
                           TI.DependsInfo.Inouts,
                           TI.DependsInfo.UnpackInstructions,
                           TI.DependsInfo.UnpackConstants,
                           LLVMContext::OB_oss_dep_inout);
 
-  gatherDependsInfoWithID(I, OI, TI.DSAInfo, TI.VLADimsInfo, TAI,
+  gatherDependsInfoWithID(I, OI, TI.DSAInfo, TI.CapturedInfo, TAI,
                           TI.DependsInfo.WeakIns,
                           TI.DependsInfo.UnpackInstructions,
                           TI.DependsInfo.UnpackConstants,
                           LLVMContext::OB_oss_dep_weakin);
 
-  gatherDependsInfoWithID(I, OI, TI.DSAInfo, TI.VLADimsInfo, TAI,
+  gatherDependsInfoWithID(I, OI, TI.DSAInfo, TI.CapturedInfo, TAI,
                           TI.DependsInfo.WeakOuts,
                           TI.DependsInfo.UnpackInstructions,
                           TI.DependsInfo.UnpackConstants,
                           LLVMContext::OB_oss_dep_weakout);
 
-  gatherDependsInfoWithID(I, OI, TI.DSAInfo, TI.VLADimsInfo, TAI,
+  gatherDependsInfoWithID(I, OI, TI.DSAInfo, TI.CapturedInfo, TAI,
                           TI.DependsInfo.WeakInouts,
                           TI.DependsInfo.UnpackInstructions,
                           TI.DependsInfo.UnpackConstants,
@@ -383,6 +400,20 @@ static void gatherIfFinalInfo(const IntrinsicInst *I, TaskInfo &TI) {
   getValueFromOperandBundleWithID(I, TI.If, LLVMContext::OB_oss_if);
 }
 
+// It's expected to have VLA dims info before calling this
+static void gatherCapturedInfo(const IntrinsicInst *I, TaskInfo &TI) {
+  getValueListFromOperandBundlesWithID(I, TI.CapturedInfo, LLVMContext::OB_oss_captured);
+  if (!DisableChecks) {
+    // VLA Dims that are not Captured is an error
+    for (auto &VLAWithDimsMap : TI.VLADimsInfo) {
+      for (Value *const &V : VLAWithDimsMap.second) {
+        if (!valueInCapturedBundle(TI.CapturedInfo, V))
+          llvm_unreachable("VLA dimension has not been captured");
+      }
+    }
+  }
+}
+
 void OmpSsRegionAnalysisPass::getOmpSsFunctionInfo(
     Function &F, DominatorTree &DT, FunctionInfo &FI,
     TaskFunctionAnalysisInfo &TFAI,
@@ -417,6 +448,7 @@ void OmpSsRegionAnalysisPass::getOmpSsFunctionInfo(
 
           gatherDSAInfo(II, T.Info);
           gatherVLADimsInfo(II, T.Info);
+          gatherCapturedInfo(II, T.Info);
           gatherDependsInfo(II, T.Info, T.AnalysisInfo, OI);
           gatherIfFinalInfo(II, T.Info);
 
@@ -448,7 +480,7 @@ void OmpSsRegionAnalysisPass::getOmpSsFunctionInfo(
               T.AnalysisInfo.UsesBeforeEntry.insert(I2);
             if (!DisableChecks
                 && !valueInDSABundles(T.Info.DSAInfo, I2)
-                && !valueInVLADimsBundles(T.Info.VLADimsInfo, I2)) {
+                && !valueInCapturedBundle(T.Info.CapturedInfo, I2)) {
                 llvm_unreachable("Value supposed to be inside task entry "
                                  "OperandBundle not found.");
               }
@@ -457,7 +489,7 @@ void OmpSsRegionAnalysisPass::getOmpSsFunctionInfo(
             T.AnalysisInfo.UsesBeforeEntry.insert(A);
             if (!DisableChecks
                 && !valueInDSABundles(T.Info.DSAInfo, A)
-                && !valueInVLADimsBundles(T.Info.VLADimsInfo, A)) {
+                && !valueInCapturedBundle(T.Info.CapturedInfo, A)) {
               llvm_unreachable("Value supposed to be inside task entry "
                                "OperandBundle not found.");
             }
diff --git a/llvm/lib/IR/LLVMContext.cpp b/llvm/lib/IR/LLVMContext.cpp
@@ -122,6 +122,11 @@ LLVMContext::LLVMContext() : pImpl(new LLVMContextImpl(*this)) {
   assert(OSSIfEntry->second == LLVMContext::OB_oss_if &&
          "oss_if operand bundle id drifted!");
   (void)OSSIfEntry;
+
+  auto *OSSCapturedEntry = pImpl->getOrInsertBundleTag("QUAL.OSS.CAPTURED");
+  assert(OSSCapturedEntry->second == LLVMContext::OB_oss_captured &&
+         "oss_captured operand bundle id drifted!");
+  (void)OSSCapturedEntry;
   // END OmpSs IDs
 
   SyncScope::ID SingleThreadSSID =
diff --git a/llvm/test/Analysis/OmpSsRegionAnalysis/task_layout_with_unpack.ll b/llvm/test/Analysis/OmpSsRegionAnalysis/task_layout_with_unpack.ll
@@ -105,7 +105,7 @@ entry:
   %9 = mul i64 %4, 4
   %10 = mul i64 %7, 4
 ; to here
-  %11 = call token @llvm.directive.region.entry() [ "DIR.OSS"([5 x i8] c"TASK\00"), "QUAL.OSS.SHARED"(i32* %vla), "QUAL.OSS.VLA.DIMS"(i32* %vla, i64 %1), "QUAL.OSS.FIRSTPRIVATE"(i32* %k.addr), "QUAL.OSS.FIRSTPRIVATE"(i32* %j.addr), "QUAL.OSS.DEP.IN"(i32* %vla, i64 %8, i64 %9, i64 %10) ]
+  %11 = call token @llvm.directive.region.entry() [ "DIR.OSS"([5 x i8] c"TASK\00"), "QUAL.OSS.SHARED"(i32* %vla), "QUAL.OSS.VLA.DIMS"(i32* %vla, i64 %1), "QUAL.OSS.CAPTURED"(i64 %1), "QUAL.OSS.FIRSTPRIVATE"(i32* %k.addr), "QUAL.OSS.FIRSTPRIVATE"(i32* %j.addr), "QUAL.OSS.DEP.IN"(i32* %vla, i64 %8, i64 %9, i64 %10) ]
   call void @llvm.directive.region.exit(token %11)
   %12 = load i8*, i8** %saved_stack, align 8
   call void @llvm.stackrestore(i8* %12)
diff --git a/llvm/test/Analysis/OmpSsRegionAnalysis/task_layout_with_vla_dims_capture_missing.ll b/llvm/test/Analysis/OmpSsRegionAnalysis/task_layout_with_vla_dims_capture_missing.ll
@@ -0,0 +1,36 @@
+; RUN: opt -ompss-2-regions -analyze -disable-checks -print-verbosity=vla_dims_capture_missing < %s 2>&1 | FileCheck %s
+
+define dso_local void @vla_section_dep(i32 %n) {
+entry:
+  %n.addr = alloca i32, align 4
+  %saved_stack = alloca i8*, align 8
+  %__vla_expr0 = alloca i64, align 8
+  %__vla_expr1 = alloca i64, align 8
+  store i32 %n, i32* %n.addr, align 4
+  %0 = load i32, i32* %n.addr, align 4
+  %add = add nsw i32 %0, 1
+  %1 = zext i32 %add to i64
+  %2 = load i32, i32* %n.addr, align 4
+  %add1 = add nsw i32 %2, 2
+  %3 = zext i32 %add1 to i64
+  %4 = call i8* @llvm.stacksave()
+  store i8* %4, i8** %saved_stack, align 8
+  %5 = mul nuw i64 %1, %3
+  %vla = alloca i32, i64 %5, align 16
+  store i64 %1, i64* %__vla_expr0, align 8
+  store i64 %3, i64* %__vla_expr1, align 8
+  %6 = call token @llvm.directive.region.entry() [ "DIR.OSS"([5 x i8] c"TASK\00"), "QUAL.OSS.SHARED"(i32* %vla), "QUAL.OSS.VLA.DIMS"(i32* %vla, i64 %1, i64 %3) ]
+  call void @llvm.directive.region.exit(token %6)
+  %7 = load i8*, i8** %saved_stack, align 8
+  call void @llvm.stackrestore(i8* %7)
+  ret void
+}
+
+; CHECK: [0] %6
+; CHECK-NEXT:   %1
+; CHECK-NEXT:   %3
+
+declare i8* @llvm.stacksave()
+declare token @llvm.directive.region.entry()
+declare void @llvm.directive.region.exit(token)
+declare void @llvm.stackrestore(i8*)
