Merge branch 'hidet-org:main' into main

.github/workflows/launch.yaml

@@ -0,0 +1,37 @@
+name: Launch CI
+
+
+on:
+  workflow_dispatch:
+    inputs:
+      shutdown_instances:
+        description: 'Shut down GPU instances when finished.'
+        required: true
+        type: boolean
+        default: true
+  issue_comment:
+    types: [created]
+
+
+jobs:
+  trigger:
+    if: |
+      github.event_name == 'workflow_dispatch' || 
+      github.event_name == 'issue_comment' && github.event.issue.pull_request != '' &&
+      contains(fromJSON('["MEMBER", "OWNER", "COLLABORATOR"]'), github.event.comment.author_association) &&
+      contains(github.event.comment.body, '$hidet-ci launch')
+    runs-on: ubuntu-latest
+    steps:
+      - name: Trigger workflow in internal repository
+        run: |
+         curl -L \
+         -X POST \
+         -H "Accept: application/vnd.github+json" \
+         -H "Authorization: Bearer ${{ secrets.GH_PAT }}" \
+         -H "X-GitHub-Api-Version: 2022-11-28" \
+         https://api.github.com/repos/${{ secrets.REPO_NAME }}/actions/workflows/regression.yaml/dispatches \
+         -d "{\"ref\": \"main\", \"inputs\": {\"shutdown_instances\": \"${{ env.SHUTDOWN }}\", \"source_repo\": \"${{ env.SOURCE_REPO }}\", \"source_ref\": \"${{ env.SOURCE_REF }}\"}}"
+        env:
+          SHUTDOWN: ${{ github.event_name == 'workflow_dispatch' && inputs.shutdown_instances || !contains(github.event.comment.body, '--keep') }}
+          SOURCE_REPO: ${{ github.repository }}
+          SOURCE_REF: ${{ github.event_name == 'workflow_dispatch' && github.ref_name || github.event.issue.number }}

.github/workflows/regression.yaml

@@ -8,16 +8,19 @@ on:
         required: true
         type: boolean
         default: true
-  issue_comment:
-    types: [created]
+      source_repo:
+        description: 'Source Repository Name. E.g, hidet-org/hidet'
+        required: true
+        type: string
+        default: 'this'
+      source_ref:
+        description: 'Source repository ref (Branch name or PR number).'
+        required: true
+        type: string
+        default: 'this'
 
 jobs:
   start_instances:
-    if: |
-      github.event_name == 'workflow_dispatch' || 
-      github.event_name == 'issue_comment' && github.event.issue.pull_request != '' &&
-      contains(fromJSON('["MEMBER", "OWNER", "COLLABORATOR"]'), github.event.comment.author_association) &&
-      contains(github.event.comment.body, '$hidet-ci launch')
     runs-on: ubuntu-latest
     outputs:
       started_instances: ${{ steps.run_py_script.outputs.started_instances }}
@@ -61,13 +64,20 @@ jobs:
     container:
       image: nvcr.io/nvidia/pytorch:23.10-py3
       options: --gpus all
+    outputs:
+      commit_time: ${{ steps.get_commit_info.outputs.commit_time }}
+      commit_author: ${{ steps.get_commit_info.outputs.commit_author }}
+      commit_sha: ${{ steps.get_commit_info.outputs.commit_sha }}
     steps:
       - name: Checkout repo
         uses: actions/checkout@v4
         with:
+          repository: |
+            ${{ inputs.source_repo == 'this' && github.repository ||
+            inputs.source_repo }}
           ref: |
-            ${{ github.event_name == 'workflow_dispatch' && github.ref_name ||
-            format('refs/pull/{0}/head', github.event.issue.number) }}
+            ${{ inputs.source_repo == 'this' && github.ref_name || 
+            format('refs/pull/{0}/head', inputs.source_ref) }}
           path: hidet
 
       - name: Checkout models
@@ -114,10 +124,8 @@ jobs:
           python hidet/.github/scripts/run_tests.py
         env:
           HW_CONFIG: ${{ matrix.hw_configs }}
-          REPO_NAME: ${{ github.repository }}
-          REPO_BRANCH: |
-            ${{ github.event_name == 'workflow_dispatch' && github.ref_name ||
-            format('pull/{0}', github.event.issue.number) }}
+          REPO_NAME: ${{ inputs.source_repo == 'this' && github.repository || inputs.source_repo }}
+          REPO_BRANCH: ${{ inputs.source_repo == 'this' && github.ref_name || format('pull/{0}', inputs.source_ref) }}
           CI_CS_HOSTNAME: ${{ secrets.CI_CS_HOSTNAME }}
           CI_CS_PORT: ${{ secrets.CI_CS_PORT }}
           CI_CS_USERNAME: ${{ secrets.CI_CS_USERNAME }}
@@ -130,6 +138,17 @@ jobs:
           name: run_configs_${{ matrix.hw_configs }}
           path: run_configs.json
           retention-days: 1
+
+      - name: Retrieve commit properties
+        id: get_commit_info
+        run: |
+          cd hidet
+          COMMIT_TIME=$(git log -1 --format=%cd --date=format:'%Y-%m-%d %H:%M:%S')
+          COMMIT_AUTHOR=$(git log -1 --format=%an)
+          COMMIT_SHA=$(git log -1 --format=%H)
+          echo "commit_time=$COMMIT_TIME" >> $GITHUB_OUTPUT
+          echo "commit_author=$COMMIT_AUTHOR" >> $GITHUB_OUTPUT
+          echo "commit_sha=$COMMIT_SHA" >> $GITHUB_OUTPUT
   
   upload_results:
     runs-on: ubuntu-latest
@@ -143,37 +162,23 @@ jobs:
 
       - name: Download run configs
         uses: actions/download-artifact@v3
-
-      - name: Setup ENV
-        run: |
-          COMMIT_TIME=$(git log -1 --format=%cd --date=format:'%Y-%m-%d %H:%M:%S')
-          COMMIT_AUTHOR=$(git log -1 --format=%an)
-          COMMIT_SHA=$(git log -1 --format=%H)
-          echo "COMMIT_TIME=$COMMIT_TIME" >> $GITHUB_ENV
-          echo "COMMIT_AUTHOR=$COMMIT_AUTHOR" >> $GITHUB_ENV
-          echo "COMMIT_SHA=$COMMIT_SHA" >> $GITHUB_ENV
 
       - name: Run main Python script
         run: python ./.github/scripts/upload_results.py
         env:
-          REPO_NAME: ${{ github.repository }}
-          REPO_BRANCH: |
-            ${{ github.event_name == 'workflow_dispatch' && github.ref_name ||
-            format('pull/{0}', github.event.issue.number) }}
-          COMMIT_SHA: ${{ env.COMMIT_SHA }}
-          COMMIT_TIME: ${{ env.COMMIT_TIME }}
-          COMMIT_AUTHOR: ${{ env.COMMIT_AUTHOR }}
+          REPO_NAME: ${{ inputs.source_repo == 'this' && github.repository || inputs.source_repo }}
+          REPO_BRANCH: ${{ inputs.source_repo == 'this' && github.ref_name || format('pull/{0}', inputs.source_ref) }}
+          COMMIT_SHA: ${{ needs.run_tests.outputs.commit_sha }}
+          COMMIT_TIME: ${{ needs.run_tests.outputs.commit_time }}
+          COMMIT_AUTHOR: ${{ needs.run_tests.outputs.commit_author }}
           HW_CONFIGS: ${{ needs.start_instances.outputs.hw_configs }}
           CI_DB_HOSTNAME: ${{ secrets.CI_DB_HOSTNAME }}
           CI_DB_PORT: ${{ secrets.CI_DB_PORT }}
           CI_DB_USERNAME: ${{ secrets.CI_DB_USERNAME }}
           CI_DB_PASSWORD: ${{ secrets.CI_DB_PASSWORD }}
 
   stop_instances:
-    if: |
-      github.event_name == 'workflow_dispatch' && inputs.shutdown_instances || 
-      github.event_name == 'issue_comment' && github.event.issue.pull_request != '' &&
-      !contains(github.event.comment.body, '--keep')
+    if: inputs.shutdown_instances
     runs-on: ubuntu-latest
     needs: [start_instances, run_tests]
     steps:

python/hidet/drivers/build_graph.py

@@ -46,6 +46,9 @@ def get_graph_weights(graph):
 
 
 def get_graph_intermediates(graph):
+    """
+    Get the intermediate tensors of the graph: {output tensors of nodes} - {output tensors of the graph}
+    """
     intermediates: List[Tensor] = []
     for node in graph.nodes:
         for y in node.outputs:
@@ -145,7 +148,12 @@ def get_graph_meta_data(graph: FlowGraph, num_kernels, space: int) -> GraphMetaD
     graph_hash = sha256('\n'.join(lines).encode('utf-8')).hexdigest()[:16]
 
     return GraphMetaData(
-        inputs=inputs, outputs=outputs, hidet_version=hidet.__version__, num_kernels=num_kernels, graph_hash=graph_hash
+        inputs=inputs,
+        outputs=outputs,
+        hidet_version=hidet.__version__,
+        num_kernels=num_kernels,
+        graph_hash=graph_hash,
+        share_map=graph.share_map,
     )
 
 
@@ -191,11 +199,15 @@ def get_workspace_size_impl(cpu_size: Var, cuda_size: Var):
             for idx in [cpu_idx, cuda_idx]:
                 sb += memory_planner_init(idx)
             for node in graph_nodes:
-                for y in node.outputs:
+                for output_idx, y in enumerate(node.outputs):
                     if y in graph_intermediates:
-                        sb += DeclareStmt(
-                            tensor_ptr[y], init=memory_planner_allocate(device2idx[y.device.kind], tensor_size[y])
-                        )
+                        if node.share_map and y in node.share_map:
+                            # share the memory with input tensor
+                            input_idx: int = node.share_map[output_idx]
+                            init_addr = tensor_ptr[node.inputs[input_idx]]
+                        else:
+                            init_addr = memory_planner_allocate(device2idx[y.device.kind], tensor_size[y])
+                        sb += DeclareStmt(tensor_ptr[y], init=init_addr)
                 sb += AssignStmt(cpu_size, primitives.max(cpu_size, memory_planner_used(cpu_idx)))
                 sb += AssignStmt(cuda_size, primitives.max(cuda_size, memory_planner_used(cuda_idx)))
                 for x in node.inputs:
@@ -247,11 +259,15 @@ def launch_impl(inputs: List[Var], outputs: List[Var], p_kernels: Var):
                         node_params.append(d2w[x.device.kind] + tensor_ptr[x])
                     else:
                         raise RuntimeError("Unknown tensor {}".format(x))
-                for y in node.outputs:
+                for output_idx, y in enumerate(node.outputs):
                     if y in graph_intermediates:
-                        sb += DeclareStmt(
-                            tensor_ptr[y], init=memory_planner_allocate(d2i[y.device.kind], tensor_size[y])
-                        )
+                        if node.share_map and y in node.share_map:
+                            # share the memory with input tensor
+                            input_idx: int = node.share_map[output_idx]
+                            init_addr = tensor_ptr[node.inputs[input_idx]]
+                        else:
+                            init_addr = memory_planner_allocate(d2i[y.device.kind], tensor_size[y])
+                        sb += DeclareStmt(tensor_ptr[y], init=init_addr)
                         node_params.append(d2w[y.device.kind] + tensor_ptr[y])
                     elif y in graph.outputs:
                         node_params.append(outputs[graph.outputs.index(y)])

python/hidet/drivers/build_task.py

@@ -202,6 +202,7 @@ def get_signature(t: TensorNode, device: str) -> TensorSignature:
         symbols=[v.name for v in task.symbols],
         inputs=[get_signature(t, input_device) for t in task.inputs],
         outputs=[get_signature(t, output_device) for t in task.outputs],
+        share_map=task.share_map,
         target=build_target,
         num_candidates=num_candidates,
         hidet_version=hidet.__version__,