wehub-resource-sync
97 lines
3.7 KiB
Python
97 lines
3.7 KiB
Python
# 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.
|
|
"""Tests for tirx.analysis.undefined_vars (VarUseDefAnalyzer)."""
|
|
|
|
import tvm
|
|
import tvm.testing
|
|
from tvm import tirx
|
|
|
|
|
|
def test_decl_buffer_data_is_use():
|
|
"""DeclBuffer's data var should be reported as undefined (USE), not defined.
|
|
|
|
When UndefinedVars encounters a DeclBuffer, the data pointer references
|
|
an existing variable from the enclosing scope. It must appear in the
|
|
undefined list so that callers (e.g., CreateComputeScope) capture it.
|
|
"""
|
|
n = tirx.Var("n", "int32")
|
|
from tvm.ir import PointerType, PrimType
|
|
|
|
data_ptr = tirx.Var("buf_data", PointerType(PrimType("float32")))
|
|
buf = tirx.decl_buffer((n,), "float32", "buf", data=data_ptr)
|
|
|
|
body = tirx.Evaluate(tirx.BufferLoad(buf, [0]))
|
|
decl = tirx.DeclBuffer(buf)
|
|
stmt = tirx.SeqStmt([decl, body])
|
|
|
|
undef = tvm.tirx.analysis.undefined_vars(stmt, [])
|
|
undef_names = {v.name for v in undef}
|
|
# data_ptr must be undefined (it comes from outside the DeclBuffer)
|
|
assert "buf_data" in undef_names, f"Expected buf_data in undefined vars, got {undef_names}"
|
|
|
|
|
|
def test_decl_buffer_elem_offset_is_use():
|
|
"""DeclBuffer's elem_offset var should be reported as undefined (USE).
|
|
|
|
After FlattenBuffer, DeclBuffer nodes carry elem_offset vars from
|
|
match_buffer entries. These must appear in the undefined list.
|
|
"""
|
|
from tvm.ir import PointerType, PrimType
|
|
|
|
n = tirx.Var("n", "int32")
|
|
data_ptr = tirx.Var("buf_data", PointerType(PrimType("float32")))
|
|
elem_off = tirx.Var("buf_elem_offset", "int32")
|
|
buf = tirx.decl_buffer((n,), "float32", "buf", data=data_ptr, elem_offset=elem_off)
|
|
|
|
body = tirx.Evaluate(tirx.BufferLoad(buf, [0]))
|
|
decl = tirx.DeclBuffer(buf)
|
|
stmt = tirx.SeqStmt([decl, body])
|
|
|
|
undef = tvm.tirx.analysis.undefined_vars(stmt, [])
|
|
undef_names = {v.name for v in undef}
|
|
assert "buf_data" in undef_names, f"Expected buf_data in undefined vars, got {undef_names}"
|
|
assert "buf_elem_offset" in undef_names, (
|
|
f"Expected buf_elem_offset in undefined vars, got {undef_names}"
|
|
)
|
|
|
|
|
|
def test_alloc_buffer_data_is_def():
|
|
"""AllocBuffer's data var should NOT be reported as undefined (it's a DEF).
|
|
|
|
AllocBuffer allocates new storage — the data pointer is a new definition,
|
|
not a reference to an external variable.
|
|
"""
|
|
n = tirx.Var("n", "int32")
|
|
buf = tirx.decl_buffer((n,), "float32", "buf")
|
|
|
|
body = tirx.Evaluate(tirx.BufferLoad(buf, [0]))
|
|
alloc = tirx.AllocBuffer(buf)
|
|
stmt = tirx.SeqStmt([alloc, body])
|
|
|
|
undef = tvm.tirx.analysis.undefined_vars(stmt, [])
|
|
undef_names = {v.name for v in undef}
|
|
# data should NOT be undefined — AllocBuffer defines it
|
|
assert buf.data.name not in undef_names, (
|
|
f"AllocBuffer data should be defined, but found {buf.data.name} in {undef_names}"
|
|
)
|
|
# shape var n should be undefined (comes from enclosing scope)
|
|
assert "n" in undef_names, f"Expected shape var 'n' in undefined vars, got {undef_names}"
|
|
|
|
|
|
if __name__ == "__main__":
|
|
tvm.testing.main()
|