chore: fixup pylint complaints

Signed-off-by: Maryam Tahhan <mtahhan@redhat.com>

Author: maryamtahhan

.github/workflows/build-triton-util.yml

@@ -0,0 +1,38 @@
+name: Build triton-util
+
+on:
+  push:
+    branches: [main]
+    paths:
+      - triton_util/**
+      - .github/workflows/build.yml
+  pull_request:
+    paths:
+      - triton_util/**
+
+jobs:
+  build:
+    runs-on: ubuntu-latest
+
+    steps:
+      - name: Checkout code
+        uses: actions/checkout@v4
+
+      - name: Set up Python
+        uses: actions/setup-python@v5
+        with:
+          python-version: 3.12
+
+      - name: Upgrade pip and install tools
+        run: |
+          python -m pip install --upgrade pip
+          pip install build setuptools wheel
+
+      - name: Build the package
+        run: |
+          cd triton_util
+          python -m build
+
+      - name: Install the package
+        run: |
+          pip install ./triton_util[dev]

.github/workflows/pylint.yml

@@ -3,10 +3,12 @@ name: Lint Python code with pylint
 on: # yamllint disable-line rule:truthy
   pull_request:
     paths:
-      - "**/*.py"
+      - tcm/**/*.py
+      - triton_util/**/*.py
   push:
     paths:
-      - "**/*.py"
+      - tcm/**/*.py
+      - triton_util/**/*.py
 
 jobs:
   pylint:
@@ -26,6 +28,7 @@ jobs:
           python -m pip install --upgrade pip
           pip install pylint
           pip install -r ./tcm/requirements.txt
+          pip install -e ./triton_util[dev]
 
       - name: Run pylint
         run: |

triton_util/setup.py

@@ -1,20 +1,28 @@
+"""Setup script for the triton-util package."""
+
 from setuptools import setup, find_packages
 
+with open("README.md", encoding="utf-8") as f:
+    long_description = f.read()
+
 setup(
-    name='triton-util',
-    version='0.0.2',
+    name="triton-util",
+    version="0.0.2",
     packages=find_packages(),
-    install_requires=['triton'],
-    author='Umer Adil',
-    author_email='umer.hayat.adil@gmail.com',
-    description='Make Triton easier - A utility package for OpenAI Triton',
-    long_description=open('README.md').read(),
-    long_description_content_type='text/markdown',
-    url='https://github.com/umerHA/triton_util',
+    install_requires=["triton"],
+    extras_require={
+        "dev": ["pytest", "pylint", "torch", "ipython"],
+    },
+    author="Umer Adil",
+    author_email="umer.hayat.adil@gmail.com",
+    description="Make Triton easier - A utility package for OpenAI Triton",
+    long_description=long_description,
+    long_description_content_type="text/markdown",
+    url="https://github.com/redhat-et/TKDK/triton_util",
     classifiers=[
-        'Programming Language :: Python :: 3',
-        'License :: OSI Approved :: MIT License',
-        'Operating System :: OS Independent',
+        "Programming Language :: Python :: 3",
+        "License :: OSI Approved :: MIT License",
+        "Operating System :: OS Independent",
     ],
-    python_requires='>=3.6',
+    python_requires=">=3.12",
 )

triton_util/tests/conftest.py

@@ -1,3 +1,5 @@
+"""Pytest configuration file with Triton interpreter mode fixture."""
+
 import os
 import importlib
 
@@ -6,9 +8,14 @@
 import triton
 import triton.language as tl
 
-@pytest.fixture(scope='class', params=['0', '1']) # Run tests in regular mode (TRITON_INTERPRET=0) and in interpreter mode (TRITON_INTERPRET=1)
+@pytest.fixture(scope='class', params=['0', '1'])
 def triton_interpret(request):
-    '''Set env var TRITON_INTERPRET and reload triton'''
+    """
+    Test Triton in both regular mode (TRITON_INTERPRET=0) and interpreter mode (TRITON_INTERPRET=1)
+
+    Sets the TRITON_INTERPRET environment variable to either "0" or "1",
+    reloads the Triton modules, and ensures the env var is cleaned up afterward.
+    """
     os.environ['TRITON_INTERPRET'] = request.param
     importlib.reload(triton)
     importlib.reload(tl)

triton_util/tests/test_coding.py

@@ -1,5 +1,5 @@
-import inspect
-
+"""Test suite for triton_util coding utilities."""
+# pylint: disable=missing-function-docstring,unused-argument,undefined-variable
 import pytest
 
 import torch
@@ -10,7 +10,9 @@
 import triton_util as tu
 
 class TestCodingUtils:
+    """Unit tests for triton_util offset, mask, load, and store utilities."""
     def test_cdiv(self):
+        """Test cdiv (ceiling division) function."""
         assert tu.cdiv(10, 2)==5
         assert tu.cdiv(10, 3)==4
 
@@ -60,7 +62,7 @@ def test_mask_1d(self, triton_interpret):
 
         @triton.jit
         def partial_copy(i_ptr, o_ptr, n):
-            offs = n*2 + tl.arage(0,2)
+            offs = n*2 + tl.arange(0,2)
             mask = mask_1d(offs, 4)
             vals = tl.load(i_ptr + offs, mask)
             tl.store(o_ptr + offs, vals, mask)
@@ -115,7 +117,7 @@ def test_load_full_1d(self, triton_interpret):
 
         @triton.jit
         def copy(i_ptr, o_ptr):
-            offs = tl.arage(0,4)
+            offs = tl.arange(0,4)
             mask = offs < 4
             vals = tu.load_full_1d(i_ptr, 4)
             tl.store(o_ptr + offs, vals, mask)

triton_util/tests/test_debugging.py

@@ -1,3 +1,8 @@
+"""
+Debugging utilities for Triton GPU kernels.
+"""
+# pylint: disable=multiple-statements,unused-argument,no-value-for-parameter,missing-class-docstring,missing-module-docstring,missing-function-docstring
+
 import pytest
 
 import torch

triton_util/tests/test_loading.py

@@ -1,3 +1,4 @@
+# pylint: disable=line-too-long,missing-module-docstring,missing-function-docstring,multiple-statements,too-few-public-methods,unused-import
 import os
 
 import pytest

triton_util/triton_util/__init__.py

@@ -1,2 +1,36 @@
+"""
+Triton-Util 🔱 – Utility functions for writing Triton GPU kernels with less friction.
+
+This package provides high-level abstractions and helpers for writing
+fast and readable Triton code, reducing the need for repetitive and error-prone
+index calculations.
+
+Features:
+- **Coding utilities**: Chunking, masking, offset computation, and bulk load/store helpers.
+- **Debugging utilities**: Convenient debugging hooks like `print_once` and `breakpoint_if`
+  that make inspecting kernel behavior simpler and less intrusive.
+
+The utilities are designed to be:
+- Minimal and interoperable: fully compatible with native Triton code.
+- Expressive: match how you actually think about GPU data access patterns.
+- Progressive: use as little or as much of the library as needed.
+
+Example usage:
+    >>> load_2d(ptr, sz0, sz1, n0, n1, max0, max1, stride0)
+
+instead of:
+    >>> offs0 = n0 * sz0 + tl.arange(0, sz0)
+    >>> offs1 = n1 * sz1 + tl.arange(0, sz1)
+    >>> offs = offs0[:, None] * stride0 + offs1[None, :] * stride1
+    >>> mask = (offs0[:, None] < max0) & (offs1[None, :] < max1)
+    >>> tl.load(ptr + offs, mask)
+
+For documentation, examples, and community support, see:
+- GitHub: https://github.com/cuda-mode/triton-util
+- Discord: https://discord.gg/cudamode (Triton channel)
+
+Author: Umer Hadil
+"""
+
 from .debugging import *
 from .coding import *

triton_util/triton_util/coding.py

@@ -1,30 +1,46 @@
+"""Triton utility functions for offset calculation, masking, and load/store operations."""
+# pylint: disable=too-many-arguments,too-many-positional-arguments,redefined-builtin,unused-argument
+
 import triton
 import triton.language as tl
 from triton.language import constexpr as const
 
-def cdiv(a,b): return (a + b - 1) // b
+def cdiv(a, b):
+    """Ceiling division."""
+    return (a + b - 1) // b
 
 # # offsets
 
 @triton.jit
-def offset_1d(sz: const, n_prev_chunks=0): return n_prev_chunks * sz + tl.arange(0, sz)
+def offset_1d(sz: const, n_prev_chunks=0):
+    """Compute 1D offset based on chunk size and previous chunks."""
+    return n_prev_chunks * sz + tl.arange(0, sz)
 
 @triton.jit
-def offset_2d(offs0, offs1, stride0, stride1=1):  return tl.expand_dims(offs0, 1)*stride0 + tl.expand_dims(offs1, 0)*stride1
+def offset_2d(offs0, offs1, stride0, stride1=1):
+    """Compute 2D offset using strides."""
+    return tl.expand_dims(offs0, 1)*stride0 + tl.expand_dims(offs1, 0)*stride1
 
 # # masks
 
 @triton.jit
-def mask_1d(offs, max): return offs < max
+def mask_1d(offs, max):
+    """Create a 1D mask based on a max bound."""
+    return offs < max
 
 @triton.jit
-def mask_2d(offs0, offs1, max0, max1): return (tl.expand_dims(offs0, 1) < max0) & (tl.expand_dims(offs1, 0) < max1)
+def mask_2d(offs0, offs1, max0, max1):
+    """Create a 2D mask using upper bounds for each axis."""
+    return (tl.expand_dims(offs0, 1) < max0) & (tl.expand_dims(offs1, 0) < max1)
 
 # # load
 
 @triton.jit
 def load_1d(ptr, sz: const, n, max, stride=1):
-    '''Chunk 1d vector (defined by ptr) into 1d grid, where each chunk has size sz. Load the nth chunk. Ie, load [n*sz,...,(n+1)*sz-1].'''
+    """
+    Chunk 1d vector (defined by ptr) into 1d grid, where each chunk has size sz.
+    Load the nth chunk. Ie, load [n*sz,...,(n+1)*sz-1].
+    """
     offs = offset_1d(sz, n)
     mask = mask_1d(offs, max)
     return tl.load(ptr + offs, mask)
@@ -38,7 +54,10 @@ def load_full_1d(ptr, sz: const, stride=1):
 
 @triton.jit
 def load_2d(ptr, sz0: const, sz1: const, n0, n1, max0, max1, stride0=None, stride1=1):
-    '''Chunk 2d matrix (defined by ptr) into 2d grid, where each chunk has size (sz0,sz1). Load the (n0,n1)th chunk. Ie, load [n0*sz0,...,(n0+1)*sz0-1] x [n1*sz1,...,(n1+1)*sz1-1].'''
+    """
+    Chunk 2d matrix (defined by ptr) into 2d grid, where each chunk has size (sz0,sz1).
+    Load the (n0,n1)th chunk. Ie, load [n0*sz0,...,(n0+1)*sz0-1] x [n1*sz1,...,(n1+1)*sz1-1].
+    """
     stride0 = stride0 or sz1
     offs0 = offset_1d(sz0, n0)
     offs1 = offset_1d(sz1, n1)
@@ -72,7 +91,10 @@ def store_full_1d(vals, ptr, sz: const, stride=1):
 
 @triton.jit
 def store_2d(vals, ptr, sz0: const, sz1: const, n0, n1, max0, max1, stride0=None, stride1=1):
-    '''Store 2d block into (n0,n1)th chunk of matrix (defined by ptr), where each chunk has size (sz0, sz1)'''
+    """
+    Store 2d block into (n0,n1)th chunk of matrix (defined by ptr), where each chunk has size
+    (sz0, sz1)
+    """
     stride0 = stride0 or sz1
     offs0 = offset_1d(sz0, n0)
     offs1 = offset_1d(sz1, n1)

triton_util/triton_util/debugging.py

@@ -1,3 +1,9 @@
+"""Debugging utilities for Triton kernels.
+
+Includes conditional breakpoints and printing based on thread identifiers,
+plus tensor readiness checks for CUDA or interpreted environments.
+"""
+# pylint: disable=multiple-statements,line-too-long,import-outside-toplevel,eval-used,fixme,unused-variable
 import os
 import triton
 import triton.language as tl
@@ -42,12 +48,17 @@ def print_if(*txt, conds):
     if test_pid_conds(conds): print(*txt)
 
 @triton.jit
-def breakpoint_once(): breakpoint_if('=0,=0,=0')
+def breakpoint_once():
+    """Trigger a breakpoint."""
+    breakpoint_if('=0,=0,=0')
 
 @triton.jit
-def print_once(*txt): print_if(*txt,conds='=0,=0,=0')
+def print_once(*txt):
+    """Print a message."""
+    print_if(*txt,conds='=0,=0,=0')
 
 def assert_tensors_gpu_ready(*tensors):
+    """Assert that each tensor is contiguous and on the GPU (unless TRITON_INTERPRET=1)."""
     for t in tensors:
         assert t.is_contiguous(), "A tensor is not contiguous"
         if not os.environ.get('TRITON_INTERPRET') == '1': assert t.is_cuda, "A tensor is not on cuda"