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