This is an automated email from the ASF dual-hosted git repository.
tqchen pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm-tirx.git
The following commit(s) were added to refs/heads/main by this push:
new ba96d293 [REFACTOR] Keep Triton constexpr integer conversion local (#4)
ba96d293 is described below
commit ba96d293999dcab3b11b99f410f5a6bf2768cf94
Author: Tianqi Chen <[email protected]>
AuthorDate: Tue Oct 20 16:55:57 2026 -0400
[REFACTOR] Keep Triton constexpr integer conversion local (#4)
Keep Triton constexpr integer conversion in its sole caller while
preserving constant extraction and error behavior.
Remove the unused public symbolic utility module and its reference
documentation.
---
docs/reference/python/sym.rst | 8 -----
docs/reference/python/sym.utils.rst | 32 -----------------
python/tvm_tirx/sym/utils.py | 46 ------------------------
python/tvm_tirx/tirx/script/ir_builder/triton.py | 14 +++++---
4 files changed, 10 insertions(+), 90 deletions(-)
diff --git a/docs/reference/python/sym.rst b/docs/reference/python/sym.rst
index 703371cd..0798c82a 100644
--- a/docs/reference/python/sym.rst
+++ b/docs/reference/python/sym.rst
@@ -59,14 +59,6 @@ Functions
subspace_divide
inverse_affine_iter_map
-Namespaces
-----------
-
-.. toctree::
- :maxdepth: 1
-
- sym.utils
-
Supporting classes
------------------
diff --git a/docs/reference/python/sym.utils.rst
b/docs/reference/python/sym.utils.rst
deleted file mode 100644
index 93b4d4b2..00000000
--- a/docs/reference/python/sym.utils.rst
+++ /dev/null
@@ -1,32 +0,0 @@
-.. 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.
-
-tvm_tirx.sym.utils
-==================
-
-.. automodule:: tvm_tirx.sym.utils
- :no-members:
-
-.. currentmodule:: tvm_tirx.sym.utils
-
-Functions
----------
-
-.. autosummary::
- :toctree: generated/
-
- get_const_int
diff --git a/python/tvm_tirx/sym/utils.py b/python/tvm_tirx/sym/utils.py
deleted file mode 100644
index 4cf1457a..00000000
--- a/python/tvm_tirx/sym/utils.py
+++ /dev/null
@@ -1,46 +0,0 @@
-# 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.
-"""Utilities for extracting values from symbolic expressions."""
-
-from numbers import Integral
-
-from tvm_tirx.ir import IntImm
-
-from .analyzer import Analyzer
-
-
-def get_const_int(expr):
- """Verifies expr is integer and get the constant value.
-
- Parameters
- ----------
- expr : tvm.Expr or int
- The input expression.
-
- Returns
- -------
- out_value : int
- The output.
- """
- if isinstance(expr, Integral):
- return expr
- if not isinstance(expr, IntImm):
- ana = Analyzer()
- expr = ana.simplify(expr)
- if not isinstance(expr, IntImm):
- raise ValueError("Expect value to be constant int")
- return int(expr.value)
diff --git a/python/tvm_tirx/tirx/script/ir_builder/triton.py
b/python/tvm_tirx/tirx/script/ir_builder/triton.py
index ba7582bf..53cd6447 100644
--- a/python/tvm_tirx/tirx/script/ir_builder/triton.py
+++ b/python/tvm_tirx/tirx/script/ir_builder/triton.py
@@ -17,6 +17,7 @@
# ruff: noqa: RUF005
"""Triton kernel integration with TIR"""
+from numbers import Integral
from typing import Any
import triton
@@ -25,9 +26,8 @@ from packaging import version
from triton.runtime.jit import type_canonicalisation_dict
import tvm_tirx
-from tvm_tirx import tirx
-from tvm_tirx.ir import PrimType, PtrType, is_prim_expr
-from tvm_tirx.sym.utils import get_const_int
+from tvm_tirx import sym, tirx
+from tvm_tirx.ir import IntImm, PrimType, PtrType, is_prim_expr
from .external_kernel import BaseKernel
@@ -117,7 +117,13 @@ class TritonKernel(BaseKernel):
kernel_args = [] # Arguments to invoke the kernel
for i, arg in enumerate(args):
if kernel_params[i].is_constexpr:
- constants[kernel_params[i].name] = get_const_int(arg)
+ if not isinstance(arg, Integral):
+ if not isinstance(arg, IntImm):
+ arg = sym.Analyzer().simplify(arg)
+ if not isinstance(arg, IntImm):
+ raise ValueError("Expect value to be constant int")
+ arg = int(arg.value)
+ constants[kernel_params[i].name] = arg
signature[kernel_params[i].name] = "constexpr"
continue
if isinstance(arg.ty, PtrType):
