395 lines · cpp
1//===- Invoke.cpp ------------------------------------*- C++ -*-===//2//3// This file is licensed under the Apache License v2.0 with LLVM Exceptions.4// See https://llvm.org/LICENSE.txt for license information.5// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception6//7//===----------------------------------------------------------------------===//8 9#include "mlir/Conversion/ArithToLLVM/ArithToLLVM.h"10#include "mlir/Conversion/FuncToLLVM/ConvertFuncToLLVMPass.h"11#include "mlir/Conversion/MemRefToLLVM/MemRefToLLVM.h"12#include "mlir/Conversion/ReconcileUnrealizedCasts/ReconcileUnrealizedCasts.h"13#include "mlir/Conversion/VectorToLLVM/ConvertVectorToLLVM.h"14#include "mlir/Conversion/VectorToSCF/VectorToSCF.h"15#include "mlir/Dialect/Func/IR/FuncOps.h"16#include "mlir/Dialect/Linalg/Passes.h"17#include "mlir/ExecutionEngine/CRunnerUtils.h"18#include "mlir/ExecutionEngine/ExecutionEngine.h"19#include "mlir/ExecutionEngine/MemRefUtils.h"20#include "mlir/ExecutionEngine/RunnerUtils.h"21#include "mlir/IR/MLIRContext.h"22#include "mlir/InitAllDialects.h"23#include "mlir/Parser/Parser.h"24#include "mlir/Pass/PassManager.h"25#include "mlir/Target/LLVMIR/Dialect/Builtin/BuiltinToLLVMIRTranslation.h"26#include "mlir/Target/LLVMIR/Dialect/LLVMIR/LLVMToLLVMIRTranslation.h"27#include "mlir/Target/LLVMIR/Export.h"28#include "llvm/Support/TargetSelect.h"29#include "llvm/Support/raw_ostream.h"30 31#include "gmock/gmock.h"32 33// SPARC currently lacks JIT support.34#ifdef __sparc__35#define SKIP_WITHOUT_JIT(x) DISABLED_##x36#else37#define SKIP_WITHOUT_JIT(x) x38#endif39 40using namespace mlir;41 42// The JIT isn't supported on Windows at that time43#ifndef _WIN3244 45static struct LLVMInitializer {46 LLVMInitializer() {47 llvm::InitializeNativeTarget();48 llvm::InitializeNativeTargetAsmPrinter();49 }50} initializer;51 52/// Simple conversion pipeline for the purpose of testing sources written in53/// dialects lowering to LLVM Dialect.54static LogicalResult lowerToLLVMDialect(ModuleOp module) {55 PassManager pm(module->getName());56 pm.addPass(mlir::createFinalizeMemRefToLLVMConversionPass());57 pm.addNestedPass<func::FuncOp>(mlir::createArithToLLVMConversionPass());58 pm.addPass(mlir::createConvertFuncToLLVMPass());59 pm.addPass(mlir::createReconcileUnrealizedCastsPass());60 return pm.run(module);61}62 63TEST(MLIRExecutionEngine, SKIP_WITHOUT_JIT(AddInteger)) {64#ifdef __s390__65 std::string moduleStr = R"mlir(66 func.func @foo(%arg0 : i32 {llvm.signext}) -> (i32 {llvm.signext}) attributes { llvm.emit_c_interface } {67 %res = arith.addi %arg0, %arg0 : i3268 return %res : i3269 }70 )mlir";71#else72 std::string moduleStr = R"mlir(73 func.func @foo(%arg0 : i32) -> i32 attributes { llvm.emit_c_interface } {74 %res = arith.addi %arg0, %arg0 : i3275 return %res : i3276 }77 )mlir";78#endif79 DialectRegistry registry;80 registerAllDialects(registry);81 registerBuiltinDialectTranslation(registry);82 registerLLVMDialectTranslation(registry);83 MLIRContext context(registry);84 OwningOpRef<ModuleOp> module =85 parseSourceString<ModuleOp>(moduleStr, &context);86 ASSERT_TRUE(!!module);87 ASSERT_TRUE(succeeded(lowerToLLVMDialect(*module)));88 auto jitOrError = ExecutionEngine::create(*module);89 ASSERT_TRUE(!!jitOrError);90 std::unique_ptr<ExecutionEngine> jit = std::move(jitOrError.get());91 // The result of the function must be passed as output argument.92 int result = 0;93 llvm::Error error =94 jit->invoke("foo", 42, ExecutionEngine::Result<int>(result));95 ASSERT_TRUE(!error);96 ASSERT_EQ(result, 42 + 42);97}98 99TEST(MLIRExecutionEngine, SKIP_WITHOUT_JIT(SubtractFloat)) {100 std::string moduleStr = R"mlir(101 func.func @foo(%arg0 : f32, %arg1 : f32) -> f32 attributes { llvm.emit_c_interface } {102 %res = arith.subf %arg0, %arg1 : f32103 return %res : f32104 }105 )mlir";106 DialectRegistry registry;107 registerAllDialects(registry);108 registerBuiltinDialectTranslation(registry);109 registerLLVMDialectTranslation(registry);110 MLIRContext context(registry);111 OwningOpRef<ModuleOp> module =112 parseSourceString<ModuleOp>(moduleStr, &context);113 ASSERT_TRUE(!!module);114 ASSERT_TRUE(succeeded(lowerToLLVMDialect(*module)));115 auto jitOrError = ExecutionEngine::create(*module);116 ASSERT_TRUE(!!jitOrError);117 std::unique_ptr<ExecutionEngine> jit = std::move(jitOrError.get());118 // The result of the function must be passed as output argument.119 float result = -1;120 llvm::Error error =121 jit->invoke("foo", 43.0f, 1.0f, ExecutionEngine::result(result));122 ASSERT_TRUE(!error);123 ASSERT_EQ(result, 42.f);124}125 126TEST(NativeMemRefJit, SKIP_WITHOUT_JIT(ZeroRankMemref)) {127 OwningMemRef<float, 0> a({});128 a[{}] = 42.;129 ASSERT_EQ(*a->data, 42);130 a[{}] = 0;131 std::string moduleStr = R"mlir(132 func.func @zero_ranked(%arg0 : memref<f32>) attributes { llvm.emit_c_interface } {133 %cst42 = arith.constant 42.0 : f32134 memref.store %cst42, %arg0[] : memref<f32>135 return136 }137 )mlir";138 DialectRegistry registry;139 registerAllDialects(registry);140 registerBuiltinDialectTranslation(registry);141 registerLLVMDialectTranslation(registry);142 MLIRContext context(registry);143 auto module = parseSourceString<ModuleOp>(moduleStr, &context);144 ASSERT_TRUE(!!module);145 ASSERT_TRUE(succeeded(lowerToLLVMDialect(*module)));146 auto jitOrError = ExecutionEngine::create(*module);147 ASSERT_TRUE(!!jitOrError);148 auto jit = std::move(jitOrError.get());149 150 llvm::Error error = jit->invoke("zero_ranked", &*a);151 ASSERT_TRUE(!error);152 EXPECT_EQ((a[{}]), 42.);153 for (float &elt : *a)154 EXPECT_EQ(&elt, &(a[{}]));155}156 157TEST(NativeMemRefJit, SKIP_WITHOUT_JIT(RankOneMemref)) {158 int64_t shape[] = {9};159 OwningMemRef<float, 1> a(shape);160 int count = 1;161 for (float &elt : *a) {162 EXPECT_EQ(&elt, &(a[{count - 1}]));163 elt = count++;164 }165 166 std::string moduleStr = R"mlir(167 func.func @one_ranked(%arg0 : memref<?xf32>) attributes { llvm.emit_c_interface } {168 %cst42 = arith.constant 42.0 : f32169 %cst5 = arith.constant 5 : index170 memref.store %cst42, %arg0[%cst5] : memref<?xf32>171 return172 }173 )mlir";174 DialectRegistry registry;175 registerAllDialects(registry);176 registerBuiltinDialectTranslation(registry);177 registerLLVMDialectTranslation(registry);178 MLIRContext context(registry);179 auto module = parseSourceString<ModuleOp>(moduleStr, &context);180 ASSERT_TRUE(!!module);181 ASSERT_TRUE(succeeded(lowerToLLVMDialect(*module)));182 auto jitOrError = ExecutionEngine::create(*module);183 ASSERT_TRUE(!!jitOrError);184 auto jit = std::move(jitOrError.get());185 186 llvm::Error error = jit->invoke("one_ranked", &*a);187 ASSERT_TRUE(!error);188 count = 1;189 for (float &elt : *a) {190 if (count == 6)191 EXPECT_EQ(elt, 42.);192 else193 EXPECT_EQ(elt, count);194 count++;195 }196}197 198TEST(NativeMemRefJit, SKIP_WITHOUT_JIT(BasicMemref)) {199 constexpr int k = 3;200 constexpr int m = 7;201 // Prepare arguments beforehand.202 auto init = [=](float &elt, ArrayRef<int64_t> indices) {203 assert(indices.size() == 2);204 elt = m * indices[0] + indices[1];205 };206 int64_t shape[] = {k, m};207 int64_t shapeAlloc[] = {k + 1, m + 1};208 // Use a large alignment to stress the case where the memref data/basePtr are209 // disjoint.210 int alignment = 8192;211 OwningMemRef<float, 2> a(shape, shapeAlloc, init, alignment);212 ASSERT_EQ(213 (void *)(((uintptr_t)a->basePtr + alignment - 1) & ~(alignment - 1)),214 a->data);215 ASSERT_EQ(a->sizes[0], k);216 ASSERT_EQ(a->sizes[1], m);217 ASSERT_EQ(a->strides[0], m + 1);218 ASSERT_EQ(a->strides[1], 1);219 for (int i = 0; i < k; ++i) {220 for (int j = 0; j < m; ++j) {221 EXPECT_EQ((a[{i, j}]), i * m + j);222 EXPECT_EQ(&(a[{i, j}]), &((*a)[i][j]));223 }224 }225 std::string moduleStr = R"mlir(226 func.func @rank2_memref(%arg0 : memref<?x?xf32>, %arg1 : memref<?x?xf32>) attributes { llvm.emit_c_interface } {227 %x = arith.constant 2 : index228 %y = arith.constant 1 : index229 %cst42 = arith.constant 42.0 : f32230 memref.store %cst42, %arg0[%y, %x] : memref<?x?xf32>231 memref.store %cst42, %arg1[%x, %y] : memref<?x?xf32>232 return233 }234 )mlir";235 DialectRegistry registry;236 registerAllDialects(registry);237 registerBuiltinDialectTranslation(registry);238 registerLLVMDialectTranslation(registry);239 MLIRContext context(registry);240 OwningOpRef<ModuleOp> module =241 parseSourceString<ModuleOp>(moduleStr, &context);242 ASSERT_TRUE(!!module);243 ASSERT_TRUE(succeeded(lowerToLLVMDialect(*module)));244 auto jitOrError = ExecutionEngine::create(*module);245 ASSERT_TRUE(!!jitOrError);246 std::unique_ptr<ExecutionEngine> jit = std::move(jitOrError.get());247 248 llvm::Error error = jit->invoke("rank2_memref", &*a, &*a);249 ASSERT_TRUE(!error);250 EXPECT_EQ(((*a)[1][2]), 42.);251 EXPECT_EQ((a[{2, 1}]), 42.);252}253 254TEST(NativeMemRefJit, SKIP_WITHOUT_JIT(OwningMemrefZeroInit)) {255 constexpr int k = 3;256 constexpr int m = 7;257 int64_t shape[] = {k, m};258 // Use a large alignment to stress the case where the memref data/basePtr are259 // disjoint.260 int alignment = 8192;261 OwningMemRef<float, 2> a(shape, {}, {}, alignment);262 ASSERT_EQ(263 (void *)(((uintptr_t)a->basePtr + alignment - 1) & ~(alignment - 1)),264 a->data);265 for (int i = 0; i < k; ++i) {266 for (int j = 0; j < m; ++j) {267 EXPECT_EQ((a[{i, j}]), 0.);268 }269 }270}271 272// A helper function that will be called from the JIT273static void memrefMultiply(::StridedMemRefType<float, 2> *memref,274 int32_t coefficient) {275 for (float &elt : *memref)276 elt *= coefficient;277}278 279// MSAN does not work with JIT.280#if __has_feature(memory_sanitizer)281#define MAYBE_JITCallback DISABLED_JITCallback282#else283#define MAYBE_JITCallback SKIP_WITHOUT_JIT(JITCallback)284#endif285TEST(NativeMemRefJit, MAYBE_JITCallback) {286 constexpr int k = 2;287 constexpr int m = 2;288 int64_t shape[] = {k, m};289 int64_t shapeAlloc[] = {k + 1, m + 1};290 OwningMemRef<float, 2> a(shape, shapeAlloc);291 int count = 1;292 for (float &elt : *a)293 elt = count++;294 295#ifdef __s390__296 std::string moduleStr = R"mlir(297 func.func private @callback(%arg0: memref<?x?xf32>, %coefficient: i32 {llvm.signext}) attributes { llvm.emit_c_interface }298 func.func @caller_for_callback(%arg0: memref<?x?xf32>, %coefficient: i32 {llvm.signext}) attributes { llvm.emit_c_interface } {299 %unranked = memref.cast %arg0: memref<?x?xf32> to memref<*xf32>300 call @callback(%arg0, %coefficient) : (memref<?x?xf32>, i32) -> ()301 return302 }303 )mlir";304#else305 std::string moduleStr = R"mlir(306 func.func private @callback(%arg0: memref<?x?xf32>, %coefficient: i32) attributes { llvm.emit_c_interface }307 func.func @caller_for_callback(%arg0: memref<?x?xf32>, %coefficient: i32) attributes { llvm.emit_c_interface } {308 %unranked = memref.cast %arg0: memref<?x?xf32> to memref<*xf32>309 call @callback(%arg0, %coefficient) : (memref<?x?xf32>, i32) -> ()310 return311 }312 )mlir";313#endif314 315 DialectRegistry registry;316 registerAllDialects(registry);317 registerBuiltinDialectTranslation(registry);318 registerLLVMDialectTranslation(registry);319 MLIRContext context(registry);320 auto module = parseSourceString<ModuleOp>(moduleStr, &context);321 ASSERT_TRUE(!!module);322 ASSERT_TRUE(succeeded(lowerToLLVMDialect(*module)));323 auto jitOrError = ExecutionEngine::create(*module);324 ASSERT_TRUE(!!jitOrError);325 auto jit = std::move(jitOrError.get());326 // Define any extra symbols so they're available at runtime.327 jit->registerSymbols([&](llvm::orc::MangleAndInterner interner) {328 llvm::orc::SymbolMap symbolMap;329 symbolMap[interner("_mlir_ciface_callback")] = {330 llvm::orc::ExecutorAddr::fromPtr(memrefMultiply),331 llvm::JITSymbolFlags::Exported};332 return symbolMap;333 });334 335 int32_t coefficient = 3.;336 llvm::Error error = jit->invoke("caller_for_callback", &*a, coefficient);337 ASSERT_TRUE(!error);338 count = 1;339 for (float elt : *a)340 ASSERT_EQ(elt, coefficient * count++);341}342 343static int initCnt = 0;344// A helper function that will be called during the JIT's initialization.345static void initCallback() { initCnt += 1; }346 347TEST(MLIRExecutionEngine, SKIP_WITHOUT_JIT(CallbackInGlobalCtor)) {348 auto tmBuilderOrError = llvm::orc::JITTargetMachineBuilder::detectHost();349 ASSERT_TRUE(!!tmBuilderOrError);350 if (tmBuilderOrError->getTargetTriple().isAArch64()) {351 GTEST_SKIP() << "Skipping global ctor initialization test on Aarch64 "352 "because of bug #71963";353 return;354 }355 std::string moduleStr = R"mlir(356 llvm.mlir.global_ctors ctors = [@ctor], priorities = [0 : i32], data = [#llvm.zero]357 llvm.func @ctor() {358 func.call @init_callback() : () -> ()359 llvm.return360 }361 func.func private @init_callback() attributes { llvm.emit_c_interface }362 )mlir";363 364 DialectRegistry registry;365 registerAllDialects(registry);366 registerBuiltinDialectTranslation(registry);367 registerLLVMDialectTranslation(registry);368 MLIRContext context(registry);369 auto module = parseSourceString<ModuleOp>(moduleStr, &context);370 ASSERT_TRUE(!!module);371 ASSERT_TRUE(succeeded(lowerToLLVMDialect(*module)));372 ExecutionEngineOptions jitOptions;373 auto jitOrError = ExecutionEngine::create(*module, jitOptions);374 ASSERT_TRUE(!!jitOrError);375 // validate initialization is not run on construction376 ASSERT_EQ(initCnt, 0);377 auto jit = std::move(jitOrError.get());378 // Define any extra symbols so they're available at initialization.379 jit->registerSymbols([&](llvm::orc::MangleAndInterner interner) {380 llvm::orc::SymbolMap symbolMap;381 symbolMap[interner("_mlir_ciface_init_callback")] = {382 llvm::orc::ExecutorAddr::fromPtr(initCallback),383 llvm::JITSymbolFlags::Exported};384 return symbolMap;385 });386 jit->initialize();387 // validate the side effect of initialization388 ASSERT_EQ(initCnt, 1);389 // next initialization should be noop390 jit->initialize();391 ASSERT_EQ(initCnt, 1);392}393 394#endif // _WIN32395