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-exception
6 //
7 //===----------------------------------------------------------------------===//
8 
9 #include "mlir/Conversion/ArithmeticToLLVM/ArithmeticToLLVM.h"
10 #include "mlir/Conversion/FuncToLLVM/ConvertFuncToLLVMPass.h"
11 #include "mlir/Conversion/LinalgToLLVM/LinalgToLLVM.h"
12 #include "mlir/Conversion/MemRefToLLVM/MemRefToLLVM.h"
13 #include "mlir/Conversion/ReconcileUnrealizedCasts/ReconcileUnrealizedCasts.h"
14 #include "mlir/Conversion/VectorToLLVM/ConvertVectorToLLVM.h"
15 #include "mlir/Conversion/VectorToSCF/VectorToSCF.h"
16 #include "mlir/Dialect/Func/IR/FuncOps.h"
17 #include "mlir/Dialect/Linalg/Passes.h"
18 #include "mlir/ExecutionEngine/CRunnerUtils.h"
19 #include "mlir/ExecutionEngine/ExecutionEngine.h"
20 #include "mlir/ExecutionEngine/MemRefUtils.h"
21 #include "mlir/ExecutionEngine/RunnerUtils.h"
22 #include "mlir/IR/MLIRContext.h"
23 #include "mlir/InitAllDialects.h"
24 #include "mlir/Parser/Parser.h"
25 #include "mlir/Pass/PassManager.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 using namespace mlir;
34 
35 // The JIT isn't supported on Windows at that time
36 #ifndef _WIN32
37 
38 static struct LLVMInitializer {
39   LLVMInitializer() {
40     llvm::InitializeNativeTarget();
41     llvm::InitializeNativeTargetAsmPrinter();
42   }
43 } initializer;
44 
45 /// Simple conversion pipeline for the purpose of testing sources written in
46 /// dialects lowering to LLVM Dialect.
47 static LogicalResult lowerToLLVMDialect(ModuleOp module) {
48   PassManager pm(module.getContext());
49   pm.addPass(mlir::createMemRefToLLVMPass());
50   pm.addNestedPass<FuncOp>(mlir::arith::createConvertArithmeticToLLVMPass());
51   pm.addPass(mlir::createConvertFuncToLLVMPass());
52   pm.addPass(mlir::createReconcileUnrealizedCastsPass());
53   return pm.run(module);
54 }
55 
56 TEST(MLIRExecutionEngine, AddInteger) {
57   std::string moduleStr = R"mlir(
58   func @foo(%arg0 : i32) -> i32 attributes { llvm.emit_c_interface } {
59     %res = arith.addi %arg0, %arg0 : i32
60     return %res : i32
61   }
62   )mlir";
63   DialectRegistry registry;
64   registerAllDialects(registry);
65   registerLLVMDialectTranslation(registry);
66   MLIRContext context(registry);
67   OwningOpRef<ModuleOp> module =
68       parseSourceString<ModuleOp>(moduleStr, &context);
69   ASSERT_TRUE(!!module);
70   ASSERT_TRUE(succeeded(lowerToLLVMDialect(*module)));
71   auto jitOrError = ExecutionEngine::create(*module);
72   ASSERT_TRUE(!!jitOrError);
73   std::unique_ptr<ExecutionEngine> jit = std::move(jitOrError.get());
74   // The result of the function must be passed as output argument.
75   int result = 0;
76   llvm::Error error =
77       jit->invoke("foo", 42, ExecutionEngine::Result<int>(result));
78   ASSERT_TRUE(!error);
79   ASSERT_EQ(result, 42 + 42);
80 }
81 
82 TEST(MLIRExecutionEngine, SubtractFloat) {
83   std::string moduleStr = R"mlir(
84   func @foo(%arg0 : f32, %arg1 : f32) -> f32 attributes { llvm.emit_c_interface } {
85     %res = arith.subf %arg0, %arg1 : f32
86     return %res : f32
87   }
88   )mlir";
89   DialectRegistry registry;
90   registerAllDialects(registry);
91   registerLLVMDialectTranslation(registry);
92   MLIRContext context(registry);
93   OwningOpRef<ModuleOp> module =
94       parseSourceString<ModuleOp>(moduleStr, &context);
95   ASSERT_TRUE(!!module);
96   ASSERT_TRUE(succeeded(lowerToLLVMDialect(*module)));
97   auto jitOrError = ExecutionEngine::create(*module);
98   ASSERT_TRUE(!!jitOrError);
99   std::unique_ptr<ExecutionEngine> jit = std::move(jitOrError.get());
100   // The result of the function must be passed as output argument.
101   float result = -1;
102   llvm::Error error =
103       jit->invoke("foo", 43.0f, 1.0f, ExecutionEngine::result(result));
104   ASSERT_TRUE(!error);
105   ASSERT_EQ(result, 42.f);
106 }
107 
108 TEST(NativeMemRefJit, ZeroRankMemref) {
109   OwningMemRef<float, 0> a({});
110   a[{}] = 42.;
111   ASSERT_EQ(*a->data, 42);
112   a[{}] = 0;
113   std::string moduleStr = R"mlir(
114   func @zero_ranked(%arg0 : memref<f32>) attributes { llvm.emit_c_interface } {
115     %cst42 = arith.constant 42.0 : f32
116     memref.store %cst42, %arg0[] : memref<f32>
117     return
118   }
119   )mlir";
120   DialectRegistry registry;
121   registerAllDialects(registry);
122   registerLLVMDialectTranslation(registry);
123   MLIRContext context(registry);
124   auto module = parseSourceString<ModuleOp>(moduleStr, &context);
125   ASSERT_TRUE(!!module);
126   ASSERT_TRUE(succeeded(lowerToLLVMDialect(*module)));
127   auto jitOrError = ExecutionEngine::create(*module);
128   ASSERT_TRUE(!!jitOrError);
129   auto jit = std::move(jitOrError.get());
130 
131   llvm::Error error = jit->invoke("zero_ranked", &*a);
132   ASSERT_TRUE(!error);
133   EXPECT_EQ((a[{}]), 42.);
134   for (float &elt : *a)
135     EXPECT_EQ(&elt, &(a[{}]));
136 }
137 
138 TEST(NativeMemRefJit, RankOneMemref) {
139   int64_t shape[] = {9};
140   OwningMemRef<float, 1> a(shape);
141   int count = 1;
142   for (float &elt : *a) {
143     EXPECT_EQ(&elt, &(a[{count - 1}]));
144     elt = count++;
145   }
146 
147   std::string moduleStr = R"mlir(
148   func @one_ranked(%arg0 : memref<?xf32>) attributes { llvm.emit_c_interface } {
149     %cst42 = arith.constant 42.0 : f32
150     %cst5 = arith.constant 5 : index
151     memref.store %cst42, %arg0[%cst5] : memref<?xf32>
152     return
153   }
154   )mlir";
155   DialectRegistry registry;
156   registerAllDialects(registry);
157   registerLLVMDialectTranslation(registry);
158   MLIRContext context(registry);
159   auto module = parseSourceString<ModuleOp>(moduleStr, &context);
160   ASSERT_TRUE(!!module);
161   ASSERT_TRUE(succeeded(lowerToLLVMDialect(*module)));
162   auto jitOrError = ExecutionEngine::create(*module);
163   ASSERT_TRUE(!!jitOrError);
164   auto jit = std::move(jitOrError.get());
165 
166   llvm::Error error = jit->invoke("one_ranked", &*a);
167   ASSERT_TRUE(!error);
168   count = 1;
169   for (float &elt : *a) {
170     if (count == 6)
171       EXPECT_EQ(elt, 42.);
172     else
173       EXPECT_EQ(elt, count);
174     count++;
175   }
176 }
177 
178 TEST(NativeMemRefJit, BasicMemref) {
179   constexpr int k = 3;
180   constexpr int m = 7;
181   // Prepare arguments beforehand.
182   auto init = [=](float &elt, ArrayRef<int64_t> indices) {
183     assert(indices.size() == 2);
184     elt = m * indices[0] + indices[1];
185   };
186   int64_t shape[] = {k, m};
187   int64_t shapeAlloc[] = {k + 1, m + 1};
188   OwningMemRef<float, 2> a(shape, shapeAlloc, init);
189   ASSERT_EQ(a->sizes[0], k);
190   ASSERT_EQ(a->sizes[1], m);
191   ASSERT_EQ(a->strides[0], m + 1);
192   ASSERT_EQ(a->strides[1], 1);
193   for (int i = 0; i < k; ++i) {
194     for (int j = 0; j < m; ++j) {
195       EXPECT_EQ((a[{i, j}]), i * m + j);
196       EXPECT_EQ(&(a[{i, j}]), &((*a)[i][j]));
197     }
198   }
199   std::string moduleStr = R"mlir(
200   func @rank2_memref(%arg0 : memref<?x?xf32>, %arg1 : memref<?x?xf32>) attributes { llvm.emit_c_interface } {
201     %x = arith.constant 2 : index
202     %y = arith.constant 1 : index
203     %cst42 = arith.constant 42.0 : f32
204     memref.store %cst42, %arg0[%y, %x] : memref<?x?xf32>
205     memref.store %cst42, %arg1[%x, %y] : memref<?x?xf32>
206     return
207   }
208   )mlir";
209   DialectRegistry registry;
210   registerAllDialects(registry);
211   registerLLVMDialectTranslation(registry);
212   MLIRContext context(registry);
213   OwningOpRef<ModuleOp> module =
214       parseSourceString<ModuleOp>(moduleStr, &context);
215   ASSERT_TRUE(!!module);
216   ASSERT_TRUE(succeeded(lowerToLLVMDialect(*module)));
217   auto jitOrError = ExecutionEngine::create(*module);
218   ASSERT_TRUE(!!jitOrError);
219   std::unique_ptr<ExecutionEngine> jit = std::move(jitOrError.get());
220 
221   llvm::Error error = jit->invoke("rank2_memref", &*a, &*a);
222   ASSERT_TRUE(!error);
223   EXPECT_EQ(((*a)[1][2]), 42.);
224   EXPECT_EQ((a[{2, 1}]), 42.);
225 }
226 
227 // A helper function that will be called from the JIT
228 static void memrefMultiply(::StridedMemRefType<float, 2> *memref,
229                            int32_t coefficient) {
230   for (float &elt : *memref)
231     elt *= coefficient;
232 }
233 
234 TEST(NativeMemRefJit, JITCallback) {
235   constexpr int k = 2;
236   constexpr int m = 2;
237   int64_t shape[] = {k, m};
238   int64_t shapeAlloc[] = {k + 1, m + 1};
239   OwningMemRef<float, 2> a(shape, shapeAlloc);
240   int count = 1;
241   for (float &elt : *a)
242     elt = count++;
243 
244   std::string moduleStr = R"mlir(
245   func private @callback(%arg0: memref<?x?xf32>, %coefficient: i32)  attributes { llvm.emit_c_interface }
246   func @caller_for_callback(%arg0: memref<?x?xf32>, %coefficient: i32) attributes { llvm.emit_c_interface } {
247     %unranked = memref.cast %arg0: memref<?x?xf32> to memref<*xf32>
248     call @callback(%arg0, %coefficient) : (memref<?x?xf32>, i32) -> ()
249     return
250   }
251   )mlir";
252   DialectRegistry registry;
253   registerAllDialects(registry);
254   registerLLVMDialectTranslation(registry);
255   MLIRContext context(registry);
256   auto module = parseSourceString<ModuleOp>(moduleStr, &context);
257   ASSERT_TRUE(!!module);
258   ASSERT_TRUE(succeeded(lowerToLLVMDialect(*module)));
259   auto jitOrError = ExecutionEngine::create(*module);
260   ASSERT_TRUE(!!jitOrError);
261   auto jit = std::move(jitOrError.get());
262   // Define any extra symbols so they're available at runtime.
263   jit->registerSymbols([&](llvm::orc::MangleAndInterner interner) {
264     llvm::orc::SymbolMap symbolMap;
265     symbolMap[interner("_mlir_ciface_callback")] =
266         llvm::JITEvaluatedSymbol::fromPointer(memrefMultiply);
267     return symbolMap;
268   });
269 
270   int32_t coefficient = 3.;
271   llvm::Error error = jit->invoke("caller_for_callback", &*a, coefficient);
272   ASSERT_TRUE(!error);
273   count = 1;
274   for (float elt : *a)
275     ASSERT_EQ(elt, coefficient * count++);
276 }
277 
278 #endif // _WIN32
279