chore: import upstream snapshot with attribution

python/tvm/tirx/analysis/__init__.py

@@ -0,0 +1,20 @@
+# 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.
+"""Namespace of all TIR analysis utils."""
+# pylint: disable=wildcard-import, invalid-name
+
+from .analysis import *

python/tvm/tirx/analysis/_ffi_api.py

@@ -0,0 +1,21 @@
+# 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.
+"""FFI APIs for tvm.tirx.analysis"""
+
+import tvm_ffi
+
+tvm_ffi.init_ffi_api("tirx.analysis", __name__)

python/tvm/tirx/analysis/analysis.py

@@ -0,0 +1,160 @@
+# 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.
+"""Wrapping existing analysis utils."""
+
+# pylint: disable=invalid-name
+
+from tvm.ir import IRModule
+from tvm.tirx.expr import Var
+from tvm.tirx.stmt import Expr
+
+from .. import Stmt
+from ..function import PrimFunc
+from . import _ffi_api
+
+
+def expr_deep_equal(lhs: Expr, rhs: Expr) -> bool:
+    """Deeply compare two nested expressions.
+
+    Parameters
+    ----------
+    lhs : Expr
+        The left operand.
+
+    rhs : Expr
+        The right operand.
+
+    Returns
+    -------
+    result : bool
+        The comparison result
+
+    Note
+    ----
+
+    This function does not remap variable bindings, it will not
+    return true for (let x = 1 in x + 1) vs (let y = 1 in y + 1), unless x.same_as(y).
+    Use py:func:`tvm_ffi.structural_equal` to handle structural variable remapping.
+
+    Due to the restriction of not remapping variables, this function can run
+    faster than StructuralEqual and can be used as a utility function during arithmetic
+    simplifications.
+
+    Always consider py:func:`tvm_ffi.structural_equal` first, which handles
+    the structural remapping.
+
+    See Also
+    --------
+    tvm_ffi.structural_equal
+    """
+    return _ffi_api.expr_deep_equal(lhs, rhs)  # type: ignore
+
+
+def verify_ssa(func: PrimFunc) -> bool:
+    """Verify if the func is in SSA form.
+
+    Parameters
+    ----------
+    func: tvm.tirx.PrimFunc
+        The module to be verified.
+
+    Returns
+    -------
+    result : bool
+        The result of verification.
+    """
+    return _ffi_api.verify_ssa(func)  # type: ignore
+
+
+def verify_memory(func: PrimFunc) -> bool:
+    """Verify if func contains illegal host side direct memory access.
+
+    Parameters
+    ----------
+    func: tvm.tirx.PrimFunc
+        The module to be verified.
+
+    Returns
+    -------
+    result : bool
+        The result of verification.
+    """
+    return _ffi_api.verify_memory(func)  # type: ignore
+
+
+def undefined_vars(node: Stmt | Expr, defs: list[Var] | None = None) -> list[Var]:
+    """Find undefined vars in a TIR statement or expression.
+
+    Parameters
+    ----------
+    node: Union[Stmt, Expr]
+        The TIR statement or expression to be checked.
+
+    defs: Optional[List[Var]]
+        The vars that is defined
+
+    Returns
+    -------
+    result : List[Var]
+        The undefined vars.
+    """
+    defs = defs or []
+    return _ffi_api.UndefinedVars(node, defs)  # type: ignore # pylint: disable=no-member
+
+
+def verify_well_formed(obj: PrimFunc | IRModule, assert_mode: bool = True) -> bool:
+    """Verify if the given TIR is well-formed. The verification includes:
+        - Check if expressions not contain vars that is defined outside the block.
+
+    Parameters
+    ----------
+    obj: Union[tvm.tirx.PrimFunc, tvm.ir.IRModule]
+        The function or module to be verified.
+
+    assert_mode: bool
+        The indicator if it raises an error when the function is not well-formed.
+
+    Returns
+    -------
+    result: bool
+        Whether it is a well-formed TIR function.
+    """
+    return _ffi_api.VerifyWellFormed(obj, assert_mode)  # type: ignore # pylint: disable=no-member
+
+
+def verify_tirx_well_formed(
+    obj: PrimFunc | IRModule, assert_mode: bool = True, device_func: bool = False
+) -> bool:
+    """Verify if the given TIRX is well-formed.
+
+    Parameters
+    ----------
+    obj: Union[tvm.tirx.PrimFunc, tvm.ir.IRModule]
+        The function or module to be verified.
+
+    assert_mode: bool
+        The indicator if it raises an error when the function is not well-formed.
+
+    device_func: bool
+        The indicator if it is a device function.
+
+    Returns
+    -------
+    result: bool
+        Whether it is a well-formed TIRX function.
+    """
+    return _ffi_api.VerifyTIRxWellFormed(obj, assert_mode, device_func)  # type: ignore # pylint: disable=no-member