irlyngaas / ckext

Building ckExt

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

Building Dependencies

composable_kernel

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

composable_kernel extensions

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

onnxruntime

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

Running Test scripts

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