chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,244 @@
|
||||
# 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.
|
||||
# ruff: noqa: F841
|
||||
|
||||
|
||||
import tvm
|
||||
import tvm.testing
|
||||
from tvm.script import relax as R
|
||||
from tvm.script import tirx as T
|
||||
|
||||
|
||||
def _analyze_func(func: tvm.relax.Function) -> list[str]:
|
||||
return [var.name_hint for var in tvm.relax.analysis.computable_at_compile_time(func)]
|
||||
|
||||
|
||||
def test_no_num_input_attribute():
|
||||
"""Without the "num_input" attribute, all params are runtime"""
|
||||
|
||||
@R.function
|
||||
def func(A: R.Tensor([16], "int32"), B: R.Tensor([16], "int32")):
|
||||
C = R.add(A, B)
|
||||
return C
|
||||
|
||||
assert _analyze_func(func) == []
|
||||
|
||||
|
||||
def test_compile_time_param():
|
||||
"""Parameters after "num_input" are known at compile-time"""
|
||||
|
||||
@R.function
|
||||
def func(A: R.Tensor([16], "int32"), B: R.Tensor([16], "int32")):
|
||||
R.func_attr({"num_input": 1})
|
||||
return ()
|
||||
|
||||
assert _analyze_func(func) == ["B"]
|
||||
|
||||
|
||||
def test_binding_using_one_param():
|
||||
"""Bindings may be computable at compile-time"""
|
||||
|
||||
@R.function
|
||||
def func(A: R.Tensor([16], "int32"), B: R.Tensor([16], "int32")):
|
||||
R.func_attr({"num_input": 1})
|
||||
C = R.add(B, B)
|
||||
D = R.add(A, C)
|
||||
return D
|
||||
|
||||
assert _analyze_func(func) == ["B", "C"]
|
||||
|
||||
|
||||
def test_binding_using_multiple_params():
|
||||
"""Compile-time bindings may use multiple parameters"""
|
||||
|
||||
@R.function
|
||||
def func(A: R.Tensor([16], "int32"), B: R.Tensor([16], "int32"), C: R.Tensor([16], "int32")):
|
||||
R.func_attr({"num_input": 1})
|
||||
D = R.add(B, C)
|
||||
E = R.add(A, D)
|
||||
return E
|
||||
|
||||
assert _analyze_func(func) == ["B", "C", "D"]
|
||||
|
||||
|
||||
def test_compile_time_binding_after_run_time():
|
||||
"""Compile-time bindings may occur after run-time
|
||||
|
||||
A binding being computable at compile-time only depends on the
|
||||
arguments used for it. A value that is computable at compile-time
|
||||
may occur after a value that is only computable at run-time.
|
||||
"""
|
||||
|
||||
@R.function
|
||||
def func(A: R.Tensor([16], "int32"), B: R.Tensor([16], "int32")):
|
||||
R.func_attr({"num_input": 1})
|
||||
C = R.add(A, A)
|
||||
D = R.add(B, B)
|
||||
E = R.add(C, D)
|
||||
return E
|
||||
|
||||
assert _analyze_func(func) == ["B", "D"]
|
||||
|
||||
|
||||
def test_sequential_compile_time_bindings():
|
||||
"""Compile-time bindings may occur after run-time
|
||||
|
||||
A compile-time value may depend on variables defined within the
|
||||
function, so long as those variables are themselves computable at
|
||||
compile-time.
|
||||
"""
|
||||
|
||||
@R.function
|
||||
def func(A: R.Tensor([16], "int32"), B: R.Tensor([16], "int32")):
|
||||
R.func_attr({"num_input": 1})
|
||||
C = R.add(B, B)
|
||||
D = R.add(C, C)
|
||||
E = R.add(D, D)
|
||||
F = R.add(E, A)
|
||||
return F
|
||||
|
||||
assert _analyze_func(func) == ["B", "C", "D", "E"]
|
||||
|
||||
|
||||
def test_dataflow_vars():
|
||||
"""Compile-time bindings may occur in dataflow blocks"""
|
||||
|
||||
@R.function
|
||||
def func(A: R.Tensor([16], "int32"), B: R.Tensor([16], "int32")):
|
||||
R.func_attr({"num_input": 1})
|
||||
with R.dataflow():
|
||||
C = R.add(B, B)
|
||||
D = R.add(C, C)
|
||||
E = R.add(D, D)
|
||||
F = R.add(E, A)
|
||||
R.output(F)
|
||||
return F
|
||||
|
||||
assert _analyze_func(func) == ["B", "C", "D", "E"]
|
||||
|
||||
|
||||
def test_compile_time_symbolic_shape():
|
||||
"""Compile-time bindings may contain symbolic shapes"""
|
||||
|
||||
@R.function
|
||||
def func(A: R.Tensor([1], "int32"), B: R.Tensor(["n"], "int32")):
|
||||
R.func_attr({"num_input": 1})
|
||||
n = T.int64()
|
||||
|
||||
C: R.Tensor([n], "int32") = R.add(B, B)
|
||||
D: R.Tensor([], "int32") = R.max(C, axis=0)
|
||||
E: R.Tensor([1], "int32") = R.add(A, D)
|
||||
return E
|
||||
|
||||
assert _analyze_func(func) == ["B", "C", "D"]
|
||||
|
||||
|
||||
def test_symbolic_variables_from_match_binding():
|
||||
"""Symbolic vars may be inferred from compile-time bindings"""
|
||||
|
||||
@R.function
|
||||
def func(A: R.Tensor(ndim=1, dtype="int32"), B: R.Tensor(ndim=1, dtype="int32")):
|
||||
R.func_attr({"num_input": 1})
|
||||
n = T.int64()
|
||||
m = T.int64()
|
||||
|
||||
A2 = R.match_cast(A, R.Tensor([n], "int32"))
|
||||
B2 = R.match_cast(B, R.Tensor([m], "int32"))
|
||||
|
||||
C = R.add(B2, B2)
|
||||
D = R.max(C, axis=0)
|
||||
E = R.max(A2, axis=0)
|
||||
F = R.add(D, E)
|
||||
return F
|
||||
|
||||
assert _analyze_func(func) == ["B", "B2", "C", "D"]
|
||||
|
||||
|
||||
def test_compile_time_expressions_may_not_use_runtime_symbolic_variables():
|
||||
"""Symbolic vars may be inferred from compile-time bindings
|
||||
|
||||
Here, `C` uses the symbolic variable `m`, which can be inferred
|
||||
from the shape of `B` and is known at compile-time. However, `D`
|
||||
uses the symbolic variable `n`, which cannot be inferred without
|
||||
first knowing `A`, and is therefore unknown at compile-time.
|
||||
"""
|
||||
|
||||
@R.function
|
||||
def func(A: R.Tensor(["n"], "int32"), B: R.Tensor(["m"], "int32")):
|
||||
R.func_attr({"num_input": 1})
|
||||
n = T.int64()
|
||||
m = T.int64()
|
||||
|
||||
C = R.ones([m], "int32")
|
||||
D = R.ones([n], "int32")
|
||||
|
||||
E = (C, D)
|
||||
return E
|
||||
|
||||
assert _analyze_func(func) == ["B", "C"]
|
||||
|
||||
|
||||
def test_compile_time_expressions_may_infer_same_variable_as_run_time():
|
||||
"""Symbolic vars may be inferred from compile-time bindings
|
||||
|
||||
A symbolic variable may be inferrable from multiple sources.
|
||||
Here, while `n` can be inferred from the runtime parameter `A`, it
|
||||
can also be inferred from the compile-time parameter `B`.
|
||||
"""
|
||||
|
||||
@R.function
|
||||
def func(A: R.Tensor(["n"], "int32"), B: R.Tensor(["n"], "int32")):
|
||||
R.func_attr({"num_input": 1})
|
||||
n = T.int64()
|
||||
|
||||
C = R.ones([n], "int32")
|
||||
D = R.ones([n], "int32")
|
||||
|
||||
E = (C, D)
|
||||
return E
|
||||
|
||||
assert _analyze_func(func) == ["B", "C", "D", "E"]
|
||||
|
||||
|
||||
def test_compile_time_expressions_may_use_variables_from_match_cast():
|
||||
"""Symbolic vars may be inferred from compile-time bindings
|
||||
|
||||
Here, `C` uses the symbolic variable `m`, which can be inferred
|
||||
from the shape of `B` and is known at compile-time. However, `D`
|
||||
uses the symbolic variable `n`, which cannot be inferred without
|
||||
first knowing `A`, and is therefore unknown at compile-time.
|
||||
"""
|
||||
|
||||
@R.function
|
||||
def func(A: R.Tensor(["n"], "int32"), B: R.Tensor(ndim=1, dtype="int32")):
|
||||
R.func_attr({"num_input": 1})
|
||||
n = T.int64()
|
||||
m = T.int64()
|
||||
|
||||
B2 = R.match_cast(B, R.Tensor([m], "int32"))
|
||||
|
||||
C = R.ones([m], "int32")
|
||||
D = R.ones([n], "int32")
|
||||
|
||||
E = (C, D)
|
||||
return E
|
||||
|
||||
assert _analyze_func(func) == ["B", "B2", "C"]
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
tvm.testing.main()
|
||||
Reference in New Issue
Block a user