mlir_loader_test.cc 1.7 KB
Newer Older
Y
Yan Chunwei 已提交
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19
// Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved.
//
// 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
//
//     http://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.

#include "paddle/infrt/dialect/mlir_loader.h"

#include <glog/logging.h>
#include <gtest/gtest.h>
#include <llvm/Support/SourceMgr.h>
20
#include <mlir/IR/BuiltinTypes.h>
Y
Yan Chunwei 已提交
21 22 23 24 25 26
#include <mlir/Parser.h>

#include <string>

#include "paddle/infrt/dialect/init_infrt_dialects.h"

27 28
namespace infrt {
namespace dialect {
Y
Yan Chunwei 已提交
29 30 31 32 33 34

TEST(MlirLoader, basic) {
  mlir::MLIRContext context;

  auto source = R"ROC(
func @main() -> f32 {
35 36 37
  %v0 = Infrt.constant.f32 1.0
  %v1 = Infrt.constant.f32 2.0
  %value = "Infrt.add.f32"(%v0, %v1) : (f32, f32) -> f32
Y
Yan Chunwei 已提交
38

39
  "Infrt.print.f32"(%v0) : (f32) -> ()
Y
Yan Chunwei 已提交
40

41
  Infrt.return %value : f32
Y
Yan Chunwei 已提交
42 43 44 45
}
)ROC";

  auto module = LoadMlirSource(&context, source);
46
  EXPECT_TRUE(mlir::succeeded(module->verify()));
Y
Yan Chunwei 已提交
47 48 49 50 51 52 53 54 55 56
  LOG(INFO) << "module name: " << module->getOperationName().data();
  for (auto func : module->getOps<mlir::FuncOp>()) {
    LOG(INFO) << "get func " << func.getName().str();
    int num_args = func.getNumArguments();
    for (int i = 0; i < num_args; i++) {
      LOG(INFO) << "arg: " << func.getArgument(i).getArgNumber();
    }
  }
}

57 58
}  // namespace dialect
}  // namespace infrt