| /* |
| * Licensed to the Apache Software Foundation (ASF) under one |
| * or more contributor license agreements. See the NOTICE file |
| * distributed with this work for additional information |
| * regarding copyright ownership. The ASF licenses this file |
| * to you 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 <dmlc/logging.h> |
| #include <gtest/gtest.h> |
| #include <tvm/node/object_path.h> |
| #include <tvm/runtime/logging.h> |
| #include <tvm/script/printer/frame.h> |
| #include <tvm/script/printer/var_table.h> |
| #include <tvm/tir/var.h> |
| |
| using namespace tvm; |
| using namespace tvm::script::printer; |
| |
| TEST(PrinterVarTableTest, Define) { |
| VarTable vars; |
| MetadataFrame frame; |
| tir::Var x("x"); |
| ObjectPath object_path = ObjectPath::Root(); |
| |
| IdDoc doc = vars->Define(x, "x", object_path, frame); |
| |
| ICHECK_EQ(doc->name, "x"); |
| |
| IdDoc second_doc = Downcast<IdDoc>(vars->GetVarDoc(x, object_path).value()); |
| |
| ICHECK_EQ(second_doc->name, "x"); |
| } |
| |
| TEST(PrinterVarTableTest, DefineByDoc) { |
| VarTable vars; |
| MetadataFrame frame; |
| tir::Var x("x"); |
| ObjectPath object_path = ObjectPath::Root(); |
| |
| auto doc_factory = []() { return LiteralDoc::Str("x"); }; |
| |
| vars->DefineByDoc(x, doc_factory, frame); |
| |
| ExprDoc doc = vars->GetVarDoc(x, object_path).value(); |
| |
| ICHECK_EQ(Downcast<String>(Downcast<LiteralDoc>(doc)->value), "x"); |
| } |
| |
| TEST(PrinterVarTableTest, GetVarDocWithUnknownVariable) { |
| VarTable vars; |
| MetadataFrame frame; |
| tir::Var x("x"); |
| tir::Var y("y"); |
| ObjectPath object_path = ObjectPath::Root(); |
| |
| Doc doc = vars->Define(x, "x", object_path, frame); |
| ICHECK(!vars->GetVarDoc(y, object_path).defined()); |
| } |
| |
| TEST(PrinterVarTableTest, GetVarDocWithObjectPath) { |
| VarTable vars; |
| MetadataFrame frame; |
| tir::Var x("x"); |
| ObjectPath object_path = ObjectPath::Root(); |
| ObjectPath second_object_path = ObjectPath::Root()->Attr("x"); |
| |
| IdDoc doc = vars->Define(x, "x", object_path, frame); |
| ICHECK_EQ(doc->source_paths[0], object_path); |
| ICHECK_EQ(doc->source_paths.size(), 1); |
| |
| Doc second_doc = vars->GetVarDoc(x, second_object_path).value(); |
| ICHECK_EQ(second_doc->source_paths[0], second_object_path); |
| ICHECK_EQ(second_doc->source_paths.size(), 1); |
| } |
| |
| TEST(PrinterVarTableTest, IsVarDefined) { |
| VarTable vars; |
| MetadataFrame frame; |
| tir::Var x("x"); |
| tir::Var y("y"); |
| ObjectPath object_path = ObjectPath::Root(); |
| |
| vars->Define(x, "x", object_path, frame); |
| ICHECK(vars->IsVarDefined(x)); |
| ICHECK(!vars->IsVarDefined(y)); |
| } |
| |
| TEST(PrinterVarTableTest, VarRemovedAfterFrameOutOfScope) { |
| VarTable vars; |
| MetadataFrame frame; |
| tir::Var x("x"); |
| ObjectPath object_path = ObjectPath::Root(); |
| |
| vars->Define(x, "x", object_path, frame); |
| ICHECK(vars->IsVarDefined(x)); |
| |
| frame->ExitWithScope(); |
| ICHECK(!vars->IsVarDefined(x)); |
| } |
| |
| TEST(PrinterVarTableTest, DefineDuplicateName) { |
| VarTable vars; |
| MetadataFrame frame; |
| tir::Var x("x"); |
| tir::Var y("y"); |
| ObjectPath object_path = ObjectPath::Root(); |
| |
| IdDoc x_doc = vars->Define(x, "x", object_path, frame); |
| IdDoc y_doc = vars->Define(y, "x", object_path, frame); |
| |
| ICHECK_NE(x_doc->name, y_doc->name); |
| } |
| |
| TEST(PrinterVarTableTest, DefineDuplicateVariable) { |
| VarTable vars; |
| MetadataFrame frame; |
| tir::Var x("x"); |
| ObjectPath object_path = ObjectPath::Root(); |
| |
| vars->Define(x, "x", object_path, frame); |
| |
| bool failed = false; |
| try { |
| vars->Define(x, "x", object_path, frame); |
| } catch (...) { |
| failed = true; |
| } |
| ASSERT_EQ(failed, true); |
| } |
| |
| TEST(PrinterVarTableTest, DefineByDocWithIdDoc) { |
| VarTable vars; |
| MetadataFrame frame; |
| tir::Var x("x"); |
| ObjectPath object_path = ObjectPath::Root(); |
| |
| bool failed = false; |
| try { |
| // User has to use `Define` if variable needs to be mapped to IdDoc |
| vars->DefineByDoc( |
| x, []() { return IdDoc("x"); }, frame); |
| } catch (...) { |
| failed = true; |
| } |
| ASSERT_EQ(failed, true); |
| } |