Merge branch 'hidet-org:main' into main

.github/requirements-ci.txt

@@ -1,4 +1,4 @@
 mysql-connector-python
-transformers
+transformers==4.37
 accelerate
 sentencepiece

.github/workflows/lint.yaml

@@ -7,7 +7,7 @@ on:
 
 jobs:
   format-and-lint:
-    if: github.repository == 'hidet-org/hidet'
+    if: github.repository == 'hidet-org/hidet' || github.repository == 'CentML/hidet'
     concurrency:
       group: ${{ github.workflow }}-${{ github.head_ref || github.run_id }}
       cancel-in-progress: true

.github/workflows/regression.yaml

@@ -47,12 +47,12 @@ jobs:
           AWS_SECRET_ACCESS_KEY: ${{ secrets.AWS_SECRET_ACCESS_KEY }}
           AWS_DEFAULT_REGION: us-east-1
 
-      - name: Upload run configs
-        uses: actions/upload-artifact@v3
-        with:
-          name: run_configs
-          path: run_configs.json
-          retention-days: 1
+      #- name: Upload run configs
+      #  uses: actions/upload-artifact@v3
+      #  with:
+      #    name: run_configs
+      #    path: run_configs.json
+      #    retention-days: 1
 
   run_tests:
     needs: start_instances
@@ -85,12 +85,16 @@ jobs:
           path: models
           ref: ci
 
-      - name: Download run configs
-        uses: actions/download-artifact@v3
-        with:
-          name: run_configs
-          path: ./mount
-
+      #- name: Download run configs
+      #  uses: actions/download-artifact@v3
+      #  with:
+      #    name: run_configs
+      #    path: ./mount
+
+      # Put run_configs.json in shared folder. Intup and output of tests is saved in it. 
+      - name: Copy run_config.json
+        run: cp hidet/tests/benchmarks/run_configs.json ./mount
+
       # Build the image
       - name: Build docker image from base image
         run: docker build -t hidet-ci -f hidet/.github/Dockerfile .
@@ -105,7 +109,7 @@ jobs:
           -e HW_CONFIG -e REPO_NAME -e REPO_BRANCH -e CI_CS_HOSTNAME -e CI_CS_PORT
           -e CI_CS_USERNAME -e CI_CS_PASSWORD -e HF_TOKEN
           -v ./mount:/workspace/mount
-          hidet-ci python hidet/.github/scripts/run_tests.py --configs /workspace/mount/run_configs.json'
+          hidet-ci python hidet/tests/benchmarks/run_tests.py --configs /workspace/mount/run_configs.json'
         env:
           HW_CONFIG: ${{ matrix.hw_configs }}
           REPO_NAME: ${{ inputs.source_repo == 'this' && github.repository || inputs.source_repo }}
@@ -162,7 +166,7 @@ jobs:
           CI_DB_PASSWORD: ${{ secrets.CI_DB_PASSWORD }}
 
   stop_instances:
-    if: inputs.shutdown_instances
+    if: ${{ always() && inputs.shutdown_instances }}
     runs-on: ubuntu-latest
     needs: [start_instances, run_tests]
     steps:

.github/workflows/tests.yaml

@@ -60,6 +60,7 @@ jobs:
 
       - name: Run tests
         run: |
+          rm -rf ~/.config/hidet
           python -m pytest -v --durations=20 --clear-cache ./tests
 
       # Build the docs

.gitignore

@@ -184,6 +184,9 @@ cython_debug/
 #  option (not recommended) you can uncomment the following to ignore the entire idea folder.
 .idea/
 
+# Ignore VSCode files 
+.vscode/*
+
 # pycharm line profiler result
 **/*.pclprof
 
@@ -213,3 +216,6 @@ build-release
 
 # experiments folder
 /experiments
+
+# vscode
+.vscode

apps/compile_server/resources/compile_worker.py

@@ -1,4 +1,4 @@
-from typing import Dict, Any, List, Tuple
+from typing import Dict, Any, List, Tuple, Sequence, Union
 import os
 import traceback
 import argparse
@@ -51,7 +51,7 @@ def compile_job(job_id: str):
 
             # load the workload
             workload: Dict[str, Any] = pickle.loads(job['workload'])
-            ir_module: hidet.ir.IRModule = workload['ir_module']
+            ir_module: Union[hidet.ir.IRModule, Sequence[hidet.ir.IRModule]] = workload['ir_module']
             target: str = workload['target']
             output_kind: str = workload['output_kind']
 

python/hidet/apps/__init__.py

@@ -9,3 +9,10 @@
 # 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.
+
+from .pretrained import PretrainedModel
+
+# from .pipeline import Pipeline
+from .registry import Registry, RegistryEntry
+
+# from .processing import BaseProcessor

python/hidet/apps/compile_server/compilation.py

@@ -9,6 +9,7 @@
 # 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.
+from typing import Sequence, Union
 import zipfile
 import shutil
 import tempfile
@@ -21,7 +22,9 @@
 from .core import api_url, access_token
 
 
-def remote_build(ir_module: IRModule, output_dir: str, *, target: str, output_kind: str = '.so'):
+def remote_build(
+    ir_module: Union[IRModule, Sequence[IRModule]], output_dir: str, *, target: str, output_kind: str = '.so'
+):
     # upload the IRModule
     if 'cuda' in target:
         if 'arch' not in target:

python/hidet/apps/diffusion/modeling/stable_diffusion/downsample.py

@@ -0,0 +1,19 @@
+from hidet.graph import nn
+from hidet.graph.tensor import Tensor
+
+
+class Downsample2D(nn.Module):
+    def __init__(self, channels: int, **kwargs):
+        super().__init__()
+
+        self.channels = channels
+        self.out_channels = kwargs.get("output_channels", None) or channels
+
+        self.conv = nn.Conv2d(
+            self.channels, self.out_channels, kernel_size=3, stride=2, padding=kwargs["downsample_padding"], bias=True
+        )
+
+    def forward(self, hidden_states: Tensor) -> Tensor:
+        assert hidden_states.shape[1] == self.channels
+
+        return self.conv(hidden_states)