fix var updates, fix test ci

EnzymeAD · Oct 20, 2023 · 30073cc · 30073cc
1 parent 036f3f6
commit 30073cc
Show file tree

Hide file tree

Showing 6 changed files with 53 additions and 8 deletions.
diff --git a/enzyme/Enzyme/Clang/EnzymePassLoader.cpp b/enzyme/Enzyme/Clang/EnzymePassLoader.cpp
@@ -77,7 +77,8 @@ static RegisterStandardPasses
     clangtoolLoader_OEarly(PassManagerBuilder::EP_EarlyAsPossible,
                            loadNVVMPass);
 static RegisterStandardPasses
-    clangtoolLoader_Ox(PassManagerBuilder::EP_VectorizerStart, loadBLASPass);
+    clangtoolBlasLoader_Ox(PassManagerBuilder::EP_VectorizerStart,
+                           loadBLASPass);
 
 static void loadLTOPass(const PassManagerBuilder &Builder,
                         legacy::PassManagerBase &PM) {

diff --git a/enzyme/Enzyme/OptBlas.cpp b/enzyme/Enzyme/OptBlas.cpp
@@ -99,8 +99,11 @@ bool cmp_or_set(llvm::CallInst *CI, std::vector<llvm::Value *> values) {
     if (values[i] == nullptr) {
       continue;
     }
-    if (CI->getArgOperand(i) != values[i])
+    if (CI->getArgOperand(i) != values[i]) {
+      llvm::errs() << "cmp_or_set: " << i << " " << *CI->getArgOperand(i)
+                   << " != " << *values[i] << "\n";
       return false;
+    }
   }
   for (size_t i = 0; i < values.size(); ++i) {
     if (values[i] == nullptr) {

diff --git a/enzyme/Enzyme/Utils.cpp b/enzyme/Enzyme/Utils.cpp
@@ -2354,7 +2354,7 @@ std::optional<BlasInfo> extractBLAS(llvm::StringRef in)
 llvm::Optional<BlasInfo> extractBLAS(llvm::StringRef in)
 #endif
 {
-  const char *extractable[] = {"dot", "scal", "axpy", "gemv", "gemm", "spmv"};
+  const char *extractable[] = {"dot", "scal", "axpy", "gemv", "gemm", "spmv", "ger"};
   const char *floatType[] = {"s", "d"}; // c, z
   const char *prefixes[] = {"" /*Fortran*/, "cblas_"};
   const char *suffixes[] = {"", "_", "64_", "_64_"};

diff --git a/enzyme/test/Enzyme/BlasOpt/CMakeLists.txt b/enzyme/test/Enzyme/BlasOpt/CMakeLists.txt
@@ -0,0 +1,12 @@
+# Run regression and unit tests
+add_lit_testsuite(check-enzyme-blasopt "Running enzyme blas optimization regression tests"
+    ${CMAKE_CURRENT_BINARY_DIR}
+    DEPENDS ${ENZYME_TEST_DEPS}
+    ARGS -v
+)
+
+set_target_properties(check-enzyme-blasopt PROPERTIES FOLDER "Tests")
+
+# add_lit_testsuites(ENZYME ${CMAKE_CURRENT_SOURCE_DIR}
+# DEPENDS ${ENZYME_TEST_DEPS}
+# )
diff --git a/enzyme/test/Enzyme/CMakeLists.txt b/enzyme/test/Enzyme/CMakeLists.txt
@@ -6,6 +6,7 @@ add_subdirectory(ForwardModeSplit)
 add_subdirectory(ForwardModeVector)
 add_subdirectory(BatchMode)
 add_subdirectory(ProbProg)
+add_subdirectory(BlasOpt)
 
 # Run regression and unit tests
 add_lit_testsuite(check-enzyme "Running enzyme regression tests"

diff --git a/enzyme/tools/enzyme-tblgen/opts.cpp b/enzyme/tools/enzyme-tblgen/opts.cpp
@@ -100,13 +100,23 @@ void emitBlasOpt(StringRef name, std::vector<DagInit *> inputs,
 
   os << "  // create a vector of calls to delete\n";
   os << "  std::vector<CallInst *> todelete;\n";
+  os << "  int num_calls = 0;\n";
 
   os << "  for (auto &BB : *F) {\n"
      << "    for (auto &I : BB) {\n"
      << "      if (auto *CI = dyn_cast<CallInst>(&I)) {\n"
-     << "        auto name = CI->getCalledFunction()->getName();\n";
+     << "        num_calls++;\n"
+     << "        auto CIname = CI->getCalledFunction()->getName();\n"
+     << "        auto blasOption = extractBLAS(CIname);\n"
+     << "#if LLVM_VERSION_MAJOR >= 16\n"
+     << "        if (!blasOption.has_value()) continue;\n"
+     << "        auto blas = blasOption.value();\n"
+     << "#else\n"
+     << "        if (!blasOption.hasValue()) continue;\n"
+     << "        auto blas = blasOption.getValue();\n"
+     << "#endif\n";
   for (auto fnc : unique_functions.keys()) {
-    os << "        if (name == \"" << fnc << "\") {\n";
+    os << "        if (blas.function == \"" << fnc << "\") {\n";
     std::string tab = "          ";
     auto fnc_vec = unique_functions[fnc];
     bool multiple = fnc_vec.size() > 1;
@@ -127,12 +137,28 @@ void emitBlasOpt(StringRef name, std::vector<DagInit *> inputs,
       os << "};\n";
     }
     os << tab << "bool set = cmp_or_set(CI, values);\n";
-    os << tab << "if (!set) continue;\n";
+    os << tab << "if (!set) {\n";
+    os << tab << "  llvm::errs() << \"args missmatch: " << fnc << "\";\n";
+    os << tab << "  continue;\n";
+    os << tab << "}\n";
+    for (size_t i = 0; i < fnc_vec[0]->getNumArgs(); i++) {
+      os << tab << "values[" << i << "] = CI->getArgOperand(" << i << ");\n";
+    }
+    for (size_t i = 0; i < fnc_vec.size(); ++i) {
+      ArrayRef<StringInit *> args = fnc_vec[i]->getArgNames();
+      size_t pos = 0;
+      for (auto arg : args) {
+        os << tab << arg->getValue() << " = CI->getArgOperand(" << i << ");\n";
+        pos++;
+      }
+    }
     os << tab << "llvm::errs() << \"found " << fnc << "\\n\";\n";
     os << tab << "idx_" << fnc << "++;\n"
        << tab << "todelete.push_back(CI);\n"
+       << tab << "continue;\n"
        << "        }\n";
   }
+  os << "        llvm::errs() << \"unhandled: \" << blas.function << \"\\n\";\n";
   os << "      }\n";
   os << "    }\n";
   os << "  }\n";
@@ -143,8 +169,10 @@ void emitBlasOpt(StringRef name, std::vector<DagInit *> inputs,
     os << "  if (idx_" << fnc << " != " << unique_functions[fnc].size() << ")\n"
        << "    found = false;\n";
   }
-  os << "  if (!found)\n"
-     << "    return false;\n";
+  os << "  if (!found) {\n"
+     << "    llvm::errs() << \"num calls: \" << num_calls << \"\\n\";\n"
+     << "    return false;\n"
+     << "  }\n";
 
   os << "  llvm::errs() << \"found optimization " << name << "\\n\";\n";