Optimize imports

.github/workflows/tests.yml

@@ -17,29 +17,29 @@ jobs:
       matrix:
         python-version: [ "3.8" ]
     steps:
-    -   uses: actions/checkout@v3
-    -   name: Set up Python ${{ matrix.python-version }}
+      - uses: actions/checkout@v3
+      - name: Set up Python ${{ matrix.python-version }}
         uses: actions/setup-python@v3
         with:
           python-version: ${{ matrix.python-version }}
-    -   name: Install dependencies
+      - name: Install dependencies
         run: |
           python -m pip install --upgrade pip
           python -m pip install flake8
           python -m pip install isort
           python -m pip install black
           python -m pip install black[jupyter]
           if [ -f requirements.txt ]; then pip install -r requirements.txt; fi
-    -   name: Check Python syntax errors
+      - name: Check Python syntax errors
         run: |
           flake8 . --count --select=E9,F63,F7,F82 --show-source --statistics
-    -   name: Validate the coding style
+      - name: Validate the coding style
         run: |
           python -m isort . --check
           python -m black . --check
-    -   name: Install Trident
+      - name: Install Trident
         run: |
           bash install_package.sh
-    -   name: Test with pytest
+      - name: Test with pytest
         run: |
           python -m pytest -n 2 .

.pre-commit-config.yaml

@@ -2,15 +2,15 @@ repos:
   - repo: https://github.com/pre-commit/pre-commit-hooks
     rev: v2.3.0
     hooks:
-    -   id: check-yaml
-    -   id: end-of-file-fixer
-    -   id: trailing-whitespace
+      - id: check-yaml
+      - id: end-of-file-fixer
+      - id: trailing-whitespace
   - repo: https://github.com/pycqa/isort
     rev: 5.12.0
     hooks:
-    -   id: isort
+      - id: isort
   - repo: https://github.com/psf/black
     rev: 22.10.0
     hooks:
-    -   id: black
-    -   id: black-jupyter
+      - id: black
+      - id: black-jupyter

benchmarks/benchmark_shift_gelu.py

@@ -14,7 +14,6 @@
 
 import torch
 import triton
-import triton.language as tl
 import util
 
 import trident

examples/playground.py

@@ -11,10 +11,3 @@
 # 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.
-
-import torch
-import torch.nn as nn
-import triton
-import triton.language as tl
-
-import trident

tests/test_geglu.py

@@ -24,7 +24,6 @@ def geglu(input, weight, bias: torch.Tensor = None):
     return state * torch.nn.functional.gelu(gate)
 
 
-# @pytest.mark.skip
 @pytest.mark.parametrize("num_batches, m_size, n_size, k_size", [(2, 4, 4, 4)])
 def test_forward(num_batches, m_size, n_size, k_size, device):
     input = torch.randn(num_batches, m_size, k_size, device=device)

trident/kernel/argmax.py

@@ -15,8 +15,6 @@
 import triton
 import triton.language as tl
 
-from trident import language
-
 
 class Argmax:
     @staticmethod

trident/kernel/group_norm.py

@@ -15,8 +15,6 @@
 import triton
 import triton.language as tl
 
-from trident import language
-
 
 class GroupNorm:
     @staticmethod

trident/kernel/layer_norm.py

@@ -15,8 +15,6 @@
 import triton
 import triton.language as tl
 
-from trident import kernel, language
-
 
 class LayerNorm:
     @staticmethod

trident/kernel/max.py

@@ -15,8 +15,6 @@
 import triton
 import triton.language as tl
 
-from trident import language
-
 
 class Max:
     @staticmethod

trident/kernel/rms_norm.py

@@ -15,8 +15,6 @@
 import triton
 import triton.language as tl
 
-from trident import language
-
 
 class RMSNorm:
     @staticmethod

trident/kernel/silu.py

@@ -15,8 +15,6 @@
 import triton
 import triton.language as tl
 
-from trident import language
-
 
 def silu_configs():
     configs = []

trident/language/combine.py

@@ -15,8 +15,6 @@
 import triton
 import triton.language as tl
 
-from trident import language
-
 
 @triton.jit
 def combine_welford(m2_a, mean_a, count_a, m2_b, mean_b, count_b):

trident/language/constant.py

@@ -12,7 +12,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import triton
 import triton.language as tl
 
 dim = [tl.constexpr(i) for i in range(3)]

trident/language/function.py

@@ -15,8 +15,6 @@
 import triton
 import triton.language as tl
 
-from trident import language
-
 
 @triton.jit
 def batch(index, num_channels, num_rows, num_cols):

trident/language/var.py

@@ -15,8 +15,6 @@
 import triton
 import triton.language as tl
 
-from trident import language
-
 
 class Var:
     @staticmethod

trident/operation/cosine_similarity.py

@@ -15,7 +15,6 @@
 from typing import Any
 
 import torch
-import triton
 
 from trident import kernel, util
 

trident/operation/layer_norm.py

@@ -18,7 +18,7 @@
 import torch
 import triton
 
-from trident import function, kernel, util
+from trident import kernel, util
 
 
 class LayerNorm(torch.autograd.Function):

trident/operation/max_pool2d.py

@@ -17,7 +17,7 @@
 import torch
 import triton
 
-from trident import kernel, math
+from trident import kernel
 
 
 class MaxPool2d(torch.autograd.Function):
@@ -39,11 +39,7 @@ def __forward(inp, knl_sz):
         assert out.is_contiguous()
 
         grid = lambda meta: (inp_bt * inp_ch * out_h * triton.cdiv(out_w, meta["grp_sz"]),)
-        grp_sz = math.clamp(
-            128 // triton.next_power_of_2(knl_sz),
-            1,
-            triton.next_power_of_2(out_w),
-        )
+        grp_sz = max(1, min(triton.next_power_of_2(out_w), 128 // triton.next_power_of_2(knl_sz)))
 
         kernel.MaxPool2d.forward[grid](
             inp,

trident/operation/mean.py

@@ -17,7 +17,7 @@
 import torch
 import triton
 
-from trident import kernel, math, util
+from trident import kernel, util
 
 
 class Mean(torch.autograd.Function):

trident/operation/prelu.py

@@ -17,7 +17,7 @@
 import torch
 import triton
 
-from trident import kernel, util
+from trident import kernel
 
 
 class PReLU(torch.autograd.Function):

trident/operation/sum.py

@@ -15,9 +15,8 @@
 from typing import Any
 
 import torch
-import triton
 
-from trident import kernel, math, util
+from trident import kernel, util
 
 
 class Sum(torch.autograd.Function):

trident/operation/var.py

@@ -17,7 +17,7 @@
 import torch
 import triton
 
-from trident import kernel, math, util
+from trident import kernel, util
 
 
 class Var(torch.autograd.Function):

trident/operation/var_mean.py

@@ -17,7 +17,7 @@
 import torch
 import triton
 
-from trident import kernel, math, util
+from trident import kernel, util
 
 
 class VarMean(torch.autograd.Function):

trident/util/util.py

@@ -13,10 +13,9 @@
 # limitations under the License.
 
 import torch
-import triton
 import triton.language as tl
 
-from trident import math, module, operation
+from trident import module
 
 
 def fill(inp, val):