Need to build dependencies in the following sections and set the following paths to build
- CK_PATH: Path to composable_kernel repo
- CK_EXT_PATH: Path to composable_kernel extension build dir
- CK_LIB_PATH: Path to composable_kernel install dir
- ORT_PATH: Path to onnxruntime repo
- ORT_BUILD_PATH: Path to onnxruntime build dir
I use a conda environment I already have setup with pytorch, not the same environment used to build onnxruntime below
self_ck_attn: Uses deviceArray ptr
python setup.py install --self_ck_attn
self_ck_attn_torch: Uses torchTensor ptr
python setup.py install --self_ck_attn_torch
self_ck_attn_torchTensor: Uses torchTensor
python setup.py install --self_ck_attn_torchTensor
self_ck_attn_deviceAray: Uses deviceArray Object
python setup.py install --self_ck_attn_deviceArray
Any combinations of these is fine
module load rocm/5.4.0
git clone [email protected]:irlyngaas/composable_kernel.git
cd composable_kernel
export CK_PATH=$PWD
export CMAKE_INSTALL_PREFIX=/PATH/TO/CKLIB
export CK_LIB_PATH=/PATH/TO/CKLIB
mkdir build && cd build
cmake \
-D CMAKE_PREFIX_PATH=/opt/rocm-5.4.0 \
-D CMAKE_CXX_COMPILER=/opt/rocm-5.4.0/bin/hipcc \
-D CMAKE_CXX_FLAGS="-O3" \
-D CMAKE_BUILD_TYPE=Release m \
-D GPU_TARGETS="gfx908;gfx90a" \
..
make install
May take awhile to build if on crusher feel free to use my build at
/gpfs/alpine/med106/world-shared/irl1/CKREDO/cklib/usr/local
cd composable_kernel/extension
export CMAKE_PREFIX_PATH=$CMAKE_PREFIX_PATH:/PATH/TO/CKLIB
mkdir build && cd build && export CK_EXT_PATH=$PWD && cmake .. && make -j
Set up python
module load PrgEnv-gnu rocm/5.4.0 cray-python
python3 -m venv onnx_env
source onnx_env/bin/activate
pip3 install --upgrade pip; pip3 install numpy; pip3 install cmake
Build onnx
git clone [email protected]:irlyngaas/onnxruntime.git
cd onnxruntime
git checkout crusher
bash build_ke.sh
export ORT_PATH=$PWD
export ORT_BUILD_PATH=$PWD/build/Release
Set environment variables needed to use kernelExplorer through python
export KERNEL_EXPLORER_BUILD_DIR=$ORT_BUILD_PATH
export PYTHONPATH=$PYTHONPATH:$ORTPATH/onnxruntime/python/tools/kernel_explorer/kernels
Run onnxruntime gemm_softmax_gemm
salloc -A PROJ -t TIME -p batch -N 1
python exp_gemm_softmax_gemm.py
srun -n1 python ckExt/test/run_rocm_extensions.py
Modify ckExt/test/run_rocm_extensions.py:L6 to whichever unit test you want to run, for example
test_dirs = ["self_ck_attn"] #Runs only deviceArray w/ pointer Test
test_dirs = ["self_ck_attn", "self_ck_attn_torch"] #Runs 2 unit tests
Run unit test on Crusher
salloc -A PROJ -t TIME -p batch -N 1
cd ckExt/test/self_ck_attn
srun -n1 python ckExt/test/run_rocm_extensions.py
Run simple test on Crusher (Mimics exp_gemm_softmax_gemm.py in onnxruntime repo; only exists for options using deviceArray)
srun -n1 python ckExt/test/self_ck_attn/simpletest.py