diff --git a/legion/CMakeLists.txt b/legion/CMakeLists.txt index 545e9b934..24a1ddb70 100644 --- a/legion/CMakeLists.txt +++ b/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) diff --git a/legion/cannonMM/taco-generated.cpp b/legion/cannonMM/taco-generated.cpp index 712954be2..4e1015e8e 100644 --- a/legion/cannonMM/taco-generated.cpp +++ b/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> AccessorROdouble2; typedef FieldAccessor> AccessorRWdouble2; @@ -433,10 +433,7 @@ void task_6(const Task* task, const std::vector& regions, Contex if (bDomain.get_volume() == 0 || cDomain.get_volume() == 0) return ; - cblas_dgemm( - CblasRowMajor, - CblasNoTrans, - CblasNoTrans, + BLAS_GEMM( (1 + (bDomain.hi()[0] - bDomain.lo()[0])), (1 + (cDomain.hi()[1] - cDomain.lo()[1])), (1 + (cDomain.hi()[0] - cDomain.lo()[0])), diff --git a/legion/include/distal_blas.h b/legion/include/distal_blas.h new file mode 100644 index 000000000..eafeba243 --- /dev/null +++ b/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 +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(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 \ No newline at end of file diff --git a/legion/include/leaf_kernels.h b/legion/include/leaf_kernels.h index 0ec4dd4db..a9c2972fc 100644 --- a/legion/include/leaf_kernels.h +++ b/legion/include/leaf_kernels.h @@ -7,9 +7,9 @@ #include #include #include -#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( C1_dimension, D2_dimension, D1_dimension, diff --git a/legion/mttkrp/taco-generated.cpp b/legion/mttkrp/taco-generated.cpp index 492072ae4..2806b47d0 100644 --- a/legion/mttkrp/taco-generated.cpp +++ b/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>(2); computePartitions.APartition.denseLevelRunPartitions = std::vector(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>(3); computePartitions.BPartition.denseLevelRunPartitions = std::vector(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>(2); computePartitions.CPartition.denseLevelRunPartitions = std::vector(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>(2); computePartitions.DPartition.denseLevelRunPartitions = std::vector(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>(2); computePartitions.APartition.denseLevelRunPartitions = std::vector(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>(3); computePartitions.BPartition.denseLevelRunPartitions = std::vector(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>(2); computePartitions.CPartition.denseLevelRunPartitions = std::vector(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>(2); computePartitions.DPartition.denseLevelRunPartitions = std::vector(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; } diff --git a/src/codegen/codegen_legion_c.cpp b/src/codegen/codegen_legion_c.cpp index 5df6c23f9..464918f89 100644 --- a/src/codegen/codegen_legion_c.cpp +++ b/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"; diff --git a/src/index_notation/transformations.cpp b/src/index_notation/transformations.cpp index 18f8f760e..e3b48682a 100644 --- a/src/index_notation/transformations.cpp +++ b/src/index_notation/transformations.cpp @@ -1409,8 +1409,6 @@ ir::Stmt GEMM::replaceValidStmt(IndexStmt stmt, std::vector results; auto ctx = ir::Symbol::make("ctx"); - auto rowMajor = ir::Symbol::make("CblasRowMajor"); - auto noTrans = ir::Symbol::make("CblasNoTrans"); std::vector tvars; for (auto var : this->content->tensorVars) { @@ -1466,9 +1464,6 @@ ir::Stmt GEMM::replaceValidStmt(IndexStmt stmt, ); std::vector 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 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"; + 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)));