File: MLModelRunnerTest.cpp

package info (click to toggle)
llvm-toolchain-15 1%3A15.0.6-4
  • links: PTS, VCS
  • area: main
  • in suites: bookworm
  • size: 1,554,644 kB
  • sloc: cpp: 5,922,452; ansic: 1,012,136; asm: 674,362; python: 191,568; objc: 73,855; f90: 42,327; lisp: 31,913; pascal: 11,973; javascript: 10,144; sh: 9,421; perl: 7,447; ml: 5,527; awk: 3,523; makefile: 2,520; xml: 885; cs: 573; fortran: 567
file content (119 lines) | stat: -rw-r--r-- 4,200 bytes parent folder | download | duplicates (3)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
//===- MLModelRunnerTest.cpp - test for MLModelRunner ---------------------===//
//
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
// See https://llvm.org/LICENSE.txt for license information.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
//
//===----------------------------------------------------------------------===//

#include "llvm/Analysis/MLModelRunner.h"
#include "llvm/Analysis/NoInferenceModelRunner.h"
#include "llvm/Analysis/ReleaseModeModelRunner.h"
#include "gtest/gtest.h"

using namespace llvm;

namespace llvm {
// This is a mock of the kind of AOT-generated model evaluator. It has 2 tensors
// of shape {1}, and 'evaluation' adds them.
// The interface is the one expected by ReleaseModelRunner.
class MockAOTModel final {
  int64_t A = 0;
  int64_t B = 0;
  int64_t R = 0;

public:
  MockAOTModel() = default;
  int LookupArgIndex(const std::string &Name) {
    if (Name == "prefix_a")
      return 0;
    if (Name == "prefix_b")
      return 1;
    return -1;
  }
  int LookupResultIndex(const std::string &) { return 0; }
  void Run() { R = A + B; }
  void *result_data(int RIndex) {
    if (RIndex == 0)
      return &R;
    return nullptr;
  }
  void *arg_data(int Index) {
    switch (Index) {
    case 0:
      return &A;
    case 1:
      return &B;
    default:
      return nullptr;
    }
  }
};
} // namespace llvm

TEST(NoInferenceModelRunner, AccessTensors) {
  const std::vector<TensorSpec> Inputs{
      TensorSpec::createSpec<int64_t>("F1", {1}),
      TensorSpec::createSpec<int64_t>("F2", {10}),
      TensorSpec::createSpec<float>("F2", {5}),
  };
  LLVMContext Ctx;
  NoInferenceModelRunner NIMR(Ctx, Inputs);
  NIMR.getTensor<int64_t>(0)[0] = 1;
  std::memcpy(NIMR.getTensor<int64_t>(1),
              std::vector<int64_t>{1, 2, 3, 4, 5, 6, 7, 8, 9, 10}.data(),
              10 * sizeof(int64_t));
  std::memcpy(NIMR.getTensor<float>(2),
              std::vector<float>{0.1f, 0.2f, 0.3f, 0.4f, 0.5f}.data(),
              5 * sizeof(float));
  ASSERT_EQ(NIMR.getTensor<int64_t>(0)[0], 1);
  ASSERT_EQ(NIMR.getTensor<int64_t>(1)[8], 9);
  ASSERT_EQ(NIMR.getTensor<float>(2)[1], 0.2f);
}

TEST(ReleaseModeRunner, NormalUse) {
  LLVMContext Ctx;
  std::vector<TensorSpec> Inputs{TensorSpec::createSpec<int64_t>("a", {1}),
                                 TensorSpec::createSpec<int64_t>("b", {1})};
  auto Evaluator = std::make_unique<ReleaseModeModelRunner<MockAOTModel>>(
      Ctx, Inputs, "", "prefix_");
  *Evaluator->getTensor<int64_t>(0) = 1;
  *Evaluator->getTensor<int64_t>(1) = 2;
  EXPECT_EQ(Evaluator->evaluate<int64_t>(), 3);
  EXPECT_EQ(*Evaluator->getTensor<int64_t>(0), 1);
  EXPECT_EQ(*Evaluator->getTensor<int64_t>(1), 2);
}

TEST(ReleaseModeRunner, ExtraFeatures) {
  LLVMContext Ctx;
  std::vector<TensorSpec> Inputs{TensorSpec::createSpec<int64_t>("a", {1}),
                                 TensorSpec::createSpec<int64_t>("b", {1}),
                                 TensorSpec::createSpec<int64_t>("c", {1})};
  auto Evaluator = std::make_unique<ReleaseModeModelRunner<MockAOTModel>>(
      Ctx, Inputs, "", "prefix_");
  *Evaluator->getTensor<int64_t>(0) = 1;
  *Evaluator->getTensor<int64_t>(1) = 2;
  *Evaluator->getTensor<int64_t>(2) = -3;
  EXPECT_EQ(Evaluator->evaluate<int64_t>(), 3);
  EXPECT_EQ(*Evaluator->getTensor<int64_t>(0), 1);
  EXPECT_EQ(*Evaluator->getTensor<int64_t>(1), 2);
  EXPECT_EQ(*Evaluator->getTensor<int64_t>(2), -3);
}

TEST(ReleaseModeRunner, ExtraFeaturesOutOfOrder) {
  LLVMContext Ctx;
  std::vector<TensorSpec> Inputs{
      TensorSpec::createSpec<int64_t>("a", {1}),
      TensorSpec::createSpec<int64_t>("c", {1}),
      TensorSpec::createSpec<int64_t>("b", {1}),
  };
  auto Evaluator = std::make_unique<ReleaseModeModelRunner<MockAOTModel>>(
      Ctx, Inputs, "", "prefix_");
  *Evaluator->getTensor<int64_t>(0) = 1;         // a
  *Evaluator->getTensor<int64_t>(1) = 2;         // c
  *Evaluator->getTensor<int64_t>(2) = -3;        // b
  EXPECT_EQ(Evaluator->evaluate<int64_t>(), -2); // a + b
  EXPECT_EQ(*Evaluator->getTensor<int64_t>(0), 1);
  EXPECT_EQ(*Evaluator->getTensor<int64_t>(1), 2);
  EXPECT_EQ(*Evaluator->getTensor<int64_t>(2), -3);
}