8000 Allow setting target cpu and vector width from command line · srcarroll/mmperf@9cd38b0 · GitHub
[go: up one dir, main page]
Skip to content

Commit 9cd38b0

Browse files
author
harsh-nod
committed
Allow setting target cpu and vector width from command line
1 parent c7418ae commit 9cd38b0

File tree

3 files changed

+27
-6
lines changed

3 files changed

+27
-6
lines changed

‎Codegen/CMakeLists.txt

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -37,6 +37,9 @@ option(USE_OPENBLAS "Enable OpenBLAS" OFF)
3737
option(USE_HALIDE "Enable Halide" OFF)
3838

3939
set(SIZE_FILE ${CMAKE_CURRENT_LIST_DIR}/benchmark_sizes.txt CACHE INTERNAL "File containing matrix sizes to be benchmarked")
40+
set(TARGET_CPU "skylake-avx512" CACHE INTERNAL "Target CPU for MLIR")
41+
set(VECTOR_WIDTH "512" CACHE INTERNAL "Vector width for MLIR")
42+
4043
ExternalProject_Add(matmul
4144
DEPENDS mlir
4245
PREFIX ${CMAKE_BINARY_DIR}/matmul
@@ -55,9 +58,14 @@ ExternalProject_Add(matmul
5558
-DUSE_OPENBLAS=${USE_OPENBLAS}
5659
-DUSE_HALIDE=${USE_HALIDE}
5760
-DHALIDE_DIR=${HALIDE_DIR}
61+
-DTARGET_CPU=${TARGET_CPU}
62+
-DVECTOR_WIDTH=${VECTOR_WIDTH}
5863
)
5964

6065
unset(USE_MLIR CACHE)
6166
unset(USE_MKL CACHE)
6267
unset(USE_OPENBLAS CACHE)
6368
unset(USE_HALIDE CACHE)
69+
unset(SIZE_FILE CACHE)
70+
unset(TARGET_CPU CACHE)
71+
unset(VECTOR_WIDTH CACHE)

‎Codegen/matmul/matmul-compile/CMakeLists.txt

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -31,3 +31,5 @@ add_llvm_executable(matmul-compile matmul-compile.cpp)
3131
llvm_update_compile_flags(matmul-compile)
3232
target_link_libraries(matmul-compile PRIVATE ${LIBS})
3333
mlir_check_all_link_libraries(matmul-compile)
34+
target_compile_definitions(matmul-compile PRIVATE TARGET_CPU=${TARGET_CPU})
35+
target_compile_definitions(matmul-compile PRIVATE VECTOR_WIDTH=${VECTOR_WIDTH})

‎Codegen/matmul/matmul-compile/matmul-compile.cpp

Lines changed: 17 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -32,6 +32,9 @@
3232
#include <string>
3333
#include <unistd.h>
3434

35+
#define STRING(s) #s
36+
#define TO_STRING(x) STRING(x)
37+
3538
using namespace mlir;
3639
using namespace mlir::linalg;
3740
using llvm::Error;
@@ -53,26 +56,31 @@ struct LinalgCodegenPass : public PassWrapper<LinalgCodegenPass, FunctionPass> {
5356
registry.insert<linalg::LinalgDialect, AffineDialect, scf::SCFDialect>();
5457
}
5558
LinalgCodegenPass() = default;
56-
LinalgCodegenPass(int M, int N, int K) : M(M), N(N), K(K) {}
59+
LinalgCodegenPass(int M, int N, int K, const std::string &target_cpu,
60+
const std::string &vector_width) : M(M), N(N), K(K),
61+
target_cpu(target_cpu), vector_width(vector_width) {}
5762
LinalgCodegenPass(const LinalgCodegenPass &pass) {
5863
M = pass.M;
5964
N = pass.N;
6065
K = pass.K;
66+
target_cpu = pass.target_cpu;
67+
vector_width = pass.vector_width;
6168
}
6269
void runOnFunction() override;
6370

6471
int M, N, K;
72+
std::string target_cpu, vector_width;
6573
};
6674
} // namespace
6775

6876
void LinalgCodegenPass::runOnFunction() {
6977
MLIRContext *ctx = getFunction().getContext();
7078
SmallVector<Attribute, 4> attrs;
7179
attrs.push_back(ArrayAttr::get({StringAttr::get("prefer-vector-width", ctx),
72-
StringAttr::get("512", ctx)},
80+
StringAttr::get(vector_width, ctx)},
7381
ctx));
7482
attrs.push_back(ArrayAttr::get({StringAttr::get("target-cpu", ctx),
75-
StringAttr::get("skylake-avx512", ctx)},
83+
StringAttr::get(target_cpu, ctx)},
7684
ctx));
7785
getFunction()->setAttr("passthrough", ArrayAttr::get(attrs, ctx));
7886

@@ -181,8 +189,9 @@ void LinalgCodegenPass::runOnFunction() {
181189
//getFunction().dump();
182190
}
183191

184-
std::unique_ptr<OperationPass<FuncOp>> createLinalgCodegenPass(int M, int N, int K) {
185-
return std::make_unique<LinalgCodegenPass>(M, N, K);
192+
std::unique_ptr<OperationPass<FuncOp>> createLinalgCodegenPass(int M, int N, int K,
193+
const std::string &target_cpu, const std::string &vector_width) {
194+
return std::make_unique<LinalgCodegenPass>(M, N, K, target_cpu, vector_width);
186195
}
187196

188197
}
@@ -237,7 +246,9 @@ Error compile(Options &options, mlir::DialectRegistry &registry) {
237246
int M, N, K;
238247
get_dimensions(options.inputFile, M, N, K);
239248
pm.addPass(createCanonicalizerPass());
240-
pm.addPass(createLinalgCodegenPass(M, N, K));
249+
std::string target_cpu = TO_STRING(TARGET_CPU);
250+
std::string vector_width = TO_STRING(VECTOR_WIDTH);
251+
pm.addPass(createLinalgCodegenPass(M, N, K, target_cpu, vector_width));
241252

242253
// Lower to LLVM
243254
pm.addPass(createConvertVectorToSCFPass());

0 commit comments

Comments
 (0)
0