[ModelBuilder] new example that test MLIR building and execution

PiperOrigin-RevId: 303201274

experimental/ModelBuilder/test/BUILD

diff --git a/experimental/ModelBuilder/test/BUILD b/experimental/ModelBuilder/test/BUILD
index f34fb38..acd39e2 100644
--- a/experimental/ModelBuilder/test/BUILD
+++ b/experimental/ModelBuilder/test/BUILD
@@ -30,6 +30,7 @@
     data = [
         ":runtime-support.so",
         # Tests.
+        ":test-dot-prod",
         ":test-mnist-jit",
         ":test-simple-jit",
         ":test-simple-mlir",
@@ -42,6 +43,24 @@
 )
 
 cc_binary(
+    name = "test-dot-prod",
+    srcs = ["TestDotProdJIT.cpp"],
+    tags = [
+        "noga",
+    ],
+    deps = [
+        ":runtime-support.so",
+        "//experimental/ModelBuilder",
+        "//experimental/ModelBuilder:ModelRunner",
+        "@llvm-project//llvm:support",
+        "@llvm-project//mlir:AllPassesAndDialects",
+        "@llvm-project//mlir:EDSC",
+        "@llvm-project//mlir:IR",
+        "@llvm-project//mlir:LoopOpsTransforms",
+    ],
+)
+
+cc_binary(
     name = "test-mnist-jit",
     srcs = ["TestMNISTJIT.cpp"],
     tags = [

experimental/ModelBuilder/test/TestDotProdJIT.cpp

diff --git a/experimental/ModelBuilder/test/TestDotProdJIT.cpp b/experimental/ModelBuilder/test/TestDotProdJIT.cpp
new file mode 100644
index 0000000..2751a6c
--- /dev/null
+++ b/experimental/ModelBuilder/test/TestDotProdJIT.cpp
@@ -0,0 +1,111 @@
+// Copyright 2020 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      https://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+// clang-format off
+
+// NOLINTNEXTLINE
+// RUN: test-dot-prod -runtime-support=$(dirname %s)/runtime-support.so 2>&1 | IreeFileCheck %s
+
+// clang-format on
+
+#include "experimental/ModelBuilder/MemRefUtils.h"
+#include "experimental/ModelBuilder/ModelBuilder.h"
+#include "experimental/ModelBuilder/ModelRunner.h"
+#include "llvm/Support/CommandLine.h"
+#include "llvm/Support/InitLLVM.h"
+
+using namespace mlir;                    // NOLINT
+using namespace mlir::edsc;              // NOLINT
+using namespace mlir::edsc::intrinsics;  // NOLINT
+
+static llvm::cl::opt<std::string> runtimeSupport(
+    "runtime-support", llvm::cl::desc("Runtime support library filename"),
+    llvm::cl::value_desc("filename"), llvm::cl::init("-"));
+
+void DotProdOnVectors() {
+  constexpr unsigned N = 4;
+
+  ModelBuilder modelBuilder;
+  // Build a func "dot_prod".
+  constexpr StringLiteral funcName = "dot-prod";
+  auto f32 = modelBuilder.f32;
+  auto vectorType = modelBuilder.getVectorType(N, f32);
+  auto refType = modelBuilder.getMemRefType(1, vectorType);
+
+  auto func = modelBuilder.makeFunction(funcName, {}, {refType, refType});
+
+  SmallVector<AffineMap, 3> accesses;
+  accesses.push_back(modelBuilder.getDimIdentityMap());
+  accesses.push_back(accesses[0]);
+  accesses.push_back(AffineMap::get(1, 0, modelBuilder.getContext()));
+
+  SmallVector<Attribute, 1> iterator_types;
+  iterator_types.push_back(modelBuilder.getStringAttr("reduction"));
+
+  OpBuilder b(&func.getBody());
+  ScopedContext scope(b, func.getLoc());
+  ValueHandle A(func.getArgument(0)), B(func.getArgument(1));
+  Value zero = ValueBuilder<ConstantIndexOp>(0);
+  Value A_val = std_load(A, ValueRange(zero));
+  Value B_val = std_load(B, ValueRange(zero));
+  Value zeroF = std_constant_float(APFloat(0.0f), f32);
+  Value res_val = (vector_contract(A_val, B_val, zeroF,
+                                   modelBuilder.getAffineMapArrayAttr(accesses),
+                                   modelBuilder.getArrayAttr(iterator_types)));
+
+  (vector_print(A_val));
+  (vector_print(B_val));
+  (vector_print(res_val));
+
+  std_ret();
+
+  // Compile the function, pass in runtime support library
+  //    to the execution engine for vector.print.
+  ModelRunner runner(modelBuilder.getModuleRef());
+  runner.compile(CompilationOptions(), runtimeSupport);
+
+  // initialize data by interoperating with the MLIR ABI by codegen.
+  auto inputInit1 = [](unsigned idx, Vector1D<N, float> *ptr) {
+    for (unsigned i = 0; i < N; ++i) ptr[idx][i] = 3.0 * i;
+  };
+  auto inputInit2 = [](unsigned idx, Vector1D<N, float> *ptr) {
+    for (unsigned i = 0; i < N; ++i) ptr[idx][i] = 2.0 * i;
+  };
+
+  auto _A = makeInitializedStridedMemRefDescriptor<Vector1D<N, float>, 1>(
+      {N}, inputInit1);
+  auto _B = makeInitializedStridedMemRefDescriptor<Vector1D<N, float>, 1>(
+      {N}, inputInit2);
+
+  // Call the funcOp
+  const std::string funAdapterName =
+      (llvm::Twine("_mlir_ciface_") + funcName).str();
+  auto *bufferA = _A.get();
+  auto *bufferB = _B.get();
+  void *args[3] = {&bufferA, &bufferB};
+
+  // CHECK: ( 0, 3, 6, 9 )
+  // CHECK: ( 0, 2, 4, 6 )
+  // CHECK: 84
+  auto err =
+      runner.engine->invoke(funAdapterName, MutableArrayRef<void *>{args});
+
+  if (err) llvm_unreachable("Error running function.");
+}
+
+int main(int argc, char **argv) {
+  llvm::InitLLVM y(argc, argv);
+  llvm::cl::ParseCommandLineOptions(argc, argv, "TestDotProd\n");
+  DotProdOnVectors();
+}