Adding python wheels for the package and for the kernels

.github/workflows/build_flashfftconv.yaml

@@ -0,0 +1,25 @@
+name: Build Wheels
+
+on:
+  push:
+    # push on github tags that start with "wheel"
+    tags:
+      - "wheel*"
+
+env:
+  MAX_JOBS: 2
+
+jobs:
+  monarch_cuda_wheels_cuda:
+    uses:
+      ./.github/workflows/workflow_build_cuda_wheels.yaml
+    with:
+      setuppypath: "csrc/flashfftconv"
+    secrets: inherit
+
+  flashfftconv_nocuda_wheels:
+    uses:
+      ./.github/workflows/workflow_build_cuda_wheels.yaml
+    with:
+      setuppypath: ""
+    secrets: inherit

.github/workflows/workflow_build_cuda_wheels.yaml

@@ -0,0 +1,120 @@
+name: install_wheels_cuda
+
+on:
+  workflow_call:
+    inputs:
+      setuppypath:
+        required: true
+        type: string
+        description: "path to the dir of the setup.py that should be built."
+
+env:
+  SETUPPYPATH: ${{ inputs.setuppypath }}
+  MAX_JOBS: "2"
+
+jobs:
+  build_wheels:
+    name: Build Cuda Wheels for torch
+    runs-on: ${{ matrix.os }}
+
+    strategy:
+      matrix:
+        os: [ubuntu-20.04] #, windows-latest]
+        pyver: ["3.8", "3.9", "3.10", "3.11"]
+        cuda: ["11.8.0", "12.1.1"]
+    defaults:
+      run:
+        shell: pwsh
+    env:
+      CUDA_VERSION: ${{ matrix.cuda }}
+      # define where the setup.py is located
+
+    steps:
+      - name: Free Disk Space
+        uses: jlumbroso/free-disk-space@v1.3.0
+        if: runner.os == 'Linux'
+        with:
+          tool-cache: false
+          android: true
+          dotnet: true
+          haskell: true
+          large-packages: false
+          docker-images: true
+          swap-storage: false
+
+      - uses: actions/checkout@v3
+
+      - uses: actions/setup-python@v3
+        with:
+          python-version: ${{ matrix.pyver }}
+
+      - name: Setup Mamba
+        uses: conda-incubator/setup-miniconda@v2.2.0
+        with:
+          activate-environment: "build"
+          python-version: ${{ matrix.pyver }}
+          miniforge-variant: Mambaforge
+          miniforge-version: latest
+          use-mamba: true
+          add-pip-as-python-dependency: true
+          auto-activate-base: false
+
+      - name: Install Dependencies
+        run: |
+          # Install CUDA toolkit
+          mamba install -y 'cuda' -c "nvidia/label/cuda-${env:CUDA_VERSION}"
+
+          # Env variables
+          $env:CUDA_PATH = $env:CONDA_PREFIX
+          $env:CUDA_HOME = $env:CONDA_PREFIX
+
+          # Install torch
+          $cudaVersion = $env:CUDA_VERSION.Replace('.', '')
+          $cudaVersionPytorch = $cudaVersion.Substring(0, $cudaVersion.Length - 1)
+          if ([int]$cudaVersionPytorch -gt 118) { $pytorchVersion = "torch==2.1.0" } else {$pytorchVersion = "torch==2.0.1"}
+          python3 -m pip install --upgrade --no-cache-dir $pytorchVersion+cu$cudaVersionPytorch --index-url https://download.pytorch.org/whl/cu$cudaVersionPytorch
+          python3 -m pip install build setuptools wheel ninja numpy
+
+          # Print version information
+          python3 --version
+          python3 -c "import torch; print('PyTorch:', torch.__version__)"
+          python3 -c "import torch; print('CUDA:', torch.version.cuda)"
+          python3 -c "import os; print('CUDA_HOME:', os.getenv('CUDA_HOME', None))"
+          python3 -c "from torch.utils import cpp_extension; print (cpp_extension.CUDA_HOME)"
+
+      - name: Build Wheel
+        run: |
+          $env:CUDA_PATH = $env:CONDA_PREFIX
+          $env:CUDA_HOME = $env:CONDA_PREFIX
+
+          # build the wheel
+          cd ${{ github.workspace }}/$env:SETUPPYPATH
+          python3 setup.py sdist bdist_wheel 
+
+          # rename wheel to reflect cuda version and torch version bash: dist/*.whl
+          $WHEEL_PATH = python3 -c "if True: \
+          import os, shutil, glob, re, torch;\
+          torch_name  = '+torch' + (torch.__version__.split('+')[0] + 'cu' + torch.version.cuda).replace('.', '');\
+          file = glob.glob(os.path.join('dist', '*.whl'))[0];\
+          current_dir = os.getcwd();\
+          insert_pos = re.search('-cp[0-9]+-cp[0-9]+-', file);\
+          torch_name = '' if insert_pos is None else torch_name;\
+          insert_pos = re.search('-py3', file).start() if insert_pos is None else insert_pos.start();\
+          new_name = f'{file[:insert_pos]}{torch_name}{file[insert_pos:]}';\
+          print(os.path.join(os.getcwd(), new_name));\
+          shutil.move(file, new_name);"         
+
+          # append to powershell env
+          echo "WHEEL_PATH=$WHEEL_PATH" | Out-File -FilePath $env:GITHUB_ENV -Encoding utf8 -Append
+
+      - name: Release with Notes
+        uses: softprops/action-gh-release@v1
+        if: startsWith(github.ref, 'refs/tags/')
+        with:
+          files: |
+            ${{ env.WHEEL_PATH }}
+        env:
+          GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
+          # from previous step
+          WHEEL_PATH: ${{ env.WHEEL_PATH }}
+

csrc/flashfftconv/setup.py

@@ -24,9 +24,9 @@ def append_nvcc_threads(nvcc_extra_args):
         return nvcc_extra_args + ["--threads", "4"]
     return nvcc_extra_args
 
-arch = get_last_arch_torch()
+# arch = get_last_arch_torch()
 # [MP] make install more flexible here
-sm_num = arch[-2:]
+# sm_num = arch[-2:]
 cc_flag = ['--generate-code=arch=compute_80,code=compute_80']
 
 
@@ -61,8 +61,11 @@ def append_nvcc_threads(nvcc_extra_args):
             'conv1d/conv1d_bwd_cuda_blh.cu',
         ],
         extra_compile_args={'cxx': ['-O3'],
-                             'nvcc': append_nvcc_threads(['-O3', '-lineinfo', '--use_fast_math', '-std=c++17'] + cc_flag)
-                            })
+                            'nvcc': ['-O3', '--threads', '4', '-lineinfo', '--use_fast_math', '-std=c++17', '--generate-code=arch=compute_80,code=compute_80']}
+        # extra_compile_args={'cxx': ['-O3'],
+        #                      'nvcc': append_nvcc_threads(['-O3', '-lineinfo', '--use_fast_math', '-std=c++17'] + cc_flag)
+        #                     }
+        )
     ],
     cmdclass={
         'build_ext': BuildExtension