[Mlir-commits] [mlir] [mlir][ExecutionEngine] Reserve one allocation per object on COFF (PR #233658)

[Mehdi Amini]

================
@@ -60,6 +62,74 @@ static LogicalResult lowerToLLVMDialect(ModuleOp module) {
   return pm.run(module);
 }
 
+namespace {
+/// Forwards to sys::Memory and records what each mapping was requested for.
+class RecordingMemoryMapper : public llvm::SectionMemoryManager::MemoryMapper {
+public:
+  llvm::sys::MemoryBlock
+  allocateMappedMemory(llvm::SectionMemoryManager::AllocationPurpose purpose,
+                       size_t numBytes, const llvm::sys::MemoryBlock *nearBlock,
+                       unsigned flags, std::error_code &ec) override {
+    purposes.push_back(purpose);
+    return llvm::sys::Memory::allocateMappedMemory(numBytes, nearBlock, flags,
+                                                   ec);
+  }
+
+  std::error_code protectMappedMemory(const llvm::sys::MemoryBlock &block,
+                                      unsigned flags) override {
+    return llvm::sys::Memory::protectMappedMemory(block, flags);
+  }
+
+  std::error_code releaseMappedMemory(llvm::sys::MemoryBlock &block) override {
+    return llvm::sys::Memory::releaseMappedMemory(block);
+  }
+
+  std::vector<llvm::SectionMemoryManager::AllocationPurpose> purposes;
+};
+} // namespace
+
+// On AArch64 and COFF hosts the JIT carves each object's code and data from
+// one reserved mapping, because their relocations limit the distance between
+// sections. Code mapped on its own could land anywhere relative to the data.
+// Unlike the tests below, this one also runs on Windows.
+TEST(MLIRExecutionEngine, SKIP_WITHOUT_JIT(CodeAndDataShareOneMapping)) {
+  llvm::Triple host(llvm::sys::getProcessTriple());
+  if (!host.isAArch64() && !host.isOSBinFormatCOFF())
+    GTEST_SKIP() << "only AArch64 and COFF hosts reserve one mapping";
+  std::string moduleStr = R"mlir(
+  func.func @foo(%arg0 : i32) -> i32 attributes { llvm.emit_c_interface } {
+    %res = arith.addi %arg0, %arg0 : i32
+    return %res : i32
+  }
+  )mlir";
+  DialectRegistry registry;
+  registerAllDialects(registry);
+  registerBuiltinDialectTranslation(registry);
+  registerLLVMDialectTranslation(registry);
+  MLIRContext context(registry);
+  OwningOpRef<ModuleOp> module =
+      parseSourceString<ModuleOp>(moduleStr, &context);
+  ASSERT_TRUE(!!module);
+  ASSERT_TRUE(succeeded(lowerToLLVMDialect(*module)));
+  RecordingMemoryMapper mapper;
+  ExecutionEngineOptions options;
+  options.sectionMemoryMapper = &mapper;
+  auto jitOrError = ExecutionEngine::create(*module, options);
+  ASSERT_TRUE(!!jitOrError);
+  std::unique_ptr<ExecutionEngine> jit = std::move(jitOrError.get());
+  int result = 0;
+  llvm::Error error =
+      jit->invoke("foo", 42, ExecutionEngine::Result<int>(result));
+  ASSERT_TRUE(!error);
+  ASSERT_EQ(result, 42 + 42);
+  ASSERT_FALSE(mapper.purposes.empty());
+  EXPECT_FALSE(llvm::is_contained(
+      mapper.purposes, llvm::SectionMemoryManager::AllocationPurpose::Code));
+}
+
+// The JIT isn't supported on Windows at that time
----------------
joker-eph wrote:

The comment says that the JIT is't supported on Windows but the test above runs?

https://github.com/llvm/llvm-project/pull/233658