Add debug build support in bash scripts

scripts/gen_gpu.sh

@@ -7,10 +7,18 @@ LLVM_DIR=`dirname $LLVM_DIR`
 LLVM_DIR=`dirname $LLVM_DIR`
 echo "Using LLVM rooted at : $LLVM_DIR"
 
+if [ "$BUILD_MODE" = "debug" ]; then
+    MLIR_BUILD="debug_build"
+    CONFIG="Debug"
+else
+    MLIR_BUILD="build"
+    CONFIG="Release"
+fi
+
+# Default cmake command
 CMAKE="cmake"
-MLIR_BUILD="build"
-CONFIG="Release"
 
+# Process options
 while getopts "d:m:c:" opt
 do
    case "$opt" in
@@ -24,7 +32,7 @@ echo "Using cmake command : $CMAKE"
 echo "Using MLIR_BUILD : $MLIR_BUILD"
 echo "Using configuration : $CONFIG"
 
-# run this from the build directory
+# Run the cmake command from the build directory
 $CMAKE -G Ninja .. -DCMAKE_EXPORT_COMPILE_COMMANDS=1 \
        -DMLIR_DIR=$LLVM_DIR/$MLIR_BUILD/lib/cmake/mlir \
        -DLLVM_BUILD_DIRECTORY=$LLVM_DIR/$MLIR_BUILD/ \
@@ -33,4 +41,4 @@ $CMAKE -G Ninja .. -DCMAKE_EXPORT_COMPILE_COMMANDS=1 \
 #      -DAMD_GPU_SUPPORT=ON
 #      -DCMAKE_C_COMPILER=clang -DCMAKE_CXX_COMPILER=clang++
 #      -DLLVM_ENABLE_LLD=ON
-#      -DCMAKE_CXX_FLAGS="-std=c++17"
+#      -DCMAKE_CXX_FLAGS="-std=c++17"

scripts/setupPythonEnv.sh

@@ -4,7 +4,11 @@ SCRIPTPATH=`dirname $SCRIPT`
 TREEBEARD_DIR=`dirname $SCRIPTPATH`
 echo "Using Treebeard rooted at : $TREEBEARD_DIR"
 
-BUILD_DIR="build"
+if [ "$BUILD_MODE" = "debug" ]; then
+    BUILD_DIR="debug_build"
+else
+     BUILD_DIR="build"
+fi
 
 while getopts "m:" opt
 do

scripts/setupSilvanForge.sh

@@ -2,29 +2,48 @@
 git clone https://github.com/asprasad/llvm-project.git
 cd llvm-project/
 git checkout release/16.x
-mkdir build
-cd build
-cmake -G Ninja ../llvm -DLLVM_ENABLE_PROJECTS="llvm;clang;lld;mlir;openmp" -DLLVM_BUILD_EXAMPLES=ON -DLLVM_TARGETS_TO_BUILD="X86;NVPTX;AMDGPU" -DCMAKE_BUILD_TYPE=Release -DLLVM_ENABLE_ASSERTIONS=ON -DCMAKE_C_COMPILER=clang -DCMAKE_CXX_COMPILER=clang++ -DLLVM_ENABLE_LLD=ON -DMLIR_ENABLE_CUDA_RUNNER=ON -DMLIR_INCLUDE_INTEGRATION_TESTS=ON
+
+if [ "$BUILD_MODE" = "debug" ]; then
+    BUILD_DIR="debug_build"
+    CONFIG="Debug"
+else
+    BUILD_DIR="build"
+    CONFIG="Release"
+fi
+
+mkdir -p $BUILD_DIR
+cd $BUILD_DIR
+cmake -G Ninja ../llvm -DLLVM_ENABLE_PROJECTS="llvm;clang;lld;mlir;openmp" -DLLVM_BUILD_EXAMPLES=ON -DLLVM_TARGETS_TO_BUILD="X86;NVPTX;AMDGPU" -DCMAKE_BUILD_TYPE=$CONFIG -DLLVM_ENABLE_ASSERTIONS=ON -DCMAKE_C_COMPILER=clang -DCMAKE_CXX_COMPILER=clang++ -DLLVM_ENABLE_LLD=ON -DMLIR_ENABLE_CUDA_RUNNER=ON -DMLIR_INCLUDE_INTEGRATION_TESTS=ON
 cmake --build .
 
 # Clone and build Treebeard
 cd ../mlir/examples
-git clone https://github.com/asprasad/treebeard.git
+
+if [ ! -d "treebeard" ]; then
+    git clone https://github.com/asprasad/treebeard.git
+else
+    echo "'treebeard' directory already exists. Skipping clone."
+fi
+
 cd treebeard
 
 # Checkout branch with changes made for silvanforge
 git checkout silvanforge
 
-mkdir build
-cd build
+mkdir -p $BUILD_DIR
+cd $BUILD_DIR
 bash ../scripts/gen_gpu.sh
 cmake --build .
 
 # Clone and build Tahoe
 cd ../../
-git clone https://github.com/sampathrg/Tahoe.git
+
+if [ ! -d "Tahoe" ]; then
+    git clone https://github.com/sampathrg/Tahoe.git
+else
+    echo "'Tahoe' directory already exists. Skipping clone."
+fi
+
 cd Tahoe
 git checkout tahoe-expts
 make
-
-

src/gpu/GPUSimtPass.cpp

@@ -160,7 +160,28 @@ void ConvertTraverseToSimtTraverse(mlir::MLIRContext &context,
                                    mlir::ModuleOp module) {
   // llvm::DebugFlag = true;
   // Lower from high-level IR to mid-level IR
+
+  // Check if the environment variable PRINT_AFTER_ALL is set
+  const char *printAfterAllEnv = std::getenv("PRINT_AFTER_ALL");
+  bool printAfterAll = printAfterAllEnv && std::string(printAfterAllEnv) == "true";
+
+  if(printAfterAll)
+    context.disableMultithreading();
+
+
   mlir::PassManager pm(&context);
+
+  // If PRINT_AFTER_ALL is set to "true", enable IR printing
+  if (printAfterAll) {
+    /* Enable Print After All For Debugging */
+    pm.enableIRPrinting(
+        [=](mlir::Pass *a, Operation *b) { return false; },  // Don't print before passes
+        [=](mlir::Pass *a, Operation *b) { return true; },   // Print after every pass
+        true,  // Print at module scope
+        false  // Print after every pass, regardless of changes
+    );
+  }
+
   pm.addPass(std::make_unique<ConvertTraverseToCooperativeTraverse>());
 
   if (mlir::failed(pm.run(module))) {

src/gpu/GPUSupportUtils.cpp

@@ -416,7 +416,28 @@ bool isThreadLoop(scf::ParallelOp parallelOp) {
 }
 
 void GreedilyMapParallelLoopsToGPU(mlir::ModuleOp module) {
-  mlir::PassManager pm(module.getContext());
+
+   // Check if the environment variable PRINT_AFTER_ALL is set
+  const char *printAfterAllEnv = std::getenv("PRINT_AFTER_ALL");
+  bool printAfterAll = printAfterAllEnv && std::string(printAfterAllEnv) == "true";
+
+  if(printAfterAll)
+    module.getContext()->disableMultithreading();
+
+
+    mlir::PassManager pm(module.getContext());
+
+  // If PRINT_AFTER_ALL is set to "true", enable IR printing
+  if (printAfterAll) {
+    /* Enable Print After All For Debugging */
+    pm.enableIRPrinting(
+        [=](mlir::Pass *a, Operation *b) { return false; },  // Don't print before passes
+        [=](mlir::Pass *a, Operation *b) { return true; },   // Print after every pass
+        true,  // Print at module scope
+        false  // Print after every pass, regardless of changes
+    );
+  }
+
   mlir::OpPassManager &optPM = pm.nest<mlir::func::FuncOp>();
   optPM.addPass(createGpuMapParallelLoopsPass());
 
@@ -427,7 +448,26 @@ void GreedilyMapParallelLoopsToGPU(mlir::ModuleOp module) {
 
 void ConvertParallelLoopsToGPU(mlir::MLIRContext &context,
                                mlir::ModuleOp module) {
+  // Check if the environment variable PRINT_AFTER_ALL is set
+  const char *printAfterAllEnv = std::getenv("PRINT_AFTER_ALL");
+  bool printAfterAll = printAfterAllEnv && std::string(printAfterAllEnv) == "true";
+
+  if(printAfterAll)
+    context.disableMultithreading();
+
+
   mlir::PassManager pm(&context);
+
+  // If PRINT_AFTER_ALL is set to "true", enable IR printing
+  if (printAfterAll) {
+    /* Enable Print After All For Debugging */
+    pm.enableIRPrinting(
+        [=](mlir::Pass *a, Operation *b) { return false; },  // Don't print before passes
+        [=](mlir::Pass *a, Operation *b) { return true; },   // Print after every pass
+        true,  // Print at module scope
+        false  // Print after every pass, regardless of changes
+    );
+  }
   mlir::OpPassManager &optPM = pm.nest<mlir::func::FuncOp>();
   // optPM.addPass(std::make_unique<MakeGPULoopsPerfectlyNestedPass>());
   optPM.addPass(createParallelLoopToGpuPass());
@@ -441,7 +481,26 @@ void ConvertParallelLoopsToGPU(mlir::MLIRContext &context,
 }
 
 void OutlineGPUKernels(mlir::MLIRContext &context, mlir::ModuleOp module) {
+  // Check if the environment variable PRINT_AFTER_ALL is set
+  const char *printAfterAllEnv = std::getenv("PRINT_AFTER_ALL");
+  bool printAfterAll = printAfterAllEnv && std::string(printAfterAllEnv) == "true";
+
+  if(printAfterAll)
+    context.disableMultithreading();
+
+
   mlir::PassManager pm(&context);
+
+  // If PRINT_AFTER_ALL is set to "true", enable IR printing
+  if (printAfterAll) {
+    /* Enable Print After All For Debugging */
+    pm.enableIRPrinting(
+        [=](mlir::Pass *a, Operation *b) { return false; },  // Don't print before passes
+        [=](mlir::Pass *a, Operation *b) { return true; },   // Print after every pass
+        true,  // Print at module scope
+        false  // Print after every pass, regardless of changes
+    );
+  }
   pm.addPass(createGpuKernelOutliningPass());
 
   if (mlir::failed(pm.run(module))) {
@@ -450,7 +509,26 @@ void OutlineGPUKernels(mlir::MLIRContext &context, mlir::ModuleOp module) {
 }
 
 void RunCanonicalizerPass(mlir::MLIRContext &context, mlir::ModuleOp module) {
+  // Check if the environment variable PRINT_AFTER_ALL is set
+  const char *printAfterAllEnv = std::getenv("PRINT_AFTER_ALL");
+  bool printAfterAll = printAfterAllEnv && std::string(printAfterAllEnv) == "true";
+
+  if(printAfterAll)
+    context.disableMultithreading();
+
+
   mlir::PassManager pm(&context);
+
+  // If PRINT_AFTER_ALL is set to "true", enable IR printing
+  if (printAfterAll) {
+    /* Enable Print After All For Debugging */
+    pm.enableIRPrinting(
+        [=](mlir::Pass *a, Operation *b) { return false; },  // Don't print before passes
+        [=](mlir::Pass *a, Operation *b) { return true; },   // Print after every pass
+        true,  // Print at module scope
+        false  // Print after every pass, regardless of changes
+    );
+  }
   mlir::GreedyRewriteConfig config;
   std::vector<std::string> disabledPatterns = {
       "(anonymous namespace)::MergeNestedParallelLoops"};

src/gpu/LowerGPUToLLVM.cpp

@@ -561,7 +561,27 @@ void LowerGPUToLLVM(
   InitializeGPUTarget(compileInfo);
   // llvm::DebugFlag = true;
   // Lower from high-level IR to mid-level IR
+
+  // Check if the environment variable PRINT_AFTER_ALL is set
+  const char *printAfterAllEnv = std::getenv("PRINT_AFTER_ALL");
+  bool printAfterAll = printAfterAllEnv && std::string(printAfterAllEnv) == "true";
+
+  if(printAfterAll)
+    context.disableMultithreading();
+
+
   mlir::PassManager pm(&context);
+
+  // If PRINT_AFTER_ALL is set to "true", enable IR printing
+  if (printAfterAll) {
+    /* Enable Print After All For Debugging */
+    pm.enableIRPrinting(
+        [=](mlir::Pass *a, Operation *b) { return false; },  // Don't print before passes
+        [=](mlir::Pass *a, Operation *b) { return true; },   // Print after every pass
+        true,  // Print at module scope
+        false  // Print after every pass, regardless of changes
+    );
+  }
   // pm.addPass(createConvertSCFToCFPass());
   pm.addPass(createGpuKernelOutliningPass());
   // pm.addPass(std::make_unique<PrintModulePass>());

src/main.cpp

@@ -9,6 +9,8 @@
 #include "json/xgboostparser.h"
 #include <iostream>
 #include <string>
+#include "llvm/Support/CommandLine.h"
+#include "llvm/Support/InitLLVM.h"
 
 namespace TreeBeard {
 namespace test {
@@ -17,6 +19,19 @@ void generateRandomXGBoostModels(const std::string &dirName);
 } // namespace test
 } // namespace TreeBeard
 
+// Commmand line options
+
+using namespace llvm;
+static cl::opt<bool> printAfterAll("print-treebeard-ir-after-all",
+    cl::desc("Print IR after each pass"), cl::init(false));
+
+static cl::opt<bool> individual("individual", cl::desc("Enable individual mode"), cl::init(false));
+
+static cl::opt<std::string> testName("testname", cl::desc("Test name"),
+    cl::value_desc("name"), cl::init(""), cl::Hidden);
+
+
+
 bool EqualsString(char *arg, const std::string &str) {
   return (std::string(arg) == str);
 }
@@ -387,6 +402,21 @@ bool GenerateRandomXGBoostBenchmarksIfNeeded(int argc, char *argv[]) {
 }
 
 int main(int argc, char *argv[]) {
+  cl::ParseCommandLineOptions(argc, argv, "TreeBread Runner");
+
+  // Check if 'printAfterAll' is used without 'individual'
+  if (printAfterAll && !individual) {
+    llvm::errs() << "Error: 'print-treebeard-ir-after-all' can only be used with "
+                    "'individual' flag.\n";
+    return 1;
+  }
+
+  // Check if 'testName' is used without 'individual'
+  if (!testName.empty() && !individual) {
+    llvm::errs() << "Error: 'testname' can only be used with 'individual' flag.\n";
+    return 1;
+  }
+
   SetInsertDebugHelpers(argc, argv);
   SetInsertPrintVectors(argc, argv);
   SetPerfNotificationListener(argc, argv);
@@ -406,7 +436,21 @@ int main(int argc, char *argv[]) {
     return 0;
   else if (GenerateRandomXGBoostBenchmarksIfNeeded(argc, argv))
     return 0;
-  else {
+  else if (individual) {
+
+    // If printAfterAll is true, set the environment variable
+    if (printAfterAll) {
+      setenv("PRINT_AFTER_ALL", "true", 1); // Set the environment variable
+    }
+
+    std::string individualTestName = testName;
+    TreeBeard::test::RunIndividualTests(individualTestName);
+
+    // Unset the environment variable after the tests are run
+    if (printAfterAll) {
+      unsetenv("PRINT_AFTER_ALL"); // Unset the environment variable
+    }
+  } else {
     std::cout
         << "TreeBeard: A compiler for gradient boosting tree inference.\n";
     TreeBeard::test::RunTests();

src/mlir/ConvertNodeTypeToIndexType.cpp

@@ -173,7 +173,27 @@ namespace decisionforest
 void ConvertNodeTypeToIndexType(mlir::MLIRContext& context, mlir::ModuleOp module) {
   // llvm::DebugFlag = true;
   // Lower from high-level IR to mid-level IR
+
+  // Check if the environment variable PRINT_AFTER_ALL is set
+  const char *printAfterAllEnv = std::getenv("PRINT_AFTER_ALL");
+  bool printAfterAll = printAfterAllEnv && std::string(printAfterAllEnv) == "true";
+
+  if(printAfterAll)
+    context.disableMultithreading();
+
+
   mlir::PassManager pm(&context);
+
+  // If PRINT_AFTER_ALL is set to "true", enable IR printing
+  if (printAfterAll) {
+    /* Enable Print After All For Debugging */
+    pm.enableIRPrinting(
+        [=](mlir::Pass *a, Operation *b) { return false; },  // Don't print before passes
+        [=](mlir::Pass *a, Operation *b) { return true; },   // Print after every pass
+        true,  // Print at module scope
+        false  // Print after every pass, regardless of changes
+    );
+  }
   pm.addPass(std::make_unique<ConvertNodeTypeToIndexTypePass>());
 
   if (mlir::failed(pm.run(module))) {