[WIP]: Enable MKL support for DISTAL

legion/CMakeLists.txt

@@ -7,9 +7,23 @@ endif(OPENMP)
 # a non-system provided HDF5 installation.
 find_package(HDF5 REQUIRED)
 
+# Allow for choosing the BLAS library for CPU builds.
+SET (DISTAL_CPU_BLAS "OpenBLAS" CACHE STRING "Choose the CPU BLAS library for DISTAL. Either OpenBLAS or MKL.")
+SET (DISTAL_BLAS_ROOT "OpenBLAS/install" CACHE STRING "Root of the BLAS installation for DISTAL.")
+if ( ${DISTAL_CPU_BLAS} STREQUAL "OpenBLAS")
+    SET (DISTAL_BLAS_LIBRARIES openblas)
+    add_compile_definitions(DISTAL_USE_OPENBLAS)
+elseif( ${DISTAL_CPU_BLAS} STREQUAL "MKL")
+    SET (DISTAL_BLAS_LIBRARIES mkl_intel_lp64 mkl_intel_thread mkl_core)
+    add_compile_definitions(DISTAL_USE_MKL)
+else()
+    message(FATAL_ERROR "DISTAL_CPU_BLAS must be either OpenBLAS or MKL")
+endif()
+include_directories("${DISTAL_BLAS_ROOT}/include")
+link_directories("${DISTAL_BLAS_ROOT}/lib")
+
 # Include the headers and supporting source files.
-include_directories(include/ OpenBLAS/install/include "${CMAKE_BINARY_DIR}/include" cub/)
-link_directories(OpenBLAS/install/lib)
+include_directories(include/ "${CMAKE_BINARY_DIR}/include" cub/)
 file(GLOB LG_SOURCES src/*.cpp)
 file(GLOB LG_CU_SOURCES src/*.cu)
 if (Legion_USE_CUDA)
@@ -62,7 +76,7 @@ endif()
 function(add_app_folder folder)
     add_executable("${folder}" "${folder}/main.cpp" "${folder}/taco-generated.cpp")
     target_include_directories("${folder}" PRIVATE "${folder}")
-    target_link_libraries("${folder}" distal-runtime-lib Legion::Legion openblas)
+    target_link_libraries("${folder}" distal-runtime-lib Legion::Legion ${DISTAL_BLAS_LIBRARIES})
     if (OPENMP)
       target_compile_options("${folder}" PRIVATE ${OpenMP_CXX_FLAGS})
     endif(OPENMP)

legion/cannonMM/taco-generated.cpp

@@ -5,7 +5,7 @@
 using namespace Legion;
 
 #include "taco-generated.h"
-#include "cblas.h"
+#include "distal_blas.h"
 typedef FieldAccessor<READ_ONLY,double,2,coord_t,Realm::AffineAccessor<double,2,coord_t>> AccessorROdouble2;
 typedef FieldAccessor<READ_WRITE,double,2,coord_t,Realm::AffineAccessor<double,2,coord_t>> AccessorRWdouble2;
 
@@ -433,10 +433,7 @@ void task_6(const Task* task, const std::vector<PhysicalRegion>& regions, Contex
   if (bDomain.get_volume() == 0 || cDomain.get_volume() == 0)
     return ;
 
-  cblas_dgemm(
-    CblasRowMajor,
-    CblasNoTrans,
-    CblasNoTrans,
+  BLAS_GEMM<double>(
     (1 + (bDomain.hi()[0] - bDomain.lo()[0])),
     (1 + (cDomain.hi()[1] - cDomain.lo()[1])),
     (1 + (cDomain.hi()[0] - cDomain.lo()[0])),

legion/include/distal_blas.h

@@ -0,0 +1,60 @@
+#ifndef DISTAL_BLAS_H
+#define DISTAL_BLAS_H
+
+// Provide an indirection around CPU BLAS operations to allow for swapping in
+// an out of different libraries. For convenience sake, we will omit the arguments
+// that control whether the inputs are transposed / row major format etc.
+template<typename T>
+void BLAS_GEMM(size_t m,
+               size_t n,
+               size_t k,
+               const T alpha,
+               const T* A,
+               const int lda,
+               const T* B,
+               const int ldb,
+               const T beta,
+               T* C,
+               const int ldc);
+
+
+// TODO (rohany): Gaurd this behind a set of defines so that we can swap the BLAS vendor that we are using. Ideally,
+//   all that will change is that we just define including a different header set.
+
+#ifdef DISTAL_USE_OPENBLAS
+#include "cblas.h"
+#elif DISTAL_USE_MKL
+#include "mkl.h"
+#endif
+
+template<>
+void BLAS_GEMM<double>(size_t m,
+                       size_t n,
+                       size_t k,
+                       const double alpha,
+                       const double* A,
+                       const int lda,
+                       const double* B,
+                       const int ldb,
+                       const double beta,
+                       double* C,
+                       const int ldc) {
+  cblas_dgemm(
+    CblasRowMajor,
+    CblasNoTrans,
+    CblasNoTrans,
+    m,
+    n,
+    k,
+    alpha,
+    A,
+    lda,
+    B,
+    ldb,
+    beta,
+    C,
+    ldc
+  );
+}
+
+#endif // DISTAL_BLAS_H

legion/include/leaf_kernels.h

@@ -7,9 +7,9 @@
 #include <iostream>
 #include <cstdlib>
 #include <cassert>
-#include "cblas.h"
 #include "legion.h"
 #include "realm_defines.h"
+#include "distal_blas.h"
 
 // An argument pack for MTTKRP.
 struct MTTKRPPack {
@@ -60,10 +60,7 @@ void mttkrp(MTTKRPPack pack, T* A_vals, const T* B_vals, const T* C_vals, const
 
   // Perform T(i, j, l) = B(i, j, k) * D(k, l) as a series of GEMM calls.
   for (size_t i = 0; i < B1_dimension; i++) {
-    cblas_dgemm(
-      CblasRowMajor,
-      CblasNoTrans,
-      CblasNoTrans,
+    BLAS_GEMM<T>(
       C1_dimension,
       D2_dimension,
       D1_dimension,

legion/mttkrp/taco-generated.cpp

@@ -70,8 +70,8 @@ Legion::LogicalPartition partitionLegionA(Legion::Context ctx, Legion::Runtime*
     }
     AColoring[(*itr)] = ARect;
   }
-  auto A_dense_run_0_Partition = runtime->create_index_partition(ctx, A_dense_run_0, domain, AColoring, LEGION_DISJOINT_COMPLETE_KIND);
-  auto A_vals_partition = copyPartition(ctx, runtime, A_dense_run_0_Partition, get_logical_region(A_vals));
+  IndexPartition A_dense_run_0_Partition_0 = runtime->create_index_partition(ctx, A_dense_run_0, domain, AColoring, LEGION_DISJOINT_COMPLETE_KIND);
+  auto A_vals_partition = copyPartition(ctx, runtime, A_dense_run_0_Partition_0, get_logical_region(A_vals));
   return A_vals_partition;
 }
 
@@ -100,8 +100,8 @@ Legion::LogicalPartition partitionLegionB(Legion::Context ctx, Legion::Runtime*
     }
     BColoring[(*itr)] = BRect;
   }
-  auto B_dense_run_0_Partition = runtime->create_index_partition(ctx, B_dense_run_0, domain, BColoring, LEGION_DISJOINT_COMPLETE_KIND);
-  auto B_vals_partition = copyPartition(ctx, runtime, B_dense_run_0_Partition, get_logical_region(B_vals));
+  IndexPartition B_dense_run_0_Partition_0 = runtime->create_index_partition(ctx, B_dense_run_0, domain, BColoring, LEGION_DISJOINT_COMPLETE_KIND);
+  auto B_vals_partition = copyPartition(ctx, runtime, B_dense_run_0_Partition_0, get_logical_region(B_vals));
   return B_vals_partition;
 }
 
@@ -127,8 +127,8 @@ Legion::LogicalPartition partitionLegionC(Legion::Context ctx, Legion::Runtime*
     }
     CColoring[(*itr)] = CRect;
   }
-  auto C_dense_run_0_Partition = runtime->create_index_partition(ctx, C_dense_run_0, domain, CColoring, LEGION_DISJOINT_COMPLETE_KIND);
-  auto C_vals_partition = copyPartition(ctx, runtime, C_dense_run_0_Partition, get_logical_region(C_vals));
+  IndexPartition C_dense_run_0_Partition_0 = runtime->create_index_partition(ctx, C_dense_run_0, domain, CColoring, LEGION_DISJOINT_COMPLETE_KIND);
+  auto C_vals_partition = copyPartition(ctx, runtime, C_dense_run_0_Partition_0, get_logical_region(C_vals));
   return C_vals_partition;
 }
 
@@ -154,8 +154,8 @@ Legion::LogicalPartition partitionLegionD(Legion::Context ctx, Legion::Runtime*
     }
     DColoring[(*itr)] = DRect;
   }
-  auto D_dense_run_0_Partition = runtime->create_index_partition(ctx, D_dense_run_0, domain, DColoring, LEGION_DISJOINT_COMPLETE_KIND);
-  auto D_vals_partition = copyPartition(ctx, runtime, D_dense_run_0_Partition, get_logical_region(D_vals));
+  IndexPartition D_dense_run_0_Partition_0 = runtime->create_index_partition(ctx, D_dense_run_0, domain, DColoring, LEGION_DISJOINT_COMPLETE_KIND);
+  auto D_vals_partition = copyPartition(ctx, runtime, D_dense_run_0_Partition_0, get_logical_region(D_vals));
   return D_vals_partition;
 }
 
@@ -183,12 +183,12 @@ partitionPackForplaceLegionA partitionForplaceLegionA(Legion::Context ctx, Legio
     }
     AColoring[(*itr)] = ARect;
   }
-  auto A_dense_run_0_Partition = runtime->create_index_partition(ctx, A_dense_run_0, domain, AColoring, LEGION_COMPUTE_KIND);
-  auto A_vals_partition = copyPartition(ctx, runtime, A_dense_run_0_Partition, get_logical_region(A_vals));
+  IndexPartition A_dense_run_0_Partition_0 = runtime->create_index_partition(ctx, A_dense_run_0, domain, AColoring, LEGION_COMPUTE_KIND);
+  auto A_vals_partition = copyPartition(ctx, runtime, A_dense_run_0_Partition_0, get_logical_region(A_vals));
   computePartitions.APartition.indicesPartitions = std::vector<std::vector<Legion::LogicalPartition>>(2);
   computePartitions.APartition.denseLevelRunPartitions = std::vector<IndexPartition>(2);
   computePartitions.APartition.valsPartition = A_vals_partition;
-  computePartitions.APartition.denseLevelRunPartitions[0] = A_dense_run_0_Partition;
+  computePartitions.APartition.denseLevelRunPartitions[0] = A_dense_run_0_Partition_0;
 
   return computePartitions;
 }
@@ -260,12 +260,12 @@ partitionPackForplaceLegionB partitionForplaceLegionB(Legion::Context ctx, Legio
     }
     BColoring[(*itr)] = BRect;
   }
-  auto B_dense_run_0_Partition = runtime->create_index_partition(ctx, B_dense_run_0, domain, BColoring, LEGION_COMPUTE_KIND);
-  auto B_vals_partition = copyPartition(ctx, runtime, B_dense_run_0_Partition, get_logical_region(B_vals));
+  IndexPartition B_dense_run_0_Partition_0 = runtime->create_index_partition(ctx, B_dense_run_0, domain, BColoring, LEGION_COMPUTE_KIND);
+  auto B_vals_partition = copyPartition(ctx, runtime, B_dense_run_0_Partition_0, get_logical_region(B_vals));
   computePartitions.BPartition.indicesPartitions = std::vector<std::vector<Legion::LogicalPartition>>(3);
   computePartitions.BPartition.denseLevelRunPartitions = std::vector<IndexPartition>(3);
   computePartitions.BPartition.valsPartition = B_vals_partition;
-  computePartitions.BPartition.denseLevelRunPartitions[0] = B_dense_run_0_Partition;
+  computePartitions.BPartition.denseLevelRunPartitions[0] = B_dense_run_0_Partition_0;
 
   return computePartitions;
 }
@@ -327,12 +327,12 @@ partitionPackForplaceLegionC partitionForplaceLegionC(Legion::Context ctx, Legio
     }
     CColoring[(*itr)] = CRect;
   }
-  auto C_dense_run_0_Partition = runtime->create_index_partition(ctx, C_dense_run_0, domain, CColoring, LEGION_COMPUTE_KIND);
-  auto C_vals_partition = copyPartition(ctx, runtime, C_dense_run_0_Partition, get_logical_region(C_vals));
+  IndexPartition C_dense_run_0_Partition_0 = runtime->create_index_partition(ctx, C_dense_run_0, domain, CColoring, LEGION_COMPUTE_KIND);
+  auto C_vals_partition = copyPartition(ctx, runtime, C_dense_run_0_Partition_0, get_logical_region(C_vals));
   computePartitions.CPartition.indicesPartitions = std::vector<std::vector<Legion::LogicalPartition>>(2);
   computePartitions.CPartition.denseLevelRunPartitions = std::vector<IndexPartition>(2);
   computePartitions.CPartition.valsPartition = C_vals_partition;
-  computePartitions.CPartition.denseLevelRunPartitions[0] = C_dense_run_0_Partition;
+  computePartitions.CPartition.denseLevelRunPartitions[0] = C_dense_run_0_Partition_0;
 
   return computePartitions;
 }
@@ -401,12 +401,12 @@ partitionPackForplaceLegionD partitionForplaceLegionD(Legion::Context ctx, Legio
     }
     DColoring[(*itr)] = DRect;
   }
-  auto D_dense_run_0_Partition = runtime->create_index_partition(ctx, D_dense_run_0, domain, DColoring, LEGION_COMPUTE_KIND);
-  auto D_vals_partition = copyPartition(ctx, runtime, D_dense_run_0_Partition, get_logical_region(D_vals));
+  IndexPartition D_dense_run_0_Partition_0 = runtime->create_index_partition(ctx, D_dense_run_0, domain, DColoring, LEGION_COMPUTE_KIND);
+  auto D_vals_partition = copyPartition(ctx, runtime, D_dense_run_0_Partition_0, get_logical_region(D_vals));
   computePartitions.DPartition.indicesPartitions = std::vector<std::vector<Legion::LogicalPartition>>(2);
   computePartitions.DPartition.denseLevelRunPartitions = std::vector<IndexPartition>(2);
   computePartitions.DPartition.valsPartition = D_vals_partition;
-  computePartitions.DPartition.denseLevelRunPartitions[0] = D_dense_run_0_Partition;
+  computePartitions.DPartition.denseLevelRunPartitions[0] = D_dense_run_0_Partition_0;
 
   return computePartitions;
 }
@@ -512,30 +512,30 @@ partitionPackForcomputeLegion partitionForcomputeLegion(Legion::Context ctx, Leg
     }
     DColoring[(*itr)] = DRect;
   }
-  auto A_dense_run_0_Partition = runtime->create_index_partition(ctx, A_dense_run_0, domain, AColoring, LEGION_COMPUTE_KIND);
-  auto A_vals_partition = copyPartition(ctx, runtime, A_dense_run_0_Partition, get_logical_region(A_vals));
-  auto B_dense_run_0_Partition = runtime->create_index_partition(ctx, B_dense_run_0, domain, BColoring, LEGION_COMPUTE_KIND);
-  auto B_vals_partition = copyPartition(ctx, runtime, B_dense_run_0_Partition, get_logical_region(B_vals));
-  auto C_dense_run_0_Partition = runtime->create_index_partition(ctx, C_dense_run_0, domain, CColoring, LEGION_COMPUTE_KIND);
-  auto C_vals_partition = copyPartition(ctx, runtime, C_dense_run_0_Partition, get_logical_region(C_vals));
-  auto D_dense_run_0_Partition = runtime->create_index_partition(ctx, D_dense_run_0, domain, DColoring, LEGION_COMPUTE_KIND);
-  auto D_vals_partition = copyPartition(ctx, runtime, D_dense_run_0_Partition, get_logical_region(D_vals));
+  IndexPartition A_dense_run_0_Partition_0 = runtime->create_index_partition(ctx, A_dense_run_0, domain, AColoring, LEGION_COMPUTE_KIND);
+  auto A_vals_partition = copyPartition(ctx, runtime, A_dense_run_0_Partition_0, get_logical_region(A_vals));
+  IndexPartition B_dense_run_0_Partition_0 = runtime->create_index_partition(ctx, B_dense_run_0, domain, BColoring, LEGION_COMPUTE_KIND);
+  auto B_vals_partition = copyPartition(ctx, runtime, B_dense_run_0_Partition_0, get_logical_region(B_vals));
+  IndexPartition C_dense_run_0_Partition_0 = runtime->create_index_partition(ctx, C_dense_run_0, domain, CColoring, LEGION_COMPUTE_KIND);
+  auto C_vals_partition = copyPartition(ctx, runtime, C_dense_run_0_Partition_0, get_logical_region(C_vals));
+  IndexPartition D_dense_run_0_Partition_0 = runtime->create_index_partition(ctx, D_dense_run_0, domain, DColoring, LEGION_COMPUTE_KIND);
+  auto D_vals_partition = copyPartition(ctx, runtime, D_dense_run_0_Partition_0, get_logical_region(D_vals));
   computePartitions.APartition.indicesPartitions = std::vector<std::vector<Legion::LogicalPartition>>(2);
   computePartitions.APartition.denseLevelRunPartitions = std::vector<IndexPartition>(2);
   computePartitions.APartition.valsPartition = A_vals_partition;
-  computePartitions.APartition.denseLevelRunPartitions[0] = A_dense_run_0_Partition;
+  computePartitions.APartition.denseLevelRunPartitions[0] = A_dense_run_0_Partition_0;
   computePartitions.BPartition.indicesPartitions = std::vector<std::vector<Legion::LogicalPartition>>(3);
   computePartitions.BPartition.denseLevelRunPartitions = std::vector<IndexPartition>(3);
   computePartitions.BPartition.valsPartition = B_vals_partition;
-  computePartitions.BPartition.denseLevelRunPartitions[0] = B_dense_run_0_Partition;
+  computePartitions.BPartition.denseLevelRunPartitions[0] = B_dense_run_0_Partition_0;
   computePartitions.CPartition.indicesPartitions = std::vector<std::vector<Legion::LogicalPartition>>(2);
   computePartitions.CPartition.denseLevelRunPartitions = std::vector<IndexPartition>(2);
   computePartitions.CPartition.valsPartition = C_vals_partition;
-  computePartitions.CPartition.denseLevelRunPartitions[0] = C_dense_run_0_Partition;
+  computePartitions.CPartition.denseLevelRunPartitions[0] = C_dense_run_0_Partition_0;
   computePartitions.DPartition.indicesPartitions = std::vector<std::vector<Legion::LogicalPartition>>(2);
   computePartitions.DPartition.denseLevelRunPartitions = std::vector<IndexPartition>(2);
   computePartitions.DPartition.valsPartition = D_vals_partition;
-  computePartitions.DPartition.denseLevelRunPartitions[0] = D_dense_run_0_Partition;
+  computePartitions.DPartition.denseLevelRunPartitions[0] = D_dense_run_0_Partition_0;
 
   return computePartitions;
 }

src/codegen/codegen_legion_c.cpp

@@ -281,7 +281,7 @@ void CodegenLegionC::emitHeaders(std::ostream &o) {
 
   struct BLASFinder : public IRVisitor {
     void visit(const Call* node) {
-      if (node->func.find("blas") != std::string::npos) {
+      if (node->func.find("BLAS") != std::string::npos) {
         this->usesBLAS = true;
       }
       if (node->func.find("mttkrp") != std::string::npos || node->func.find("ttv") != std::string::npos) {
@@ -294,7 +294,7 @@ void CodegenLegionC::emitHeaders(std::ostream &o) {
   BLASFinder bs;
   this->stmt.accept(&bs);
   if (bs.usesBLAS) {
-    o << "#include \"cblas.h\"\n";
+    o << "#include \"distal_blas.h\"\n";
   }
   if (bs.usesLeafKernels) {
     o << "#include \"leaf_kernels.h\"\n";

src/index_notation/transformations.cpp

@@ -1409,8 +1409,6 @@ ir::Stmt GEMM::replaceValidStmt(IndexStmt stmt,
 
   std::vector<ir::Stmt> results;
   auto ctx = ir::Symbol::make("ctx");
-  auto rowMajor = ir::Symbol::make("CblasRowMajor");
-  auto noTrans = ir::Symbol::make("CblasNoTrans");
 
   std::vector<ir::Expr> tvars;
   for (auto var : this->content->tensorVars) {
@@ -1466,9 +1464,6 @@ ir::Stmt GEMM::replaceValidStmt(IndexStmt stmt,
   );
 
   std::vector<ir::Expr> args = {
-      rowMajor,
-      noTrans,
-      noTrans,
       ir::Add::make(1, ir::Sub::make(ir::Load::make(getBounds(bDomain, "hi"), 0), ir::Load::make(getBounds(bDomain, "lo"), 0))),
       ir::Add::make(1, ir::Sub::make(ir::Load::make(getBounds(cDomain, "hi"), 1), ir::Load::make(getBounds(cDomain, "lo"), 1))),
       ir::Add::make(1, ir::Sub::make(ir::Load::make(getBounds(cDomain, "hi"), 0), ir::Load::make(getBounds(cDomain, "lo"), 0))),
@@ -1482,8 +1477,8 @@ ir::Stmt GEMM::replaceValidStmt(IndexStmt stmt,
       ldA,
   };
 
-  // TODO (rohany): Pick the right call between double and float here.
-  results.push_back(ir::SideEffect::make(ir::Call::make("cblas_dgemm", args, Auto)));
+  auto funcName = "BLAS_GEMM<" + util::toString(type.getDataType()) + ">";
+  results.push_back(ir::SideEffect::make(ir::Call::make(funcName, args, Auto)));
   return ir::Block::make(results);
 }
 
@@ -1864,9 +1859,6 @@ ir::Stmt TTMC::replaceValidStmt(IndexStmt stmt, ProvenanceGraph pg, std::map<Ten
   };
 
   std::vector<ir::Expr> args {
-    rowMajor,
-    noTrans,
-    noTrans,
     ir::Add::make(1, ir::Sub::make(ir::Load::make(getBounds(bDomain, "hi"), 1), ir::Load::make(getBounds(bDomain, "lo"), 1))),
     ir::Add::make(1, ir::Sub::make(ir::Load::make(getBounds(cDomain, "hi"), 1), ir::Load::make(getBounds(cDomain, "lo"), 1))),
     ir::Add::make(1, ir::Sub::make(ir::Load::make(getBounds(cDomain, "hi"), 0), ir::Load::make(getBounds(cDomain, "lo"), 0))),
@@ -1879,7 +1871,8 @@ ir::Stmt TTMC::replaceValidStmt(IndexStmt stmt, ProvenanceGraph pg, std::map<Ten
     getBasePtr(aAccess, aDomain),
     getLD(aAccess, 1),
   };
-  loopStmts.push_back(ir::SideEffect::make(ir::Call::make("cblas_dgemm", args, Auto)));
+  auto funcName = "BLAS_GEMM<" + util::toString(type.getDataType()) + ">";
+  loopStmts.push_back(ir::SideEffect::make(ir::Call::make(funcName, args, Auto)));
 
   // The loop bounds are from 0 to aBounds.hi()[0] - aBounds.lo()[0].
   auto bounds = ir::Add::make(1, ir::Sub::make(ir::Load::make(getBounds(aDomain, "hi"), 0), ir::Load::make(getBounds(aDomain, "lo"), 0)));