diff --git a/.github/workflows/Ubuntu-latest.yml b/.github/workflows/Ubuntu-latest.yml index 47dd4c5b..b7634019 100644 --- a/.github/workflows/Ubuntu-latest.yml +++ b/.github/workflows/Ubuntu-latest.yml @@ -1,20 +1,21 @@ # This starter workflow is for a CMake project running on a single platform. There is a different starter workflow if you need cross-platform coverage. # See: https://github.com/actions/starter-workflows/blob/main/ci/cmake-multi-platform.yml name: Test on Ubuntu Linux - on: push: - branches: [ "master" ] + branches: [ "master", "dev-new"] pull_request: - branches: [ "master" ] + branches: [ "master", "dev-new"] + workflow_dispatch: env: # Customize the CMake build type here (Release, Debug, RelWithDebInfo, etc.) BUILD_TYPE: Release COMET_SRC: ${{github.workspace}} - + COMETPY_COMET_PATH: ${{github.workspace}}/build/comet/ + COMETPY_LLVM_PATH: ${{github.workspace}}/build/llvm/src/llvm-build jobs: - build-and-test-comet: + build-comet-and-dependencies: # The CMake configure and build commands are platform agnostic and should work equally well on Windows or Mac. # You can convert this to a matrix build if you need cross-platform coverage. # See: https://docs.github.com/en/free-pro-team@latest/actions/learn-github-actions/managing-complex-workflows#using-a-build-matrix @@ -33,51 +34,46 @@ jobs: with: path: | ${{github.workspace}}/llvm - ${{github.workspace}}/blis/ - ${{github.workspace}}/install/ - ${{github.workspace}}/build/ + ${{github.workspace}}/build/llvm + ${{github.workspace}}/build/blis key: ${{ runner.os }}-submodules - - name: Update git submodules + - name: Build clean if: steps.cache-submodule.outputs.cache-hit != 'true' - uses: actions/checkout@v3 - with: - submodules: recursive + run: mkdir ${{github.workspace}}/build && cd ${{github.workspace}}/build && cmake ../ -DCMAKE_BUILD_TYPE=Release && make - - name: Build llvm - if: steps.cache-submodule.outputs.cache-hit != 'true' - run: mkdir ${{github.workspace}}/llvm/build && cd llvm/build/ && cmake -G Ninja ../llvm -DLLVM_ENABLE_PROJECTS="mlir;openmp;clang" -DLLVM_TARGETS_TO_BUILD="X86" -DLLVM_ENABLE_ASSERTIONS=ON -DCMAKE_BUILD_TYPE=Release && ninja + - name: Build from cached + if: steps.cache-submodule.outputs.cache-hit == 'true' + run: cd ${{github.workspace}}/build && cmake ../ -DCMAKE_BUILD_TYPE=Release -DLLVM_CUSTOM_BUILD_PATH=${{github.workspace}}/build/llvm/src/llvm-build -DCMAKE_BUILD_TYPE=Release && make - - name: Build blis - if: steps.cache-submodule.outputs.cache-hit != 'true' - run: cd ${{github.workspace}} && patch -s -p0 < comet-blis.patch && cd blis && ./configure --prefix=$COMET_SRC/install --disable-shared auto && make && make install - - name: Build COMET - # Build your program with the given configuration - run: rm -rf ${{github.workspace}}/build && mkdir ${{github.workspace}}/build && cd ${{github.workspace}}/build && cmake -G Ninja .. -DMLIR_DIR=${{github.workspace}}/llvm/build/lib/cmake/mlir -DLLVM_DIR=${{github.workspace}}/llvm/build/lib/cmake/llvm -DLLVM_ENABLE_ASSERTIONS=ON -DCMAKE_BUILD_TYPE=Release && ninja + test-comet-backend: - # test-comet-backend: - - # needs: build-comet - # runs-on: ubuntu-latest + needs: build-comet-and-dependencies + runs-on: ubuntu-latest - # steps: - # - uses: actions/checkout@v4 - # # - uses: lukka/get-cmake@latest - # - name: Install CMake - # run: sudo apt-get install cmake && sudo apt-get install ninja-build + steps: + - uses: actions/checkout@v4 + # - uses: lukka/get-cmake@latest + - name: Install CMake + run: sudo apt-get install cmake && sudo apt-get install ninja-build - # - name: Cache Submodules - # id: cache-submodule - # uses: actions/cache@v4 - # # if: always() - # with: - # path: | - # ${{github.workspace}}/llvm - # ${{github.workspace}}/blis/ - # ${{github.workspace}}/install/ - # ${{github.workspace}}/build/ - # key: ${{ runner.os }}-submodules + + - name: Cache Submodules + id: cache-submodule + uses: actions/cache@v4 + # if: always() + with: + path: | + ${{github.workspace}}/llvm + ${{github.workspace}}/build/llvm + ${{github.workspace}}/build/blis + key: ${{ runner.os }}-submodules + + - name: Build COMET + working-directory: ${{github.workspace}}/build/ + # Build your program with the given configuration + run: cmake .. -DLLVM_CUSTOM_BUILD_PATH=${{github.workspace}}/build/llvm/src/llvm-build -DCMAKE_BUILD_TYPE=Release && make - name: Initialize Python 3.11 uses: actions/setup-python@v4 @@ -91,7 +87,7 @@ jobs: sudo apt-get install -y python3-pip - name: Test - working-directory: ${{github.workspace}}/build + working-directory: ${{github.workspace}}/build/comet/ shell: bash run: ninja check-comet-integration @@ -129,7 +125,7 @@ jobs: # cargo test test-cometpy: - needs: build-and-test-comet + needs: build-comet-and-dependencies runs-on: ubuntu-latest steps: - uses: actions/checkout@v4 @@ -141,9 +137,8 @@ jobs: with: path: | ${{github.workspace}}/llvm - ${{github.workspace}}/blis/ - ${{github.workspace}}/install/ - ${{github.workspace}}/build/ + ${{github.workspace}}/build/llvm + ${{github.workspace}}/build/blis key: ${{ runner.os }}-submodules - name: Initialize Python 3.11 @@ -152,10 +147,23 @@ jobs: with: python-version: 3.11 + - name: Install python dependencies + run: | + sudo apt-get install -y python3-pip + + - name: Build COMET + working-directory: ${{github.workspace}}/build/ + # Build your program with the given configuration + run: cmake .. -DLLVM_CUSTOM_BUILD_PATH=${{github.workspace}}/build/llvm/src/llvm-build -DCMAKE_BUILD_TYPE=Release && make + + - name: Setup cometPy run: cd ${{github.workspace}}/frontends/numpy-scipy/ && pip3 install -e . - + + - name: Install PyTest + run: pip3 install -U pytest + - name: Test CometPy working-directory: ${{github.workspace}}/frontends/numpy-scipy/integration_tests/ - run: python3 numpy_integration.py -v + run: pytest diff --git a/.gitignore b/.gitignore index a941a229..9f3825ee 100644 --- a/.gitignore +++ b/.gitignore @@ -19,4 +19,6 @@ lit.site.cfg.py .vscode /.idea -# External software \ No newline at end of file +# Clangd +/.cache + diff --git a/.gitmodules b/.gitmodules index b6b04e46..32fd02ec 100644 --- a/.gitmodules +++ b/.gitmodules @@ -7,3 +7,18 @@ [submodule "triton"] path = triton url = https://github.com/triton-lang/triton.git +[submodule "tools/llvm-spirv"] + path = tools/llvm-spirv + url = https://github.com/llvm/llvm-project.git + branch = release/7.x +[submodule "tools/spirv-llvm-tranlate"] + path = tools/spirv-llvm-tranlate + url = https://github.com/KhronosGroup/SPIRV-LLVM-Translator.git + branch = llvm_release_70 +[submodule "tools/spirv-llvm-translate"] + path = tools/spirv-llvm-translate + url = https://github.com/KhronosGroup/SPIRV-LLVM-Translator.git + branch = llvm_release_70 +[submodule "runtimes/mcl"] + path = runtimes/mcl + url = https://github.com/pnnl/mcl.git diff --git a/CMakeLists.txt b/CMakeLists.txt index 6f39c52c..520285dd 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -1,204 +1,143 @@ -##===- CMakeLists.txt - COMET cmake root ----------------------*- cmake -*-===// -## -## Configure the COMET build. -## -##===----------------------------------------------------------------------===// - cmake_minimum_required(VERSION 3.25) +project(comet_super_build LANGUAGES CXX C) +include(ExternalProject) + +# Adopted from https://gist.github.com/scivision/bb1d47a9529e153617414e91ff5390af +find_package(Git REQUIRED) + +function(init_git_submodule dir) +# Init a Git submodule +# +# include(AddGitSubmodule.cmake) +# add_git_submodule(mysubmod_dir) + +if(NOT EXISTS ${dir}/README.md) + execute_process(COMMAND ${GIT_EXECUTABLE} submodule update --init --recursive --depth 1 -- ${dir} + WORKING_DIRECTORY ${PROJECT_SOURCE_DIR} + COMMAND_ERROR_IS_FATAL ANY) +endif() -# If we are not building as a part of LLVM, build COMET as an -# standalone project, using LLVM as an external library: -if( CMAKE_SOURCE_DIR STREQUAL CMAKE_CURRENT_SOURCE_DIR ) - +endfunction(init_git_submodule) #------------------------------------------------------------------------------- -# Project setup and globals -#------------------------------------------------------------------------------- -project(comet LANGUAGES CXX C) - -set(CMAKE_CXX_STANDARD 17) -set(CMAKE_CXX_STANDARD_REQUIRED YES) -set(CMAKE_MODULE_PATH ${CMAKE_SOURCE_DIR}/cmake/modules) -#------------------------------------------------------------------------------- -# Options and settings -#------------------------------------------------------------------------------- - - option(LLVM_INCLUDE_TOOLS "Generate build targets for the LLVM tools." ON) - option(LLVM_BUILD_TOOLS "Build the LLVM tools. If OFF, just generate build targets." ON) +message(STATUS "Running on ${CMAKE_SYSTEM_NAME} ${CMAKE_SYSTEM_PROCESSOR}") -if (MSVC) - set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} /EHs-c- /GR-") -else () - set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -fno-exceptions -fno-rtti") -endif () - -#------------------------------------------------------------------------------- -# MLIR/LLVM Configuration -#------------------------------------------------------------------------------- - - find_package(MLIR REQUIRED CONFIG) - - message(STATUS "Using MLIRConfig.cmake in: ${MLIR_DIR}") - message(STATUS "Using LLVMConfig.cmake in: ${LLVM_DIR}") - - set(LLVM_RUNTIME_OUTPUT_INTDIR ${CMAKE_BINARY_DIR}/bin) - set(LLVM_LIBRARY_OUTPUT_INTDIR ${CMAKE_BINARY_DIR}/lib) - - list(APPEND CMAKE_MODULE_PATH "${MLIR_CMAKE_DIR}") - list(APPEND CMAKE_MODULE_PATH "${LLVM_CMAKE_DIR}") - - include(TableGen) - include(AddLLVM) - include(AddMLIR) - include(HandleLLVMOptions) - - set(COMET_BUILT_STANDALONE 1) - set(BACKEND_PACKAGE_STRING "LLVM ${LLVM_PACKAGE_VERSION}") +if(CMAKE_SYSTEM_PROCESSOR MATCHES "arm|aarch64") + set(DETECTED_CPU_ARCH "AArch64") +elseif(CMAKE_SYSTEM_PROCESSOR MATCHES "x86_64|amd64") + set(DETECTED_CPU_ARCH "X86") else() - set(MLIR_MAIN_SRC_DIR ${LLVM_MAIN_SRC_DIR}/../mlir ) # --src-root - set(MLIR_INCLUDE_DIR ${MLIR_MAIN_SRC_DIR}/include ) # --includedir - set(MLIR_TABLEGEN_OUTPUT_DIR ${CMAKE_BINARY_DIR}/tools/mlir/include) - set(MLIR_TABLEGEN_EXE $) - include_directories(SYSTEM ${MLIR_INCLUDE_DIR}) - include_directories(SYSTEM ${MLIR_TABLEGEN_OUTPUT_DIR}) - - set(BACKEND_PACKAGE_STRING "${PACKAGE_STRING}") + message(FATAL_ERROR "Unsupported CPU architecture: ${CMAKE_SYSTEM_PROCESSOR}") endif() -# Define the default arguments to use with 'lit', and an option for the user to -# override. -set(LIT_ARGS_DEFAULT "-sv") -if (MSVC_IDE OR XCODE) - set(LIT_ARGS_DEFAULT "${LIT_ARGS_DEFAULT} --no-progress-bar") +init_git_submodule(${CMAKE_SOURCE_DIR}/blis) +ExternalProject_Add( + blis + PREFIX ${CMAKE_BINARY_DIR}/blis + BINARY_DIR ${CMAKE_BINARY_DIR}/blis + SOURCE_DIR ${CMAKE_SOURCE_DIR}/blis + CONFIGURE_COMMAND cd ${CMAKE_SOURCE_DIR}/blis && ./configure --prefix=${CMAKE_BINARY_DIR}/blis --disable-shared auto + PATCH_COMMAND cd ${CMAKE_SOURCE_DIR} && patch -s -p0 < ${CMAKE_SOURCE_DIR}/comet-blis.patch + BUILD_COMMAND cd ${CMAKE_SOURCE_DIR}/blis && make -j && make install + INSTALL_COMMAND "" +) + +set(LLVM_CUSTOM_BUILD_PATH "" CACHE PATH "Path to existing LLVM build directory") +set(LLVM_PROJECTS "mlir,openmp,clang" CACHE STRING "LLVM projects to build") +set(LLVM_TARGETS "${DETECTED_CPU_ARCH}" CACHE STRING "LLVM targets to build") +option(ENABLE_AMD_GPU_BACKEND OFF) +option(ENABLE_NVIDIA_GPU_BACKEND OFF) +option(ENABLE_FPGA_TARGET OFF) +if(ENABLE_AMD_GPU_BACKEND) + set(LLVM_TARGETS "${LLVM_TARGETS},AMDGPU") endif() -set(LLVM_LIT_ARGS "${LIT_ARGS_DEFAULT}" CACHE STRING "Default options for lit") - - -#------------------------------------------------------------------------------- -# BLIS Configuration -#------------------------------------------------------------------------------- - set(BLA_STATIC ON) #linking with BLIS static library - set(BLA_VENDOR FLAME) - set(BLA_PREFER_PKGCONFIG TRUE) - set(BLA_PKGCONFIG_BLAS blis) - set(CMAKE_PREFIX_PATH "${CMAKE_CURRENT_SOURCE_DIR}/install") - find_package(BLAS REQUIRED) - -#------------------------------------------------------------------------------- -# COMET configuration -#------------------------------------------------------------------------------- - -# COMET project. -set(COMET_MAIN_SRC_DIR ${CMAKE_CURRENT_SOURCE_DIR} ) # --src-root -set(COMET_MAIN_INCLUDE_DIR ${COMET_MAIN_SRC_DIR}/include) - -set(COMET_SOURCE_DIR ${CMAKE_CURRENT_SOURCE_DIR}) -set(COMET_BINARY_DIR ${CMAKE_CURRENT_BINARY_DIR}) -set(COMET_INCLUDE_DIR ${CMAKE_CURRENT_BINARY_DIR}/include ) -set(COMET_TOOLS_DIR ${CMAKE_BINARY_DIR}/bin) - -list(APPEND CMAKE_MODULE_PATH "${MLIR_MAIN_SRC_DIR}/cmake/modules") -list(APPEND CMAKE_MODULE_PATH "${CMAKE_CURRENT_SOURCE_DIR}/cmake/modules") - -# Installing the headers and docs needs to depend on generating any public -# tablegen'd targets. -add_custom_target(comet-headers) -set_target_properties(comet-headers PROPERTIES FOLDER "Misc") -add_custom_target(comet-doc) - -# Add MLIR, LLVM and BLIS headers to the include path -include_directories(${LLVM_INCLUDE_DIRS}) -include_directories(${MLIR_INCLUDE_DIRS}) -include_directories("${CMAKE_CURRENT_SOURCE_DIR}/install/include/blis/") - -# Add COMET files to the include path -include_directories(${COMET_MAIN_INCLUDE_DIR}) -include_directories(${COMET_INCLUDE_DIR}) - -#------------------------------------------------------------------------------- -# Directory setup -#------------------------------------------------------------------------------- - -option(ENABLE_GPU_TARGET OFF) - -if(${ENABLE_GPU_TARGET}) - set(TRITON_PATH "" CACHE PATH "Path to Triton") - set(TRITON_BUILD_PATH "${CMAKE_BINARY_DIR}/triton" CACHE INTERNAL "Path to Triton Build") - if(NOT DEFINED CUDA_COMPUTE_CAPABILITY) - message(FATAL_ERROR "Please specify cuda compute capability requested") - endif() - add_compile_definitions(CUDA_COMPUTE_CAPABILITY=${CUDA_COMPUTE_CAPABILITY}) - - add_subdirectory(${TRITON_PATH} ${TRITON_BUILD_PATH}) - get_property(triton_libs GLOBAL PROPERTY TRITON_LIBS) - include_directories("${TRITON_PATH}") - include_directories("${TRITON_PATH}/include/") - include_directories("${TRITON_BUILD_PATH}/include/") - add_definitions(-DENABLE_GPU_TARGET) +if(ENABLE_NVIDIA_GPU_BACKEND) + set(LLVM_TARGETS "${LLVM_TARGETS},NVPTX") endif() - -add_subdirectory(include/comet) -add_subdirectory(lib) -add_subdirectory(frontends/comet_dsl) -add_subdirectory(integration_test) - - -option(COMET_INCLUDE_DOCS "Generate build targets for the COMET docs.") -if (COMET_INCLUDE_DOCS) - add_subdirectory(docs) +set(DEPENDENCIES llvm blis) +if(NOT LLVM_CUSTOM_BUILD_PATH) + init_git_submodule(${CMAKE_SOURCE_DIR}/llvm) + message(STATUS "Building LLVM from source with projects: ${LLVM_PROJECTS} and targets: ${LLVM_TARGETS}") + ExternalProject_Add( + llvm + LIST_SEPARATOR "," + PREFIX ${CMAKE_BINARY_DIR}/llvm + SOURCE_DIR ${CMAKE_CURRENT_SOURCE_DIR}/llvm/ + CONFIGURE_COMMAND cmake -G Ninja ${CMAKE_CURRENT_SOURCE_DIR}/llvm/llvm -DLLVM_ENABLE_PROJECTS=${LLVM_PROJECTS} -DLLVM_TARGETS_TO_BUILD=${LLVM_TARGETS} -DLLVM_ENABLE_ASSERTIONS=ON -DCMAKE_BUILD_TYPE=Release + BUILD_COMMAND cmake --build . -- -j 8 + INSTALL_COMMAND "" + ) + set(LLVM_BUILD_PATH ${CMAKE_BINARY_DIR}/llvm/src/llvm-build) +else() + message(STATUS "Using existing LLVM build at: ${LLVM_CUSTOM_BUILD_PATH}") + ExternalProject_Add( + llvm + BINARY_DIR ${LLVM_CUSTOM_BUILD_PATH} + SOURCE_DIR ${LLVM_CUSTOM_BUILD_PATH} + CONFIGURE_COMMAND "" + BUILD_COMMAND "" + INSTALL_COMMAND "" + ) + + set(LLVM_BUILD_PATH ${LLVM_CUSTOM_BUILD_PATH}) endif() - - -install(DIRECTORY include/comet - DESTINATION include - COMPONENT comet-headers - FILES_MATCHING - PATTERN "*.def" - PATTERN "*.h" - PATTERN "*.inc" - PATTERN "*.td" - PATTERN "*.sv" - PATTERN "LICENSE.TXT" - ) - -install(DIRECTORY ${COMET_INCLUDE_DIR}/comet - DESTINATION include - COMPONENT comet-headers - FILES_MATCHING - PATTERN "*.def" - PATTERN "*.h" - PATTERN "*.gen" - PATTERN "*.inc" - PATTERN "*.td" - PATTERN "CMakeFiles" EXCLUDE - PATTERN "config.h" EXCLUDE - ) - -if (NOT LLVM_ENABLE_IDE) - add_llvm_install_targets(install-comet-headers - DEPENDS comet-headers - COMPONENT comet-headers) +if(ENABLE_AMD_GPU_BACKEND OR ENABLE_NVIDIA_GPU_BACKEND) + init_git_submodule(${CMAKE_SOURCE_DIR}/triton) + set(DEVICE_COMPUTE_CAPABILITY "" CACHE STRING "Device compute capability for Triton codegen") + ExternalProject_Add( + triton + SOURCE_DIR ${CMAKE_SOURCE_DIR}/triton + BINARY_DIR ${CMAKE_BINARY_DIR}/triton + CONFIGURE_COMMAND "" + BUILD_COMMAND "" + INSTALL_COMMAND "" + PATCH_COMMAND git apply ${CMAKE_SOURCE_DIR}/triton.patch + ) + set(TRITON_PATH ${CMAKE_SOURCE_DIR}/triton) + set(TRITON_BUILD_PATH ${CMAKE_BINARY_DIR}/triton) + list(APPEND DEPENDENCIES triton) endif() -set(CMAKE_INSTALL_RPATH "${CMAKE_CURRENT_SOURCE_DIR}/install/lib/") -set(CMAKE_INSTALL_RPATH_USE_LINK_PATH TRUE) - -# Set RPATH to $ORIGIN on all targets. -function(set_rpath_all_targets dir) - get_property(subdirectories DIRECTORY ${dir} PROPERTY SUBDIRECTORIES) - foreach(subdir ${subdirectories}) - set_rpath_all_targets(${subdir}) - endforeach() +if(ENABLE_FPGA_TARGET) + init_git_submodule(${CMAKE_SOURCE_DIR}/tools/llvm-spirv) + init_git_submodule(${CMAKE_SOURCE_DIR}/tools/spirv-llvm-translate) + file(CREATE_LINK ${CMAKE_SOURCE_DIR}/tools/spirv-llvm-translate ${CMAKE_SOURCE_DIR}/tools/llvm-spirv/llvm/projects/spirv-llvm-translate SYMBOLIC) + # ADD_CUSTOM_TARGET(link_spirv_translate ALL + # COMMAND ${CMAKE_COMMAND} -E create_symlink ${target} ${link}) + ExternalProject_Add( + llvm-spirv + SOURCE_DIR ${CMAKE_SOURCE_DIR}/tools/llvm-spirv + BINARY_DIR ${CMAKE_BINARY_DIR}/tools/llvm-spirv + CONFIGURE_COMMAND cmake -G Ninja ${CMAKE_SOURCE_DIR}/tools/llvm-spirv/llvm/ -DLLVM_ENABLE_PROJECTS=clang -DLLVM_TARGETS_TO_BUILD=host + BUILD_COMMAND ninja llvm-spirv && ninja llvm-dis + INSTALL_COMMAND "" + ) + + init_git_submodule(${CMAKE_SOURCE_DIR}/runtimes/mcl) + ExternalProject_Add( + mcl + SOURCE_DIR ${CMAKE_SOURCE_DIR}/runtimes/mcl + BINARY_DIR ${CMAKE_BINARY_DIR}/runtimes/mcl + CONFIGURE_COMMAND cd ${CMAKE_SOURCE_DIR}/runtimes/mcl && autoreconf --install && ./configure --prefix=${CMAKE_BINARY_DIR}/runtimes/mcl + BUILD_COMMAND cd ${CMAKE_SOURCE_DIR}/runtimes/mcl && make -j && make install + INSTALL_COMMAND "" + ) + set(MCL_BUILD_PATH ${CMAKE_BINARY_DIR}/runtimes/mcl) + list(APPEND DEPENDENCIES llvm-spirv mcl) +endif() - get_directory_property(LCL_TARGETS DIRECTORY ${dir} BUILDSYSTEM_TARGETS) - set_property(TARGET ${LCL_TARGETS} PROPERTY INSTALL_RPATH "$ORIGIN/../lib") -endfunction() -option(STANDALONE_INSTALL "Create an 'install' for packaging which doesn't \ - require installation" off) -if (STANDALONE_INSTALL) - message(STATUS "Setting an $ORIGIN-based RPATH on all executables") - set_rpath_all_targets(${CMAKE_CURRENT_SOURCE_DIR}) -endif() +ExternalProject_Add( + comet + BINARY_DIR ${CMAKE_BINARY_DIR}/comet + LIST_SEPARATOR "," + PREFIX ${CMAKE_BINARY_DIR} + SOURCE_DIR ${CMAKE_CURRENT_SOURCE_DIR} + CONFIGURE_COMMAND cmake -G Ninja ${CMAKE_CURRENT_SOURCE_DIR}/cmake -DMLIR_DIR=${LLVM_BUILD_PATH}/lib/cmake/mlir -DLLVM_DIR=${LLVM_BUILD_PATH}/lib/cmake/llvm -DCMAKE_PREFIX_PATH=${CMAKE_BINARY_DIR}/blis/ -DBLAS_INCLUDE_DIRS=${CMAKE_BINARY_DIR}/blis/include/blis/ -DENABLE_AMD_GPU_BACKEND=${ENABLE_AMD_GPU_BACKEND} -DENABLE_NVIDIA_GPU_BACKEND=${ENABLE_NVIDIA_GPU_BACKEND} -DENABLE_FPGA_TARGET=${ENABLE_FPGA_TARGET} -DDEVICE_COMPUTE_CAPABILITY=${DEVICE_COMPUTE_CAPABILITY} -DTRITON_PATH=${TRITON_PATH} -DTRITON_BUILD_PATH=${TRITON_BUILD_PATH} -DMCL_BUILD_PATH=${MCL_BUILD_PATH} -DCMAKE_BUILD_TYPE=${CMAKE_BUILD_TYPE} + BUILD_COMMAND cmake --build . -- -j 8 + BUILD_ALWAYS TRUE + INSTALL_COMMAND "" + DEPENDS ${DEPENDENCIES} +) \ No newline at end of file diff --git a/README.md b/README.md index 179118b2..e77852ca 100644 --- a/README.md +++ b/README.md @@ -11,7 +11,8 @@ Comprehensive documentation of the COMET compiler can be found [here](https://pn These commands can be used to setup COMET project: -1) **Install Dependencies** To install COMET and LLVM/MLIR, the following dependencies need to be installed: +1) **Requirements.** +To install COMET and LLVM/MLIR, the following dependencies need to be already installed: * [CMake (3.25 or later)](https://cmake.org/download) * [Ninja (1.5 or later)](https://ninja-build.org/) * C++ compiler toolchain as [mentioned here](https://llvm.org/docs/GettingStarted.html#requirements) @@ -19,6 +20,8 @@ These commands can be used to setup COMET project: * [Git (1.8.4 or later)](https://www.git-scm.com/) * [pkg-config (0.29.2 or later)](https://www.freedesktop.org/wiki/Software/pkg-config/) +When targeting GPUs or/and FPGAs you will also need the drivers and runtimes of the respective vendors (Nvidia/CUDA, AMD/ROCm, Xilinx/XRT,Vitis). + 1.a **[Optional but recommended] Create a new python environment** ```bash $ export PYTHON_EXECUTABLE=$(which python3.x) # Replace 3.x with your version @@ -26,97 +29,70 @@ These commands can be used to setup COMET project: $ source comet/bin/activate ``` -2) **Get submodules required for COMET.** COMET contains LLVM and blis as a git -submodule. The LLVM repo here includes staged changes to MLIR which -may be necessary to support COMET. It also represents the version of -LLVM that has been tested. MLIR is still changing relatively rapidly, -so feel free to use the current version of LLVM, but APIs may have -changed. BLIS is an award-winning portable software framework for instantiating high-performance +2) **Build COMET.** +LLVM and blis are dependencies included in this repo as git submodules that point to the respective versions of the libraries that COMET has been tested with. LLVM/MLIR are changing relatively rapidly, so feel free to use the current version of LLVM, but APIs may have changed. + +BLIS is an award-winning portable software framework for instantiating high-performance BLAS-like dense linear algebra libraries. COMET generates a call to BLIS microkernel -after some optimizations. +after some optimizations. Also, blis is patched with changes specific to COMET, so an existing installation may not be used. +To build COMET for CPU execution only, run the following commands: ```bash -$ git clone https://github.com/pnnl/COMET.git -$ export COMET_SRC=`pwd`/COMET -$ cd $COMET_SRC -$ git submodule init -$ git submodule update --depth=1 # --depth=1 requires git>=1.8.4 +$ mkdir build +$ cd build +$ cmake ../ +$ make ``` -3) **Build and test LLVM/MLIR:** - +This will fetch and build the LLVM and blis dependencies automatically, and build COMET. Once the command completes COMET will be installed in `build/comet/`. +You may also specify a custom LLVM installation, instead of downloading a fresh copy, by passing its path to the `cmake` command: ```bash -$ export PYTHON_EXECUTABLE=$(which python3.x) # Replace 3.x with your version. Skip if already run in step 1.a -$ cd $COMET_SRC -$ mkdir llvm/build -$ cd llvm/build -# NVPTX is only required if targetting Nvidia GPUs -$ cmake -G Ninja ../llvm \ - -DLLVM_ENABLE_PROJECTS="mlir;openmp;clang" \ - -DLLVM_TARGETS_TO_BUILD="AArch64;X86;NVPTX" \ - -DCMAKE_OSX_ARCHITECTURES="arm64" \ - -DPython3_EXECUTABLE=${PYTHON_EXECUTABLE} \ - -DLLVM_ENABLE_ASSERTIONS=ON \ - -DCMAKE_BUILD_TYPE=Release -$ ninja -$ ninja check-mlir +$ cmake ../ -DLLVM_CUSTOM_BUILD_PATH=/path/to/llvm/build/ ``` -4) **Apply BLIS patch to meet COMET requirements:** - +*Note*: The LLVM installation should have enabled the `mlir, openmp, clang` projects. +Once complete, you can run the integration tests using the following commands: ```bash -$ cd $COMET_SRC -$ patch -s -p0 < comet-blis.patch +$ cd comet +$ ninja check-comet-integration ``` -5) **Build and test BLIS:** - -```bash -$ cd $COMET_SRC -$ cd blis -$ ./configure --prefix=$COMET_SRC/install --disable-shared auto -$ make [-j] -$ make check [-j] -$ make install [-j] -``` +The `-DCMAKE_BUILD_TYPE=DEBUG` flag enables debug information, which makes the +whole tree compile slower, but allows you to step through code into COMET. -6) **(If targetting GPUs) Patch Triton:** +To get something that runs fast, use `-DCMAKE_BUILD_TYPE=Release` or +`-DCMAKE_BUILD_TYPE=RelWithDebInfo` if you want to go fast and optionally if +you want debug info to go with it. `Release` mode makes a very large difference +in performance. +3) **Enabling GPU Support.** +To enable support for Nvidia, AMD GPUs you need to set the respective option in `cmake`: ```bash -$ cd $COMET_SRC -$ cd triton -$ git apply ${COMET_SRC}/triton.patch -``` +# NVIDIA GPU +$ cmake ../ -DENABLE_NVIDIA_GPU_BACKEND=ON +$ make -7) **Build and test COMET:** +# AMD GPU +$ cmake ../ -DENABLE_AMD_GPU_BACKEND=ON +$ make +``` +This will download and install [Triton](https://github.com/triton-lang/triton), a MLIR dialect for targeting GPUs used by COMET as a backend, and enable the GPU-related options in `comet-opt`. +You can also specify a default target device capability by passing the option +`-DDEVICE_COMPUTE_CAPABILITY=` in cmake. You can specify the same attribute later at the `comet-opt` command using the flag `--gpu-compute-capability` +Example options include `sm_80, sm_90` for Nvidia and `gfx908` for AMD. For example: ```bash -$ cd $COMET_SRC -$ mkdir build -$ cd build -# Omit -DENABLE_GPU_TARGET, -DCUDA_COMPUTE_CAPABILITY and -DTRITON_PATH -# if only targetting CPUs -# In -DCUDA_COMPUTE_CAPABILITY=70 replace 70 with the desired value -$ cmake -G Ninja .. \ - -DMLIR_DIR=$PWD/../llvm/build/lib/cmake/mlir \ - -DLLVM_DIR=$PWD/../llvm/build/lib/cmake/llvm \ - -DENABLE_GPU_TARGET=ON \ - -DCUDA_COMPUTE_CAPABILITY=70 \ - -DTRITON_PATH=$PWD/../triton/ \ - -DLLVM_ENABLE_ASSERTIONS=ON \ - -DCMAKE_BUILD_TYPE=Release -$ ninja -$ ninja check-comet-integration # Run the integration tests. -``` +# Example for NVIDIA GPU +$ cmake ../ -DENABLE_NVIDIA_GPU_BACKEND=ON -DDEVICE_COMPUTE_CAPABILITY=sm_90 +$ make -The `-DCMAKE_BUILD_TYPE=DEBUG` flag enables debug information, which makes the -whole tree compile slower, but allows you to step through code into the LLVM -and MLIR frameworks. +# Example for AMD GPU +$ cmake ../ -DENABLE_AMD_GPU_BACKEND=ON -DDEVICE_COMPUTE_CAPABILITY=gfx908 +$ make +``` -To get something that runs fast, use `-DCMAKE_BUILD_TYPE=Release` or -`-DCMAKE_BUILD_TYPE=RelWithDebInfo` if you want to go fast and optionally if -you want debug info to go with it. `Release` mode makes a very large difference -in performance. +4) **Enabling FPGA Support.** +To enable support for FPGAs, (currently only Xilinx/AMD), you need to set the flag `-DENABLE_FPGA_TARGET=ON` in `cmake`. The FPGA support relies on other dependencies including an older version of LLVM found as a submodule in `tools/llvm-spirv` and a LLVM-SPIRV translator found in `tools/spriv-llvm-translate`. Setting the above flag will automatically download and install these dependencies, as well as [MCL](https://minos-computing.github.io/) the runtime system used to issue interact with the FPGA. For more information see [here](tools/README.md). ## License diff --git a/cmake/CMakeLists.txt b/cmake/CMakeLists.txt new file mode 100644 index 00000000..b3ecc8ca --- /dev/null +++ b/cmake/CMakeLists.txt @@ -0,0 +1,256 @@ +##===- CMakeLists.txt - COMET cmake root ----------------------*- cmake -*-===// +## +## Configure the COMET build. +## +##===----------------------------------------------------------------------===// + +cmake_minimum_required(VERSION 3.25) + +# If we are not building as a part of LLVM, build COMET as an +# standalone project, using LLVM as an external library: +if( CMAKE_SOURCE_DIR STREQUAL CMAKE_CURRENT_SOURCE_DIR ) + +#------------------------------------------------------------------------------- +# Project setup and globals +#------------------------------------------------------------------------------- +project(comet LANGUAGES CXX C) + +set(CMAKE_CXX_STANDARD 17) +set(CMAKE_CXX_STANDARD_REQUIRED YES) +set(CMAKE_MODULE_PATH ${CMAKE_SOURCE_DIR}/cmake/modules) + +#------------------------------------------------------------------------------- +# Options and settings +#------------------------------------------------------------------------------- +option(ENABLE_AMD_GPU_BACKEND OFF) +option(ENABLE_NVIDIA_GPU_BACKEND OFF) +option(ENABLE_FPGA_TARGET OFF) + +option(LLVM_INCLUDE_TOOLS "Generate build targets for the LLVM tools." ON) +option(LLVM_BUILD_TOOLS "Build the LLVM tools. If OFF, just generate build targets." ON) + +set(CODEGEN_BACKENDS "" CACHE INTERNAL "TRITON BACKENDS") +if (MSVC) + set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} /EHs-c- /GR-") +else () + set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -fno-exceptions -fno-rtti") +endif () + +if(${ENABLE_AMD_GPU_BACKEND}) + add_definitions(-DENABLE_AMD_GPU_BACKEND) + set(CODEGEN_BACKENDS ${CODEGEN_BACKENDS}amd;nvidia) # Triton is looking for some NVGPUIR library + # even if we only enable AMDGPU + set(ENABLE_GPU_TARGET ON) +endif() + +if(${ENABLE_NVIDIA_GPU_BACKEND}) + if(NOT ENABLE_AMD_GPU_BACKEND) + set(CODEGEN_BACKENDS ${CODEGEN_BACKENDS}nvidia;) + endif() + add_definitions(-DENABLE_NVIDIA_GPU_BACKEND) + set(ENABLE_GPU_TARGET ON) +endif() + +if(${ENABLE_FPGA_TARGET}) + add_definitions(-DENABLE_FPGA_TARGET) +endif() + +if(${ENABLE_GPU_TARGET}) + add_definitions(-DENABLE_GPU_TARGET) +endif() + + +#------------------------------------------------------------------------------- +# MLIR/LLVM Configuration +#------------------------------------------------------------------------------- + + find_package(MLIR REQUIRED CONFIG) + + message(STATUS "Using MLIRConfig.cmake in: ${MLIR_DIR}") + message(STATUS "Using LLVMConfig.cmake in: ${LLVM_DIR}") + + set(LLVM_RUNTIME_OUTPUT_INTDIR ${CMAKE_BINARY_DIR}/bin) + set(LLVM_LIBRARY_OUTPUT_INTDIR ${CMAKE_BINARY_DIR}/lib) + + list(APPEND CMAKE_MODULE_PATH "${MLIR_CMAKE_DIR}") + list(APPEND CMAKE_MODULE_PATH "${LLVM_CMAKE_DIR}") + + include(TableGen) + include(AddLLVM) + include(AddMLIR) + include(HandleLLVMOptions) + + set(COMET_BUILT_STANDALONE 1) + set(BACKEND_PACKAGE_STRING "LLVM ${LLVM_PACKAGE_VERSION}") +else() + set(MLIR_MAIN_SRC_DIR ${LLVM_MAIN_SRC_DIR}/../mlir ) # --src-root + set(MLIR_INCLUDE_DIR ${MLIR_MAIN_SRC_DIR}/include ) # --includedir + set(MLIR_TABLEGEN_OUTPUT_DIR ${CMAKE_BINARY_DIR}/tools/mlir/include) + set(MLIR_TABLEGEN_EXE $) + include_directories(SYSTEM ${MLIR_INCLUDE_DIR}) + include_directories(SYSTEM ${MLIR_TABLEGEN_OUTPUT_DIR}) + + set(BACKEND_PACKAGE_STRING "${PACKAGE_STRING}") +endif() + +# Define the default arguments to use with 'lit', and an option for the user to +# override. +set(LIT_ARGS_DEFAULT "-sv") +if (MSVC_IDE OR XCODE) + set(LIT_ARGS_DEFAULT "${LIT_ARGS_DEFAULT} --no-progress-bar") +endif() +set(LLVM_LIT_ARGS "${LIT_ARGS_DEFAULT}" CACHE STRING "Default options for lit") + +# message(STATUS "BLIS and BLAS configuration: ${CMAKE_BINARY_DIR}/blis") +# message(STATUS "BLIS and BLAS configuration: ${CMAKE_CURRENT_BINARY_DIR}/blis") +#------------------------------------------------------------------------------- +# BLIS Configuration +#------------------------------------------------------------------------------- + set(BLA_STATIC ON) #linking with BLIS static library + set(BLA_VENDOR FLAME) + set(BLA_PREFER_PKGCONFIG TRUE) + set(BLA_PKGCONFIG_BLAS blis) + # set(CMAKE_PREFIX_PATH "${CMAKE_BINARY_DIR}/blis/") + find_package(BLAS REQUIRED) + if(BLAS_FOUND) + message(STATUS "BLAS found: ${BLAS_LIBRARIES}") + message(STATUS "BLAS include dirs: ${BLAS_INCLUDE_DIR}") + else() + message(FATAL_ERROR "BLAS not found!") + endif() + +#------------------------------------------------------------------------------- +# MCL include paths +#------------------------------------------------------------------------------- +if(${ENABLE_FPGA_TARGET}) + include_directories("${MCL_BUILD_PATH}/include/") +endif() + +#------------------------------------------------------------------------------- +# Triton Setup +#------------------------------------------------------------------------------- +if(${ENABLE_GPU_TARGET}) + set(TRITON_PATH "" CACHE PATH "Path to Triton source") + set(TRITON_BUILD_UT OFF CACHE BOOL "Disable triton unittests" FORCE) + set(TRITON_CODEGEN_BACKENDS "${CODEGEN_BACKENDS}" CACHE STRING "Enable different codegen backends" FORCE) + if(NOT DEFINED DEVICE_COMPUTE_CAPABILITY) + message(FATAL_ERROR "Please specify the default compute capability for your GPU device") + endif() + add_compile_definitions(DEVICE_COMPUTE_CAPABILITY="${DEVICE_COMPUTE_CAPABILITY}") + + add_subdirectory(${TRITON_PATH} ${TRITON_BUILD_PATH}) + get_property(triton_libs GLOBAL PROPERTY TRITON_LIBS) + include_directories("${TRITON_PATH}") + include_directories("${TRITON_PATH}/include/") + include_directories("${TRITON_PATH}/third_party/") + include_directories("${TRITON_BUILD_PATH}/include/") + include_directories("${TRITON_BUILD_PATH}/third_party/") +endif() + +#------------------------------------------------------------------------------- +# COMET configuration +#------------------------------------------------------------------------------- + +# COMET project. +set(COMET_MAIN_SRC_DIR ${CMAKE_SOURCE_DIR}/../ ) # --src-root +set(COMET_MAIN_INCLUDE_DIR ${COMET_MAIN_SRC_DIR}/include) + +set(COMET_SOURCE_DIR ${CMAKE_SOURCE_DIR}/../) +# set(COMET_BINARY_DIR ${CMAKE_CURRENT_BINARY_DIR}) +set(COMET_INCLUDE_DIR ${CMAKE_CURRENT_BINARY_DIR}/include ) +set(COMET_TOOLS_DIR ${CMAKE_BINARY_DIR}/bin) + +list(APPEND CMAKE_MODULE_PATH "${MLIR_MAIN_SRC_DIR}/cmake/modules") +list(APPEND CMAKE_MODULE_PATH "${CMAKE_CURRENT_SOURCE_DIR}/cmake/modules") + + +# Installing the headers and docs needs to depend on generating any public +# tablegen'd targets. +add_custom_target(comet-headers) +set_target_properties(comet-headers PROPERTIES FOLDER "Misc") +add_custom_target(comet-doc) + +set(CMAKE_INCLUDE_CURRENT_DIR ON) +message(STATUS "aasaaaaa Using BLAS_INCLUDE_LIBRARIES: ${BLAS_INCLUDE_DIRS}") +# Add MLIR, LLVM and BLIS headers to the include path +include_directories(${LLVM_INCLUDE_DIRS}) +include_directories(${MLIR_INCLUDE_DIRS}) +include_directories(${BLAS_INCLUDE_DIRS}) + +# Add COMET files to the include path +include_directories(${COMET_MAIN_INCLUDE_DIR}) +include_directories(${COMET_INCLUDE_DIR}) + +#------------------------------------------------------------------------------- +# Directory setup +#------------------------------------------------------------------------------- +add_subdirectory(${CMAKE_SOURCE_DIR}/../include/comet/ ${CMAKE_BINARY_DIR}/include/comet/) +add_subdirectory(${CMAKE_SOURCE_DIR}/../lib ${CMAKE_BINARY_DIR}/lib) +add_subdirectory(${CMAKE_SOURCE_DIR}/../frontends/comet_dsl ${CMAKE_BINARY_DIR}/frontends/comet_dsl) +add_subdirectory(${CMAKE_SOURCE_DIR}/../test/integration ${CMAKE_BINARY_DIR}/test/integration) + + +option(COMET_INCLUDE_DOCS "Generate build targets for the COMET docs.") +if (COMET_INCLUDE_DOCS) + add_subdirectory(docs) +endif() + + + +install(DIRECTORY include/comet + DESTINATION include + COMPONENT comet-headers + FILES_MATCHING + PATTERN "*.def" + PATTERN "*.h" + PATTERN "*.inc" + PATTERN "*.td" + PATTERN "*.sv" + PATTERN "LICENSE.TXT" + ) + +install(DIRECTORY ${COMET_INCLUDE_DIR}/comet + DESTINATION include + COMPONENT comet-headers + FILES_MATCHING + PATTERN "*.def" + PATTERN "*.h" + PATTERN "*.gen" + PATTERN "*.inc" + PATTERN "*.td" + PATTERN "CMakeFiles" EXCLUDE + PATTERN "config.h" EXCLUDE + ) + +if (NOT LLVM_ENABLE_IDE) + add_llvm_install_targets(install-comet-headers + DEPENDS comet-headers + COMPONENT comet-headers) +endif() + +set(CMAKE_INSTALL_RPATH "${CMAKE_CURRENT_SOURCE_DIR}/install/lib/") +set(CMAKE_INSTALL_RPATH_USE_LINK_PATH TRUE) + +# Set RPATH to $ORIGIN on all targets. +function(set_rpath_all_targets dir) + get_property(subdirectories DIRECTORY ${dir} PROPERTY SUBDIRECTORIES) + foreach(subdir ${subdirectories}) + set_rpath_all_targets(${subdir}) + endforeach() + + get_directory_property(LCL_TARGETS DIRECTORY ${dir} BUILDSYSTEM_TARGETS) + set_property(TARGET ${LCL_TARGETS} PROPERTY INSTALL_RPATH "$ORIGIN/../lib") +endfunction() + +option(STANDALONE_INSTALL "Create an 'install' for packaging which doesn't \ + require installation" off) +if (STANDALONE_INSTALL) + message(STATUS "Setting an $ORIGIN-based RPATH on all executables") + set_rpath_all_targets(${CMAKE_CURRENT_SOURCE_DIR}) +endif() + +option(DEBUG_MODE "Create a installation with debug information" off) +if (DEBUG_MODE) + message(STATUS "Building comet in debug mode") + add_compile_options(-DCOMET_DEBUG_MODE) +endif() diff --git a/frontends/comet_dsl/CMakeLists.txt b/frontends/comet_dsl/CMakeLists.txt index 0447035e..27e17d9d 100644 --- a/frontends/comet_dsl/CMakeLists.txt +++ b/frontends/comet_dsl/CMakeLists.txt @@ -22,27 +22,69 @@ get_property(conversion_libs GLOBAL PROPERTY MLIR_CONVERSION_LIBS) llvm_update_compile_flags(comet-opt) set(LIBS -MLIRAnalysis -MLIRIR -MLIRParser -MLIRPass -MLIRTransforms -COMETUtils -COMETTensorAlgebraDialect -COMETIndexTreeDialect -COMETIndexTreeToSCF + MLIRAnalysis + MLIRIR + MLIRParser + MLIRPass + MLIRFuncAllExtensions + MLIRTransforms + COMETUtils + COMETTensorAlgebraDialect + COMETIndexTreeDialect + # COMETIndexTreeToSCF ) +if(ENABLE_FPGA_TARGET) + set(LIBS + ${LIBS} + COMETParallelLoopsToGpuFPGA + ) +endif() + +if(ENABLE_GPU_TARGET OR ENABLE_FPGA_TARGET) + set(LIBS + ${LIBS} + MLIRGPUToLLVMIRTranslation + ) +endif() + + if(ENABLE_GPU_TARGET) set(LIBS ${LIBS} - COMETParallelLoopsToGpu - COMETGpuToTriton - COMETTritonToCuda + COMETGPUUtils + COMETForallToGpu + COMETGpuToBlockedGpu + COMETBlockedGpuToTriton + #COMETGpuToTriton + COMETPrepareGpuHost ${triton_libs} ) endif() + +if(ENABLE_NVIDIA_GPU_BACKEND) + set(LIBS + ${LIBS} + COMETTritonToCuda + ) +endif() + +if(ENABLE_AMD_GPU_BACKEND) + set(LIBS + ${LIBS} + COMETTritonToHIP + ) +endif() + +if(ENABLE_FPGA_TARGET) + set(LIBS + ${LIBS} + COMETGpuToOCLSPIRV + COMETGpuHostToMCLRT + ) +endif() + target_link_libraries(comet-opt PRIVATE MLIRIR ${LIBS} diff --git a/frontends/comet_dsl/comet.cpp b/frontends/comet_dsl/comet.cpp index 83f093f6..92181855 100644 --- a/frontends/comet_dsl/comet.cpp +++ b/frontends/comet_dsl/comet.cpp @@ -25,30 +25,50 @@ /// //===----------------------------------------------------------------------===// + +#ifdef ENABLE_FPGA_TARGET +#include "comet/Conversion/ParallelLoopsToGpuFPGA/ParallelLoopsToGpuFPGA.h" +#include "comet/Conversion/GpuToOCLSPIRV/GpuToOCLSPIRVPass.h" +#include "comet/Conversion/GpuHostToMCLRT/GpuHostToMCLRTPass.h" +#include "mlir/Conversion/GPUToSPIRV/GPUToSPIRVPass.h" +#include "mlir/Dialect/SPIRV/IR/SPIRVOps.h" +#endif + + +#include "comet/Dialect/IndexTree/IR/IndexTreeDialect.h" +#include "comet/Dialect/IndexTree/Passes.h" #include "comet/Dialect/TensorAlgebra/IR/TADialect.h" #include "comet/Dialect/TensorAlgebra/Passes.h" #include "comet/Dialect/Utils/Utils.h" -#include "comet/Dialect/IndexTree/IR/IndexTreeDialect.h" -#include "comet/Dialect/IndexTree/Passes.h" #include "comet/Conversion/Passes.h" #include "MLIRGen.h" #include "Parser.h" -#include "mlir/Support/TypeID.h" +#include "mlir/Conversion/ConvertToLLVM/ToLLVMPass.h" +#include "mlir/Conversion/GPUCommon/GPUCommonPass.h" +#include "mlir/Dialect/GPU/IR/GPUDialect.h" +#include "mlir/Dialect/SPIRV/Transforms/Passes.h" +#include "mlir/Conversion/SCFToOpenMP/SCFToOpenMP.h" +#include "mlir/Dialect/Arith/IR/Arith.h" +#include "mlir/Dialect/Func/Extensions/InlinerExtension.h" +#include "mlir/Dialect/Func/Transforms/Passes.h" +#include "mlir/Dialect/GPU/IR/GPUDialect.h" +#include "mlir/Dialect/LLVMIR/Transforms/Passes.h" #include "mlir/Dialect/Linalg/IR/Linalg.h" #include "mlir/Dialect/Linalg/Passes.h" -#include "mlir/Dialect/LLVMIR/Transforms/Passes.h" -#include "mlir/Dialect/Arith/IR/Arith.h" #include "mlir/Dialect/MemRef/Transforms/Passes.h" -#include "mlir/Dialect/Func/Transforms/Passes.h" #include "mlir/Dialect/Tensor/Transforms/Passes.h" +#include "mlir/Dialect/SCF/Transforms/Passes.h" +#include "mlir/Dialect/Bufferization/Transforms/Passes.h" + +#include "mlir/Support/TypeID.h" -#include "mlir/Conversion/Passes.h" #include "mlir/Conversion/AffineToStandard/AffineToStandard.h" #include "mlir/Conversion/FuncToLLVM/ConvertFuncToLLVMPass.h" #include "mlir/Conversion/IndexToLLVM/IndexToLLVM.h" +#include "mlir/Conversion/Passes.h" #include "mlir/Conversion/VectorToLLVM/ConvertVectorToLLVMPass.h" // #include "mlir/Conversion/LinalgToLLVM/LinalgToLLVM.h" #include "mlir/Conversion/MathToLLVM/MathToLLVM.h" @@ -58,17 +78,21 @@ // #include "mlir/Conversion/VectorToLLVM/ConvertVectorToLLVM.h" #include "mlir/Conversion/VectorToSCF/VectorToSCF.h" -#include "mlir/IR/Verifier.h" #include "mlir/ExecutionEngine/ExecutionEngine.h" #include "mlir/ExecutionEngine/OptUtils.h" -#include "mlir/IR/MLIRContext.h" #include "mlir/IR/BuiltinOps.h" +#include "mlir/IR/MLIRContext.h" +#include "mlir/IR/Verifier.h" #include "mlir/InitAllDialects.h" #include "mlir/Parser/Parser.h" #include "mlir/Pass/Pass.h" #include "mlir/Pass/PassManager.h" +#include "mlir/Target/LLVMIR/Dialect/AMX/AMXToLLVMIRTranslation.h" +#include "mlir/Target/LLVMIR/Dialect/GPU/GPUToLLVMIRTranslation.h" +#include "mlir/Target/LLVMIR/Dialect/ROCDL/ROCDLToLLVMIRTranslation.h" #include "mlir/Target/LLVMIR/Export.h" #include "mlir/Transforms/Passes.h" +#include "mlir/IR/DialectRegistry.h" #include "llvm/ADT/StringRef.h" #include "llvm/IR/Module.h" #include "llvm/Support/CommandLine.h" @@ -79,7 +103,14 @@ #include "llvm/Support/TargetSelect.h" #include "llvm/Support/raw_ostream.h" #ifdef ENABLE_GPU_TARGET -#include "comet/Conversion/ParallelLoopsToGpu/ParallelLoopsToGpu.h" +#include "comet/Conversion/PrepareGpuHost/PrepareGpuHostPass.h" +#include "comet/Conversion/BlockedGpuToTriton/BlockedGpuToTriton.h" +#include "comet/Conversion/GpuToBlockedGpu/GpuToBlockedGpu.h" +#include "comet/Conversion/ForallToGpu/ForallToGpu.h" +#include "comet/Conversion/TritonToHIP/TritonToHIPPass.h" +#include "triton/Dialect/TritonGPU/IR/Dialect.h" +#include "mlir/Conversion/SCFToGPU/SCFToGPUPass.h" +#include "mlir/Dialect/GPU/Transforms/Passes.h" #include "comet/Conversion/GpuToTriton/GpuToTritonPass.h" #include "comet/Conversion/TritonToCuda/TritonToCudaPass.h" #include "triton/Dialect/Triton/IR/Dialect.h" @@ -88,12 +119,11 @@ #include "mlir/Dialect/GPU/Transforms/Passes.h" #endif -#include "mlir/Target/LLVMIR/Dialect/All.h" -#include "mlir/Target/LLVMIR/Dialect/All.h" -#include "mlir/Dialect/LLVMIR/NVVMDialect.h" -#include "mlir/InitAllPasses.h" #include "mlir/Conversion/NVVMToLLVM/NVVMToLLVM.h" #include "mlir/Conversion/Passes.h" +#include "mlir/Dialect/LLVMIR/NVVMDialect.h" +#include "mlir/InitAllPasses.h" +#include "mlir/Target/LLVMIR/Dialect/All.h" // #ifdef ENABLE_GPU_TARGET // #include "comet/TritonConfig.h" // #endif @@ -129,18 +159,12 @@ using namespace mlir::indexTree; #define DEBUG_TYPE "comet_dsl" namespace cl = llvm::cl; -static cl::opt inputFilename(cl::Positional, - cl::desc(""), - cl::init("-"), - cl::value_desc("filename")); +static cl::opt + inputFilename(cl::Positional, cl::desc(""), + cl::init("-"), cl::value_desc("filename")); -namespace -{ - enum InputType - { - TensorAlgebra, - MLIR - }; +namespace { +enum InputType { TensorAlgebra, MLIR }; } static cl::opt inputType( @@ -167,20 +191,41 @@ static cl::opt emitLLVM("emit-llvm", cl::desc("output the LLVM dialect dum /// ============================================================================= static cl::opt CodegenTarget("target", cl::init(CPU), cl::desc("Code generation target"), - cl::values( + cl::values( clEnumVal(CPU, "Codegen target is CPU") #ifdef ENABLE_GPU_TARGET , clEnumVal(GPU, "Codegen target is GPU") #endif - ) - ); - + #ifdef ENABLE_FPGA_TARGET + , + clEnumVal(FPGA, "Codegen target is FPGA") + #endif + ) +); + +#if defined(ENABLE_GPU_TARGET) | defined(ENABLE_FPGA_TARGET) +static cl::opt GPUBlockSizeX("kernel-block-x-size", cl::init(32), cl::desc("Kernel Block size in X direction")); +static cl::opt GPUBlockSizeY("kernel-block-y-size", cl::init(8), cl::desc("Kernel Block size in Y direction")); +#if defined(ENABLE_GPU_TARGET) +static cl::opt GPUBlockSizeR("kernel-block-r-size", cl::init(32), cl::desc("Kernel Block size in R direction")); +#else +static cl::opt GPUBlockSizeR("kernel-block-r-size", cl::init(1), cl::desc("Kernel Block size in R direction")); +#endif +#endif +#ifdef ENABLE_FPGA_TARGET +static cl::opt xclbinPath("xclbin_path", cl::init("-"), cl::desc("Path to xclbin")); +static cl::opt sprirvBinOutPath("spirv_bin_path", cl::init("-"), cl::desc("Path to output SPIRV binary")); +#endif #ifdef ENABLE_GPU_TARGET -static cl::opt GPUBlockSizeX("gpu-block-x-size", cl::init(32), cl::desc("GPU Block size in X direction")); -static cl::opt GPUBlockSizeY("gpu-block-y-size", cl::init(8), cl::desc("GPU Block size in Y direction")); -static cl::opt GPUBlockSizeR("gpu-block-r-size", cl::init(32), cl::desc("GPU Block size in R direction")); -static cl::opt GPUComputeCapability("gpu-compute-capability", cl::init(CUDA_COMPUTE_CAPABILITY), cl::desc("GPU compute capability")); +static cl::opt GPUTargetCompilationFormat( + "gpu-code-format", cl::init(Binary), + cl::desc("GPU target code generation format"), + cl::values(clEnumVal(Assembly, "GPU target format is assembly"), + clEnumVal(Binary, "GPU target format is binary"), + clEnumVal(Fatbin, "GPU target format is fat binary") + )); +static cl::opt GPUComputeCapability("gpu-compute-capability", cl::init(DEVICE_COMPUTE_CAPABILITY), cl::desc("GPU target architecture")); static cl::opt GPUNumWarps("gpu-num-warps", cl::init(4), cl::desc("GPU number of warps")); static cl::opt GPUThreadsPerWarp("gpu-threads-per-warp", cl::init(32), cl::desc("GPU threads per warp")); static cl::opt GPUNumCTAs("gpu-num-ctas", cl::init(1), cl::desc("GPU num CTAs")); @@ -193,10 +238,7 @@ static cl::opt GPUNumStages("gpu-num-stages", cl::init(1), cl::desc("GPU nu static cl::opt OptMultiOpFactorization("opt-multiop-factorize", cl::desc("Multi operations factorization optimization")); -static cl::opt IsSelectBestPermTTGT("opt-bestperm-ttgt", - cl::desc("Select the best index permutation for TTGT, otherwise the first appropriate permutation")); - -static cl::opt selectedPermNum("perm-num", cl::init(1), +static cl::opt selectedPermNum("perm-num", cl::init(-1), cl::ZeroOrMore, cl::desc("Select the permutation number to choose")); /// ============================================================================= @@ -226,6 +268,8 @@ static cl::opt OptWorkspace("opt-comp-workspace", cl::init(false), /// ============================================================================= static cl::opt OptKernelFusion("opt-fusion", cl::init(false), cl::desc("Output IT dialect after redundancy-aware fusion")); +static cl::opt OptDimensionReduction("opt-dimension-reduction", cl::init(false), + cl::desc("Reduce intermediate tensors' dimension after kernel fusion")); /// ============================================================================= /// TTGT reformulation for tensor contraction operations @@ -250,8 +294,12 @@ static cl::opt IsLoweringtoSCF("convert-to-loops", /// ============================================================================= /// Lowering loops to Triton /// ============================================================================= -static cl::opt IsLoweringtoTriton("convert-to-triton", +static cl::opt IsLoweringToTriton("convert-to-triton", cl::desc("Output Triton dialect after lowering all operations")); + +static cl::opt IsGeneratingGpuAllocsAndTransfers("gpu-generate-allocs-transfers", cl::init(true), + + cl::desc("Whether to generate GPU allocations and transfers")); #endif /// ============================================================================= @@ -260,6 +308,15 @@ static cl::opt IsLoweringtoTriton("convert-to-triton", static cl::opt isLoweringToLLVM("convert-to-llvm", cl::desc("Output LLVM IR")); +/// ============================================================================= +/// Debug Options +/// ============================================================================= +static cl::opt GenTensorAlgebraLabelsInAlphabeticalOrder( + "debug-ta-labels-alphabet-order", + cl::desc("If turned on, when generating Tensor Algebra dialect, the order of the affine maps' dimensions will be " + "the alphabetical order of index labels. " + "For example, for C[i,k] = A[i,h] * B[h,k], the d0, d1, d2 will be h, i, k.")); + /// ============================================================================= /// Utility functions /// ============================================================================= @@ -283,7 +340,7 @@ std::unique_ptr parseInputFile(llvm::StringRef filenam } int loadMLIR(mlir::MLIRContext &context, - mlir::OwningOpRef &module) + mlir::OwningOpRef &module, bool useI64) { /// Handle '.ta' input to the compiler. if (inputType != InputType::MLIR && @@ -292,7 +349,7 @@ int loadMLIR(mlir::MLIRContext &context, auto moduleAST = parseInputFile(inputFilename); if (!moduleAST) return 6; - module = mlirGen(context, *moduleAST); + module = mlirGen(context, *moduleAST, useI64); return !module ? 1 : 0; } @@ -318,7 +375,7 @@ int loadMLIR(mlir::MLIRContext &context, } int loadAndProcessMLIR(mlir::MLIRContext &context, - mlir::OwningOpRef &module) + mlir::OwningOpRef &module, bool useI64) { #ifdef ENABLE_GPU_TARGET bool emitTriton_ = emitTriton && CodegenTarget == TargetDevice::GPU; @@ -326,7 +383,13 @@ int loadAndProcessMLIR(mlir::MLIRContext &context, bool emitTriton_ = false; #endif - if (int error = loadMLIR(context, module)) + /// Load debug options + if (GenTensorAlgebraLabelsInAlphabeticalOrder) + { + tensorAlgebra::debugOptions.insert("debug-ta-labels-alphabet-order"); + } + /// end Load debug options + if (int error = loadMLIR(context, module, useI64)) return error; mlir::PassManager pm(module.get()->getName()); @@ -339,6 +402,14 @@ int loadAndProcessMLIR(mlir::MLIRContext &context, mlir::OpPassManager &optPM = pm.nest(); + /// Check to see if we are dumping to TA dialect. + if (emitTA) + { + if (mlir::failed(pm.run(*module))) + return 4; + return 0; + } + /// ============================================================================= /// High-level optimization at the TA dialect /// Such as finding the optimal ordering of dense tensor contractions, or reformulating tensor contractions @@ -372,20 +443,17 @@ int loadAndProcessMLIR(mlir::MLIRContext &context, /// =================================================================================== if (IsLoweringtoIndexTree || emitIT || emitLoops || emitTriton_ || emitLLVM) { - /// Generate the index tree IR - optPM.addPass(mlir::comet::createLowerTensorAlgebraToIndexTreePass(CodegenTarget)); + /// Generate the index tree IR + optPM.addPass(mlir::comet::createLowerTensorAlgebraToIndexTreePass(CodegenTarget)); if (OptKernelFusion) { - /// Apply partial fusion on index tree dialect for some compound expressions. + /// Apply kernel fusion on index tree dialect for some compound expressions. optPM.addPass(mlir::comet::createIndexTreeKernelFusionPass()); } - if (OptWorkspace) - { - /// Optimized workspace transformations, reduce iteration space for nonzero elements - optPM.addPass(mlir::comet::createIndexTreeWorkspaceTransformationsPass()); - } + // Create new pass manager to optimize the index tree dialect + optPM.addPass(mlir::comet::createIndexTreeDomainInferencePass()); /// Dump index tree dialect. if (emitIT) @@ -396,6 +464,15 @@ int loadAndProcessMLIR(mlir::MLIRContext &context, } } + /// Concretize the domains of all the index variables + optPM.addPass(mlir::comet::createIndexTreeDomainConcretizationPass()); + + if (OptKernelFusion || OptDimensionReduction) + { + /// Reduce intermediate tensors' dimension after kernel fusion + optPM.addPass(mlir::comet::createIndexTreeDimensionReductionPass()); + } + /// ============================================================================= /// ============================================================================= @@ -408,49 +485,42 @@ int loadAndProcessMLIR(mlir::MLIRContext &context, /// input and output sparse tensor declaration lowering are distant and need different information optPM.addPass(mlir::comet::createSparseTensorDeclLoweringPass()); optPM.addPass(mlir::comet::createDenseTensorDeclLoweringPass()); + optPM.addPass(mlir::comet::createSparseTempOutputTensorDeclLoweringPass()); + optPM.addPass(mlir::comet::createSparseOutputTensorDeclLoweringPass()); optPM.addPass(mlir::comet::createTensorFillLoweringPass()); + optPM.addPass(mlir::comet::createDimOpLoweringPass()); /// ============================================================================= + optPM.addPass(mlir::comet::createIndexTreeDomainConcretizationPass()); + + if (OptWorkspace) { + /// Optimized workspace transformations, reduce iteration space for nonzero elements + optPM.addPass(mlir::comet::createIndexTreeWorkspaceTransformationsPass()); + } + /// TTGT reformulation for dense tensor contraction operations if (IsLoweringTCtoTTGT) { /// Sparse input and dense input/output tensor declarations needed be lowered before for TTGT pass - optPM.addPass(mlir::comet::createLoweringTTGTPass(IsSelectBestPermTTGT, selectedPermNum, IsPrintFlops)); + optPM.addPass(mlir::comet::createLoweringTTGTDynPass(selectedPermNum, IsPrintFlops)); + // optPM.addPass(mlir::comet::createLoweringTTGTPass(true, selectedPermNum, IsPrintFlops)); } - /// ============================================================================= - /// Operation based optimizations - /// ============================================================================= + // /// ============================================================================= + // /// Operation based optimizations + // /// ============================================================================= if (OptMatmulTiling) { optPM.addPass(mlir::comet::createLinAlgMatmulTilingPass()); } - if (OptCallToMatMulMicroKernel) - { - optPM.addPass(mlir::comet::createLinAlgMatmulMicroKernelPass()); - } /// ============================================================================= /// Lowering all the operations to loops /// ============================================================================= - if (IsLoweringtoSCF || emitLoops || emitTriton_ || emitLLVM ) - { - - /// Workspace transformations will create new dense tensor declarations, so we need to call createDenseTensorDeclLoweringPass - optPM.addPass(mlir::comet::createDenseTensorDeclLoweringPass()); /// lowers dense input/output tensor declaration - optPM.addPass(mlir::comet::createSparseTempOutputTensorDeclLoweringPass()); /// Temporary sparse output tensor declarations introduced by compound expressions - /// should be lowered before sparse output tensor declarations - optPM.addPass(mlir::comet::createSparseOutputTensorDeclLoweringPass()); /// lowering for sparse output tensor declarations - //(sparse_output_tensor_decl and temp_sparse_output_tensor_decl) - - optPM.addPass(mlir::comet::createDimOpLoweringPass()); - - /// The partial Fusion pass might add new tensor.fill operations - optPM.addPass(mlir::comet::createTensorFillLoweringPass()); - optPM.addPass(mlir::comet::createPCToLoopsLoweringPass()); - + if (IsLoweringtoSCF || emitLoops || emitLLVM || emitTriton_) + { /// ============================================================================= /// Lowering of other operations such as transpose, sum, etc. to SCF dialect /// ============================================================================= @@ -458,11 +528,17 @@ int loadAndProcessMLIR(mlir::MLIRContext &context, /// If it is a transpose of sparse tensor, it lowers the code to make a runtime call to specific sorting algorithm optPM.addPass(mlir::comet::createLowerTensorAlgebraToSCFPass()); + optPM.addPass(mlir::comet::createIndexTreeSymbolicComputePass()); + /// Finally lowering index tree to SCF dialect optPM.addPass(mlir::comet::createLowerIndexTreeToSCFPass()); - optPM.addPass(mlir::tensor::createTensorBufferizePass()); - pm.addPass(mlir::func::createFuncBufferizePass()); /// Needed for func - pm.addPass(mlir::createConvertLinalgToLoopsPass()); + optPM.addPass(mlir::comet::createWorkspaceOptimizationsPass()); + optPM.addPass(mlir::comet::createConvertSymbolicDomainsPass()); + optPM.addPass(mlir::comet::createSparseTensorConversionPass()); + optPM.addPass(mlir::comet::createIndexTreeInliningPass()); + optPM.addPass(mlir::createCanonicalizerPass()); + optPM.addPass(mlir::createLoopInvariantCodeMotionPass()); + optPM.addPass(mlir::createLoopInvariantSubsetHoistingPass()); if (OptDenseTransposeOp) /// Optimize Dense Transpose operation { @@ -487,34 +563,121 @@ int loadAndProcessMLIR(mlir::MLIRContext &context, - /// ============================================================================= - /// Late lowering passes - /// ============================================================================= + // /// ============================================================================= + // /// Late lowering passes + // /// ============================================================================= + // pm.addPass(mlir::bufferization::createEmptyTensorToAllocTensorPass()); + pm.addPass(mlir::comet::createTABufferizeFunc()); + pm.addPass(mlir::createCanonicalizerPass()); + pm.addPass(mlir::createLowerAffinePass()); - optPM.addPass(mlir::comet::createSTCRemoveDeadOpsPass()); - optPM.addPass(mlir::comet::createLateLoweringPass()); - // pm.addPass(mlir::createCanonicalizerPass()); - optPM.addPass(mlir::createCSEPass()); + mlir::bufferization::OneShotBufferizationOptions opts; + opts.allowUnknownOps = true; + pm.addPass(mlir::bufferization::createOneShotBufferizePass(opts)); -#ifdef ENABLE_GPU_TARGET - if (CodegenTarget == TargetDevice::GPU && (emitTriton_ || emitLLVM || IsLoweringtoTriton)) + mlir::OpPassManager &late_lowering_pm = pm.nest(); + late_lowering_pm.addPass(mlir::comet::createSTCRemoveDeadOpsPass()); + late_lowering_pm.addPass(mlir::comet::createLateLoweringPass()); + + pm.addPass(mlir::createCanonicalizerPass()); + pm.addPass(mlir::createCSEPass()); + pm.addPass(mlir::memref::createFoldMemRefAliasOpsPass()); + pm.addPass(mlir::createCanonicalizerPass()); + if (OptCallToMatMulMicroKernel) + { + pm.addNestedPass(mlir::comet::createLinAlgMatmulMicroKernelPass()); + pm.addNestedPass(mlir::comet::createMatvecToParallelLoopsPass()); + + } + pm.addNestedPass(mlir::createConvertVectorToSCFPass()); + /// Blanket-convert any remaining linalg ops to loops if any remain. + pm.addNestedPass( + mlir::createConvertLinalgToLoopsPass()); + /// Blanket-convert any remaining affine ops if any remain. + pm.addPass(mlir::createLowerAffinePass()); + /// Convert SCF to CF (always needed). + #ifdef ENABLE_GPU_TARGET + if(CodegenTarget != TargetDevice::GPU) + { + pm.addPass(mlir::createForallToParallelLoopPass()); + pm.addPass(mlir::createLoopInvariantCodeMotionPass()); + pm.addPass(mlir::createCanonicalizerPass()); + } + #else + pm.addPass(mlir::createForallToParallelLoopPass()); + pm.addPass(mlir::createLoopInvariantCodeMotionPass()); + pm.addPass(mlir::createCanonicalizerPass()); + #endif + + +#ifndef ENABLE_GPU_TARGET + [[maybe_unused]] bool IsLoweringToTriton = false; +#endif +#if defined(ENABLE_GPU_TARGET) | defined(ENABLE_FPGA_TARGET) + if ((CodegenTarget == TargetDevice::GPU || CodegenTarget == TargetDevice::FPGA) && (emitTriton_ || emitLLVM || isLoweringToLLVM || IsLoweringToTriton)) { - pm.addNestedPass(mlir::comet::createConvertParallelLoopsToGpuPass(GPUBlockSizeX, GPUBlockSizeY, GPUBlockSizeR)); + #ifdef ENABLE_FPGA_TARGET + if (CodegenTarget == TargetDevice::FPGA) + { + pm.addNestedPass(mlir::comet::createConvertParallelLoopsToGpuFPGAPass(GPUBlockSizeX, GPUBlockSizeY, GPUBlockSizeR, CodegenTarget)); + } + #endif + #ifdef ENABLE_GPU_TARGET + if (CodegenTarget == TargetDevice::GPU) + { + pm.addNestedPass(mlir::comet::createConvertForallToGpuPass(GPUBlockSizeX, GPUBlockSizeY, GPUBlockSizeR)); + } + #endif + pm.addPass(mlir::createLoopInvariantCodeMotionPass()); pm.addPass(mlir::createParallelLoopToGpuPass()); pm.addPass(mlir::createGpuKernelOutliningPass()); pm.addPass(mlir::createCanonicalizerPass()); - pm.addPass(mlir::comet::createConvertGpuKernelToTritonPass()); + } + + #ifdef ENABLE_GPU_TARGET + if(CodegenTarget == TargetDevice::GPU && (emitTriton_ || emitLLVM || IsLoweringToTriton || isLoweringToLLVM)) + { + + pm.nest().addNestedPass(mlir::comet::createConvertGpuToBlockedGpuPass()); + pm.addPass(mlir::comet::createConvertBlockedGpuToTritonPass()); + pm.addPass(mlir::createLoopInvariantCodeMotionPass()); + pm.addPass(mlir::createCSEPass());; + pm.addPass(mlir::createSymbolDCEPass());; + pm.addPass(mlir::createCanonicalizerPass()); + // pm.addPass(mlir::comet::createConvertGpuKernelToTritonPass()); if (emitTriton_) { + if (mlir::failed(pm.run(*module))) return 4; return 0; } } + #endif + +#endif + +#ifdef ENABLE_FPGA_TARGET + + if (CodegenTarget == TargetDevice::FPGA && (isLoweringToLLVM || emitLLVM)) + { + pm.addPass(mlir::createLowerAffinePass()); + pm.addPass(mlir::comet::createConvertGpuHostToMCLRTPass(xclbinPath.c_str())); + if(sprirvBinOutPath == "-") + { + pm.addPass(mlir::comet::createConvertGPUKernelToOCLSPIRVPass(GPUBlockSizeX, GPUBlockSizeY, GPUBlockSizeR, inputFilename.c_str())); + } + else + { + pm.addPass(mlir::comet::createConvertGPUKernelToOCLSPIRVPass(GPUBlockSizeX, GPUBlockSizeY, GPUBlockSizeR, sprirvBinOutPath.c_str())); + } + + } #endif + pm.addPass(mlir::createCanonicalizerPass()); /// ============================================================================= @@ -522,27 +685,44 @@ int loadAndProcessMLIR(mlir::MLIRContext &context, if (isLoweringToLLVM || emitLLVM) { #ifdef ENABLE_GPU_TARGET - if (CodegenTarget == GPU) + if ((isLoweringToLLVM || emitLLVM) && CodegenTarget == TargetDevice::GPU) { - pm.addPass(mlir::comet::createLowerTritonDeviceToCudaPass(GPUNumWarps, GPUThreadsPerWarp, GPUNumCTAs, GPUNumStages, GPUComputeCapability)); - pm.addPass(mlir::comet::createLowerGpuHostToCudaPass()); + if(GPUComputeCapability.getValue().find("sm_") != std::string::npos || GPUComputeCapability.getValue().find("compute_") != std::string::npos) + { + #ifdef ENABLE_NVIDIA_GPU_BACKEND + int32_t cudaCC = std::stoi(GPUComputeCapability.substr(GPUComputeCapability.find("_")+1)); + pm.addPass(mlir::comet::createLowerTritonDeviceToCudaPass(GPUNumWarps, GPUThreadsPerWarp, GPUNumCTAs, GPUNumStages, cudaCC, GPUTargetCompilationFormat)); + pm.addPass(mlir::comet::createPrepareGpuHostPass(IsGeneratingGpuAllocsAndTransfers)); + pm.addPass(mlir::comet::createLowerGpuHostToCudaPass()); + #else + llvm::errs() << "Trying to lower to NVIDIA(?) device without enabling the NVIDIA backend \n"; + return 6; + #endif + } + else { + #ifdef ENABLE_AMD_GPU_BACKEND + pm.addPass(mlir::comet::createLowerTritonDeviceToHIPPass(GPUNumWarps, GPUThreadsPerWarp, GPUNumCTAs, GPUNumStages, GPUComputeCapability, GPUTargetCompilationFormat)); + pm.addPass(mlir::comet::createPrepareGpuHostPass(IsGeneratingGpuAllocsAndTransfers)); + pm.addPass(mlir::comet::createLowerGpuHostToHIPPass()); + #else + llvm::errs() << "Trying to lower to AMDGPU(?) device without enabling the NVIDIA backend \n"; + return 6; + #endif + } } -#endif +#endif optPM.addPass(mlir::createCanonicalizerPass()); /// Blanket-convert any remaining high-level vector ops to loops if any remain. - pm.addNestedPass(mlir::createConvertVectorToSCFPass()); - /// Blanket-convert any remaining linalg ops to loops if any remain. - pm.addNestedPass(mlir::createConvertLinalgToLoopsPass()); - /// Blanket-convert any remaining affine ops if any remain. - pm.addPass(mlir::createLowerAffinePass()); - /// Convert SCF to CF (always needed). - pm.addPass(mlir::createConvertSCFToCFPass()); - /// Sprinkle some cleanups. + + pm.addPass(mlir::createConvertSCFToOpenMPPass()); pm.addPass(mlir::createCanonicalizerPass()); pm.addPass(mlir::createCSEPass()); + pm.addPass(mlir::createConvertSCFToCFPass()); + /// Sprinkle some cleanups. /// Convert vector to LLVM (always needed). - pm.addPass(mlir::createConvertVectorToLLVMPass()); // TODO: add more options on a per-need basis. + pm.addPass(mlir::createConvertVectorToLLVMPass()); // TODO: add more options + // on a per-need basis. //// Convert Math to LLVM (always needed). pm.addNestedPass(mlir::createConvertMathToLLVMPass()); /// Expand complicated MemRef operations before lowering them. @@ -550,11 +730,13 @@ int loadAndProcessMLIR(mlir::MLIRContext &context, /// The expansion may create affine expressions. Get rid of them. pm.addPass(mlir::createLowerAffinePass()); /// Convert MemRef to LLVM (always needed). - pm.addPass(mlir::createFinalizeMemRefToLLVMConversionPass()); /// Convert Func to LLVM (always needed). + pm.addPass(mlir::createConvertControlFlowToLLVMPass()); + pm.addPass(mlir::createFinalizeMemRefToLLVMConversionPass()); pm.addPass(mlir::createConvertFuncToLLVMPass()); /// Convert Index to LLVM (always needed). pm.addPass(mlir::createConvertIndexToLLVMPass()); + pm.addPass(mlir::createConvertOpenMPToLLVMPass()); /// Convert remaining unrealized_casts (always needed). pm.addPass(mlir::createReconcileUnrealizedCastsPass()); @@ -589,6 +771,9 @@ int main(int argc, char **argv) mlir::MLIRContext context; mlir::registerAllDialects(context); + mlir::DialectRegistry registry; + mlir::tensorAlgebra::registerBufferizableOpInterfaceExternalModels(registry); + context.appendDialectRegistry(registry); mlir::registerPassManagerCLOptions(); cl::ParseCommandLineOptions(argc, argv, "Tensor Algebra compiler\n"); @@ -600,15 +785,25 @@ int main(int argc, char **argv) /// Register our Dialect with MLIR. #ifdef ENABLE_GPU_TARGET context.loadDialect(); - registerLLVMDialectTranslation(context); + mlir::func::registerInlinerExtension(registry); + context.appendDialectRegistry(registry); registerLLVMDialectTranslation(context); registerBuiltinDialectTranslation(context); + mlir::registerGPUDialectTranslation(context); + #endif + + #ifdef ENABLE_AMD_GPU_BACKEND + mlir::registerROCDLDialectTranslation(context); + #endif + + #ifdef ENABLE_NVIDIA_GPU_BACKEND registerNVVMDialectTranslation(context); LLVMInitializeNVPTXTargetInfo(); LLVMInitializeNVPTXTarget(); LLVMInitializeNVPTXTargetMC(); LLVMInitializeNVPTXAsmPrinter(); #endif + context.loadDialect(); context.loadDialect(); context.loadDialect(); @@ -616,10 +811,19 @@ int main(int argc, char **argv) context.loadDialect(); context.loadDialect(); context.loadDialect(); + context.loadDialect(); mlir::OwningOpRef module; + bool useI64 = true; + #ifdef ENABLE_GPU_TARGET + if(CodegenTarget == TargetDevice::GPU) + { + useI64 = false; + } + #endif + - if (int error = loadAndProcessMLIR(context, module)) + if (int error = loadAndProcessMLIR(context, module, useI64)) return error; /// If we aren't exporting to non-mlir, then we are done. diff --git a/frontends/comet_dsl/include/MLIRGen.h b/frontends/comet_dsl/include/MLIRGen.h index b3b76dc2..8ade7992 100644 --- a/frontends/comet_dsl/include/MLIRGen.h +++ b/frontends/comet_dsl/include/MLIRGen.h @@ -30,6 +30,8 @@ #define COMET_DSL_MLIRGEN_H_ #include +#include +#include namespace mlir { @@ -46,7 +48,9 @@ namespace tensorAlgebra /// Emit IR for the given Tensor Algebra moduleAST, returns a newly created MLIR module /// or nullptr on failure. mlir::OwningOpRef mlirGen(mlir::MLIRContext &context, - ModuleAST &moduleAST); + ModuleAST &moduleAST, bool useI64); + + extern std::unordered_set debugOptions; } /// namespace tensorAlgebra #endif /// COMET_DSL_MLIRGEN_H_ diff --git a/frontends/comet_dsl/mlir/MLIRGen.cpp b/frontends/comet_dsl/mlir/MLIRGen.cpp index 45778086..2a6a5ce0 100644 --- a/frontends/comet_dsl/mlir/MLIRGen.cpp +++ b/frontends/comet_dsl/mlir/MLIRGen.cpp @@ -36,6 +36,8 @@ #include "mlir/Dialect/Arith/IR/Arith.h" #include "mlir/Dialect/Tensor/IR/Tensor.h" #include "mlir/Dialect/Func/IR/FuncOps.h" +#include "mlir/IR/BuiltinAttributes.h" +#include "mlir/IR/BuiltinTypeInterfaces.h" #include "mlir/IR/Verifier.h" #include "mlir/IR/Attributes.h" #include "mlir/IR/Builders.h" @@ -43,18 +45,23 @@ #include "mlir/IR/MLIRContext.h" #include "mlir/IR/Operation.h" +#include "llvm/ADT/ArrayRef.h" #include "llvm/ADT/STLExtras.h" #include "llvm/ADT/ScopedHashTable.h" +#include "llvm/ADT/StringRef.h" #include "llvm/Support/raw_ostream.h" #include +#include #include #include #include /// for random num generation #include /// for seed of random num generation +#include #include "mlir/IR/Types.h" #include "mlir/IR/BuiltinTypes.h" #include "mlir/IR/Verifier.h" +#include "mlir/Support/LLVM.h" using namespace mlir::tensorAlgebra; using namespace tensorAlgebra; @@ -68,11 +75,12 @@ using llvm::ScopedHashTableScope; using llvm::SmallVector; using llvm::StringRef; using llvm::Twine; +int32_t defaultSpTensorIndiceBitWidth = 64; // TODO: We should be able to pass this from the DSL using StringSet = std::set; // *********** For debug purpose *********// -// #define COMET_DEBUG_MODE +//#define COMET_DEBUG_MODE #include "comet/Utils/debug.h" #undef COMET_DEBUG_MODE // *********** For debug purpose *********// @@ -471,9 +479,9 @@ namespace mlir::StringAttr opAttr = builder.getStringAttr(op); mlir::RankedTensorType returnDataType; - if(lhs.getType().cast().getShape() != rhs.getType().cast().getShape()) + if(mlir::cast(lhs.getType()).getShape() != mlir::cast(rhs.getType()).getShape()) { - returnDataType = lhs.getType().cast(); + returnDataType = mlir::cast(lhs.getType()); auto bcastRhs = builder.create(location, returnDataType, mlir::cast(rhs.getDefiningOp()).getValueAttr()); comet_vdump(bcastRhs); rhs.replaceAllUsesWith(bcastRhs); @@ -582,24 +590,42 @@ namespace } comet_debug() << "\n"; - auto lhs_tensor = lhs.getDefiningOp()->getOpResult(0).getType(); - assert(lhs_tensor.isa()); + auto lhs_tensor = lhs.getType(); comet_pdump(lhs.getDefiningOp()); + auto lhs_labeledtensor = lhs.getDefiningOp()->getOpResult(0); comet_vdump(lhs_labeledtensor); // ta.labeled_tensor - auto lhs_el_type = lhs_tensor.cast().getElementType(); + mlir::Type lhs_el_type; + if(auto tensor_type = llvm::dyn_cast(lhs_tensor)){ + lhs_el_type = tensor_type.getElementType(); + } + else if(auto tensor_type = llvm::dyn_cast(lhs_tensor)){ + lhs_el_type = tensor_type.getElementType(); + } + else { + assert(false && "Expected a tensor input"); + } auto rhs_tensor = rhs.getDefiningOp()->getOpResult(0).getType(); comet_pdump(rhs.getDefiningOp()); - assert(rhs_tensor.isa()); auto rhs_labeledtensor = rhs.getDefiningOp()->getOpResult(0); comet_vdump(rhs_labeledtensor); - auto rhs_el_type = rhs_tensor.cast().getElementType(); + mlir::Type rhs_el_type; + if(auto tensor_type = llvm::dyn_cast(rhs_tensor)){ + rhs_el_type = tensor_type.getElementType(); + } + else if(auto tensor_type = llvm::dyn_cast(rhs_tensor)){ + rhs_el_type = tensor_type.getElementType(); + } + else { + assert(false && "Expected a tensor input"); + } + auto result_type = getBinOpResultType(lhs_el_type, rhs_el_type); comet_debug() << __LINE__ << " "; comet_vdump(result_type); @@ -789,9 +815,10 @@ namespace ret_lbls_value.push_back(all_lbls_value[n]); } + // std::vector result_dims = getDimSizes(ret_lbls_value); auto affineMapArrayAttr = builder.getAffineMapArrayAttr(affine_maps); - auto res_map = affineMapArrayAttr[affineMapArrayAttr.size() - 1].cast().getValue(); + auto res_map = mlir::cast(affineMapArrayAttr[affineMapArrayAttr.size() - 1]).getValue(); /// get return-type based on affine-maps std::vector result_dims; @@ -800,96 +827,33 @@ namespace { for (size_t i = 0; i < affineMapArrayAttr.size() - 1; i++) { - auto map = affineMapArrayAttr[i].cast().getValue(); + auto map = mlir::cast(affineMapArrayAttr[i]).getValue(); if (auto pos = map.getResultPosition(v)) { - result_dims.push_back((i == 0 ? lhs_labeledtensor : rhs_labeledtensor).getType().cast().getDimSize(*pos)); + mlir::Value operand = i == 0 ? lhs_labeledtensor : rhs_labeledtensor; + if(auto spTensorType = mlir::dyn_cast(operand.getType())) + { + result_dims.push_back(spTensorType.getDimSize(*pos)); + } + else{ + assert(false && "Unexpected Input type"); + } break; } } } - auto ret_tensor_type = mlir::RankedTensorType::get(result_dims, result_type); - SmallVector formats; + SmallVector formats; std::vector exprs{lhs_labeledtensor, rhs_labeledtensor}; std::vector tensors; // TODO(gkestor): URGENT refactor the following code - too much repetition for (auto e : exprs) { - if (isa(e.getDefiningOp())) - { - comet_debug() << " is TensorDeclOp\n"; - /// infer the format - auto lhs_format = dyn_cast(e.getDefiningOp()).getFormat(); - comet_debug() << " lhs_format: " << lhs_format << "\n"; - formats.push_back(lhs_format); - - tensors.push_back(dyn_cast(e.getDefiningOp())); - } - else if (isa(e.getDefiningOp())) - { - comet_debug() << " is TensorDeclOp\n"; - /// infer the format - auto lhs_format = dyn_cast(e.getDefiningOp()).getFormat(); - comet_debug() << " lhs_format: " << lhs_format << "\n"; - formats.push_back(lhs_format); - - tensors.push_back(dyn_cast(e.getDefiningOp())); - } - else if (isa(e.getDefiningOp())) - { - comet_debug() << " is TensorMultOp\n"; - /// infer the format - mlir::ArrayAttr opFormatsArrayAttr = dyn_cast(e.getDefiningOp()).getFormats(); - unsigned int i = opFormatsArrayAttr.size() - 1; - mlir::StringRef lhs_format = opFormatsArrayAttr[i].cast().getValue(); - comet_debug() << __LINE__ << " lhs_format: " << lhs_format << "\n"; - comet_debug() << " lhs_format: " << lhs_format << "\n"; - formats.push_back(lhs_format); - - tensors.push_back(dyn_cast(e.getDefiningOp()).getOperation()->getResult(0)); - } - else if (isa(e.getDefiningOp())) - { - comet_debug() << " is TensorElewsMultOp\n"; - /// infer the format - mlir::ArrayAttr opFormatsArrayAttr = dyn_cast(e.getDefiningOp()).getFormats(); - unsigned int i = opFormatsArrayAttr.size() - 1; - mlir::StringRef lhs_format = opFormatsArrayAttr[i].cast().getValue(); - comet_debug() << __LINE__ << " lhs_format: " << lhs_format << "\n"; - - comet_debug() << " lhs_format: " << lhs_format << "\n"; - formats.push_back(lhs_format); - - tensors.push_back(dyn_cast(e.getDefiningOp()).getOperation()->getResult(0)); - } - else if (isa(e.getDefiningOp())) - { - comet_debug() << " is TensorAddOp\n"; - /// infer the format - mlir::ArrayAttr opFormatsArrayAttr = dyn_cast(e.getDefiningOp()).getFormats(); - unsigned int i = opFormatsArrayAttr.size() - 1; - mlir::StringRef lhs_format = opFormatsArrayAttr[i].cast().getValue(); - comet_debug() << __LINE__ << " lhs_format: " << lhs_format << "\n"; - - comet_debug() << " lhs_format: " << lhs_format << "\n"; - formats.push_back(lhs_format); - - tensors.push_back(dyn_cast(e.getDefiningOp()).getOperation()->getResult(0)); - } - else if (isa(e.getDefiningOp())) + if (isa(e.getDefiningOp())) { - comet_debug() << " is TensorSubstract Op\n"; - /// infer the format - mlir::ArrayAttr opFormatsArrayAttr = dyn_cast(e.getDefiningOp()).getFormats(); - unsigned int i = opFormatsArrayAttr.size() - 1; - mlir::StringRef lhs_format = opFormatsArrayAttr[i].cast().getValue(); - comet_debug() << __LINE__ << " lhs_format: " << lhs_format << "\n"; - - comet_debug() << " lhs_format: " << lhs_format << "\n"; + auto lhs_format = getTensorFormatString(e.getType()); formats.push_back(lhs_format); - - tensors.push_back(dyn_cast(e.getDefiningOp()).getOperation()->getResult(0)); + tensors.push_back(e); } else if (isa(e.getDefiningOp())) { @@ -909,20 +873,7 @@ namespace } /// get the format of transposeOut tensor - if (isa(transposeOut.getDefiningOp())) - { - auto denseFormat = dyn_cast(transposeOut.getDefiningOp()).getFormat(); - formats.push_back(denseFormat); - } - else if (isa(transposeOut.getDefiningOp())) - { - auto sparseFormat = dyn_cast(transposeOut.getDefiningOp()).getFormat(); - formats.push_back(sparseFormat); - } - else - { - llvm::errs() << __FILE__ << ":" << __LINE__ << " ERROR: Can not determine tensor format with transpose op\n"; - } + formats.push_back(getTensorFormatString(transposeOut.getType())); tensors.push_back(transposeOut); } else @@ -932,20 +883,42 @@ namespace } comet_debug() << __LINE__ << " formats.size(): " << formats.size() << "\n"; assert(formats.size() == 2 && " less than 2 input tensors\n"); + mlir::Type ret_tensor_type; + // if (formats[0].compare("CSR") == 0 && formats[1].compare("CSR") == 0) + // { + // formats.push_back("CSR"); + // std::vector format_array = getFormats("CSR", result_dims.size(), builder.getContext()); + // ret_tensor_type = SparseTensorType::get(builder.getContext(), result_type, builder.getIntegerType(defaultSpTensorIndiceBitWidth), result_dims, format_array); + // } + // else if (formats[0].compare("Dense") == 0 && formats[1].compare("Dense") == 0) + // { + // formats.push_back("Dense"); + // ret_tensor_type = mlir::RankedTensorType::get(result_dims, result_type); + // } if (out_format.length() > 0) // non-empty format string provided. { comet_debug() << " Output Format: " << out_format << "\n"; formats.push_back(out_format); + if(out_format.compare("Dense") == 0) + { + ret_tensor_type = mlir::RankedTensorType::get(result_dims, result_type); + } else { + std::vector format_array = getFormats(out_format, result_dims.size(), builder.getContext()); + ret_tensor_type = SparseTensorType::get(builder.getContext(), result_type, builder.getIntegerType(defaultSpTensorIndiceBitWidth), result_dims, format_array); + } } else if (formats[0].compare("CSR") == 0) { if(formats[1].compare("CSR") == 0) { formats.push_back("CSR"); + std::vector format_array = getFormats("CSR", result_dims.size(), builder.getContext()); + ret_tensor_type = SparseTensorType::get(builder.getContext(), result_type, builder.getIntegerType(defaultSpTensorIndiceBitWidth), result_dims, format_array); } else if(formats[1].compare("Dense") == 0) { formats.push_back("Dense"); + ret_tensor_type = mlir::RankedTensorType::get(result_dims, result_type); } } else if (formats[0].compare("Dense") == 0) @@ -953,10 +926,12 @@ namespace if(formats[1].compare("CSR") == 0) // Redundant but shows the intention { formats.push_back("Dense"); + ret_tensor_type = mlir::RankedTensorType::get(result_dims, result_type); } else if(formats[1].compare("Dense") == 0) // Redundant but shows the intention { formats.push_back("Dense"); + ret_tensor_type = mlir::RankedTensorType::get(result_dims, result_type); } } else @@ -964,7 +939,6 @@ namespace llvm::errs() << __FILE__ << ":" << __LINE__ << " ERROR: the format of output tensor could not be determined during generation of binOp\n"; } comet_debug() << " formats.size(): " << formats.size() << "\n"; - auto strAttr = builder.getStrArrayAttr(formats); assert(tensors.size() == 2 && " less than 2 input tensors for ta.mul or ta.elews_mul\n"); @@ -993,14 +967,14 @@ namespace SemiringAttr = builder.getStringAttr("noop_plusxy"); // this is for standard elementwise addition MaskingAttr = builder.getStringAttr("none"); // default for standard elementwise addition return builder.create(location, ret_tensor_type, tensors[0], tensors[1], - labels, affineMapArrayAttr, strAttr, SemiringAttr, + labels, affineMapArrayAttr, SemiringAttr, MaskingAttr); case '-': comet_debug() << "creating TensorSubtractOp\n"; SemiringAttr = builder.getStringAttr("noop_minus"); // this is for standard elementwise subtraction MaskingAttr = builder.getStringAttr("none"); // default for standard elementwise subtraction return builder.create(location, ret_tensor_type, tensors[0], tensors[1], - labels, affineMapArrayAttr, strAttr, SemiringAttr, + labels, affineMapArrayAttr, SemiringAttr, MaskingAttr); case '*': { @@ -1012,7 +986,7 @@ namespace SemiringAttr = builder.getStringAttr("plusxy_times"); // this is for standard matrix multiplication MaskingAttr = builder.getStringAttr("none"); // default for standard matrix multiplication mlir::Value tcop = builder.create(location, ret_tensor_type, tensors[0], tensors[1], - labels, affineMapArrayAttr, strAttr, SemiringAttr, + labels, affineMapArrayAttr, SemiringAttr, MaskingAttr, nullptr); // TODO: masking is an optional operand tcop.getDefiningOp()->setAttr("__alpha__", builder.getF64FloatAttr(1.0)); tcop.getDefiningOp()->setAttr("__beta__", builder.getF64FloatAttr(0.0)); @@ -1030,7 +1004,7 @@ namespace auto SemiringAttr = builder.getStringAttr("noop_times"); /// this is for standard element-wise multiplication MaskingAttr = builder.getStringAttr("none"); /// default for standard element-wise multiplication mlir::Value tcop = builder.create(location, ret_tensor_type, tensors[0], tensors[1], labels, - affineMapArrayAttr, strAttr, SemiringAttr, + affineMapArrayAttr, SemiringAttr, MaskingAttr); comet_vdump(tcop); @@ -1145,9 +1119,10 @@ namespace auto *rhsLT = llvm::cast(expr); auto name = rhsLT->getTensorName(); mlir::Value tensorValue = symbolTable.lookup(name); + mlir::ShapedType shapedT = mlir::cast(tensorValue.getType()); comet_debug() << " generate ta.sum op\n"; /// TODO(gkestor): look at reduceOp in linalg - sumVal = builder.create(location, builder.getF64Type(), tensorValue); + sumVal = builder.create(location, shapedT.getElementType(), tensorValue); } /// Case 2: SUM(A[i,j]*B[j,k]) @@ -1242,25 +1217,7 @@ namespace return nullptr; comet_debug() << " get lhs\n"; - mlir::Value lhs_decl; - std::string out_format; - if (isa(lhs.getDefiningOp())) - { - lhs_decl = lhs; - out_format = "Dense"; - } - else if (isa(lhs.getDefiningOp())) - { - lhs_decl = lhs; - auto lhs_decl_op = cast(lhs_decl.getDefiningOp()); - std::string formatStr(lhs_decl_op.getFormatAttr().getValue()); - out_format = formatStr; - } - else - { - emitError(loc(tensor_op.loc()), - "error: Tensor Decl Not Found! '"); - } + std::string out_format = getTensorFormatString(lhs.getType()); auto out_lbls_vec = cast(*tensor_op.getLHS()) @@ -1281,7 +1238,7 @@ namespace auto tens_beta = tensor_op.getBeta(); - auto ret_op = builder.create(loc(tensor_op.loc()), rhs, lhs_decl); + auto ret_op = builder.create(loc(tensor_op.loc()), rhs, lhs); ret_op.getOperation()->setAttr("__beta__", builder.getF64FloatAttr(tens_beta)); return rhs; @@ -1456,9 +1413,24 @@ namespace if (isDense(formats_str, ", ") == false) { /// BoolAttr is false because there is explicit sparse densor declaration. - /// SparseTensorDeclOp is not for temporaries in compound expressions + /// SparseTensorDeclOp is not for temporaries in compound expression + std::vector format = mlir::tensorAlgebra::getFormats(tensor_format, dims_sizes.size(), builder.getContext()); + mlir::Type element_type; + switch (vartype.elt_ty) + { + case VarType::TY_FLOAT: + element_type = builder.getF32Type(); + break; + case VarType::TY_DOUBLE: + element_type = builder.getF64Type(); + break; + case VarType::TY_INT: + element_type = builder.getIntegerType(64); + break; + } + auto sp_tensor_type = SparseTensorType::get(builder.getContext(), element_type, builder.getIntegerType(defaultSpTensorIndiceBitWidth), dims_sizes, format); value = builder.create(loc(tensordecl.loc()), - tensor_type, labels, tensor_format, false); + sp_tensor_type, labels, false); comet_debug() << "MLIRGen SparseTensorDeclaration creation\n"; comet_vdump(value); @@ -1475,7 +1447,7 @@ namespace else { value = builder.create(loc(tensordecl.loc()), - tensor_type, labels, tensor_format); + tensor_type, labels); comet_debug() << "MLIRGen DenseTensorDeclaration creation\n"; comet_vdump(value); } @@ -1535,72 +1507,12 @@ namespace auto affineMapArrayAttr = builder.getAffineMapArrayAttr(affine_maps); - SmallVector formats; - - /// Firstly, Look at the rhs - if (isa(rhs_tensor.getDefiningOp())) - { - comet_debug() << " is TensorDeclOp\n"; - /// infer the format - auto rhs_format = dyn_cast(rhs_tensor.getDefiningOp()).getFormat(); - comet_debug() << " rhs_format: " << rhs_format << "\n"; - formats.push_back(rhs_format); - } - else if (isa(rhs_tensor.getDefiningOp())) - { - comet_debug() << " is TensorDeclOp\n"; - /// infer the format - auto rhs_format = dyn_cast(rhs_tensor.getDefiningOp()).getFormat(); - comet_debug() << " rhs_format: " << rhs_format << "\n"; - formats.push_back(rhs_format); - } - else - { - comet_debug() << " not TensorDeclOp\n"; - } - - /// Secondly, Look at the lhs - std::vector exprs{&lhsLT}; - std::vector lhs_lbls_value; - - for (auto e : exprs) - { - auto lhsLT_tensor_name = e->getTensorName(); - mlir::Value lhsLT_op; - if ((lhsLT_op = symbolTable.lookup(lhsLT_tensor_name)) != NULL) - { - if (isa(lhsLT_op.getDefiningOp())) - { - comet_debug() << " is TensorDeclOp\n"; - /// infer the format - auto lhs_format = dyn_cast(lhsLT_op.getDefiningOp()).getFormat(); - comet_debug() << " lhs_format: " << lhs_format << "\n"; - formats.push_back(lhs_format); - } - else if (isa(lhsLT_op.getDefiningOp())) - { - comet_debug() << " is TensorDeclOp\n"; - /// infer the format - auto lhs_format = dyn_cast(lhsLT_op.getDefiningOp()).getFormat(); - comet_debug() << " lhs_format: " << lhs_format << "\n"; - formats.push_back(lhs_format); - } - else - { - comet_debug() << " not TensorDeclOp\n"; - } - } - } - - comet_debug() << " formats.size(): " << formats.size() << "\n"; - auto strAttr = builder.getStrArrayAttr(formats); - auto lhs_tensor = symbolTable.lookup(lhsLT.getTensorName()); comet_debug() << " create TransposeOp\n"; mlir::Value t = builder.create(loc(transpose.loc()), lhs_tensor.getType(), - rhs_tensor, all_lbls_value, affineMapArrayAttr, strAttr); - builder.create(loc(transpose.loc()), t.getDefiningOp()->getResult(0), lhs_tensor); + rhs_tensor, all_lbls_value, affineMapArrayAttr); + builder.create(loc(transpose.loc()), t, lhs_tensor); comet_vdump(t); return t; @@ -1651,30 +1563,6 @@ namespace auto affineMapArrayAttr = builder.getAffineMapArrayAttr(affine_maps); - SmallVector formats; - - /// Firstly, Look at the rhs - if (isa(rhs_tensor.getDefiningOp())) - { - comet_debug() << " is TensorDeclOp\n"; - /// infer the format - auto rhs_format = dyn_cast(rhs_tensor.getDefiningOp()).getFormat(); - comet_debug() << " rhs_format: " << rhs_format << "\n"; - formats.push_back(rhs_format); - } - else if (isa(rhs_tensor.getDefiningOp())) - { - comet_debug() << " is TensorDeclOp\n"; - /// infer the format - auto rhs_format = dyn_cast(rhs_tensor.getDefiningOp()).getFormat(); - comet_debug() << " rhs_format: " << rhs_format << "\n"; - formats.push_back(rhs_format); - } - else - { - comet_debug() << " not TensorDeclOp\n"; - } - /// Secondly, Look at the lhs /// Collect labels values std::vector lhs_labels_val; @@ -1724,17 +1612,17 @@ namespace } /// get return-type based on affine-maps - auto res_map = affineMapArrayAttr[1].cast().getValue(); + auto res_map = cast(affineMapArrayAttr[1]).getValue(); std::vector indices; std::vector shape; for (auto v : res_map.getResults()) { - auto map = affineMapArrayAttr[0].cast().getValue(); + auto map = cast(affineMapArrayAttr[0]).getValue(); if (auto pos = map.getResultPosition(v)) { - if (isa(rhs_tensor.getType()) && !rhs_tensor.getType().cast().isDynamicDim(*pos)) + if (auto tensorT = dyn_cast(rhs_tensor.getType()); tensorT && !tensorT.isDynamicDim(*pos)) { - shape.push_back(rhs_tensor.getType().cast().getDimSize(*pos)); + shape.push_back(tensorT.getDimSize(*pos)); } else { @@ -1744,45 +1632,37 @@ namespace } } + mlir::Type return_type = getType(shape); + /// Create Tensor Declarations Ops and populate formats (for lhs) mlir::Value lhs_tensor; - if (isa(rhs_tensor.getDefiningOp())) + if (auto tensorT = dyn_cast(rhs_tensor.getType())) { - /// for DenseTensorDeclOp create - mlir::StringRef format_strref = dyn_cast(rhs_tensor.getDefiningOp()).getFormat(); - mlir::StringAttr formatAttr = builder.getStringAttr(format_strref); - lhs_tensor = builder.create(loc(transpose.loc()), mlir::RankedTensorType::get(shape, builder.getF64Type()), indices, formatAttr); - - /// populate formats - /// assumes lhs and rhs formats are same - auto lhs_format = dyn_cast(rhs_tensor.getDefiningOp()).getFormat(); - formats.push_back(lhs_format); + auto declOp = builder.create(loc(transpose.loc()), mlir::RankedTensorType::get(shape, tensorT.getElementType()), indices); + + lhs_tensor = declOp; } - else if (isa(rhs_tensor.getDefiningOp())) + else if (auto SparseTensorT = dyn_cast(rhs_tensor.getType())) { - /// for SparseTensorDeclOp create - mlir::StringRef format_strref = dyn_cast(rhs_tensor.getDefiningOp()).getFormat(); - mlir::StringAttr formatAttr = builder.getStringAttr(format_strref); - /// no lhs_LabeledTensor has been created. The output tensor of tranpose doesn't have explicit declaration, + ArrayRef format = SparseTensorT.getFormat(); + mlir::ShapedType shapedT = mlir::cast(rhs_tensor.getType()); + mlir::Type element_type = shapedT.getElementType(); + return_type = SparseTensorType::get(builder.getContext(), element_type, builder.getIntegerType(defaultSpTensorIndiceBitWidth), shape, format); + auto sp_tensor_type = SparseTensorType::get(builder.getContext(), element_type, builder.getIntegerType(defaultSpTensorIndiceBitWidth), shape, format); + /// BoolAttr is true to speficy SparseTensorDeclOp is for temporaries - lhs_tensor = builder.create(loc(transpose.loc()), mlir::RankedTensorType::get(shape, builder.getF64Type()), indices, formatAttr, builder.getBoolAttr(true)); + lhs_tensor = builder.create(loc(transpose.loc()), sp_tensor_type, indices, builder.getBoolAttr(true)); comet_debug() << "MLIRGen SparseTensorDeclaration creation\n"; comet_vdump(lhs_tensor); - - /// populate formats - /// assumes lhs and rhs formats are same - auto lhs_format = dyn_cast(rhs_tensor.getDefiningOp()).getFormat(); - formats.push_back(lhs_format); } - comet_debug() << " formats.size(): " << formats.size() << "\n"; - auto strAttr = builder.getStrArrayAttr(formats); comet_debug() << " create TransposeOp\n"; - mlir::Value t = builder.create(loc(transpose.loc()), mlir::RankedTensorType::get(shape, builder.getF64Type()), - rhs_tensor, all_labels_val, affineMapArrayAttr, strAttr); - builder.create(loc(transpose.loc()), t.getDefiningOp()->getResult(0), lhs_tensor); + mlir::ShapedType shapedT = mlir::cast(rhs_tensor.getType()); + mlir::Value t = builder.create(loc(transpose.loc()), mlir::RankedTensorType::get(shape, shapedT.getElementType()), + rhs_tensor, all_labels_val, affineMapArrayAttr); + builder.create(loc(transpose.loc()), t, lhs_tensor); comet_vdump(t); return t; @@ -2204,7 +2084,6 @@ namespace std::vector tensors; auto binop = tensor_op.getOp(); - SmallVector formats; std::vector exprs{rhsLT, lhsLT}; for (auto e : exprs) { @@ -2212,23 +2091,14 @@ namespace mlir::Value lhsLT_op; if ((lhsLT_op = symbolTable.lookup(lhsLT_tensor_name)) != NULL) { - if (isa(lhsLT_op.getDefiningOp())) + if (isa(lhsLT_op.getType())) { /// infer the format - auto lhs_format = dyn_cast(lhsLT_op.getDefiningOp()).getFormat(); - formats.push_back(lhs_format); - tensors.push_back(dyn_cast(lhsLT_op.getDefiningOp())); - } - else if (isa(lhsLT_op.getDefiningOp())) - { - /// infer the format - auto lhs_format = dyn_cast(lhsLT_op.getDefiningOp()).getFormat(); - formats.push_back(lhs_format); - tensors.push_back(dyn_cast(lhsLT_op.getDefiningOp())); + tensors.push_back(lhsLT_op); } else { - comet_debug() << " not TensorDeclOp\n"; + return mlir::failure(); /// this should not happen, we expect a tensor type here. } } } @@ -2363,12 +2233,59 @@ namespace all_lbls_value.push_back(symbolTable.lookup(n)); } + /// determine the map. The order is rhs1's dimensions, rhs2's, then lhs'. + /// TODO: the order can be determined by autotuning. std::map expr_map; unsigned dim = 0; - for (const auto &lbl : all_lbls) - { - expr_map[lbl] = getAffineDimExpr(dim++, builder.getContext()); + if (debugOptions.find("debug-ta-labels-alphabet-order") != debugOptions.end()) + {/// Use alphabet order + for (const auto &lbl : all_lbls) + { + expr_map[lbl] = getAffineDimExpr(dim++, builder.getContext()); + } } + else + {/// Use order of rhs1 and rhs2. + std::unordered_set labels_set; + llvm::SmallVector labels_ordered; + for (const auto &label : rhs1_lbls) { + if (labels_set.find(label) == labels_set.end()) { + /// A new label + labels_set.insert(label); + labels_ordered.push_back(label); + } + } + for (const auto &label : rhs2_lbls) { + if (labels_set.find(label) == labels_set.end()) { + /// A new label + labels_set.insert(label); + labels_ordered.push_back(label); + } + } + for (const auto &label : lhs_lbls) { + if (labels_set.find(label) == labels_set.end()) { + /// A new label + labels_set.insert(label); + labels_ordered.push_back(label); + } + } + + for (const auto &label : labels_ordered) { + expr_map[label] = getAffineDimExpr(dim++, builder.getContext()); + } + } + +// std::map expr_map; +// unsigned dim = 0; +// for (const auto &lbl : all_lbls) +// { +// expr_map[lbl] = getAffineDimExpr(dim++, builder.getContext()); +// {/// test +// comet_debug() << lbl << "\n"; +// comet_vdump(expr_map[lbl]); +// } +// } + std::vector rhs1_exprs; std::vector rhs2_exprs; std::vector lhs_exprs; @@ -2396,7 +2313,6 @@ namespace auto affineMapArrayAttr = builder.getAffineMapArrayAttr(affine_maps); - SmallVector formats; std::vector exprs{rhs1LT, rhs2LT, lhsLT}; std::vector tensors; for (auto e : exprs) @@ -2405,27 +2321,9 @@ namespace mlir::Value lhsLT_op; if ((lhsLT_op = symbolTable.lookup(lhsLT_tensor_name)) != NULL) { - if (isa(lhsLT_op.getDefiningOp())) + if (isa(lhsLT_op.getDefiningOp())) { - comet_debug() << " is TensorDeclOp\n"; - - /// infer the format - auto lhs_format = dyn_cast(lhsLT_op.getDefiningOp()).getFormat(); - comet_debug() << " lhs_format: " << lhs_format << "\n"; - formats.push_back(lhs_format); - - tensors.push_back(dyn_cast(lhsLT_op.getDefiningOp())); - } - else if (isa(lhsLT_op.getDefiningOp())) - { - comet_debug() << " is TensorDeclOp\n"; - - /// infer the format - auto lhs_format = dyn_cast(lhsLT_op.getDefiningOp()).getFormat(); - comet_debug() << " lhs_format: " << lhs_format << "\n"; - formats.push_back(lhs_format); - - tensors.push_back(dyn_cast(lhsLT_op.getDefiningOp())); + tensors.push_back(lhsLT_op); } else { @@ -2433,8 +2331,6 @@ namespace } } } - comet_debug() << " formats.size(): " << formats.size() << "\n"; - auto strAttr = builder.getStrArrayAttr(formats); assert(tensors.size() == 3 && "Not 3 tensors for ta.tc or ta.elews_mul\n"); @@ -2455,7 +2351,7 @@ namespace tensors[0], tensors[1], all_lbls_value, affineMapArrayAttr, - strAttr, SemiringAttr, + SemiringAttr, MaskingAttr); comet_vdump(op); @@ -2470,7 +2366,7 @@ namespace tensors[0], tensors[1], all_lbls_value, affineMapArrayAttr, - strAttr, SemiringAttr, + SemiringAttr, MaskingAttr); comet_vdump(op); /// source is 1st parameter, dest is the second @@ -2483,7 +2379,7 @@ namespace tensors[0], tensors[1], all_lbls_value, affineMapArrayAttr, - strAttr, SemiringAttr, + SemiringAttr, MaskingAttr, maskVal); op.getOperation()->setAttr("__alpha__", builder.getF64FloatAttr(1.0)); op.getOperation()->setAttr("__beta__", builder.getF64FloatAttr(tens_beta)); @@ -2494,7 +2390,7 @@ namespace } else if (binop == tok_elews || binop == tok_monoid) { - auto op = builder.create(loc(tensor_op.loc()), tensors[2].getType(), tensors[0], tensors[1], all_lbls_value, affineMapArrayAttr, strAttr, SemiringAttr, MaskingAttr); + auto op = builder.create(loc(tensor_op.loc()), tensors[2].getType(), tensors[0], tensors[1], all_lbls_value, affineMapArrayAttr, SemiringAttr, MaskingAttr); op.getOperation()->setAttr("__alpha__", builder.getF64FloatAttr(1.0)); op.getOperation()->setAttr("__beta__", builder.getF64FloatAttr(tens_beta)); @@ -2514,8 +2410,7 @@ namespace StringRef tensor_name, double value) { mlir::Value tensorValue = symbolTable.lookup(tensor_name); - auto tensorType = tensorValue.getDefiningOp()->getOpResult(0).getType(); - auto tensorElType = tensorType.cast().getElementType(); + auto tensorElType = cast(tensorValue.getType()).getElementType(); mlir::FloatAttr valueAttr; if (tensorElType.isF64()) @@ -2544,23 +2439,24 @@ namespace comet_debug() << " in mlirGenTensorFillRandom\n"; mlir::Value tensorValue = symbolTable.lookup(tensor_name); - auto lhs_labeledtensor = tensorValue.getDefiningOp()->getOpResult(0); + auto lhs_labeledtensor = tensorValue; comet_debug() << "\n"; comet_vdump(lhs_labeledtensor); std::vector result_dims; std::vector lhs_lbls_value; - if (isa(lhs_labeledtensor.getDefiningOp())) + if (mlir::TensorType tensorT = dyn_cast(lhs_labeledtensor.getType())) { - mlir::TensorType tensor = cast(lhs_labeledtensor.getType()); - result_dims = tensor.getShape(); + // mlir::TensorType tensor = cast(lhs_labeledtensor.getType()); + result_dims = tensorT.getShape(); } else { - if (isa(lhs_labeledtensor.getDefiningOp())) + if (isa(lhs_labeledtensor.getType())) llvm::errs() << __FILE__ << ":" << __LINE__ << " ERROR: random initialization is currently not supported for sparse tensors.\n"; llvm::errs() << __FILE__ << ":" << __LINE__ << " ERROR: Not supported format encountered during random initialization of tensor.\n"; + return mlir::failure(); } comet_debug() << " dims size: " << result_dims.size() << "\n"; @@ -2611,15 +2507,20 @@ namespace } }; -} /// namespace +} /// anonymous namespace namespace tensorAlgebra { + std::unordered_set debugOptions; /// The public API for codegen. mlir::OwningOpRef mlirGen(mlir::MLIRContext &context, - ModuleAST &moduleAST) + ModuleAST &moduleAST, bool useI64) { + if(!useI64) + { + defaultSpTensorIndiceBitWidth = 32; + } return MLIRGenImpl(context).mlirGen(moduleAST); } diff --git a/frontends/numpy-scipy/README.md b/frontends/numpy-scipy/README.md index 03cecca3..f69429b0 100644 --- a/frontends/numpy-scipy/README.md +++ b/frontends/numpy-scipy/README.md @@ -43,10 +43,13 @@ More information about the COMET compiler can be found at: ``` 2) **Testing:** Run the integration tests to make sure the installation was successfull - + * If not already installed, install pytest + ``` + pip install pytest + ``` + ``` - cd integration_tests - python3 numpy_integration.py -v + pytest ``` # How to use cometpy in a program: diff --git a/frontends/numpy-scipy/cometpy/MLIRGen/builders.py b/frontends/numpy-scipy/cometpy/MLIRGen/builders.py deleted file mode 100644 index 29c108c4..00000000 --- a/frontends/numpy-scipy/cometpy/MLIRGen/builders.py +++ /dev/null @@ -1,633 +0,0 @@ -import jinja2 -import functools -import itertools - -from typing import Dict, List, Tuple, Sequence, Union -from collections import OrderedDict -from ast import operator -from cometpy.MLIRGen import types_mlir -from cometpy.MLIRGen.types import * - - -class Dialect: - def __init__(self, name): - self.name = name - - def __repr__(self): - return f"{self.name} dialect" - - -class MLIRFunctionBuilder: - _ops = {} - - default_indentation_size = 4 - indentation_delta_size = 2 - module_wrapper_text = jinja2.Template( - "{{ aliases }}module {\n {{ body }}\nfunc.func private @quick_sort(memref<*xindex>, index)\n}\n", - undefined=jinja2.StrictUndefined, - ) - function_wrapper_text = jinja2.Template( - ("" * default_indentation_size) - + "func.func {% if private_func %}private {% endif %}@{{func_name}}({{signature}}) -> {{return_type}} {" - + "\n" - + "{{statements}}" - + "\n" - + (" " * default_indentation_size) - + "}", - undefined=jinja2.StrictUndefined, - ) - - def __init__( - self, - func_name: str, - input_types, - return_types: Sequence[Union[str, types_mlir.Type]], - aliases: types_mlir.AliasMap = None, - ) -> None: - # TODO mlir functions can return zero or more results https://mlir.llvm.org/docs/LangRef/#operations - # handle the cases where the number of return types is not 1 - if aliases is None: - - aliases = types_mlir.AliasMap() - self.aliases = aliases - - # Build input vars and ensure all types are proper Types - # inputs = [] - # for i, it in enumerate(input_types): - # it = Type.find(it, aliases) - # # Create initialized MLIRVar - # iv = MLIRVar(f"arg{i}", it) - # iv._initialized = True - # inputs.append(iv) - return_types = [types_mlir.Type.find(rt, aliases) for rt in return_types] - - self.func_name = func_name - self.inputs = input_types - self.return_types = return_types - - self.var_name_counter = itertools.count() - self.function_body_statements: List[str] = [] - self.temporary_statement_lists: List[List[str]] = [] - - # function_name -> (function_mlir_definition, input_mlir_types, return_mlir_type) - self.needed_function_table: Dict[ - str, Tuple[str, List[str], str] - ] = OrderedDict() - - self.indentation_level = 1 - self._initialize_ops() - - def _initialize_ops(self): - for dialect, ops in self._ops.items(): - if dialect is None: - attach_point = self - else: - attach_point = getattr(self, dialect, None) - if attach_point is None: - attach_point = Dialect(dialect) - setattr(self, dialect, attach_point) - - for opclass in ops.values(): - - def op(opclass, *args, **kwargs): - ret_val, mlir = opclass.call(self, *args, **kwargs) - self.add_statement(mlir) - return ret_val - - func = functools.partial(op, opclass) - setattr(attach_point, opclass.name, func) - - ####################################### - # MLIR Generation/Compilation Methods # - ####################################### - - def get_mlir_module(self, make_private=False): - """Get the MLIR text for this function wrapped in a MLIR module with - declarations of external helper functions.""" - aliases = "\n".join( - f"#{name} = {typ.to_pretty_string()}" for name, typ in self.aliases.items() - ) - body = self.get_mlir(make_private=make_private) - return self.module_wrapper_text.render(aliases=aliases, body=body) - - def get_mlir(self, make_private=True, include_func_defs=True) -> str: - if include_func_defs: - needed_function_definitions = "\n ".join( - func_def for func_def, _, _ in self.needed_function_table.values() - ) - else: - needed_function_definitions = "" - - if len(self.temporary_statement_lists) > 0: - raise RuntimeError( - "Cannot get MLIR code while using temporary statement storage." - ) - joined_statements = "\n".join(self.function_body_statements) - - return_type = ", ".join(str(rt) for rt in self.return_types) - if len(self.return_types) != 1: - return_type = f"({return_type})" - signature = ", ".join(f"{ var[0].replace('%','%arg_')}: {var[1].replace('tensor', 'memref')}" if 'tensor' in var[1] else f"{ var[0].replace('%', '%arg_') }: memref<1xf64>" if var[1] =="f64" else f"{ var[0]}: {var[1]}" for var in self.inputs) - - return needed_function_definitions + self.function_wrapper_text.render( - private_func=make_private, - func_name=self.func_name, - signature=signature, - return_type=return_type, - statements=joined_statements, - ) - - - def compile(self): - # if engine is None: - # engine = self.engine - - # Force recompilation if name is already registered - # if self.func_name in engine.name_to_callable: - # del engine.name_to_callable[self.func_name] - - mlir = self.get_mlir_module() - - # engine.add(mlir, passes) - # func = engine[self.func_name] - # func.builder = self - - return mlir - - - ################################ - # MLIR Building Method Helpers # - ################################ - - @property - def current_statement_list(self) -> List[str]: - return ( - self.temporary_statement_lists[-1] - if len(self.temporary_statement_lists) > 0 - else self.function_body_statements - ) - - def add_statement(self, statement: str) -> None: - """In an ideal world, no human would ever call this method.""" - if not statement: - return - - for line in map(str.strip, statement.split("\n")): - self.current_statement_list.append( - " " * self.default_indentation_size - + " " * self.indentation_delta_size * self.indentation_level - + line - ) - - -class TensorSumBuilder: - indentation_size = 4 - tensor_sum_wrapper_text = jinja2.Template( - ("" * indentation_size) - + '%t{{lhs}} = "ta.reduce"{{operators}}' - +' : ({{inputtype}})' - +"-> {{output_types}}" - + "\n" , - undefined=jinja2.StrictUndefined, - ) - - def __init__(self, data): # lhs, operators, tensors_shapes, label_map): - self.lhs = data["out_id"] - self.input_type = "tensor<{}xf64>".format("x".join(str(v) for v in data["shapes"][0])) - - self.operators = "({})".format(",".join("%t"+str(v) for v in data["operands"])) - - def build_op(self): - output_type = "f64" - - return self.tensor_sum_wrapper_text.render( - lhs = self.lhs, - operators = self.operators, - inputtype = self.input_type, - output_types = output_type - ) - -class SetOp_Builder: - indentation_size = 4 - beta_val = 0.0 - - set_op_wrapper_text = jinja2.Template( - ("" * indentation_size) - + '"ta.set_op"(%t{{input}},%t{{dest}}) {__beta__ = {{beta}} : f64} : ({{outputtype}}, {{outputtype}}) -> ()\n', - undefined=jinja2.StrictUndefined, - ) - - def __init__(self, data):# in_tensor, target, tensors_shapes, label_map, beta) : - self.target = data["lhs"] - self.in_tensor = data["rhs"] - self.tensors_shapes = data["shapes"] - self.beta = "{:e}".format(data["beta"]) - - - def build_op(self): - output_type = "tensor<{}xf64>".format("x".join(str(v) for v in self.tensors_shapes[-1])) - # input_type = [] - # for t in self.tensors_shapes[:-1]: - # input_type.append("tensor<{}xf64>".format("x".join(str(v) for v in t))) - - return self.set_op_wrapper_text.render( - dest = self.target, - input = self.in_tensor, - beta = self.beta, - outputtype= output_type, - - ) - -class ScalarOp_Builder: - indentation_size = 4 - - scalar_op_wrapper_text = jinja2.Template ( - ("" * indentation_size) - +'%t{{dest}} = "ta.scalar"({{operators}})' - +' <{op = "{{op}}"}> ' - +' : ({{inputtype}})' - +' -> ({{outputtype}})' - + "\n", - undefined=jinja2.StrictUndefined, - ) - - def __init__(self, data): - - self.dest = data["out_id"] - self.operators = "{}".format(",".join("%t"+str(v) for v in data["operands"])) - self.tensors_shapes =[] - for l in data["shapes"]: - if isinstance(l, int): - self.tensors_shapes.append('f64') - else: - self.tensors_shapes.append('tensor<1xf64>') - - self.op = data["op"] - - def build_op(self): - input_type = [] - - return self.scalar_op_wrapper_text.render( - dest = self.dest, - operators = self.operators, - op = self.op, - inputtype = ",".join(self.tensors_shapes[:2]), - outputtype = self.tensors_shapes[-1], - ) - -class ArithOp_Builder: - formats_str = ['Dense', 'CSR', 'COO', 'CSC'] - indentation_size = 4 - beta_val = 0.0 - - - - tc_decl_wrapper_text = jinja2.Template( - ("" * indentation_size) - + '%t{{dest}} = "ta.mul"({{operators}})' - + '{MaskType = "{{mask_type}}", ' - + '__alpha__ = 1.000000e+00 : f64, ' - +"__beta__ = {{beta}} : f64," - + 'formats = [{{formats}}],' - +'indexing_maps = {{indexing_maps}}, ' - +'operand_segment_sizes = array, ' #[TODO] operand_segment_sizes should not be static - +'semiring = "{{semiring}}"} : ' - +"({{inputtype}})" - +"-> {{outputtype}}" - + "\n" , - # + '"ta.set_op"(%t{{dest}},%t{{dest}}) {__beta__ = {{beta}} : f64} : ({{outputtype}}, {{outputtype}}) -> ()\n', - undefined=jinja2.StrictUndefined, - ) - - - - tensor_add_wrapper_text = jinja2.Template( - ("" * indentation_size) - +'%t{{dest}} = "ta.add"({{operators}})' - +' {' - +' Masktype = "none",' - +' formats = [{{formats}}],' - +' indexing_maps = {{indexing_maps}},' - +' semiring = "noop_plusxy"' - +' }' - +' : ({{inputtype}})' - +"-> {{outputtype}}" - + "\n", - # + '"ta.set_op"(%t{{dest}},%t{{dest}}): ({{outputtype}}, {{outputtype}}) -> ()\n', - undefined=jinja2.StrictUndefined, - ) - - tensor_sub_wrapper_text = jinja2.Template( - ("" * indentation_size) - +'%t{{dest}} = "ta.subtract"({{operators}})' - +' {' - +' Masktype = "none",' - +' formats = [{{formats}}],' - +' indexing_maps = {{indexing_maps}},' - +' semiring = "noop_minus"' - +' }' - +' : ({{inputtype}})' - +"-> {{outputtype}}" - + "\n", - # + '"ta.set_op"(%t{{dest}},%t{{dest}}): ({{outputtype}}, {{outputtype}}) -> ()\n', - undefined=jinja2.StrictUndefined, - ) - - elewisemult_wrapper_text = jinja2.Template( - ("" * indentation_size) - +'%t{{dest}} = "ta.elews_mul"({{operators}})' - +' {__alpha__ = 1.000000e+00 : f64, ' - +"__beta__ = {{beta}}: f64," - + 'formats = [{{formats}}],' - +'indexing_maps = {{indexing_maps}}, semiring = "{{semiring}}"} : ' - +"({{inputtype}})" - +"-> {{outputtype}}" - + "\n" , - # + '"ta.set_op"(%t{{dest}},%t{{dest}}): ({{outputtype}}, {{outputtype}}) -> ()\n', - undefined=jinja2.StrictUndefined, - ) - - tranpose_wrapper_text = jinja2.Template( - ("" * indentation_size) - + '%t{{dest}} = "ta.transpose"({{operators}})' - + '{__alpha__ = 1.000000e+00 : f64, ' - +"__beta__ = 0.000000e+00 : f64," - + 'formats = [{{formats}}],' - +'indexing_maps = {{indexing_maps}},semiring = "plusxy_times"} : ' - +"({{inputtype}})" - +"-> {{outputtype}}" - + "\n" , - # + '"ta.set_op"(%temp_{{dest}},%t{{dest}}): ({{outputtype}}, {{outputtype}}) -> ()\n', - undefined=jinja2.StrictUndefined, - ) - - def __init__(self, data): - - self.mask = None - self.mask_type = "None" - self.mask_shape = None - self.semiring = None - - - self.dest = data["out_id"] - self.operators = "{}".format(",".join("%t"+str(v) for v in data["operands"])+","+",".join("%i"+str(vv) for v in data["op_ilabels"] for vv in v)) - self.tensors_shapes =[] - self.op_ilabels = data["op_ilabels"] - for l in data["shapes"]: - self.tensors_shapes.append([ str(lbl) for lbl in l ] ) - # self.tensors_shapes = [label_map[lbl][0] for lbl in tensors_shapes] - self.opr_type = data["op_type"] - # self.op = op - self.formats = data["formats"] - if "mask" in data: - self.mask = data["mask"][0] - self.mask_type = data["mask"][1] - if data["mask"][2] != None: - self.mask_shape = [ str(lbl) for lbl in data["mask"][2] ] - self.operators+=",%t"+str(self.mask) - if "semiring" in data: - self.semiring = data["semiring"] - self.beta = "{:e}".format(data["beta"]) - - def build_op(self): - input_type = [] - for t in self.tensors_shapes[:-1]: - input_type.append("tensor<{}xf64>".format("x".join(str(v) for v in t))) - for t in self.tensors_shapes: - for v in t: - input_type.append("!ta.indexlabel") - input_type = ",".join(input_type) - if self.mask_shape != None: - input_type += ",tensor<{}xf64>".format("x".join(str(v) for v in self.mask_shape)) - # beta_val = ArithOp_Builder.get_beta_val(self.op) - - iMap = {} - vMap = {} - indexing_map = [] - i = 0 - temp = [] - for k, l in enumerate(self.op_ilabels[0]): - iMap[l] = i - vMap[l] = self.tensors_shapes[0][k] - temp.append(i) - i+=1 - - indexing_map.append(temp) - if len(self.op_ilabels) > 2: - temp = [] - for k, l in enumerate(self.op_ilabels[1]): - if l not in iMap: - iMap[l] = i - vMap[l] = self.tensors_shapes[1][k] - temp.append(i) - i+=1 - else: - temp.append(iMap[l]) - indexing_map.append(temp) - temp = [] - - for l in self.op_ilabels[-1]: - temp.append(iMap[l]) - indexing_map.append(temp) - indexing_maps = [] - - output_type = "tensor<{}xf64>".format("x".join(str(vMap[v]) for v in self.op_ilabels[-1])) - - for imap in indexing_map: - indexing_maps.append("affine_map<({})->({})>".format(",".join(["d"+str(l) for l in range(i)]) , ",".join(["d"+str(l) for l in imap]))) - - indexing_maps = str(indexing_maps).replace("'","") - - # Tensor contraction - if self.opr_type == 'c': - semiring = "plusxy_times" - if self.semiring != None: - s1, s2 = self.semiring.split(",") - if s1 == "+": - semiring = "plusxy_" - elif s1 == "any": - semiring = "any_" - elif s1 == "min": - semiring = "minxy_" - if s2 == "*": - semiring += "times" - elif s2 == "pair": - semiring += "pairxy" - elif s2 == "first": - semiring += "first" - elif s2 == "+": - semiring += "plusxy" - elif s2 == "second": - semiring += "second" - - - return self.tc_decl_wrapper_text.render( - dest = self.dest, - operators = self.operators, - indexing_maps = indexing_maps, - inputtype = input_type, - outputtype = output_type, - # beta = self.beta_val, - formats = '"{}", "{}", "{}"'.format(*[self.formats_str[x] for x in self.formats]), - lhs_dims = sum([len(t) for t in self.tensors_shapes ]), - semiring = semiring, - mask=self.mask, - mask_type = self.mask_type, - num_masks = 0 if self.mask == None else 1, - beta = self.beta, - ) - # Add - elif(self.opr_type == '+'): - return self.tensor_add_wrapper_text.render( - dest = self.dest, - operators = self.operators, - inputtype = input_type, - outputtype = output_type, - formats = '"{}", "{}", "{}"'.format(*[self.formats_str[x] for x in self.formats]), - indexing_maps = indexing_maps - ) - # Subtract - elif(self.opr_type == '-'): - return self.tensor_sub_wrapper_text.render( - dest = self.dest, - operators = self.operators, - inputtype = input_type, - outputtype = output_type, - formats = '"{}", "{}", "{}"'.format(*[self.formats_str[x] for x in self.formats]), - indexing_maps = indexing_maps - ) - # Elementwise mult - elif(self.opr_type == '*'): - semiring = "noop_times" - if self.semiring != None: - if self.semiring == "min": - semiring = "noop_minxy" - elif self.semiring == "-": - semiring = "noop_minus" - elif self.semiring == "+": - semiring = "noop_plusxy" - elif self.semiring == "*": - semiring = "noop_times" - return self.elewisemult_wrapper_text.render( - dest = self.dest, - operators = self.operators, - indexing_maps = indexing_maps, - inputtype = input_type, - outputtype = output_type, - formats = '"{}", "{}", "{}"'.format(*[self.formats_str[x] for x in self.formats]), - semiring = semiring, - beta = self.beta_val - ) - # Transpose - elif(self.opr_type == "t"): - return self.tranpose_wrapper_text.render( - dest = self.dest, - operators = self.operators, - indexing_maps = indexing_maps, - inputtype = input_type, - outputtype = output_type, - # beta = self.beta_val, - formats = '"{}", "{}"'.format(*[self.formats_str[x] for x in self.formats]), - ) - # def get_beta_val(op): - # if(op == '='): - # beta_val = '0.000000e+00' - # elif(op == '+='): - # beta_val = '1.000000e+00' - # elif(op == '-='): - # beta_val = '-1.000000e+00' - # elif(op == '+' or op == '-'): - # beta_val = '0.000000e+00' - # return beta_val - - -class Tensor_Decl_Builder: - formats = ['Dense', 'CSR', 'COO', 'CSC'] - indentation_size = 4 - - tensor_decl_wrapper_text = jinja2.Template( - ("" * indentation_size) - + '%t{{lhs}} = "ta.{{decl}}ensor_decl"{{dims_tuple}}' - + '{format = {{format}}} : ' - +"{{ranges_tuple}} -> " - + "{{inputtype}}" - + "\n" - +'"ta.fill{{where}}"(%t{{lhs}}) {{value}} : ({{inputtype}}) -> ()\n', - undefined=jinja2.StrictUndefined, - ) - - tensor_decl_wrapper_text_no_fill = jinja2.Template( - ("" * indentation_size) - + '%t{{lhs}} = "ta.{{decl}}ensor_decl"{{dims_tuple}}' - + '{format = {{format}}} : ' - +"{{ranges_tuple}} -> " - + "{{inputtype}}" - + "\n" , - undefined=jinja2.StrictUndefined, - ) - - def __init__(self, data)->None: - self.lhs = data["id"] - self.inputtype = "tensor<{}x{}>".format("x".join(str(v) for v in data["shape"]), data["value_type"]) - # self.decl_vars = data["dimsSSA"] - self.decl_vars = [] - self.format = data["format"] - self.is_input = data["is_input"] - - - def build_tensor(self): - dims_tuple = "({})".format(",".join("%d"+str(v) for v in self.decl_vars)) - ranges_tuple = "({})".format(",".join(["index"]* len(self.decl_vars))) - - if not self.format == DENSE: - where = "_from_file" - format = '"{}" , temporal_tensor = false'.format(self.formats[self.format]) - value = '{filename = "SPARSE_FILE_NAME0", readMode = 1 : i32}' - else: - where = "" - format = '"{}"'.format(self.formats[self.format]) - value = '{value = 0.0 : f64}' - - if self.is_input or self.format == DENSE: - return self.tensor_decl_wrapper_text.render( - lhs = self.lhs, - dims_tuple = dims_tuple, - ranges_tuple = ranges_tuple, - format = format, - inputtype = self.inputtype, - decl = "dense_t" if self.format == DENSE else "spT", - where = where, - value = value - ) - else: - return self.tensor_decl_wrapper_text_no_fill.render( - lhs = self.lhs, - dims_tuple = dims_tuple, - ranges_tuple = ranges_tuple, - format = format, - inputtype = self.inputtype, - decl = "dense_t" if self.format == DENSE else "spT", - ) - -class PrintBuilder: - indentation_size = 4 - - tensor_print_text = jinja2.Template( - ("" * indentation_size) - +'"ta.print"(%t{{tensor}}) : ({{outtype}}) -> ()\n', - undefined=jinja2.StrictUndefined, - ) - - def __init__(self, data): #operand, input_labels, dtype, label_map): - self.operand = data["operands"][0] - if data["shapes"] == 1 or data["shapes"] == [1]: - self.outtype = data["value_type"] - else: - self.outtype = "x".join(str(v) for v in data["shapes"][0]) - self.outtype = "tensor<{}x{}>".format(self.outtype, data["value_type"]) - - def build_op(self): - return self.tensor_print_text.render( - tensor = self.operand, - outtype = self.outtype, - ) \ No newline at end of file diff --git a/frontends/numpy-scipy/cometpy/MLIRGen/lowering.py b/frontends/numpy-scipy/cometpy/MLIRGen/lowering.py index 8c6962b4..18eb4a75 100644 --- a/frontends/numpy-scipy/cometpy/MLIRGen/lowering.py +++ b/frontends/numpy-scipy/cometpy/MLIRGen/lowering.py @@ -29,11 +29,52 @@ import ctypes from ctypes import * import cometpy.cfg as cfg +from cometpy.MLIRGen import types +from cometpy.MLIRGen import ops import atexit import uuid import scipy as scp +from cometpy.MLIRGen.utils import * + debug = False temp_dir = '.cometpy/' +class KernelCache: + + def __init__(self): + self.map = {} + + def find(self, func_name, input_types): + cached_kernels = self.map.get(func_name) + if cached_kernels: + cached_kernel = cached_kernels.get("_".join(input_types)) + if not cached_kernel: + return None + if cached_kernel and cached_kernel.timestamp < os.path.getmtime(func_name.split(':')[0]): + return None + return cached_kernel + else: + return None + + def insert(self, func_name, cached_kernel): + cached_kernels = self.map.get(func_name) + if cached_kernels: + cached_kernels["_".join(cached_kernel.input_types)] = cached_kernel + else: + self.map[func_name] = { "_".join(cached_kernel.input_types) : cached_kernel} + +cache = KernelCache() + +class CachedKernelInfo: + + def __init__(self, name, src_path, lib_path, timestamp, kernel_name, input_types, output_types): + self.name = name + self.src_path = src_path + self.lib_path = lib_path + self.timestamp = timestamp + self.kernel_name = kernel_name + self.input_types = input_types + self.output_types = output_types + if not os.path.exists(temp_dir): try: @@ -63,33 +104,6 @@ def cleanup(): # pass atexit.register(cleanup) -class memref_i64(Structure): - _fields_ = [ ('mem_aligned', POINTER(c_longlong)), ('mem', POINTER(c_longlong)), ('offset', c_longlong), ('dim', c_longlong), ('stride', c_longlong)] - -class memref_f64(Structure): - _fields_ = [ ('mem_aligned', POINTER(c_double)), ('mem', POINTER(c_double)), ('offset', c_longlong), ('dim', c_longlong), ('stride', c_longlong)] - - -def np_array_to_memref(np_array): - ctype = ctypes.c_longlong - if np_array.dtype == 'int32': - ctype = c_int32 - elif np_array.dtype == 'float32': - ctype = c_float - elif np_array.dtype == 'float64': - ctype = c_double - return np_array.ctypes.data_as(ctypes.POINTER(ctype)), np_array.ctypes.data_as(ctypes.POINTER(ctype)), 0, np_array.shape[0], 1 - -def expand_memref_ptr(memref): - return byref(memref), byref(memref), 0, 1, 1 - -def len_dense(vals): - num = 0 - for v in vals: - if not scp.sparse.issparse(v): - num+=1 - return num - @@ -119,236 +133,7 @@ def lower_ta_to_mlir(mlir_in, mlir_lower_flags, uuid_s): return scf_out_file -def all_dense(arg_vals) -> bool : - for v in arg_vals: - if scp.sparse.issparse(v): - return False; - return True - -def comment_unneeded_dense(input_, arg_vals): - input = input_.splitlines() - outs = [] - indexes = [] - for i, v in enumerate(arg_vals): - if not scp.sparse.issparse(v): - indexes.append(i) - replace = {} - fill_remove = [] - - allocs_needed = len_dense(arg_vals) - - for i in range(len(input)): - if "call @comet_print_memref_f64" in input[i]: - cast = input[i][input[i].find("(") + 1 : input[i].find(")")] - # input[i] = "// from dense " + input[i] - input[i] = "" - for j in range(len(input[:i])): - if cast + " = memref.cast" in input[j]: - out = input[j][input[j].find("%alloc") : input[j].find(":")].lstrip().strip() - outs.append(out) - start = input[j].find(":") - end = input[j][start:].find("to") - replace['%arg'+str(len(arg_vals))] = out +" " + input[j][start:start+end].lstrip().strip() - for k in range(len(input[:j])): - if out+" = memref.alloc(" in input[k]: - # input[k] = "//from dense" + input[k] - input[k] = "" - elif allocs_needed > 0 and "memref.alloc" in input[i]: - allocs_needed = allocs_needed - 1 - a = input[i][input[i].find("%") : input[i].find("=")].lstrip().strip() - start = input[i].rfind(":") - replace['%arg'+str(indexes[0])] = a +" " + input[i][start:].lstrip().strip() - indexes = indexes[1:] - # input[i] = "//from dense" + input[i] - input[i] = "" - fill_remove.append(a) - - elif "linalg.fill" in input[i]: - for k in range(len(fill_remove)): - if fill_remove[k] in input[i]: - # input[i] = "//from dense" + input[i] - input[i] = "" - fill_remove.remove(fill_remove[k]) - break - - for v in replace: - start = input[1].find(v) - end = input[1][start:].find(",") - if end == -1 : - end = input[1][start:].find(")") - repl = input[1][start:start+end] - input[1] = input[1].replace(repl, replace[v]) - - output = "" - arg_vals_init = {} - input_arg_names = [] - collected = False - for i in range(len(input)): - l = input[i] - if not collected and l.lstrip().strip().startswith("func.func @") and l.strip().endswith("{"): - for arg in l[l.find("(")+1:l.find(")")].split(","): - if arg.split(":")[0].strip().lstrip() not in outs: - input_arg_names.append(arg.split(":")[0].strip().lstrip()) - for v in input_arg_names: - arg_vals_init[v] = False - collected = True - elif "memref.store" in input[i] : - for arg in input_arg_names: - if arg+"[" in input[i] and arg_vals_init[arg] == False: - init = input[i].lstrip().split(" ")[1][:-1].lstrip().strip() - for l in input: - if l.lstrip().strip().startswith(init) and "arith.constant" in l: - if float(l.strip().lstrip().split(" ")[3]) == 0.0 : - input[i] = "" - arg_vals_init[arg] = True - break - - - for line in input: - if line: - output += line +"\n" - - return output - -def comment_unneeded_sparse(input_, arg_vals): - output = "" - input = input_.splitlines() - indexes = [] - indexes = [] - allocs = [] - alloc_lines = [] - returns = [] - return_found = False - out_len = len(input) - for i in range(len(input)): - line = input[i] - - if "call @read_input_sizes" in input[i]: - - # cast = line[line.find("(") : line.find(")")].split(",")[3].lstrip().strip() - # With tiles - cast = line[line.find("(") : line.find(")")].split(",")[5].lstrip().strip() - - input[i] = "//from sparse" +input[i] - # input[i] = "" - alloc = "" - for j in range(len(input[:i])): - lline = input[j] - if cast +" = memref.cast" in lline: - alloc = lline.split()[3].lstrip().strip() - # input[j] = "// from sparse" + input[j] - # input[j] = "" - for k in range(len(input[:j])): - if alloc + " = memref.alloc" in input[k]: - # input[k] = "// from sparse"+input[k] - input[k] = "" - allocs.append(alloc) - i+=1 - found = 0 - - # while(found != 7): - # With tiles - while(found != 9): - if "memref.load " + alloc in input[i]: - idx = input[i].split('=')[0].lstrip().strip() - indexes.append(idx) - found += 1 - i+=1 - idx = 0 - elif 'memref.alloc(' in input[i]: - line = input[i] - idx = line[line.find('(') + 1: line.find(')')] - # if idx in indexes[:-1]: - if idx in indexes: - indexes.remove(idx) - allocs.append(line.split('=')[0].rstrip().strip()) - cur_lines = [] - cur_lines.append(i) - while "scf.for" not in input[i]: - i+=1 - cur_lines.append(i) - cur_lines.append(i+1) - cur_lines.append(i+2) - cur_lines.append(i+3) - alloc_lines.append(cur_lines) - elif "call @read_input_2D" in input[i]: - # casts = input[i].split(",")[3:8] - # With tiles - casts = input[i].split(",")[5:14] - for lines in alloc_lines: - if "memref.cast" in input[lines[-1]]: - for c in casts: - if c in input[lines[-1]]: - for l in lines[:-1]: - # input[l] = "// from sparse new " + input[l] - input[l] = "" - input[i] = '// from sparse' + input[i] - # input[i] = "" - elif "call @comet_print_memref_i64" in input[i] or "call @comet_print_memref_f64" in input[i]: - cast = input[i][input[i].find("(") + 1 : input[i].find(")")] - for j in range(len(input[:i])): - lline = input[j] - if cast + " = memref.cast" in lline: - alloc = lline.split()[3].lstrip().strip() - type = lline.split(":")[1].split("to")[0].strip() - returns.append((alloc, i, type)) - elif ("return" in input[i]) and len(returns) > 1 and not return_found: - return_found = True - add = "" - for k, r in enumerate(returns[:-1]): - add += "\t\tmemref.store {}, %marg{}[%c0] : memref<1x{}>\n".format(r[0], k, r[2] ) - # input[r[1]] = "//from sparse" + input[r[1]] - input[r[1]] = "" - # input[returns[-1][1]] = "//from sparse" + input[returns[-1][1]] - input[returns[-1][1]] = "" - add += "\t\tmemref.store {}, %marg{}[%c0] : memref<1x{}>\n".format(returns[-1][0], len(returns)-1, returns[-1][2]) - add += "\t\treturn" - input[i] = add - - args = input[1][input[1].find("(") + 1: input[1].find(")")].split(",") - ai = 0 - for i, v in enumerate(arg_vals): - if scp.sparse.issparse(v): - - # input[1] = input[1].replace(args[i], allocs[ai] +" : memref<7xindex>, " + " : memref, ".join([s for s in allocs[ai+1:ai+6]]) + " : memref") - # ai += 6 - # With tiles - input[1] = input[1].replace(args[i], allocs[ai] +" : memref<13xindex>, " + " : memref, ".join([s for s in allocs[ai+1:ai+10]]) + " : memref") - ai += 10 - - if len(returns) > 1: -# input[1] = input[1].replace(")", ", %marg0: memref<1xmemref>, %marg1: memref<1xmemref>, %marg2: memref<1xmemref>, %marg3: memref<1xmemref>, %marg4: memref<1xmemref>)") -# input[-1] = '\n func.func @dealloc(%to_dealloc: memref, %to_dealloc1: memref, %to_dealloc2: memref, %to_dealloc3: memref, %to_dealloc4: memref){\n \ -# \t\tmemref.dealloc %to_dealloc : memref\n \ -# \t\tmemref.dealloc %to_dealloc1 : memref\n \ -# \t\tmemref.dealloc %to_dealloc2 : memref\n \ -# \t\tmemref.dealloc %to_dealloc3 : memref\n \ -# \t\tmemref.dealloc %to_dealloc4 : memref\n \ -# \t\treturn\n \ -# \t}\n\ -# }\n' - # With tiles - input[1] = input[1].replace(")", ", %marg0: memref<1x{}>, %marg1: memref<1x{}>, %marg2: memref<1x{}>, %marg3: memref<1x{}>, %marg4: memref<1x{}>, %marg5: memref<1x{}>, %marg6: memref<1x{}>, %marg7: memref<1x{}>, %marg8: memref<1x{}>)".format(*[x[2] for x in returns])) - input[-1] = '\n func.func @dealloc(%to_dealloc: memref, %to_dealloc1: memref, %to_dealloc2: memref, %to_dealloc3: memref, %to_dealloc4: memref, %to_dealloc5: memref, %to_dealloc6: memref, %to_dealloc7: memref, %to_dealloc8: memref){\n \ -\t\tmemref.dealloc %to_dealloc : memref\n \ -\t\tmemref.dealloc %to_dealloc1 : memref\n \ -\t\tmemref.dealloc %to_dealloc2 : memref\n \ -\t\tmemref.dealloc %to_dealloc3 : memref\n \ -\t\tmemref.dealloc %to_dealloc4 : memref\n \ -\t\tmemref.dealloc %to_dealloc5 : memref\n \ -\t\tmemref.dealloc %to_dealloc6 : memref\n \ -\t\tmemref.dealloc %to_dealloc7 : memref\n \ -\t\tmemref.dealloc %to_dealloc8 : memref\n \ -\t\treturn\n \ -\t}\n\ -}\n' - for line in input: - if line: - output += line +"\n" - return output - def lower_ta_to_mlir_with_jit(mlir_in, mlir_lower_flags, arg_vals, uuid_s): - path_to_comet = cfg.comet_path+"/bin/comet-opt -x mlir " command = path_to_comet + mlir_lower_flags p = subprocess.run(shlex.split(command), input = mlir_in.encode('utf-8') , stdout=subprocess.PIPE, stderr=subprocess.PIPE, shell=False,close_fds=False) @@ -367,9 +152,8 @@ def lower_ta_to_mlir_with_jit(mlir_in, mlir_lower_flags, arg_vals, uuid_s): # files_to_cleanup.append(os.path.join( os.getcwd(), scf_out_file)) scf_out = p.stderr.decode() - - scf_out = comment_unneeded_sparse(scf_out, arg_vals) - scf_out = comment_unneeded_dense(scf_out, arg_vals) + # scf_out = comment_unneeded_sparse(scf_out, arg_vals) + # scf_out = comment_unneeded_dense(scf_out, arg_vals) # f.write(scf_out) # f.close() @@ -419,183 +203,262 @@ def lower_scf_to_llvm(scf_in, scf_lower_flags, uuid_s): # result = p.stdout # return result +def memref_index_type(): + return [ctypes.POINTER(c_int64), ctypes.POINTER(c_int64), c_int64, c_int64, c_int64] + +def memref_f64_type(): + return [ctypes.POINTER(c_double), ctypes.POINTER(c_double), c_int64, c_int64, c_int64] + +def generate_llvm_args_from_ndarrays(inputs, output_types): + llvm_args_types_len = 0 + for out_type in output_types: + if isinstance(out_type, types.ShapedType): + llvm_args_types_len += 1 + llvm_args_len = llvm_args_types_len + for input in inputs: + if scp.sparse.issparse(input): + llvm_args_len += 7 + else: + llvm_args_len += 1 + + llvm_args = [None] * llvm_args_len + aux = [] # Used to make sure data created in this function are not freed prematurely. + + for i, out_type in enumerate(output_types): + if isinstance(out_type, types.ShapedType): + if out_type.format == types.DENSE: + memref = memref_from_shaped_type(out_type) + llvm_args[i] = memref + else: + dim_sizes = memref_from_shaped_type(types.TensorType([1], 'index') ) + aux.append(dim_sizes) + Aval = memref_from_shaped_type(types.TensorType([out_type.shape[0]], out_type.element_type)) + if out_type.format == types.CSR: + A1pos_temp = np.array([out_type.shape[0]], dtype= ops.mlir_type_to_dtype(out_type.indices_type)) + aux.append(A1pos_temp) # Make sure the array we created persists until we actually call the mlir-generated function + A1pos = memref_from_np_array(A1pos_temp) #memref_from_shaped_type(types.TensorType([out_type.shape[0]], out_type.indices_type)) + A2pos = memref_from_shaped_type(types.TensorType([out_type.shape[0]], out_type.indices_type)) + A2crd = memref_from_shaped_type(types.TensorType([out_type.shape[0]], out_type.indices_type)) + + + if out_type.element_type == 'f64': + if out_type.indices_type == 'i32': + out_ctor = output_csr_f64_i32 + else: + out_ctor = output_csr_f64_i64 + elif out_type.element_type == 'f32': + if out_type.indices_type == 'i32': + out_ctor = output_csr_f32_i32 + else: + out_ctor = output_csr_f32_i64 + + out = out_ctor(dim_sizes, 0, A1pos, 0, A2pos, A2crd, Aval) + elif out_type.format == types.COO: + A1pos_temp = np.array([0, 1], dtype=ops.mlir_type_to_dtype(out_type.indices_type)) + aux.append(A1pos_temp) # Make sure the array we created persists until we actually call the mlir-generated function + A1pos = memref_from_np_array(A1pos_temp) #memref_from_shaped_type(types.TensorType([out_type.shape[0]], out_type.indices_type)) + A1crd = memref_from_shaped_type(types.TensorType([out_type.shape[0]], out_type.indices_type)) + A2crd = memref_from_shaped_type(types.TensorType([out_type.shape[0]], out_type.indices_type)) + + if out_type.element_type == 'f64': + if out_type.indices_type == 'i32': + out_ctor = output_coo_f64_i32 + else: + out_ctor = output_coo_f64_i64 + elif out_type.element_type == 'f32': + if out_type.indices_type == 'i32': + out_ctor = output_coo_f32_i32 + else: + out_ctor = output_coo_f32_i64 + + out = out_ctor(dim_sizes, 0, A1pos, A1crd, 0, A2crd, Aval) + else : + raise Exception("Unsupported sparse matrix type") + + llvm_args[i] = out + offset = llvm_args_types_len + for i, ndarray in enumerate(inputs): + -def generate_llvm_args_from_ndarrays(num_in, *ndargs): - llvm_args = [] - llvm_args_types = [] - all_outputs = [] - for i, ndarray in enumerate(ndargs): - # Ndarray is dense if not scp.sparse.issparse(ndarray): - llvm_args.append(ndarray.ctypes.data_as(ctypes.POINTER(ctypes.c_double))) - llvm_args.append(ndarray.ctypes.data_as(ctypes.POINTER(ctypes.c_double))) - llvm_args.append(0) - llvm_args_types.append(ctypes.POINTER(ctypes.c_double)) - llvm_args_types.append(ctypes.POINTER(ctypes.c_double)) - llvm_args_types.append(ctypes.c_longlong) - for s in ndarray.shape: - llvm_args.append(s) - llvm_args_types.append(ctypes.c_longlong) - - for s in ndarray.strides: - llvm_args.append(s//8) - llvm_args_types.append(ctypes.c_longlong) - if i >= num_in: - all_outputs.append(ndarray) - # Ndarray is sparse + # ndarray = np.array(out_type.shape, dtype=ops.mlir_type_to_dtype(out_type.element_type)) ## [TODO] No need to allocate + if hasattr(ndarray, 'shape'): + memref = memref_from_np_array(ndarray) + llvm_args[offset + i] = memref + else: + llvm_args[offset + i] = ndarray else: - # Working on the output matrix - if i >= num_in: - A1pos = memref_i64() - A1crd = memref_i64() - A1tile_pos = memref_i64() - A1tile_crd = memref_i64() - A2pos = memref_i64() - A2crd = memref_i64() - A2tile_pos = memref_i64() - A2tile_crd = memref_i64() - Aval = memref_f64() - - # llvm_args += [*expand_memref_ptr(A1pos), *expand_memref_ptr(A1crd), *expand_memref_ptr(A2pos), *expand_memref_ptr(A2crd), *expand_memref_ptr(Aval)] - # llvm_args_types += [POINTER(memref_i64), POINTER(memref_i64), c_longlong, c_longlong, c_longlong] * 4 + [POINTER(memref_f64), POINTER(memref_f64), c_longlong, c_longlong, c_longlong] - # all_outputs.append((A1pos, A1crd, A2pos, A2crd, Aval)) - - # With tiles - llvm_args += [*expand_memref_ptr(A1pos), *expand_memref_ptr(A1crd), *expand_memref_ptr(A1tile_pos), *expand_memref_ptr(A1tile_crd), *expand_memref_ptr(A2pos), *expand_memref_ptr(A2crd), *expand_memref_ptr(A2tile_pos), *expand_memref_ptr(A2tile_crd), *expand_memref_ptr(Aval)] - llvm_args_types += [POINTER(memref_i64), POINTER(memref_i64), c_longlong, c_longlong, c_longlong] * 8 + [POINTER(memref_f64), POINTER(memref_f64), c_longlong, c_longlong, c_longlong] - all_outputs.append((A1pos, A1crd, A1tile_pos, A1tile_crd, A2pos, A2crd, A2tile_pos, A2tile_crd, Aval)) - else: - # [TODO] The arrays used as inputs for the comet generated code need to be updated to take into account the extra tile component - A1tile_pos = np.array([-1], dtype=np.int64) - A1tile_crd = np.array([-1], dtype=np.int64) - A2tile_pos = np.array([-1], dtype=np.int64) - A2tile_crd = np.array([-1], dtype=np.int64) - # CSR - if ndarray.format == 'csr': - A1pos = np.array([ndarray.shape[0]], dtype=np.int64) - A1crd = np.array([-1], dtype=np.int64) - A2pos = ndarray.indptr.astype('int64') - A2crd = ndarray.indices.astype('int64') - - # Based on the desc_sizes array in SparseUtils.cpp:read_input_sizes_2D - - # llvm_args += [*np_array_to_memref(np.array([1, 1, ndarray.shape[0] + 1, ndarray.nnz, ndarray.nnz, ndarray.shape[0], ndarray.shape[1]], dtype='int64'))] - # With tiles - llvm_args += [*np_array_to_memref(np.array([1, 1, 0, 0, ndarray.shape[0] + 1, ndarray.nnz, 0, 0, ndarray.nnz, ndarray.shape[0], ndarray.shape[1]], dtype='int64'))] - # COO - elif ndarray.format == 'coo': - A1pos = np.array([0, ndarray.nnz], dtype=np.int64) - A1crd = ndarray.row.astype('int64') - A2pos = np.array([-1], dtype=np.int64) - A2crd = ndarray.col.astype('int64') - - # Based on the desc_sizes array in SparseUtils.cpp:read_input_sizes_2D - - # llvm_args += [*np_array_to_memref(np.array([2, ndarray.nnz, 1, ndarray.nnz, ndarray.nnz, ndarray.shape[0], ndarray.shape[1]], dtype='int64'))] - # With tiles - llvm_args += [*np_array_to_memref(np.array([2, ndarray.nnz, 0, 0, 1, ndarray.nnz, 0, 0, ndarray.nnz, ndarray.shape[0], ndarray.shape[1]], dtype='int64'))] - - # CSC - elif ndarray.format == 'csc': - A1pos = ndarray.indptr.astype('int64') - A1crd = ndarray.indices.astype('int64') - A2pos = np.array([ndarray.shape[1]], dtype=np.int64) - - # Based on the desc_sizes array in SparseUtils.cpp:read_input_sizes_2D - - # llvm_args += [*np_array_to_memref(np.array([ndarray.shape[1] + 1, ndarray.nnz, 1, 1, ndarray.nnz, ndarray.shape[0], ndarray.shape[1]], dtype='int64'))] - # With tiles - llvm_args += [*np_array_to_memref(np.array([ndarray.shape[1] + 1, ndarray.nnz, 0, 0, 1, 1, 0, 0, ndarray.nnz, ndarray.shape[0], ndarray.shape[1]], dtype='int64'))] + dims = np.array(ndarray.shape, dtype=np.int64) + dim_sizes = memref_from_np_array(dims) + aux.append(dims) + Aval = memref_from_np_array(ndarray.data) # memref_from_shaped_type(types.TensorType([ndarray.shape[0]], out_type.element_type)) + insert_pos_1 = 0 + insert_pos_2 = 0 + # CSR + if ndarray.format == 'csr': + # A1tile_pos = np.array([-1], dtype=ndarray.indptr.dtype) + # A1tile_crd = np.array([-1], dtype=ndarray.indptr.dtype) + # A2tile_pos = np.array([-1], dtype=ndarray.indptr.dtype) + # A2tile_crd = np.array([-1], dtype=ndarray.indptr.dtype) + A1pos_temp = np.array([ndarray.shape[0]], dtype=ndarray.indptr.dtype) + aux.append(A1pos_temp) # Make sure the array we created persists until we actually call the mlir-generated function + A1pos = memref_from_np_array(A1pos_temp) + A2pos = memref_from_np_array(ndarray.indptr) + A2crd = memref_from_np_array(ndarray.indices) + + # Based on the desc_sizes array in SparseUtils.cpp:read_input_sizes_2D + llvm_args[offset + i] = dim_sizes + offset += 1 + llvm_args[offset + i] = insert_pos_1 + offset += 1 + llvm_args[offset + i] = A1pos + offset += 1 + llvm_args[offset + i] = insert_pos_2 + offset += 1 + llvm_args[offset + i] = A2pos + offset += 1 + llvm_args[offset + i] = A2crd + offset += 1 + # COO + elif ndarray.format == 'coo': + # A1tile_pos = np.array([-1], dtype=ndarray.row.dtype) + # A1tile_crd = np.array([-1], dtype=ndarray.row.dtype) + # A2tile_pos = np.array([-1], dtype=ndarray.row.dtype) + # A2tile_crd = np.array([-1], dtype=ndarray.row.dtype) + A1pos_temp = np.array([0, ndarray.nnz], dtype=ndarray.row.dtype) + aux.append(A1pos_temp) # Make sure the array we created persists until we actually call the mlir-generated function + A1pos = memref_from_np_array(A1pos_temp) + A1crd = memref_from_np_array(ndarray.row) + A2crd = memref_from_np_array(ndarray.col) + llvm_args[offset + i] = dim_sizes + offset += 1 + llvm_args[offset + i] = insert_pos_1 + offset += 1 + llvm_args[offset + i] = A1pos + offset += 1 + llvm_args[offset + i] = A1crd + offset += 1 + llvm_args[offset + i] = insert_pos_2 + offset += 1 + llvm_args[offset + i] = A2crd + offset += 1 + # # CSC + # elif ndarray.format == 'csc': + # A1pos = ndarray.indptr.astype('int64') + # A1crd = ndarray.indices.astype('int64') + # A2pos = np.array([ndarray.shape[1]], dtype=np.int64) - Aval = ndarray.data.astype('float64') - # Based on the desc_A1pos/crd, desc_A2pos/crd, desc_Aval arrays in SparseUtils.cpp: read_input_2D - # Expand to memrefs llvmir implementation - - # llvm_args += [*np_array_to_memref(A1pos), *np_array_to_memref(A1crd), *np_array_to_memref(A2pos), *np_array_to_memref(A2crd), *np_array_to_memref(Aval)] - # With tiles - llvm_args += [*np_array_to_memref(A1pos), *np_array_to_memref(A1crd), *np_array_to_memref(A1tile_pos), *np_array_to_memref(A1tile_crd), *np_array_to_memref(A2pos), *np_array_to_memref(A2crd), *np_array_to_memref(A2tile_pos), *np_array_to_memref(A2tile_crd), *np_array_to_memref(Aval)] - - # Set the datatypes expected from the function in the shared library. - # If we don't define this the data are not passed correctly - - # llvm_args_types += [ctypes.POINTER(c_longlong), ctypes.POINTER(c_longlong), c_longlong, c_longlong, c_longlong] * 5 + [ctypes.POINTER(c_double), ctypes.POINTER(c_double), c_longlong, c_longlong, c_longlong] - # With tiles - llvm_args_types += [ctypes.POINTER(c_longlong), ctypes.POINTER(c_longlong), c_longlong, c_longlong, c_longlong] * 9 + [ctypes.POINTER(c_double), ctypes.POINTER(c_double), c_longlong, c_longlong, c_longlong] - - return llvm_args, llvm_args_types, all_outputs + # # Based on the desc_sizes array in SparseUtils.cpp:read_input_sizes_2D + + # # llvm_args += [*np_array_to_memref(np.array([ndarray.shape[1] + 1, ndarray.nnz, 1, 1, ndarray.nnz, ndarray.shape[0], ndarray.shape[1]], dtype='int64'))] + # # With tiles + # llvm_args += [*np_array_to_memref(np.array([ndarray.shape[1] + 1, ndarray.nnz, 0, 0, 1, 1, 0, 0, ndarray.nnz, ndarray.shape[0], ndarray.shape[1]], dtype='int64'))] + + # Based on the desc_A1pos/crd, desc_A2pos/crd, desc_Aval arrays in SparseUtils.cpp: read_input_2D + # Expand to memrefs llvmir implementation + llvm_args[offset + i] = Aval + return llvm_args, aux #Translating llvm dialect to llvm IR using mlir-translate and then executing the IR using lli -def translate_and_exec_llvm_with_jit(llvm_in,scf_lower_flags, func_name, inputs, outputs, uuid_s): - - - llvmir_file = uuid_s+'.ll' - - # path_to_cometopt = cfg.comet_path+"/bin/comet-opt" - path_to_cometopt = cfg.comet_path+"/bin/comet-opt -x mlir" - to_llvm_command = path_to_cometopt + scf_lower_flags #+ llvm_in - translate_mlir_command = cfg.llvm_path+"/bin/mlir-translate --mlir-to-llvmir -- " - libname = "./lib"+llvmir_file+func_name+".so" - gcc_command = cfg.llvm_path+"/bin/clang -march=native -mtune=native -x ir -Wno-everything --shared -O3 "+platform_args+ " -o "+ temp_dir + libname+" -fpic -L {0}/lib/ -Wl,-rpath,{0}/lib/ -lcomet_runner_utils -".format(cfg.comet_path) - - # We merge all several calls in a single call to the shell in order to only pay the overhead of process creation once. - # 1. Call comet to lower scf code to llvm - # 2. Call mlir-translate to convert llvm to llvmir - # 3. Call clang to generate library - p = subprocess.run(to_llvm_command +' 2>&1 | '+ translate_mlir_command +' | ' + gcc_command , input=llvm_in.encode('utf-8'), stdout=subprocess.PIPE, stderr=subprocess.PIPE, shell=True) - if(p.returncode != 0): - cleanup() - raise AssertionError("gcc failed with error code: {}. Error: {}".format(p.returncode, p.stderr)) - files_to_cleanup.append(os.path.join( os.getcwd(), temp_dir +libname)) - - # Load code generated from COMET - lib = ctypes.cdll.LoadLibrary(temp_dir +libname) - func = lib.__getattr__(func_name) - - # Get the inputs, input types and the output containers - args, arg_types, all_output = generate_llvm_args_from_ndarrays(len(inputs),*(inputs), *(outputs)) - func.argtypes = arg_types - - # Uncomment to measure execution time without the compilation process - +def translate_and_exec_llvm_with_jit(llvm_in,scf_lower_flags, kernel_name, inputs, output_types, uuid_s, func_name, input_types, cached_kernel): + llvmir_file = None + if not cached_kernel: + llvmir_file = uuid_s+'.ll' + llvm_in = llvm_in.replace('call @comet_print_memref_', '//call @comet_print_memref_') + # path_to_cometopt = cfg.comet_path+"/bin/comet-opt" + path_to_cometopt = cfg.comet_path+"/bin/comet-opt -x mlir" + to_llvm_command = path_to_cometopt + scf_lower_flags #+ llvm_in + translate_mlir_command = cfg.llvm_path+"/bin/mlir-translate --mlir-to-llvmir -- " + libname = "./lib"+llvmir_file+kernel_name+".so" + gcc_command = cfg.llvm_path+"/bin/clang -march=native -mtune=native -x ir -Wno-everything --shared -O3 "+platform_args+ " -o "+ temp_dir + libname+" -fpic -L {0}/lib/ -Wl,-rpath,{0}/lib/ -lcomet_runner_utils -L {1}/lib/ -Wl,-rpath,{1}/lib/ -fopenmp -".format(cfg.comet_path, cfg.llvm_path) + + # We merge all several calls in a single call to the shell in order to only pay the overhead of process creation once. + # 1. Call comet to lower scf code to llvm + # 2. Call mlir-translate to convert llvm to llvmir + # 3. Call clang to generate library + # p = subprocess.run(to_llvm_command, input=llvm_in.encode('utf-8'), stdout=subprocess.PIPE, stderr=subprocess.PIPE, shell=True) + # start = time.time() + p = subprocess.run(to_llvm_command +' 2>&1 | '+ translate_mlir_command +' | ' + gcc_command , input=llvm_in.encode('utf-8'), stdout=subprocess.PIPE, stderr=subprocess.PIPE, shell=True) + if(p.returncode != 0): + cleanup() + raise AssertionError("gcc failed with error code: {}. Error: {} {}".format(p.returncode, p.stdout, p.stderr)) + files_to_cleanup.append(os.path.join( os.getcwd(), temp_dir +libname)) + # end = time.time() + # print(f'To llvm time: {end-start}') + # start = time.time() + + # Load code generated from COMET + lib = ctypes.cdll.LoadLibrary(temp_dir +libname) + func = lib.__getattr__("_mlir_ciface_"+kernel_name) + timestamp = os.path.getmtime(temp_dir +libname) + + # Get the inputs, input types and the output containers + args, aux = generate_llvm_args_from_ndarrays(inputs, output_types) + # func.argtypes = arg_types + + # end = time.time() + # print(f'dlopen time: {end-start}') + + # Uncomment to measure execution time without the compilation process + # start = time.time() + cache.insert(func_name, CachedKernelInfo(func_name, func_name.split(':')[0], func, timestamp, kernel_name, input_types, output_types)) + else: + args, aux = generate_llvm_args_from_ndarrays(inputs, cached_kernel.output_types) + output_types = cached_kernel.output_types + func = cached_kernel.lib_path + + + if len(output_types) == 1: + if not isinstance(output_types[0], types.ShapedType): + type_str = str(output_types[0]) + if type_str == 'i64': + func.restype = c_int64 + elif type_str == 'i32': + func.restype = c_int32 + elif type_str == 'f32': + func.restype = c_float + elif type_str == 'f64': + func.restype = c_double + else: + raise Exception("Unexpeted return type") + # end = time.time() + # print("create args: {}".format(end-start)) # start = time.time() - func(*(args)) + ret = func(*[byref(arg) if not isinstance(arg, int) else arg for arg in args]) # end = time.time() # print("Kernel execution time JIT: {}".format(end-start)) - out = None ret_outputs = [] - for v0, v1 in zip(all_output, outputs): - if scp.sparse.issparse(v1): - # A1pos, A1crd, A2pos, A2crd, Aval = v0 - # With tiles - A1pos, A1crd, A1tile_pos, A1tile_crd, A2pos, A2crd, A2tile_pos, A2tile_crd, Aval = v0 - if v1.format == 'csr': - np_r_indices = np.ctypeslib.as_array(A2pos.mem, [A2pos.dim]) - np_c_indices = np.ctypeslib.as_array(A2crd.mem, [A2crd.dim]) - np_values = np.ctypeslib.as_array(Aval.mem, [Aval.dim]) - ret_outputs.append(scp.sparse.csr_array((np_values, np_c_indices, np_r_indices), copy=False)) - elif v1.format == 'coo': - rows = np.ctypeslib.as_array(A1crd.mem, [A1crd.dim]) - cols = np.ctypeslib.as_array(A2crd.mem, [A2crd.dim]) - np_values = np.ctypeslib.as_array(Aval.mem, [Aval.dim]) - ret_outputs.append(scp.sparse.coo_array((np_values, (rows, cols)), copy=False)) - else: - if v1.shape == (1,): - ret_outputs.append(np.squeeze(v1)) - else: - ret_outputs.append(v1) - + for v0, v1 in zip(args, output_types): + if isinstance(v1, types.ShapedType): + if v1.format == types.CSR: + out_csr = v0 + np_r_indices = np.ctypeslib.as_array(out_csr.A2pos.mem, [out_csr.A2pos.dim[0]]) + np_c_indices = np.ctypeslib.as_array(out_csr.A2crd.mem, [out_csr.A2crd.dim[0]]) + np_values = np.ctypeslib.as_array(out_csr.Aval.mem, [out_csr.Aval.dim[0]]) + ret_outputs.append(scp.sparse.csr_array((np_values, np_c_indices, np_r_indices), (out_csr.dims_sizes.mem[0], out_csr.dims_sizes.mem[1]) , copy=False)) + elif v1.format == types.COO: + out_coo = v0 + rows = np.ctypeslib.as_array(out_coo.A1crd.mem, [out_coo.A1crd.dim[0]]) + cols = np.ctypeslib.as_array(out_coo.A2crd.mem, [out_coo.A2crd.dim[0]]) + np_values = np.ctypeslib.as_array(out_coo.Aval.mem, [out_coo.Aval.dim[0]]) + ret_outputs.append(scp.sparse.coo_array((np_values, (rows, cols)), (out_coo.dims_sizes.mem[0], out_coo.dims_sizes.mem[1]) , copy=False)) + elif v1.format == types.DENSE: + ret_outputs.append(np.ctypeslib.as_array(v0.mem, v1.shape)) + + if output_types and not ret_outputs and not isinstance(output_types[0], types.ShapedType): + ret_outputs.append(ret) + if len(ret_outputs) == 1: out = ret_outputs.pop() - else: + elif len(ret_outputs) > 1: out = ret_outputs - + else: + out = None + # print("Kernel execution time JIT: {}".format(end-start)) return out, llvmir_file -def func_execute(func, args): - func(*(args)) - def translate_and_exec_llvm(llvm_in,func_name, out_dims, uuid_s): translate_mlir_command = "../llvm/build/bin/mlir-translate --mlir-to-llvmir " + llvm_in @@ -643,112 +506,54 @@ def translate_and_exec_llvm(llvm_in,func_name, out_dims, uuid_s): else: return output_arrays_list.pop(),llvmir_file -def lower_dialect(ta_dialect_rep, out_dims, compile_with_flags,func_name): - - #lower TA dialect to the SCF dialect - mlir_lower_flags = "" - - if isinstance(compile_with_flags,tuple): - for i in range(len(compile_with_flags)): - mlir_lower_flags += compile_with_flags[i] + " " - - elif(isinstance(compile_with_flags,str)): - mlir_lower_flags += compile_with_flags - - mlir_lower_flags += " --convert-ta-to-it --convert-to-loops " - - # scf_lower_flags = " --lower-affine --convert-linalg-to-loops --convert-scf-to-std --convert-linalg-to-llvm --convert-std-to-llvm " - scf_lower_flags = " --convert-to-llvm " - - if("-emit-ta" in mlir_lower_flags): - print(ta_dialect_rep) - return - - #write the TA dialect rep to file - uuid_s = str(uuid.uuid4()) - ta_dialect_file = uuid_s+'.mlir' - # print("uuid_s: ", uuid_s) - if(os.path.exists(ta_dialect_file) is False): - f = open(os.path.join( os.getcwd(), ta_dialect_file), 'w') - files_to_cleanup.append(os.path.join( os.getcwd(), ta_dialect_file)) - else: - f = open(ta_dialect_file, 'w') - - f.write(ta_dialect_rep) - f.close() +def lower_dialect_with_jit(ta_dialect_rep, target: str, out_dims, compile_with_flags,kernel_name, args_vals, outputs_types, func_name, input_types, cached_kernel): - # Uncomment for debugging pusposes - scf_out_file = lower_ta_to_mlir(ta_dialect_file, mlir_lower_flags, uuid_s) - # scf_out_file = lower_ta_to_mlir_with_jit(ta_dialect_file, mlir_lower_flags, args_vals) - - # Running --convert-ta-to-it --convert-to-loops and --convert-to-llvm in separate steps - # does not produce correct output. This is an issue with the backend. - - #lower the SCF dialect to first STD dialect and then to the llvm dialect - llvm_out_file = lower_scf_to_llvm(scf_out_file, scf_lower_flags, uuid_s) - # llvm_out_file = lower_scf_to_llvm(ta_dialect_file, mlir_lower_flags + scf_lower_flags) - - result,llvmir_file = translate_and_exec_llvm(llvm_out_file,func_name, out_dims, uuid_s) - - return result - - -def lower_dialect_with_jit(ta_dialect_rep, target: str, out_dims, compile_with_flags,func_name, args_vals, outputs): - - mlir_lower_flags = " " - - if compile_with_flags != None: - if "--convert-tc-to-ttgt" not in compile_with_flags: - mlir_lower_flags += " --convert-ta-to-it " - if "--opt-fusion" in compile_with_flags: - mlir_lower_flags += "--opt-fusion" - compile_with_flags = compile_with_flags.replace("--opt-fusion","") - compile_with_flags = compile_with_flags.replace("--opt-comp-workspace","") - if "-opt-matmul-tiling" not in compile_with_flags: - mlir_lower_flags += " --convert-to-loops " - mlir_lower_flags =" "+compile_with_flags + mlir_lower_flags - else: - mlir_lower_flags = " --convert-ta-to-it --convert-to-loops " - # scf_lower_flags = " --lower-affine --convert-linalg-to-loops --convert-scf-to-std --convert-linalg-to-llvm --convert-std-to-llvm " - scf_lower_flags = " --convert-to-llvm " - - if target != "cpu": - if target.startswith("sm_") or target.startswith("compute_") or target.startswith("lto_"): - if not cfg.gpu_target_enabled: - raise "COMET gpu target is not enabled" + scf_lower_flags = None + scf_out_file = None + uuid_s = None + if not cached_kernel: + mlir_lower_flags = " " - scf_lower_flags += " " + " --convert-to-triton --target=GPU --gpu-compute-capability="+target.split("_")[1] - mlir_lower_flags += " " + "--target=GPU" - elif target == "gpu": - if not cfg.gpu_target_enabled: - raise "COMET gpu target is not enabled" + if compile_with_flags != None: + if "--convert-tc-to-ttgt" not in compile_with_flags: + mlir_lower_flags += " --convert-ta-to-it " + if "--opt-fusion" in compile_with_flags: + mlir_lower_flags += "--opt-fusion" + compile_with_flags = compile_with_flags.replace("--opt-fusion","") + compile_with_flags = compile_with_flags.replace("--opt-comp-workspace","") + # if "-opt-matmul-tiling" not in compile_with_flags: + mlir_lower_flags += " --convert-to-loops " + mlir_lower_flags =" "+compile_with_flags + mlir_lower_flags + else: + mlir_lower_flags = " --convert-ta-to-it --convert-to-loops " + # scf_lower_flags = " --lower-affine --convert-linalg-to-loops --convert-scf-to-std --convert-linalg-to-llvm --convert-std-to-llvm " + scf_lower_flags = " --convert-to-llvm " - scf_lower_flags += " " + " --convert-to-triton --target=GPU" - mlir_lower_flags += " " + "--target=GPU" - else : - raise "Expected target formats:\ - cpu, compute_, sm_, lto_" - - if("-emit-ta" in mlir_lower_flags): - print(ta_dialect_rep) - return - - uuid_s = str(uuid.uuid4()) - ta_dialect_file = temp_dir+uuid_s+'.mlir' - # if(os.path.exists(ta_dialect_file) == False): - # f = open(os.path.join( os.getcwd(), ta_dialect_file), 'w') - # files_to_cleanup.append(os.path.join( os.getcwd(), ta_dialect_file)) - # else: - # f = open(ta_dialect_file, 'w') - - # f.write(ta_dialect_rep) - # f.close() + if target != "cpu": + if target.startswith("sm_") or target.startswith("compute_") or target.startswith("lto_"): + if not cfg.gpu_target_enabled: + raise Exception("COMET gpu target is not enabled") + scf_lower_flags += " " + "--convert-to-triton --target=GPU --gpu-compute-capability="+target.split("_")[1] + mlir_lower_flags += " " + "--target=GPU" + elif target == "gpu": + if not cfg.gpu_target_enabled: + raise Exception("COMET gpu target is not enabled") + scf_lower_flags += " " + "--convert-to-triton --target=GPU" + mlir_lower_flags += " " + "--target=GPU" + else : + raise "Expected target formats:\ + cpu, compute_, sm_, lto_" + + if("-emit-ta" in mlir_lower_flags): + print(ta_dialect_rep) + return - # Convert TA to SCF - # scf_out_file = lower_ta_to_mlir_with_jit(ta_dialect_file, mlir_lower_flags, args_vals, uuid_s) - scf_out_file = lower_ta_to_mlir_with_jit(ta_dialect_rep, mlir_lower_flags, args_vals, uuid_s) + uuid_s = str(uuid.uuid4()) + scf_out_file = lower_ta_to_mlir_with_jit(ta_dialect_rep, mlir_lower_flags, args_vals, uuid_s) + # end = time.time() + # print(f"To SCF time: {end-start}") #lower the SCF dialect to LLVMIR and execute - result,llvmir_file = translate_and_exec_llvm_with_jit(scf_out_file, scf_lower_flags, func_name, args_vals, outputs, uuid_s) + result,llvmir_file = translate_and_exec_llvm_with_jit(scf_out_file, scf_lower_flags, kernel_name, args_vals, outputs_types, uuid_s, func_name, input_types, cached_kernel) return result \ No newline at end of file diff --git a/frontends/numpy-scipy/cometpy/MLIRGen/ops.py b/frontends/numpy-scipy/cometpy/MLIRGen/ops.py new file mode 100644 index 00000000..645481fd --- /dev/null +++ b/frontends/numpy-scipy/cometpy/MLIRGen/ops.py @@ -0,0 +1,626 @@ +import jinja2 + +from cometpy.MLIRGen.types import * +from cometpy.MLIRGen.utils import * +import scipy as sp + + +class Dialect: + def __init__(self, name): + self.name = name + + def __repr__(self): + return f"{self.name} dialect" + + +class Operation: + + def __init__(self, operands, result_types, startsBody = False, endsBody = False): + self.startsBody = startsBody + self.endsBody = endsBody + self.operands = operands + self.results = [Symbol(t) for t in result_types] + + +class BinaryElementwiseOp(Operation): + + def __init__(self, lhs, rhs, _type = None): + super().__init__([lhs, rhs], [_type if _type else lhs.type]) + self.lhs = lhs + self.rhs = rhs + + def dump(self, name): + return self.binary_op_text.render( + ssa = self.results[0].ssa, + name = name, + lhs = self.lhs.ssa, + rhs = self.rhs.ssa, + type = self.results[0].type + ) + + + binary_op_text = jinja2.Template( + '%{{ssa}} = {{name}} %{{lhs}}, %{{rhs}}: {{type}}', + undefined=jinja2.StrictUndefined, + ) + +class SubOp(BinaryElementwiseOp): + + def __init__(self, lhs, rhs): + super().__init__(lhs, rhs) + + def dump(self): + name = self.text.render( + type_suffix = 'i' if self.lhs.type == 'index' else 'f', + ) + return super().dump(name) + + text = jinja2.Template( + 'arith.sub{{type_suffix}}', + undefined=jinja2.StrictUndefined, + ) + +class MulOp(BinaryElementwiseOp): + + def __init__(self, lhs, rhs): + super().__init__(lhs, rhs) + + def dump(self): + name = self.text.render( + type_suffix = 'i' if self.lhs.type == 'index' else 'f', + ) + return super().dump(name) + + text = jinja2.Template( + 'arith.mul{{type_suffix}}', + undefined=jinja2.StrictUndefined, + ) + +class TensorIndexBasedOp(Operation): + def __init__(self, name, inputs, inputs_indices, res_indices, res_format, alpha = None, beta = None , mask = None, mask_type = None, semiring = None): + self.name = name + self.inputs = inputs + self.inputs_indices = inputs_indices + self.mask_type = mask_type + self.semiring = semiring + self.alpha = alpha + self.beta = beta + self.mask = mask + + self.res_indices = res_indices + all_indices = inputs_indices[0][:] + for indices in inputs_indices: + for index in indices: + if not index in all_indices: + all_indices.append(index) + + indices_to_dims = {} + for i, index in enumerate(all_indices): + indices_to_dims[index] = i + + res_shape = [] + for res_index in self.res_indices: + found = False + for input_id, indices in enumerate(inputs_indices): + for dim, index in enumerate(indices): + if res_index == index: + res_shape.append(inputs[input_id].type.shape[dim]) + found = True + break + if found: + break + + if found: + continue + + affine_map_all = '{}'.format(",".join([f'd{d}' for d in range(len(all_indices))])) + + affine_map_res = 'affine_map<({}) -> ({})>'.format(affine_map_all, ",".join([f'd{indices_to_dims[index]}' for index in self.res_indices])) + affine_maps = [] + for indices in inputs_indices: + affine_map_input = 'affine_map<({}) -> ({})>'.format(affine_map_all, ",".join([f'd{indices_to_dims[index]}' for index in indices])) + affine_maps.append(affine_map_input) + + affine_maps.append(affine_map_res) + self.indexing_maps = ",".join(affine_maps) + + if res_format != DENSE: + for input in inputs: + if isinstance(input.type, TASparseTensorType): + indices_type = input.type.indices_type + res_type = TASparseTensorType(res_shape, inputs[0].type.element_type, indices_type, res_format) + else: + res_type = TensorType(res_shape, inputs[0].type.element_type) + + super().__init__([*inputs, *inputs_indices, res_indices], [res_type]) + + def dump(self): + num_indices = len(self.res_indices) + for indices in self.inputs_indices : + num_indices += len(indices) + input_types = ", ".join([ + ",".join([f'{input.type}' for input in self.inputs]) + , ", ".join([f'{index.type}' for indices in self.inputs_indices for index in indices]) + , ", ".join([f'{oper.type}' for oper in self.res_indices])]) + if self.mask != None: + input_types += f", {self.mask.type}" + return self.tensor_binary_op_text.render( + + ssa = self.results[0].ssa, + name = self.name, + inputs = ",".join([f'%{input.ssa}' for input in self.inputs]), + inputs_indices = ",".join([f'%{index.ssa}' for indices in self.inputs_indices for index in indices]), + res_indices = ",".join([f'%{input.ssa}' for input in self.res_indices]), + mask_type = self.mask_type, + mask = self.mask, + semiring = self.semiring, + input_types = input_types, + res_type = self.results[0].type, + affine_maps = self.indexing_maps, + operand_segment_sizes = ", ".join([", ".join('1' * len(self.inputs)), str(num_indices), '1' if self.mask_type and self.mask_type != 'none' else '0']), + alpha = self.alpha, + beta = self.beta, + ) + + tensor_binary_op_text = jinja2.Template( + '%{{ssa}} = "{{name}}" ({{inputs}}, {{inputs_indices}}, {{res_indices}} {%if mask!=None%}, %{{mask.ssa}}{%endif%} ) <{ {% if mask_type%} MaskType = "{{mask_type}}", {% endif%} indexing_maps = [{{affine_maps}}] {% if semiring%}, operandSegmentSizes = array, semiring="{{semiring}}" {%endif%}}> {%if alpha != None or beta != None %} { {%endif%} {%if alpha!= None%} __alpha__ = {{alpha}} : f64 {%endif%} {%if alpha!= None and beta!= None %}, {%endif%} {%if beta!= None %} __beta__ = {{beta}}: f64 {%endif%} {%if alpha!= None or beta!= None %}} {%endif%} : ({{input_types}}) -> {{res_type}}', + undefined=jinja2.StrictUndefined, + ) + +class TensorBinaryOp(TensorIndexBasedOp): + + def __init__(self, name, lhs, rhs, lhs_indices, rhs_indices, res_indices, alpha, beta, mask, masktype, semiring, res_format): + super().__init__(name, [lhs, rhs], [lhs_indices, rhs_indices], res_indices, res_format, alpha, beta, mask, masktype, semiring) + self.lhs = lhs + self.lhs_indices = lhs_indices + self.rhs = rhs + self.rhs_indices = rhs_indices + self.res_indices = res_indices + self.mask_type = masktype + self.semiring = semiring + + + def dump(self): + return super().dump() + +class TensorElewiseBinaryOp(TensorBinaryOp): + + def __init__(self, name, lhs, rhs, lhs_indices, alpha, beta, semiring, format): + super().__init__(name, lhs, rhs, lhs_indices, lhs_indices, lhs_indices, alpha, beta, None, "none", semiring, format) + + def dump(self): + return super().dump() + + +class TensorIndexLabelOp(Operation): + + def __init__(self): + super().__init__([], [TensorIndexType()]) + + def dump(self): + return self.text.render( + ssa = self.results[0].ssa, + res_type = self.results[0].type + ) + + text = jinja2.Template ( + '%{{ssa}} = "ta.index_label" () : () -> {{res_type}}' + ) + +class TensorAddOp(TensorElewiseBinaryOp): + + def __init__(self, lhs, rhs, res_indices, format, alpha = 1.0, beta= 0.0, semiring= "noop_plusxy"): + super().__init__('ta.add', lhs, rhs, res_indices, alpha, beta, semiring, format) + + def dump(self): + return super().dump() + +class TensorSubOp(TensorElewiseBinaryOp): + + def __init__(self, lhs, rhs, res_indices, format, alpha = 1.0, beta= 0.0, semiring= "noop_minus"): + super().__init__('ta.subtract', lhs, rhs, res_indices, alpha, beta, semiring, format) + + def dump(self): + return super().dump() + +class TensorMulOp(TensorElewiseBinaryOp): + + def __init__(self, lhs, rhs, res_indices, format, alpha = 1.0, beta= 0.0, semiring= "noop_times"): + super().__init__('ta.elews_mul', lhs, rhs, res_indices, alpha, beta, semiring, format) + + def dump(self): + return super().dump() + +class TensorTransposeOp(TensorIndexBasedOp): + + def __init__(self, input, input_indices, res_indices, res_format): + super().__init__('ta.transpose', [input], [input_indices], res_indices, res_format) + self.input = input + self.input_indices = input_indices + self.res_indices = res_indices + + def dump(self): + return super().dump() + + +class TensorSumOp(Operation): + + def __init__(self, value): + super().__init__([value], [value.type.element_type]) + self.value = value + + def dump(self): + return self.text.render( + ssa = self.results[0].ssa, + value = self.value.ssa, + val_type = self.value.type, + res_type = self.results[0].type + ) + + text = jinja2.Template( + '%{{ssa}} = "ta.reduce"(%{{value}}) : ({{val_type}}) -> {{res_type}}' + ) + +class TensorSetOp(Operation): + + def __init__(self, src, dst, beta): + super().__init__([src, dst], []) + self.src = src + self.dst = dst + self.beta = beta + + def dump(self): + return self.text.render ( + src = self.src.ssa, + dst = self.dst.ssa, + beta = self.beta, + beta_type = 'f64', + src_type = self.src.type, + dst_type = self.dst.type, + ) + + + text = jinja2.Template( + '"ta.set_op" (%{{src}}, %{{dst}}) {__beta__ = {{beta}} : {{beta_type}} } : ({{src_type}}, {{dst_type}}) -> ()' + ) + +class TensorPrintOp(Operation): + + def __init__(self, input): + super().__init__([input], []) + self.input = input + + def dump(self): + return self.text.render ( + input = self.input.ssa, + input_type = self.input.type + ) + + + text = jinja2.Template( + '"ta.print" (%{{input}}) : ({{input_type}}) -> ()' + ) + + +class TensorMatMultOp(TensorBinaryOp): + + def __init__(self, lhs, rhs, lhs_indices, rhs_indices, res_indices, res_format, alpha = 1.0, beta = 0.0, mask = None, masktype='none', semiring='plusxy_times'): + super().__init__('ta.mul', lhs, rhs, lhs_indices, rhs_indices, res_indices, alpha, beta, mask, masktype, semiring, res_format) + + +class AddOp(BinaryElementwiseOp): + + def __init__(self, lhs, rhs): + super().__init__(lhs, rhs) + + def dump(self): + name = self.text.render( + type_suffix = 'i' if self.lhs.type == 'index' else 'f', + ) + return super().dump(name) + + text = jinja2.Template( + 'arith.add{{type_suffix}}', + undefined=jinja2.StrictUndefined, + ) + + +class ConstantOp(Operation): + + def __init__(self, val): + self.value = val + + if isinstance(val, int): + _type = 'index' + elif isinstance(val, float): + _type = 'f64' + + super().__init__([], [_type]) + + def dump(self): + return self.text.render( + ssa = self.results[0].ssa, + value = self.value, + type = self.results[0].type + ) + + text = jinja2.Template( + '%{{ssa}} = arith.constant {{value}}: {{type}}', + undefined=jinja2.StrictUndefined, + ) + +class ToMemrefOp(Operation): + + def __init__(self, src): + super().__init__([src], [MemrefType(src.type.shape, src.type.element_type)]) + self.src = src + + + def dump(self): + return self.text.render( + ssa = self.results[0].ssa, + src = self.src.ssa, + res_type = self.results[0].type, + ) + + + text = jinja2.Template ( + '%{{ssa}} = bufferization.to_memref %{{src}}: {{res_type}} ' + ) + + +class LoadOp(Operation): + + def __init__(self, src, indices): + super().__init__([src] + indices, [src.type.element_type]) + self.src = src + self.indices = indices + + def dump(self): + return self.text.render( + ssa = self.results[0].ssa, + src = self.src.ssa, + indices = ",".join(["%{}".format(index.ssa) for index in self.indices]), + src_type = self.src.type + ) + + text = jinja2.Template( + '%{{ssa}} = memref.load %{{src}}[{{indices}}]: {{src_type}}', + undefined=jinja2.StrictUndefined, + ) + +class StoreOp(Operation): + + def __init__(self, value, dst, indices): + super().__init__([value, dst] + indices, []) + self.value = value + self.dst = dst + self.indices = indices + + def dump(self): + return self.text.render( + value = self.value.ssa, + dst = self.dst.ssa, + indices = ",".join(["%{}".format(index.ssa) for index in self.indices]), + dst_type = self.dst.type + ) + + text = jinja2.Template( + 'memref.store %{{value}}, %{{dst}}[{{indices}}]: {{dst_type}}', + undefined=jinja2.StrictUndefined, + ) + + +class Symbol: + curr_ssa = 0 + + def __init__(self, _type): + self.ssa = Symbol.curr_ssa + self.type = _type + Symbol.curr_ssa += 1 + +class SymbolTable: + + def __init__(self): + self.table = {} + Symbol.curr_ssa = 0 + + def insert(self, id, symbol): + self.table[id] = symbol + + def find(self, id): + if id in self.table: + return self.table[id] + else: + return None + +class YieldOp(Operation): + + def __init__(self, values): + super().__init__(values, [v.type for v in values], endsBody=True) + self. values = values + + def dump(self): + return self.text.render( + values = ",".join([f'%{v}' for v in self.values]), + ) + + + text = jinja2.Template( + 'scf.yield {% if values %} ({{values}}) -> () {% endif%}', + undefined=jinja2.StrictUndefined, + ) + +class ReduceOp(Operation): + + def __init__(self, values): + super().__init__(values, [v.type for v in values], endsBody=True) + self. values = values + + def dump(self): + return self.text.render( + values = ",".join([f'%{v}' for v in self.values]), + ) + + + text = jinja2.Template( + 'scf.reduce {% if values %} ({{values}}) -> () {% endif%}', + undefined=jinja2.StrictUndefined, + ) + + +class ForallInParallel(Operation): + + def __init__(self, values): + super().__init__(values, [v.type for v in values], endsBody=True) + self. values = values + + def dump(self): + return self.text.render( + values = ",".join([f'%{v}' for v in self.values]), + ) + + + text = jinja2.Template( + 'scf.forall.in_parallel {}', + undefined=jinja2.StrictUndefined, + ) + + +class ReturnOp(Operation): + + def __init__(self, values): + super().__init__(values, [v.type for v in values], endsBody=True) + self. values = values + + def dump(self): + return self.text.render( + values = ",".join([f'%{v.ssa}' for v in self.values]), + ret_types = ",".join([f'{v.type}' for v in self.values]), + ) + + + text = jinja2.Template( + 'func.return {% if values %} {{values}} : {{ret_types}} {% endif%}', + undefined=jinja2.StrictUndefined, + ) + +class OperationWithBody(Operation): + + def __init__(self, operands, return_types): + super().__init__(operands, return_types, True) + self.body = [] + self.identation = 0 + + def dump(self) -> str: + statements = "\n".join([" "*self.identation + stmt.dump() for stmt in self.body]) + return self.body_text.render( + identation = " "* (self.identation - 2), + statements = statements + ) + + body_text = jinja2.Template( + " {" + + "\n" + + "{{statements}}" + + "\n" + + "{{identation}}" + + "}", + undefined=jinja2.StrictUndefined, + ) + + +class ModuleOp(OperationWithBody): + + def __init__(self): + super().__init__([], []) + + def dump(self) -> str: + op = self.text.render() + statements = super().dump() + return op + statements + + text = jinja2.Template( + "module", + undefined=jinja2.StrictUndefined, + ) + +class FuncOp(OperationWithBody): + + function_text = jinja2.Template( + "func.func {% if private %} private {% endif %}@{{func_name}}({{inputs}}) -> ({{return_types}}) attributes {llvm.emit_c_interface}", + undefined=jinja2.StrictUndefined, + ) + + def __init__( + self, + func_name: str, + inputs, + return_types, + private: bool, + ) -> None: + super().__init__([], return_types) + self.func_name = func_name + self.inputs = inputs + self.return_types = return_types + self.private = private + + + def dump(self) -> str: + return self.function_text.render( + func_name = self.func_name, + inputs = ", ".join(["%{}: {}".format(input.ssa, input.type) for input in self.inputs]), + return_types = ", ".join(["{}".format(ret) for ret in self.return_types]), + private = self.private + ) + super().dump() + + +class ForOp(OperationWithBody): + + def __init__(self, lb, ub, step): + super().__init__([lb, ub, step], []) + self.lb = lb + self.ub = ub + self.step = step + self.iv = Symbol('index') + + def dump(self): + op = self.text.render( + iv = self.iv.ssa, + lb = self.lb.ssa, + ub = self.ub.ssa, + step = self.step.ssa, + ) + return op + super().dump() + + + text = jinja2.Template( + 'scf.for %{{iv}} = %{{lb}} to %{{ub}} step %{{step}}', + undefined=jinja2.StrictUndefined, + ) + + +class ForAllOp(OperationWithBody): + + def __init__(self, lb, ub, step): + super().__init__([lb, ub, step], []) + self.lb = lb + self.ub = ub + self.step = step + self.iv = Symbol('index') + + def dump(self): + my_render = self.text.render( + iv = self.iv.ssa, + lb = self.lb.ssa, + ub = self.ub.ssa, + step = self.step.ssa, + ) + body_render = super().dump() + return my_render + body_render + + + text = jinja2.Template( + 'scf.forall (%{{iv}}) = (%{{lb}}) to (%{{ub}}) step (%{{step}}) ', + undefined=jinja2.StrictUndefined, + ) diff --git a/frontends/numpy-scipy/cometpy/MLIRGen/types.py b/frontends/numpy-scipy/cometpy/MLIRGen/types.py index 7814668a..015f36be 100644 --- a/frontends/numpy-scipy/cometpy/MLIRGen/types.py +++ b/frontends/numpy-scipy/cometpy/MLIRGen/types.py @@ -2,4 +2,48 @@ CSR = 1 COO = 2 CSC = 3 -UNSUPPORTED_FORMAT = -1 \ No newline at end of file +UNSUPPORTED_FORMAT = -1 + + +class ShapedType: + def __init__(self, shape, element_type, format): + self.shape = shape + self.element_type = element_type + self.format = format + + +class TensorIndexType: + def __init__(self): + pass + + def __str__(self): + return '!ta.index' + +class MemrefType(ShapedType): + + def __init__(self, shape, element_type): + super().__init__(shape, element_type, DENSE) + + def __str__(self): + return "memref<{}x{}>".format("x".join([str(d) for d in self.shape]), self.element_type) + +class TensorType(ShapedType): + + def __init__(self, shape, element_type): + super().__init__(shape, element_type, DENSE) + + def __str__(self): + return "tensor<{}x{}>".format("x".join([str(d) for d in self.shape]), self.element_type) + + +class TASparseTensorType(ShapedType): + def __init__(self, shape, element_type, indices_type, format): + super().__init__(shape, element_type, format) + self.indices_type = indices_type + if format == CSR: + self.sp_format = 'd, unk, cu, unk' + elif format == COO: + self.sp_format = 'cn, unk, s, unk' + + def __str__(self): + return "!ta.sparse_tensor<{}, {}, {}, {}>".format(self.element_type, self.indices_type, "x".join(['?' for s in self.shape]), self.sp_format) diff --git a/frontends/numpy-scipy/cometpy/MLIRGen/types_mlir.py b/frontends/numpy-scipy/cometpy/MLIRGen/types_mlir.py deleted file mode 100644 index a016ff72..00000000 --- a/frontends/numpy-scipy/cometpy/MLIRGen/types_mlir.py +++ /dev/null @@ -1,324 +0,0 @@ -# -# Copyright 2022 Battelle Memorial Institute -# -# Redistribution and use in source and binary forms, with or without modification, -# are permitted provided that the following conditions are met: -# -# 1. Redistributions of source code must retain the above copyright notice, this list of conditions -# and the following disclaimer. -# -# 2. Redistributions in binary form must reproduce the above copyright notice, this list of conditions -# and the following disclaimer in the documentation and/or other materials provided with the distribution. -# -# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED -# WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR -# PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, -# INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE -# GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, -# WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE -# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. -# - -import re -from collections import OrderedDict -from typing import Sequence - - -class AliasMap(OrderedDict): - def __setitem__(self, name, type): - type = Type.find(type, aliases=self) - super().__setitem__(name, type) - - -class Type: - _subtypes = [] - - def __repr__(self): - return f"{self.__class__.__name__}({self})" - - def __eq__(self, other): - if not isinstance(other, Type): - return NotImplemented - return self.__class__ is other.__class__ and str(self) == str(other) - - def __init_subclass__(cls): - Type._subtypes.append(cls) - - def to_pretty_string(self): - return str(self) - - @staticmethod - def find(text: str, aliases: AliasMap = None): - if isinstance(text, Type): - return text - if aliases is not None and text[:1] == "#": - if alias := aliases.get(text[1:]): - return alias - - for klass in Type._subtypes: - result = klass.parse(text, aliases=aliases) - if result is not None: - return result - - raise TypeError(f"Unknown type: {text}") - - -class IndexType(Type): - def __init__(self): - pass - - def __str__(self): - return "index" - - @classmethod - def parse(cls, text: str, aliases: AliasMap = None): - if text == "index": - return IndexType() - - -class FloatType(Type): - _patt = re.compile(r"^f(\d+)$") - - def __init__(self, num: int): - self.num = num - - def __str__(self): - return f"f{self.num}" - - @classmethod - def parse(cls, text: str, aliases: AliasMap = None): - if m := cls._patt.match(text): - return FloatType(int(m.group(1))) - - -class IntType(Type): - _patt = re.compile(r"^i(\d+)$") - - def __init__(self, num: int): - self.num = num - - def __str__(self): - return f"i{self.num}" - - @classmethod - def parse(cls, text: str, aliases: AliasMap = None): - if m := cls._patt.match(text): - return IntType(int(m.group(1))) - - -class MemrefType(Type): - _patt = re.compile(r"^memref<\s*((?:(?:[^,])*x)*)([^, ]+)\s*>$") - - def __init__(self, shape: Sequence[int], value_type: Type): - shape = tuple(shape) - if not isinstance(value_type, Type): - raise TypeError(f"value_type must be a Type, not {type(value_type)}") - elif not all(isinstance(dim, int) for dim in shape): - raise TypeError(f"shape must be a sequence of ints, not {type(shape)}") - self.shape = shape - self.value_type = value_type - - def __str__(self): - shape_string = "x".join("?" if dim == -1 else str(dim) for dim in self.shape) - return f"memref<{shape_string}x{self.value_type}>" - - @classmethod - def parse(cls, text: str, aliases: AliasMap = None): - if m := cls._patt.match(text): - dim_strings = m.group(1).split("x")[:-1] - if dim_strings == ["*"]: - raise NotImplementedError(f"Unranked memrefs not currently supported.") - elif ( - any(not s.isdigit() for s in dim_strings if s != "?") - or len(dim_strings) == 0 - ): - raise ValueError(f"{repr(text)} does not have a valid shape.") - shape = [int(dim) if dim.isdigit() else -1 for dim in dim_strings] - value_type = Type.find(m.group(2), aliases=aliases) - return MemrefType(shape, value_type) - - -class TensorType(Type): - _patt = re.compile(r"^tensor<\s*((?:(?:[^,])*x)*)([^, ]+)\s*>$") - - def __init__(self, shape: Sequence[int], value_type: Type): - shape = tuple(shape) - if not isinstance(value_type, Type): - raise TypeError(f"value_type must be a Type, not {type(value_type)}") - elif not all(isinstance(dim, int) for dim in shape): - raise TypeError(f"shape must be a sequence of ints, not {type(shape)}") - self.shape = shape - self.value_type = value_type - - def __str__(self): - shape_string = "x".join("?" if dim == -1 else str(dim) for dim in self.shape) - return f"tensor<{shape_string}x{self.value_type}>" - - @classmethod - def parse(cls, text: str, aliases: AliasMap = None): - if m := cls._patt.match(text): - dim_strings = m.group(1).split("x")[:-1] - if dim_strings == ["*"]: - raise NotImplementedError(f"Unranked tensors not currently supported.") - elif ( - any(not s.isdigit() for s in dim_strings if s != "?") - or len(dim_strings) == 0 - ): - raise ValueError(f"{repr(text)} does not have a valid shape.") - shape = [int(dim) if dim.isdigit() else -1 for dim in dim_strings] - value_type = Type.find(m.group(2), aliases=aliases) - return TensorType(shape, value_type) - - -class SparseTensorType(Type): - _patt = re.compile(r"^tensor<\s*((?:(?:.)*x)*)(.+),\s*(#.+)\s*>$") - - def __init__( - self, shape: Sequence[int], value_type: Type, encoding: "SparseEncodingType" - ): - shape = tuple(shape) - if not isinstance(value_type, Type): - raise TypeError(f"value_type must be a Type, not {type(value_type)}") - elif not isinstance(encoding, SparseEncodingType): - raise TypeError( - f"encoding must be a SparseEncodingType, not {type(encoding)}" - ) - elif not all(isinstance(dim, int) for dim in shape): - raise TypeError(f"shape must be a sequence of ints, not {type(shape)}") - self.shape = shape - self.value_type = value_type - self.encoding = encoding - - def __str__(self): - shape_string = "x".join("?" if dim == -1 else str(dim) for dim in self.shape) - return f"tensor<{shape_string}x{self.value_type}, {self.encoding}>" - - def to_short_string(self): - ret = [] - if self.encoding.rank == 2: - ret.append("matrix") - ret.append("csc" if self.encoding.ordering == [1, 0] else "csr") - elif self.encoding.rank == 1: - ret.append("vector") - else: - raise ValueError(f"Invalid rank: {self.encoding.rank}") - ret.append(str(self.value_type)) - ret.append( - f"p{self.encoding.pointer_bit_width}i{self.encoding.index_bit_width}" - ) - return "_".join(ret) - - @classmethod - def parse(cls, text: str, aliases: AliasMap = None): - if m := cls._patt.match(text): - dim_strings = m.group(1).split("x")[:-1] - if dim_strings == ["*"]: - raise NotImplementedError(f"Unranked tensors not currently supported.") - elif ( - any(not s.isdigit() for s in dim_strings if s != "?") - or len(dim_strings) == 0 - ): - raise ValueError(f"{repr(text)} does not have a valid shape.") - shape = [int(dim) if dim.isdigit() else -1 for dim in dim_strings] - - value_type = Type.find(m.group(2), aliases=aliases) - encoding = Type.find(m.group(3), aliases=aliases) - - return SparseTensorType(shape, value_type, encoding) - - -class SparseEncodingType(Type): - _patt = re.compile( - r"^#sparse_tensor.encoding<\{" - r"\s*,?\s*(?:dimLevelType\s*=\s*\[(?P.+)\])?" - r"\s*,?\s*(?:dimOrdering\s*=\s*affine_map<(?P.+)>)?" - r"\s*,?\s*(?:pointerBitWidth\s*=\s*(?P\d+))?" - r"\s*,?\s*(?:indexBitWidth\s*=\s*(?P\d+))?" - r"\s*,?\s*\}>$" - ) - - def __init__(self, levels, ordering=None, pointer_bit_width=64, index_bit_width=64): - if levels is None: - if ordering is not None: - raise TypeError("Cannot provide ordering without levels") - self.levels = levels - self.ordering = ordering - self.pointer_bit_width = pointer_bit_width - self.index_bit_width = index_bit_width - - @property - def rank(self): - if self.levels is None and self.ordering is None: - return -1 - return len(self.levels) - - def __str__(self): - return self.to_pretty_string(multiline=False) - - def to_pretty_string(self, multiline=True): - internals = [] - if self.levels is not None: - lvl_str = ", ".join(f'"{lvl}"' for lvl in self.levels) - internals.append(f"dimLevelType = [ {lvl_str} ]") - if self.ordering is not None: - lhs = [f"d{i}" for i in range(len(self.levels))] - rhs = [lhs[idx] for idx in self.ordering] - internals.append( - f"dimOrdering = affine_map<({', '.join(lhs)}) -> ({', '.join(rhs)})>" - ) - internals.append(f"pointerBitWidth = {self.pointer_bit_width}") - internals.append(f"indexBitWidth = {self.index_bit_width}") - if multiline: - internals = ",\n ".join(internals) - return f"#sparse_tensor.encoding<{{\n {internals}\n}}>" - else: - return f"#sparse_tensor.encoding<{{ {', '.join(internals)} }}>" - - @classmethod - def parse(cls, text: str, aliases: AliasMap = None): - if m := cls._patt.match(text): - if levels := m["levels"]: - levels = [lvl.strip().strip("\"'") for lvl in levels.split(",")] - if ordering := m["ordering"]: - lhs, rhs = ordering.split("->") - lhs = [x.strip() for x in lhs.strip().strip("()").split(",")] - rhs = [x.strip() for x in rhs.strip().strip("()").split(",")] - ordering = [lhs.index(x) for x in rhs] - pointer_bit_width = int(m["pointer"]) - index_bit_width = int(m["index"]) - return SparseEncodingType( - levels, ordering, pointer_bit_width, index_bit_width - ) - - -class LlvmPtrType(Type): - _patt = re.compile(r"^!llvm\.ptr<\s*(.+)\s*>$") - - def __init__(self, internal_type: Type): - if not isinstance(internal_type, Type): - raise TypeError(f"internal_type must be a Type, not {type(internal_type)}") - self.internal_type = internal_type - - def __str__(self): - return f"!llvm.ptr<{self.internal_type}>" - - @classmethod - def parse(cls, text: str, aliases: AliasMap = None): - if m := cls._patt.match(text): - internal_type = Type.find(m.group(1), aliases=aliases) - return LlvmPtrType(internal_type) - - -class AffineMap(Type): - _patt = re.compile(r"^affine_map<([^>]+)>$") - - def __init__(self, text_within_angle_brackets): - self.text_within_angle_brackets = text_within_angle_brackets - - def __str__(self): - return f"affine_map<{self.text_within_angle_brackets}>" - - @classmethod - def parse(cls, text: str, aliases: AliasMap = None): - if m := cls._patt.match(text): - return AffineMap(m.group(1)) diff --git a/frontends/numpy-scipy/cometpy/MLIRGen/utils.py b/frontends/numpy-scipy/cometpy/MLIRGen/utils.py new file mode 100644 index 00000000..d6c0ab60 --- /dev/null +++ b/frontends/numpy-scipy/cometpy/MLIRGen/utils.py @@ -0,0 +1,250 @@ +from cometpy.MLIRGen.lowering import types +import scipy as sp +import ctypes +from ctypes import * +def create_memref_type(type, rank): + class memref_template(Structure): + _fields_ = [ ('mem_aligned', POINTER(type)), ('mem', POINTER(type)), ('offset', c_int64), ('dims', c_int64 * rank), ('stride', c_int64 * rank)] + + return memref_template + +class memref_i64(Structure): + _fields_ = [ ('mem_aligned', POINTER(c_int64)), ('mem', POINTER(c_int64)), ('offset', c_int64), ('dim', c_int64 * 1), ('stride', c_int64 * 1)] + +class memref_i32(Structure): + _fields_ = [ ('mem_aligned', POINTER(c_int32)), ('mem', POINTER(c_int32)), ('offset', c_int64), ('dim', c_int64 * 1), ('stride', c_int64 * 1)] + +class memref_f64(Structure): + _fields_ = [ ('mem_aligned', POINTER(c_double)), ('mem', POINTER(c_double)), ('offset', c_int64), ('dim', c_int64 * 1), ('stride', c_int64 * 1)] + +class memref_f32(Structure): + _fields_ = [ ('mem_aligned', POINTER(c_float)), ('mem', POINTER(c_float)), ('offset', c_int64), ('dim', c_int64 * 1), ('stride', c_int64 * 1)] + + +# llvm_args += [*expand_memref_ptr(dim_sizes), 0, *expand_memref_ptr(A1pos), 0, *expand_memref_ptr(A2pos), *expand_memref_ptr(A2crd), *expand_memref_ptr(Aval)] + +class output_csr_f64_i64(Structure): + _fields_ = [('dims_sizes', memref_i64), ('insert_1', c_int64), ('A1pos', memref_i64), ('insert_2', c_int64), ('A2pos', memref_i64), ('A2crd', memref_i64), ('Aval', memref_f64)] + +class output_csr_f32_i64(Structure): + _fields_ = [('dims_sizes', memref_i64), ('insert_1', c_int64), ('A1pos', memref_i64), ('insert_2', c_int64), ('A2pos', memref_i64), ('A2crd', memref_i64), ('Aval', memref_f32)] + +class output_coo_f64_i64(Structure): + _fields_ = [('dims_sizes', memref_i64), ('insert_1', c_int64), ('A1pos', memref_i64), ('A1crd', memref_i64), ('insert_2', c_int64), ('A2crd', memref_i64), ('Aval', memref_f64)] + +class output_coo_f32_i64(Structure): + _fields_ = [('dims_sizes', memref_i64), ('insert_1', c_int64), ('A1pos', memref_i64), ('A1crd', memref_i64), ('insert_2', c_int64), ('A2crd', memref_i64), ('Aval', memref_f32)] + +class output_csr_f64_i32(Structure): + _fields_ = [('dims_sizes', memref_i64), ('insert_1', c_int64), ('A1pos', memref_i32), ('insert_2', c_int64), ('A2pos', memref_i32), ('A2crd', memref_i32), ('Aval', memref_f64)] + +class output_csr_f32_i32(Structure): + _fields_ = [('dims_sizes', memref_i64), ('insert_1', c_int64), ('A1pos', memref_i32), ('insert_2', c_int64), ('A2pos', memref_i32), ('A2crd', memref_i32), ('Aval', memref_f32)] + +class output_coo_f64_i32(Structure): + _fields_ = [('dims_sizes', memref_i64), ('insert_1', c_int64), ('A1pos', memref_i32), ('A1crd', memref_i32), ('insert_2', c_int64), ('A2crd', memref_i32), ('Aval', memref_f64)] + +class output_coo_f32_i32(Structure): + _fields_ = [('dims_sizes', memref_i64), ('insert_1', c_int64), ('A1pos', memref_i32), ('A1crd', memref_i32), ('insert_2', c_int64), ('A2crd', memref_i32), ('Aval', memref_f32)] + + + +def get_tensor_type(datatype, shape, format, indices_type): + if format != types.DENSE: + tensor_formats = [] + if format == types.CSR: + tensor_formats.append("d") + tensor_formats.append("unk") + tensor_formats.append("cu") + tensor_formats.append("unk") + elif format == types.COO: + tensor_formats.append("cn") + tensor_formats.append("unk") + tensor_formats.append("s") + tensor_formats.append("unk") + return "!ta.sparse_tensor<{}, {}, {}, {}>".format(datatype, indices_type,"x".join(str(v) for v in shape), ",".join(f for f in tensor_formats)) + else: + return "tensor<{}x{}>".format("x".join(str(v) for v in shape), datatype) + + + +def memref_type_from_dense_ndarray(ndarray): + return types.MemrefType(ndarray.shape, dtype_to_mlir_type(ndarray.dtype)) + +def tensor_type_from_dense_ndarray(ndarray): + return types.TensorType(ndarray.shape, dtype_to_mlir_type(ndarray.dtype)) + +def mlir_type_from_sparse_ndarray(sp_ndarray): + format = None + if sp_ndarray.format == 'csr': + format = types.CSR + return types.TASparseTensorType(sp_ndarray.shape, dtype_to_mlir_type(sp_ndarray.dtype), dtype_to_mlir_type(sp_ndarray.indices.dtype), format) + elif sp_ndarray.format == 'coo': + format = types.COO + return types.TASparseTensorType(sp_ndarray.shape, dtype_to_mlir_type(sp_ndarray.dtype), dtype_to_mlir_type(sp_ndarray.row.dtype), format) + else: + raise Exception("Unsupported format") + +def mlir_type_from_ndarray(A): + if not sp.sparse.issparse(A): + return tensor_type_from_dense_ndarray(A) + else: + return mlir_type_from_sparse_ndarray(A) + +def mlir_type_from_python_type(value): + if isinstance(value, int): + mlir_type = 'index' + elif isinstance(value, float): + mlir_type = 'f64' + else: + raise Exception(f'Unsupported Python type {type(value)}') + return mlir_type + + + + +def dtype_to_mlir_type(dtype): + if dtype == 'int64': + return 'i64' + elif dtype == 'int32': + return 'i32' + elif dtype == 'float32': + return 'f32' + elif dtype == 'float64': + return 'f64' + else : + raise Exception("YBE") + +def mlir_type_to_dtype(mlir_type): + if mlir_type == 'index': + return 'int64' + elif mlir_type == 'i64': + return 'int64' + elif mlir_type == 'i32': + return 'int32' + elif mlir_type == 'f32': + return 'float32' + elif mlir_type == 'f64': + return 'float64' + else : + raise Exception("YBE") + + +def get_format(A): + if not sp.sparse.issparse(A): + return types.DENSE + elif A.format == 'csr': + return types.CSR + elif A.format == 'coo': + return types.COO + elif A.format == 'csc': + return types.CSC + else: + raise RuntimeError('Unsupported sparse format') + + +def format_to_string(format): + if format == types.DENSE: + return "Dense" + elif format == types.CSR: + return "CSR" + elif format == types.COO: + return "COO" + else: + raise RuntimeError('Unsupported sparse format') + + +def python_type_to_ctype(t): + if isinstance(t, int): + ctype = c_int32 + elif isinstance(t, float): + ctype = c_double + return ctype + +def np_array_to_memref(np_array): + ctype = c_int64 + if np_array.dtype == 'int32': + ctype = c_int32 + elif np_array.dtype == 'float32': + ctype = c_float + elif np_array.dtype == 'float64': + ctype = c_double + return np_array.ctypes.data_as(POINTER(ctype)), np_array.ctypes.data_as(POINTER(ctype)), 0, np_array.shape[0], 1 + +def expand_memref_ptr(memref): + return byref(memref), byref(memref), 0, 1, 1 + +def len_dense(vals): + num = 0 + for v in vals: + if not sp.sparse.issparse(v): + num+=1 + return num + + +def memref_from_shaped_type(tensor_type): + ctype = None + shape = tensor_type.shape + if tensor_type.element_type == 'index': + ctype = c_int64 + if(len(shape) == 1): + constructor = memref_i64 + else: + constructor = create_memref_type(ctype,len(shape)) + if tensor_type.element_type == 'i64': + ctype = c_int64 + if(len(shape) == 1): + constructor = memref_i64 + else: + constructor = create_memref_type(ctype,len(shape)) + elif tensor_type.element_type == 'i32': + ctype = c_int32 + if(len(shape) == 1): + constructor = memref_i32 + else: + constructor = create_memref_type(ctype,len(shape)) + elif tensor_type.element_type == 'f64': + ctype = c_double + if(len(shape) == 1): + constructor = memref_f64 + else: + constructor = create_memref_type(ctype,len(shape)) + elif tensor_type.element_type == 'f32': + ctype = c_float + if(len(shape) == 1): + constructor = memref_f32 + else: + constructor = create_memref_type(ctype,len(shape)) + return constructor(ctypes.cast(None, ctypes.POINTER(ctype)), ctypes.cast(None, ctypes.POINTER(ctype)), 0, (c_int64*len(shape))(*shape), (c_int64*len(shape))(*[0]*len(shape))) + +def memref_from_np_array(np_array): + ctype = None + if np_array.dtype == 'int64': + ctype = c_int64 + if(len(np_array.shape) == 1): + constructor = memref_i64 + else: + constructor = create_memref_type(ctype,len(np_array.shape)) + elif np_array.dtype == 'int32': + ctype = c_int32 + if(len(np_array.shape) == 1): + constructor = memref_i32 + else: + constructor = create_memref_type(ctype,len(np_array.shape)) + elif np_array.dtype == 'float32': + ctype = c_float + if(len(np_array.shape) == 1): + constructor = memref_f32 + else: + constructor = create_memref_type(ctype,len(np_array.shape)) + elif np_array.dtype == 'float64': + ctype = c_double + if(len(np_array.shape) == 1): + constructor = memref_f64 + else: + constructor = create_memref_type(ctype,len(np_array.shape)) + if hasattr(np_array, '__cuda_array_interface__'): + ptr = np_array.__cuda_array_interface__['data'][0] + return constructor(ctypes.cast(ptr, ctypes.POINTER(ctype)), ctypes.cast(ptr, ctypes.POINTER(ctype)), 0, (c_int64*len(np_array.shape))(*np_array.shape), (c_int64*len(np_array.shape))(*[s//8 for s in np_array.strides])) + else: + return constructor(np_array.ctypes.data_as(ctypes.POINTER(ctype)), np_array.ctypes.data_as(ctypes.POINTER(ctype)), 0, (c_int64*len(np_array.shape))(*np_array.shape), (c_int64*len(np_array.shape))(*[s//8 for s in np_array.strides])) diff --git a/frontends/numpy-scipy/cometpy/comet.py b/frontends/numpy-scipy/cometpy/comet.py index 3e917976..b13b804f 100644 --- a/frontends/numpy-scipy/cometpy/comet.py +++ b/frontends/numpy-scipy/cometpy/comet.py @@ -29,31 +29,484 @@ # OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. # -from _ast import Attribute, BinOp, Call, Constant -import ast +# from _ast import Attribute, BinOp, Call, Constant +import os +import time +import ast_comments as ast import inspect # import comet +import jinja2 import numpy as np import scipy as sp from cometpy.MLIRGen import lowering -from cometpy.MLIRGen import builders +from cometpy.MLIRGen import ops +from cometpy.MLIRGen.utils import * +from cometpy.MLIRGen import types from cometpy.MLIRGen.types import * #import time -def get_format(A): - if not sp.sparse.issparse(A): - return DENSE - elif A.format == 'csr': - return CSR - elif A.format == 'coo': - return COO - elif A.format == 'csc': - return CSC - else: - raise RuntimeError('Unsupported sparse format') +def lower_einsum_expresion(args, kwargs, visitor): + char_index_to_symbol_index = {} + all_input_indices, res_indices = args[0].value.split('->') + inputs_indices = all_input_indices.split(',') + inputs_symbol_indices = [] + for input_indices in inputs_indices: + input_symbol_indices = [] + for c in input_indices: + if c not in char_index_to_symbol_index: + char_index_to_symbol_index[c] = visitor.build(ops.TensorIndexLabelOp()).results[0] + input_symbol_indices.append(char_index_to_symbol_index[c]) + inputs_symbol_indices.append(input_symbol_indices) + + operands = [visitor.visit(arg) for arg in args[1:len(inputs_indices)+1]] + semiring = None + masktype = "none" + mask = None + if len(inputs_indices) < len(args[1:]): + semiring = args[-1].value + + for kwarg in kwargs: + if kwarg.arg == 'semiring': + semiring = kwarg.value.value + elif kwarg.arg == 'mask_type': + masktype = kwarg.value.value + elif kwarg.arg == 'mask': + mask = visitor.visit(kwarg.value) + + all_same = True + for indices in inputs_indices[1:]: + if indices != inputs_indices[0]: + all_same = False + if len(inputs_indices) == 1 : # Tranpose + res_symbol_indices = [] + + for res_indice in res_indices: + res_symbol_indices.append(char_index_to_symbol_index[res_indice]) + input = operands[0] + input_indices = inputs_symbol_indices[0] + res_format = visitor.sp_mattr_conversions[input.type.format] + res = visitor.build(ops.TensorTransposeOp(input, input_indices, res_symbol_indices, res_format)).results[0] + lhs = res + + elif all_same and inputs_indices[0] == res_indices: # Elementwise multiplication + if semiring: + if semiring == "min": + semiring = "noop_minxy" + elif semiring == "-": + semiring = "noop_minus" + elif semiring == "+": + semiring = "noop_plusxy" + elif semiring == "*": + semiring = "noop_times" + + lhs, lhs_indices = operands[0], inputs_symbol_indices[0] + for rhs in operands[1:]: + res_format = visitor.sp_elw_mult_conversions[lhs.type.format][rhs.type.format] + if semiring: + res = visitor.build(ops.TensorMulOp(lhs, rhs, lhs_indices, res_format, semiring=semiring)).results[0] + else: + res = visitor.build(ops.TensorMulOp(lhs, rhs, lhs_indices, res_format)).results[0] + lhs = res + + else: # Tensor contraction + if semiring: + s1, s2 = semiring.split(",") + if s1 == "+": + semiring = "plusxy_" + elif s1 == "any": + semiring = "any_" + elif s1 == "min": + semiring = "minxy_" + if s2 == "*": + semiring += "times" + elif s2 == "pair": + semiring += "pairxy" + elif s2 == "first": + semiring += "first" + elif s2 == "+": + semiring += "plusxy" + elif s2 == "second": + semiring += "second" + + lhs, lhs_indices = operands[0], inputs_symbol_indices[0] + for (i, (rhs, rhs_indices)) in enumerate(zip(operands[1:], inputs_symbol_indices[1:])): + if i == len(operands) - 2: + res_symbol_indices = [] + for symbol in res_indices: + res_symbol_indices.append(char_index_to_symbol_index[symbol]) + else: + res_symbol_indices = [] + for indice in lhs_indices: + if indice not in rhs_indices: + res_symbol_indices.append(indice) + for indice in rhs_indices: + if indice not in lhs_indices: + res_symbol_indices.append(indice) + + res_format = visitor.sp_matmult_conversions[lhs.type.format][rhs.type.format] + if semiring: + res = visitor.build(ops.TensorMatMultOp(lhs, rhs, lhs_indices, rhs_indices, res_symbol_indices, res_format, 1.0, 0.0, mask, masktype, semiring=semiring)).results[0] + else: + res = visitor.build(ops.TensorMatMultOp(lhs, rhs, lhs_indices, rhs_indices, res_symbol_indices, res_format, 1.0, 0.0, mask, masktype)).results[0] + lhs, lhs_indices = res, res_symbol_indices + + if lhs.type.format != DENSE: + visitor.need_opt_comp_workspace = True + return lhs + # for c in res_indices: + # res_symbol_indices.append(visitor.build(TensorIndexLabelOp())) + + + +class NewAstParser(ast.NodeVisitor): + + identation_default = 2 + def __init__(self,inputs): + self.symbol_table = ops.SymbolTable() + self.inputs = inputs + self.insertion_point = [self] + self.identation = 0 + self.return_type = [] + self.pragma = '' + self.body = [] + self.need_opt_comp_workspace = False + self.build(ops.ModuleOp()) + + + # Output formats when multiply matrices of different formats + self.sp_matmult_conversions = { + CSR: { + CSR : CSR, + COO : CSR, + DENSE : DENSE, + }, + COO: { + CSR : CSR, + COO : CSR, + DENSE : DENSE + }, + DENSE: { + CSR : DENSE, + COO : DENSE, + DENSE : DENSE, + } + } + + self.sp_indices_results = { + 'i64': { + 'i64': 'i64', + 'i32': 'i64', + None: 'i64', + }, + + 'i32': { + 'i64': 'i64', + 'i32': 'i32', + None: 'i32', + } + , + None: { + 'i64': 'i64', + 'i32': 'i32', + None: None, + } + } + + # Output formats when transposing matrices of different formats + # self.sp_mattr_conversions = { + # CSR : CSC, + # COO : COO, + # DENSE : DENSE + # } + + self.sp_mattr_conversions = { + CSR : CSR, + COO : COO, + DENSE : DENSE + } + + # Output formats when elwise mult matrices of different formats + # self.sp_elw_mult_conversions = { + # CSR: { + # CSR : CSR, + # COO : CSR, + # DENSE : COO, + # }, + # COO: { + # CSR : CSR, + # COO : CSR, + # DENSE : COO + # }, + # DENSE: { + # CSR : DENSE, + # COO : DENSE, + # DENSE : DENSE, + # } + # } + self.sp_elw_mult_conversions = { + CSR: { + CSR : CSR, + COO : CSR, + DENSE : DENSE, + }, + COO: { + CSR : CSR, + COO : CSR, + DENSE : COO + }, + DENSE: { + CSR : DENSE, + COO : DENSE, + DENSE : DENSE, + } + } + # Output formats when elwise add or subtract matrices of different formats + # Add and subtract is almost the same as elwise mult with two differences + self.sp_elw_add_sub_conversions = self.sp_elw_mult_conversions + + + # self.sp_elw_add_sub_conversions[CSR][DENSE] = DENSE + # self.sp_elw_add_sub_conversions[COO][DENSE] = DENSE + + self.sp_elw_add_sub_conversions[CSR][DENSE] = CSR + self.sp_elw_add_sub_conversions[COO][DENSE] = COO + def dump(self): + return "\n".join([stmt.dump() for stmt in self.body]) + def build(self, op): + # self.mlir_code.append(" " * self.identation + op.dump()) + self.insertion_point[-1].body.append(op) + if op.startsBody: + self.identation += NewAstParser.identation_default + op.identation = self.identation + self.insertion_point.append(op) + elif op.endsBody: + self.identation -= NewAstParser.identation_default + # self.mlir_code.append(" " * self.identation + "}") + self.insertion_point.pop() + return op + + def visit_FunctionDef(self, node): + args = [] + for arg, val in zip(node.args.args, self.inputs): + if hasattr(val, 'shape'): + new_arg = ops.Symbol(mlir_type_from_ndarray(val)) + else: + new_arg = ops.Symbol(mlir_type_from_python_type(val)) + + self.symbol_table.insert(arg.arg, new_arg) + args.append(new_arg) + funcOp = self.build(ops.FuncOp(node.name, args, [], False)) + # for arg in args: + # self.build(ops.TensorPrintOp(arg)) + for stmt in node.body: + self.visit(stmt) + + + if self.return_type: + funcOp.return_types = self.return_type + else: + self.build(ops.ReturnOp([])) + # self.mlir_code.append("" *self.identation + "}") + + def visit_Comment(self, node): + if node.value.startswith('#pragma'): + self.pragma = node.value + + def visit_Assign(self, node): + self.pragma = '' + rhs = self.visit(node.value) + if isinstance(node.targets[0], ast.Name): + self.symbol_table.insert(node.targets[0].id, rhs) + elif isinstance(node.targets[0], ast.Subscript): + print(node.targets[0].slice) + mem = self.visit(node.targets[0].value) + indices = self.visit(node.targets[0].slice) + if indices: # + if not isinstance(indices, list): + indices = [indices] + if isinstance(mem.type, TensorType): + mem = self.build(ops.ToMemrefOp(mem)) + self.build(ops.StoreOp(rhs, mem.results[0], indices)) + else: + self.build(ops.TensorSetOp(rhs, mem, 0.0)) + self.build(ops.TensorPrintOp(mem)) + # self.mlir_code.append(store.dump()) + # lhs = self.visit(node.targets[0].) + # if indices: + # store = StoreOp(load.results[0], lhs) + + + def visit_BinOp(self, node): + self.pragma = '' + + left = self.visit(node.left) + right = self.visit(node.right) + val = None + res_format = DENSE + if isinstance(node.op, ast.Add): + if isinstance(left.type, ShapedType): + indices = [self.build(ops.TensorIndexLabelOp()).results[0] for d in left.type.shape] + res_format = self.sp_elw_add_sub_conversions[left.type.format][right.type.format] + val = self.build(ops.TensorAddOp(left, right, indices, res_format)) + else: + val = self.build(ops.AddOp(left, right)) + elif isinstance(node.op, ast.Sub): + if isinstance(left.type, ShapedType): + indices = [self.build(ops.TensorIndexLabelOp()).results[0] for d in left.type.shape] + res_format = self.sp_elw_add_sub_conversions[left.type.format][right.type.format] + val = self.build(ops.TensorSubOp(left, right, indices, res_format)) + else: + val = self.build(ops.SubOp(left, right)) + elif isinstance(node.op, ast.Mult): + if isinstance(left.type, ShapedType): + indices = [self.build(ops.TensorIndexLabelOp()).results[0] for d in left.type.shape] + res_format = self.sp_elw_mult_conversions[left.type.format][right.type.format] + val = self.build(ops.TensorMulOp(left, right, indices, res_format)) + else: + val = self.build(ops.MulOp(left, right)) + elif isinstance(node.op, ast.MatMult): + assert(isinstance(left.type, ShapedType) and isinstance(right.type, ShapedType)) + lhs_indices = [self.build(ops.TensorIndexLabelOp()).results[0] for d in left.type.shape] + rhs_indices = [lhs_indices[-1]] + [self.build(ops.TensorIndexLabelOp()).results[0] for d in right.type.shape[1:]] + if len(lhs_indices) ==2 and len(rhs_indices) == 2: + res_indices = [lhs_indices[0], rhs_indices[1]] + elif len(lhs_indices) == 2 and len(rhs_indices) == 1: + res_indices = [lhs_indices[0]] + elif len(lhs_indices) == 1 and len(rhs_indices) == 2: + res_indices = [rhs_indices[1]] + res_format = self.sp_matmult_conversions[left.type.format][right.type.format] + val = self.build(ops.TensorMatMultOp(left, right, lhs_indices, rhs_indices, res_indices, res_format)) + if res_format != DENSE: + self.need_opt_comp_workspace = True + if val: + return val.results[0] + + + + def visit_For(self, node): + pragma = self.pragma + assert(isinstance(node.iter, ast.Call) and node.iter.func.id == 'range') + if len(node.iter.args) == 1: + lb = self.build(ops.ConstantOp(0)) + + lb = lb.results[0] + ub = self.visit(node.iter.args[0]) + step = self.build(ops.ConstantOp(1)) + step = step.results[0] + elif len(node.iter.args) == 2: + lb = self.visit(node.iter.args[0]) + ub = self.visit(node.iter.args[1]) + step = self.build(ops.ConstantOp(1)) + step = step.results[0] + else: + lb = self.visit(node.iter.args[0]) + ub = self.visit(node.iter.args[1]) + step = self.visit(node.iter.args[2]) + if pragma == "#pragma parallel": + forOp = self.build(ops.ForAllOp(lb, ub, step)) + else: + forOp = self.build(ops.ForOp(lb, ub, step)) + self.symbol_table.insert(node.target.id, forOp.iv) + for stmt in node.body: + self.visit(stmt) + if pragma == "#pragma parallel": + self.build(ops.ForallInParallel([])) + else: + self.build(ops.YieldOp([])) + + + + def visit_Subscript(self, node): + self.pragma = '' + + mem = self.visit(node.value) + assert(mem != None) + indices = self.visit(node.slice) + if not isinstance(indices, list): + indices = [indices] + if isinstance(node.value.ctx, ast.Load): + if isinstance(mem.type, TensorType): + mem = self.build(ops.ToMemrefOp(mem)).results[0] + load = self.build(ops.LoadOp(mem, indices)) + return load.results[0] + # elif isinstance(node.value.ctx, ast.Store): + # store = StoreOp(mem, indices) + # self.mlir_code.append(store.dump()) + + def visit_Tuple(self, node): + elts = [self.visit(elt) for elt in node.elts] + return elts + + + + def visit_Call(self, node): + if isinstance(node.func, ast.Attribute): + value = self.visit(node.func.value) + if isinstance(value, ops.Symbol): + values = [value] + else : + assert(value == 'comet') + attr = node.func.attr + elif isinstance(node.func, ast.Name): + values = [] + attr = self.visit(node.func.id) + + if attr == 'sum': + res = self.build(ops.TensorSumOp(*values)).results[0] + return res + elif attr == 'transpose': + indices = [self.build(ops.TensorIndexLabelOp()).results[0] for d in values[0].type.shape] + transpose_indices = [indices[1], indices[0]] + format = None + if isinstance(values[0].type, TensorType) or isinstance(values[0].type, MemrefType): + format = DENSE + elif isinstance(values[0].type, types.TASparseTensorType): + format = values[0].type.format + res = self.build(ops.TensorTransposeOp(values[0], indices, transpose_indices, format)).results[0] + return res + elif attr == 'einsum': + res = lower_einsum_expresion(node.args, node.keywords, self) + return res + elif attr == 'multiply': + indices = [self.build(ops.TensorIndexLabelOp()).results[0] for d in values[0].type.shape] + for arg in node.args: + values.append(self.visit(arg)) + res_format = self.sp_elw_mult_conversions[values[0].type.format][values[1].type.format] + res = self.build(ops.TensorMulOp(values[0], values[1], indices, res_format)).results[0] + return res + + + + def visit_Name(self, node): + self.pragma = '' + + symbol = self.symbol_table.find(node.id) + if symbol: + return symbol + else: + return node.id + + def visit_Constant(self, node): + self.pragma = '' + + val = None + if isinstance(node.value, int): + val = self.build(ops.ConstantOp(node.value)) + elif isinstance(node.value, float): + val = self.build(ops.ConstantOp(node.value)) + if val: + return val.results[0] + + def visit_Return(self, node): + ret = self.visit(node.value) + # self.build(ops.TensorPrintOp(ret)) + if not isinstance(ret, list): + ret = [ret] + retOp = self.build(ops.ReturnOp(ret)) + self.return_type = [ret.type for ret in retOp.results] class NewVisitor(ast.NodeVisitor): @@ -96,6 +549,26 @@ def __init__(self,inputs): } } + self.sp_indices_results = { + 'i64': { + 'i64': 'i64', + 'i32': 'i64', + None: 'i64', + }, + + 'i32': { + 'i64': 'i64', + 'i32': 'i32', + None: 'i32', + } + , + None: { + 'i64': 'i64', + 'i32': 'i32', + None: None, + } + } + # Output formats when transposing matrices of different formats # self.sp_mattr_conversions = { # CSR : CSC, @@ -189,28 +662,30 @@ def get_index_constant(self,v) : def visit_FunctionDef(self, node): for i, arg in enumerate(node.args.args): self.tsymbols[arg.arg] = self.tcurr - format = get_format(self.inputs[i]) + format = ops.get_format(self.inputs[i]) for s in self.inputs[i].shape : + indices_type = None + if format == DENSE : + indices_type = None + elif format == COO: + indices_type = ops.dtype_to_mlir_type(self.inputs[i].coords[0].dtype) + elif format == CSR: + indices_type = ops.dtype_to_mlir_type(self.inputs[i].indices.dtype) + else: + raise Exception("Unexpected sparse matrix type") self.get_next_indexlabel_with_val(s) self.reset_indexlabel_with_val(s) self.tsemantics[self.tcurr] = { + 'indices_type': indices_type, + 'value_type' : ops.dtype_to_mlir_type(self.inputs[i].dtype), 'shape': list(self.inputs[i].shape), + 'in_device': hasattr(self.inputs[i], '__cuda_array_interface__'), 'format': format, 'dimsSSA': [self.get_index_constant(d) for d in self.inputs[i].shape], 'scalar': False, } - self.declarations.append( - { - "type": "T", - "is_input": True, - "arg_num": i, - "format": format, - "shape": self.inputs[i].shape, - "dimsSSA": [self.get_index_constant(d) for d in self.inputs[i].shape], - "id": self.tcurr, - }) self.in_args.append(self.tcurr) self.tcurr += 1 for stmt in node.body: @@ -219,23 +694,74 @@ def visit_FunctionDef(self, node): def visit_Assign(self, node): # We do not support multiple targets currently - if isinstance(node.targets[0], ast.Subscript): + if isinstance(node.targets[0], ast.Subscript): id = node.targets[0].value.id - mask = node.targets[0].slice - self.mask = mask - v = NewVisitor.visit(self, node.value) - self.tsymbols[id] = v else: - # vals = [NewVisitor.visit(self, node.value)] - # for l,v in zip(node.targets, vals): - # self.tsymbols[l.id] = v - id = node.targets[0].id - # if id in self.tsymbols: - # self.no_assign = True + + + if isinstance(node.targets[0], ast.Subscript) and isinstance(node.targets[0].ctx, ast.Store) and not isinstance(node.targets[0].slice, ast.Tuple): + + v = NewVisitor.visit(self, node.value) + # self.tsymbols[id] = v + if not isinstance(node.targets[0].slice, ast.Tuple): + self.ops.append( + { + "indices_type": [self.tsemantics[v]['indices_type']]*2, + "value_type": self.tsemantics[v]['value_type'], + "op_type": "=", + "shapes": [self.tsemantics[v]['shape']]*2, + "lhs": self.tsymbols[id], + "rhs": v, + "beta": True, + }) + if self.tsemantics[self.tsymbols[id]]['format'] == DENSE and self.tsemantics[self.tsymbols[id]]['shape'] != 1: + self.ops.append( + { + "indices_type": [None], + "value_type": self.tsemantics[self.tsymbols[id]]['value_type'], + "op_type": "p", + "shapes": [self.tsemantics[self.tsymbols[id]]["shape"]], + "value_type": self.tsemantics[self.tsymbols[id]]["value_type"], + "operands": [self.tsymbols[id]] + }) + + else: + + if isinstance(node.targets[0], ast.Subscript) and isinstance(node.targets[0].ctx, ast.Store) and (isinstance(node.targets[0].slice, ast.Tuple)) : + mask = node.targets[0].slice + self.mask = mask + print(slice) v = NewVisitor.visit(self, node.value) - self.tsymbols[id] = v + if self.tsemantics[v]['shape'] != 1: + self.declarations.append({ + 'indices_type': self.tsemantics[v]['indices_type'], + "value_type": self.tsemantics[v]['value_type'], + "type": "T", + "is_input": False, + "todo": "l", + "format": self.tsemantics[v]['format'], + "shape": self.tsemantics[v]['shape'], + "dimsSSA": [self.get_index_constant(d) for d in self.tsemantics[v]['shape']], + "id": self.tcurr, + }) + self.tsemantics[self.tcurr] = self.tsemantics[v] + self.ops.append( + { + "indices_type": [self.tsemantics[v]['indices_type']]*2, + "value_type": self.tsemantics[v]['value_type'], + "op_type": "=", + "shapes": [self.tsemantics[v]['shape']]*2, + "lhs": self.tcurr, + "rhs": v, + "beta": True, + }) + self.tsymbols[id] = self.tcurr + self.tcurr += 1 + else: + self.tsymbols[id] = v + @@ -262,6 +788,8 @@ def visit_AugAssign(self, node): if no_assign: self.ops.append( { + "indices_type": [self.tsemantics[v]['indices_type'], self.tsemantics[self.tsymbols[id]]['indices_type']], + "value_type": self.tsemantics[v]['value_type'], "op_type": "=", "beta": beta, "lhs": self.tsymbols[id], @@ -276,6 +804,8 @@ def visit_AugAssign(self, node): if no_assign: self.ops.append( { + "indices_type": [self.tsemantics[v]['indices_type'], self.tsemantics[self.tsymbols[id]]['indices_type']], + "value_type": self.tsemantics[v]['value_type'], "op_type": "=", "beta": beta, "lhs": self.tsymbols[id], @@ -289,6 +819,8 @@ def visit_AugAssign(self, node): if no_assign: self.ops.append( { + "indices_type": [self.tsemantics[v]['indices_type'], self.tsemantics[self.tsymbols[id]]['indices_type']], + "value_type": self.tsemantics[v]['value_type'], "op_type": "=", "beta": beta, "lhs": self.tsymbols[id], @@ -303,8 +835,7 @@ def visit_AugAssign(self, node): res = self.create_binOp(node, [self.tsymbols[id], v], no_assign) self.tsymbols[id] = res - - def visit_Call(self, node: Call) : + def visit_Call(self, node: ast.Call) : obj = None obj = NewVisitor.visit(self,node.func) if obj is not None: @@ -314,13 +845,13 @@ def visit_Call(self, node: Call) : return self.visit_Einsum_Call(node) - def visit_Attribute(self, node: Attribute) : + def visit_Attribute(self, node: ast.Attribute) : return NewVisitor.visit(self, node.value) def visit_Name(self, node: ast.Name): return self.tsymbols[node.id] - def visit_Constant(self, node: Constant) : + def visit_Constant(self, node: ast.Constant) : out_id = self.tcurr self.tsemantics[self.tcurr] = {'shape': [1,], 'format': DENSE, 'scalar': True} self.declarations.append( @@ -350,6 +881,8 @@ def create_binOp(self, node, operands, no_assign): raise "Unexpected operator {}".format(node.op) self.ops.append( { + "indices_type": [op1_sems['indices_type'], op0_sems['indices_type'], None], + "value_type": op0_sems['value_type'], "op_type": "scalar", "op": op, "operands": operands[::-1], @@ -357,10 +890,12 @@ def create_binOp(self, node, operands, no_assign): "out_id": self.tcurr, } ) - self.tsemantics[self.tcurr] = {'shape': [1,], 'format': DENSE, 'scalar': True} + self.tsemantics[self.tcurr] = {"value_type": op_semantics['value_type'], 'indices_type': None, 'shape': [1,], 'format': DENSE, 'scalar': True} self.tcurr += 1 self.declarations.append({ + 'indices_type': None, + "value_type": op_semantics['value_type'], "type": "V", "value": f"{0:e}", "is_input": False, @@ -372,6 +907,8 @@ def create_binOp(self, node, operands, no_assign): }) self.ops.append( { + 'indices_type': [None, None], + "value_type": op_semantics['value_type'], "op_type": "=", "shapes": [[1,]]*2, "lhs": self.tcurr, @@ -384,6 +921,7 @@ def create_binOp(self, node, operands, no_assign): return self.tcurr-1 format = self.sp_elw_add_sub_conversions[op0_sems['format']][op1_sems['format']] + indices_type = self.sp_indices_results[op0_sems['indices_type']][op1_sems['indices_type']] if self.tsemantics[operands[0]]['format'] != DENSE: op_semantics = self.tsemantics[operands[0]] else: @@ -393,6 +931,8 @@ def create_binOp(self, node, operands, no_assign): if isinstance(node.op, ast.Add): self.ops.append( { + 'indices_type': [op0_sems['indices_type'], op1_sems['indices_type'], indices_type], + "value_type": op_semantics['value_type'], "op_type": "+", "shapes": [op_semantics['shape']] * 3, "operands": operands, @@ -404,30 +944,12 @@ def create_binOp(self, node, operands, no_assign): for d in op_semantics['shape']: self.reset_indexlabel_with_val(d) - self.tsemantics[self.tcurr] = {'shape': op_semantics['shape'], 'format': format, 'dimsSSA': [self.get_index_constant(d) for d in op_semantics['shape']], 'scalar': False} - # if not no_assign: - self.tcurr += 1 - self.declarations.append({ - "type": "T", - "is_input": False, - "todo": "l", - "format": format, - "shape": op_semantics['shape'], - "dimsSSA": [self.get_index_constant(d) for d in op_semantics['shape']], - "id": self.tcurr, - }) - self.ops.append( - { - "op_type": "=", - "shapes": [op_semantics['shape']]*2, - "lhs": self.tcurr, - "rhs": self.tcurr-1, - "beta": no_assign, - }) - self.tsemantics[self.tcurr] = self.tsemantics[self.tcurr-1] + self.tsemantics[self.tcurr] = {'value_type': op_semantics['value_type'], 'indices_type': indices_type, 'shape': op_semantics['shape'], 'format': format, 'dimsSSA': [self.get_index_constant(d) for d in op_semantics['shape']], 'scalar': False} elif isinstance(node.op, ast.Sub): self.ops.append( { + 'indices_type': [op0_sems['indices_type'], op1_sems['indices_type'], indices_type], + "value_type": op_semantics['value_type'], "op_type": "-", "shapes": [op_semantics['shape']]* 3, "operands": operands, @@ -438,27 +960,8 @@ def create_binOp(self, node, operands, no_assign): for d in op_semantics['shape']: self.reset_indexlabel_with_val(d) - self.tsemantics[self.tcurr] = {'shape': op_semantics['shape'], 'format': format, 'dimsSSA': [self.get_index_constant(d) for d in op_semantics['shape']], 'scalar': False} - # if not no_assign: - self.tcurr += 1 - self.declarations.append({ - "type": "T", - "is_input": False, - "todo": "l", - "format": format, - "shape": op_semantics['shape'], - "dimsSSA": [self.get_index_constant(d) for d in op_semantics['shape']], - "id": self.tcurr, - }) - self.ops.append( - { - "op_type": "=", - "shapes": [op_semantics['shape']]*2, - "lhs": self.tcurr, - "rhs": self.tcurr-1, - "beta": no_assign, - }) - self.tsemantics[self.tcurr] = self.tsemantics[self.tcurr-1] + self.tsemantics[self.tcurr] = {"value_type": op_semantics['value_type'], 'indices_type': indices_type, 'shape': op_semantics['shape'], 'format': format, 'dimsSSA': [self.get_index_constant(d) for d in op_semantics['shape']], 'scalar': False} + elif isinstance(node.op, ast.Mult): in_ilabels = [self.get_next_indexlabel_with_val(d) for d in op_semantics['shape']] for d in op_semantics['shape']: @@ -466,6 +969,8 @@ def create_binOp(self, node, operands, no_assign): self.ops.append( { + 'indices_type': [op0_sems['indices_type'], op1_sems['indices_type'], indices_type], + "value_type": op_semantics['value_type'], "op_type": "*", "shapes": [op_semantics['shape']] * 3, "op_ilabels": [in_ilabels] * 3, @@ -477,28 +982,8 @@ def create_binOp(self, node, operands, no_assign): format = self.sp_elw_mult_conversions[op0_sems['format']][op1_sems['format']] for d in op_semantics['shape']: self.reset_indexlabel_with_val(d) - self.tsemantics[self.tcurr] = {'shape': op_semantics['shape'], 'format': format, 'dimsSSA': [self.get_index_constant(d) for d in op_semantics['shape']], 'scalar': False} - # if not no_assign: - self.tcurr += 1 - - self.declarations.append({ - "type": "T", - "is_input": False, - "todo": "l", - "format": format, - "shape": op_semantics['shape'], - "dimsSSA": [self.get_index_constant(d) for d in op_semantics['shape']], - "id": self.tcurr, - }) - self.ops.append( - { - "op_type": "=", - "shapes": [op_semantics['shape']]*2, - "lhs": self.tcurr, - "rhs": self.tcurr-1, - "beta": no_assign, - }) - self.tsemantics[self.tcurr] = self.tsemantics[self.tcurr-1] + self.tsemantics[self.tcurr] = {'value_type': op_semantics['value_type'], 'indices_type': indices_type,'shape': op_semantics['shape'], 'format': format, 'dimsSSA': [self.get_index_constant(d) for d in op_semantics['shape']], 'scalar': False} + elif isinstance(node.op, ast.MatMult): # def visit_Bin_Einsum_Call(self, operands, llabels, mask,semiring, no_assing): @@ -515,18 +1000,18 @@ def create_binOp(self, node, operands, no_assign): elif len(op1['shape']) == 1 and len(op2['shape']) == 2: indices = 'j,jk->k' - return self.visit_Bin_Einsum_Call(operands, indices, mask, None, 0, False) + return self.visit_Bin_Einsum_Call(operands, indices, mask, None, 0, False, {}) self.tcurr +=1 return self.tcurr-1 - def visit_BinOp(self, node: BinOp) : + def visit_BinOp(self, node: ast.BinOp) : no_assign = self.no_assign self.no_assign = False operands = [NewVisitor.visit(self, node.left), NewVisitor.visit(self, node.right)] return self.create_binOp(node, operands, no_assign) - def visit_Method_Call(self, node: Call, obj): + def visit_Method_Call(self, node: ast.Call, obj): no_assign = self.no_assign self.no_assign = False # operands = [] @@ -536,7 +1021,7 @@ def visit_Method_Call(self, node: Call, obj): if node.func.attr == "transpose": out_format = self.sp_mattr_conversions[op_semantics['format']] - self.tsemantics[self.tcurr] = {'shape': op_semantics['shape'][::-1], 'format': out_format, 'dimsSSA': [self.get_index_constant(d) for d in op_semantics['shape'][::-1]], 'scalar': False} + self.tsemantics[self.tcurr] = {"value_type": op_semantics['value_type'], 'indices_type': op_semantics['indices_type'], 'shape': op_semantics['shape'][::-1], 'format': out_format, 'dimsSSA': [self.get_index_constant(d) for d in op_semantics['shape'][::-1]], 'scalar': False} in_ilabels = [self.get_next_indexlabel_with_val(d) for d in op_semantics['shape']] for d in op_semantics['shape']: @@ -548,6 +1033,8 @@ def visit_Method_Call(self, node: Call, obj): self.ops.append( { + "indices_type": [op_semantics['indices_type']] * 2, + "value_type": op_semantics['value_type'], "op_type": "t", "shapes": [op_semantics['shape'], op_semantics['shape'][::-1]], "operands": [obj], @@ -555,45 +1042,19 @@ def visit_Method_Call(self, node: Call, obj): "out_id": self.tcurr, "beta": 0, }) - self.tcurr += 1 - # if not no_assign: - self.declarations.append( - { - "type": "T", - "is_input": False, - "todo": "l", - "format": out_format, - "shape": op_semantics['shape'][::-1], - "dimsSSA": [self.get_index_constant(d) for d in op_semantics['shape'][::-1]], - "id": self.tcurr, - }) - self.ops.append( - { - "op_type": "=", - "shapes": [op_semantics['shape'][::-1]]*2, - "lhs": self.tcurr, - "rhs": self.tcurr-1, - "beta": no_assign, - }) - self.tsemantics[self.tcurr] = self.tsemantics[self.tcurr-1] elif node.func.attr == "sum": - self.tsemantics[self.tcurr] = {'shape': 1, 'format': DENSE, 'scalar': True} + self.tsemantics[self.tcurr] = {"value_type": op_semantics['value_type'], 'indices_type': None, 'shape': 1, 'format': DENSE, 'scalar': True} self.ops.append( { + "value_type": op_semantics['value_type'], + 'indices_type': [op_semantics["indices_type"], None], "op_type": "s", "shapes": [op_semantics["shape"], [1,]], "operands": [obj], "out_id": self.tcurr, }) - # if not no_assign: - # self.declarations.append( - # { - # "type": "V", - # "todo": "l", - # "format": DENSE, - # "id": self.tcurr, - # }) + elif node.func.attr == "multiply": op1 = NewVisitor.visit(self, node.args[0]) op1_sems = self.tsemantics[op1] @@ -603,8 +1064,11 @@ def visit_Method_Call(self, node: Call, obj): for d in op_semantics['shape']: self.reset_indexlabel_with_val(d) + indices_type = self.sp_indices_results[op_semantics['indices_type']][op1_sems['indices_type']] self.ops.append( { + "value_type": op_semantics['value_type'], + "indices_type": [op_semantics['indices_type'], op1_sems['indices_type'], indices_type], "op_type": "*", "shapes": [op_semantics['shape']] * 3, "operands": [obj, op1], @@ -615,28 +1079,7 @@ def visit_Method_Call(self, node: Call, obj): }) format = self.sp_elw_mult_conversions[op_semantics['format']][op1_sems['format']] - self.tsemantics[self.tcurr] = {'shape': op_semantics['shape'], 'format': format, 'dimsSSA': [self.get_index_constant(d) for d in op_semantics['shape'][::-1]], 'scalar': False} - # if not no_assign: - self.tcurr += 1 - self.declarations.append( - { - "type": "T", - "is_input": False, - "todo": "l", - "format": format, - "shape": op_semantics['shape'][::-1], - "dimsSSA": [self.get_index_constant(d) for d in op_semantics['shape'][::-1]], - "id": self.tcurr, - }) - self.ops.append( - { - "op_type": "=", - "shapes": [op_semantics['shape'][::-1]] * 2, - "rhs": self.tcurr-1, - "lhs": self.tcurr, - "beta": no_assign, - }) - self.tsemantics[self.tcurr] = self.tsemantics[self.tcurr-1] + self.tsemantics[self.tcurr] = {"value_type": op_semantics['value_type'], "indices_type": indices_type, 'shape': op_semantics['shape'], 'format': format, 'dimsSSA': [self.get_index_constant(d) for d in op_semantics['shape'][::-1]], 'scalar': False} self.tcurr +=1 @@ -645,17 +1088,30 @@ def visit_Method_Call(self, node: Call, obj): def visit_Return(self, node): obj = NewVisitor.visit(self, node.value) self.returns.append(obj) - if self.tsemantics[obj]['format'] == DENSE: + if self.tsemantics[obj]['format'] == DENSE and self.tsemantics[obj]['shape'] != 1: self.in_args.append(obj) - self.ops.append( - { - "op_type": "p", - "shapes": [self.tsemantics[obj]["shape"]], - "value_type": "f64", - "operands": [obj] - }) + self.ops.append( + { + "indices_type": [None], + "value_type": self.tsemantics[obj]['value_type'], + "op_type": "p", + "shapes": [self.tsemantics[obj]["shape"]], + "value_type": self.tsemantics[obj]["value_type"], + "operands": [obj] + }) + if self.tsemantics[obj]['format'] != DENSE or self.tsemantics[obj]['shape'] == 1: + self.ops.append( + + { + "indices_type": [self.tsemantics[obj]['indices_type']], + "value_type": self.tsemantics[obj]['value_type'], + "op_type": "r", + "shapes": [self.tsemantics[obj]["shape"]], + "value_type": self.tsemantics[obj]["value_type"], + "operands": [obj] + }) - def visit_Bin_Einsum_Call(self, operands, llabels, mask,semiring, beta, no_assign): + def visit_Bin_Einsum_Call(self, operands, llabels, mask,semiring, beta, no_assign, lbls_to_ilbls): ops = llabels.split('->')[0].split(',') res = list(llabels.split('->')[1]) @@ -670,12 +1126,12 @@ def visit_Bin_Einsum_Call(self, operands, llabels, mask,semiring, beta, no_assig all_dims.append(self.tsemantics[operands[1]]['shape'][i]) lbls_to_dims = {} - lbls_to_ilbls = {} ilbls_seen = set() for d,l in zip(all_dims, all_lbls): lbls_to_dims[l] = d - lbls_to_ilbls[l] = self.get_next_indexlabel_with_val(d) + if l not in lbls_to_ilbls: + lbls_to_ilbls[l] = self.get_next_indexlabel_with_val(d) for d,l in zip(all_dims, all_lbls): self.reset_indexlabel_with_val(d) @@ -688,6 +1144,7 @@ def visit_Bin_Einsum_Call(self, operands, llabels, mask,semiring, beta, no_assig in1_ilabels = [lbls_to_ilbls[l] for l in list(ops[1])] format = self.tsemantics[operands[0]]['format'] + indices_type = self.tsemantics[operands[0]]['indices_type'] if format != DENSE: self.need_opt_comp_workspace = True @@ -705,11 +1162,15 @@ def visit_Bin_Einsum_Call(self, operands, llabels, mask,semiring, beta, no_assig if len(ops) > 1: for op in operands[1:]: format = self.sp_matmult_conversions[format][self.tsemantics[op]['format']] + indices_type = self.sp_indices_results[indices_type][self.tsemantics[op]['indices_type']] + if format != DENSE: self.need_opt_comp_workspace = True self.ops.append( { + "indices_type": [self.tsemantics[operands[0]]['indices_type'], self.tsemantics[operands[1]]['indices_type'],indices_type], + "value_type": self.tsemantics[operands[0]]['value_type'], "op_type": "c", "operands": operands, "op_ilabels": [in0_ilabels, in1_ilabels, labels], @@ -731,7 +1192,9 @@ def visit_Bin_Einsum_Call(self, operands, llabels, mask,semiring, beta, no_assig self.ops.append( { + "value_type": self.tsemantics[operands[0]]['value_type'], "op_type": "t", + "indices_type": [self.tsemantics[operands[0]]['indices_type'], indices_type], "shapes": [self.tsemantics[operands[0]]['shape'], shape], "operands": operands, "out_id": self.tcurr, @@ -739,34 +1202,12 @@ def visit_Bin_Einsum_Call(self, operands, llabels, mask,semiring, beta, no_assig "op_ilabels": [in_ilabels, out_ilabels] }) - self.tsemantics[self.tcurr] = {'shape': shape, 'format': format, 'dimsSSA': [self.get_index_constant(d) for d in shape], 'scalar': False} - if not no_assign: - self.tcurr += 1 - self.declarations.append( - { - "type": "T", - "is_input": False, - "todo": "l", - "format": format, - "shape": shape, - "dimsSSA": [self.get_index_constant(d) for d in shape], - "id": self.tcurr, - }) - self.ops.append( - { - "op_type": "=", - "shapes": [shape] * 2, - "lhs": self.tcurr, - "rhs": self.tcurr-1, - "beta": beta, - }) - self.tsemantics[self.tcurr] = self.tsemantics[self.tcurr-1] - + self.tsemantics[self.tcurr] = {'value_type': self.tsemantics[operands[0]]['value_type'],"indices_type": indices_type, 'shape': shape, 'format': format, 'dimsSSA': [self.get_index_constant(d) for d in shape], 'scalar': False} self.tcurr += 1 return self.tcurr - 1 - def visit_Einsum_Call(self, node: Call): + def visit_Einsum_Call(self, node: ast.Call): no_assign = self.no_assign self.no_assign = False out_id = self.tcurr @@ -802,6 +1243,7 @@ def visit_Einsum_Call(self, node: Call): format = self.sp_elw_add_sub_conversions[op0_sems['format']][op1_sems['format']] + indices_type = self.sp_indices_results[op0_sems['indices_type']][op1_sems['indices_type']] if self.tsemantics[operands[0]]['format'] != DENSE: op_semantics = self.tsemantics[operands[0]] else: @@ -813,6 +1255,8 @@ def visit_Einsum_Call(self, node: Call): self.ops.append( { + "indices_type": [op0_sems['indices_type'], op1_sems['indices_type'], indices_type], + "value_type": op_semantics['value_type'], "op_type": "*", "shapes": [op_semantics['shape']] * 3, "operands": operands, @@ -824,27 +1268,8 @@ def visit_Einsum_Call(self, node: Call): # ("*", operands, indices+','+indices+'->'+indices, self.tcurr, semiring)) format = self.sp_elw_mult_conversions[op0_sems['format']][op1_sems['format']] - self.tsemantics[self.tcurr] = {'shape': op_semantics['shape'], 'format': format, 'dimsSSA': [self.get_index_constant(d) for d in op_semantics['shape']], 'scalar': False } - self.tcurr += 1 - self.declarations.append( - { - "type": "T", - "is_input": False, - "todo": "l", - "format": format, - "shape": op_semantics['shape'], - "dimsSSA": [self.get_index_constant(d) for d in op_semantics['shape']], - "id": self.tcurr, - }) - self.ops.append( - { - "op_type": "=", - "shapes": [op_semantics['shape']] * 2, - "rhs": self.tcurr-1, - "lhs": self.tcurr, - "beta": no_assign, - }) - self.tsemantics[self.tcurr] = self.tsemantics[self.tcurr-1] + indices_type = self.sp_indices_results[op0_sems['indices_type']][op1_sems['indices_type']] + self.tsemantics[self.tcurr] = {'value_type' : op_semantics['value_type'], 'indices_type': indices_type, 'shape': op_semantics['shape'], 'format': format, 'dimsSSA': [self.get_index_constant(d) for d in op_semantics['shape']], 'scalar': False } self.tcurr += 1 return self.tcurr-1 @@ -853,6 +1278,7 @@ def visit_Einsum_Call(self, node: Call): lops = ops[0] saved_no_assign = no_assign no_assign = True + lbls_to_ilbls = {} for i in range(1, len(operands)): all_lbls = list(lops) for l in list(ops[i]): @@ -879,7 +1305,7 @@ def visit_Einsum_Call(self, node: Call): if i == len(operands) -1 : ret = res no_assign = saved_no_assign - tid = self.visit_Bin_Einsum_Call([loperand,operands[i]], lops+","+ops[i]+"->"+"".join(ret), mask, semiring, saved_no_assign, no_assign) + tid = self.visit_Bin_Einsum_Call([loperand,operands[i]], lops+","+ops[i]+"->"+"".join(ret), mask, semiring, saved_no_assign, no_assign, lbls_to_ilbls) loperand = tid lops = "".join(ret) return tid @@ -888,7 +1314,7 @@ def visit_Einsum_Call(self, node: Call): loperand = operands[0] lops = ops[0] ret = res - tid = self.visit_Bin_Einsum_Call([loperand], lops+"->"+"".join(ret), mask, semiring, no_assign, no_assign) + tid = self.visit_Bin_Einsum_Call([loperand], lops+"->"+"".join(ret), mask, semiring, no_assign, no_assign, {}) lops = "".join(ret) return tid @@ -898,98 +1324,52 @@ def compile(flags, target:str = "cpu", with_jit=True): def innerfunc(func): def wrapper(*pos_args, **kwargs): + # start = time.time() func_str = ast.parse(inspect.getsource(func)) parsed_func = ast.parse(func_str) func_def = parsed_func.body[0] - v = NewVisitor([*pos_args]) - v.visit(parsed_func) - in_types = [] - for arg in v.in_args: - if isinstance(v.tsemantics[arg]['shape'], int): - in_types.append(("%t"+str(arg), "tensor<1xf64>")) + func_name = ':'.join((os.path.abspath(inspect.getfile(func)), func_def.name)) + arg_vals = [*pos_args] + input_types = [None] * len(arg_vals) + for i, arg in enumerate(arg_vals): + if hasattr(arg, 'shape'): + type = mlir_type_from_ndarray(arg) + input_types[i] = f'{type}' else: - in_types.append(("%t"+str(arg), "tensor<{}xf64>".format("x".join(str(d) for d in v.tsemantics[arg]['shape'])))) - irb = builders.MLIRFunctionBuilder( - func_def.name, - input_types=in_types, - return_types=[], - ) - - for i in range(v.currIndexLabel): - irb.add_statement('%i{} = "ta.index_label"() : () -> !ta.indexlabel'.format(i)) - - - dense_tensors = [] - scalars = [] - for dec in v.declarations: + type = mlir_type_from_python_type(arg) + input_types[i] = f'{type}' - if dec["type"] == "T": - dec["value_type"] = "f64" - t = builders.Tensor_Decl_Builder(dec) - if dec["format"] == DENSE: - dense_tensors.append(t) - else: - irb.add_statement(t.build_tensor()) - elif dec["type"] == "C": - irb.add_statement('%d{} = arith.constant {} : index '.format(dec["id"], dec["value"])) - elif dec["type"] == "V": - scalars.append('%t{} = ta.constant dense<{}> : tensor<1xf64> '.format(dec["id"], dec["value"])) - # irb.add_statement('%t{} = ta.constant dense<{}> : tensor<1xf64> '.format(dec["id"], dec["value"])) - - - for t in dense_tensors: - irb.add_statement(t.build_tensor()) - - for t in scalars: - irb.add_statement(t) - - for op in v.ops: - if op["op_type"] == 'c': - op["formats"] = [v.tsemantics[t]['format'] for t in op["operands"]] + [v.tsemantics[op["out_id"]]['format']] - if op["mask"][0] is not None : - op["mask"] = (op["mask"][0], op["mask"][1], v.tsemantics[op["mask"][0]]['shape']) - irb.add_statement(builders.ArithOp_Builder(op).build_op()) + cached_kernel = lowering.cache.find(func_name, input_types) + code = None + new_flags = None + kernel_name = None + return_type = None + + if not cached_kernel: + new_v = NewAstParser([*pos_args]) + new_v.visit(parsed_func) + + new_flags = flags + if new_v.need_opt_comp_workspace: + if new_flags: + new_flags = new_flags + ' --opt-comp-workspace' else: - op["mask"] = (op["mask"][0], op["mask"][1], None) - irb.add_statement(builders.ArithOp_Builder(op).build_op()) - elif op["op_type"] == 'scalar': - irb.add_statement(builders.ScalarOp_Builder(op).build_op()) - elif op["op_type"] == 's': - irb.add_statement(builders.TensorSumBuilder(op).build_op()) - elif op["op_type"] == 'p': - irb.add_statement(builders.PrintBuilder(op).build_op()) - elif op["op_type"] == '*': - op["formats"] = [v.tsemantics[t]['format'] for t in op["operands"]] + [v.tsemantics[op["out_id"]]['format']] - irb.add_statement(builders.ArithOp_Builder(op).build_op()) - elif op["op_type"] == '=': - irb.add_statement(builders.SetOp_Builder(op).build_op()) - else: - op["formats"] = [v.tsemantics[t]['format'] for t in op["operands"]] + [v.tsemantics[op["out_id"]]['format']] - irb.add_statement(builders.ArithOp_Builder(op).build_op()) - irb.add_statement("return") - - outputs = [] - ret = v.tsemantics[v.returns[0]] - format = ret['format'] - if format == DENSE: - outputs.append(np.empty(ret['shape'])) - elif format == CSR: - outputs.append(sp.sparse.csr_array(np.empty(ret['shape']))) - elif format == COO: - outputs.append(sp.sparse.coo_array(np.empty(ret['shape']))) - elif format == CSC: - outputs.append(sp.sparse.csc_array(np.empty(ret['shape']))) - - arg_vals = v.inputs - new_flags = flags - if v.need_opt_comp_workspace: + new_flags = ' --opt-comp-workspace' if new_flags: - new_flags = new_flags + ' --opt-comp-workspace' + kernel_name = func_def.name + new_flags.replace('-','_').replace(' ','').replace('=','_') else: - new_flags = ' --opt-comp-workspace' - code = irb.compile() + kernel_name = func_def.name + # code = irb.compile() + moduleOp = new_v.body[0] + for op in moduleOp.body: + if isinstance(op, ops.FuncOp): + op.func_name = kernel_name + code = new_v.dump() + return_type = new_v.return_type + # end = time.time() + # print(f"Parsing time: {end-start}") # start = time.time() - lowering_result = lowering.lower_dialect_with_jit(code, target, None, new_flags,func_def.name, arg_vals, outputs) + lowering_result = lowering.lower_dialect_with_jit(code, target, None, new_flags, kernel_name, arg_vals, return_type, func_name, input_types, cached_kernel) # end = time.time() # print("Time for JIT", end-start) return lowering_result diff --git a/frontends/numpy-scipy/integration_tests/compound_exps/test_CSR_mult_dTranspose.py b/frontends/numpy-scipy/integration_tests/compound_exps/test_CSR_mult_dTranspose.py deleted file mode 100644 index c95faf31..00000000 --- a/frontends/numpy-scipy/integration_tests/compound_exps/test_CSR_mult_dTranspose.py +++ /dev/null @@ -1,26 +0,0 @@ -import time -import numpy as np -import scipy as sp -from cometpy import comet - -def run_numpy(A,B): - C = A @ B.transpose() - - return C - -@comet.compile(flags=None) -def run_comet_with_jit(A, B): - C = A @ B.transpose() - - return C - -A = sp.sparse.csr_array(sp.io.mmread("../../../../integration_test/data/test_rank2.mtx")) -B = np.full([5, A.shape[1]], 3.2, dtype=float) -C = np.full([A.shape[0], 5], 0.0, dtype=float) -expected_result = run_numpy(A,B) -result_with_jit = run_comet_with_jit(A,B) -if sp.sparse.issparse(expected_result): - expected_result = expected_result.todense() -if sp.sparse.issparse(result_with_jit): - result_with_jit = result_with_jit.todense() -np.testing.assert_almost_equal(result_with_jit, expected_result) diff --git a/frontends/numpy-scipy/integration_tests/compound_exps/test_CSR_mult_spTranspose_CSR.py b/frontends/numpy-scipy/integration_tests/compound_exps/test_CSR_mult_spTranspose_CSR.py deleted file mode 100644 index 0f5a4399..00000000 --- a/frontends/numpy-scipy/integration_tests/compound_exps/test_CSR_mult_spTranspose_CSR.py +++ /dev/null @@ -1,24 +0,0 @@ -import time -import numpy as np -import scipy as sp -from cometpy import comet - -def run_numpy(A,B): - C = A @ B.transpose() - - return C - -@comet.compile(flags=None) -def run_comet_with_jit(A,B): - C = A @ B.transpose() - - return C - -A = sp.sparse.csr_array(sp.io.mmread("../../../../integration_test/data/test_rank2.mtx")) -B = sp.sparse.csr_array(sp.io.mmread("../../../../integration_test/data/test_rank2.mtx")) -expected_result = run_numpy(A,B) -result_with_jit = run_comet_with_jit(A,B) -if sp.sparse.issparse(expected_result): - expected_result = expected_result.todense() - result_with_jit = result_with_jit.todense() -np.testing.assert_almost_equal(result_with_jit, expected_result) diff --git a/frontends/numpy-scipy/integration_tests/compound_exps/test_Dense_chain_mult.py b/frontends/numpy-scipy/integration_tests/compound_exps/test_Dense_chain_mult.py index 43bf07c4..52b6dcbe 100644 --- a/frontends/numpy-scipy/integration_tests/compound_exps/test_Dense_chain_mult.py +++ b/frontends/numpy-scipy/integration_tests/compound_exps/test_Dense_chain_mult.py @@ -15,15 +15,13 @@ def numpy_chain_multiply(A,B,C,D): return res +def test_Dense_chain_mult(): + A = np.full([2,4], 1.0) + B = np.full([4,4], 1.0) + C = np.full([4,4], 1.0) + D = np.full([4,6], 1.0) + Dn = numpy_chain_multiply(A,B,C,D) + Dc = comet_chain_multiply(A,B,C,D) - -A = np.full([2,4], 1.0) -B = np.full([4,4], 1.0) -C = np.full([4,4], 1.0) -D = np.full([4,6], 1.0) - -Dn = numpy_chain_multiply(A,B,C,D) -Dc = comet_chain_multiply(A,B,C,D) - -np.testing.assert_almost_equal(Dn,Dc) \ No newline at end of file + np.testing.assert_almost_equal(Dn,Dc) \ No newline at end of file diff --git a/frontends/numpy-scipy/integration_tests/compound_exps/test_Dense_eltwise_dTranspose.py b/frontends/numpy-scipy/integration_tests/compound_exps/test_Dense_eltwise_dTranspose.py index 1661a431..56ec0fda 100644 --- a/frontends/numpy-scipy/integration_tests/compound_exps/test_Dense_eltwise_dTranspose.py +++ b/frontends/numpy-scipy/integration_tests/compound_exps/test_Dense_eltwise_dTranspose.py @@ -14,12 +14,13 @@ def run_comet_with_jit(A,B): return C -A = np.full([4, 4], 3.2, dtype=float) -B = np.full([4, 4], 2.0, dtype=float) -C = np.full([4, 4], 0.0, dtype=float) -expected_result = run_numpy(A, B) -result_with_jit = run_comet_with_jit(A, B) -if sp.sparse.issparse(expected_result): - expected_result = expected_result.todense() - result_with_jit = result_with_jit.todense() -np.testing.assert_almost_equal(result_with_jit, expected_result) +def test_Dense_eltwise_dTranspose(): + A = np.full([4, 4], 3.2, dtype=float) + B = np.full([4, 4], 2.0, dtype=float) + C = np.full([4, 4], 0.0, dtype=float) + expected_result = run_numpy(A, B) + result_with_jit = run_comet_with_jit(A, B) + if sp.sparse.issparse(expected_result): + expected_result = expected_result.todense() + result_with_jit = result_with_jit.todense() + np.testing.assert_almost_equal(result_with_jit, expected_result) diff --git a/frontends/numpy-scipy/integration_tests/compound_exps/test_Dense_eltwise_sTranspose_CSR.py b/frontends/numpy-scipy/integration_tests/compound_exps/test_Dense_eltwise_sTranspose_CSR.py deleted file mode 100644 index 5516271b..00000000 --- a/frontends/numpy-scipy/integration_tests/compound_exps/test_Dense_eltwise_sTranspose_CSR.py +++ /dev/null @@ -1,25 +0,0 @@ -import time -import numpy as np -import scipy as sp -from cometpy import comet - -def run_numpy(A,B): - C = A * B.transpose() - - return C - -@comet.compile(flags=None) -def run_comet_with_jit(A,B): - C = A * B.transpose() - - return C - -B = sp.sparse.csr_array(sp.io.mmread("../../../../integration_test/data/test_rank2.mtx")) -A = np.full([B.shape[1], B.shape[0]], 3.2, dtype=float) -expected_result = run_numpy(A,B) -result_with_jit = run_comet_with_jit(A,B) -if sp.sparse.issparse(expected_result): - expected_result = expected_result.todense() -if sp.sparse.issparse(result_with_jit): - result_with_jit = result_with_jit.todense() -np.testing.assert_almost_equal(result_with_jit, expected_result) diff --git a/frontends/numpy-scipy/integration_tests/compound_exps/test_Dense_mult_dTranspose.py b/frontends/numpy-scipy/integration_tests/compound_exps/test_Dense_mult_dTranspose.py index 2b63f321..75833abc 100644 --- a/frontends/numpy-scipy/integration_tests/compound_exps/test_Dense_mult_dTranspose.py +++ b/frontends/numpy-scipy/integration_tests/compound_exps/test_Dense_mult_dTranspose.py @@ -14,12 +14,13 @@ def run_comet_with_jit(A,B): return C -B = np.full([5, 5], 3.2, dtype=float) -A = np.full([5, 5], 2.3, dtype=float) -C = np.full([5, 5], 0.0, dtype=float) -expected_result = run_numpy(A,B) -result_with_jit = run_comet_with_jit(A,B) -if sp.sparse.issparse(expected_result): - expected_result = expected_result.todense() - result_with_jit = result_with_jit.todense() -np.testing.assert_almost_equal(result_with_jit, expected_result) +def test_Dense_mult_dTranspose(): + B = np.full([5, 5], 3.2, dtype=float) + A = np.full([5, 5], 2.3, dtype=float) + C = np.full([5, 5], 0.0, dtype=float) + expected_result = run_numpy(A,B) + result_with_jit = run_comet_with_jit(A,B) + if sp.sparse.issparse(expected_result): + expected_result = expected_result.todense() + result_with_jit = result_with_jit.todense() + np.testing.assert_almost_equal(result_with_jit, expected_result) diff --git a/frontends/numpy-scipy/integration_tests/compound_exps/test_Dense_mult_spTranspose_CSR.py b/frontends/numpy-scipy/integration_tests/compound_exps/test_Dense_mult_spTranspose_CSR.py deleted file mode 100644 index f6947b19..00000000 --- a/frontends/numpy-scipy/integration_tests/compound_exps/test_Dense_mult_spTranspose_CSR.py +++ /dev/null @@ -1,25 +0,0 @@ -import time -import numpy as np -import scipy as sp -from cometpy import comet - -def run_numpy(A,B): - C = A @ B.transpose() - - return C - -@comet.compile(flags=None) -def run_comet_with_jit(A,B): - C = A @ B.transpose() - - return C - -B = sp.sparse.csr_array(sp.io.mmread("../../../../integration_test/data/test_rank2.mtx")) -A = np.full([5, B.shape[1]], 2.3, dtype=float) -C = np.full([5, B.shape[0]], 0.0, dtype=float) -expected_result = run_numpy(A,B) -result_with_jit = run_comet_with_jit(A,B) -if sp.sparse.issparse(expected_result): - expected_result = expected_result.todense() - result_with_jit = result_with_jit.todense() -np.testing.assert_almost_equal(result_with_jit, expected_result) diff --git a/frontends/numpy-scipy/integration_tests/compound_exps/test_dTranspose_eltwise_CSR.py b/frontends/numpy-scipy/integration_tests/compound_exps/test_dTranspose_eltwise_CSR.py index eaa5ce91..1db73d56 100644 --- a/frontends/numpy-scipy/integration_tests/compound_exps/test_dTranspose_eltwise_CSR.py +++ b/frontends/numpy-scipy/integration_tests/compound_exps/test_dTranspose_eltwise_CSR.py @@ -4,22 +4,23 @@ from cometpy import comet def run_numpy(A, B): - C = A * B.transpose() + C = A.transpose() * B return C @comet.compile(flags=None) def run_comet_with_jit(A, B): - C = A * B.transpose() + C = A.transpose() * B return C -B = sp.sparse.csr_array(sp.io.mmread("../../../../integration_test/data/test_rank2.mtx")) -A = np.full([B.shape[1], B.shape[0]], 3.2, dtype=float) -expected_result = run_numpy(A, B) -result_with_jit = run_comet_with_jit(A, B) -if sp.sparse.issparse(expected_result): - expected_result = expected_result.todense() -if sp.sparse.issparse(result_with_jit): - result_with_jit = result_with_jit.todense() -np.testing.assert_almost_equal(result_with_jit, expected_result) +def test_dTranspose_eltwise_CSR(data_rank2_path): + B = sp.sparse.csr_array(sp.io.mmread(data_rank2_path)) + A = np.full([B.shape[1], B.shape[0]], 3.2, dtype=float) + expected_result = run_numpy(A, B) + result_with_jit = run_comet_with_jit(A, B) + if sp.sparse.issparse(expected_result): + expected_result = expected_result.todense() + if sp.sparse.issparse(result_with_jit): + result_with_jit = result_with_jit.todense() + np.testing.assert_almost_equal(result_with_jit, expected_result) diff --git a/frontends/numpy-scipy/integration_tests/compound_exps/test_dTranspose_eltwise_Dense.py b/frontends/numpy-scipy/integration_tests/compound_exps/test_dTranspose_eltwise_Dense.py index c24c8763..3319257b 100644 --- a/frontends/numpy-scipy/integration_tests/compound_exps/test_dTranspose_eltwise_Dense.py +++ b/frontends/numpy-scipy/integration_tests/compound_exps/test_dTranspose_eltwise_Dense.py @@ -14,12 +14,13 @@ def run_comet_with_jit(A, B): return C -A = np.full([4, 4], 3.2, dtype=float) -B = np.full([4, 4], 2.0, dtype=float) -C = np.full([4, 4], 0.0, dtype=float) -expected_result = run_numpy(A, B) -result_with_jit = run_comet_with_jit(A, B) -if sp.sparse.issparse(expected_result): - expected_result = expected_result.todense() - result_with_jit = result_with_jit.todense() -np.testing.assert_almost_equal(result_with_jit, expected_result) +def test_dTranspose_eltwise_Dense(): + A = np.full([4, 4], 3.2, dtype=float) + B = np.full([4, 4], 2.0, dtype=float) + C = np.full([4, 4], 0.0, dtype=float) + expected_result = run_numpy(A, B) + result_with_jit = run_comet_with_jit(A, B) + if sp.sparse.issparse(expected_result): + expected_result = expected_result.todense() + result_with_jit = result_with_jit.todense() + np.testing.assert_almost_equal(result_with_jit, expected_result) diff --git a/frontends/numpy-scipy/integration_tests/compound_exps/test_dTranspose_mult_CSR.py b/frontends/numpy-scipy/integration_tests/compound_exps/test_dTranspose_mult_CSR.py index 4031fa59..37aa54c5 100644 --- a/frontends/numpy-scipy/integration_tests/compound_exps/test_dTranspose_mult_CSR.py +++ b/frontends/numpy-scipy/integration_tests/compound_exps/test_dTranspose_mult_CSR.py @@ -14,12 +14,13 @@ def run_comet_with_jit(A, B): return C -B = sp.sparse.csr_array(sp.io.mmread("../../../../integration_test/data/test_rank2.mtx")) -A = np.full([B.shape[0], 4], 1.7, dtype=float) -C = np.full([4, B.shape[1]], 0.0, dtype=float) -expected_result = run_numpy(A, B) -result_with_jit = run_comet_with_jit(A, B) -if sp.sparse.issparse(expected_result): - expected_result = expected_result.todense() - result_with_jit = result_with_jit.todense() -np.testing.assert_almost_equal(result_with_jit, expected_result) +def test_dTranspose_mult_CSR(data_rank2_path): + B = sp.sparse.csr_array(sp.io.mmread(data_rank2_path)) + A = np.full([B.shape[0], 4], 1.7, dtype=float) + C = np.full([4, B.shape[1]], 0.0, dtype=float) + expected_result = run_numpy(A, B) + result_with_jit = run_comet_with_jit(A, B) + if sp.sparse.issparse(expected_result): + expected_result = expected_result.todense() + result_with_jit = result_with_jit.todense() + np.testing.assert_almost_equal(result_with_jit, expected_result) diff --git a/frontends/numpy-scipy/integration_tests/compound_exps/test_dTranspose_mult_Dense.py b/frontends/numpy-scipy/integration_tests/compound_exps/test_dTranspose_mult_Dense.py index 9779d597..67aeaa22 100644 --- a/frontends/numpy-scipy/integration_tests/compound_exps/test_dTranspose_mult_Dense.py +++ b/frontends/numpy-scipy/integration_tests/compound_exps/test_dTranspose_mult_Dense.py @@ -14,12 +14,13 @@ def run_comet_with_jit(A, B): return C -A = np.full([4, 4], 3.2, dtype=float) -B = np.full([4, 4], 1.0, dtype=float) -C = np.full([4, 4], 0.0, dtype=float) -expected_result = run_numpy(A, B) -result_with_jit = run_comet_with_jit(A, B) -if sp.sparse.issparse(expected_result): - expected_result = expected_result.todense() - result_with_jit = result_with_jit.todense() -np.testing.assert_almost_equal(result_with_jit, expected_result) +def test_dTranspose_mult_Dense(): + A = np.full([4, 4], 3.2, dtype=float) + B = np.full([4, 4], 1.0, dtype=float) + C = np.full([4, 4], 0.0, dtype=float) + expected_result = run_numpy(A, B) + result_with_jit = run_comet_with_jit(A, B) + if sp.sparse.issparse(expected_result): + expected_result = expected_result.todense() + result_with_jit = result_with_jit.todense() + np.testing.assert_almost_equal(result_with_jit, expected_result) diff --git a/frontends/numpy-scipy/integration_tests/compound_exps/test_spTranspose_CSR_eltwise_CSR.py b/frontends/numpy-scipy/integration_tests/compound_exps/test_spTranspose_CSR_eltwise_CSR.py deleted file mode 100644 index 7d646c68..00000000 --- a/frontends/numpy-scipy/integration_tests/compound_exps/test_spTranspose_CSR_eltwise_CSR.py +++ /dev/null @@ -1,24 +0,0 @@ -import time -import numpy as np -import scipy as sp -from cometpy import comet - -def run_numpy(A,B): - C = A.transpose() * B - - return C - -@comet.compile(flags=None) -def run_comet_with_jit(A,B): - C = A.transpose() * B - - return C - -A = sp.sparse.csr_array(sp.io.mmread("../../../../integration_test/data/test_rank2.mtx")) -B = sp.sparse.csr_array(sp.io.mmread("../../../../integration_test/data/test_rank2.mtx")) -expected_result = run_numpy(A,B) -result_with_jit = run_comet_with_jit(A,B) -if sp.sparse.issparse(expected_result): - expected_result = expected_result.todense() - result_with_jit = result_with_jit.todense() -np.testing.assert_almost_equal(result_with_jit, expected_result) diff --git a/frontends/numpy-scipy/integration_tests/compound_exps/test_spTranspose_CSR_eltwise_Dense.py b/frontends/numpy-scipy/integration_tests/compound_exps/test_spTranspose_CSR_eltwise_Dense.py deleted file mode 100644 index 15034220..00000000 --- a/frontends/numpy-scipy/integration_tests/compound_exps/test_spTranspose_CSR_eltwise_Dense.py +++ /dev/null @@ -1,25 +0,0 @@ -import time -import numpy as np -import scipy as sp -from cometpy import comet - -def run_numpy(A,B): - C = A.transpose() * B - - return C - -@comet.compile(flags=None) -def run_comet_with_jit(A,B): - C = A.transpose() * B - - return C - -A = sp.sparse.csr_array(sp.io.mmread("../../../../integration_test/data/test_rank2.mtx")) -B = np.full([A.shape[1], A.shape[0]], 2.3, dtype=float) -expected_result = run_numpy(A,B) -result_with_jit = run_comet_with_jit(A,B) -if sp.sparse.issparse(expected_result): - expected_result = expected_result.todense() -if sp.sparse.issparse(result_with_jit): - result_with_jit = result_with_jit.todense() -np.testing.assert_almost_equal(result_with_jit, expected_result) diff --git a/frontends/numpy-scipy/integration_tests/compound_exps/test_spTranspose_CSR_mult_CSR.py b/frontends/numpy-scipy/integration_tests/compound_exps/test_spTranspose_CSR_mult_CSR.py deleted file mode 100644 index fe29bfa8..00000000 --- a/frontends/numpy-scipy/integration_tests/compound_exps/test_spTranspose_CSR_mult_CSR.py +++ /dev/null @@ -1,24 +0,0 @@ -import time -import numpy as np -import scipy as sp -from cometpy import comet - -def run_numpy(A, B): - C = A.transpose() @ B - - return C - -@comet.compile(flags=None) -def run_comet_with_jit(A, B): - C = A.transpose() @ B - - return C - -A = sp.sparse.csr_array(sp.io.mmread("../../../../integration_test/data/test_rank2.mtx")) -B = sp.sparse.csr_array(sp.io.mmread("../../../../integration_test/data/test_rank2.mtx")) -expected_result = run_numpy(A, B) -result_with_jit = run_comet_with_jit(A, B) -if sp.sparse.issparse(expected_result): - expected_result = expected_result.todense() - result_with_jit = result_with_jit.todense() -np.testing.assert_almost_equal(result_with_jit, expected_result) diff --git a/frontends/numpy-scipy/integration_tests/conftest.py b/frontends/numpy-scipy/integration_tests/conftest.py new file mode 100644 index 00000000..5ca3f92a --- /dev/null +++ b/frontends/numpy-scipy/integration_tests/conftest.py @@ -0,0 +1,17 @@ +import pytest +import os +import cometpy.cfg + +@pytest.fixture +def data_rank2_path(): + return os.path.join(os.path.dirname(__file__), 'data', 'test_rank2.mtx') + +@pytest.fixture +def data_rank2_transpose_path(): + return os.path.join(os.path.dirname(__file__), 'data', 'test_rank2_transpose.mtx') + +@pytest.fixture +def data_tc_path(): + return os.path.join(os.path.dirname(__file__), 'data', 'tc.mtx') + +gpu = pytest.mark.skipif(not cometpy.cfg.gpu_target_enabled, reason="GPU support not enabled") \ No newline at end of file diff --git a/integration_test/data/tc.mtx b/frontends/numpy-scipy/integration_tests/data/tc.mtx similarity index 100% rename from integration_test/data/tc.mtx rename to frontends/numpy-scipy/integration_tests/data/tc.mtx diff --git a/integration_test/data/test_8x6.mtx b/frontends/numpy-scipy/integration_tests/data/test_8x6.mtx similarity index 100% rename from integration_test/data/test_8x6.mtx rename to frontends/numpy-scipy/integration_tests/data/test_8x6.mtx diff --git a/integration_test/data/test_rank2.mtx b/frontends/numpy-scipy/integration_tests/data/test_rank2.mtx similarity index 100% rename from integration_test/data/test_rank2.mtx rename to frontends/numpy-scipy/integration_tests/data/test_rank2.mtx diff --git a/integration_test/data/test_rank2_denser.mtx b/frontends/numpy-scipy/integration_tests/data/test_rank2_denser.mtx similarity index 100% rename from integration_test/data/test_rank2_denser.mtx rename to frontends/numpy-scipy/integration_tests/data/test_rank2_denser.mtx diff --git a/integration_test/data/test_rank2_small.mtx b/frontends/numpy-scipy/integration_tests/data/test_rank2_small.mtx similarity index 100% rename from integration_test/data/test_rank2_small.mtx rename to frontends/numpy-scipy/integration_tests/data/test_rank2_small.mtx diff --git a/integration_test/data/test_rank2_transpose.mtx b/frontends/numpy-scipy/integration_tests/data/test_rank2_transpose.mtx similarity index 100% rename from integration_test/data/test_rank2_transpose.mtx rename to frontends/numpy-scipy/integration_tests/data/test_rank2_transpose.mtx diff --git a/integration_test/data/test_rank3.tns b/frontends/numpy-scipy/integration_tests/data/test_rank3.tns similarity index 100% rename from integration_test/data/test_rank3.tns rename to frontends/numpy-scipy/integration_tests/data/test_rank3.tns diff --git a/integration_test/data/test_rank8.tns b/frontends/numpy-scipy/integration_tests/data/test_rank8.tns similarity index 100% rename from integration_test/data/test_rank8.tns rename to frontends/numpy-scipy/integration_tests/data/test_rank8.tns diff --git a/integration_test/data/wide.mtx b/frontends/numpy-scipy/integration_tests/data/wide.mtx similarity index 100% rename from integration_test/data/wide.mtx rename to frontends/numpy-scipy/integration_tests/data/wide.mtx diff --git a/frontends/numpy-scipy/integration_tests/kernels/test_ccsd_t1_21_loops.py b/frontends/numpy-scipy/integration_tests/kernels/test_ccsd_t1_21_loops.py index 9cd3e7a5..afa8f8e0 100644 --- a/frontends/numpy-scipy/integration_tests/kernels/test_ccsd_t1_21_loops.py +++ b/frontends/numpy-scipy/integration_tests/kernels/test_ccsd_t1_21_loops.py @@ -14,12 +14,13 @@ def run_comet_with_jit(v,t2): return i0 -v = np.full([2, 2, 4, 4], 2.3, dtype=float) -t2 = np.full([4, 4, 2, 4], 3.4, dtype=float) -i0 = np.full([2, 4], 0.0, dtype=float) -expected_result = run_numpy(v,t2) -result_with_jit = run_comet_with_jit(v,t2) -if sp.sparse.issparse(expected_result): - expected_result = expected_result.todense() - result_with_jit = result_with_jit.todense() -np.testing.assert_almost_equal(result_with_jit, expected_result) +def test_ccsd_t1_21_loops(): + v = np.full([2, 2, 4, 4], 2.3, dtype=float) + t2 = np.full([4, 4, 2, 4], 3.4, dtype=float) + i0 = np.full([2, 4], 0.0, dtype=float) + expected_result = run_numpy(v,t2) + result_with_jit = run_comet_with_jit(v,t2) + if sp.sparse.issparse(expected_result): + expected_result = expected_result.todense() + result_with_jit = result_with_jit.todense() + np.testing.assert_almost_equal(result_with_jit, expected_result) diff --git a/frontends/numpy-scipy/integration_tests/kernels/test_ccsd_t1_3_loops.py b/frontends/numpy-scipy/integration_tests/kernels/test_ccsd_t1_3_loops.py index 6e16c8bb..9f76b4a5 100644 --- a/frontends/numpy-scipy/integration_tests/kernels/test_ccsd_t1_3_loops.py +++ b/frontends/numpy-scipy/integration_tests/kernels/test_ccsd_t1_3_loops.py @@ -14,13 +14,14 @@ def run_comet_with_jit(f,t1): return i0 -f = np.full([2, 4], 2.3, dtype=float) -t1 = np.full([2, 2], 3.4, dtype=float) -i0 = np.full([2, 4], 0.0, dtype=float) -expected_result = run_numpy(f,t1) -result_with_jit = run_comet_with_jit(f,t1) -if sp.sparse.issparse(expected_result): - expected_result = expected_result.todense() -if sp.sparse.issparse(result_with_jit): - result_with_jit = result_with_jit.todense() -np.testing.assert_almost_equal(result_with_jit, expected_result) +def test_ccsd_t1_3_loops(): + f = np.full([2, 4], 2.3, dtype=float) + t1 = np.full([2, 2], 3.4, dtype=float) + i0 = np.full([2, 4], 0.0, dtype=float) + expected_result = run_numpy(f,t1) + result_with_jit = run_comet_with_jit(f,t1) + if sp.sparse.issparse(expected_result): + expected_result = expected_result.todense() + if sp.sparse.issparse(result_with_jit): + result_with_jit = result_with_jit.todense() + np.testing.assert_almost_equal(result_with_jit, expected_result) diff --git a/frontends/numpy-scipy/integration_tests/kernels/test_ccsd_t1_4_loops.py b/frontends/numpy-scipy/integration_tests/kernels/test_ccsd_t1_4_loops.py index f7401022..5d1476fa 100644 --- a/frontends/numpy-scipy/integration_tests/kernels/test_ccsd_t1_4_loops.py +++ b/frontends/numpy-scipy/integration_tests/kernels/test_ccsd_t1_4_loops.py @@ -14,12 +14,13 @@ def run_comet_with_jit(v,t1): return i0 -v = np.full([2, 2, 4, 4], 2.3, dtype=float) -t1 = np.full([4, 2], 3.4, dtype=float) -i0 = np.full([2, 4], 0.0, dtype=float) -expected_result = run_numpy(v,t1) -result_with_jit = run_comet_with_jit(v,t1) -if sp.sparse.issparse(expected_result): - expected_result = expected_result.todense() - result_with_jit = result_with_jit.todense() -np.testing.assert_almost_equal(result_with_jit, expected_result) +def test_ccsd_t1_4_loops(): + v = np.full([2, 2, 4, 4], 2.3, dtype=float) + t1 = np.full([4, 2], 3.4, dtype=float) + i0 = np.full([2, 4], 0.0, dtype=float) + expected_result = run_numpy(v,t1) + result_with_jit = run_comet_with_jit(v,t1) + if sp.sparse.issparse(expected_result): + expected_result = expected_result.todense() + result_with_jit = result_with_jit.todense() + np.testing.assert_almost_equal(result_with_jit, expected_result) diff --git a/frontends/numpy-scipy/integration_tests/kernels/test_gnn.py b/frontends/numpy-scipy/integration_tests/kernels/test_gnn.py index f2d02d9d..24aacbe1 100644 --- a/frontends/numpy-scipy/integration_tests/kernels/test_gnn.py +++ b/frontends/numpy-scipy/integration_tests/kernels/test_gnn.py @@ -15,13 +15,14 @@ def run_comet_with_jit(B,C,D): return A -B = sp.sparse.csr_array(sp.io.mmread("../../../../integration_test/data/test_rank2.mtx")) -C = np.full([B.shape[0], 4], 1.2, dtype=float) -D = np.full([4, 4], 3.4, dtype=float) -expected_result = run_numpy(B,C,D) -result_with_jit = run_comet_with_jit(B,C,D) -if sp.sparse.issparse(expected_result): - expected_result = expected_result.todense() -if sp.sparse.issparse(result_with_jit): - result_with_jit = result_with_jit.todense() -np.testing.assert_almost_equal(result_with_jit, expected_result) +def test_gnn(data_rank2_path): + B = sp.sparse.csr_array(sp.io.mmread(data_rank2_path)) + C = np.full([B.shape[0], 4], 1.2, dtype=float) + D = np.full([4, 4], 3.4, dtype=float) + expected_result = run_numpy(B,C,D) + result_with_jit = run_comet_with_jit(B,C,D) + if sp.sparse.issparse(expected_result): + expected_result = expected_result.todense() + if sp.sparse.issparse(result_with_jit): + result_with_jit = result_with_jit.todense() + np.testing.assert_almost_equal(result_with_jit, expected_result) diff --git a/frontends/numpy-scipy/integration_tests/kernels/test_triangleCount_SandiaLL.py b/frontends/numpy-scipy/integration_tests/kernels/test_triangleCount_SandiaLL.py index 66d79f8f..7f4218dd 100644 --- a/frontends/numpy-scipy/integration_tests/kernels/test_triangleCount_SandiaLL.py +++ b/frontends/numpy-scipy/integration_tests/kernels/test_triangleCount_SandiaLL.py @@ -1,23 +1,26 @@ import time import numpy as np import scipy as sp +import pytest from cometpy import comet -def run_numpy(L0,L1,L2): - C = ((L1 @ L2) * L0).sum() +def run_numpy(L0): + C = ((L0 @ L0) * L0).sum() return C @comet.compile(flags=None) -def run_comet_with_jit(L0,L1,L2): - C = ((L1 @ L2) * L0).sum() +def run_comet_with_jit(L0): + C = ((L0 @ L0) * L0).sum() return C -A = sp.sparse.csr_array(sp.io.mmread("../../../../integration_test/data/tc.mtx")) -L0 = sp.sparse.csr_array(sp.sparse.tril(A, format='csr')) -expected_result = run_numpy(L0,L0,L0) -result_with_jit = run_comet_with_jit(L0,L0,L0) -if sp.sparse.issparse(expected_result): - expected_result = expected_result.todense() -if sp.sparse.issparse(result_with_jit): - result_with_jit = result_with_jit.todense() -np.testing.assert_almost_equal(result_with_jit, expected_result) +@pytest.mark.skip(reason="Triangle counting currently fails") +def test_triangleCount_SandiaLL(data_tc_path): + A = sp.sparse.csr_array(sp.io.mmread(data_tc_path)) + L0 = sp.sparse.csr_array(sp.sparse.tril(A, format='csr')) + expected_result = run_numpy(L0) + result_with_jit = run_comet_with_jit(L0) + if sp.sparse.issparse(expected_result): + expected_result = expected_result.todense() + if sp.sparse.issparse(result_with_jit): + result_with_jit = result_with_jit.todense() + np.testing.assert_almost_equal(result_with_jit, expected_result) diff --git a/frontends/numpy-scipy/integration_tests/kernels/test_triangleCount_SandiaLLmask.py b/frontends/numpy-scipy/integration_tests/kernels/test_triangleCount_SandiaLLmask.py index 4c3ea665..16862250 100644 --- a/frontends/numpy-scipy/integration_tests/kernels/test_triangleCount_SandiaLLmask.py +++ b/frontends/numpy-scipy/integration_tests/kernels/test_triangleCount_SandiaLLmask.py @@ -10,19 +10,18 @@ def run_numpy(L0,L1,L2): @comet.compile(flags=None) def run_comet_with_jit(L0,L1,L2): - C[L0,"push"] = L1 @ L2 # Performs masking. Currently, only works on a single matmul operation - #or C = comet.einsum('ij,jk->ik', L1,L2, mask=L0, mask_type="push") - #or C[L0,"push"] = comet.einsum('ij,jk->ik', L1,L2) + C = comet.einsum('ij,jk->ik', L1,L2, mask=L0, mask_type="push") D = C.sum() return D -A = sp.sparse.csr_array(sp.io.mmread("../../../../integration_test/data/tc.mtx")) -L0 = sp.sparse.csr_array(sp.sparse.tril(A, format='csr')) -expected_result = run_numpy(L0,L0,L0) -result_with_jit = run_comet_with_jit(L0,L0,L0) -if sp.sparse.issparse(expected_result): - expected_result = expected_result.todense() -if sp.sparse.issparse(result_with_jit): - result_with_jit = result_with_jit.todense() -np.testing.assert_almost_equal(result_with_jit, expected_result) +def test_triangleCount_SandiaLLmask(data_tc_path): + A = sp.sparse.csr_array(sp.io.mmread(data_tc_path)) + L0 = sp.sparse.csr_array(sp.sparse.tril(A, format='csr')) + expected_result = run_numpy(L0,L0,L0) + result_with_jit = run_comet_with_jit(L0,L0,L0) + if sp.sparse.issparse(expected_result): + expected_result = expected_result.todense() + if sp.sparse.issparse(result_with_jit): + result_with_jit = result_with_jit.todense() + np.testing.assert_almost_equal(result_with_jit, expected_result) diff --git a/frontends/numpy-scipy/integration_tests/not_supported/test_CSR_mult_dTranspose.py b/frontends/numpy-scipy/integration_tests/not_supported/test_CSR_mult_dTranspose.py new file mode 100644 index 00000000..642ccc3b --- /dev/null +++ b/frontends/numpy-scipy/integration_tests/not_supported/test_CSR_mult_dTranspose.py @@ -0,0 +1,27 @@ +import time +import numpy as np +import scipy as sp +from cometpy import comet + +def run_numpy(A,B): + C = A @ B.transpose() + + return C + +@comet.compile(flags=None) +def run_comet_with_jit(A, B): + C = A @ B.transpose() + + return C + +def test_CSR_mult_dTranspose(data_rank2_path): + A = sp.sparse.csr_array(sp.io.mmread(data_rank2_path)) + B = np.full([5, A.shape[1]], 3.2, dtype=float) + C = np.full([A.shape[0], 5], 0.0, dtype=float) + expected_result = run_numpy(A,B) + result_with_jit = run_comet_with_jit(A,B) + if sp.sparse.issparse(expected_result): + expected_result = expected_result.todense() + if sp.sparse.issparse(result_with_jit): + result_with_jit = result_with_jit.todense() + np.testing.assert_almost_equal(result_with_jit, expected_result) diff --git a/frontends/numpy-scipy/integration_tests/not_supported/test_CSR_mult_spTranspose_CSR.py b/frontends/numpy-scipy/integration_tests/not_supported/test_CSR_mult_spTranspose_CSR.py new file mode 100644 index 00000000..c9ad52ff --- /dev/null +++ b/frontends/numpy-scipy/integration_tests/not_supported/test_CSR_mult_spTranspose_CSR.py @@ -0,0 +1,25 @@ +import time +import numpy as np +import scipy as sp +from cometpy import comet + +def run_numpy(A,B): + C = A @ B.transpose() + + return C + +@comet.compile(flags=None) +def run_comet_with_jit(A,B): + C = A @ B.transpose() + + return C + +def test_CSR_mult_spTranspose_CSR(data_rank2_path): + A = sp.sparse.csr_array(sp.io.mmread(data_rank2_path)) + B = sp.sparse.csr_array(sp.io.mmread(data_rank2_path)) + expected_result = run_numpy(A,B) + result_with_jit = run_comet_with_jit(A,B) + if sp.sparse.issparse(expected_result): + expected_result = expected_result.todense() + result_with_jit = result_with_jit.todense() + np.testing.assert_almost_equal(result_with_jit, expected_result) diff --git a/frontends/numpy-scipy/integration_tests/not_supported/test_Dense_eltwise_sTranspose_CSR.py b/frontends/numpy-scipy/integration_tests/not_supported/test_Dense_eltwise_sTranspose_CSR.py new file mode 100644 index 00000000..7615f96c --- /dev/null +++ b/frontends/numpy-scipy/integration_tests/not_supported/test_Dense_eltwise_sTranspose_CSR.py @@ -0,0 +1,26 @@ +import time +import numpy as np +import scipy as sp +from cometpy import comet + +def run_numpy(A,B): + C = A * B.transpose() + + return C + +@comet.compile(flags=None) +def run_comet_with_jit(A,B): + C = A * B.transpose() + + return C + +def test_Dense_eltwise_sTranspose_CSR(data_rank2_path): + B = sp.sparse.csr_array(sp.io.mmread(data_rank2_path)) + A = np.full([B.shape[1], B.shape[0]], 3.2, dtype=float) + expected_result = run_numpy(A,B) + result_with_jit = run_comet_with_jit(A,B) + if sp.sparse.issparse(expected_result): + expected_result = expected_result.todense() + if sp.sparse.issparse(result_with_jit): + result_with_jit = result_with_jit.todense() + np.testing.assert_almost_equal(result_with_jit, expected_result) diff --git a/frontends/numpy-scipy/integration_tests/not_supported/test_Dense_mult_spTranspose_CSR.py b/frontends/numpy-scipy/integration_tests/not_supported/test_Dense_mult_spTranspose_CSR.py new file mode 100644 index 00000000..1491bd1c --- /dev/null +++ b/frontends/numpy-scipy/integration_tests/not_supported/test_Dense_mult_spTranspose_CSR.py @@ -0,0 +1,26 @@ +import time +import numpy as np +import scipy as sp +from cometpy import comet + +def run_numpy(A,B): + C = A @ B.transpose() + + return C + +@comet.compile(flags=None) +def run_comet_with_jit(A,B): + C = A @ B.transpose() + + return C + +def test_Dense_mult_spTranspose_CSR(data_rank2_path): + B = sp.sparse.csr_array(sp.io.mmread(data_rank2_path)) + A = np.full([5, B.shape[1]], 2.3, dtype=float) + C = np.full([5, B.shape[0]], 0.0, dtype=float) + expected_result = run_numpy(A,B) + result_with_jit = run_comet_with_jit(A,B) + if sp.sparse.issparse(expected_result): + expected_result = expected_result.todense() + result_with_jit = result_with_jit.todense() + np.testing.assert_almost_equal(result_with_jit, expected_result) diff --git a/frontends/numpy-scipy/integration_tests/not_supported/test_eltwise_CSRxCSR_oCSR_sameSpPattern.py b/frontends/numpy-scipy/integration_tests/not_supported/test_eltwise_CSRxCSR_oCSR_sameSpPattern.py new file mode 100644 index 00000000..ed72180a --- /dev/null +++ b/frontends/numpy-scipy/integration_tests/not_supported/test_eltwise_CSRxCSR_oCSR_sameSpPattern.py @@ -0,0 +1,26 @@ +import time +import numpy as np +import scipy as sp +from cometpy import comet + +def run_numpy(A,B): + C = A * B + + return C + +@comet.compile(flags=None) +def run_comet_with_jit(A,B): + C = A * B + + return C + +def test_eltwise_CSRxCSR_oCSR_sameSpPattern(data_rank2_path): + A = sp.sparse.csr_array(sp.io.mmread(data_rank2_path)) + B = sp.sparse.csr_array(sp.io.mmread(data_rank2_path)) + expected_result = run_numpy(A,B) + result_with_jit = run_comet_with_jit(A,B) + if sp.sparse.issparse(expected_result): + expected_result = expected_result.todense() + if sp.sparse.issparse(result_with_jit): + result_with_jit = result_with_jit.todense() + np.testing.assert_almost_equal(result_with_jit, expected_result) diff --git a/frontends/numpy-scipy/integration_tests/not_supported/test_eltwise_add_CSRxCSR_oCSR_sameSpPattern.py b/frontends/numpy-scipy/integration_tests/not_supported/test_eltwise_add_CSRxCSR_oCSR_sameSpPattern.py new file mode 100644 index 00000000..667c2d46 --- /dev/null +++ b/frontends/numpy-scipy/integration_tests/not_supported/test_eltwise_add_CSRxCSR_oCSR_sameSpPattern.py @@ -0,0 +1,25 @@ +import time +import numpy as np +import scipy as sp +from cometpy import comet + +def run_numpy(A,B): + C = A + B + + return C + +@comet.compile(flags=None) +def run_comet_with_jit(A,B): + C = A + B + + return C + +def test_eltwise_add_CSRxCSR_oCSR_sameSpPattern(data_rank2_path, data_rank2_transpose_path): + A = sp.sparse.csr_array(sp.io.mmread(data_rank2_path)) + B = sp.sparse.csr_array(sp.io.mmread(data_rank2_transpose_path)) + expected_result = run_numpy(A,B) + result_with_jit = run_comet_with_jit(A,B) + if sp.sparse.issparse(expected_result): + expected_result = expected_result.todense() + result_with_jit = result_with_jit.todense() + np.testing.assert_almost_equal(result_with_jit, expected_result) diff --git a/frontends/numpy-scipy/integration_tests/not_supported/test_eltwise_subtract_CSRxCSR_oCSR_sameSpPattern.py b/frontends/numpy-scipy/integration_tests/not_supported/test_eltwise_subtract_CSRxCSR_oCSR_sameSpPattern.py new file mode 100644 index 00000000..274fc083 --- /dev/null +++ b/frontends/numpy-scipy/integration_tests/not_supported/test_eltwise_subtract_CSRxCSR_oCSR_sameSpPattern.py @@ -0,0 +1,25 @@ +import time +import numpy as np +import scipy as sp +from cometpy import comet + +def run_numpy(A,B): + C = A - B + + return C + +@comet.compile(flags=None) +def run_comet_with_jit(A,B): + C = A - B + + return C + +def test_eltwise_subtract_CSRxCSR_oCSR_sameSpPattern(data_rank2_path, data_rank2_transpose_path): + A = sp.sparse.csr_array(sp.io.mmread(data_rank2_path)) + B = sp.sparse.csr_array(sp.io.mmread(data_rank2_transpose_path)) + expected_result = run_numpy(A,B) + result_with_jit = run_comet_with_jit(A,B) + if sp.sparse.issparse(expected_result): + expected_result = expected_result.todense() + result_with_jit = result_with_jit.todense() + np.testing.assert_almost_equal(result_with_jit, expected_result) diff --git a/frontends/numpy-scipy/integration_tests/not_supported/test_gpu_sum_dense_matrix.py b/frontends/numpy-scipy/integration_tests/not_supported/test_gpu_sum_dense_matrix.py new file mode 100644 index 00000000..1a593297 --- /dev/null +++ b/frontends/numpy-scipy/integration_tests/not_supported/test_gpu_sum_dense_matrix.py @@ -0,0 +1,26 @@ +import time +import numpy as np +import scipy as sp +from cometpy import comet +from conftest import gpu + +def run_numpy(A): + var = A.sum() + + return var + +@comet.compile(flags=None, target="gpu") +def run_comet_with_jit(A): + var = A.sum() + + return var + +@gpu +def test_gpu_sum_dense_matrix(): + A = np.full([4, 4], 3.7, dtype=float) + expected_result = run_numpy(A) + result_with_jit = run_comet_with_jit(A) + if sp.sparse.issparse(expected_result): + expected_result = expected_result.todense() + result_with_jit = result_with_jit.todense() + np.testing.assert_almost_equal(result_with_jit, expected_result) diff --git a/frontends/numpy-scipy/integration_tests/not_supported/test_gpu_transpose_dense_matrix.py b/frontends/numpy-scipy/integration_tests/not_supported/test_gpu_transpose_dense_matrix.py new file mode 100644 index 00000000..5a880483 --- /dev/null +++ b/frontends/numpy-scipy/integration_tests/not_supported/test_gpu_transpose_dense_matrix.py @@ -0,0 +1,26 @@ +import time +import numpy as np +import scipy as sp +from cometpy import comet +from conftest import gpu + +def run_numpy(A): + B = A.transpose() + + return B + +@comet.compile(flags=None, target="gpu") +def run_comet_with_jit(A): + B = A.transpose() + + return B + +@gpu +def test_gpu_transpose_dense_matrix(): + A = np.full([4, 4], 3.2, dtype=float) + expected_result = run_numpy(A) + result_with_jit = run_comet_with_jit(A) + if sp.sparse.issparse(expected_result): + expected_result = expected_result.todense() + result_with_jit = result_with_jit.todense() + np.testing.assert_almost_equal(result_with_jit, expected_result) diff --git a/frontends/numpy-scipy/integration_tests/not_supported/test_spTranspose_CSR_eltwise_CSR.py b/frontends/numpy-scipy/integration_tests/not_supported/test_spTranspose_CSR_eltwise_CSR.py new file mode 100644 index 00000000..51f0e271 --- /dev/null +++ b/frontends/numpy-scipy/integration_tests/not_supported/test_spTranspose_CSR_eltwise_CSR.py @@ -0,0 +1,25 @@ +import time +import numpy as np +import scipy as sp +from cometpy import comet + +def run_numpy(A,B): + C = A.transpose() * B + + return C + +@comet.compile(flags=None) +def run_comet_with_jit(A,B): + C = A.transpose() * B + + return C + +def test_spTranspose_CSR_eltwise_CSR(data_rank2_path): + A = sp.sparse.csr_array(sp.io.mmread(data_rank2_path)) + B = sp.sparse.csr_array(sp.io.mmread(data_rank2_path)) + expected_result = run_numpy(A,B) + result_with_jit = run_comet_with_jit(A,B) + if sp.sparse.issparse(expected_result): + expected_result = expected_result.todense() + result_with_jit = result_with_jit.todense() + np.testing.assert_almost_equal(result_with_jit, expected_result) diff --git a/frontends/numpy-scipy/integration_tests/not_supported/test_spTranspose_CSR_eltwise_Dense.py b/frontends/numpy-scipy/integration_tests/not_supported/test_spTranspose_CSR_eltwise_Dense.py new file mode 100644 index 00000000..648fdc88 --- /dev/null +++ b/frontends/numpy-scipy/integration_tests/not_supported/test_spTranspose_CSR_eltwise_Dense.py @@ -0,0 +1,26 @@ +import time +import numpy as np +import scipy as sp +from cometpy import comet + +def run_numpy(A,B): + C = A.transpose() * B + + return C + +@comet.compile(flags=None) +def run_comet_with_jit(A,B): + C = A.transpose() * B + + return C + +def test_spTranspose_CSR_eltwise_Dense(data_rank2_path): + A = sp.sparse.csr_array(sp.io.mmread(data_rank2_path)) + B = np.full([A.shape[1], A.shape[0]], 2.3, dtype=float) + expected_result = run_numpy(A,B) + result_with_jit = run_comet_with_jit(A,B) + if sp.sparse.issparse(expected_result): + expected_result = expected_result.todense() + if sp.sparse.issparse(result_with_jit): + result_with_jit = result_with_jit.todense() + np.testing.assert_almost_equal(result_with_jit, expected_result) diff --git a/frontends/numpy-scipy/integration_tests/not_supported/test_spTranspose_CSR_mult_CSR.py b/frontends/numpy-scipy/integration_tests/not_supported/test_spTranspose_CSR_mult_CSR.py new file mode 100644 index 00000000..3df7d738 --- /dev/null +++ b/frontends/numpy-scipy/integration_tests/not_supported/test_spTranspose_CSR_mult_CSR.py @@ -0,0 +1,25 @@ +import time +import numpy as np +import scipy as sp +from cometpy import comet + +def run_numpy(A, B): + C = A.transpose() @ B + + return C + +@comet.compile(flags=None) +def run_comet_with_jit(A, B): + C = A.transpose() @ B + + return C + +def test_spTranspose_CSR_mult_CSR(data_rank2_path): + A = sp.sparse.csr_array(sp.io.mmread(data_rank2_path)) + B = sp.sparse.csr_array(sp.io.mmread(data_rank2_path)) + expected_result = run_numpy(A, B) + result_with_jit = run_comet_with_jit(A, B) + if sp.sparse.issparse(expected_result): + expected_result = expected_result.todense() + result_with_jit = result_with_jit.todense() + np.testing.assert_almost_equal(result_with_jit, expected_result) diff --git a/frontends/numpy-scipy/integration_tests/not_supported/test_transpose_COO_matrix.py b/frontends/numpy-scipy/integration_tests/not_supported/test_transpose_COO_matrix.py new file mode 100644 index 00000000..b7ed436e --- /dev/null +++ b/frontends/numpy-scipy/integration_tests/not_supported/test_transpose_COO_matrix.py @@ -0,0 +1,24 @@ +import time +import numpy as np +import scipy as sp +from cometpy import comet + +def run_numpy(A): + B = A.transpose() + + return B + +@comet.compile(flags=None) +def run_comet_with_jit(A): + B = A.transpose() + + return B + +def test_transpose_COO_matrix(data_rank2_path): + A = sp.sparse.coo_array(sp.io.mmread(data_rank2_path)) + expected_result = run_numpy(A) + result_with_jit = run_comet_with_jit(A) + if sp.sparse.issparse(expected_result): + expected_result = expected_result.todense() + result_with_jit = result_with_jit.todense() + np.testing.assert_almost_equal(result_with_jit, expected_result) diff --git a/frontends/numpy-scipy/integration_tests/not_supported/test_transpose_CSR_matrix.py b/frontends/numpy-scipy/integration_tests/not_supported/test_transpose_CSR_matrix.py new file mode 100644 index 00000000..d7f94b68 --- /dev/null +++ b/frontends/numpy-scipy/integration_tests/not_supported/test_transpose_CSR_matrix.py @@ -0,0 +1,24 @@ +import time +import numpy as np +import scipy as sp +from cometpy import comet + +def run_numpy(A): + B = A.transpose() + + return B + +@comet.compile(flags=None) +def run_comet_with_jit(A): + B = A.transpose() + + return B + +def test_transpose_CSR_matrix(data_rank2_path): + A = sp.sparse.csr_array(sp.io.mmread(data_rank2_path)) + expected_result = run_numpy(A) + result_with_jit = run_comet_with_jit(A) + if sp.sparse.issparse(expected_result): + expected_result = expected_result.todense() + result_with_jit = result_with_jit.todense() + np.testing.assert_almost_equal(result_with_jit, expected_result) diff --git a/frontends/numpy-scipy/integration_tests/not_supported/test_transpose_dense_matrix.py b/frontends/numpy-scipy/integration_tests/not_supported/test_transpose_dense_matrix.py new file mode 100644 index 00000000..8ae20199 --- /dev/null +++ b/frontends/numpy-scipy/integration_tests/not_supported/test_transpose_dense_matrix.py @@ -0,0 +1,24 @@ +import time +import numpy as np +import scipy as sp +from cometpy import comet + +def run_numpy(A): + B = A.transpose() + + return B + +@comet.compile(flags=None) +def run_comet_with_jit(A): + B = A.transpose() + + return B + +def test_transpose_dense_matrix(): + A = np.full([4, 4], 3.2, dtype=float) + expected_result = run_numpy(A) + result_with_jit = run_comet_with_jit(A) + if sp.sparse.issparse(expected_result): + expected_result = expected_result.todense() + result_with_jit = result_with_jit.todense() + np.testing.assert_almost_equal(result_with_jit, expected_result) diff --git a/frontends/numpy-scipy/integration_tests/not_supported/test_transpose_dense_tensor.py b/frontends/numpy-scipy/integration_tests/not_supported/test_transpose_dense_tensor.py new file mode 100644 index 00000000..ce3c2b08 --- /dev/null +++ b/frontends/numpy-scipy/integration_tests/not_supported/test_transpose_dense_tensor.py @@ -0,0 +1,24 @@ +import time +import numpy as np +import scipy as sp +from cometpy import comet + +def run_numpy(A): + B = np.einsum('ijk->kij', A) + + return B + +@comet.compile(flags=None) +def run_comet_with_jit(A): + B = comet.einsum('ijk->kij', A) + + return B + +def test_transpose_dense_tensor(): + A = np.full([4, 4, 4], 3.7, dtype=float) + expected_result = run_numpy(A) + result_with_jit = run_comet_with_jit(A) + if sp.sparse.issparse(expected_result): + expected_result = expected_result.todense() + result_with_jit = result_with_jit.todense() + np.testing.assert_almost_equal(result_with_jit, expected_result) diff --git a/frontends/numpy-scipy/integration_tests/numpy_integration.py b/frontends/numpy-scipy/integration_tests/numpy_integration.py deleted file mode 100644 index 97ff6f78..00000000 --- a/frontends/numpy-scipy/integration_tests/numpy_integration.py +++ /dev/null @@ -1,72 +0,0 @@ - -import glob - -import subprocess -import multiprocessing -import os -import sys -import cometpy.cfg - - -def run_test_case(test_file): - print("Running", test_file, end=" ") - p = subprocess.run( 'cd '+'/'.join(test_file.split('/')[:-1])+' && python3 '+test_file.split('/')[-1], stdout=subprocess.PIPE, stderr=subprocess.PIPE, shell=True) - if(p.returncode != 0): - ret = "FAILED" - # print(ret) - # print("=======================================================") - # print(test_file) - # failed_tests = failed_tests + 1 - # list_failed_tests.append(test_file) - else: - ret = "PASSED" - print(ret) - - return (test_file, ret, p.stderr.decode()) - - -if __name__ == '__main__': - categories = ['ops', 'opts', 'kernels', 'compound_exps', 'semiring'] - files = [] - - for c in categories: - files = files + glob.glob("./"+c+"/test_*.py") - if cometpy.cfg.gpu_target_enabled: - files = files + glob.glob("./"+c+"/gpu/test_*.py") - - - print("\nFound" , len(files), "test cases") - - print("Running the tests with up to {} cores......".format(os.cpu_count())) - - failed_tests = 0 - list_failed_tests = [] - - with multiprocessing.Pool() as p: - results = p.map(run_test_case, files) - for res in results: - if(res[1] == "FAILED"): - list_failed_tests.append((res[0],res[2])) - - print("Passed = ", len(files) - len(list_failed_tests)) - - print("Failed = " , len(list_failed_tests)) - - if(list_failed_tests): - print("The following tests failed:") - - for failed_test in list_failed_tests: - print(' ' , failed_test[0]) - - if len(sys.argv) == 2: - if sys.argv[1] == '-v': - print() - print("Error messages of failed tests:") - for failed_test in list_failed_tests: - print(' ' , failed_test[0]) - print('='*40) - print(failed_test[1]) - print('*'*40) - - if len(list_failed_tests) > 0: - exit(127) diff --git a/frontends/numpy-scipy/integration_tests/ops/gpu/test_eltwise_add_dense_matrix.py b/frontends/numpy-scipy/integration_tests/ops/gpu/test_eltwise_add_dense_matrix.py deleted file mode 100644 index 10d138a3..00000000 --- a/frontends/numpy-scipy/integration_tests/ops/gpu/test_eltwise_add_dense_matrix.py +++ /dev/null @@ -1,25 +0,0 @@ -import time -import numpy as np -import scipy as sp -from cometpy import comet - -def run_numpy(A,B): - C = A+B - - return C - -@comet.compile(flags=None, target="gpu") -def run_comet_with_jit(A,B): - C = A+B - - return C - -A = np.full([4, 4], 2.2, dtype=float) -B = np.full([4, 4], 3.4, dtype=float) -C = np.full([4, 4], 0.0, dtype=float) -expected_result = run_numpy(A,B) -result_with_jit = run_comet_with_jit(A,B) -if sp.sparse.issparse(expected_result): - expected_result = expected_result.todense() - result_with_jit = result_with_jit.todense() -np.testing.assert_almost_equal(result_with_jit, expected_result) diff --git a/frontends/numpy-scipy/integration_tests/ops/gpu/test_eltwise_mult_DensexDense_oDense.py b/frontends/numpy-scipy/integration_tests/ops/gpu/test_eltwise_mult_DensexDense_oDense.py deleted file mode 100644 index 9dadc145..00000000 --- a/frontends/numpy-scipy/integration_tests/ops/gpu/test_eltwise_mult_DensexDense_oDense.py +++ /dev/null @@ -1,25 +0,0 @@ -import time -import numpy as np -import scipy as sp -from cometpy import comet - -def run_numpy(A,B): - C = A * B - - return C - -@comet.compile(flags=None, target="gpu") -def run_comet_with_jit(A,B): - C = A * B - - return C - -A = np.full([4, 4], 2.7, dtype=float) -B = np.full([4, 4], 3.2, dtype=float) -C = np.full([4, 4], 0.0, dtype=float) -expected_result = run_numpy(A,B) -result_with_jit = run_comet_with_jit(A,B) -if sp.sparse.issparse(expected_result): - expected_result = expected_result.todense() - result_with_jit = result_with_jit.todense() -np.testing.assert_almost_equal(result_with_jit, expected_result) diff --git a/frontends/numpy-scipy/integration_tests/ops/gpu/test_eltwise_subtract_dense_matrix.py b/frontends/numpy-scipy/integration_tests/ops/gpu/test_eltwise_subtract_dense_matrix.py deleted file mode 100644 index ca1d6150..00000000 --- a/frontends/numpy-scipy/integration_tests/ops/gpu/test_eltwise_subtract_dense_matrix.py +++ /dev/null @@ -1,25 +0,0 @@ -import time -import numpy as np -import scipy as sp -from cometpy import comet - -def run_numpy(A,B): - C = A - B - - return C - -@comet.compile(flags=None, target="gpu") -def run_comet_with_jit(A,B): - C = A - B - - return C - -A = np.full([4, 4], 3.4, dtype=float) -B = np.full([4, 4], 2.2, dtype=float) -C = np.full([4, 4], 0.0, dtype=float) -expected_result = run_numpy(A,B) -result_with_jit = run_comet_with_jit(A,B) -if sp.sparse.issparse(expected_result): - expected_result = expected_result.todense() - result_with_jit = result_with_jit.todense() -np.testing.assert_almost_equal(result_with_jit, expected_result) diff --git a/frontends/numpy-scipy/integration_tests/ops/gpu/test_gpu_eltwise_add_dense_matrix.py b/frontends/numpy-scipy/integration_tests/ops/gpu/test_gpu_eltwise_add_dense_matrix.py new file mode 100644 index 00000000..96d729de --- /dev/null +++ b/frontends/numpy-scipy/integration_tests/ops/gpu/test_gpu_eltwise_add_dense_matrix.py @@ -0,0 +1,28 @@ +import time +import numpy as np +import scipy as sp +from cometpy import comet +from conftest import gpu + +def run_numpy(A,B): + C = A+B + + return C + +@comet.compile(flags=None, target="gpu") +def run_comet_with_jit(A,B): + C = A+B + + return C + +@gpu +def test_gpu_eltwise_add_dense_matrix(): + A = np.full([4, 4], 2.2, dtype=float) + B = np.full([4, 4], 3.4, dtype=float) + C = np.full([4, 4], 0.0, dtype=float) + expected_result = run_numpy(A,B) + result_with_jit = run_comet_with_jit(A,B) + if sp.sparse.issparse(expected_result): + expected_result = expected_result.todense() + result_with_jit = result_with_jit.todense() + np.testing.assert_almost_equal(result_with_jit, expected_result) diff --git a/frontends/numpy-scipy/integration_tests/ops/gpu/test_gpu_eltwise_add_dense_matrix_interop.py b/frontends/numpy-scipy/integration_tests/ops/gpu/test_gpu_eltwise_add_dense_matrix_interop.py new file mode 100644 index 00000000..0db5cb52 --- /dev/null +++ b/frontends/numpy-scipy/integration_tests/ops/gpu/test_gpu_eltwise_add_dense_matrix_interop.py @@ -0,0 +1,25 @@ +import time +import numpy as np +import scipy as sp +from cometpy import comet +from conftest import gpu + +def run_numpy(A,B,C): + C[:] = A+B + +@comet.compile(flags=None, target="gpu") +def run_comet_with_jit(A,B,C): + C[:] = A+B + +@gpu +def test_gpu_eltwise_add_dense_matrix_interop(): + import cupy + A = np.full([4, 4], 2.2, dtype=float) + B = np.full([4, 4], 3.4, dtype=float) + C = np.full([4, 4], 0.0, dtype=float) + Ad = cupy.asarray(A) + Bd = cupy.asarray(B) + Cd = cupy.asarray(C) + run_numpy(A,B,C) + run_comet_with_jit(Ad,Bd,Cd) + np.testing.assert_almost_equal(Cd.get(), C) diff --git a/frontends/numpy-scipy/integration_tests/ops/gpu/test_gpu_eltwise_mult_DensexDense_oDense.py b/frontends/numpy-scipy/integration_tests/ops/gpu/test_gpu_eltwise_mult_DensexDense_oDense.py new file mode 100644 index 00000000..4fec0178 --- /dev/null +++ b/frontends/numpy-scipy/integration_tests/ops/gpu/test_gpu_eltwise_mult_DensexDense_oDense.py @@ -0,0 +1,28 @@ +import time +import numpy as np +import scipy as sp +from cometpy import comet +from conftest import gpu + +def run_numpy(A,B): + C = A * B + + return C + +@comet.compile(flags=None, target="gpu") +def run_comet_with_jit(A,B): + C = A * B + + return C + +@gpu +def test_gpu_eltwise_mult_DensexDense_oDense(): + A = np.full([4, 4], 2.7, dtype=float) + B = np.full([4, 4], 3.2, dtype=float) + C = np.full([4, 4], 0.0, dtype=float) + expected_result = run_numpy(A,B) + result_with_jit = run_comet_with_jit(A,B) + if sp.sparse.issparse(expected_result): + expected_result = expected_result.todense() + result_with_jit = result_with_jit.todense() + np.testing.assert_almost_equal(result_with_jit, expected_result) diff --git a/frontends/numpy-scipy/integration_tests/ops/gpu/test_gpu_eltwise_subtract_dense_matrix.py b/frontends/numpy-scipy/integration_tests/ops/gpu/test_gpu_eltwise_subtract_dense_matrix.py new file mode 100644 index 00000000..cae3bf39 --- /dev/null +++ b/frontends/numpy-scipy/integration_tests/ops/gpu/test_gpu_eltwise_subtract_dense_matrix.py @@ -0,0 +1,28 @@ +import time +import numpy as np +import scipy as sp +from cometpy import comet +from conftest import gpu + +def run_numpy(A,B): + C = A - B + + return C + +@comet.compile(flags=None, target="gpu") +def run_comet_with_jit(A,B): + C = A - B + + return C + +@gpu +def test_gpu_eltwise_subtract_dense_matrix(): + A = np.full([4, 4], 3.4, dtype=float) + B = np.full([4, 4], 2.2, dtype=float) + C = np.full([4, 4], 0.0, dtype=float) + expected_result = run_numpy(A,B) + result_with_jit = run_comet_with_jit(A,B) + if sp.sparse.issparse(expected_result): + expected_result = expected_result.todense() + result_with_jit = result_with_jit.todense() + np.testing.assert_almost_equal(result_with_jit, expected_result) diff --git a/frontends/numpy-scipy/integration_tests/ops/gpu/test_gpu_mult_dense_ij-ikj-kj.py b/frontends/numpy-scipy/integration_tests/ops/gpu/test_gpu_mult_dense_ij-ikj-kj.py new file mode 100644 index 00000000..feffb616 --- /dev/null +++ b/frontends/numpy-scipy/integration_tests/ops/gpu/test_gpu_mult_dense_ij-ikj-kj.py @@ -0,0 +1,28 @@ +import time +import numpy as np +import scipy as sp +from cometpy import comet +from conftest import gpu + +def run_numpy(A,B): + C = np.einsum('ikj,kj->ij', A,B) + + return C + +@comet.compile(flags=None, target="gpu") +def run_comet_with_jit(A,B): + C = comet.einsum('ikj,kj->ij', A,B) + + return C + +@gpu +def test_gpu_mult_dense_ij_ikj_kj(): + A = np.full([4, 4, 4], 3.2, dtype=float) + B = np.full([4, 4], 1.7, dtype=float) + C = np.full([4, 4], 0.0, dtype=float) + expected_result = run_numpy(A,B) + result_with_jit = run_comet_with_jit(A,B) + if sp.sparse.issparse(expected_result): + expected_result = expected_result.todense() + result_with_jit = result_with_jit.todense() + np.testing.assert_almost_equal(result_with_jit, expected_result) diff --git a/frontends/numpy-scipy/integration_tests/ops/gpu/test_gpu_mult_dense_matrix.py b/frontends/numpy-scipy/integration_tests/ops/gpu/test_gpu_mult_dense_matrix.py new file mode 100644 index 00000000..35502af1 --- /dev/null +++ b/frontends/numpy-scipy/integration_tests/ops/gpu/test_gpu_mult_dense_matrix.py @@ -0,0 +1,28 @@ +import time +import numpy as np +import scipy as sp +from cometpy import comet +from conftest import gpu + +def run_numpy(A,B): + C = np.einsum('ij,jk->ik', A,B) + + return C + +@comet.compile(flags=None, target="gpu") +def run_comet_with_jit(A,B): + C = comet.einsum('ij,jk->ik', A,B) + + return C + +@gpu +def test_gpu_mult_dense_matrix(): + A = np.full([8, 4], 2.2, dtype=float) + B = np.full([4, 2], 3.4, dtype=float) + C = np.full([8, 2], 0.0, dtype=float) + expected_result = run_numpy(A,B) + result_with_jit = run_comet_with_jit(A,B) + if sp.sparse.issparse(expected_result): + expected_result = expected_result.todense() + result_with_jit = result_with_jit.todense() + np.testing.assert_almost_equal(result_with_jit, expected_result) diff --git a/frontends/numpy-scipy/integration_tests/ops/gpu/test_gpu_mult_dense_matrix_vector.py b/frontends/numpy-scipy/integration_tests/ops/gpu/test_gpu_mult_dense_matrix_vector.py new file mode 100644 index 00000000..0706a590 --- /dev/null +++ b/frontends/numpy-scipy/integration_tests/ops/gpu/test_gpu_mult_dense_matrix_vector.py @@ -0,0 +1,28 @@ +import time +import numpy as np +import scipy as sp +from cometpy import comet +from conftest import gpu + +def run_numpy(A,B): + C = np.einsum('ij,j->i', A,B) + + return C + +@comet.compile(flags=None, target="gpu") +def run_comet_with_jit(A,B): + C = comet.einsum('ij,j->i', A,B) + + return C + +@gpu +def test_gpu_mult_dense_matrix_vector(): + A = np.full([8, 16], 2.3, dtype=float) + B = np.full([16], 3.7, dtype=float) + C = np.full([8], 0.0, dtype=float) + expected_result = run_numpy(A,B) + result_with_jit = run_comet_with_jit(A,B) + if sp.sparse.issparse(expected_result): + expected_result = expected_result.todense() + result_with_jit = result_with_jit.todense() + np.testing.assert_almost_equal(result_with_jit, expected_result) diff --git a/frontends/numpy-scipy/integration_tests/ops/gpu/test_mult_dense_ij-ikj-kj.py b/frontends/numpy-scipy/integration_tests/ops/gpu/test_mult_dense_ij-ikj-kj.py deleted file mode 100644 index 1d6961fe..00000000 --- a/frontends/numpy-scipy/integration_tests/ops/gpu/test_mult_dense_ij-ikj-kj.py +++ /dev/null @@ -1,25 +0,0 @@ -import time -import numpy as np -import scipy as sp -from cometpy import comet - -def run_numpy(A,B): - C = np.einsum('ikj,kj->ij', A,B) - - return C - -@comet.compile(flags=None, target="gpu") -def run_comet_with_jit(A,B): - C = comet.einsum('ikj,kj->ij', A,B) - - return C - -A = np.full([4, 4, 4], 3.2, dtype=float) -B = np.full([4, 4], 1.7, dtype=float) -C = np.full([4, 4], 0.0, dtype=float) -expected_result = run_numpy(A,B) -result_with_jit = run_comet_with_jit(A,B) -if sp.sparse.issparse(expected_result): - expected_result = expected_result.todense() - result_with_jit = result_with_jit.todense() -np.testing.assert_almost_equal(result_with_jit, expected_result) diff --git a/frontends/numpy-scipy/integration_tests/ops/gpu/test_mult_dense_matrix.py b/frontends/numpy-scipy/integration_tests/ops/gpu/test_mult_dense_matrix.py deleted file mode 100644 index 891eef96..00000000 --- a/frontends/numpy-scipy/integration_tests/ops/gpu/test_mult_dense_matrix.py +++ /dev/null @@ -1,24 +0,0 @@ -import time -import numpy as np -import scipy as sp -from cometpy import comet -def run_numpy(A,B): - C = np.einsum('ij,jk->ik', A,B) - - return C - -@comet.compile(flags=None, target="gpu") -def run_comet_with_jit(A,B): - C = comet.einsum('ij,jk->ik', A,B) - - return C - -A = np.full([8, 4], 2.2, dtype=float) -B = np.full([4, 2], 3.4, dtype=float) -C = np.full([8, 2], 0.0, dtype=float) -expected_result = run_numpy(A,B) -result_with_jit = run_comet_with_jit(A,B) -if sp.sparse.issparse(expected_result): - expected_result = expected_result.todense() - result_with_jit = result_with_jit.todense() -np.testing.assert_almost_equal(result_with_jit, expected_result) diff --git a/frontends/numpy-scipy/integration_tests/ops/gpu/test_mult_dense_matrix_vector.py b/frontends/numpy-scipy/integration_tests/ops/gpu/test_mult_dense_matrix_vector.py deleted file mode 100644 index e8a46b49..00000000 --- a/frontends/numpy-scipy/integration_tests/ops/gpu/test_mult_dense_matrix_vector.py +++ /dev/null @@ -1,25 +0,0 @@ -import time -import numpy as np -import scipy as sp -from cometpy import comet - -def run_numpy(A,B): - C = np.einsum('ij,j->i', A,B) - - return C - -@comet.compile(flags=None, target="gpu") -def run_comet_with_jit(A,B): - C = comet.einsum('ij,j->i', A,B) - - return C - -A = np.full([8, 16], 2.3, dtype=float) -B = np.full([16], 3.7, dtype=float) -C = np.full([8], 0.0, dtype=float) -expected_result = run_numpy(A,B) -result_with_jit = run_comet_with_jit(A,B) -if sp.sparse.issparse(expected_result): - expected_result = expected_result.todense() - result_with_jit = result_with_jit.todense() -np.testing.assert_almost_equal(result_with_jit, expected_result) diff --git a/frontends/numpy-scipy/integration_tests/ops/gpu/test_sum_dense_matrix.py b/frontends/numpy-scipy/integration_tests/ops/gpu/test_sum_dense_matrix.py deleted file mode 100644 index 0e3fb6d6..00000000 --- a/frontends/numpy-scipy/integration_tests/ops/gpu/test_sum_dense_matrix.py +++ /dev/null @@ -1,23 +0,0 @@ -import time -import numpy as np -import scipy as sp -from cometpy import comet - -def run_numpy(A): - var = A.sum() - - return var - -@comet.compile(flags=None, target="gpu") -def run_comet_with_jit(A): - var = A.sum() - - return var - -A = np.full([4, 4], 3.7, dtype=float) -expected_result = run_numpy(A) -result_with_jit = run_comet_with_jit(A) -if sp.sparse.issparse(expected_result): - expected_result = expected_result.todense() - result_with_jit = result_with_jit.todense() -np.testing.assert_almost_equal(result_with_jit, expected_result) diff --git a/frontends/numpy-scipy/integration_tests/ops/gpu/test_transpose_dense_matrix.py b/frontends/numpy-scipy/integration_tests/ops/gpu/test_transpose_dense_matrix.py deleted file mode 100644 index cfb9149a..00000000 --- a/frontends/numpy-scipy/integration_tests/ops/gpu/test_transpose_dense_matrix.py +++ /dev/null @@ -1,23 +0,0 @@ -import time -import numpy as np -import scipy as sp -from cometpy import comet - -def run_numpy(A): - B = A.transpose() - - return B - -@comet.compile(flags=None, target="gpu") -def run_comet_with_jit(A): - B = A.transpose() - - return B - -A = np.full([4, 4], 3.2, dtype=float) -expected_result = run_numpy(A) -result_with_jit = run_comet_with_jit(A) -if sp.sparse.issparse(expected_result): - expected_result = expected_result.todense() - result_with_jit = result_with_jit.todense() -np.testing.assert_almost_equal(result_with_jit, expected_result) diff --git a/frontends/numpy-scipy/integration_tests/ops/test_eltwise_CSRxCSR_oCSR_sameSpPattern.py b/frontends/numpy-scipy/integration_tests/ops/test_eltwise_CSRxCSR_oCSR_sameSpPattern.py deleted file mode 100644 index 05483444..00000000 --- a/frontends/numpy-scipy/integration_tests/ops/test_eltwise_CSRxCSR_oCSR_sameSpPattern.py +++ /dev/null @@ -1,25 +0,0 @@ -import time -import numpy as np -import scipy as sp -from cometpy import comet - -def run_numpy(A,B): - C = A * B - - return C - -@comet.compile(flags=None) -def run_comet_with_jit(A,B): - C = A * B - - return C - -A = sp.sparse.csr_array(sp.io.mmread("../../../../integration_test/data/test_rank2.mtx")) -B = sp.sparse.csr_array(sp.io.mmread("../../../../integration_test/data/test_rank2.mtx")) -expected_result = run_numpy(A,B) -result_with_jit = run_comet_with_jit(A,B) -if sp.sparse.issparse(expected_result): - expected_result = expected_result.todense() -if sp.sparse.issparse(result_with_jit): - result_with_jit = result_with_jit.todense() -np.testing.assert_almost_equal(result_with_jit, expected_result) diff --git a/frontends/numpy-scipy/integration_tests/ops/test_eltwise_add_CSRxCSR_oCSR_sameSpPattern.py b/frontends/numpy-scipy/integration_tests/ops/test_eltwise_add_CSRxCSR_oCSR_sameSpPattern.py deleted file mode 100644 index 66615990..00000000 --- a/frontends/numpy-scipy/integration_tests/ops/test_eltwise_add_CSRxCSR_oCSR_sameSpPattern.py +++ /dev/null @@ -1,24 +0,0 @@ -import time -import numpy as np -import scipy as sp -from cometpy import comet - -def run_numpy(A,B): - C = A + B - - return C - -@comet.compile(flags=None) -def run_comet_with_jit(A,B): - C = A + B - - return C - -A = sp.sparse.csr_array(sp.io.mmread("../../../../integration_test/data/test_rank2.mtx")) -B = sp.sparse.csr_array(sp.io.mmread("../../../../integration_test/data/test_rank2_transpose.mtx")) -expected_result = run_numpy(A,B) -result_with_jit = run_comet_with_jit(A,B) -if sp.sparse.issparse(expected_result): - expected_result = expected_result.todense() - result_with_jit = result_with_jit.todense() -np.testing.assert_almost_equal(result_with_jit, expected_result) diff --git a/frontends/numpy-scipy/integration_tests/ops/test_eltwise_add_dense_matrix.py b/frontends/numpy-scipy/integration_tests/ops/test_eltwise_add_dense_matrix.py index 57268a30..018cb6b9 100644 --- a/frontends/numpy-scipy/integration_tests/ops/test_eltwise_add_dense_matrix.py +++ b/frontends/numpy-scipy/integration_tests/ops/test_eltwise_add_dense_matrix.py @@ -14,12 +14,13 @@ def run_comet_with_jit(A,B): return C -A = np.full([4, 4], 2.2, dtype=float) -B = np.full([4, 4], 3.4, dtype=float) -C = np.full([4, 4], 0.0, dtype=float) -expected_result = run_numpy(A,B) -result_with_jit = run_comet_with_jit(A,B) -if sp.sparse.issparse(expected_result): - expected_result = expected_result.todense() - result_with_jit = result_with_jit.todense() -np.testing.assert_almost_equal(result_with_jit, expected_result) +def test_eltwise_add_dense_matrix(): + A = np.full([4, 4], 2.2, dtype=float) + B = np.full([4, 4], 3.4, dtype=float) + C = np.full([4, 4], 0.0, dtype=float) + expected_result = run_numpy(A,B) + result_with_jit = run_comet_with_jit(A,B) + if sp.sparse.issparse(expected_result): + expected_result = expected_result.todense() + result_with_jit = result_with_jit.todense() + np.testing.assert_almost_equal(result_with_jit, expected_result) diff --git a/frontends/numpy-scipy/integration_tests/ops/test_eltwise_mult_COOxDense_oCOO.py b/frontends/numpy-scipy/integration_tests/ops/test_eltwise_mult_COOxDense_oCOO.py index e5f6faf9..e4d2d163 100644 --- a/frontends/numpy-scipy/integration_tests/ops/test_eltwise_mult_COOxDense_oCOO.py +++ b/frontends/numpy-scipy/integration_tests/ops/test_eltwise_mult_COOxDense_oCOO.py @@ -14,12 +14,13 @@ def run_comet_with_jit(A,B): return C -A = sp.sparse.coo_array(sp.io.mmread("../../../../integration_test/data/test_rank2.mtx")) -B = np.full([A.shape[0], A.shape[1]], 2.7, dtype=float) -expected_result = run_numpy(A,B) -result_with_jit = run_comet_with_jit(A,B) -if sp.sparse.issparse(expected_result): - expected_result = expected_result.todense() -if sp.sparse.issparse(result_with_jit): - result_with_jit = result_with_jit.todense() -np.testing.assert_almost_equal(result_with_jit, expected_result) +def test_eltwise_mult_COOxDense_oCOO(data_rank2_path): + A = sp.sparse.coo_array(sp.io.mmread(data_rank2_path)) + B = np.full([A.shape[0], A.shape[1]], 2.7, dtype=float) + expected_result = run_numpy(A,B) + result_with_jit = run_comet_with_jit(A,B) + if sp.sparse.issparse(expected_result): + expected_result = expected_result.todense() + if sp.sparse.issparse(result_with_jit): + result_with_jit = result_with_jit.todense() + np.testing.assert_almost_equal(result_with_jit, expected_result) diff --git a/frontends/numpy-scipy/integration_tests/ops/test_eltwise_mult_CSRxDense_oCSR.py b/frontends/numpy-scipy/integration_tests/ops/test_eltwise_mult_CSRxDense_oCSR.py index 2426eadb..c13265ed 100644 --- a/frontends/numpy-scipy/integration_tests/ops/test_eltwise_mult_CSRxDense_oCSR.py +++ b/frontends/numpy-scipy/integration_tests/ops/test_eltwise_mult_CSRxDense_oCSR.py @@ -14,12 +14,14 @@ def run_comet_with_jit(A,B): return C -A = sp.sparse.csr_array(sp.io.mmread("../../../../integration_test/data/test_rank2.mtx")) -B = np.full([A.shape[0], A.shape[1]], 2.7, dtype=float) -expected_result = run_numpy(A,B) -result_with_jit = run_comet_with_jit(A,B) -if sp.sparse.issparse(expected_result): - expected_result = expected_result.todense() -if sp.sparse.issparse(result_with_jit): - result_with_jit = result_with_jit.todense() -np.testing.assert_almost_equal(result_with_jit, expected_result) + +def test_eltwise_mult_CSRxDense_oCSR(data_rank2_path): + A = sp.sparse.csr_array(sp.io.mmread(data_rank2_path)) + B = np.full([A.shape[0], A.shape[1]], 2.7, dtype=float) + expected_result = run_numpy(A,B) + result_with_jit = run_comet_with_jit(A,B) + if sp.sparse.issparse(expected_result): + expected_result = expected_result.todense() + if sp.sparse.issparse(result_with_jit): + result_with_jit = result_with_jit.todense() + np.testing.assert_almost_equal(result_with_jit, expected_result) diff --git a/frontends/numpy-scipy/integration_tests/ops/test_eltwise_mult_CSRxDense_oDense.py b/frontends/numpy-scipy/integration_tests/ops/test_eltwise_mult_CSRxDense_oDense.py index 3e6af3d8..8f5fc467 100644 --- a/frontends/numpy-scipy/integration_tests/ops/test_eltwise_mult_CSRxDense_oDense.py +++ b/frontends/numpy-scipy/integration_tests/ops/test_eltwise_mult_CSRxDense_oDense.py @@ -14,13 +14,14 @@ def run_comet_with_jit(A,B): return C -A = sp.sparse.csr_array(sp.io.mmread("../../../../integration_test/data/test_rank2.mtx")) -B = np.full([A.shape[0], A.shape[1]], 2.7, dtype=float) -C = np.full([A.shape[0], A.shape[1]], 0.0, dtype=float) -expected_result = run_numpy(A,B) -result_with_jit = run_comet_with_jit(A,B) -if sp.sparse.issparse(expected_result): - expected_result = expected_result.todense() -if sp.sparse.issparse(result_with_jit): - result_with_jit = result_with_jit.todense() -np.testing.assert_almost_equal(result_with_jit, expected_result) +def test_eltwise_mult_CSRxDense_oDense(data_rank2_path): + A = sp.sparse.csr_array(sp.io.mmread(data_rank2_path)) + B = np.full([A.shape[0], A.shape[1]], 2.7, dtype=float) + C = np.full([A.shape[0], A.shape[1]], 0.0, dtype=float) + expected_result = run_numpy(A,B) + result_with_jit = run_comet_with_jit(A,B) + if sp.sparse.issparse(expected_result): + expected_result = expected_result.todense() + if sp.sparse.issparse(result_with_jit): + result_with_jit = result_with_jit.todense() + np.testing.assert_almost_equal(result_with_jit, expected_result) diff --git a/frontends/numpy-scipy/integration_tests/ops/test_eltwise_mult_DensexCSR_oDense.py b/frontends/numpy-scipy/integration_tests/ops/test_eltwise_mult_DensexCSR_oDense.py index 9c811ff9..871428d0 100644 --- a/frontends/numpy-scipy/integration_tests/ops/test_eltwise_mult_DensexCSR_oDense.py +++ b/frontends/numpy-scipy/integration_tests/ops/test_eltwise_mult_DensexCSR_oDense.py @@ -14,13 +14,14 @@ def run_comet_with_jit(A,B): return C -B = sp.sparse.csr_array(sp.io.mmread("../../../../integration_test/data/test_rank2.mtx")) -A = np.full([B.shape[0], B.shape[1]], 2.7, dtype=float) -C = np.full([B.shape[0], B.shape[1]], 0.0, dtype=float) -expected_result = run_numpy(A,B) -result_with_jit = run_comet_with_jit(A,B) -if sp.sparse.issparse(expected_result): - expected_result = expected_result.todense() -if sp.sparse.issparse(result_with_jit): - result_with_jit = result_with_jit.todense() -np.testing.assert_almost_equal(result_with_jit, expected_result) +def test_eltwise_mult_DensexCSR_oDense(data_rank2_path): + B = sp.sparse.csr_array(sp.io.mmread(data_rank2_path)) + A = np.full([B.shape[0], B.shape[1]], 2.7, dtype=float) + C = np.full([B.shape[0], B.shape[1]], 0.0, dtype=float) + expected_result = run_numpy(A,B) + result_with_jit = run_comet_with_jit(A,B) + if sp.sparse.issparse(expected_result): + expected_result = expected_result.todense() + if sp.sparse.issparse(result_with_jit): + result_with_jit = result_with_jit.todense() + np.testing.assert_almost_equal(result_with_jit, expected_result) diff --git a/frontends/numpy-scipy/integration_tests/ops/test_eltwise_mult_DensexDense_oDense.py b/frontends/numpy-scipy/integration_tests/ops/test_eltwise_mult_DensexDense_oDense.py index 3aafd796..33c44c05 100644 --- a/frontends/numpy-scipy/integration_tests/ops/test_eltwise_mult_DensexDense_oDense.py +++ b/frontends/numpy-scipy/integration_tests/ops/test_eltwise_mult_DensexDense_oDense.py @@ -14,12 +14,13 @@ def run_comet_with_jit(A,B): return C -A = np.full([4, 4], 2.7, dtype=float) -B = np.full([4, 4], 3.2, dtype=float) -C = np.full([4, 4], 0.0, dtype=float) -expected_result = run_numpy(A,B) -result_with_jit = run_comet_with_jit(A,B) -if sp.sparse.issparse(expected_result): - expected_result = expected_result.todense() - result_with_jit = result_with_jit.todense() -np.testing.assert_almost_equal(result_with_jit, expected_result) +def test_eltwise_mult_DensexDense_oDense(): + A = np.full([4, 4], 2.7, dtype=float) + B = np.full([4, 4], 3.2, dtype=float) + C = np.full([4, 4], 0.0, dtype=float) + expected_result = run_numpy(A,B) + result_with_jit = run_comet_with_jit(A,B) + if sp.sparse.issparse(expected_result): + expected_result = expected_result.todense() + result_with_jit = result_with_jit.todense() + np.testing.assert_almost_equal(result_with_jit, expected_result) diff --git a/frontends/numpy-scipy/integration_tests/ops/test_eltwise_mult_dense_4Dtensors.py b/frontends/numpy-scipy/integration_tests/ops/test_eltwise_mult_dense_4Dtensors.py index 5b689014..18165bcf 100644 --- a/frontends/numpy-scipy/integration_tests/ops/test_eltwise_mult_dense_4Dtensors.py +++ b/frontends/numpy-scipy/integration_tests/ops/test_eltwise_mult_dense_4Dtensors.py @@ -14,12 +14,13 @@ def run_comet_with_jit(A,B): return C -A = np.full([2, 2, 2, 2], 2.2, dtype=float) -B = np.full([2, 2, 2, 2], 3.6, dtype=float) -C = np.full([2, 2, 2, 2], 0.0, dtype=float) -expected_result = run_numpy(A,B) -result_with_jit = run_comet_with_jit(A,B) -if sp.sparse.issparse(expected_result): - expected_result = expected_result.todense() - result_with_jit = result_with_jit.todense() -np.testing.assert_almost_equal(result_with_jit, expected_result) +def test_eltwise_mult_dense_4Dtensors(): + A = np.full([2, 2, 2, 2], 2.2, dtype=float) + B = np.full([2, 2, 2, 2], 3.6, dtype=float) + C = np.full([2, 2, 2, 2], 0.0, dtype=float) + expected_result = run_numpy(A,B) + result_with_jit = run_comet_with_jit(A,B) + if sp.sparse.issparse(expected_result): + expected_result = expected_result.todense() + result_with_jit = result_with_jit.todense() + np.testing.assert_almost_equal(result_with_jit, expected_result) diff --git a/frontends/numpy-scipy/integration_tests/ops/test_eltwise_subtract_CSRxCSR_oCSR_sameSpPattern.py b/frontends/numpy-scipy/integration_tests/ops/test_eltwise_subtract_CSRxCSR_oCSR_sameSpPattern.py deleted file mode 100644 index 07dca683..00000000 --- a/frontends/numpy-scipy/integration_tests/ops/test_eltwise_subtract_CSRxCSR_oCSR_sameSpPattern.py +++ /dev/null @@ -1,24 +0,0 @@ -import time -import numpy as np -import scipy as sp -from cometpy import comet - -def run_numpy(A,B): - C = A - B - - return C - -@comet.compile(flags=None) -def run_comet_with_jit(A,B): - C = A - B - - return C - -A = sp.sparse.csr_array(sp.io.mmread("../../../../integration_test/data/test_rank2.mtx")) -B = sp.sparse.csr_array(sp.io.mmread("../../../../integration_test/data/test_rank2_transpose.mtx")) -expected_result = run_numpy(A,B) -result_with_jit = run_comet_with_jit(A,B) -if sp.sparse.issparse(expected_result): - expected_result = expected_result.todense() - result_with_jit = result_with_jit.todense() -np.testing.assert_almost_equal(result_with_jit, expected_result) diff --git a/frontends/numpy-scipy/integration_tests/ops/test_eltwise_subtract_dense_matrix.py b/frontends/numpy-scipy/integration_tests/ops/test_eltwise_subtract_dense_matrix.py index 78b20051..2e26fb53 100644 --- a/frontends/numpy-scipy/integration_tests/ops/test_eltwise_subtract_dense_matrix.py +++ b/frontends/numpy-scipy/integration_tests/ops/test_eltwise_subtract_dense_matrix.py @@ -14,12 +14,13 @@ def run_comet_with_jit(A,B): return C -A = np.full([4, 4], 3.4, dtype=float) -B = np.full([4, 4], 2.2, dtype=float) -C = np.full([4, 4], 0.0, dtype=float) -expected_result = run_numpy(A,B) -result_with_jit = run_comet_with_jit(A,B) -if sp.sparse.issparse(expected_result): - expected_result = expected_result.todense() - result_with_jit = result_with_jit.todense() -np.testing.assert_almost_equal(result_with_jit, expected_result) +def test_eltwise_subtract_dense_matrix(): + A = np.full([4, 4], 3.4, dtype=float) + B = np.full([4, 4], 2.2, dtype=float) + C = np.full([4, 4], 0.0, dtype=float) + expected_result = run_numpy(A,B) + result_with_jit = run_comet_with_jit(A,B) + if sp.sparse.issparse(expected_result): + expected_result = expected_result.todense() + result_with_jit = result_with_jit.todense() + np.testing.assert_almost_equal(result_with_jit, expected_result) diff --git a/frontends/numpy-scipy/integration_tests/ops/test_mult_DenseMatxCOO.py b/frontends/numpy-scipy/integration_tests/ops/test_mult_DenseMatxCOO.py index af2e56b4..43fa3341 100644 --- a/frontends/numpy-scipy/integration_tests/ops/test_mult_DenseMatxCOO.py +++ b/frontends/numpy-scipy/integration_tests/ops/test_mult_DenseMatxCOO.py @@ -14,12 +14,13 @@ def run_comet_with_jit(A,B): return C -B = sp.sparse.coo_array(sp.io.mmread("../../../../integration_test/data/test_rank2.mtx")) -A = np.full([4, B.shape[0]], 1.7, dtype=float) -C = np.full([4, B.shape[1]], 0.0, dtype=float) -expected_result = run_numpy(A,B) -result_with_jit = run_comet_with_jit(A,B) -if sp.sparse.issparse(expected_result): - expected_result = expected_result.todense() - result_with_jit = result_with_jit.todense() -np.testing.assert_almost_equal(result_with_jit, expected_result) +def test_mult_DenseMatxCOO(data_rank2_path): + B = sp.sparse.coo_array(sp.io.mmread(data_rank2_path)) + A = np.full([4, B.shape[0]], 1.7, dtype=float) + C = np.full([4, B.shape[1]], 0.0, dtype=float) + expected_result = run_numpy(A,B) + result_with_jit = run_comet_with_jit(A,B) + if sp.sparse.issparse(expected_result): + expected_result = expected_result.todense() + result_with_jit = result_with_jit.todense() + np.testing.assert_almost_equal(result_with_jit, expected_result) diff --git a/frontends/numpy-scipy/integration_tests/ops/test_mult_DenseMatxCSR.py b/frontends/numpy-scipy/integration_tests/ops/test_mult_DenseMatxCSR.py index e31d76e4..75ab03dd 100644 --- a/frontends/numpy-scipy/integration_tests/ops/test_mult_DenseMatxCSR.py +++ b/frontends/numpy-scipy/integration_tests/ops/test_mult_DenseMatxCSR.py @@ -14,12 +14,13 @@ def run_comet_with_jit(A,B): return C -B = sp.sparse.csr_array(sp.io.mmread("../../../../integration_test/data/test_rank2.mtx")) -A = np.full([4, B.shape[0]], 1.7, dtype=float) -C = np.full([4, B.shape[1]], 0.0, dtype=float) -expected_result = run_numpy(A,B) -result_with_jit = run_comet_with_jit(A,B) -if sp.sparse.issparse(expected_result): - expected_result = expected_result.todense() - result_with_jit = result_with_jit.todense() -np.testing.assert_almost_equal(result_with_jit, expected_result) +def test_mult_DenseMatxCSR(data_rank2_path): + B = sp.sparse.csr_array(sp.io.mmread(data_rank2_path)) + A = np.full([4, B.shape[0]], 1.7, dtype=float) + C = np.full([4, B.shape[1]], 0.0, dtype=float) + expected_result = run_numpy(A,B) + result_with_jit = run_comet_with_jit(A,B) + if sp.sparse.issparse(expected_result): + expected_result = expected_result.todense() + result_with_jit = result_with_jit.todense() + np.testing.assert_almost_equal(result_with_jit, expected_result) diff --git a/frontends/numpy-scipy/integration_tests/ops/test_mult_DenseVecxCOO.py b/frontends/numpy-scipy/integration_tests/ops/test_mult_DenseVecxCOO.py index 3190c6ef..76f17895 100644 --- a/frontends/numpy-scipy/integration_tests/ops/test_mult_DenseVecxCOO.py +++ b/frontends/numpy-scipy/integration_tests/ops/test_mult_DenseVecxCOO.py @@ -14,12 +14,13 @@ def run_comet_with_jit(A,B): return C -B = sp.sparse.coo_array(sp.io.mmread("../../../../integration_test/data/test_rank2.mtx")) -A = np.full([B.shape[0]], 1.7, dtype=float) -C = np.full([B.shape[1]], 0.0, dtype=float) -expected_result = run_numpy(A,B) -result_with_jit = run_comet_with_jit(A,B) -if sp.sparse.issparse(expected_result): - expected_result = expected_result.todense() - result_with_jit = result_with_jit.todense() -np.testing.assert_almost_equal(result_with_jit, expected_result) +def test_mult_DenseVecxCOO(data_rank2_path): + B = sp.sparse.coo_array(sp.io.mmread(data_rank2_path)) + A = np.full([B.shape[0]], 1.7, dtype=float) + C = np.full([B.shape[1]], 0.0, dtype=float) + expected_result = run_numpy(A,B) + result_with_jit = run_comet_with_jit(A,B) + if sp.sparse.issparse(expected_result): + expected_result = expected_result.todense() + result_with_jit = result_with_jit.todense() + np.testing.assert_almost_equal(result_with_jit, expected_result) diff --git a/frontends/numpy-scipy/integration_tests/ops/test_mult_DenseVecxCSR.py b/frontends/numpy-scipy/integration_tests/ops/test_mult_DenseVecxCSR.py index 7badf457..87e53f16 100644 --- a/frontends/numpy-scipy/integration_tests/ops/test_mult_DenseVecxCSR.py +++ b/frontends/numpy-scipy/integration_tests/ops/test_mult_DenseVecxCSR.py @@ -14,12 +14,13 @@ def run_comet_with_jit(A,B): return C -B = sp.sparse.csr_array(sp.io.mmread("../../../../integration_test/data/test_rank2.mtx")) -A = np.full([B.shape[0]], 1.7, dtype=float) -C = np.full([B.shape[1]], 0.0, dtype=float) -expected_result = run_numpy(A,B) -result_with_jit = run_comet_with_jit(A,B) -if sp.sparse.issparse(expected_result): - expected_result = expected_result.todense() - result_with_jit = result_with_jit.todense() -np.testing.assert_almost_equal(result_with_jit, expected_result) +def test_mult_DenseVecxCSR(data_rank2_path): + B = sp.sparse.csr_array(sp.io.mmread(data_rank2_path)) + A = np.full([B.shape[0]], 1.7, dtype=float) + C = np.full([B.shape[1]], 0.0, dtype=float) + expected_result = run_numpy(A,B) + result_with_jit = run_comet_with_jit(A,B) + if sp.sparse.issparse(expected_result): + expected_result = expected_result.todense() + result_with_jit = result_with_jit.todense() + np.testing.assert_almost_equal(result_with_jit, expected_result) diff --git a/frontends/numpy-scipy/integration_tests/ops/test_mult_dense_4Dtensors.py b/frontends/numpy-scipy/integration_tests/ops/test_mult_dense_4Dtensors.py index 942666ef..21837b9f 100644 --- a/frontends/numpy-scipy/integration_tests/ops/test_mult_dense_4Dtensors.py +++ b/frontends/numpy-scipy/integration_tests/ops/test_mult_dense_4Dtensors.py @@ -14,12 +14,13 @@ def run_comet_with_jit(A,B): return C -A = np.full([2, 2, 2, 2], 2.2, dtype=float) -B = np.full([2, 2, 2, 2], 3.6, dtype=float) -C = np.full([2, 2, 2, 2], 0.0, dtype=float) -expected_result = run_numpy(A,B) -result_with_jit = run_comet_with_jit(A,B) -if sp.sparse.issparse(expected_result): - expected_result = expected_result.todense() - result_with_jit = result_with_jit.todense() -np.testing.assert_almost_equal(result_with_jit, expected_result) +def test_mult_dense_4Dtensors(): + A = np.full([2, 2, 2, 2], 2.2, dtype=float) + B = np.full([2, 2, 2, 2], 3.6, dtype=float) + C = np.full([2, 2, 2, 2], 0.0, dtype=float) + expected_result = run_numpy(A,B) + result_with_jit = run_comet_with_jit(A,B) + if sp.sparse.issparse(expected_result): + expected_result = expected_result.todense() + result_with_jit = result_with_jit.todense() + np.testing.assert_almost_equal(result_with_jit, expected_result) diff --git a/frontends/numpy-scipy/integration_tests/ops/test_mult_dense_ij-ikj-kj.py b/frontends/numpy-scipy/integration_tests/ops/test_mult_dense_ij-ikj-kj.py index dc041881..f23c10bd 100644 --- a/frontends/numpy-scipy/integration_tests/ops/test_mult_dense_ij-ikj-kj.py +++ b/frontends/numpy-scipy/integration_tests/ops/test_mult_dense_ij-ikj-kj.py @@ -14,12 +14,13 @@ def run_comet_with_jit(A,B): return C -A = np.full([4, 4, 4], 3.2, dtype=float) -B = np.full([4, 4], 1.7, dtype=float) -C = np.full([4, 4], 0.0, dtype=float) -expected_result = run_numpy(A,B) -result_with_jit = run_comet_with_jit(A,B) -if sp.sparse.issparse(expected_result): - expected_result = expected_result.todense() - result_with_jit = result_with_jit.todense() -np.testing.assert_almost_equal(result_with_jit, expected_result) +def test_mult_dense_ij_ikj_kj(): + A = np.full([4, 4, 4], 3.2, dtype=float) + B = np.full([4, 4], 1.7, dtype=float) + C = np.full([4, 4], 0.0, dtype=float) + expected_result = run_numpy(A,B) + result_with_jit = run_comet_with_jit(A,B) + if sp.sparse.issparse(expected_result): + expected_result = expected_result.todense() + result_with_jit = result_with_jit.todense() + np.testing.assert_almost_equal(result_with_jit, expected_result) diff --git a/frontends/numpy-scipy/integration_tests/ops/test_mult_dense_matrix.py b/frontends/numpy-scipy/integration_tests/ops/test_mult_dense_matrix.py index 268358cf..e34c8e19 100644 --- a/frontends/numpy-scipy/integration_tests/ops/test_mult_dense_matrix.py +++ b/frontends/numpy-scipy/integration_tests/ops/test_mult_dense_matrix.py @@ -13,12 +13,13 @@ def run_comet_with_jit(A,B): return C -A = np.full([8, 4], 2.2, dtype=float) -B = np.full([4, 2], 3.4, dtype=float) -C = np.full([8, 2], 0.0, dtype=float) -expected_result = run_numpy(A,B) -result_with_jit = run_comet_with_jit(A,B) -if sp.sparse.issparse(expected_result): - expected_result = expected_result.todense() - result_with_jit = result_with_jit.todense() -np.testing.assert_almost_equal(result_with_jit, expected_result) +def test_mult_dense_matrix(): + A = np.full([8, 4], 2.2, dtype=float) + B = np.full([4, 2], 3.4, dtype=float) + C = np.full([8, 2], 0.0, dtype=float) + expected_result = run_numpy(A,B) + result_with_jit = run_comet_with_jit(A,B) + if sp.sparse.issparse(expected_result): + expected_result = expected_result.todense() + result_with_jit = result_with_jit.todense() + np.testing.assert_almost_equal(result_with_jit, expected_result) diff --git a/frontends/numpy-scipy/integration_tests/ops/test_mult_dense_matrix_vector.py b/frontends/numpy-scipy/integration_tests/ops/test_mult_dense_matrix_vector.py index ec97b611..9de25bf5 100644 --- a/frontends/numpy-scipy/integration_tests/ops/test_mult_dense_matrix_vector.py +++ b/frontends/numpy-scipy/integration_tests/ops/test_mult_dense_matrix_vector.py @@ -14,12 +14,13 @@ def run_comet_with_jit(A,B): return C -A = np.full([8, 16], 2.3, dtype=float) -B = np.full([16], 3.7, dtype=float) -C = np.full([8], 0.0, dtype=float) -expected_result = run_numpy(A,B) -result_with_jit = run_comet_with_jit(A,B) -if sp.sparse.issparse(expected_result): - expected_result = expected_result.todense() - result_with_jit = result_with_jit.todense() -np.testing.assert_almost_equal(result_with_jit, expected_result) +def test_mult_dense_matrix_vector(): + A = np.full([8, 16], 2.3, dtype=float) + B = np.full([16], 3.7, dtype=float) + C = np.full([8], 0.0, dtype=float) + expected_result = run_numpy(A,B) + result_with_jit = run_comet_with_jit(A,B) + if sp.sparse.issparse(expected_result): + expected_result = expected_result.todense() + result_with_jit = result_with_jit.todense() + np.testing.assert_almost_equal(result_with_jit, expected_result) diff --git a/frontends/numpy-scipy/integration_tests/ops/test_mult_spgemm_CSRxCSR_oCSR.py b/frontends/numpy-scipy/integration_tests/ops/test_mult_spgemm_CSRxCSR_oCSR.py index 1b2b07b5..6c516352 100644 --- a/frontends/numpy-scipy/integration_tests/ops/test_mult_spgemm_CSRxCSR_oCSR.py +++ b/frontends/numpy-scipy/integration_tests/ops/test_mult_spgemm_CSRxCSR_oCSR.py @@ -14,11 +14,12 @@ def run_comet_with_jit(A,B): return C -A = sp.sparse.csr_array(sp.io.mmread("../../../../integration_test/data/test_rank2.mtx")) -B = sp.sparse.csr_array(sp.io.mmread("../../../../integration_test/data/test_rank2.mtx")) -expected_result = run_numpy(A,B) -result_with_jit = run_comet_with_jit(A,B) -if sp.sparse.issparse(expected_result): - expected_result = expected_result.todense() - result_with_jit = result_with_jit.todense() -np.testing.assert_almost_equal(result_with_jit, expected_result) +def test_mult_spgemm_CSRxCSR_oCSR(data_rank2_path): + A = sp.sparse.csr_array(sp.io.mmread(data_rank2_path)) + B = sp.sparse.csr_array(sp.io.mmread(data_rank2_path)) + expected_result = run_numpy(A,B) + result_with_jit = run_comet_with_jit(A,B) + if sp.sparse.issparse(expected_result): + expected_result = expected_result.todense() + result_with_jit = result_with_jit.todense() + np.testing.assert_almost_equal(result_with_jit, expected_result) \ No newline at end of file diff --git a/frontends/numpy-scipy/integration_tests/ops/test_mult_spmm_COOxDense.py b/frontends/numpy-scipy/integration_tests/ops/test_mult_spmm_COOxDense.py index 4b02d806..3a522225 100644 --- a/frontends/numpy-scipy/integration_tests/ops/test_mult_spmm_COOxDense.py +++ b/frontends/numpy-scipy/integration_tests/ops/test_mult_spmm_COOxDense.py @@ -14,12 +14,13 @@ def run_comet_with_jit(A,B): return C -A = sp.sparse.coo_array(sp.io.mmread("../../../../integration_test/data/test_rank2.mtx")) -B = np.full([A.shape[1], 4], 1.7, dtype=float) -C = np.full([A.shape[0], 4], 0.0, dtype=float) -expected_result = run_numpy(A,B) -result_with_jit = run_comet_with_jit(A,B) -if sp.sparse.issparse(expected_result): - expected_result = expected_result.todense() - result_with_jit = result_with_jit.todense() -np.testing.assert_almost_equal(result_with_jit, expected_result) +def test_mult_spmm_COOxDense(data_rank2_path): + A = sp.sparse.coo_array(sp.io.mmread(data_rank2_path)) + B = np.full([A.shape[1], 4], 1.7, dtype=float) + C = np.full([A.shape[0], 4], 0.0, dtype=float) + expected_result = run_numpy(A,B) + result_with_jit = run_comet_with_jit(A,B) + if sp.sparse.issparse(expected_result): + expected_result = expected_result.todense() + result_with_jit = result_with_jit.todense() + np.testing.assert_almost_equal(result_with_jit, expected_result) diff --git a/frontends/numpy-scipy/integration_tests/ops/test_mult_spmm_CSRxDense.py b/frontends/numpy-scipy/integration_tests/ops/test_mult_spmm_CSRxDense.py index 5a98d2cd..e4d2837e 100644 --- a/frontends/numpy-scipy/integration_tests/ops/test_mult_spmm_CSRxDense.py +++ b/frontends/numpy-scipy/integration_tests/ops/test_mult_spmm_CSRxDense.py @@ -14,12 +14,13 @@ def run_comet_with_jit(A,B): return C -A = sp.sparse.csr_array(sp.io.mmread("../../../../integration_test/data/test_rank2.mtx")) -B = np.full([A.shape[1], 4], 1.7, dtype=float) -C = np.full([A.shape[0], 4], 0.0, dtype=float) -expected_result = run_numpy(A,B) -result_with_jit = run_comet_with_jit(A,B) -if sp.sparse.issparse(expected_result): - expected_result = expected_result.todense() - result_with_jit = result_with_jit.todense() -np.testing.assert_almost_equal(result_with_jit, expected_result) +def test_mult_spmm_CSRxDense(data_rank2_path): + A = sp.sparse.csr_array(sp.io.mmread(data_rank2_path)) + B = np.full([A.shape[1], 4], 1.7, dtype=float) + C = np.full([A.shape[0], 4], 0.0, dtype=float) + expected_result = run_numpy(A,B) + result_with_jit = run_comet_with_jit(A,B) + if sp.sparse.issparse(expected_result): + expected_result = expected_result.todense() + result_with_jit = result_with_jit.todense() + np.testing.assert_almost_equal(result_with_jit, expected_result) diff --git a/frontends/numpy-scipy/integration_tests/ops/test_mult_spmv_COOxDense.py b/frontends/numpy-scipy/integration_tests/ops/test_mult_spmv_COOxDense.py index 135cee53..8f52bd00 100644 --- a/frontends/numpy-scipy/integration_tests/ops/test_mult_spmv_COOxDense.py +++ b/frontends/numpy-scipy/integration_tests/ops/test_mult_spmv_COOxDense.py @@ -14,12 +14,13 @@ def run_comet_with_jit(A,B): return C -A = sp.sparse.coo_array(sp.io.mmread("../../../../integration_test/data/test_rank2.mtx")) -B = np.full([A.shape[1]], 1.7, dtype=float) -C = np.full([A.shape[0]], 0.0, dtype=float) -expected_result = run_numpy(A,B) -result_with_jit = run_comet_with_jit(A,B) -if sp.sparse.issparse(expected_result): - expected_result = expected_result.todense() - result_with_jit = result_with_jit.todense() -np.testing.assert_almost_equal(result_with_jit, expected_result) +def test_mult_spmv_COOxDense(data_rank2_path): + A = sp.sparse.coo_array(sp.io.mmread(data_rank2_path)) + B = np.full([A.shape[1]], 1.7, dtype=float) + C = np.full([A.shape[0]], 0.0, dtype=float) + expected_result = run_numpy(A,B) + result_with_jit = run_comet_with_jit(A,B) + if sp.sparse.issparse(expected_result): + expected_result = expected_result.todense() + result_with_jit = result_with_jit.todense() + np.testing.assert_almost_equal(result_with_jit, expected_result) diff --git a/frontends/numpy-scipy/integration_tests/ops/test_mult_spmv_CSRxDense.py b/frontends/numpy-scipy/integration_tests/ops/test_mult_spmv_CSRxDense.py index 11be7077..abb07a2f 100644 --- a/frontends/numpy-scipy/integration_tests/ops/test_mult_spmv_CSRxDense.py +++ b/frontends/numpy-scipy/integration_tests/ops/test_mult_spmv_CSRxDense.py @@ -14,12 +14,13 @@ def run_comet_with_jit(A,B): return C -A = sp.sparse.csr_array(sp.io.mmread("../../../../integration_test/data/test_rank2.mtx")) -B = np.full([A.shape[1]], 1.7, dtype=float) -C = np.full([A.shape[0]], 0.0, dtype=float) -expected_result = run_numpy(A,B) -result_with_jit = run_comet_with_jit(A,B) -if sp.sparse.issparse(expected_result): - expected_result = expected_result.todense() - result_with_jit = result_with_jit.todense() -np.testing.assert_almost_equal(result_with_jit, expected_result) +def test_mult_spmv_COOxDense(data_rank2_path): + A = sp.sparse.csr_array(sp.io.mmread(data_rank2_path)) + B = np.full([A.shape[1]], 1.7, dtype=float) + C = np.full([A.shape[0]], 0.0, dtype=float) + expected_result = run_numpy(A,B) + result_with_jit = run_comet_with_jit(A,B) + if sp.sparse.issparse(expected_result): + expected_result = expected_result.todense() + result_with_jit = result_with_jit.todense() + np.testing.assert_almost_equal(result_with_jit, expected_result) diff --git a/frontends/numpy-scipy/integration_tests/ops/test_sum_COO.py b/frontends/numpy-scipy/integration_tests/ops/test_sum_COO.py index cad64925..03143863 100644 --- a/frontends/numpy-scipy/integration_tests/ops/test_sum_COO.py +++ b/frontends/numpy-scipy/integration_tests/ops/test_sum_COO.py @@ -14,10 +14,11 @@ def run_comet_with_jit(A): return var -A = sp.sparse.coo_array(sp.io.mmread("../../../../integration_test/data/test_rank2.mtx")) -expected_result = run_numpy(A) -result_with_jit = run_comet_with_jit(A) -if sp.sparse.issparse(expected_result): - expected_result = expected_result.todense() - result_with_jit = result_with_jit.todense() -np.testing.assert_almost_equal(result_with_jit, expected_result) +def test_sum_COO(data_rank2_path): + A = sp.sparse.coo_array(sp.io.mmread(data_rank2_path)) + expected_result = run_numpy(A) + result_with_jit = run_comet_with_jit(A) + if sp.sparse.issparse(expected_result): + expected_result = expected_result.todense() + result_with_jit = result_with_jit.todense() + np.testing.assert_almost_equal(result_with_jit, expected_result) diff --git a/frontends/numpy-scipy/integration_tests/ops/test_sum_CSR.py b/frontends/numpy-scipy/integration_tests/ops/test_sum_CSR.py index fa517d58..5e5fa902 100644 --- a/frontends/numpy-scipy/integration_tests/ops/test_sum_CSR.py +++ b/frontends/numpy-scipy/integration_tests/ops/test_sum_CSR.py @@ -14,10 +14,11 @@ def run_comet_with_jit(A): return var -A = sp.sparse.csr_array(sp.io.mmread("../../../../integration_test/data/test_rank2.mtx")) -expected_result = run_numpy(A) -result_with_jit = run_comet_with_jit(A) -if sp.sparse.issparse(expected_result): - expected_result = expected_result.todense() - result_with_jit = result_with_jit.todense() -np.testing.assert_almost_equal(result_with_jit, expected_result) +def test_sum_CSR(data_rank2_path): + A = sp.sparse.csr_array(sp.io.mmread(data_rank2_path)) + expected_result = run_numpy(A) + result_with_jit = run_comet_with_jit(A) + if sp.sparse.issparse(expected_result): + expected_result = expected_result.todense() + result_with_jit = result_with_jit.todense() + np.testing.assert_almost_equal(result_with_jit, expected_result) diff --git a/frontends/numpy-scipy/integration_tests/ops/test_sum_dense_matrix.py b/frontends/numpy-scipy/integration_tests/ops/test_sum_dense_matrix.py index 4f88db39..f10ab30b 100644 --- a/frontends/numpy-scipy/integration_tests/ops/test_sum_dense_matrix.py +++ b/frontends/numpy-scipy/integration_tests/ops/test_sum_dense_matrix.py @@ -14,10 +14,11 @@ def run_comet_with_jit(A): return var -A = np.full([4, 4], 3.7, dtype=float) -expected_result = run_numpy(A) -result_with_jit = run_comet_with_jit(A) -if sp.sparse.issparse(expected_result): - expected_result = expected_result.todense() - result_with_jit = result_with_jit.todense() -np.testing.assert_almost_equal(result_with_jit, expected_result) +def test_sum_dense_matrix(): + A = np.full([4, 4], 3.7, dtype=float) + expected_result = run_numpy(A) + result_with_jit = run_comet_with_jit(A) + if sp.sparse.issparse(expected_result): + expected_result = expected_result.todense() + result_with_jit = result_with_jit.todense() + np.testing.assert_almost_equal(result_with_jit, expected_result) diff --git a/frontends/numpy-scipy/integration_tests/ops/test_sum_dense_tensor.py b/frontends/numpy-scipy/integration_tests/ops/test_sum_dense_tensor.py index 7da0abe4..599ef764 100644 --- a/frontends/numpy-scipy/integration_tests/ops/test_sum_dense_tensor.py +++ b/frontends/numpy-scipy/integration_tests/ops/test_sum_dense_tensor.py @@ -14,10 +14,11 @@ def run_comet_with_jit(A): return var -A = np.full([4, 4, 4], 3.7, dtype=float) -expected_result = run_numpy(A) -result_with_jit = run_comet_with_jit(A) -if sp.sparse.issparse(expected_result): - expected_result = expected_result.todense() - result_with_jit = result_with_jit.todense() -np.testing.assert_almost_equal(result_with_jit, expected_result) +def test_sum_dense_tensor(): + A = np.full([4, 4, 4], 3.7, dtype=float) + expected_result = run_numpy(A) + result_with_jit = run_comet_with_jit(A) + if sp.sparse.issparse(expected_result): + expected_result = expected_result.todense() + result_with_jit = result_with_jit.todense() + np.testing.assert_almost_equal(result_with_jit, expected_result) diff --git a/frontends/numpy-scipy/integration_tests/ops/test_transpose_COO_matrix.py b/frontends/numpy-scipy/integration_tests/ops/test_transpose_COO_matrix.py deleted file mode 100644 index 465610d1..00000000 --- a/frontends/numpy-scipy/integration_tests/ops/test_transpose_COO_matrix.py +++ /dev/null @@ -1,23 +0,0 @@ -import time -import numpy as np -import scipy as sp -from cometpy import comet - -def run_numpy(A): - B = A.transpose() - - return B - -@comet.compile(flags=None) -def run_comet_with_jit(A): - B = A.transpose() - - return B - -A = sp.sparse.coo_array(sp.io.mmread("../../../../integration_test/data/test_rank2.mtx")) -expected_result = run_numpy(A) -result_with_jit = run_comet_with_jit(A) -if sp.sparse.issparse(expected_result): - expected_result = expected_result.todense() - result_with_jit = result_with_jit.todense() -np.testing.assert_almost_equal(result_with_jit, expected_result) diff --git a/frontends/numpy-scipy/integration_tests/ops/test_transpose_CSR_matrix.py b/frontends/numpy-scipy/integration_tests/ops/test_transpose_CSR_matrix.py deleted file mode 100644 index 3632d33d..00000000 --- a/frontends/numpy-scipy/integration_tests/ops/test_transpose_CSR_matrix.py +++ /dev/null @@ -1,23 +0,0 @@ -import time -import numpy as np -import scipy as sp -from cometpy import comet - -def run_numpy(A): - B = A.transpose() - - return B - -@comet.compile(flags=None) -def run_comet_with_jit(A): - B = A.transpose() - - return B - -A = sp.sparse.csr_array(sp.io.mmread("../../../../integration_test/data/test_rank2.mtx")) -expected_result = run_numpy(A) -result_with_jit = run_comet_with_jit(A) -if sp.sparse.issparse(expected_result): - expected_result = expected_result.todense() - result_with_jit = result_with_jit.todense() -np.testing.assert_almost_equal(result_with_jit, expected_result) diff --git a/frontends/numpy-scipy/integration_tests/ops/test_transpose_dense_matrix.py b/frontends/numpy-scipy/integration_tests/ops/test_transpose_dense_matrix.py deleted file mode 100644 index cb727122..00000000 --- a/frontends/numpy-scipy/integration_tests/ops/test_transpose_dense_matrix.py +++ /dev/null @@ -1,23 +0,0 @@ -import time -import numpy as np -import scipy as sp -from cometpy import comet - -def run_numpy(A): - B = A.transpose() - - return B - -@comet.compile(flags=None) -def run_comet_with_jit(A): - B = A.transpose() - - return B - -A = np.full([4, 4], 3.2, dtype=float) -expected_result = run_numpy(A) -result_with_jit = run_comet_with_jit(A) -if sp.sparse.issparse(expected_result): - expected_result = expected_result.todense() - result_with_jit = result_with_jit.todense() -np.testing.assert_almost_equal(result_with_jit, expected_result) diff --git a/frontends/numpy-scipy/integration_tests/ops/test_transpose_dense_tensor.py b/frontends/numpy-scipy/integration_tests/ops/test_transpose_dense_tensor.py deleted file mode 100644 index e014e727..00000000 --- a/frontends/numpy-scipy/integration_tests/ops/test_transpose_dense_tensor.py +++ /dev/null @@ -1,23 +0,0 @@ -import time -import numpy as np -import scipy as sp -from cometpy import comet - -def run_numpy(A): - B = np.einsum('ijk->kij', A) - - return B - -@comet.compile(flags=None) -def run_comet_with_jit(A): - B = comet.einsum('ijk->kij', A) - - return B - -A = np.full([4, 4, 4], 3.7, dtype=float) -expected_result = run_numpy(A) -result_with_jit = run_comet_with_jit(A) -if sp.sparse.issparse(expected_result): - expected_result = expected_result.todense() - result_with_jit = result_with_jit.todense() -np.testing.assert_almost_equal(result_with_jit, expected_result) diff --git a/frontends/numpy-scipy/integration_tests/opts/test_ccsd_t1_21_ttgt.py b/frontends/numpy-scipy/integration_tests/opts/test_ccsd_t1_21_ttgt.py index 63862a91..72456f60 100644 --- a/frontends/numpy-scipy/integration_tests/opts/test_ccsd_t1_21_ttgt.py +++ b/frontends/numpy-scipy/integration_tests/opts/test_ccsd_t1_21_ttgt.py @@ -14,12 +14,13 @@ def run_comet_with_jit(v,t2): return i0 -v = np.full([2, 2, 4, 4], 2.3, dtype=float) -t2 = np.full([4, 4, 2, 4], 3.4, dtype=float) -i0 = np.full([2, 4], 0.0, dtype=float) -expected_result = run_numpy(v,t2) -result_with_jit = run_comet_with_jit(v,t2) -if sp.sparse.issparse(expected_result): - expected_result = expected_result.todense() - result_with_jit = result_with_jit.todense() -np.testing.assert_almost_equal(result_with_jit, expected_result) +def test_ccsd_t1_21_ttgt(): + v = np.full([2, 2, 4, 4], 2.3, dtype=float) + t2 = np.full([4, 4, 2, 4], 3.4, dtype=float) + i0 = np.full([2, 4], 0.0, dtype=float) + expected_result = run_numpy(v,t2) + result_with_jit = run_comet_with_jit(v,t2) + if sp.sparse.issparse(expected_result): + expected_result = expected_result.todense() + result_with_jit = result_with_jit.todense() + np.testing.assert_almost_equal(result_with_jit, expected_result) diff --git a/frontends/numpy-scipy/integration_tests/opts/test_ccsd_t1_21_ttgt_all_opts.py b/frontends/numpy-scipy/integration_tests/opts/test_ccsd_t1_21_ttgt_all_opts.py index 1552cfc7..df9c3538 100644 --- a/frontends/numpy-scipy/integration_tests/opts/test_ccsd_t1_21_ttgt_all_opts.py +++ b/frontends/numpy-scipy/integration_tests/opts/test_ccsd_t1_21_ttgt_all_opts.py @@ -8,18 +8,19 @@ def run_numpy(v,t2): return i0 -@comet.compile(flags="-opt-bestperm-ttgt -opt-matmul-tiling -opt-matmul-mkernel -opt-dense-transpose --convert-tc-to-ttgt") +@comet.compile(flags="-opt-matmul-tiling -opt-matmul-mkernel -opt-dense-transpose --convert-tc-to-ttgt") def run_comet_with_jit(v,t2): i0 = comet.einsum('icmn,mnca->ia', v,t2) return i0 -v = np.full([16, 16, 16, 16], 2.3, dtype=float) -t2 = np.full([16, 16, 16, 16], 3.4, dtype=float) -i0 = np.full([16, 16], 0.0, dtype=float) -expected_result = run_numpy(v,t2) -result_with_jit = run_comet_with_jit(v,t2) -if sp.sparse.issparse(expected_result): - expected_result = expected_result.todense() - result_with_jit = result_with_jit.todense() -np.testing.assert_almost_equal(result_with_jit, expected_result) +def test_ccsd_t1_21_ttgt_all_opts(): + v = np.full([16, 16, 16, 16], 2.3, dtype=float) + t2 = np.full([16, 16, 16, 16], 3.4, dtype=float) + i0 = np.full([16, 16], 0.0, dtype=float) + expected_result = run_numpy(v,t2) + result_with_jit = run_comet_with_jit(v,t2) + if sp.sparse.issparse(expected_result): + expected_result = expected_result.todense() + result_with_jit = result_with_jit.todense() + np.testing.assert_almost_equal(result_with_jit, expected_result) diff --git a/frontends/numpy-scipy/integration_tests/opts/test_ccsd_t1_21_ttgt_tilling.py b/frontends/numpy-scipy/integration_tests/opts/test_ccsd_t1_21_ttgt_tilling.py index a690bbbd..d92a4a6a 100644 --- a/frontends/numpy-scipy/integration_tests/opts/test_ccsd_t1_21_ttgt_tilling.py +++ b/frontends/numpy-scipy/integration_tests/opts/test_ccsd_t1_21_ttgt_tilling.py @@ -14,12 +14,13 @@ def run_comet_with_jit(v,t2): return i0 -v = np.full([2, 2, 4, 4], 2.3, dtype=float) -t2 = np.full([4, 4, 2, 4], 3.4, dtype=float) -i0 = np.full([2, 4], 0.0, dtype=float) -expected_result = run_numpy(v,t2) -result_with_jit = run_comet_with_jit(v,t2) -if sp.sparse.issparse(expected_result): - expected_result = expected_result.todense() - result_with_jit = result_with_jit.todense() -np.testing.assert_almost_equal(result_with_jit, expected_result) +def test_ccsd_t1_21_ttgt_tilling(): + v = np.full([2, 2, 4, 4], 2.3, dtype=float) + t2 = np.full([4, 4, 2, 4], 3.4, dtype=float) + i0 = np.full([2, 4], 0.0, dtype=float) + expected_result = run_numpy(v,t2) + result_with_jit = run_comet_with_jit(v,t2) + if sp.sparse.issparse(expected_result): + expected_result = expected_result.todense() + result_with_jit = result_with_jit.todense() + np.testing.assert_almost_equal(result_with_jit, expected_result) diff --git a/frontends/numpy-scipy/integration_tests/opts/test_ccsd_t1_3_ttgt.py b/frontends/numpy-scipy/integration_tests/opts/test_ccsd_t1_3_ttgt.py index ebf75c7e..f0fd339b 100644 --- a/frontends/numpy-scipy/integration_tests/opts/test_ccsd_t1_3_ttgt.py +++ b/frontends/numpy-scipy/integration_tests/opts/test_ccsd_t1_3_ttgt.py @@ -14,12 +14,13 @@ def run_comet_with_jit(f,t1): return i0 -f = np.full([2, 4], 2.3, dtype=float) -t1 = np.full([2, 2], 3.4, dtype=float) -i0 = np.full([2, 4], 0.0, dtype=float) -expected_result = run_numpy(f,t1) -result_with_jit = run_comet_with_jit(f,t1) -if sp.sparse.issparse(expected_result): - expected_result = expected_result.todense() - result_with_jit = result_with_jit.todense() -np.testing.assert_almost_equal(result_with_jit, expected_result) +def test_ccsd_t1_3_ttgt(): + f = np.full([2, 4], 2.3, dtype=float) + t1 = np.full([2, 2], 3.4, dtype=float) + i0 = np.full([2, 4], 0.0, dtype=float) + expected_result = run_numpy(f,t1) + result_with_jit = run_comet_with_jit(f,t1) + if sp.sparse.issparse(expected_result): + expected_result = expected_result.todense() + result_with_jit = result_with_jit.todense() + np.testing.assert_almost_equal(result_with_jit, expected_result) diff --git a/frontends/numpy-scipy/integration_tests/opts/test_ccsd_t1_4_ttgt.py b/frontends/numpy-scipy/integration_tests/opts/test_ccsd_t1_4_ttgt.py index cd5cd52c..a8c13056 100644 --- a/frontends/numpy-scipy/integration_tests/opts/test_ccsd_t1_4_ttgt.py +++ b/frontends/numpy-scipy/integration_tests/opts/test_ccsd_t1_4_ttgt.py @@ -14,12 +14,13 @@ def run_comet_with_jit(v,t1): return i0 -v = np.full([2, 2, 4, 4], 2.3, dtype=float) -t1 = np.full([4, 2], 3.4, dtype=float) -i0 = np.full([2, 4], 0.0, dtype=float) -expected_result = run_numpy(v,t1) -result_with_jit = run_comet_with_jit(v,t1) -if sp.sparse.issparse(expected_result): - expected_result = expected_result.todense() - result_with_jit = result_with_jit.todense() -np.testing.assert_almost_equal(result_with_jit, expected_result) +def test_ccsd_t1_4_ttgt(): + v = np.full([2, 2, 4, 4], 2.3, dtype=float) + t1 = np.full([4, 2], 3.4, dtype=float) + i0 = np.full([2, 4], 0.0, dtype=float) + expected_result = run_numpy(v,t1) + result_with_jit = run_comet_with_jit(v,t1) + if sp.sparse.issparse(expected_result): + expected_result = expected_result.todense() + result_with_jit = result_with_jit.todense() + np.testing.assert_almost_equal(result_with_jit, expected_result) diff --git a/frontends/numpy-scipy/integration_tests/opts/test_ccsd_t1_4_ttgt_bestperm.py b/frontends/numpy-scipy/integration_tests/opts/test_ccsd_t1_4_ttgt_bestperm.py index e6b6ece0..831bf0ac 100644 --- a/frontends/numpy-scipy/integration_tests/opts/test_ccsd_t1_4_ttgt_bestperm.py +++ b/frontends/numpy-scipy/integration_tests/opts/test_ccsd_t1_4_ttgt_bestperm.py @@ -8,18 +8,19 @@ def run_numpy(v,t1): return i0 -@comet.compile(flags="--opt-bestperm-ttgt --convert-tc-to-ttgt") +@comet.compile(flags="--convert-tc-to-ttgt") def run_comet_with_jit(v,t1): i0 = comet.einsum('cima,mc->ia', v,t1) return i0 -v = np.full([2, 2, 4, 4], 2.3, dtype=float) -t1 = np.full([4, 2], 3.4, dtype=float) -i0 = np.full([2, 4], 0.0, dtype=float) -expected_result = run_numpy(v,t1) -result_with_jit = run_comet_with_jit(v,t1) -if sp.sparse.issparse(expected_result): - expected_result = expected_result.todense() - result_with_jit = result_with_jit.todense() -np.testing.assert_almost_equal(result_with_jit, expected_result) +def test_ccsd_t1_4_ttgt_bestperm(): + v = np.full([2, 2, 4, 4], 2.3, dtype=float) + t1 = np.full([4, 2], 3.4, dtype=float) + i0 = np.full([2, 4], 0.0, dtype=float) + expected_result = run_numpy(v,t1) + result_with_jit = run_comet_with_jit(v,t1) + if sp.sparse.issparse(expected_result): + expected_result = expected_result.todense() + result_with_jit = result_with_jit.todense() + np.testing.assert_almost_equal(result_with_jit, expected_result) diff --git a/frontends/numpy-scipy/integration_tests/opts/test_fusion.py b/frontends/numpy-scipy/integration_tests/opts/test_fusion.py index 3d893ec9..c24009d1 100644 --- a/frontends/numpy-scipy/integration_tests/opts/test_fusion.py +++ b/frontends/numpy-scipy/integration_tests/opts/test_fusion.py @@ -2,6 +2,7 @@ import numpy as np import scipy as sp from cometpy import comet +import pytest def run_numpy(B,C,D): T = B @ C @@ -16,14 +17,16 @@ def run_comet_with_jit(B,C,D): return A -B = sp.sparse.csr_array(sp.io.mmread("../../../../integration_test/data/test_rank2_small.mtx")) -C = np.full([B.shape[1], 4], 1.2, dtype=float) -D = np.full([4, 4], 3.4, dtype=float) -A = np.full([B.shape[0], 4], 0.0, dtype=float) -T = np.full([B.shape[0], 4], 0.0, dtype=float) -expected_result = run_numpy(B,C,D) -result_with_jit = run_comet_with_jit(B,C,D) -if sp.sparse.issparse(expected_result): - expected_result = expected_result.todense() - result_with_jit = result_with_jit.todense() -np.testing.assert_almost_equal(result_with_jit, expected_result) +@pytest.mark.skip("Fusing more than one iterations is not currently supported") +def test_fusion(data_rank2_path): + B = sp.sparse.csr_array(sp.io.mmread(data_rank2_path)) + C = np.full([B.shape[1], 4], 1.2, dtype=float) + D = np.full([4, 4], 3.4, dtype=float) + A = np.full([B.shape[0], 4], 0.0, dtype=float) + T = np.full([B.shape[0], 4], 0.0, dtype=float) + expected_result = run_numpy(B,C,D) + result_with_jit = run_comet_with_jit(B,C,D) + if sp.sparse.issparse(expected_result): + expected_result = expected_result.todense() + result_with_jit = result_with_jit.todense() + np.testing.assert_almost_equal(result_with_jit, expected_result) diff --git a/frontends/numpy-scipy/integration_tests/opts/test_opt_dense_transpose.py b/frontends/numpy-scipy/integration_tests/opts/test_opt_dense_transpose.py index e158a4dd..e22629e7 100644 --- a/frontends/numpy-scipy/integration_tests/opts/test_opt_dense_transpose.py +++ b/frontends/numpy-scipy/integration_tests/opts/test_opt_dense_transpose.py @@ -14,10 +14,11 @@ def run_comet_with_jit(A): return B -A = np.full([2, 4, 8, 16], 3.7, dtype=float) -expected_result = run_numpy(A) -result_with_jit = run_comet_with_jit(A) -if sp.sparse.issparse(expected_result): - expected_result = expected_result.todense() - result_with_jit = result_with_jit.todense() -np.testing.assert_almost_equal(result_with_jit, expected_result) +def test_opt_dense_transpose(): + A = np.full([2, 4, 8, 16], 3.7, dtype=float) + expected_result = run_numpy(A) + result_with_jit = run_comet_with_jit(A) + if sp.sparse.issparse(expected_result): + expected_result = expected_result.todense() + result_with_jit = result_with_jit.todense() + np.testing.assert_almost_equal(result_with_jit, expected_result) diff --git a/frontends/numpy-scipy/integration_tests/semiring/test_eltwise_monoidMin_DensexDense_oDense.py b/frontends/numpy-scipy/integration_tests/semiring/test_eltwise_monoidMin_DensexDense_oDense.py index f6f7ec55..232c570e 100644 --- a/frontends/numpy-scipy/integration_tests/semiring/test_eltwise_monoidMin_DensexDense_oDense.py +++ b/frontends/numpy-scipy/integration_tests/semiring/test_eltwise_monoidMin_DensexDense_oDense.py @@ -12,6 +12,7 @@ def run_comet(A,B): A = np.full((4,4), 2.7) B = np.full((4,4), 3.2) -res = run_comet(A,B) -expected = np.array([2.7,2.7,2.7,2.7,2.7,2.7,2.7,2.7,2.7,2.7,2.7,2.7,2.7,2.7,2.7,2.7]).reshape((4,4)) -np.testing.assert_almost_equal(res, expected) \ No newline at end of file +def test_eltwise_monoidMin_DensexDense_oDense(): + res = run_comet(A,B) + expected = np.array([2.7,2.7,2.7,2.7,2.7,2.7,2.7,2.7,2.7,2.7,2.7,2.7,2.7,2.7,2.7,2.7]).reshape((4,4)) + np.testing.assert_almost_equal(res, expected) \ No newline at end of file diff --git a/frontends/numpy-scipy/integration_tests/semiring/test_eltwise_monoidMinus_DensexDense_oDense.py b/frontends/numpy-scipy/integration_tests/semiring/test_eltwise_monoidMinus_DensexDense_oDense.py index d1074625..6dd88b3b 100644 --- a/frontends/numpy-scipy/integration_tests/semiring/test_eltwise_monoidMinus_DensexDense_oDense.py +++ b/frontends/numpy-scipy/integration_tests/semiring/test_eltwise_monoidMinus_DensexDense_oDense.py @@ -12,6 +12,7 @@ def run_comet(A,B): A = np.full((4,4), 4.2) B = np.full((4,4), 2.7) -res = run_comet(A,B) -expected = np.array([1.5,1.5,1.5,1.5,1.5,1.5,1.5,1.5,1.5,1.5,1.5,1.5,1.5,1.5,1.5,1.5]).reshape((4,4)) -np.testing.assert_almost_equal(res, expected) \ No newline at end of file +def test_eltwise_monoidMinus_DensexDense_oDense(): + res = run_comet(A,B) + expected = np.array([1.5,1.5,1.5,1.5,1.5,1.5,1.5,1.5,1.5,1.5,1.5,1.5,1.5,1.5,1.5,1.5]).reshape((4,4)) + np.testing.assert_almost_equal(res, expected) \ No newline at end of file diff --git a/frontends/numpy-scipy/integration_tests/semiring/test_eltwise_monoidPlus_COOxDense_oCOO.py b/frontends/numpy-scipy/integration_tests/semiring/test_eltwise_monoidPlus_COOxDense_oCOO.py index 09dcd2a8..d3f289b7 100644 --- a/frontends/numpy-scipy/integration_tests/semiring/test_eltwise_monoidPlus_COOxDense_oCOO.py +++ b/frontends/numpy-scipy/integration_tests/semiring/test_eltwise_monoidPlus_COOxDense_oCOO.py @@ -8,10 +8,10 @@ def run_comet(A,B): return C +def test_eltwise_monoidPlus_COOxDense_oCOO(data_rank2_path): + A = sp.sparse.coo_array(sp.io.mmread(data_rank2_path)) + B = np.full((A.shape), 2.7) -A = sp.sparse.coo_array(sp.io.mmread("../../../../integration_test/data/test_rank2.mtx")) -B = np.full((A.shape), 2.7) - -res = run_comet(A,B) -expected = sp.sparse.coo_array(([3.7,4.1,4.7,5.2,5.7,6.8,6.7,7.9,7.7], ([0,0,1,1,2,3,3,4,4], [0,3,1,4,2,0,3,1,4]))) -np.testing.assert_almost_equal(res.todense(), expected.todense()) \ No newline at end of file + res = run_comet(A,B) + expected = sp.sparse.coo_array(([3.7,4.1,4.7,5.2,5.7,6.8,6.7,7.9,7.7], ([0,0,1,1,2,3,3,4,4], [0,3,1,4,2,0,3,1,4]))) + np.testing.assert_almost_equal(res.todense(), expected.todense()) \ No newline at end of file diff --git a/frontends/numpy-scipy/integration_tests/semiring/test_eltwise_monoidPlus_DensexDense_oDense.py b/frontends/numpy-scipy/integration_tests/semiring/test_eltwise_monoidPlus_DensexDense_oDense.py index a25edbba..108ad333 100644 --- a/frontends/numpy-scipy/integration_tests/semiring/test_eltwise_monoidPlus_DensexDense_oDense.py +++ b/frontends/numpy-scipy/integration_tests/semiring/test_eltwise_monoidPlus_DensexDense_oDense.py @@ -8,10 +8,10 @@ def run_comet(A,B): return C +def test_eltwise_monoidPlus_DensexDense_oDense(): + A = np.full((4,4), 2.7) + B = np.full((4,4), 3.2) -A = np.full((4,4), 2.7) -B = np.full((4,4), 3.2) - -res = run_comet(A,B) -expected = np.array([5.9,5.9,5.9,5.9,5.9,5.9,5.9,5.9,5.9,5.9,5.9,5.9,5.9,5.9,5.9,5.9]).reshape((4,4)) -np.testing.assert_almost_equal(res, expected) \ No newline at end of file + res = run_comet(A,B) + expected = np.array([5.9,5.9,5.9,5.9,5.9,5.9,5.9,5.9,5.9,5.9,5.9,5.9,5.9,5.9,5.9,5.9]).reshape((4,4)) + np.testing.assert_almost_equal(res, expected) \ No newline at end of file diff --git a/frontends/numpy-scipy/integration_tests/semiring/test_eltwise_monoidTimes_COOxDense_oCOO.py b/frontends/numpy-scipy/integration_tests/semiring/test_eltwise_monoidTimes_COOxDense_oCOO.py index 65b1db8c..40c84111 100644 --- a/frontends/numpy-scipy/integration_tests/semiring/test_eltwise_monoidTimes_COOxDense_oCOO.py +++ b/frontends/numpy-scipy/integration_tests/semiring/test_eltwise_monoidTimes_COOxDense_oCOO.py @@ -8,10 +8,10 @@ def run_comet(A,B): return C +def test_eltwise_monoidTimes_COOxDense_oCOO(data_rank2_path): + A = sp.sparse.coo_array(sp.io.mmread(data_rank2_path)) + B = np.full((A.shape), 2.7) -A = sp.sparse.coo_array(sp.io.mmread("../../../../integration_test/data/test_rank2.mtx")) -B = np.full((A.shape), 2.7) - -res = run_comet(A,B) -expected = sp.sparse.coo_array(([2.7,3.78,5.4,6.75,8.1,11.07,10.8,14.04,13.5], ([0,0,1,1,2,3,3,4,4], [0,3,1,4,2,0,3,1,4]))) -np.testing.assert_almost_equal(res.todense(), expected.todense()) \ No newline at end of file + res = run_comet(A,B) + expected = sp.sparse.coo_array(([2.7,3.78,5.4,6.75,8.1,11.07,10.8,14.04,13.5], ([0,0,1,1,2,3,3,4,4], [0,3,1,4,2,0,3,1,4]))) + np.testing.assert_almost_equal(res.todense(), expected.todense()) \ No newline at end of file diff --git a/frontends/numpy-scipy/integration_tests/semiring/test_eltwise_monoidTimes_DensexDense_oDense.py b/frontends/numpy-scipy/integration_tests/semiring/test_eltwise_monoidTimes_DensexDense_oDense.py index 2e85566f..6128b241 100644 --- a/frontends/numpy-scipy/integration_tests/semiring/test_eltwise_monoidTimes_DensexDense_oDense.py +++ b/frontends/numpy-scipy/integration_tests/semiring/test_eltwise_monoidTimes_DensexDense_oDense.py @@ -8,10 +8,10 @@ def run_comet(A,B): return C +def test_eltwise_monoidTimes_DensexDense_oDense(): + A = np.full((4,4), 2.7) + B = np.full((4,4), 3.2) -A = np.full((4,4), 2.7) -B = np.full((4,4), 3.2) - -res = run_comet(A,B) -expected = np.array([8.64,8.64,8.64,8.64,8.64,8.64,8.64,8.64,8.64,8.64,8.64,8.64,8.64,8.64,8.64,8.64]).reshape((4,4)) -np.testing.assert_almost_equal(res, expected) \ No newline at end of file + res = run_comet(A,B) + expected = np.array([8.64,8.64,8.64,8.64,8.64,8.64,8.64,8.64,8.64,8.64,8.64,8.64,8.64,8.64,8.64,8.64]).reshape((4,4)) + np.testing.assert_almost_equal(res, expected) \ No newline at end of file diff --git a/frontends/numpy-scipy/integration_tests/semiring/test_eltwise_monoidTimes_dense_4Dtensors.py b/frontends/numpy-scipy/integration_tests/semiring/test_eltwise_monoidTimes_dense_4Dtensors.py index f2c3d5bc..36e027cb 100644 --- a/frontends/numpy-scipy/integration_tests/semiring/test_eltwise_monoidTimes_dense_4Dtensors.py +++ b/frontends/numpy-scipy/integration_tests/semiring/test_eltwise_monoidTimes_dense_4Dtensors.py @@ -8,10 +8,10 @@ def run_comet(A,B): return C +def test_eltwise_monoidTimes_dense_4Dtensors(): + A = np.full((2,2,2,2), 2.2) + B = np.full((2,2,2,2), 3.6) -A = np.full((2,2,2,2), 2.2) -B = np.full((2,2,2,2), 3.6) - -res = run_comet(A,B) -expected = np.array([7.92,7.92,7.92,7.92,7.92,7.92,7.92,7.92,7.92,7.92,7.92,7.92,7.92,7.92,7.92,7.92]).reshape((2,2,2,2)) -np.testing.assert_almost_equal(res, expected) \ No newline at end of file + res = run_comet(A,B) + expected = np.array([7.92,7.92,7.92,7.92,7.92,7.92,7.92,7.92,7.92,7.92,7.92,7.92,7.92,7.92,7.92,7.92]).reshape((2,2,2,2)) + np.testing.assert_almost_equal(res, expected) \ No newline at end of file diff --git a/frontends/numpy-scipy/integration_tests/semiring/test_mm_SemiringAnyPair_CSRxCSR_oCSR.py b/frontends/numpy-scipy/integration_tests/semiring/test_mm_SemiringAnyPair_CSRxCSR_oCSR.py index aa67e9d2..766e71ce 100644 --- a/frontends/numpy-scipy/integration_tests/semiring/test_mm_SemiringAnyPair_CSRxCSR_oCSR.py +++ b/frontends/numpy-scipy/integration_tests/semiring/test_mm_SemiringAnyPair_CSRxCSR_oCSR.py @@ -8,10 +8,10 @@ def run_comet(A,B): return C +def test_mm_SemiringAnyPair_CSRxCSR_oCSR(data_rank2_path): + A = sp.sparse.csr_array(sp.io.mmread(data_rank2_path)) + B = sp.sparse.csr_array(sp.io.mmread(data_rank2_path)) -A = sp.sparse.csr_array(sp.io.mmread("../../../../integration_test/data/test_rank2.mtx")) -B = sp.sparse.csr_array(sp.io.mmread("../../../../integration_test/data/test_rank2.mtx")) - -res = run_comet(A,B) -expected = sp.sparse.csr_array(([1,1,1,1,1,1,1,1,1], [0,3,1,4,2,0,3,1,4], [0,2,4,5,7,9])) -np.testing.assert_almost_equal(res.todense(), expected.todense()) \ No newline at end of file + res = run_comet(A,B) + expected = sp.sparse.csr_array(([1,1,1,1,1,1,1,1,1], [0,3,1,4,2,0,3,1,4], [0,2,4,5,7,9])) + np.testing.assert_almost_equal(res.todense(), expected.todense()) \ No newline at end of file diff --git a/frontends/numpy-scipy/integration_tests/semiring/test_mm_SemiringMinFirst_CSRxCSR_oCSR.py b/frontends/numpy-scipy/integration_tests/semiring/test_mm_SemiringMinFirst_CSRxCSR_oCSR.py index 145b7bd1..d2b32a2d 100644 --- a/frontends/numpy-scipy/integration_tests/semiring/test_mm_SemiringMinFirst_CSRxCSR_oCSR.py +++ b/frontends/numpy-scipy/integration_tests/semiring/test_mm_SemiringMinFirst_CSRxCSR_oCSR.py @@ -8,10 +8,10 @@ def run_comet(A,B): return C +def test_mm_SemiringMinFirst_CSRxCSR_oCSR(data_rank2_path): + A = sp.sparse.csr_array(sp.io.mmread(data_rank2_path)) + B = sp.sparse.csr_array(sp.io.mmread(data_rank2_path)) -A = sp.sparse.csr_array(sp.io.mmread("../../../../integration_test/data/test_rank2.mtx")) -B = sp.sparse.csr_array(sp.io.mmread("../../../../integration_test/data/test_rank2.mtx")) - -res = run_comet(A,B) -expected = sp.sparse.csr_array(([1,1,2,2,3,4,4,5,5], [0,3,1,4,2,0,3,1,4], [0,2,4,5,7,9])) -np.testing.assert_almost_equal(res.todense(), expected.todense()) \ No newline at end of file + res = run_comet(A,B) + expected = sp.sparse.csr_array(([1,1,2,2,3,4,4,5,5], [0,3,1,4,2,0,3,1,4], [0,2,4,5,7,9])) + np.testing.assert_almost_equal(res.todense(), expected.todense()) \ No newline at end of file diff --git a/frontends/numpy-scipy/integration_tests/semiring/test_mm_SemiringMinPlus_CSRxCSR_oCSR.py b/frontends/numpy-scipy/integration_tests/semiring/test_mm_SemiringMinPlus_CSRxCSR_oCSR.py index 00d82356..7ea1b477 100644 --- a/frontends/numpy-scipy/integration_tests/semiring/test_mm_SemiringMinPlus_CSRxCSR_oCSR.py +++ b/frontends/numpy-scipy/integration_tests/semiring/test_mm_SemiringMinPlus_CSRxCSR_oCSR.py @@ -8,10 +8,10 @@ def run_comet(A,B): return C +def test_mm_SemiringMinPlus_CSRxCSR_oCSR(data_rank2_path): + A = sp.sparse.csr_array(sp.io.mmread(data_rank2_path)) + B = sp.sparse.csr_array(sp.io.mmread(data_rank2_path)) -A = sp.sparse.csr_array(sp.io.mmread("../../../../integration_test/data/test_rank2.mtx")) -B = sp.sparse.csr_array(sp.io.mmread("../../../../integration_test/data/test_rank2.mtx")) - -res = run_comet(A,B) -expected = sp.sparse.csr_array(([2,2.4,4,4.5,6,5.1,5.5,7.2,7.7], [0,3,1,4,2,0,3,1,4], [0,2,4,5,7,9])) -np.testing.assert_almost_equal(res.todense(), expected.todense()) \ No newline at end of file + res = run_comet(A,B) + expected = sp.sparse.csr_array(([2,2.4,4,4.5,6,5.1,5.5,7.2,7.7], [0,3,1,4,2,0,3,1,4], [0,2,4,5,7,9])) + np.testing.assert_almost_equal(res.todense(), expected.todense()) \ No newline at end of file diff --git a/frontends/numpy-scipy/integration_tests/semiring/test_mm_SemiringMinSecond_CSRxCSR_oCSR.py b/frontends/numpy-scipy/integration_tests/semiring/test_mm_SemiringMinSecond_CSRxCSR_oCSR.py index cc7e907b..333c96f8 100644 --- a/frontends/numpy-scipy/integration_tests/semiring/test_mm_SemiringMinSecond_CSRxCSR_oCSR.py +++ b/frontends/numpy-scipy/integration_tests/semiring/test_mm_SemiringMinSecond_CSRxCSR_oCSR.py @@ -8,10 +8,10 @@ def run_comet(A,B): return C +def test_mm_SemiringMinSecond_CSRxCSR_oCSR(data_rank2_path): + A = sp.sparse.csr_array(sp.io.mmread(data_rank2_path)) + B = sp.sparse.csr_array(sp.io.mmread(data_rank2_path)) -A = sp.sparse.csr_array(sp.io.mmread("../../../../integration_test/data/test_rank2.mtx")) -B = sp.sparse.csr_array(sp.io.mmread("../../../../integration_test/data/test_rank2.mtx")) - -res = run_comet(A,B) -expected = sp.sparse.csr_array(([1,1.4,2,2.5,3,1,1.4,2,2.5], [0,3,1,4,2,0,3,1,4], [0,2,4,5,7,9])) -np.testing.assert_almost_equal(res.todense(), expected.todense()) \ No newline at end of file + res = run_comet(A,B) + expected = sp.sparse.csr_array(([1,1.4,2,2.5,3,1,1.4,2,2.5], [0,3,1,4,2,0,3,1,4], [0,2,4,5,7,9])) + np.testing.assert_almost_equal(res.todense(), expected.todense()) \ No newline at end of file diff --git a/frontends/numpy-scipy/integration_tests/semiring/test_mm_SemiringPlusFirst_CSRxCSR_oCSR.py b/frontends/numpy-scipy/integration_tests/semiring/test_mm_SemiringPlusFirst_CSRxCSR_oCSR.py index c3bf88c7..20d1f8d5 100644 --- a/frontends/numpy-scipy/integration_tests/semiring/test_mm_SemiringPlusFirst_CSRxCSR_oCSR.py +++ b/frontends/numpy-scipy/integration_tests/semiring/test_mm_SemiringPlusFirst_CSRxCSR_oCSR.py @@ -8,10 +8,10 @@ def run_comet(A,B): return C +def test_mm_SemiringPlusFirst_CSRxCSR_oCSR(data_rank2_path): + A = sp.sparse.csr_array(sp.io.mmread(data_rank2_path)) + B = sp.sparse.csr_array(sp.io.mmread(data_rank2_path)) -A = sp.sparse.csr_array(sp.io.mmread("../../../../integration_test/data/test_rank2.mtx")) -B = sp.sparse.csr_array(sp.io.mmread("../../../../integration_test/data/test_rank2.mtx")) - -res = run_comet(A,B) -expected = sp.sparse.csr_array(([2.4,2.4,4.5,4.5,3,8.1,8.1,10.2,10.2], [0,3,1,4,2,0,3,1,4], [0,2,4,5,7,9])) -np.testing.assert_almost_equal(res.todense(), expected.todense()) \ No newline at end of file + res = run_comet(A,B) + expected = sp.sparse.csr_array(([2.4,2.4,4.5,4.5,3,8.1,8.1,10.2,10.2], [0,3,1,4,2,0,3,1,4], [0,2,4,5,7,9])) + np.testing.assert_almost_equal(res.todense(), expected.todense()) \ No newline at end of file diff --git a/frontends/numpy-scipy/integration_tests/semiring/test_mm_SemiringPlusPair_CSRxCSR_oCSR.py b/frontends/numpy-scipy/integration_tests/semiring/test_mm_SemiringPlusPair_CSRxCSR_oCSR.py index 55e78965..f5a9d6b3 100644 --- a/frontends/numpy-scipy/integration_tests/semiring/test_mm_SemiringPlusPair_CSRxCSR_oCSR.py +++ b/frontends/numpy-scipy/integration_tests/semiring/test_mm_SemiringPlusPair_CSRxCSR_oCSR.py @@ -8,10 +8,10 @@ def run_comet(A,B): return C +def test_mm_SemiringPlusPair_CSRxCSR_oCSR(data_rank2_path): + A = sp.sparse.csr_array(sp.io.mmread(data_rank2_path)) + B = sp.sparse.csr_array(sp.io.mmread(data_rank2_path)) -A = sp.sparse.csr_array(sp.io.mmread("../../../../integration_test/data/test_rank2.mtx")) -B = sp.sparse.csr_array(sp.io.mmread("../../../../integration_test/data/test_rank2.mtx")) - -res = run_comet(A,B) -expected = sp.sparse.csr_array(([2,2,2,2,1,2,2,2,2], [0,3,1,4,2,0,3,1,4], [0,2,4,5,7,9])) -np.testing.assert_almost_equal(res.todense(), expected.todense()) \ No newline at end of file + res = run_comet(A,B) + expected = sp.sparse.csr_array(([2,2,2,2,1,2,2,2,2], [0,3,1,4,2,0,3,1,4], [0,2,4,5,7,9])) + np.testing.assert_almost_equal(res.todense(), expected.todense()) \ No newline at end of file diff --git a/frontends/numpy-scipy/integration_tests/semiring/test_mm_SemiringPlusSecond_CSRxCSR_oCSR.py b/frontends/numpy-scipy/integration_tests/semiring/test_mm_SemiringPlusSecond_CSRxCSR_oCSR.py index 413c8770..86eee2da 100644 --- a/frontends/numpy-scipy/integration_tests/semiring/test_mm_SemiringPlusSecond_CSRxCSR_oCSR.py +++ b/frontends/numpy-scipy/integration_tests/semiring/test_mm_SemiringPlusSecond_CSRxCSR_oCSR.py @@ -8,10 +8,10 @@ def run_comet(A,B): return C +def test_mm_SemiringPlusSecond_CSRxCSR_oCSR(data_rank2_path): + A = sp.sparse.csr_array(sp.io.mmread(data_rank2_path)) + B = sp.sparse.csr_array(sp.io.mmread(data_rank2_path)) -A = sp.sparse.csr_array(sp.io.mmread("../../../../integration_test/data/test_rank2.mtx")) -B = sp.sparse.csr_array(sp.io.mmread("../../../../integration_test/data/test_rank2.mtx")) - -res = run_comet(A,B) -expected = sp.sparse.csr_array(([5.1,5.4,7.2,7.5,3,5.1,5.4,7.2,7.5], [0,3,1,4,2,0,3,1,4], [0,2,4,5,7,9])) -np.testing.assert_almost_equal(res.todense(), expected.todense()) \ No newline at end of file + res = run_comet(A,B) + expected = sp.sparse.csr_array(([5.1,5.4,7.2,7.5,3,5.1,5.4,7.2,7.5], [0,3,1,4,2,0,3,1,4], [0,2,4,5,7,9])) + np.testing.assert_almost_equal(res.todense(), expected.todense()) \ No newline at end of file diff --git a/frontends/numpy-scipy/integration_tests/semiring/test_mm_SemiringPlusTime_CSRxDense_oDense.py b/frontends/numpy-scipy/integration_tests/semiring/test_mm_SemiringPlusTime_CSRxDense_oDense.py index 911a3c96..9b8a3ffe 100644 --- a/frontends/numpy-scipy/integration_tests/semiring/test_mm_SemiringPlusTime_CSRxDense_oDense.py +++ b/frontends/numpy-scipy/integration_tests/semiring/test_mm_SemiringPlusTime_CSRxDense_oDense.py @@ -8,10 +8,10 @@ def run_comet(A,B): return C +def test_mm_SemiringPlusTime_CSRxDense_oDense(data_rank2_path): + A = sp.sparse.csr_array(sp.io.mmread(data_rank2_path)) + B = np.full((A.shape[1]), 1.7) -A = sp.sparse.csr_array(sp.io.mmread("../../../../integration_test/data/test_rank2.mtx")) -B = np.full((A.shape[1]), 1.7) - -res = run_comet(A,B) -expected = np.array([4.08,7.65,5.1,13.77,17.34]).reshape((A.shape[0])) -np.testing.assert_almost_equal(res, expected) \ No newline at end of file + res = run_comet(A,B) + expected = np.array([4.08,7.65,5.1,13.77,17.34]).reshape((A.shape[0])) + np.testing.assert_almost_equal(res, expected) \ No newline at end of file diff --git a/frontends/numpy-scipy/integration_tests/semiring/test_mm_SemiringPlusTimes_CSRxCSR_oCSR.py b/frontends/numpy-scipy/integration_tests/semiring/test_mm_SemiringPlusTimes_CSRxCSR_oCSR.py index 3682dfd3..8ee08fd7 100644 --- a/frontends/numpy-scipy/integration_tests/semiring/test_mm_SemiringPlusTimes_CSRxCSR_oCSR.py +++ b/frontends/numpy-scipy/integration_tests/semiring/test_mm_SemiringPlusTimes_CSRxCSR_oCSR.py @@ -8,10 +8,10 @@ def run_comet(A,B): return C +def test_mm_SemiringPlusTimes_CSRxCSR_oCSR(data_rank2_path): + A = sp.sparse.csr_array(sp.io.mmread(data_rank2_path)) + B = sp.sparse.csr_array(sp.io.mmread(data_rank2_path)) -A = sp.sparse.csr_array(sp.io.mmread("../../../../integration_test/data/test_rank2.mtx")) -B = sp.sparse.csr_array(sp.io.mmread("../../../../integration_test/data/test_rank2.mtx")) - -res = run_comet(A,B) -expected = sp.sparse.csr_array(([6.74,7,17,17.5,9,20.5,21.74,36.4,38], [0,3,1,4,2,0,3,1,4], [0,2,4,5,7,9])) -np.testing.assert_almost_equal(res.todense(), expected.todense()) \ No newline at end of file + res = run_comet(A,B) + expected = sp.sparse.csr_array(([6.74,7,17,17.5,9,20.5,21.74,36.4,38], [0,3,1,4,2,0,3,1,4], [0,2,4,5,7,9])) + np.testing.assert_almost_equal(res.todense(), expected.todense()) \ No newline at end of file diff --git a/frontends/numpy-scipy/integration_tests/semiring/test_mm_SemiringPlusTimes_CSRxDense_oDense.py b/frontends/numpy-scipy/integration_tests/semiring/test_mm_SemiringPlusTimes_CSRxDense_oDense.py index afe09adb..aae0fd50 100644 --- a/frontends/numpy-scipy/integration_tests/semiring/test_mm_SemiringPlusTimes_CSRxDense_oDense.py +++ b/frontends/numpy-scipy/integration_tests/semiring/test_mm_SemiringPlusTimes_CSRxDense_oDense.py @@ -4,14 +4,14 @@ @comet.compile(flags=None) def run_comet(A,B): - C = comet.einsum('ij,jk->ik', A,B,semiring='+,*') + C = comet.einsum('ij,jk->ik', A,B, '+,*') return C +def test_mm_SemiringPlusTimes_CSRxCSR_oCSR(data_rank2_path): + A = sp.sparse.csr_array(sp.io.mmread(data_rank2_path)) + B = np.full((A.shape[1], 4), 1.7) -A = sp.sparse.csr_array(sp.io.mmread("../../../../integration_test/data/test_rank2.mtx")) -B = np.full((A.shape[1], 4), 1.7) - -res = run_comet(A,B) -expected = np.array([4.08,4.08,4.08,4.08,7.65,7.65,7.65,7.65,5.1,5.1,5.1,5.1,13.77,13.77,13.77,13.77,17.34,17.34,17.34,17.34]).reshape((A.shape[0], 4)) -np.testing.assert_almost_equal(res, expected) \ No newline at end of file + res = run_comet(A,B) + expected = np.array([4.08,4.08,4.08,4.08,7.65,7.65,7.65,7.65,5.1,5.1,5.1,5.1,13.77,13.77,13.77,13.77,17.34,17.34,17.34,17.34]).reshape((A.shape[0], 4)) + np.testing.assert_almost_equal(res, expected) \ No newline at end of file diff --git a/frontends/numpy-scipy/integration_tests/semiring/test_mm_SemiringPlusTimes_DensexCOO_oDense.py b/frontends/numpy-scipy/integration_tests/semiring/test_mm_SemiringPlusTimes_DensexCOO_oDense.py index 95859748..01523bf1 100644 --- a/frontends/numpy-scipy/integration_tests/semiring/test_mm_SemiringPlusTimes_DensexCOO_oDense.py +++ b/frontends/numpy-scipy/integration_tests/semiring/test_mm_SemiringPlusTimes_DensexCOO_oDense.py @@ -8,10 +8,10 @@ def run_comet(A,B): return C +def test_mm_SemiringPlusTimes_DensexCOO_oDense(data_rank2_path): + B = sp.sparse.coo_array(sp.io.mmread(data_rank2_path)) + A = np.full((4, B.shape[0]), 1.7) -B = sp.sparse.coo_array(sp.io.mmread("../../../../integration_test/data/test_rank2.mtx")) -A = np.full((4, B.shape[0]), 1.7) - -res = run_comet(A,B) -expected = np.array([8.67,12.24,5.1,9.18,12.75,8.67,12.24,5.1,9.18,12.75,8.67,12.24,5.1,9.18,12.75,8.67,12.24,5.1,9.18,12.75]).reshape((4, B.shape[1])) -np.testing.assert_almost_equal(res, expected) \ No newline at end of file + res = run_comet(A,B) + expected = np.array([8.67,12.24,5.1,9.18,12.75,8.67,12.24,5.1,9.18,12.75,8.67,12.24,5.1,9.18,12.75,8.67,12.24,5.1,9.18,12.75]).reshape((4, B.shape[1])) + np.testing.assert_almost_equal(res, expected) \ No newline at end of file diff --git a/frontends/numpy-scipy/integration_tests/semiring/test_mm_SemiringPlusTimes_DensexCSR_oDense.py b/frontends/numpy-scipy/integration_tests/semiring/test_mm_SemiringPlusTimes_DensexCSR_oDense.py index 6144c6a4..d3182e80 100644 --- a/frontends/numpy-scipy/integration_tests/semiring/test_mm_SemiringPlusTimes_DensexCSR_oDense.py +++ b/frontends/numpy-scipy/integration_tests/semiring/test_mm_SemiringPlusTimes_DensexCSR_oDense.py @@ -8,10 +8,10 @@ def run_comet(A,B): return C +def test_mm_SemiringPlusTimes_DensexCSR_oDense(data_rank2_path): + B = sp.sparse.csr_array(sp.io.mmread(data_rank2_path)) + A = np.full((4, B.shape[0]), 1.7) -B = sp.sparse.csr_array(sp.io.mmread("../../../../integration_test/data/test_rank2.mtx")) -A = np.full((4, B.shape[0]), 1.7) - -res = run_comet(A,B) -expected = np.array([8.67,12.24,5.1,9.18,12.75,8.67,12.24,5.1,9.18,12.75,8.67,12.24,5.1,9.18,12.75,8.67,12.24,5.1,9.18,12.75]).reshape((4, B.shape[1])) -np.testing.assert_almost_equal(res, expected) \ No newline at end of file + res = run_comet(A,B) + expected = np.array([8.67,12.24,5.1,9.18,12.75,8.67,12.24,5.1,9.18,12.75,8.67,12.24,5.1,9.18,12.75,8.67,12.24,5.1,9.18,12.75]).reshape((4, B.shape[1])) + np.testing.assert_almost_equal(res, expected) \ No newline at end of file diff --git a/frontends/numpy-scipy/integration_tests/semiring/test_mm_SemiringPlusTimes_DensexDense_oDense.py b/frontends/numpy-scipy/integration_tests/semiring/test_mm_SemiringPlusTimes_DensexDense_oDense.py index 834b449d..66c4360e 100644 --- a/frontends/numpy-scipy/integration_tests/semiring/test_mm_SemiringPlusTimes_DensexDense_oDense.py +++ b/frontends/numpy-scipy/integration_tests/semiring/test_mm_SemiringPlusTimes_DensexDense_oDense.py @@ -4,14 +4,14 @@ @comet.compile(flags=None) def run_comet(A,B): - C = comet.einsum('ij,jk->ik', A,B,semiring='+,*') + C = comet.einsum('ij,jk->ik', A,B, semiring='+,*') return C +def test_mm_SemiringPlusTimes_DensexDense_oDense(): + A = np.full((8, 4), 2.2) + B = np.full((4, 2), 3.4) -A = np.full((8, 4), 2.2) -B = np.full((4, 2), 3.4) - -res = run_comet(A,B) -expected = np.array([29.92,29.92,29.92,29.92,29.92,29.92,29.92,29.92,29.92,29.92,29.92,29.92,29.92,29.92,29.92,29.92]).reshape((8, 2)) -np.testing.assert_almost_equal(res, expected) \ No newline at end of file + res = run_comet(A,B) + expected = np.array([29.92,29.92,29.92,29.92,29.92,29.92,29.92,29.92,29.92,29.92,29.92,29.92,29.92,29.92,29.92,29.92]).reshape((8, 2)) + np.testing.assert_almost_equal(res, expected) \ No newline at end of file diff --git a/frontends/numpy-scipy/integration_tests/semiring/test_mm_SemiringPlusTimes_dense_4Dtensors.py b/frontends/numpy-scipy/integration_tests/semiring/test_mm_SemiringPlusTimes_dense_4Dtensors.py index 66691cfc..11d51ef8 100644 --- a/frontends/numpy-scipy/integration_tests/semiring/test_mm_SemiringPlusTimes_dense_4Dtensors.py +++ b/frontends/numpy-scipy/integration_tests/semiring/test_mm_SemiringPlusTimes_dense_4Dtensors.py @@ -8,10 +8,10 @@ def run_comet(A,B): return C +def test_mm_SemiringPlusTimes_dense_4Dtensors(): + A = np.full((2,2,2,2), 2.2) + B = np.full((2,2,2,2), 3.6) -A = np.full((2,2,2,2), 2.2) -B = np.full((2,2,2,2), 3.6) - -res = run_comet(A,B) -expected = np.array([31.68,31.68,31.68,31.68,31.68,31.68,31.68,31.68,31.68,31.68,31.68,31.68,31.68,31.68,31.68,31.68]).reshape((2,2,2,2)) -np.testing.assert_almost_equal(res, expected) \ No newline at end of file + res = run_comet(A,B) + expected = np.array([31.68,31.68,31.68,31.68,31.68,31.68,31.68,31.68,31.68,31.68,31.68,31.68,31.68,31.68,31.68,31.68]).reshape((2,2,2,2)) + np.testing.assert_almost_equal(res, expected) \ No newline at end of file diff --git a/frontends/numpy-scipy/integration_tests/semiring/test_mv_SemiringPlusTimes_COOxDense_oDense.py b/frontends/numpy-scipy/integration_tests/semiring/test_mv_SemiringPlusTimes_COOxDense_oDense.py index b100d86f..9c28b4e9 100644 --- a/frontends/numpy-scipy/integration_tests/semiring/test_mv_SemiringPlusTimes_COOxDense_oDense.py +++ b/frontends/numpy-scipy/integration_tests/semiring/test_mv_SemiringPlusTimes_COOxDense_oDense.py @@ -8,10 +8,10 @@ def run_comet(A,B): return C +def test_mv_SemiringPlusTimes_COOxDense_oDense(data_rank2_path): + A = sp.sparse.coo_array(sp.io.mmread(data_rank2_path)) + B = np.full((A.shape[1]), 1.7) -A = sp.sparse.coo_array(sp.io.mmread("../../../../integration_test/data/test_rank2.mtx")) -B = np.full((A.shape[1]), 1.7) - -res = run_comet(A,B) -expected = np.array([4.08,7.65,5.1,13.77,17.34]).reshape((A.shape[0])) -np.testing.assert_almost_equal(res, expected) \ No newline at end of file + res = run_comet(A,B) + expected = np.array([4.08,7.65,5.1,13.77,17.34]).reshape((A.shape[0])) + np.testing.assert_almost_equal(res, expected) \ No newline at end of file diff --git a/frontends/numpy-scipy/integration_tests/semiring/test_mv_SemiringPlusTimes_CSRxDense_oDense.py b/frontends/numpy-scipy/integration_tests/semiring/test_mv_SemiringPlusTimes_CSRxDense_oDense.py index 911a3c96..e119fba4 100644 --- a/frontends/numpy-scipy/integration_tests/semiring/test_mv_SemiringPlusTimes_CSRxDense_oDense.py +++ b/frontends/numpy-scipy/integration_tests/semiring/test_mv_SemiringPlusTimes_CSRxDense_oDense.py @@ -8,10 +8,10 @@ def run_comet(A,B): return C +def test_mv_SemiringPlusTimes_CSRxDense_oDense(data_rank2_path): + A = sp.sparse.csr_array(sp.io.mmread(data_rank2_path)) + B = np.full((A.shape[1]), 1.7) -A = sp.sparse.csr_array(sp.io.mmread("../../../../integration_test/data/test_rank2.mtx")) -B = np.full((A.shape[1]), 1.7) - -res = run_comet(A,B) -expected = np.array([4.08,7.65,5.1,13.77,17.34]).reshape((A.shape[0])) -np.testing.assert_almost_equal(res, expected) \ No newline at end of file + res = run_comet(A,B) + expected = np.array([4.08,7.65,5.1,13.77,17.34]).reshape((A.shape[0])) + np.testing.assert_almost_equal(res, expected) \ No newline at end of file diff --git a/frontends/numpy-scipy/integration_tests/semiring/test_mv_SemiringPlusTimes_DensexCOO_oDense.py b/frontends/numpy-scipy/integration_tests/semiring/test_mv_SemiringPlusTimes_DensexCOO_oDense.py index daf89f12..675f40d0 100644 --- a/frontends/numpy-scipy/integration_tests/semiring/test_mv_SemiringPlusTimes_DensexCOO_oDense.py +++ b/frontends/numpy-scipy/integration_tests/semiring/test_mv_SemiringPlusTimes_DensexCOO_oDense.py @@ -8,10 +8,10 @@ def run_comet(A,B): return C +def test_mv_SemiringPlusTimes_DensexCOO_oDense(data_rank2_path): + B = sp.sparse.coo_array(sp.io.mmread(data_rank2_path)) + A = np.full((B.shape[0]), 1.7) -B = sp.sparse.coo_array(sp.io.mmread("../../../../integration_test/data/test_rank2.mtx")) -A = np.full((B.shape[0]), 1.7) - -res = run_comet(A,B) -expected = np.array([8.67,12.24,5.1,9.18,12.75]).reshape((B.shape[1])) -np.testing.assert_almost_equal(res, expected) \ No newline at end of file + res = run_comet(A,B) + expected = np.array([8.67,12.24,5.1,9.18,12.75]).reshape((B.shape[1])) + np.testing.assert_almost_equal(res, expected) \ No newline at end of file diff --git a/frontends/numpy-scipy/integration_tests/semiring/test_mv_SemiringPlusTimes_DensexCSR_oDense.py b/frontends/numpy-scipy/integration_tests/semiring/test_mv_SemiringPlusTimes_DensexCSR_oDense.py index 93adbc47..0a7f5b19 100644 --- a/frontends/numpy-scipy/integration_tests/semiring/test_mv_SemiringPlusTimes_DensexCSR_oDense.py +++ b/frontends/numpy-scipy/integration_tests/semiring/test_mv_SemiringPlusTimes_DensexCSR_oDense.py @@ -8,10 +8,10 @@ def run_comet(A,B): return C +def test_mv_SemiringPlusTimes_DensexCSR_oDense(data_rank2_path): + B = sp.sparse.csr_array(sp.io.mmread(data_rank2_path)) + A = np.full((B.shape[0]), 1.7) -B = sp.sparse.csr_array(sp.io.mmread("../../../../integration_test/data/test_rank2.mtx")) -A = np.full((B.shape[0]), 1.7) - -res = run_comet(A,B) -expected = np.array([8.67,12.24,5.1,9.18,12.75]).reshape((B.shape[1])) -np.testing.assert_almost_equal(res, expected) \ No newline at end of file + res = run_comet(A,B) + expected = np.array([8.67,12.24,5.1,9.18,12.75]).reshape((B.shape[1])) + np.testing.assert_almost_equal(res, expected) \ No newline at end of file diff --git a/frontends/numpy-scipy/integration_tests/semiring/test_mv_SemiringPlusTimes_DensexDense_oDense.py b/frontends/numpy-scipy/integration_tests/semiring/test_mv_SemiringPlusTimes_DensexDense_oDense.py index 010be253..2fb6e100 100644 --- a/frontends/numpy-scipy/integration_tests/semiring/test_mv_SemiringPlusTimes_DensexDense_oDense.py +++ b/frontends/numpy-scipy/integration_tests/semiring/test_mv_SemiringPlusTimes_DensexDense_oDense.py @@ -8,10 +8,10 @@ def run_comet(A,B): return C +def test_mv_SemiringPlusTimes_DensexDense_oDense(): + A = np.full((8,16), 2.3) + B = np.full((16), 3.7) -A = np.full((8,16), 2.3) -B = np.full((16), 3.7) - -res = run_comet(A,B) -expected = np.array([136.16,136.16,136.16,136.16,136.16,136.16,136.16,136.16]).reshape((8)) -np.testing.assert_almost_equal(res, expected) \ No newline at end of file + res = run_comet(A,B) + expected = np.array([136.16,136.16,136.16,136.16,136.16,136.16,136.16,136.16]).reshape((8)) + np.testing.assert_almost_equal(res, expected) \ No newline at end of file diff --git a/frontends/numpy-scipy/pytest.ini b/frontends/numpy-scipy/pytest.ini new file mode 100644 index 00000000..d2aace7c --- /dev/null +++ b/frontends/numpy-scipy/pytest.ini @@ -0,0 +1,2 @@ +[pytest] +norecursedirs = frontends/numpy-scipy/integration_tests/not_supported \ No newline at end of file diff --git a/frontends/numpy-scipy/setup.py b/frontends/numpy-scipy/setup.py index 6cde2bb2..43254fa1 100644 --- a/frontends/numpy-scipy/setup.py +++ b/frontends/numpy-scipy/setup.py @@ -56,7 +56,8 @@ install_requires=[ 'jinja2', 'numpy', - 'scipy>=1.14' + 'scipy>=1.14', + 'ast-comments>=1.2.2', ], python_requires=">=3.8", ) diff --git a/include/comet/Conversion/BlockedGpuToTriton/BlockedGpuToTriton.h b/include/comet/Conversion/BlockedGpuToTriton/BlockedGpuToTriton.h new file mode 100644 index 00000000..a3ca6057 --- /dev/null +++ b/include/comet/Conversion/BlockedGpuToTriton/BlockedGpuToTriton.h @@ -0,0 +1,17 @@ +#ifndef COMET_CONVERSION_BLOCKEDGPUTOTRITON_H +#define COMET_CONVERSION_BLOCKEDGPUTOTRITON_H + +#include "mlir/Dialect/GPU/IR/GPUDialect.h" +#include +namespace mlir { +class ModuleOp; +namespace gpu{ +class GPUFuncOp; +} +template class OperationPass; +namespace comet { +std::unique_ptr > createConvertBlockedGpuToTritonPass(); +} +} + +#endif \ No newline at end of file diff --git a/include/comet/Conversion/BlockedGpuToTriton/BlockedGpuToTritonConversion.h b/include/comet/Conversion/BlockedGpuToTriton/BlockedGpuToTritonConversion.h new file mode 100644 index 00000000..249dbc56 --- /dev/null +++ b/include/comet/Conversion/BlockedGpuToTriton/BlockedGpuToTritonConversion.h @@ -0,0 +1,22 @@ + +#include "mlir/IR/MLIRContext.h" +#include "mlir/Transforms/DialectConversion.h" +namespace mlir { +namespace comet{ + +class TritonTypeConverter : public TypeConverter { +public: + TritonTypeConverter(MLIRContext *context); +private: + MLIRContext *context; +}; + +class TritonConversionTarget : public ConversionTarget { + +public: + explicit TritonConversionTarget(MLIRContext &ctx, + TritonTypeConverter &typeConverter); +}; + +} +} \ No newline at end of file diff --git a/include/comet/Conversion/BlockedGpuToTriton/CMakeLists.txt b/include/comet/Conversion/BlockedGpuToTriton/CMakeLists.txt new file mode 100644 index 00000000..e64f348e --- /dev/null +++ b/include/comet/Conversion/BlockedGpuToTriton/CMakeLists.txt @@ -0,0 +1,3 @@ +set(LLVM_TARGET_DEFINITIONS Passes.td) +mlir_tablegen(Passes.h.inc -gen-pass-decls --name BlockedGpuToTriton) +add_public_tablegen_target(BlockedGpuToTritonPassIncGen) \ No newline at end of file diff --git a/include/comet/Conversion/BlockedGpuToTriton/Passes.h b/include/comet/Conversion/BlockedGpuToTriton/Passes.h new file mode 100644 index 00000000..a9474079 --- /dev/null +++ b/include/comet/Conversion/BlockedGpuToTriton/Passes.h @@ -0,0 +1,13 @@ +#ifndef COMET_BLOCKED_GPU_TO_TRITON_CONVERSION_PASSES +#define COMET_BLOCKED_GPU_TO_TRITON_CONVERSION_PASSES + +#include "comet/Conversion/BlockedGpuToTriton/BlockedGpuToTriton.h" + +namespace mlir { +namespace comet { +#define GEN_PASS_REGISTRATION +#include "comet/Conversion/BlockedGpuToTriton/Passes.h.inc" +} +} + +#endif \ No newline at end of file diff --git a/include/comet/Conversion/BlockedGpuToTriton/Passes.td b/include/comet/Conversion/BlockedGpuToTriton/Passes.td new file mode 100644 index 00000000..1138876e --- /dev/null +++ b/include/comet/Conversion/BlockedGpuToTriton/Passes.td @@ -0,0 +1,14 @@ +#ifndef COMET_BLOCKED_GPU_TO_TRITON_CONVERSION_PASSES +#define COMET_BLOCKED_GPU_TO_TRITON_CONVERSION_PASSES + +include "mlir/Pass/PassBase.td" + +def CometBlockedGpuToTriton : Pass<"comet-blocked-gpu-to-triton", "mlir::ModuleOp"> { + let summary = "Map blocked gpu to trtion"; + + let constructor = "mlir::comet::createConvertBlockedGpuToTriton()"; + + let dependentDialects = ["mlir::scf::SCFDialect", "mlir::gpu::GPUDialect", "mlir::tensor::TensorDialect"]; +} + +#endif diff --git a/include/comet/Conversion/CMakeLists.txt b/include/comet/Conversion/CMakeLists.txt index b866deb7..28e7ce84 100644 --- a/include/comet/Conversion/CMakeLists.txt +++ b/include/comet/Conversion/CMakeLists.txt @@ -1,7 +1,18 @@ +if(ENABLE_GPU_TARGET OR ENABLE_FPGA_TARGET) +add_subdirectory(ForallToGpu) +add_subdirectory(ParallelLoopsToGpuFPGA) +endif() +if(ENABLE_FPGA_TARGET) +add_subdirectory(GpuToOCLSPIRV) +add_subdirectory(GpuHostToMCLRT) +endif() if(ENABLE_GPU_TARGET) -add_subdirectory(ParallelLoopsToGpu) add_subdirectory(GpuToTriton) +add_subdirectory(GpuToBlockedGpu) +add_subdirectory(BlockedGpuToTriton) add_subdirectory(TritonToCuda) +add_subdirectory(TritonToHIP) +add_subdirectory(PrepareGpuHost) endif() set(LLVM_TARGET_DEFINITIONS Passes.td) mlir_tablegen(Passes.h.inc -gen-pass-decls -name Conversion) diff --git a/include/comet/Conversion/ForallToGpu/CMakeLists.txt b/include/comet/Conversion/ForallToGpu/CMakeLists.txt new file mode 100644 index 00000000..f73b9b7c --- /dev/null +++ b/include/comet/Conversion/ForallToGpu/CMakeLists.txt @@ -0,0 +1,3 @@ +set(LLVM_TARGET_DEFINITIONS Passes.td) +mlir_tablegen(Passes.h.inc -gen-pass-decls --name ForallToGpu) +add_public_tablegen_target(ForallConversionPassIncGen) \ No newline at end of file diff --git a/include/comet/Conversion/ForallToGpu/ForallToGpu.h b/include/comet/Conversion/ForallToGpu/ForallToGpu.h new file mode 100644 index 00000000..4efd767d --- /dev/null +++ b/include/comet/Conversion/ForallToGpu/ForallToGpu.h @@ -0,0 +1,38 @@ +// +// Copyright 2022 Battelle Memorial Institute +// +// Redistribution and use in source and binary forms, with or without modification, +// are permitted provided that the following conditions are met: +// +// 1. Redistributions of source code must retain the above copyright notice, this list of conditions +// and the following disclaimer. +// +// 2. Redistributions in binary form must reproduce the above copyright notice, this list of conditions +// and the following disclaimer in the documentation and/or other materials provided with the distribution. +// +// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED +// WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR +// PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, +// INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE +// GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, +// WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +// OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. +// + +#ifndef COMET_CONVERSION_FORALLTOGPU_H +#define COMET_CONVERSION_FORALLTOGPU_H + +#include +namespace mlir { +class ModuleOp; +namespace func{ +class FuncOp; +} +template class OperationPass; +namespace comet { +std::unique_ptr> createConvertForallToGpuPass(); +std::unique_ptr> createConvertForallToGpuPass(int blockX, int blockY, int blockR); +} +} + +#endif \ No newline at end of file diff --git a/include/comet/Conversion/ForallToGpu/Passes.h b/include/comet/Conversion/ForallToGpu/Passes.h new file mode 100644 index 00000000..bc5a1536 --- /dev/null +++ b/include/comet/Conversion/ForallToGpu/Passes.h @@ -0,0 +1,34 @@ +// +// Copyright 2022 Battelle Memorial Institute +// +// Redistribution and use in source and binary forms, with or without modification, +// are permitted provided that the following conditions are met: +// +// 1. Redistributions of source code must retain the above copyright notice, this list of conditions +// and the following disclaimer. +// +// 2. Redistributions in binary form must reproduce the above copyright notice, this list of conditions +// and the following disclaimer in the documentation and/or other materials provided with the distribution. +// +// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED +// WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR +// PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, +// INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE +// GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, +// WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +// OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. +// + +#ifndef COMET_FORALL_TO_GPU_CONVERSION_PASSES +#define COMET_FORALL_TO_GPU_CONVERSION_PASSES + +#include "comet/Conversion/ForallToGpu/ForallToGpu.h" + +namespace mlir { +namespace comet { +#define GEN_PASS_REGISTRATION +#include "comet/Conversion/ForallToGpu/Passes.h.inc" +} +} + +#endif \ No newline at end of file diff --git a/include/comet/Conversion/ForallToGpu/Passes.td b/include/comet/Conversion/ForallToGpu/Passes.td new file mode 100644 index 00000000..60cae3f1 --- /dev/null +++ b/include/comet/Conversion/ForallToGpu/Passes.td @@ -0,0 +1,50 @@ +// +// Copyright 2022 Battelle Memorial Institute +// +// Redistribution and use in source and binary forms, with or without modification, +// are permitted provided that the following conditions are met: +// +// 1. Redistributions of source code must retain the above copyright notice, this list of conditions +// and the following disclaimer. +// +// 2. Redistributions in binary form must reproduce the above copyright notice, this list of conditions +// and the following disclaimer in the documentation and/or other materials provided with the distribution. +// +// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED +// WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR +// PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, +// INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE +// GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, +// WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +// OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. +// + +#ifndef COMET_Forall_TO_GPU_CONVERSION_PASSES +#define COMET_Forall_TO_GPU_CONVERSION_PASSES + +include "mlir/Pass/PassBase.td" + +def CometForallToGpu : Pass<"comet-forall-to-gpu", "mlir::func::FuncOp"> { + let summary = "Map loops to gpu"; + + let constructor = "mlir::comet::createConvertForallToGpuPass()"; + + let dependentDialects = ["mlir::scf::SCFDialect", "mlir::gpu::GPUDialect", "mlir::affine::AffineDialect", "mlir::DLTIDialect"]; + + + let options = [ + Option<"blockX", "blockX", + "int32_t", /*default*/"32", + "triton block size for dim X">, + + Option<"blockY", "blockY", + "int32_t", /*default*/"8", + "triton block size for dim Y">, + + Option<"blockR", "blockR", + "int32_t", /*default*/"32", + "triton block size for reduction dim">, + ]; +} + +#endif diff --git a/include/comet/Conversion/GpuHostToMCLRT/CMakeLists.txt b/include/comet/Conversion/GpuHostToMCLRT/CMakeLists.txt new file mode 100644 index 00000000..1a9dd0f1 --- /dev/null +++ b/include/comet/Conversion/GpuHostToMCLRT/CMakeLists.txt @@ -0,0 +1,3 @@ +set(LLVM_TARGET_DEFINITIONS Passes.td) +mlir_tablegen(Passes.h.inc -gen-pass-decls --name GpuHostToMCLRT) +add_public_tablegen_target(GpuHostToMCLRTConversionPassIncGen) \ No newline at end of file diff --git a/include/comet/Conversion/GpuHostToMCLRT/GpuHostToMCLRTPass.h b/include/comet/Conversion/GpuHostToMCLRT/GpuHostToMCLRTPass.h new file mode 100644 index 00000000..b1f09448 --- /dev/null +++ b/include/comet/Conversion/GpuHostToMCLRT/GpuHostToMCLRTPass.h @@ -0,0 +1,39 @@ +// +// Copyright 2022 Battelle Memorial Institute +// +// Redistribution and use in source and binary forms, with or without modification, +// are permitted provided that the following conditions are met: +// +// 1. Redistributions of source code must retain the above copyright notice, this list of conditions +// and the following disclaimer. +// +// 2. Redistributions in binary form must reproduce the above copyright notice, this list of conditions +// and the following disclaimer in the documentation and/or other materials provided with the distribution. +// +// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED +// WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR +// PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, +// INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE +// GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, +// WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +// OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. +// + +#ifndef COMET_CONVERSION_GPUHOSTTOMCLRT_H +#define COMET_CONVERSION_GPUHOSTTOMCLRT_H + +#include +namespace mlir { +class ModuleOp; +namespace func{ +class FuncOp; +} + +template class OperationPass; +namespace comet { +std::unique_ptr> createConvertGpuHostToMCLRTPass(); +std::unique_ptr> createConvertGpuHostToMCLRTPass(const char* xclbin_path); +} +} + +#endif \ No newline at end of file diff --git a/include/comet/Conversion/GpuHostToMCLRT/Passes.td b/include/comet/Conversion/GpuHostToMCLRT/Passes.td new file mode 100644 index 00000000..65d0164b --- /dev/null +++ b/include/comet/Conversion/GpuHostToMCLRT/Passes.td @@ -0,0 +1,43 @@ +// +// Copyright 2022 Battelle Memorial Institute +// +// Redistribution and use in source and binary forms, with or without modification, +// are permitted provided that the following conditions are met: +// +// 1. Redistributions of source code must retain the above copyright notice, this list of conditions +// and the following disclaimer. +// +// 2. Redistributions in binary form must reproduce the above copyright notice, this list of conditions +// and the following disclaimer in the documentation and/or other materials provided with the distribution. +// +// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED +// WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR +// PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, +// INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE +// GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, +// WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +// OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. +// + +#ifndef GPUHOST_TO_MCLRT_CONVERSION_PASSES +#define GPUHOST_TO_MCLRT_CONVERSION_PASSES + +include "mlir/Pass/PassBase.td" + +def ConvertGpuHostToMCLRTPass: Pass<"convert-gpu-host-to-mcl-rt", "mlir::ModuleOp"> { + let summary = "Convert Gpu Host to MCR Runtime Calls"; + let description = [{ + + }]; + let constructor = "mlir::comet::createConvertGpuHostToMCLRTPass()"; + + let dependentDialects = ["mlir::gpu::GPUDialect", "mlir::LLVM::LLVMDialect", "mlir::memref::MemRefDialect", "mlir::scf::SCFDialect", + "mlir::spirv::SPIRVDialect"]; + + let options = [ + Option<"xclbin_path", "xclbin-path", + "const char*", /*default*/"\"-\"", + "Path to xclbin">, + ]; +} +#endif diff --git a/include/comet/Conversion/GpuToBlockedGpu/CMakeLists.txt b/include/comet/Conversion/GpuToBlockedGpu/CMakeLists.txt new file mode 100644 index 00000000..a5a65596 --- /dev/null +++ b/include/comet/Conversion/GpuToBlockedGpu/CMakeLists.txt @@ -0,0 +1,3 @@ +set(LLVM_TARGET_DEFINITIONS Passes.td) +mlir_tablegen(Passes.h.inc -gen-pass-decls --name GpuToBlockedGpu) +add_public_tablegen_target(GpuToBlockedGpuPassIncGen) \ No newline at end of file diff --git a/include/comet/Conversion/GpuToBlockedGpu/GpuToBlockedGpu.h b/include/comet/Conversion/GpuToBlockedGpu/GpuToBlockedGpu.h new file mode 100644 index 00000000..c6c6c5e1 --- /dev/null +++ b/include/comet/Conversion/GpuToBlockedGpu/GpuToBlockedGpu.h @@ -0,0 +1,16 @@ +#ifndef COMET_CONVERSION_GPUTOBLOCKEDGPU_H +#define COMET_CONVERSION_GPUTOBLOCKEDGPU_H + +#include +namespace mlir { +class ModuleOp; +namespace gpu{ +class GPUFuncOp; +} +template class OperationPass; +namespace comet { +std::unique_ptr > createConvertGpuToBlockedGpuPass(); +} +} + +#endif \ No newline at end of file diff --git a/include/comet/Conversion/GpuToBlockedGpu/GpuToBlockedGpuConversion.h b/include/comet/Conversion/GpuToBlockedGpu/GpuToBlockedGpuConversion.h new file mode 100644 index 00000000..e69de29b diff --git a/include/comet/Conversion/GpuToBlockedGpu/Passes.h b/include/comet/Conversion/GpuToBlockedGpu/Passes.h new file mode 100644 index 00000000..eb8d8d4d --- /dev/null +++ b/include/comet/Conversion/GpuToBlockedGpu/Passes.h @@ -0,0 +1,13 @@ +#ifndef COMET_GPU_TO_BLOCKED_GPU_CONVERSION_PASSES +#define COMET_GPU_TO_BLOCKED_GPU_CONVERSION_PASSES + +#include "comet/Conversion/GpuToBlockedGpu/GpuToBlockedGpu.h" + +namespace mlir { +namespace comet { +#define GEN_PASS_REGISTRATION +#include "comet/Conversion/GpuToBlockedGpu/Passes.h.inc" +} +} + +#endif \ No newline at end of file diff --git a/include/comet/Conversion/GpuToBlockedGpu/Passes.td b/include/comet/Conversion/GpuToBlockedGpu/Passes.td new file mode 100644 index 00000000..9579bda0 --- /dev/null +++ b/include/comet/Conversion/GpuToBlockedGpu/Passes.td @@ -0,0 +1,14 @@ +#ifndef COMET_GPU_TO_BLOCKED_GPU_CONVERSION_PASSES +#define COMET_GPU_TO_BLOCKED_GPU_CONVERSION_PASSES + +include "mlir/Pass/PassBase.td" + +def CometGpuToBlockedGpu : Pass<"comet-gpu-to-blocked-gpu", "mlir::gpu::GPUFuncOp"> { + let summary = "Map loops to gpu"; + + let constructor = "mlir::comet::createConvertGpuToBlockedGpuPass()"; + + let dependentDialects = ["mlir::scf::SCFDialect", "mlir::affine::AffineDialect", "mlir::gpu::GPUDialect", "mlir::vector::VectorDialect", "mlir::tensor::TensorDialect"]; +} + +#endif diff --git a/include/comet/Conversion/GpuToOCLSPIRV/CMakeLists.txt b/include/comet/Conversion/GpuToOCLSPIRV/CMakeLists.txt new file mode 100644 index 00000000..b8b48255 --- /dev/null +++ b/include/comet/Conversion/GpuToOCLSPIRV/CMakeLists.txt @@ -0,0 +1,3 @@ +set(LLVM_TARGET_DEFINITIONS Passes.td) +mlir_tablegen(Passes.h.inc -gen-pass-decls --name GpuToOCLSPIRV) +add_public_tablegen_target(GpuToOCLSPIRVConversionPassIncGen) \ No newline at end of file diff --git a/include/comet/Conversion/GpuToOCLSPIRV/GPUToSPIRVPass.h b/include/comet/Conversion/GpuToOCLSPIRV/GPUToSPIRVPass.h new file mode 100644 index 00000000..43a149fb --- /dev/null +++ b/include/comet/Conversion/GpuToOCLSPIRV/GPUToSPIRVPass.h @@ -0,0 +1,35 @@ +//===- GPUToSPIRVPass.h - GPU to SPIR-V Passes ------------------*- C++ -*-===// +// +// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. +// See https://llvm.org/LICENSE.txt for license information. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception +// +//===----------------------------------------------------------------------===// +// +// Provides passes to convert GPU dialect to SPIR-V dialect. +// +//===----------------------------------------------------------------------===// + +#ifndef MLIR_CONVERSION_GPUTOSPIRV_GPUTOSPIRVPASS_H +#define MLIR_CONVERSION_GPUTOSPIRV_GPUTOSPIRVPASS_H + +#include + +namespace mlir { + +class ModuleOp; +template +class OperationPass; + +#define GEN_PASS_DECL_CONVERTGPUTOSPIRV +#include "mlir/Conversion/Passes.h.inc" + +/// Creates a pass to convert GPU kernel ops to corresponding SPIR-V ops. For a +/// gpu.func to be converted, it should have a spirv.entry_point_abi attribute. +/// If `mapMemorySpace` is true, performs MemRef memory space to SPIR-V mapping +/// according to default Vulkan rules first. +std::unique_ptr> +createConvertGPUToSPIRVPass2(bool mapMemorySpace = true, bool use64bitIndex = false); + +} // namespace mlir +#endif // MLIR_CONVERSION_GPUTOSPIRV_GPUTOSPIRVPASS_H diff --git a/include/comet/Conversion/GpuToOCLSPIRV/GpuToOCLSPIRVPass.h b/include/comet/Conversion/GpuToOCLSPIRV/GpuToOCLSPIRVPass.h new file mode 100644 index 00000000..58d8e81f --- /dev/null +++ b/include/comet/Conversion/GpuToOCLSPIRV/GpuToOCLSPIRVPass.h @@ -0,0 +1,39 @@ +// +// Copyright 2022 Battelle Memorial Institute +// +// Redistribution and use in source and binary forms, with or without modification, +// are permitted provided that the following conditions are met: +// +// 1. Redistributions of source code must retain the above copyright notice, this list of conditions +// and the following disclaimer. +// +// 2. Redistributions in binary form must reproduce the above copyright notice, this list of conditions +// and the following disclaimer in the documentation and/or other materials provided with the distribution. +// +// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED +// WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR +// PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, +// INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE +// GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, +// WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +// OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. +// + +#ifndef COMET_CONVERSION_GPUKERNELTOOCLSPIRV_H +#define COMET_CONVERSION_GPUKERNELTOOCLSPIRV_H + +#include +namespace mlir { +class ModuleOp; +namespace func{ +class FuncOp; +} + +template class OperationPass; +namespace comet { +std::unique_ptr> createConvertGPUKernelToOCLSPIRVPass(); +std::unique_ptr> createConvertGPUKernelToOCLSPIRVPass(int blockX, int blockY, int blockZ, const char* spirv_bin_path); +} +} + +#endif \ No newline at end of file diff --git a/include/comet/Conversion/GpuToOCLSPIRV/Passes.td b/include/comet/Conversion/GpuToOCLSPIRV/Passes.td new file mode 100644 index 00000000..81b24c60 --- /dev/null +++ b/include/comet/Conversion/GpuToOCLSPIRV/Passes.td @@ -0,0 +1,57 @@ +// +// Copyright 2022 Battelle Memorial Institute +// +// Redistribution and use in source and binary forms, with or without modification, +// are permitted provided that the following conditions are met: +// +// 1. Redistributions of source code must retain the above copyright notice, this list of conditions +// and the following disclaimer. +// +// 2. Redistributions in binary form must reproduce the above copyright notice, this list of conditions +// and the following disclaimer in the documentation and/or other materials provided with the distribution. +// +// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED +// WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR +// PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, +// INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE +// GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, +// WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +// OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. +// + +#ifndef GPU_TO_OCLSPIRV_CONVERSION_PASSES +#define GPU_TO_OCLSPIRV_CONVERSION_PASSES + +include "mlir/Pass/PassBase.td" + +def ConvertGpuKernelToOCLSPIRVPass: Pass<"convert-gpu-kernel-to-spirv", "mlir::ModuleOp"> { + let summary = "Convert Gpu Kernel to SPIRV"; + let description = [{ + + }]; + let constructor = "mlir::triton::createConvertGpuKernelToOCLSPIRVPass()"; + + let dependentDialects = ["mlir::arith::ArithDialect", + "mlir::scf::SCFDialect", + "mlir::spirv::SPIRVDialect"]; + + let options = [ + Option<"blockX", "blockX", + "int32_t", /*default*/"32", + "opencl block size for dim X">, + + Option<"blockY", "blockY", + "int32_t", /*default*/"8", + "opencl block size for dim Y">, + + Option<"blockZ", "blockZ", + "int32_t", /*default*/"32", + "opencl block size for reduction dim Z">, + + Option<"spirv_bin_path", "spirv-bin-path", + "const char*", "\"spirv_comet_\"", + "Path to SPIRV output binary"> + + ]; +} +#endif diff --git a/include/comet/Conversion/GpuToTriton/GpuToTritonConversion.h b/include/comet/Conversion/GpuToTriton/GpuToTritonConversion.h index 250c7b89..867be81c 100644 --- a/include/comet/Conversion/GpuToTriton/GpuToTritonConversion.h +++ b/include/comet/Conversion/GpuToTriton/GpuToTritonConversion.h @@ -1,3 +1,23 @@ +// +// Copyright 2022 Battelle Memorial Institute +// +// Redistribution and use in source and binary forms, with or without modification, +// are permitted provided that the following conditions are met: +// +// 1. Redistributions of source code must retain the above copyright notice, this list of conditions +// and the following disclaimer. +// +// 2. Redistributions in binary form must reproduce the above copyright notice, this list of conditions +// and the following disclaimer in the documentation and/or other materials provided with the distribution. +// +// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED +// WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR +// PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, +// INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE +// GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, +// WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +// OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. +// #ifndef GPU_DIALECT_TRITON_TRANSFORMS_COMETCONVERSION_H_ #define GPU_DIALECT_TRITON_TRANSFORMS_COMETCONVERSION_H_ @@ -26,7 +46,7 @@ class GpuTypeConverter2 : public TypeConverter { GpuTypeConverter2(MLIRContext *context); int blockX, blockY, blockR; private: - MLIRContext *context; + [[maybe_unused]]MLIRContext *context; }; class GpuConversionTarget2 : public ConversionTarget { diff --git a/include/comet/Conversion/GpuToTriton/GpuToTritonPass.h b/include/comet/Conversion/GpuToTriton/GpuToTritonPass.h index cb30ee8b..300f7cf2 100644 --- a/include/comet/Conversion/GpuToTriton/GpuToTritonPass.h +++ b/include/comet/Conversion/GpuToTriton/GpuToTritonPass.h @@ -1,3 +1,24 @@ +// +// Copyright 2022 Battelle Memorial Institute +// +// Redistribution and use in source and binary forms, with or without modification, +// are permitted provided that the following conditions are met: +// +// 1. Redistributions of source code must retain the above copyright notice, this list of conditions +// and the following disclaimer. +// +// 2. Redistributions in binary form must reproduce the above copyright notice, this list of conditions +// and the following disclaimer in the documentation and/or other materials provided with the distribution. +// +// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED +// WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR +// PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, +// INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE +// GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, +// WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +// OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. +// + #ifndef GPU_CONVERSION_GPUTOTRITON_H #define GPU_CONVERSION_GPUTOTRITON_H diff --git a/include/comet/Conversion/GpuToTriton/Passes.h b/include/comet/Conversion/GpuToTriton/Passes.h index 5b604e73..9d4f30db 100644 --- a/include/comet/Conversion/GpuToTriton/Passes.h +++ b/include/comet/Conversion/GpuToTriton/Passes.h @@ -1,3 +1,24 @@ +// +// Copyright 2022 Battelle Memorial Institute +// +// Redistribution and use in source and binary forms, with or without modification, +// are permitted provided that the following conditions are met: +// +// 1. Redistributions of source code must retain the above copyright notice, this list of conditions +// and the following disclaimer. +// +// 2. Redistributions in binary form must reproduce the above copyright notice, this list of conditions +// and the following disclaimer in the documentation and/or other materials provided with the distribution. +// +// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED +// WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR +// PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, +// INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE +// GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, +// WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +// OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. +// + #ifndef GPU_TO_TRITON_CONVERSION_PASSES_H #define GPU_TO_TRITON_CONVERSION_PASSES_H diff --git a/include/comet/Conversion/GpuToTriton/Passes.td b/include/comet/Conversion/GpuToTriton/Passes.td index e0164fbb..4174a798 100644 --- a/include/comet/Conversion/GpuToTriton/Passes.td +++ b/include/comet/Conversion/GpuToTriton/Passes.td @@ -1,3 +1,24 @@ +// +// Copyright 2022 Battelle Memorial Institute +// +// Redistribution and use in source and binary forms, with or without modification, +// are permitted provided that the following conditions are met: +// +// 1. Redistributions of source code must retain the above copyright notice, this list of conditions +// and the following disclaimer. +// +// 2. Redistributions in binary form must reproduce the above copyright notice, this list of conditions +// and the following disclaimer in the documentation and/or other materials provided with the distribution. +// +// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED +// WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR +// PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, +// INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE +// GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, +// WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +// OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. +// + #ifndef GPU_TO_TRITON_CONVERSION_PASSES #define GPU_TO_TRITON_CONVERSION_PASSES diff --git a/include/comet/Conversion/GpuUtils/GpuUtils.h b/include/comet/Conversion/GpuUtils/GpuUtils.h new file mode 100644 index 00000000..135144ec --- /dev/null +++ b/include/comet/Conversion/GpuUtils/GpuUtils.h @@ -0,0 +1,13 @@ +#ifndef GPU_UTILS_H +#define GPU_UTILS_H + +#include "comet/Dialect/Utils/Utils.h" +#include "mlir/IR/Builders.h" +#include "mlir/IR/BuiltinOps.h" +#include "mlir/Support/LLVM.h" + +mlir::LogicalResult specializeGpuHost(mlir::OpBuilder& builder, mlir::ModuleOp modOp, std::string vendor_prefix); +void declare_vendor_funcs(mlir::OpBuilder& builder, mlir::ModuleOp modOp, std::string vendor_prefix); +mlir::LogicalResult specializeGpuKernel(mlir::OpBuilder& builder, mlir::ModuleOp modOp, mlir::tensorAlgebra::GPUCompilationFormat codeFormat, mlir::Attribute target, std::function add_ttir_passes, std::function add_ttgir_passes, std::function add_llir_passes, std::vector llvm_func_attrs={}); + +#endif \ No newline at end of file diff --git a/include/comet/Conversion/IndexTreeToSCF/IndexTreeToSCF.h b/include/comet/Conversion/IndexTreeToSCF/IndexTreeToSCF.h index 2a8f44ba..77431b64 100644 --- a/include/comet/Conversion/IndexTreeToSCF/IndexTreeToSCF.h +++ b/include/comet/Conversion/IndexTreeToSCF/IndexTreeToSCF.h @@ -34,16 +34,13 @@ namespace mlir namespace comet { #define GEN_PASS_DECL_CONVERTINDEXTREETOSCF +#define GEN_PASS_DECL_CONVERTSYMBOLICDOMAINS #include "comet/Conversion/Passes.h.inc" - - /// Collect a set of patterns to convert IndexTree operations to SCF - /// operations within the SCF dialect. - void populateIndexTreeToSCFConversionPatterns(RewritePatternSet &patterns); - /// Lowers indexTree operations (e.g., IndexTreeComputeLHSOp, IndexTreeComputeRHSOp and IndexTreeComputeOp) /// to equivalent scf constructs including basic blocks and arithmetic /// primitives). std::unique_ptr createLowerIndexTreeToSCFPass(); + std::unique_ptr createConvertSymbolicDomainsPass(); } } // namespace mlir diff --git a/include/comet/Conversion/ParallelLoopsToGpu/ParallelLoopsToGpu.h b/include/comet/Conversion/ParallelLoopsToGpu/ParallelLoopsToGpu.h deleted file mode 100644 index 5f1374e5..00000000 --- a/include/comet/Conversion/ParallelLoopsToGpu/ParallelLoopsToGpu.h +++ /dev/null @@ -1,17 +0,0 @@ -#ifndef COMET_CONVERSION_PARALLELLOOPSTOGPU_H -#define COMET_CONVERSION_PARALLELLOOPSTOGPU_H - -#include -namespace mlir { -class ModuleOp; -namespace func{ -class FuncOp; -} -template class OperationPass; -namespace comet { -std::unique_ptr> createConvertParallelLoopsToGpuPass(); -std::unique_ptr> createConvertParallelLoopsToGpuPass(int blockX, int blockY, int blockR); -} -} - -#endif \ No newline at end of file diff --git a/include/comet/Conversion/ParallelLoopsToGpu/Passes.h b/include/comet/Conversion/ParallelLoopsToGpu/Passes.h deleted file mode 100644 index 3f9a9c85..00000000 --- a/include/comet/Conversion/ParallelLoopsToGpu/Passes.h +++ /dev/null @@ -1,13 +0,0 @@ -#ifndef COMET_PARALLEL_LOOPS_TO_GPU_CONVERSION_PASSES -#define COMET_PARALLEL_LOOPS_TO_GPU_CONVERSION_PASSES - -#include "comet/Conversion/ParallelLoopsToGpu/ParallelLoopsToGpu.h" - -namespace mlir { -namespace comet { -#define GEN_PASS_REGISTRATION -#include "comet/Conversion/ParallelLoopsToGpu/Passes.h.inc" -} -} - -#endif \ No newline at end of file diff --git a/include/comet/Conversion/ParallelLoopsToGpu/Passes.td b/include/comet/Conversion/ParallelLoopsToGpu/Passes.td deleted file mode 100644 index 4f092976..00000000 --- a/include/comet/Conversion/ParallelLoopsToGpu/Passes.td +++ /dev/null @@ -1,30 +0,0 @@ -#ifndef COMET_PARALLEL_LOOPS_TO_GPU_CONVERSION_PASSES -#define COMET_PARALLEL_LOOPS_TO_GPU_CONVERSION_PASSES - -include "mlir/Pass/PassBase.td" - -def CometParallelLoopsToGpu : Pass<"comet-parallel-loops-to-gpu", "mlir::func::FuncOp"> { - let summary = "Map loops to gpu"; - - let constructor = "mlir::comet::createConvertParallelLoopsToGpuPass()"; - - let dependentDialects = ["mlir::scf::SCFDialect", "mlir::gpu::GPUDialect", "mlir::affine::AffineDialect", "mlir::DLTIDialect"]; - - - let options = [ - Option<"blockX", "blockX", - "int32_t", /*default*/"32", - "triton block size for dim X">, - - Option<"blockY", "blockY", - "int32_t", /*default*/"8", - "triton block size for dim Y">, - - Option<"blockR", "blockR", - "int32_t", /*default*/"32", - "triton block size for reduction dim">, - - ]; -} - -#endif diff --git a/include/comet/Conversion/ParallelLoopsToGpu/CMakeLists.txt b/include/comet/Conversion/ParallelLoopsToGpuFPGA/CMakeLists.txt similarity index 56% rename from include/comet/Conversion/ParallelLoopsToGpu/CMakeLists.txt rename to include/comet/Conversion/ParallelLoopsToGpuFPGA/CMakeLists.txt index 02b42919..2b6bee23 100644 --- a/include/comet/Conversion/ParallelLoopsToGpu/CMakeLists.txt +++ b/include/comet/Conversion/ParallelLoopsToGpuFPGA/CMakeLists.txt @@ -1,3 +1,3 @@ set(LLVM_TARGET_DEFINITIONS Passes.td) -mlir_tablegen(Passes.h.inc -gen-pass-decls --name ParallelLoopsToGpu) -add_public_tablegen_target(ParallelLoopsConversionPassIncGen) \ No newline at end of file +mlir_tablegen(Passes.h.inc -gen-pass-decls --name ParallelLoopsToGpuFPGA) +add_public_tablegen_target(ParallelLoopsToGpuFPGAConversionPassIncGen) \ No newline at end of file diff --git a/include/comet/Conversion/ParallelLoopsToGpuFPGA/ParallelLoopsToGpuFPGA.h b/include/comet/Conversion/ParallelLoopsToGpuFPGA/ParallelLoopsToGpuFPGA.h new file mode 100644 index 00000000..f70a8783 --- /dev/null +++ b/include/comet/Conversion/ParallelLoopsToGpuFPGA/ParallelLoopsToGpuFPGA.h @@ -0,0 +1,39 @@ +// +// Copyright 2022 Battelle Memorial Institute +// +// Redistribution and use in source and binary forms, with or without modification, +// are permitted provided that the following conditions are met: +// +// 1. Redistributions of source code must retain the above copyright notice, this list of conditions +// and the following disclaimer. +// +// 2. Redistributions in binary form must reproduce the above copyright notice, this list of conditions +// and the following disclaimer in the documentation and/or other materials provided with the distribution. +// +// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED +// WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR +// PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, +// INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE +// GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, +// WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +// OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. +// + +#ifndef COMET_CONVERSION_PARALLELLOOPSTOGPUFPGA_H +#define COMET_CONVERSION_PARALLELLOOPSTOGPUFPGA_H +#include "comet/Dialect/Utils/Utils.h" + +#include +namespace mlir { +class ModuleOp; +namespace func{ +class FuncOp; +} +template class OperationPass; +namespace comet { +std::unique_ptr> createConvertParallelLoopsToGpuFPGAPass(); +std::unique_ptr> createConvertParallelLoopsToGpuFPGAPass(int blockX, int blockY, int blockR, mlir::tensorAlgebra::TargetDevice target_device); +} +} + +#endif \ No newline at end of file diff --git a/include/comet/Conversion/ParallelLoopsToGpuFPGA/Passes.h b/include/comet/Conversion/ParallelLoopsToGpuFPGA/Passes.h new file mode 100644 index 00000000..bf3a1b14 --- /dev/null +++ b/include/comet/Conversion/ParallelLoopsToGpuFPGA/Passes.h @@ -0,0 +1,34 @@ +// +// Copyright 2022 Battelle Memorial Institute +// +// Redistribution and use in source and binary forms, with or without modification, +// are permitted provided that the following conditions are met: +// +// 1. Redistributions of source code must retain the above copyright notice, this list of conditions +// and the following disclaimer. +// +// 2. Redistributions in binary form must reproduce the above copyright notice, this list of conditions +// and the following disclaimer in the documentation and/or other materials provided with the distribution. +// +// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED +// WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR +// PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, +// INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE +// GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, +// WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +// OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. +// + +#ifndef COMET_PARALLEL_LOOPS_TO_GPU_FPGA_CONVERSION_PASSES +#define COMET_PARALLEL_LOOPS_TO_GPU_FPGA_CONVERSION_PASSES + +#include "comet/Conversion/ParallelLoopsToGpuFPGA/ParallelLoopsToGpuFPGA.h" + +namespace mlir { +namespace comet { +#define GEN_PASS_REGISTRATION +#include "comet/Conversion/ParallelLoopsToGpuFPGA/Passes.h.inc" +} +} + +#endif \ No newline at end of file diff --git a/include/comet/Conversion/ParallelLoopsToGpuFPGA/Passes.td b/include/comet/Conversion/ParallelLoopsToGpuFPGA/Passes.td new file mode 100644 index 00000000..c997b590 --- /dev/null +++ b/include/comet/Conversion/ParallelLoopsToGpuFPGA/Passes.td @@ -0,0 +1,55 @@ +// +// Copyright 2022 Battelle Memorial Institute +// +// Redistribution and use in source and binary forms, with or without modification, +// are permitted provided that the following conditions are met: +// +// 1. Redistributions of source code must retain the above copyright notice, this list of conditions +// and the following disclaimer. +// +// 2. Redistributions in binary form must reproduce the above copyright notice, this list of conditions +// and the following disclaimer in the documentation and/or other materials provided with the distribution. +// +// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED +// WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR +// PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, +// INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE +// GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, +// WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +// OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. +// + +#ifndef COMET_PARALLEL_LOOPS_TO_GPU_FPGA_CONVERSION_PASSES +#define COMET_PARALLEL_LOOPS_TO_GPU_FPGA_CONVERSION_PASSES + +include "mlir/Pass/PassBase.td" + +def CometParallelLoopsToGpuFPGA : Pass<"comet-parallel-loops-to-gpu-fpga", "mlir::func::FuncOp"> { + let summary = "Map loops to gpu"; + + let constructor = "mlir::comet::createConvertParallelLoopsToGpuFPGAPass()"; + + let dependentDialects = ["mlir::scf::SCFDialect", "mlir::gpu::GPUDialect", "mlir::affine::AffineDialect", "mlir::DLTIDialect"]; + + + let options = [ + Option<"blockX", "blockX", + "int32_t", /*default*/"32", + "triton block size for dim X">, + + Option<"blockY", "blockY", + "int32_t", /*default*/"8", + "triton block size for dim Y">, + + Option<"blockR", "blockR", + "int32_t", /*default*/"32", + "triton block size for reduction dim">, + + Option<"target_device", "target_device", + "mlir::tensorAlgebra::TargetDevice", /*default*/"mlir::tensorAlgebra::TargetDevice::GPU", + "Target device (GPU or FPGA)">, + + ]; +} + +#endif diff --git a/include/comet/Conversion/Passes.td b/include/comet/Conversion/Passes.td index 052812a1..36eb3582 100644 --- a/include/comet/Conversion/Passes.td +++ b/include/comet/Conversion/Passes.td @@ -43,6 +43,19 @@ def ConvertIndexTreeToSCF : Pass<"convert-it-to-scf"> { ]; } +def ConvertSymbolicDomains : Pass<"convert-symbolic-domains"> { + let summary = " " + ""; + let description = [{ + + }]; + let constructor = "comet::createConvertSymbolicDomainsPass()"; + let dependentDialects = [ + "memref::MemRefDialect", + "scf::SCFDialect" + ]; +} + //===----------------------------------------------------------------------===// // TensorAlgebraToIndexTree //===----------------------------------------------------------------------===// @@ -73,7 +86,22 @@ def ConvertTensorAlgebraToSCF : Pass<"convert-ta-to-scf"> { let constructor = "comet::createLowerTensorAlgebraToSCFPass()"; let dependentDialects = [ "memref::MemRefDialect", - "scf::SCFDialect" + "scf::SCFDialect", + "index::IndexDialect" + ]; +} + +def SparseTensorConversionPass : Pass<"convert-sparse-tensor"> { + let summary = "Lowers operations on Tensor Algebra sparse tensors to mlir tensor operations"; + let description = [{}]; + + let constructor = "comet::createSparseTensorConversionPass()"; + + let dependentDialects = [ + "memref::MemRefDialect", + "scf::SCFDialect", + "index::IndexDialect", + "tensor::TensorDialect" ]; } diff --git a/include/comet/Conversion/PrepareGpuHost/CMakeLists.txt b/include/comet/Conversion/PrepareGpuHost/CMakeLists.txt new file mode 100644 index 00000000..bfaf874d --- /dev/null +++ b/include/comet/Conversion/PrepareGpuHost/CMakeLists.txt @@ -0,0 +1,3 @@ +set(LLVM_TARGET_DEFINITIONS Passes.td) +mlir_tablegen(Passes.h.inc -gen-pass-decls --name PrepareGpuHost) +add_public_tablegen_target(PrepareGpuHostPassIncGen) \ No newline at end of file diff --git a/include/comet/Conversion/PrepareGpuHost/Passes.h b/include/comet/Conversion/PrepareGpuHost/Passes.h new file mode 100644 index 00000000..1cf7051a --- /dev/null +++ b/include/comet/Conversion/PrepareGpuHost/Passes.h @@ -0,0 +1,13 @@ +#ifndef PREPARE_GPU_HOST_PASSES +#define PREPARE_GPU_HOST_PASSES + +#include "comet/Conversion/PrepareGpuHost/PrepareGpuHostPass.h" + +namespace mlir { +namespace comet { +#define GEN_PASS_REGISTRATION +#include "comet/Conversion/PrepareGpuHost/Passes.h.inc" +} +} + +#endif \ No newline at end of file diff --git a/include/comet/Conversion/PrepareGpuHost/Passes.td b/include/comet/Conversion/PrepareGpuHost/Passes.td new file mode 100644 index 00000000..67f89201 --- /dev/null +++ b/include/comet/Conversion/PrepareGpuHost/Passes.td @@ -0,0 +1,19 @@ +#ifndef PREPARE_GPU_HOST_PASSES +#define PREPARE_GPU_HOST_PASSES + +include "mlir/Pass/PassBase.td" + +def PrepareGpuHost: Pass<"prepare-gpu-host", "mlir::ModuleOp"> { + let summary = "Prepare GPU host"; + let description = [{}]; + let constructor = "mlir::comet::createPrepareGpuHostPass()"; + + let options = [ + Option<"generateAllocsAndTransfers", "generate-allocs-and-transfers", + "bool", "true", + "Whether to generate allocations and transfers"> + ]; + + let dependentDialects = ["mlir::gpu::GPUDialect", "mlir::func::FuncDialect", "mlir::LLVM::LLVMDialect"]; +} +#endif diff --git a/include/comet/Conversion/PrepareGpuHost/PrepareGpuHostPass.h b/include/comet/Conversion/PrepareGpuHost/PrepareGpuHostPass.h new file mode 100644 index 00000000..ae21508b --- /dev/null +++ b/include/comet/Conversion/PrepareGpuHost/PrepareGpuHostPass.h @@ -0,0 +1,21 @@ +#ifndef COMET_PREPARE_GPU_HOST_PASS +#define COMET_PREPARE_GPU_HOST_PASS + +#include +#include "comet/Dialect/Utils/Utils.h" + +namespace mlir { + +class ModuleOp; +template class OperationPass; + +namespace comet { + + +std::unique_ptr> createPrepareGpuHostPass(); +std::unique_ptr> createPrepareGpuHostPass(bool generateAllocsAndTransfers); + +} // namespace triton +} // namespace mlir + +#endif \ No newline at end of file diff --git a/include/comet/Conversion/TensorAlgebraToSCF/TensorAlgebraToSCF.h b/include/comet/Conversion/TensorAlgebraToSCF/TensorAlgebraToSCF.h index af46bb1d..700a7633 100644 --- a/include/comet/Conversion/TensorAlgebraToSCF/TensorAlgebraToSCF.h +++ b/include/comet/Conversion/TensorAlgebraToSCF/TensorAlgebraToSCF.h @@ -25,16 +25,18 @@ #define COMET_CONVERSION_TENSORALGEBRATOSCF_H #include "mlir/Support/LLVM.h" +#include "mlir/Transforms/DialectConversion.h" namespace mlir { class Pass; - class RewritePatternSet; - + namespace comet { #define GEN_PASS_DECL_CONVERTTENSORALGEBRATOSCF #include "comet/Conversion/Passes.h.inc" + void populateSparseTensorConversionPatterns(MLIRContext *context, RewritePatternSet &patterns, TypeConverter &typeConverter); + std::unique_ptr createSparseTensorConversionPass(); /// Collect a set of patterns to convert remaining TensorAlgebra operations /// that are not converted to IndexTree operations to the operations with SCF diff --git a/include/comet/Conversion/TritonToCuda/Passes.h b/include/comet/Conversion/TritonToCuda/Passes.h index 3f6a2edf..1c1d916c 100644 --- a/include/comet/Conversion/TritonToCuda/Passes.h +++ b/include/comet/Conversion/TritonToCuda/Passes.h @@ -1,7 +1,29 @@ +// +// Copyright 2022 Battelle Memorial Institute +// +// Redistribution and use in source and binary forms, with or without modification, +// are permitted provided that the following conditions are met: +// +// 1. Redistributions of source code must retain the above copyright notice, this list of conditions +// and the following disclaimer. +// +// 2. Redistributions in binary form must reproduce the above copyright notice, this list of conditions +// and the following disclaimer in the documentation and/or other materials provided with the distribution. +// +// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED +// WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR +// PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, +// INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE +// GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, +// WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +// OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. +// + #ifndef TRITON_TO_CUDA_CONVERSION_PASSES #define TRITON_TO_CUDA_CONVERSION_PASSES #include "comet/Conversion/TritonToCuda/TritonToCudaPass.h" +#include "nvidia/include/Dialect/NVGPU/IR/Dialect.h" namespace mlir { namespace comet { diff --git a/include/comet/Conversion/TritonToCuda/Passes.td b/include/comet/Conversion/TritonToCuda/Passes.td index dc0a4862..8d97d837 100644 --- a/include/comet/Conversion/TritonToCuda/Passes.td +++ b/include/comet/Conversion/TritonToCuda/Passes.td @@ -1,3 +1,24 @@ +// +// Copyright 2022 Battelle Memorial Institute +// +// Redistribution and use in source and binary forms, with or without modification, +// are permitted provided that the following conditions are met: +// +// 1. Redistributions of source code must retain the above copyright notice, this list of conditions +// and the following disclaimer. +// +// 2. Redistributions in binary form must reproduce the above copyright notice, this list of conditions +// and the following disclaimer in the documentation and/or other materials provided with the distribution. +// +// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED +// WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR +// PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, +// INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE +// GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, +// WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +// OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. +// + #ifndef GPU_TO_TRITON_CONVERSION_PASSES #define GPU_TO_TRITON_CONVERSION_PASSES @@ -36,6 +57,10 @@ def LowerTritonDeviceToCuda: Pass<"lower-triton-device-to-cuda", "mlir::ModuleOp Option<"computeCapability", "computeCapability", "int32_t", "80", "Target compute capability">, + + Option<"codeFormat", "codeFormat", + "mlir::tensorAlgebra::GPUCompilationFormat", "mlir::tensorAlgebra::Binary", + "Target code format">, ]; let dependentDialects = ["mlir::arith::ArithDialect", "mlir::math::MathDialect", diff --git a/include/comet/Conversion/TritonToCuda/TritonToCudaPass.h b/include/comet/Conversion/TritonToCuda/TritonToCudaPass.h index 8b77c02e..b9740fc5 100644 --- a/include/comet/Conversion/TritonToCuda/TritonToCudaPass.h +++ b/include/comet/Conversion/TritonToCuda/TritonToCudaPass.h @@ -1,7 +1,29 @@ +// +// Copyright 2022 Battelle Memorial Institute +// +// Redistribution and use in source and binary forms, with or without modification, +// are permitted provided that the following conditions are met: +// +// 1. Redistributions of source code must retain the above copyright notice, this list of conditions +// and the following disclaimer. +// +// 2. Redistributions in binary form must reproduce the above copyright notice, this list of conditions +// and the following disclaimer in the documentation and/or other materials provided with the distribution. +// +// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED +// WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR +// PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, +// INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE +// GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, +// WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +// OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. +// + #ifndef TRITON_CONVERSION_TRITONTOCUDA_H #define TRITON_CONVERSION_TRITONTOCUDA_H #include +#include "comet/Dialect/Utils/Utils.h" namespace mlir { @@ -17,7 +39,8 @@ std::unique_ptr> createLowerTritonDeviceToCudaPass int threadsPerWarp, int numCTAs, int numStages, - int computeCapability); + int computeCapability, + mlir::tensorAlgebra::GPUCompilationFormat format); } // namespace triton } // namespace mlir diff --git a/include/comet/Conversion/TritonToHIP/CMakeLists.txt b/include/comet/Conversion/TritonToHIP/CMakeLists.txt new file mode 100644 index 00000000..d3374469 --- /dev/null +++ b/include/comet/Conversion/TritonToHIP/CMakeLists.txt @@ -0,0 +1,3 @@ +set(LLVM_TARGET_DEFINITIONS Passes.td) +mlir_tablegen(Passes.h.inc -gen-pass-decls --name TritonToHIP) +add_public_tablegen_target(TritonDeviceToHIPConversionPassIncGen) \ No newline at end of file diff --git a/include/comet/Conversion/TritonToHIP/Passes.h b/include/comet/Conversion/TritonToHIP/Passes.h new file mode 100644 index 00000000..82f4cec3 --- /dev/null +++ b/include/comet/Conversion/TritonToHIP/Passes.h @@ -0,0 +1,14 @@ +#ifndef TRITON_TO_HIP_CONVERSION_PASSES +#define TRITON_TO_HIP_CONVERSION_PASSES + +#include "comet/Conversion/TritonToHIP/TritonToHIPPass.h" +#include "amd/include/Dialect/TritonAMDGPU/IR/Dialect.h" + +namespace mlir { +namespace comet { +#define GEN_PASS_REGISTRATION +#include "comet/Conversion/TritonToHIP/Passes.h.inc" +} +} + +#endif \ No newline at end of file diff --git a/include/comet/Conversion/TritonToHIP/Passes.td b/include/comet/Conversion/TritonToHIP/Passes.td new file mode 100644 index 00000000..4b159086 --- /dev/null +++ b/include/comet/Conversion/TritonToHIP/Passes.td @@ -0,0 +1,50 @@ +#ifndef GPU_TO_TRITON_CONVERSION_PASSES +#define GPU_TO_TRITON_CONVERSION_PASSES + +include "mlir/Pass/PassBase.td" + +def LowerHostToHIP: Pass<"lower-gpu-host-to-hip", "mlir::ModuleOp"> { + let summary = "Lower Gpu dialect host to hip"; + let description = [{}]; + let constructor = "mlir::comet::createLowerGpuHostToHIPPass()"; + + let dependentDialects = ["mlir::func::FuncDialect", "mlir::LLVM::LLVMDialect"]; +} + +def LowerTritonDeviceToHIP: Pass<"lower-triton-device-to-hip", "mlir::ModuleOp"> { + let summary = "Lower Triton dialect host to hip"; + let description = [{}]; + let constructor = "mlir::comet::createLowerTritonDeviceToHIPPass()"; + + let options = [ + Option<"numWarps", "numWarps", + "int32_t", /*default*/ "4", + "Number of warps">, + + Option<"threadsPerWarp", "threadsPerWarp", + "int32_t", /*default*/ "32", + "Number of threads per warp">, + + Option<"numStages", "numStages", + "int32_t", /*default*/ "3", + "Number of stages">, + + Option<"numCTAs", "numCTAs", + "int32_t", /*default*/ "1", + "Number of CTAs">, + + Option<"computeCapability", "computeCapability", + "std::string", "\"gfx90\"", + "Target compute capability">, + + Option<"codeFormat", "codeFormat", + "mlir::tensorAlgebra::GPUCompilationFormat", "mlir::tensorAlgebra::Binary", + "Target code format">, + ]; + + let dependentDialects = ["mlir::arith::ArithDialect", "mlir::math::MathDialect", + "mlir::gpu::GPUDialect", + "mlir::memref::MemRefDialect", + "mlir::scf::SCFDialect", "mlir::tensor::TensorDialect","mlir::func::FuncDialect", "mlir::LLVM::LLVMDialect", "mlir::triton::TritonDialect", "mlir::triton::gpu::TritonGPUDialect", "mlir::triton::nvidia_gpu::TritonNvidiaGPUDialect", "mlir::triton::amdgpu::TritonAMDGPUDialect", "mlir::NVVM::NVVMDialect", "mlir::ROCDL::ROCDLDialect"]; +} +#endif diff --git a/include/comet/Conversion/TritonToHIP/TritonToHIPPass.h b/include/comet/Conversion/TritonToHIP/TritonToHIPPass.h new file mode 100644 index 00000000..aea98b64 --- /dev/null +++ b/include/comet/Conversion/TritonToHIP/TritonToHIPPass.h @@ -0,0 +1,27 @@ +#ifndef TRITON_CONVERSION_TRITONTOHIP_H +#define TRITON_CONVERSION_TRITONTOHIP_H + +#include +#include "comet/Dialect/Utils/Utils.h" + +namespace mlir { + +class ModuleOp; +template class OperationPass; + +namespace comet { + + +std::unique_ptr> createLowerGpuHostToHIPPass(); +std::unique_ptr> createLowerTritonDeviceToHIPPass(); +std::unique_ptr> createLowerTritonDeviceToHIPPass(int numWarps, + int threadsPerWarp, + int numCTAs, + int numStages, + std::string computeCapability, + mlir::tensorAlgebra::GPUCompilationFormat format); + +} // namespace triton +} // namespace mlir + +#endif \ No newline at end of file diff --git a/include/comet/Dialect/IndexTree/Analysis/CopiedDomainAnalysis.h b/include/comet/Dialect/IndexTree/Analysis/CopiedDomainAnalysis.h new file mode 100644 index 00000000..eff69dff --- /dev/null +++ b/include/comet/Dialect/IndexTree/Analysis/CopiedDomainAnalysis.h @@ -0,0 +1,28 @@ +#ifndef COMET_DIALECT_INDEXTREE_ANALYSIS_H +#define COMET_DIALECT_INDEXTREE_ANALYSIS_H + +#include "mlir/IR/Value.h" +#include "mlir/Transforms/DialectConversion.h" + +#include "llvm/ADT/SmallSet.h" + +#include "comet/Dialect/IndexTree/IR/IndexTreeDialect.h" + +namespace mlir { + namespace indexTree { + struct CopiedDomainAnalysis { + public: + CopiedDomainAnalysis(Operation* op); + bool isCopiedDomain(Value tensor, unsigned dim); + bool isReductionVar(IndexTreeComputeOp op, Value index_var); + + private: + llvm::SmallDenseSet> copiedDomains; + llvm::SmallDenseSet> reductionVars; + void analyzeDomains(IndexTreeComputeOp compute_op); + }; + + } +} + +#endif // COMET_DIALECT_INDEXTREE_ANALYSIS_H \ No newline at end of file diff --git a/include/comet/Dialect/IndexTree/IR/CMakeLists.txt b/include/comet/Dialect/IndexTree/IR/CMakeLists.txt index 9fb7a6d7..9095274e 100644 --- a/include/comet/Dialect/IndexTree/IR/CMakeLists.txt +++ b/include/comet/Dialect/IndexTree/IR/CMakeLists.txt @@ -1,8 +1,13 @@ set(LLVM_TARGET_DEFINITIONS IndexTreeOps.td) mlir_tablegen(IndexTreeOps.h.inc -gen-op-decls) mlir_tablegen(IndexTreeOps.cpp.inc -gen-op-defs) -mlir_tablegen(IndexTreeDialect.h.inc -gen-dialect-decls) -mlir_tablegen(IndexTreeDialect.cpp.inc -gen-dialect-defs) +mlir_tablegen(IndexTreeDialect.h.inc -gen-dialect-decls -dialect=it) +mlir_tablegen(IndexTreeDialect.cpp.inc -gen-dialect-defs -dialect=it) +mlir_tablegen(IndexTreeOpInterfaces.h.inc -gen-op-interface-decls) +mlir_tablegen(IndexTreeOpInterfaces.cpp.inc -gen-op-interface-defs) add_public_tablegen_target(COMETIndexTreeOpsIncGen) - +set(LLVM_TARGET_DEFINITIONS IndexTreeTypes.td) +mlir_tablegen(IndexTreeTypes.h.inc -gen-typedef-decls) +mlir_tablegen(IndexTreeTypes.cpp.inc -gen-typedef-defs) +add_public_tablegen_target(COMETIndexTreeTypesIncGen) \ No newline at end of file diff --git a/include/comet/Dialect/IndexTree/IR/IndexTreeBase.td b/include/comet/Dialect/IndexTree/IR/IndexTreeBase.td new file mode 100644 index 00000000..97116146 --- /dev/null +++ b/include/comet/Dialect/IndexTree/IR/IndexTreeBase.td @@ -0,0 +1,45 @@ +// +// Copyright 2022 Battelle Memorial Institute +// +// Redistribution and use in source and binary forms, with or without modification, +// are permitted provided that the following conditions are met: +// +// 1. Redistributions of source code must retain the above copyright notice, this list of conditions +// and the following disclaimer. +// +// 2. Redistributions in binary form must reproduce the above copyright notice, this list of conditions +// and the following disclaimer in the documentation and/or other materials provided with the distribution. +// +// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED +// WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR +// PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, +// INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE +// GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, +// WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +// OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. +// + +#ifndef INDEXTREE_BASE +#define INDEXTREE_BASE + +include "mlir/IR/OpBase.td" + +// Provide a definition of the 'it' dialect in the ODS framework so that we +// can define our operations. +def IndexTreeDialect : Dialect { + let name = "it"; + let cppNamespace = "::mlir::indexTree"; + + // We set this bit to generate the declarations for the dialect's type parsing + // and printing hooks. + let useDefaultTypePrinterParser = 1; +} + +class IndexTreeOpTrait : NativeOpTrait<""> { + let trait = name; + let cppNamespace = "::mlir::indexTree"; +} + +def UnknownDomain : IndexTreeOpTrait<"UnknownDomain">; + +#endif // INDEXTREE_BASE \ No newline at end of file diff --git a/include/comet/Dialect/IndexTree/IR/IndexTreeDialect.h b/include/comet/Dialect/IndexTree/IR/IndexTreeDialect.h index 771f01a1..e1f5c200 100644 --- a/include/comet/Dialect/IndexTree/IR/IndexTreeDialect.h +++ b/include/comet/Dialect/IndexTree/IR/IndexTreeDialect.h @@ -30,11 +30,25 @@ #include "mlir/IR/Dialect.h" #include "mlir/IR/BuiltinOps.h" +#include "mlir/Interfaces/ControlFlowInterfaces.h" +#include "llvm/ADT/StringSet.h" +#include "comet/Dialect/TensorAlgebra/IR/TADialect.h" /// Include the auto-generated header file containing the declaration of the index tree /// dialect. #include "comet/Dialect/IndexTree/IR/IndexTreeDialect.h.inc" +/// Include the auto-generated header file containing the declaration of the index tree +/// types. +#define GET_TYPEDEF_CLASSES +#include "comet/Dialect/IndexTree/IR/IndexTreeTypes.h.inc" + +// Include the trait definitions +#include "comet/Dialect/IndexTree/IR/IndexTreeTraits.h" + +// Include the op interface definitions +#include "comet/Dialect/IndexTree/IR/IndexTreeOpInterfaces.h.inc" + /// Include the auto-generated header file containing the declarations of the /// Index Tree operations and also the operations of the Shape Inference Op Interface. //===----------------------------------------------------------------------===// diff --git a/include/comet/Dialect/IndexTree/IR/IndexTreeOps.td b/include/comet/Dialect/IndexTree/IR/IndexTreeOps.td index f3208284..f4f58d0d 100644 --- a/include/comet/Dialect/IndexTree/IR/IndexTreeOps.td +++ b/include/comet/Dialect/IndexTree/IR/IndexTreeOps.td @@ -21,32 +21,24 @@ // // ============================================================================= // -// Defines the operations of the IT dialect. +// Defines the operations of the IndexTree dialect. // //===----------------------------------------------------------------------===// #ifndef INDEXTREE_OPS #define INDEXTREE_OPS -include "mlir/IR/OpBase.td" +include "mlir/IR/OpBase.td" +include "mlir/Interfaces/ControlFlowInterfaces.td" +include "mlir/IR/RegionKindInterface.td" include "mlir/Interfaces/FunctionInterfaces.td" include "mlir/IR/SymbolInterfaces.td" include "mlir/Interfaces/CallInterfaces.td" -include "mlir/Interfaces/CastInterfaces.td" include "mlir/Interfaces/SideEffectInterfaces.td" +include "mlir/Dialect/SCF/IR/DeviceMappingInterface.td" - -// Provide a definition of the 'it' dialect in the ODS framework so that we -// can define our operations. -def IndexTreeDialect : Dialect { - let name = "it"; - let cppNamespace = "::mlir::indexTree"; - - // We set this bit to generate the declarations for the dialect's type parsing - // and printing hooks. - let useDefaultTypePrinterParser = 1; - -} +include "comet/Dialect/IndexTree/IR/IndexTreeTypes.td" +include "comet/Dialect/TensorAlgebra/IR/TATypes.td" // Base class for ta dialect operations. This operation inherits from the base // `Op` class in OpBase.td, and provides: @@ -60,62 +52,372 @@ class IndexTree_Op traits = []> : // Index Tree Operations //===----------------------------------------------------------------------===// -def IndexTreeComputeLHSOp : IndexTree_Op<"ComputeLHS", [Pure]>{ +def IndexTreeOp : IndexTree_Op<"itree", + [SingleBlockImplicitTerminator<"indexTree::YieldOp">, AttrSizedOperandSegments]> { + + let summary = "Create a scope for index tree iteration"; + let description = [{}]; + + let arguments = (ins Variadic:$inputs, Variadic:$intermediates); + let results = (outs Variadic:$results); + let regions = (region SizedRegion<1>:$region); +} + +def YieldOp : IndexTree_Op<"yield", [Pure, ReturnLike, Terminator, + HasParent<"IndexTreeOp">]> { + let summary = "index tree yield and termination operation"; + let description = [{ + }]; + + let arguments = (ins Variadic:$results); + let builders = [OpBuilder<(ins), [{ /* nothing to do */ }]>]; + + let assemblyFormat = + [{ attr-dict ($results^ `:` type($results))? }]; +} + +def IndexTreeNodeInterface : OpInterface<"IndexTreeNode"> { + let description = [{ + Describes any operation which is attached to specific point in the tree. + }]; + let cppNamespace = "::mlir::indexTree"; + + let methods = [ + InterfaceMethod< + [{Method to get the dimension size of a concrete dimension.}], + "mlir::Value", "getParentNode", (ins), /*methodBody=*/[{}], + /*defaultImplementation=*/[{ + return $_op.getParent(); + }] + >, + ]; +} + +def IndexTreeRootOp : IndexTree_Op<"RootOp", [ + HasParent<"IndexTreeOp">, Pure]> { + let summary = "Create the base of the iteration tree"; + let description = [{}]; + + let results = (outs IndexTree_TreeType:$output); +} + +def IndexTreeIndicesOp : IndexTree_Op<"IndexOp", [Pure, IndexTreeNodeInterface]>{ + let summary = "Create an index variable bound to a specific computation."; + let description = [{ + }]; + + let arguments = (ins IndexTree_NodeType:$parent, Optional:$domain, DefaultValuedAttr:$IsParallel, OptionalAttr:$parallelDim); + + let results = (outs IndexTree_IndexNodeType:$output); +} +def IndexTreeLHSOperandOp : IndexTree_Op<"LHSOperandOp", [Pure, SameVariadicOperandSize, IndexTreeNodeInterface]>{ let summary = ""; let description = [{}]; - let arguments = (ins Variadic:$tensors, ArrayAttr:$allPerms, ArrayAttr:$allFormats); - let results = (outs AnyType:$output); + let arguments = (ins + // TA_AnyTensor:$tensor, + AnyTypeOf<[TA_AnyTensor, AnyFloat]>:$tensor, + Variadic:$pos, + Variadic:$crds); + let results = (outs IndexTree_OperandType:$result); + + let extraClassDeclaration = [{ + Value getParentNode() { + for(Operation* user : getOperation()->getUsers()) { + if(auto compute_op = llvm::dyn_cast(user)) { + return compute_op.getParentNode(); + } + } + return nullptr; + } + }]; } -def IndexTreeComputeRHSOp : IndexTree_Op<"ComputeRHS", [Pure]>{ +def IndexTreeOperandOp : IndexTree_Op<"OperandOp", [Pure, SameVariadicOperandSize, IndexTreeNodeInterface]>{ let summary = ""; let description = [{}]; - let arguments = (ins Variadic:$tensors, ArrayAttr:$allPerms, ArrayAttr:$allFormats); - let results = (outs AnyType:$output); + let arguments = (ins + // TA_AnyTensor:$tensor, + AnyTypeOf<[TA_AnyTensor, AnyFloat]>:$tensor, + Variadic:$pos, + Variadic:$crds); + let results = (outs IndexTree_OperandType:$result); + + let extraClassDeclaration = [{ + Value getParentNode() { + for(Operation* user : getOperation()->getUsers()) { + if(auto compute_op = llvm::dyn_cast(user)) { + return compute_op.getParentNode(); + } + } + return nullptr; + } + }]; } -def IndexTreeComputeOp : IndexTree_Op<"Compute", [Pure]>{ +def IndexTreeComputeOp : IndexTree_Op<"ComputeOp", [Pure, AttrSizedOperandSegments, IndexTreeNodeInterface]>{ let summary = ""; let description = [{ }]; //TODO(gkestor): rethink the use of comp_worksp_opt, should we decouple that? + //TODO(alokvk2): Rename semiring to operation to account for elementwise addition /// MaskType attribute: {push, pull, auto, none} - let arguments = (ins Variadic:$rhs, AnyType:$lhs, BoolAttr:$comp_worksp_opt, StrAttr:$semiring, StrAttr:$MaskType); + let arguments = (ins + IndexTree_NodeType:$parent, + IndexTree_OperandType:$lhs, + Variadic:$rhs, + Optional:$mask, + StrAttr:$semiring, + DefaultValuedAttr:$compute_missing + ); - let results = (outs I64:$output); + // let results = (outs TA_AnyTensor); + let results = (outs AnyTypeOf<[TA_AnyTensor, AnyFloat]>); //TODO(gkestor): add verifier //let hasVerifier = 1; - } -def IndexTreeIndicesOp : IndexTree_Op<"Indices", [Pure]>{ +class IndexTreeDomain_Op traits = []> : + IndexTree_Op{ + + let extraClassDeclaration = [{ + Value getParentNode() { + for(Operation* user : getOperation()->getUsers()) { + if(auto index_node = llvm::dyn_cast(user)) { + return index_node.getParentNode(); + } + } + return nullptr; + } + }]; +} + +def IndexTreeTensorDomainOp : IndexTreeDomain_Op<"DomainOp", [Pure, IndexTreeNodeInterface]>{ let summary = ""; let description = [{ }]; - // Added `iterator_type` to the IndexTreeIndicesOp. It is analogous the `iterator_types` in the `linalg.generic` op. - // Candidate options: parallel, reduction, window, serial, default. What options should support needs further consideration. - // References: https://mlir.llvm.org/docs/Dialects/Linalg/#linalggeneric-linalggenericop - let arguments = (ins Variadic:$children, ArrayAttr:$indices, StrAttr:$iterator_type); - let results = (outs I64:$output); - //TODO(gkestor): add verifier - //let hasVerifier = 1; + let arguments = (ins TA_AnyTensor:$tensor, UI32Attr:$dim, I32Attr:$format, Optional>:$parent); + + let results = (outs IndexTree_DomainType:$domain); } -def IndexTreeOp : IndexTree_Op<"itree", [Pure]>{ +def IndexTreeEmptyDomainOp : IndexTreeDomain_Op<"EmptyDomain", [Pure, IndexTreeNodeInterface]>{ let summary = ""; let description = [{ }]; - let arguments = (ins AnyType:$children); - let results = (outs I64:$output); + let results = (outs IndexTree_DomainType:$domain); +} - //TODO(gkestor): add verifier - //let hasVerifier = 1; +def ConcreteDomainInterface : OpInterface<"ConcreteDomain"> { + let description = [{ + Describes an operation which implements a concrete domain. That is + the domain is fully-dscribed by it's arguments and does not rely on the + the tensor definition to create the domain. + }]; + let cppNamespace = "::mlir::indexTree"; + + let methods = [ + InterfaceMethod<[{ + Method to get the dimension size of a concrete dimension. + }], + "mlir::Value", "getDimensionSize", (ins), /*methodBody=*/[{ + return $_op.getDimSize(); + }]>, + ]; +} + +def IndexTreeSparseDomainOp : IndexTreeDomain_Op<"SparseDomainOp", + [Pure, ConcreteDomainInterface, IndexTreeNodeInterface]>{ + + let summary = ""; + let description = [{ + }]; + + let arguments = (ins + TA_AnyTensor:$tensor, + UI32Attr:$dim, + I32Attr:$format, + AnyTensor:$pos, + AnyTensor:$crd, + Index:$pos_size, + Index:$crd_size, + Index:$dim_size, + Optional>:$parent); + + let results = (outs IndexTree_DomainType:$domain); +} + +def IndexTreeDenseDomainOp : IndexTreeDomain_Op<"DenseDomainOp", + [Pure, SameVariadicOperandSize, ConcreteDomainInterface, IndexTreeNodeInterface]>{ + let summary = ""; + let description = [{ + }]; + + let arguments = (ins AnyTypeOf<[UI32,Index,I32]>:$dim_size, Variadic:$tensors, I32ArrayAttr:$dims); + + let results = (outs IndexTree_DomainType:$domain); +} + +def IndexTreeWorkspaceDomainOp : IndexTreeDomain_Op<"WorkspaceDomainOp", + [Pure, ConcreteDomainInterface, IndexTreeNodeInterface]>{ + let summary = ""; + let description = [{}]; + + let arguments = (ins + WorkspaceTensor:$tensor, + Index:$dim_size, + UI32Attr:$dim, + Optional>:$parent + ); + let results = (outs IndexTree_DomainType:$domain); } +def IndexTreeDomainUnionOp : IndexTreeDomain_Op<"DomainUnionOp", + [Pure, UnknownDomain, ConcreteDomainInterface, AttrSizedOperandSegments, IndexTreeNodeInterface]>{ + let summary = ""; + let description = [{ + }]; + + let arguments = (ins Variadic:$domains, Optional>:$dim_size); + let results = (outs IndexTree_DomainType:$domain); +} + +def IndexTreeDomainIntersectionOp : IndexTreeDomain_Op<"DomainIntersectionOp", + [Pure, UnknownDomain, ConcreteDomainInterface, AttrSizedOperandSegments, IndexTreeNodeInterface]>{ + let summary = ""; + let description = [{ + }]; + + let arguments = (ins Variadic:$domains, Optional:$dim_size); + let results = (outs IndexTree_DomainType:$domain); +} + +def IndexTreeFillMaskOp : IndexTree_Op<"FillMaskOp", [Pure, IndexTreeNodeInterface]>{ + let summary = ""; + let description = [{ + }]; + + let arguments = (ins IndexTree_NodeType:$parent, IndexTree_DomainType:$domain, TensorOf<[I1]>:$init); + let results = (outs TensorOf<[I1]>:$result); +} + +def IndexTreeZeroMaskOp : IndexTree_Op<"ZeroMaskOp", [Pure, IndexTreeNodeInterface]>{ + let summary = ""; + let description = [{ + }]; + + let arguments = (ins IndexTree_NodeType:$parent, IndexTree_DomainType:$domain, TensorOf<[I1]>:$init); + let results = (outs TensorOf<[I1]>:$result); +} + +def IndexTreeMaskedDomainOp : IndexTreeDomain_Op<"MaskedDomainOp", + [Pure, UnknownDomain, ConcreteDomainInterface, IndexTreeNodeInterface]>{ + let summary = ""; + let description = [{ + }]; + + let arguments = (ins + AnyTypeOf<[TensorOf<[I1]>, IndexTree_DomainType]>:$mask, + IndexTree_DomainType:$base, + Optional:$dim_size); + + let results = (outs IndexTree_DomainType:$domain); +} + +def IndexTreeNestedDomainOp : IndexTree_Op<"NestedDomainOp", + [Pure, UnknownDomain, ConcreteDomainInterface]>{ + let summary = ""; + let description = [{ + }]; + + let arguments = (ins Variadic:$domains, Index:$dim_size); + let results = (outs IndexTree_DomainType:$domain); +} + +def DomainGetSize : IndexTree_Op<"DomainGetSize", [Pure]>{ + let summary = ""; + let description = [{}]; + + let arguments = (ins IndexTree_DomainType:$domain); + let results = (outs Index:$result); +} + +def IndexTreeIndexToTensorOp : IndexTree_Op<"IndexToTensorDim", [Pure]>{ + let summary = ""; + let description = [{}]; + + let arguments = (ins + TA_AnyTensor:$tensor, + IndexTree_IndexNodeType:$index, + UI32Attr:$dim, + Optional:$prev_dim + ); + + let results = (outs + Index:$crd, + Index:$pos + ); +} + +def DeclDomainOp : IndexTree_Op<"DeclDomainOp", [Pure]>{ + let summary = ""; + let description = [{ + }]; + + let arguments = (ins Index:$dim_size, Index:$num_rows, OptionalAttr:$is_dynamic, OptionalAttr:$indices_bitwidth); + let results = (outs IndexTree_SymbolicDomainType); +} + +def ComputeSymbolicDomainOp : IndexTree_Op<"ComputeSymbolicDomainOp", [Pure]>{ + let summary = ""; + let description = [{}]; + + let arguments = (ins IndexTree_NodeType:$parent, IndexTree_SymbolicDomainType:$domain, DefaultValuedAttr:$is_unique); + let results = (outs IndexTree_SymbolicDomainType); +} + +def ComputeSymbolicDomainRowOp : IndexTree_Op<"ComputeSymbolicDomainRowOp", [Pure]>{ + let summary = ""; + let description = [{}]; + + let arguments = (ins IndexTree_NodeType:$parent, IndexTree_SymbolicDomainType:$domain, DefaultValuedAttr:$needs_mark); + let results = (outs IndexTree_SymbolicDomainType); +} + +def SymbolicDomainInsertOp : IndexTree_Op<"SymbolicDomainInsertOp", [Pure]>{ + let summary = ""; + let description = [{}]; + + let arguments = (ins IndexTree_SymbolicDomainType:$domain, Index:$crd, DefaultValuedAttr:$is_unique); + let results = (outs IndexTree_SymbolicDomainType); +} + +def SymbolicDomainEndRowOp : IndexTree_Op<"SymbolicDomainEndRowOp", [Pure]>{ + let summary = ""; + let description = [{}]; + + let arguments = (ins IndexTree_SymbolicDomainType:$domain, DefaultValuedAttr:$needs_mark); + let results = (outs IndexTree_SymbolicDomainType); +} + +def IndexTreeSparseTensorOp : IndexTree_Op<"IndexTreeSparseTensorOp", [Pure]>{ + let summary = "Declare a sparse tensor from an index tree domain"; + let description = [{}]; + + let arguments = (ins Variadic:$domains); + let results = (outs TA_AnyTensor); +} + +def IndexTreeCleanWorkspaceOp : IndexTree_Op<"WorkspaceStartRowOp", [Pure]>{ + let summary = "Associate the loop of this index variable with a clean workspace"; + let description = [{}]; + + let arguments = (ins IndexTree_NodeType:$parent, WorkspaceTensor:$workspace); + let results = (outs WorkspaceTensor:$result); +} #endif // INDEXTREE_OPS \ No newline at end of file diff --git a/include/comet/Dialect/IndexTree/IR/IndexTreeTraits.h b/include/comet/Dialect/IndexTree/IR/IndexTreeTraits.h new file mode 100644 index 00000000..8b2deebb --- /dev/null +++ b/include/comet/Dialect/IndexTree/IR/IndexTreeTraits.h @@ -0,0 +1,36 @@ +// +// Copyright 2022 Battelle Memorial Institute +// +// Redistribution and use in source and binary forms, with or without modification, +// are permitted provided that the following conditions are met: +// +// 1. Redistributions of source code must retain the above copyright notice, this list of conditions +// and the following disclaimer. +// +// 2. Redistributions in binary form must reproduce the above copyright notice, this list of conditions +// and the following disclaimer in the documentation and/or other materials provided with the distribution. +// +// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED +// WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR +// PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, +// INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE +// GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, +// WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +// OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. +// + +#ifndef INDEXTREE_TRAITS_H_ +#define INDEXTREE_TRAITS_H_ + +#include "mlir/IR/Dialect.h" +#include "mlir/IR/OpDefinition.h" + +namespace mlir { +namespace indexTree { + template + class UnknownDomain : public ::mlir::OpTrait::TraitBase {}; + +} // indexTree +} // mlir + +#endif //INDEXTREE_TRAITS_H_ \ No newline at end of file diff --git a/include/comet/Dialect/IndexTree/IR/IndexTreeTypes.td b/include/comet/Dialect/IndexTree/IR/IndexTreeTypes.td new file mode 100644 index 00000000..abe7a01a --- /dev/null +++ b/include/comet/Dialect/IndexTree/IR/IndexTreeTypes.td @@ -0,0 +1,85 @@ +// +// Copyright 2022 Battelle Memorial Institute +// +// Redistribution and use in source and binary forms, with or without modification, +// are permitted provided that the following conditions are met: +// +// 1. Redistributions of source code must retain the above copyright notice, this list of conditions +// and the following disclaimer. +// +// 2. Redistributions in binary form must reproduce the above copyright notice, this list of conditions +// and the following disclaimer in the documentation and/or other materials provided with the distribution. +// +// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED +// WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR +// PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, +// INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE +// GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, +// WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +// OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. +// + +#ifndef INDEXTREE_TYPES +#define INDEXTREE_TYPES + +include "mlir/IR/AttrTypeBase.td" +include "IndexTreeBase.td" + +class IndexTree_Type traits = []> + : TypeDef { + let mnemonic = typeMnemonic; +} + +def IndexTree_TreeType : IndexTree_Type<"IndexTree", "index_tree"> { + let summary = "Operand for compute expression"; + let description = [{}]; +} + +def IndexTree_IndexNodeType : IndexTree_Type<"IndexNode", "index"> { + let summary = "Iteration Tree for Index Variables"; + let description = [{ + Type for storing iteration tree of index variables. + }]; +} + +def IndexTree_NodeType : AnyTypeOf<[IndexTree_TreeType, IndexTree_IndexNodeType]>; + +def IndexTree_TensorAccessType : IndexTree_Type<"TensorAccess", "tensor_access"> { + let summary = "Tensor access variabls"; + let description = [{ + Type for storing tensor access of index variables. + }]; +} + +def IndexTree_OperandType : IndexTree_Type<"Operand", "operand"> { + let summary = "Operand for compute expression"; + let description = [{}]; +} + +def IndexTree_SymbolicDomainType : IndexTree_Type<"SymbolicDomain", "symbolic_domain"> { + let summary = "Type representing a computed iteration/tensor domain"; + let description = [{}]; + let parameters = ( + ins "unsigned":$indices_bitwidth + ); + + let assemblyFormat = "`<`$indices_bitwidth`>`"; +} + +def IndexTree_DomainType : IndexTree_Type<"Domain", "domain"> { + let summary = "Type representing an iteration domain"; + let description = [{ + Type for storing iteration domain of an index variable. + }]; +} + +def IndexTree_AnyDomainType : AnyTypeOf<[IndexTree_SymbolicDomainType, IndexTree_DomainType]>; + +def IndexTree_ProvenanceGraphType : IndexTree_Type<"ProvenanceGraph", "prov_graph"> { + let summary = "Provenance graph for index tree transformations"; + let description = [{ + Type for storing provenance graph associated with a index. + }]; +} + +#endif // INDEXTREE_TYPES diff --git a/include/comet/Dialect/IndexTree/Passes.h b/include/comet/Dialect/IndexTree/Passes.h index 951c0856..48a83ecc 100644 --- a/include/comet/Dialect/IndexTree/Passes.h +++ b/include/comet/Dialect/IndexTree/Passes.h @@ -33,12 +33,26 @@ namespace mlir /// Generate the code for registering conversion passes. #define GEN_PASS_DECL #include "comet/Dialect/IndexTree/Passes.h.inc" + // Create a pass for infering the domain from the use of the index variables + std::unique_ptr createIndexTreeDomainInferencePass(); + + // Create a pass for concretizing the domain from the tensor definitions + std::unique_ptr createIndexTreeDomainConcretizationPass(); + + // Create a pass for creating the symbolic pass + std::unique_ptr createIndexTreeSymbolicComputePass(); + + // Create a pass for inlining the index tree + std::unique_ptr createIndexTreeInliningPass(); /// Create a pass for applying compressed workspace transformation into IndexTreeIR std::unique_ptr createIndexTreeWorkspaceTransformationsPass(); /// Create a pass for the redundancy-aware kernel fusion on index tree dialect for some compound expressions std::unique_ptr createIndexTreeKernelFusionPass(); + + /// Create a pass for reducing the dimensions of intermediate tensors after indextree-kernel-fusion + std::unique_ptr createIndexTreeDimensionReductionPass(); } } diff --git a/include/comet/Dialect/IndexTree/Passes.td b/include/comet/Dialect/IndexTree/Passes.td index 62eae9bd..55cb72df 100644 --- a/include/comet/Dialect/IndexTree/Passes.td +++ b/include/comet/Dialect/IndexTree/Passes.td @@ -30,14 +30,101 @@ include "mlir/Pass/PassBase.td" /// Kernel Fusion ///===----------------------------------------------------------------------===/// -def IndexTreeKernelFusion : Pass<"indextree-kernel-fusion"> { +def IndexTreeKernelFusion : Pass<"indextree-kernel-fusion", "func::FuncOp"> { let summary = "The redundancy-aware kernel fusion on index tree dialect for compound expressions"; let description = [{ + This pass fuses multiple index trees into one to reduce the memory footprint + and time complexity if possible. - }]; + For example, for GNN kernel A = B * C * D where B is sparse, C and D are + dense, the DSL may be + T[i, h] = B[i, k] * C[k, h]; + A[i, j] = T[i, h] * D[h, j]; + + The original pseudo code is + void no_fution_index_tree() + { + for (h = 0 to NH) { + for (i = 0 to NI) { + for (k = 0 to NK) { + T[i, h] += B[i, k] * C[k, h]; + } + } + } + for (h = 0 to NH) { + for (i = 0 to NI) { + for (j = 0 to NJ) { + A[i, j] += T[i, h] * D[h, j]; + } + } + } + } + + The fused version would be + void partial_fusion_index_tree() + { + for (h = 0 to NH) { + for (i = 0 to NI) { + for (k = 0 to NK) { + T[i, h] += B[i, k] * C[k, h]; + } + for (j = 0 to NJ) { + A[i, j] += T[i, h] * D[h, j]; + } + } + } + } + + }]; let constructor = "comet::createIndexTreeKernelFusionPass()"; let dependentDialects = [ - "comet::IndexTreeDialect", + "indexTree::IndexTreeDialect", + "memref::MemRefDialect", + "scf::SCFDialect" + ]; +} + +def IndexTreeDimensionReduction : Pass<"indextree-dimension-reduction", "func::FuncOp"> { + let summary = "Reduce the dimensions of intermediate tensors after indextree-kernel-fusion"; + let description = [{ + After the index tree kernel fusion, the dimensions of intermediate tensors + could be reduced. For example, for the fused GNN kernel + void partial_fusion_index_tree() + { + for (h = 0 to NH) { + for (i = 0 to NI) { + for (k = 0 to NK) { + T[i, h] += B[i, k] * C[k, h]; + } + for (j = 0 to NJ) { + A[i, j] += T[i, h] * D[h, j]; + } + } + } + } + + The dimension of T could be reduced from 2 to 0, making it a scalar. + void partial_fusion_index_tree() + { + for (h = 0 to NH) { + for (i = 0 to NI) { + for (k = 0 to NK) { + t += B[i, k] * C[k, h]; + } + for (j = 0 to NJ) { + A[i, j] += t * D[h, j]; + } + t = 0; + } + } + } + + The reduced dimensions are the shared dimensions of the two compute nodes. + + }]; + let constructor = "comet::createIndexTreeDimensionReductionPass()"; + let dependentDialects = [ + "indexTree::IndexTreeDialect", "memref::MemRefDialect", "scf::SCFDialect" ]; @@ -47,19 +134,53 @@ def IndexTreeKernelFusion : Pass<"indextree-kernel-fusion"> { /// Workspace Transformations ///===----------------------------------------------------------------------===/// -def IndexTreeWorkspaceTranformations: Pass<"indextree-workspace-transformations"> { +def IndexTreeWorkspaceTranformations: Pass<"indextree-workspace-transformations", "func::FuncOp"> { let summary = "Compressed workspace transformation on IndexTree dialect" "to produce sparse output"; - let description = [{ + let description = [{}]; + let constructor = "comet::createIndexTreeWorkspaceTransformationsPass()"; + let dependentDialects = ["indexTree::IndexTreeDialect"]; +} - }]; - let constructor = "comet::createIndexTreeWorkspaceTransformationsPass("; - let dependentDialects = [ - "comet::IndexTreeDialect", - "memref::MemRefDialect", - "scf::SCFDialect" - ]; +def IndexTreeDomainInference : Pass<"indextree-domain-inference", "func::FuncOp"> { + let summary = "Infer domain of index variables"; + let description = [{ + Propogate domain values from tensors, through compute operations to index variables. + Necessary for eventual lowering of index variables to for loops + }]; + + let constructor = "comet::createIndexTreeDomainInferencePass()"; + let dependentDialects = ["indexTree::IndexTreeDialect"]; } +def IndexTreeDomainConcretization : Pass<"indextree-domain-concretization", "func::FuncOp"> { + let summary = "Transform tensor domains into concrete descriptions of dense or sparse domains"; + let description = [{ + Transform tensor domains into concrete descriptions of dense or sparse domains + }]; + + let constructor = "comet::createIndexTreeDomainConcretizationPass()"; + let dependentDialects = ["indexTree::IndexTreeDialect"]; +} + +def IndexTreeSymbolicComputePass : Pass<"indextree-symbolic-compute", "func::FuncOp"> { + let summary = "Create index-tree for symbolic computations for sparse output"; + let description = [{ + Inserting values into a sparse output tensor cannot be done in parallel because it requires + sequential access to the row pointers and dynamic allocation. The symbolic pass + computes the size of allocations and row pointers without performing the computations + so the sparse tensors can be allocated and the computations can be done in parallel. + }]; + + let constructor = "comet::createIndexTreeSymbolicComputePass()"; + let dependentDialects = ["indexTree::IndexTreeDialect"]; +} + +def IndexTreeInliningPass : Pass <"indextree-inlining", "func::FuncOp"> { + let summary = "Inline index tree"; + let description = [{}]; + let constructor = "comet::createIndexTreeInliningPass()"; + let dependentDialects = ["indexTree::IndexTreeDialect"]; +} #endif /// COMET_DIALECT_INDEXTREE_PASSES diff --git a/include/comet/Dialect/IndexTree/Patterns.h b/include/comet/Dialect/IndexTree/Patterns.h new file mode 100644 index 00000000..b0fa1bab --- /dev/null +++ b/include/comet/Dialect/IndexTree/Patterns.h @@ -0,0 +1,43 @@ +//===- Patterns.h - Conversion Pass Construction and Registration -----------===// +// +// Copyright 2022 Battelle Memorial Institute +// +// Redistribution and use in source and binary forms, with or without modification, +// are permitted provided that the following conditions are met: +// +// 1. Redistributions of source code must retain the above copyright notice, this list of conditions +// and the following disclaimer. +// +// 2. Redistributions in binary form must reproduce the above copyright notice, this list of conditions +// and the following disclaimer in the documentation and/or other materials provided with the distribution. +// +// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED +// WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR +// PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, +// INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE +// GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, +// WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +// OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. +// +//===----------------------------------------------------------------------===// + +#ifndef COMET_DIALECT_INDEXTREE_PATTERNS_H +#define COMET_DIALECT_INDEXTREE_PATTERNS_H + +#include "mlir/Transforms/DialectConversion.h" + +#include "comet/Dialect/IndexTree/Analysis/CopiedDomainAnalysis.h" + +namespace mlir +{ + namespace indexTree + { + void populateDomainInferencePatterns(MLIRContext *context, RewritePatternSet &patterns, CopiedDomainAnalysis &copiedDomains); + void populateDomainConcretizationPatterns(MLIRContext *context, RewritePatternSet &patterns); + void populateIndexTreeTypeConversionPatterns(MLIRContext *context, RewritePatternSet &patterns, TypeConverter &typeConverter, ConversionTarget& target); + void populateIndexTreeInliningPatterns(MLIRContext *context, RewritePatternSet &patterns); + void populateMaskDomainTransformationPatterns(MLIRContext *context, RewritePatternSet &patterns); + } +} + +#endif // COMET_DIALECT_INDEXTREE_PASSES_H diff --git a/include/comet/Dialect/TensorAlgebra/IR/CMakeLists.txt b/include/comet/Dialect/TensorAlgebra/IR/CMakeLists.txt index b5a0ddeb..dc504d68 100644 --- a/include/comet/Dialect/TensorAlgebra/IR/CMakeLists.txt +++ b/include/comet/Dialect/TensorAlgebra/IR/CMakeLists.txt @@ -3,4 +3,13 @@ mlir_tablegen(TAOps.h.inc -gen-op-decls) mlir_tablegen(TAOps.cpp.inc -gen-op-defs) mlir_tablegen(TADialect.h.inc -gen-dialect-decls) mlir_tablegen(TADialect.cpp.inc -gen-dialect-defs) -add_public_tablegen_target(COMETTensorAlgebraOpsIncGen) \ No newline at end of file +mlir_tablegen(TAEnums.h.inc -gen-enum-decls) +mlir_tablegen(TAEnums.cpp.inc -gen-enum-defs) +mlir_tablegen(TAAttrs.h.inc -gen-attrdef-decls) +mlir_tablegen(TAAttrs.cpp.inc -gen-attrdef-defs) +add_public_tablegen_target(COMETTensorAlgebraOpsIncGen) + +set(LLVM_TARGET_DEFINITIONS TATypes.td) +mlir_tablegen(TATypes.h.inc -gen-typedef-decls) +mlir_tablegen(TATypes.cpp.inc -gen-typedef-defs) +add_public_tablegen_target(COMETTensorAlgebraTypesIncGen) \ No newline at end of file diff --git a/include/comet/Dialect/TensorAlgebra/IR/TAAttrs.td b/include/comet/Dialect/TensorAlgebra/IR/TAAttrs.td new file mode 100644 index 00000000..6a2378b6 --- /dev/null +++ b/include/comet/Dialect/TensorAlgebra/IR/TAAttrs.td @@ -0,0 +1,31 @@ +// +// Copyright 2022 Battelle Memorial Institute +// +// Redistribution and use in source and binary forms, with or without modification, +// are permitted provided that the following conditions are met: +// +// 1. Redistributions of source code must retain the above copyright notice, this list of conditions +// and the following disclaimer. +// +// 2. Redistributions in binary form must reproduce the above copyright notice, this list of conditions +// and the following disclaimer in the documentation and/or other materials provided with the distribution. +// +// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED +// WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR +// PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, +// INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE +// GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, +// WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +// OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. +// + +#ifndef TA_ATTRS +#define TA_ATTRS + +include "comet/Dialect/TensorAlgebra/IR/TAEnums.td" + +def TAFormatArrayAttr : TypedArrayAttrBase { + let constBuilderCall = "$_builder.getI32ArrayAttr($0)"; +} + +#endif //TA_ATTRS \ No newline at end of file diff --git a/include/comet/Dialect/TensorAlgebra/IR/TABase.td b/include/comet/Dialect/TensorAlgebra/IR/TABase.td new file mode 100644 index 00000000..d6775aed --- /dev/null +++ b/include/comet/Dialect/TensorAlgebra/IR/TABase.td @@ -0,0 +1,37 @@ +// +// Copyright 2022 Battelle Memorial Institute +// +// Redistribution and use in source and binary forms, with or without modification, +// are permitted provided that the following conditions are met: +// +// 1. Redistributions of source code must retain the above copyright notice, this list of conditions +// and the following disclaimer. +// +// 2. Redistributions in binary form must reproduce the above copyright notice, this list of conditions +// and the following disclaimer in the documentation and/or other materials provided with the distribution. +// +// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED +// WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR +// PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, +// INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE +// GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, +// WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +// OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. +// + +#ifndef TA_BASE +#define TA_BASE +include "mlir/IR/OpBase.td" + +/// Provide a definition of the 'TA' dialect in the ODS framework so that we +/// can define our operations. +def TA_Dialect : Dialect { + let name = "ta"; + let cppNamespace = "::mlir::tensorAlgebra"; + + /// We set this bit to generate the declarations for the dialect's type parsing + /// and printing hooks. + let useDefaultTypePrinterParser = 1; +} + +#endif //TA_BASE \ No newline at end of file diff --git a/include/comet/Dialect/TensorAlgebra/IR/TADialect.h b/include/comet/Dialect/TensorAlgebra/IR/TADialect.h index 114f873e..2ad8ea04 100644 --- a/include/comet/Dialect/TensorAlgebra/IR/TADialect.h +++ b/include/comet/Dialect/TensorAlgebra/IR/TADialect.h @@ -28,6 +28,7 @@ #ifndef TENSORALGEBRA_DIALECT_H_ #define TENSORALGEBRA_DIALECT_H_ + #include "mlir/IR/BuiltinTypes.h" #include "mlir/IR/Dialect.h" #include "mlir/Interfaces/FunctionInterfaces.h" @@ -35,11 +36,25 @@ #include "mlir/Interfaces/CallInterfaces.h" #include "mlir/Interfaces/CastInterfaces.h" #include "mlir/IR/PatternMatch.h" +#include "mlir/Interfaces/DestinationStyleOpInterface.h" +#include "mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h" /// Include the auto-generated header file containing the declaration of the Tensor Algebra /// dialect. #include "comet/Dialect/TensorAlgebra/IR/TADialect.h.inc" +/// Include the auto-generated enum declerations +//===---------------------------------------------------------------------===// +#include "comet/Dialect/TensorAlgebra/IR/TAEnums.h.inc" + +#define GET_TYPEDEF_CLASSES +#include "comet/Dialect/TensorAlgebra/IR/TATypes.h.inc" + +/// Include the auto-generated header file containing the declaration of the index tree +/// types. +#define GET_ATTRDEF_CLASSES +#include "comet/Dialect/TensorAlgebra/IR/TAAttrs.h.inc" + /// Include the auto-generated header file containing the declarations of the /// tensorAlgbra operations and also the operations of the Shape Inference Op Interface. //===----------------------------------------------------------------------===// @@ -50,13 +65,6 @@ namespace mlir { namespace tensorAlgebra { - std::vector getFormatsValue(std::string formats_str, int rank_size, PatternRewriter &rewriter, Location loc, IndexType indexType); - - namespace detail - { - struct SparseTensorTypeStorage; - } /// end namespace detail - void populateMultiOpFactorizationPatterns( RewritePatternSet &patterns, MLIRContext *context); @@ -69,53 +77,7 @@ namespace mlir void populateSTCRemoveDeadOpsPatterns( RewritePatternSet &patterns, MLIRContext *context); - //===----------------------------------------------------------------------===// - /// Tensor Algebra Types - //===----------------------------------------------------------------------===// - - class IndexLabelType : public mlir::Type::TypeBase - { - public: - /// Used for generic hooks in TypeBase. - using Base::Base; - - static IndexLabelType get(MLIRContext *context) - { - /// Custom, uniq'ed construction in the MLIRContext. - return Base::get(context); - } - - /// The name of this struct type. - static constexpr StringLiteral name = "ta.indexLabel"; - }; - - /// This class defines the TA sparse tensor type. It represents a collection of - /// element types for data and indices of COO format. - /// All derived types in MLIR must inherit from the CRTP class - /// 'Type::TypeBase'. It takes as template parameters the concrete type - /// (SparseTensorType), the base class to use (Type), and the storage class - /// (SparseTensorTypeStorage). - class SparseTensorType : public mlir::Type::TypeBase - { - public: - /// Inherit some necessary constructors from 'TypeBase'. - using Base::Base; - - /// Create an instance of a `SparseTensorType` with the given element types. There - /// *must* be atleast one element type. - static SparseTensorType get(llvm::ArrayRef elementTypes); - - /// Returns the element types of this sparse tensor type. - llvm::ArrayRef getElementTypes(); - - /// Returns the number of element type held by this sparse tensor. - size_t getNumElementTypes() { return getElementTypes().size(); } - - /// The name of this struct type. - static constexpr StringLiteral name = "ta.spTensor"; - }; - + void registerBufferizableOpInterfaceExternalModels(DialectRegistry ®istry); } /// end namespace tensorAlgebra } /// end namespace mlir diff --git a/include/comet/Dialect/TensorAlgebra/IR/TAEnums.td b/include/comet/Dialect/TensorAlgebra/IR/TAEnums.td new file mode 100644 index 00000000..b177d01f --- /dev/null +++ b/include/comet/Dialect/TensorAlgebra/IR/TAEnums.td @@ -0,0 +1,40 @@ +// +// Copyright 2022 Battelle Memorial Institute +// +// Redistribution and use in source and binary forms, with or without modification, +// are permitted provided that the following conditions are met: +// +// 1. Redistributions of source code must retain the above copyright notice, this list of conditions +// and the following disclaimer. +// +// 2. Redistributions in binary form must reproduce the above copyright notice, this list of conditions +// and the following disclaimer in the documentation and/or other materials provided with the distribution. +// +// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED +// WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR +// PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, +// INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE +// GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, +// WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +// OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. +// + +#ifndef TA_ENUMS +#define TA_ENUMS + +include "mlir/IR/EnumAttr.td" + +def Unknown: I32EnumAttrCase<"UNK", 0, "unk">; +def Dense: I32EnumAttrCase<"D", 1, "d">; +def Compressed: I32EnumAttrCase<"CU", 2, "cu">; +def CompressedNonunique: I32EnumAttrCase<"CN", 3, "cn">; +def Singular: I32EnumAttrCase<"S", 4, "s">; + +def TAFormatEnum: I32EnumAttr<"TensorFormatEnum", "Valid format specifiers", + [Unknown, Dense, Compressed, CompressedNonunique, Singular]> { + let cppNamespace = "::mlir::tensorAlgebra"; + let stringToSymbolFnName = "ConvertToEnum"; + let symbolToStringFnName = "ConvertToString"; +} + +#endif //TA_ENUMS \ No newline at end of file diff --git a/include/comet/Dialect/TensorAlgebra/IR/TAOps.td b/include/comet/Dialect/TensorAlgebra/IR/TAOps.td index 40398645..148dc97c 100644 --- a/include/comet/Dialect/TensorAlgebra/IR/TAOps.td +++ b/include/comet/Dialect/TensorAlgebra/IR/TAOps.td @@ -29,35 +29,18 @@ #define TA_OPS -include "mlir/IR/OpBase.td" +include "mlir/IR/OpBase.td" include "mlir/Interfaces/FunctionInterfaces.td" include "mlir/IR/SymbolInterfaces.td" include "mlir/Interfaces/CallInterfaces.td" include "mlir/Interfaces/CastInterfaces.td" include "mlir/Interfaces/SideEffectInterfaces.td" +include "mlir/Interfaces/DestinationStyleOpInterface.td" +include "mlir/Dialect/Bufferization/IR/BufferizableOpInterface.td" -/// Provide a definition of the 'TA' dialect in the ODS framework so that we -/// can define our operations. -def TA_Dialect : Dialect { - let name = "ta"; - let cppNamespace = "::mlir::tensorAlgebra"; - - /// We set this bit to generate the declarations for the dialect's type parsing - /// and printing hooks. - let useDefaultTypePrinterParser = 1; - -} - -/// An implementation of IndexLabelType. -def TA_IndexLabelType : - DialectType()">, - "IndexLabelType">; - - -/// Whether a type is a IndexLabelType. -def TAIsIndexLabelTypePred : CPred<"$_self.isa()">; -def IndexLabel : Type; +include "comet/Dialect/TensorAlgebra/IR/TABase.td" +include "comet/Dialect/TensorAlgebra/IR/TATypes.td" +include "comet/Dialect/TensorAlgebra/IR/TAAttrs.td" /// Base class for ta dialect operations. This operation inherits from the base /// `Op` class in OpBase.td, and provides: @@ -67,14 +50,33 @@ def IndexLabel : Type; class TA_Op traits = []> : Op; -/// Provide a definition for the TA SparseTensorType for use in ODS. -/// This allows for using SparseTensorType in a similar way to Tensor or MemRef. -def SparseTensor : - Type()">, "TA sparse tensor type">; +//===----------------------------------------------------------------------===// +/// Tensor Algebra Operations +//===----------------------------------------------------------------------===// +def IndexLabelStaticOp : + TA_Op<"static_index_label", [Pure]>, + Arguments<(ins Index:$min, Index:$max, Index:$step)>, + Results<(outs Range)> { + let summary = "Create an index label type value, used to create views"; + let description = [{ + The `ta.static_index_label` op creates a ta.static_index_label from 3 values of type `index` + that represent the min, max and step values of a range. + + Example: -/// Provide a definition of the types that are used within the TA dialect. -def TA_AnyTensor : AnyTypeOf<[TensorOf<[AnyType]>, SparseTensor]>; + %3 = ta.static_index_label %0:%1:%2 : !ta.index_label + }]; + + let builders = [OpBuilder<(ins "Value":$min, "Value":$max, "Value":$step), + [{ + auto rangeType = RangeType::get($_builder.getContext()); + build($_builder, $_state, rangeType, min, max, step); + }]>]; + + //TODO(gkestor): add verifier + //let hasVerifier = 1; +} //===----------------------------------------------------------------------===// /// Tensor Algebra Operations @@ -118,11 +120,11 @@ def DenseTensorDeclOp : TA_Op<"dense_tensor_decl", [Pure]> { }]; /// The constant operation takes an attribute as the only input. - let arguments = (ins Variadic:$labels, StrAttr:$format); + let arguments = (ins Variadic:$labels); /// The constant operation returns a single value of TensorType. let results = (outs AnyTensor); - + /// Invoke a static verify method to verify this constant operation. //TODO(gkestor): add verifier @@ -145,30 +147,30 @@ def SparseTensorDeclOp : TA_Op<"spTensor_decl", [Pure]> { }]; /// The constant operation takes an attribute as the only input. - let arguments = (ins Variadic:$labels, StrAttr:$format, BoolAttr:$temporal_tensor); + let arguments = (ins Variadic:$labels, BoolAttr:$temporal_tensor); /// The constant operation returns a single value of TensorType. let results = (outs TA_AnyTensor); let extraClassDeclaration = [{ unsigned int getParameterCount() { - mlir::TensorType type = getResult().getType().cast(); + mlir::tensorAlgebra::SparseTensorType type = cast(getResult().getType()); return (type.getRank() * 6) + 1; } unsigned int getDimArrayCount() { - mlir::TensorType type = getResult().getType().cast(); + mlir::tensorAlgebra::SparseTensorType type = cast(getResult().getType()); return type.getRank() * 4; } unsigned int getValueArrayPos() { - mlir::TensorType type = getResult().getType().cast(); + mlir::tensorAlgebra::SparseTensorType type = cast(getResult().getType()); return (type.getRank() * 4) + 1; } unsigned int getTotalArrayCount() { - mlir::TensorType type = getResult().getType().cast(); + mlir::tensorAlgebra::SparseTensorType type = cast(getResult().getType()); return (type.getRank() * 4) + 1; } }]; @@ -195,10 +197,10 @@ def SparseOutputTensorDeclOp : TA_Op<"sparse_output_tensor_decl", [Pure]> { }]; /// The constant operation takes an attribute as the only input. - let arguments = (ins Variadic:$labels, StrAttr:$format); + let arguments = (ins Variadic:$labels); /// The constant operation returns a single value of TensorType. - let results = (outs AnyTensor); + let results = (outs TA_AnyTensor); /// Invoke a static verify method to verify this constant operation. @@ -225,10 +227,10 @@ def TempSparseOutputTensorDeclOp : TA_Op<"temp_sparse_output_tensor_decl", [Pure }]; /// The constant operation takes an attribute as the only input. - let arguments = (ins Variadic:$labels, StrAttr:$format); + let arguments = (ins Variadic:$labels); /// The constant operation returns a single value of TensorType. - let results = (outs AnyTensor); + let results = (outs TA_AnyTensor); //TODO(gkestor): add verifier //let hasVerifier = 1; @@ -275,7 +277,7 @@ def DenseConstantOp : TA_Op<"constant", [Pure]> { let hasVerifier = 1; } -def SparseTensorConstructOp : TA_Op<"spTensor_construct", [Pure]>{ +def SparseTensorConstructOp : TA_Op<"spTensor_construct", [SameVariadicOperandSize]>{ let summary = ""; let description = [{ @@ -291,12 +293,12 @@ def SparseTensorConstructOp : TA_Op<"spTensor_construct", [Pure]>{ //Aval_size (size of value array) //dim1_size, dim2_size (size of each dimension in sparse tensor) //TODO(gkestor): might be better to have a struct with all the data elements - let arguments = (ins Variadic:$indices, I32Attr:$tensor_rank); + let arguments = (ins TensorOf<[Index]>:$dims, Variadic>:$pos_indices, Variadic>:$crd_indices, Variadic>:$block_pos_indices, Variadic>:$block_crd_indices, TA_FloatTensor:$vals, I32Attr:$tensor_rank, TAFormatArrayAttr:$dimension_formats); let results = (outs TA_AnyTensor:$output); - let assemblyFormat = [{ - `(` $indices `)` attr-dict `:` `(` type($indices) `)` `->` `(` type($output) `)` - }]; + // let assemblyFormat = [{ + // `(` $indices `)` attr-dict `:` `(` type($indices) `)` `->` `(` type($output) `)` + // }]; // TODO(pflynn) This may need to be adjusted // This works for the moment; the idea is that each rank has a block @@ -330,7 +332,7 @@ def SparseTensorConstructOp : TA_Op<"spTensor_construct", [Pure]>{ } -def TensorFillOp : TA_Op<"fill", [Pure]>{ +def TensorFillOp : TA_Op<"fill">{ let summary = ""; let description = [{ @@ -350,7 +352,7 @@ def TensorMultOp : TA_Op<"mul", [Pure, AttrSizedOperandSegments]>{ }]; let arguments = (ins TA_AnyTensor:$rhs1, TA_AnyTensor:$rhs2, Variadic:$index_labels, - AffineMapArrayAttr:$indexing_maps, StrArrayAttr:$formats, StrAttr:$semiring, + AffineMapArrayAttr:$indexing_maps, StrAttr:$semiring, OptionalAttr:$MaskType, /// TODO: try to use: DefaultValuedAttr Optional:$mask); @@ -358,13 +360,29 @@ def TensorMultOp : TA_Op<"mul", [Pure, AttrSizedOperandSegments]>{ let results = (outs TA_AnyTensor); let extraClassDeclaration = [{ + template + size_t getTensorRank(T tensor) + { + if(auto denseT = mlir::dyn_cast(tensor.getType())) + { + return denseT.getRank(); + } + else if(auto sparseT = mlir::dyn_cast(tensor.getType())) + { + return sparseT.getRank(); + } + else + { + assert(false&& "Unexpected Type"); + } + } + std::vector getRhs1IndexLabels() { std::vector labels; size_t start = 0; - mlir::TensorType type = getRhs1().getType().cast(); - size_t end = type.getRank(); + size_t end = getTensorRank(getRhs1()); for(size_t i = start; i < end; i++ ) { @@ -378,13 +396,11 @@ def TensorMultOp : TA_Op<"mul", [Pure, AttrSizedOperandSegments]>{ std::vector getRhs2IndexLabels() { std::vector labels; - - mlir::TensorType type1 = getRhs1().getType().cast(); - mlir::TensorType type2 = getRhs2().getType().cast(); - size_t start = type1.getRank(); + + size_t start =getTensorRank(getRhs1()); - size_t end = static_cast(type1.getRank() + type2.getRank()); + size_t end = getTensorRank(getRhs1()) + getTensorRank(getRhs2()); for(size_t i = start; i < end; i++ ) { @@ -400,10 +416,7 @@ def TensorMultOp : TA_Op<"mul", [Pure, AttrSizedOperandSegments]>{ std::vector labels; size_t end = getIndexLabels().size(); - mlir::TensorType type1 = getRhs1().getType().cast(); - mlir::TensorType type2 = getRhs2().getType().cast(); - - size_t start = type1.getRank() + type2.getRank(); + size_t start = getTensorRank(getRhs1()) + getTensorRank(getRhs2()); for(size_t i = start; i < end; i++ ) { @@ -429,20 +442,34 @@ def TensorElewsMultOp : TA_Op<"elews_mul", [Pure]>{ TA_AnyTensor:$rhs2, Variadic:$index_labels, AffineMapArrayAttr:$indexing_maps, - StrArrayAttr:$formats, StrAttr:$semiring, OptionalAttr:$MaskType); let results = (outs TA_AnyTensor); let extraClassDeclaration = [{ + template + size_t getTensorRank(T tensor) + { + if(auto denseT = mlir::dyn_cast(tensor.getType())) + { + return denseT.getRank(); + } + else if(auto sparseT = mlir::dyn_cast(tensor.getType())) + { + return sparseT.getRank(); + } + else + { + assert(false&& "Unexpected Type"); + } + } std::vector getRhs1IndexLabels() { std::vector labels; size_t start = 0; - mlir::TensorType type = getRhs1().getType().cast(); - size_t end = type.getRank(); + size_t end = getTensorRank(getRhs1()); for(size_t i = start; i < end; i++ ) { @@ -456,11 +483,9 @@ def TensorElewsMultOp : TA_Op<"elews_mul", [Pure]>{ std::vector getRhs2IndexLabels() { std::vector labels; - mlir::TensorType type1 = getRhs1().getType().cast(); - mlir::TensorType type2 = getRhs2().getType().cast(); - size_t start = type1.getRank(); + size_t start = getTensorRank(getRhs1()); - size_t end = static_cast(type1.getRank() + type2.getRank()); + size_t end = getTensorRank(getRhs1()) + getTensorRank(getRhs2()); for(size_t i = start; i < end; i++ ) { @@ -476,9 +501,7 @@ def TensorElewsMultOp : TA_Op<"elews_mul", [Pure]>{ std::vector labels; size_t end = getIndexLabels().size(); - mlir::TensorType type1 = getRhs1().getType().cast(); - mlir::TensorType type2 = getRhs2().getType().cast(); - size_t start = type1.getRank() + type2.getRank(); + size_t start = getTensorRank(getRhs1()) + getTensorRank(getRhs2()); for(size_t i = start; i < end; i++ ) @@ -517,7 +540,7 @@ def TensorDimOp : TA_Op<"dim", [Pure]>{ } -def TensorSetOp : TA_Op<"set_op", [Pure]>{ +def TensorSetOp : TA_Op<"set_op">{ let summary = ""; let description = [{ @@ -535,18 +558,35 @@ def TensorSetOp : TA_Op<"set_op", [Pure]>{ def TransposeOp : TA_Op<"transpose",[Pure]> { let summary = "transpose operation"; - let arguments = (ins TA_AnyTensor:$rhs, Variadic:$index_labels, AffineMapArrayAttr:$indexing_maps, StrArrayAttr:$formats); + let arguments = (ins TA_AnyTensor:$rhs, Variadic:$index_labels, AffineMapArrayAttr:$indexing_maps); let results = (outs TA_AnyTensor); let extraClassDeclaration = [{ + template + size_t getTensorRank(T tensor) + { + if(auto denseT = mlir::dyn_cast(tensor.getType())) + { + return denseT.getRank(); + } + else if(auto sparseT = mlir::dyn_cast(tensor.getType())) + { + return sparseT.getRank(); + } + else + { + assert(false&& "Unexpected Type"); + } + } + std::vector getRhsIndexLabels() { std::vector labels; size_t start = 0; - size_t end = static_cast(getRhs().getType().cast().getRank()); + size_t end = static_cast(getTensorRank(getRhs())); for(size_t i = start; i < end; i++ ) { labels.push_back(getIndexLabels()[i]); @@ -559,7 +599,7 @@ def TransposeOp : TA_Op<"transpose",[Pure]> { std::vector getResultIndexLabels() { std::vector labels; - size_t start = static_cast(getRhs().getType().cast().getRank()); + size_t start = static_cast(getTensorRank(getRhs())); size_t end = getIndexLabels().size(); for(size_t i = start; i < end; i++ ) @@ -580,7 +620,7 @@ def ReduceOp : TA_Op<"reduce",[Pure]> { let arguments = (ins TA_AnyTensor:$rhs); - let results = (outs F64:$lhs); + let results = (outs AnyTypeOf<[F32,F64,Index]>:$lhs); let builders = [OpBuilder< (ins "Value":$input)>]; @@ -601,18 +641,34 @@ def TensorAddOp : TA_Op<"add", TA_AnyTensor:$rhs2, Variadic:$index_labels, AffineMapArrayAttr:$indexing_maps, - StrArrayAttr:$formats, StrAttr:$semiring, OptionalAttr:$MaskType); let results = (outs TA_AnyTensor); let extraClassDeclaration = [{ + template + size_t getTensorRank(T tensor) + { + if(auto denseT = mlir::dyn_cast(tensor.getType())) + { + return denseT.getRank(); + } + else if(auto sparseT = mlir::dyn_cast(tensor.getType())) + { + return sparseT.getRank(); + } + else + { + assert(false&& "Unexpected Type"); + } + } + std::vector getRhs1IndexLabels() { std::vector labels; size_t start = 0; - size_t end = static_cast(getRhs1().getType().cast().getRank()); + size_t end = getTensorRank(getRhs1()); for(size_t i = start; i < end; i++ ) { labels.push_back(getIndexLabels()[i]); @@ -625,8 +681,8 @@ def TensorAddOp : TA_Op<"add", std::vector getRhs2IndexLabels() { std::vector labels; - size_t start = static_cast(getRhs1().getType().cast().getRank()); - size_t end = static_cast(getRhs1().getType().cast().getRank() + getRhs2().getType().cast().getRank()); + size_t start = getTensorRank(getRhs1()); + size_t end = getTensorRank(getRhs1()) + getTensorRank(getRhs2()); for(size_t i = start; i < end; i++ ) { @@ -640,7 +696,7 @@ def TensorAddOp : TA_Op<"add", std::vector getResultIndexLabels() { std::vector labels; - size_t start = static_cast(getRhs1().getType().cast().getRank() + getRhs2().getType().cast().getRank()); + size_t start = getTensorRank(getRhs1()) + getTensorRank(getRhs2()); size_t end = getIndexLabels().size(); for(size_t i = start; i < end; i++ ) @@ -666,18 +722,34 @@ def TensorSubtractOp : TA_Op<"subtract", TA_AnyTensor:$rhs2, Variadic:$index_labels, AffineMapArrayAttr:$indexing_maps, - StrArrayAttr:$formats, StrAttr:$semiring, OptionalAttr:$MaskType); let results = (outs TA_AnyTensor); let extraClassDeclaration = [{ + template + size_t getTensorRank(T tensor) + { + if(auto denseT = mlir::dyn_cast(tensor.getType())) + { + return denseT.getRank(); + } + else if(auto sparseT = mlir::dyn_cast(tensor.getType())) + { + return sparseT.getRank(); + } + else + { + assert(false&& "Unexpected Type"); + } + } + std::vector getRhs1IndexLabels() { std::vector labels; size_t start = 0; - size_t end = static_cast(getRhs1().getType().cast().getRank()); + size_t end = getTensorRank(getRhs1()); for(size_t i = start; i < end; i++ ) { labels.push_back(getIndexLabels()[i]); @@ -690,8 +762,8 @@ def TensorSubtractOp : TA_Op<"subtract", std::vector getRhs2IndexLabels() { std::vector labels; - size_t start = static_cast(getRhs1().getType().cast().getRank()); - size_t end = static_cast(getRhs1().getType().cast().getRank() + getRhs2().getType().cast().getRank()); + size_t start = getTensorRank(getRhs1()); + size_t end = getTensorRank(getRhs1()) + getTensorRank(getRhs2()); for(size_t i = start; i < end; i++ ) { @@ -705,7 +777,7 @@ def TensorSubtractOp : TA_Op<"subtract", std::vector getResultIndexLabels() { std::vector labels; - size_t start = static_cast(getRhs1().getType().cast().getRank() + getRhs2().getType().cast().getRank()); + size_t start = getTensorRank(getRhs1()) + getTensorRank(getRhs2()); size_t end = getIndexLabels().size(); for(size_t i = start; i < end; i++ ) @@ -830,7 +902,9 @@ def PrintOp : TA_Op<"print"> { /// The print operation takes an input tensor to print. /// We can extend the list of supported datatype for print with F64Tensor, I8MemRef, I64MemRef, F32MemRef, etc. - let arguments = (ins AnyTypeOf<[F64, + let arguments = (ins AnyTypeOf<[F32, + F64, + F32MemRef, F64MemRef, TA_AnyTensor]>:$input); } @@ -943,7 +1017,7 @@ def GenericCallOp : TA_Op<"generic_call", ]; } -def TensorFillFromFileOp : TA_Op<"fill_from_file", [Pure]>{ +def TensorFillFromFileOp : TA_Op<"fill_from_file">{ let summary = ""; let description = [{ }]; @@ -967,8 +1041,204 @@ def TensorCopyOp : TA_Op<"copy", [Pure]>{ //let hasVerifier = 1; } -#endif /// TA_OPS +def TensorInsertOp : TA_Op<"TAInsertOp", [Pure, SameVariadicOperandSize]>{ + let summary = "Insert intro a sparse tensor. Sparse Tensor equivalent of tensor.insert"; + let description = [{}]; + + let arguments = ( + ins TA_AnyTensor:$tensor, + Variadic:$pos, + Variadic:$crds, + AnyTypeOf<[F32,F64]>:$value + ); + + let results = (outs TA_AnyTensor); +} + +def TensorExtractOp : TA_Op<"TAExtractOp", [Pure, SameVariadicOperandSize]>{ + let summary = "Extract from a sparse tensor. Sparse Tensor equivalent of tensor.extract"; + let description = [{}]; + + let arguments = ( + ins TA_AnyTensor:$tensor, + Index:$pos, + Variadic:$crds, + AnyAttrOf<[F32Attr, F64Attr]>:$zero + ); + + let results = (outs AnyFloat); +} + +def SpTensorAliasOp : TA_Op<"SpTensorAlias", [Pure]>{ + let arguments = (ins TA_AnyTensor:$tensor); + let results = (outs TA_AnyTensor); +} + +def SpTensorGetDimPos : TA_Op<"SpTensorGetDimPos", [Pure]>{ + let arguments = (ins TA_AnyTensor:$tensor, I32Attr:$dim); + let results = (outs TensorOf<[AnyType]>); + + let builders = + [OpBuilder<(ins "Value":$tensor, "IntegerAttr":$dim), [{ + auto indices_type = mlir::cast(tensor.getType()).getIndicesType(); + auto outType = RankedTensorType::get({ShapedType::kDynamic}, indices_type); + build($_builder, $_state, outType, tensor, dim); + }]>, + + ]; +} + +def SpTensorGetDimCrd : TA_Op<"SpTensorGetDimCrd", [Pure]>{ + let arguments = (ins TA_AnyTensor:$tensor, I32Attr:$dim); + let results = (outs TensorOf<[AnyType]>); + + let builders = + [OpBuilder<(ins "Value":$tensor, "IntegerAttr":$dim), [{ + auto indices_type = mlir::cast(tensor.getType()).getIndicesType(); + auto outType = RankedTensorType::get({ShapedType::kDynamic}, indices_type); + build($_builder, $_state, outType, tensor, dim); + }]>]; +} + +def SpTensorGetDimBlockPos : TA_Op<"SpTensorGetDimBlockPos", [Pure]>{ + let arguments = (ins TA_AnyTensor:$tensor, I32Attr:$dim); + let results = (outs TensorOf<[AnyType]>); + + let builders = + [OpBuilder<(ins "Value":$tensor, "IntegerAttr":$dim), [{ + auto indices_type = mlir::cast(tensor.getType()).getIndicesType(); + auto outType = RankedTensorType::get({ShapedType::kDynamic}, indices_type); + build($_builder, $_state, outType, tensor, dim); + }]>]; +} + +def SpTensorGetDimBlockCrd : TA_Op<"SpTensorGetDimBlockCrd", [Pure]>{ + let arguments = (ins TA_AnyTensor:$tensor, I32Attr:$dim); + let results = (outs TensorOf<[AnyType]>); + + let builders = + [OpBuilder<(ins "Value":$tensor, "IntegerAttr":$dim), [{ + auto indices_type = mlir::cast(tensor.getType()).getIndicesType(); + auto outType = RankedTensorType::get({ShapedType::kDynamic}, indices_type); + build($_builder, $_state, outType, tensor, dim); + }]>]; +} + +def SpTensorGetVals : TA_Op<"SpTensorGetVals", [Pure]>{ + let arguments = (ins TA_AnyTensor:$tensor); + let results = (outs TensorOf<[AnyType]>); + + + let builders = + [OpBuilder<(ins "Value":$tensor), [{ + auto val_type = mlir::cast(tensor.getType()).getElementType(); + auto outType = RankedTensorType::get({ShapedType::kDynamic}, val_type); + build($_builder, $_state, outType, tensor); + }]>]; +} + +def SpTensorGetCrd : TA_Op<"SpTensorGetCrd", [Pure]>{ + let arguments = (ins TA_AnyTensor:$tensor, Index:$idx, I32Attr:$dim); + let results = (outs AnySignlessIntegerOrIndex:$crd); + + let builders = + [OpBuilder<(ins "Value":$tensor, "Value":$idx, "IntegerAttr":$dim), [{ + ::mlir::IntegerType indices_type; + if(auto spTensor = ::mlir::dyn_cast(tensor.getType())) + { + indices_type = spTensor.getIndicesType(); + } + else if(auto wsTensor = ::mlir::dyn_cast(tensor.getType())) + { + indices_type = wsTensor.getIndicesType(); + } + else + { + assert(false && "Undexpected input type"); + } + build($_builder, $_state, indices_type, tensor, idx, dim); + }]>]; +} + +def SpTensorInsertCrd : TA_Op<"SpTensorInsertCrd", [Pure]>{ + let arguments = (ins TA_AnyTensor:$tensor, I32Attr:$dim, Index:$idx, Index:$crd); + let results = (outs TA_AnyTensor:$result); +} + +def SpTensorGetDimSize : TA_Op<"SpTensorGetDimSize", [Pure]>{ + let arguments = (ins TA_AnyTensor:$tensor, I32Attr:$dim); + let results = (outs Index:$result); + let builders = + [OpBuilder<(ins "Value":$tensor, "IntegerAttr":$dim), [{ + build($_builder, $_state, $_builder.getIndexType(), tensor, dim); + }]>]; +} + +def SpTensorGetNNZ : TA_Op<"SpTensorGetNNZ", [Pure]>{ + let arguments = (ins TA_AnyTensor:$tensor, OptionalAttr:$dim); + let results = (outs Index:$result); +} +def TensorFindPos : TA_Op<"TensorFindPos", [Pure]>{ + let arguments = (ins TA_AnyTensor:$tensor, Optional:$crd, I32Attr:$dim, DefaultValuedAttr:$is_linear); + let results = (outs Index:$result); +} + +def AllocWorkspaceOp : TA_Op<"AllocWorkspace", [Pure, SameVariadicOperandSize]>{ + let summary = "Create a dense workspace to be used as an intermediate output tensor"; + let description = [{}]; + + let arguments = (ins SparseTensor:$tensor, I32ArrayAttr:$dims); + let results = (outs WorkspaceTensor:$result); +} +def WorkspaceClearOp : TA_Op<"WorkspaceClear", [Pure]>{ + let summary = "Clear the workspace for use in the next iteration"; + let description = [{}]; + let arguments = (ins WorkspaceTensor:$tensor); + let results = (outs WorkspaceTensor:$result); +} + +def SortCrdOp : TA_Op<"SortCrdOp", [Pure]> { + let summary = "Sort the coordinates of a tensor"; + let description = [{}]; + + let arguments = (ins TA_AnyTensor:$tensor); + let results = (outs TA_AnyTensor:$result); +} + +def TensorSortOp : TA_Op<"TensorSortOp", [Pure, DestinationStyleOpInterface]> { + let summary = "Sort a tensor"; + let description = [{}]; + + let arguments = (ins TensorRankOf<[AnyInteger], [1]>:$tensor, Index:$start, Index:$end); + let results = (outs TensorRankOf<[AnyInteger], [1]>:$result); + + let extraClassDeclaration = [{ + MutableOperandRange getDpsInitsMutable() { return getTensorMutable(); } + }]; +} + +def WorkspaceAccumulateOp : TA_Op<"WorkspaceAccumulateOp", [Pure]> { + let summary = "Merged extract, accumulate, insert operation for a workspace"; + let description = [{}]; + + let arguments = ( + ins WorkspaceTensor:$tensor, + Index:$pos, + Variadic:$crds, + AnyTypeOf<[F32,F64]>:$value + ); + let results = (outs WorkspaceTensor:$result); +} + +def WorkspaceReadOp : TA_Op<"WorkspaceReadOp", [Pure]> { + let summary = "Read from a workspace without checking the mark array"; + let description = [{}]; + + let arguments = (ins WorkspaceTensor:$tensor, Index:$pos, Index:$crd); + let results = (outs AnyFloat:$result); +} +#endif /// TA_OPS \ No newline at end of file diff --git a/include/comet/Dialect/TensorAlgebra/IR/TATypes.td b/include/comet/Dialect/TensorAlgebra/IR/TATypes.td new file mode 100644 index 00000000..00bdee9d --- /dev/null +++ b/include/comet/Dialect/TensorAlgebra/IR/TATypes.td @@ -0,0 +1,78 @@ +// +// Copyright 2022 Battelle Memorial Institute +// +// Redistribution and use in source and binary forms, with or without modification, +// are permitted provided that the following conditions are met: +// +// 1. Redistributions of source code must retain the above copyright notice, this list of conditions +// and the following disclaimer. +// +// 2. Redistributions in binary form must reproduce the above copyright notice, this list of conditions +// and the following disclaimer in the documentation and/or other materials provided with the distribution. +// +// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED +// WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR +// PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, +// INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE +// GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, +// WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +// OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. +// + +#ifndef TA_TYPES +#define TA_TYPES + +include "mlir/IR/AttrTypeBase.td" +include "mlir/IR/BuiltinTypeInterfaces.td" +include "comet/Dialect/TensorAlgebra/IR/TABase.td" + +class TensorAlgebra_Type traits = []> + : TypeDef { + let mnemonic = typeMnemonic; +} + +def IndexLabel : TensorAlgebra_Type<"IndexLabel","index"> { + let summary = "Type representing the index of one dimnsion of a tensor"; + let description = [{}]; +} + +def Range : TensorAlgebra_Type<"Range","range"> { + let summary = "Type representing the range of one dimnsion of a tensor"; + let description = [{}]; +} + +def SparseTensor : TensorAlgebra_Type<"SparseTensor", "sparse_tensor", [DeclareTypeInterfaceMethods]> { + let summary = "Sparse tensor to use in tensor algebra dialect"; + let description = [{}]; + + let parameters = (ins + "::mlir::Type":$element_type, + "::mlir::IntegerType":$indices_type, + ArrayRefParameter<"int64_t", "Dimensions of tensor">:$dims, + ArrayRefParameter<"TensorFormatEnum", "Format">:$format + ); + + // let assemblyFormat = "`<` $element_type `,` $indices_type `,` $dims `,` $format `>`"; + let hasCustomAssemblyFormat = 1; + // TODO: Implement custom builder from "common" format strings into format strings +} + +def WorkspaceTensor : TensorAlgebra_Type<"Workspace", "workspace", [DeclareTypeInterfaceMethods]> { + let summary = "Temporary tensor generated from a workspace transfrom"; + let description = [{ + Dense, temporary tensor generated from a workspace transformation. + Needed to represent dense row as well as mark array. + }]; + let parameters = ( + ins "::mlir::Type":$element_type, + "::mlir::IntegerType":$indices_type, + ArrayRefParameter<"int64_t", "Dimensions of workspace">:$dims + ); + // let assemblyFormat = "`<` $element_type `,` $indices_type `,` $dims `>`"; + let hasCustomAssemblyFormat = 1; +} + +/// Provide a definition of the types that are used within the TA dialect. +def TA_AnyTensor : AnyTypeOf<[TensorOf<[AnyType]>, SparseTensor, WorkspaceTensor]>; +def TA_FloatTensor: TensorOf<[F32, F64]>; +#endif //TA_TYPES \ No newline at end of file diff --git a/include/comet/Dialect/TensorAlgebra/Passes.h b/include/comet/Dialect/TensorAlgebra/Passes.h index b7c3bcf6..f21562e2 100644 --- a/include/comet/Dialect/TensorAlgebra/Passes.h +++ b/include/comet/Dialect/TensorAlgebra/Passes.h @@ -61,9 +61,17 @@ namespace mlir std::unique_ptr createLoweringTTGTPass(bool enableBestPerm, int whatPermID = 1, bool printFlops = false); + + /// Create a pass for lowering TA operations to TTGT + /// This pass selects either the best permutation among all + /// or pass can specify the iteration order of the permutation, ith permutation + std::unique_ptr createLoweringTTGTDynPass(int whatPermID = -1, + bool printFlops = false); std::unique_ptr createLinAlgMatmulTilingPass(); std::unique_ptr createLinAlgMatmulMicroKernelPass(); + std::unique_ptr createMatvecToParallelLoopsPass(); + // Optimize dense transpose (linalg.copy) based on the following paper: // HPTT: A High-Performance Tensor Transposition C++ Library @@ -76,6 +84,9 @@ namespace mlir /// Create a pass for lowering tensor fill operation to linalg.fill std::unique_ptr createTensorFillLoweringPass(); + // Create a pass for merging and optimizing operations on workspaces. + std::unique_ptr createWorkspaceOptimizationsPass(); + /// Create a pass for lowering to the rest of the operations in `Std` dialects, /// such as printOp, constantOp, ReturnOp.. std::unique_ptr createLateLoweringPass(); @@ -95,6 +106,7 @@ namespace mlir std::unique_ptr createFuncOpLoweringPass(); // Conversion std::unique_ptr createDimOpLoweringPass(); + std::unique_ptr createTABufferizeFunc(); } } diff --git a/include/comet/Dialect/TensorAlgebra/Passes.td b/include/comet/Dialect/TensorAlgebra/Passes.td index d156be70..8a6feb55 100644 --- a/include/comet/Dialect/TensorAlgebra/Passes.td +++ b/include/comet/Dialect/TensorAlgebra/Passes.td @@ -85,5 +85,24 @@ def TensorAlgebraSparseOutputTensorDeclLowering : Pass<"lower-sparse-output-tens "comet::TensorAlgebraDialect" ]; } +def TensorAlgebraBufferizeFunc : Pass<"ta-bufferize-func"> { + let summary = ""; + let description = [{ + + }]; + let constructor = "comet::createTABufferizeFunc()"; + let dependentDialects = [ + "comet::TensorAlgebraDialect", "func::FuncDialect" + ]; +} + +def TensorAlgebraWorkspaceOptimizations : Pass<"ta-workspace-optimizations", "func::FuncOp"> { + let summary = "Substitute merged workspace operations where possible"; + let description = [{}]; + + let constructor = "comet::createWorkspaceOptimizationPass()"; + let dependentDialects = ["tensorAlgebra::TADialect"]; +} + #endif /// COMET_DIALECT_TENSORALGEBRA_PASSES diff --git a/include/comet/Dialect/TensorAlgebra/Patterns.h b/include/comet/Dialect/TensorAlgebra/Patterns.h new file mode 100644 index 00000000..9cd765e8 --- /dev/null +++ b/include/comet/Dialect/TensorAlgebra/Patterns.h @@ -0,0 +1,37 @@ +//===- Patterns.h - Conversion Pass Construction and Registration -----------===// +// +// Copyright 2022 Battelle Memorial Institute +// +// Redistribution and use in source and binary forms, with or without modification, +// are permitted provided that the following conditions are met: +// +// 1. Redistributions of source code must retain the above copyright notice, this list of conditions +// and the following disclaimer. +// +// 2. Redistributions in binary form must reproduce the above copyright notice, this list of conditions +// and the following disclaimer in the documentation and/or other materials provided with the distribution. +// +// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED +// WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR +// PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, +// INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE +// GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, +// WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +// OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. +// +//===----------------------------------------------------------------------===// + +#ifndef COMET_DIALECT_TENSORALGEBRA_PATTERNS_H +#define COMET_DIALECT_TENSORALGEBRA_PATTERNS_H + +#include "mlir/Transforms/DialectConversion.h" + +namespace mlir +{ + namespace tensorAlgebra + { + void populateWorkspaceOptimizationPatterns(MLIRContext *context, RewritePatternSet &patterns); + } +} + +#endif // COMET_DIALECT_INDEXTREE_PASSES_H diff --git a/include/comet/Dialect/Utils/Utils.h b/include/comet/Dialect/Utils/Utils.h index 669d1b97..4f841241 100644 --- a/include/comet/Dialect/Utils/Utils.h +++ b/include/comet/Dialect/Utils/Utils.h @@ -24,22 +24,30 @@ #ifndef TENSORALGEBRA_UTILS_H_ #define TENSORALGEBRA_UTILS_H_ +#include "comet/Dialect/TensorAlgebra/IR/TADialect.h" +#include "mlir/Dialect/MemRef/IR/MemRef.h" +#include "mlir/Dialect/Tensor/IR/Tensor.h" +#include "mlir/IR/Builders.h" +#include "mlir/IR/BuiltinTypes.h" +#include "mlir/IR/Location.h" #include "mlir/IR/PatternMatch.h" #include "mlir/Dialect/SCF/IR/SCF.h" #include "mlir/Dialect/Linalg/IR/Linalg.h" #include "mlir/Dialect/Arith/IR/Arith.h" #include "mlir/Dialect/Affine/LoopUtils.h" +#include "llvm/ADT/APFloat.h" #include "llvm/ADT/MapVector.h" #include "llvm/ADT/SetVector.h" +#include "mlir/IR/Value.h" #include "mlir/Transforms/DialectConversion.h" #include #include #include +#include /// TODO(gkestor): supports only f64 - need generalization -extern std::string VALUETYPE; using namespace mlir::linalg; @@ -47,7 +55,8 @@ namespace mlir { namespace tensorAlgebra { - enum TargetDevice { CPU, GPU}; + enum TargetDevice { CPU, GPU, FPGA}; + enum GPUCompilationFormat { Binary, Assembly, Fatbin}; using IndexSizeMap = std::unordered_map; using IndexVector = std::vector; @@ -109,7 +118,7 @@ namespace mlir void print_vector_value(std::vector vec); std::string dump2str(Value t); - std::vector stringSplit(std::string s, std::string delimiter); + std::vector stringSplit(llvm::StringRef s, llvm::StringRef delimiter); std::vector getReverseIdentityPermutation(size_t size); std::vector getIdentityPermutation(size_t size); @@ -125,7 +134,7 @@ namespace mlir std::vector getFreeIndices(std::vector rhs_perm, std::vector lhs_perm); std::vector getSumIndices(std::vector rhs_perm, std::vector rhs_perm_free); std::vector getIndexIterateOrder(std::vector rhs1_perm, std::vector rhs2_perm); - std::vector> getAllFormats(ArrayAttr opFormatsArrayAttr, std::vector> allPerms); + std::vector> getAllFormats(ArrayAttr opFormatsArrayAttr, std::vector> allPerms); bool checkIsElementwise(std::vector> allPerms); bool checkIsMixedMode(std::vector> formats); bool checkIsDense(std::vector format); @@ -133,10 +142,13 @@ namespace mlir bool isDense(std::string s, std::string delim); bool isMergedIndex(std::vector format_vec, int cur_idx, int sumIndex); - std::vector getFormatsValue(std::string formats_str, int rank_size, + std::vector getFormatsValue(llvm::StringRef formats_str, int rank_size, PatternRewriter &rewriter, Location loc, IndexType indexType); - std::vector getFormatsValueInt(std::string formats_str, int rank_size, + std::string getTensorFormatString(Type tensorT); + + std::vector getFormatsValueInt(llvm::StringRef formats_str, int rank_size, PatternRewriter &rewriter, Location loc, IntegerType intType); + std::vector getFormats(llvm::StringRef formats_str, int rank_size, MLIRContext* ctx); double loopCostHeuristic(const std::vector &loopOrder, size_t dim_, std::vector &sourceOrder, std::vector &destOrder); @@ -221,6 +233,8 @@ namespace mlir std::vector &formats); void replaceOperands(Operation *itComputeOp, std::vector newComputeOps); + TypedValue collapseMemref(TypedValue val, mlir::OpBuilder& builder); + mlir::Value get_memref_num_elements(mlir::MLIRContext* ctx, mlir::OpBuilder& builder, mlir::Location loc, mlir::Value memref); // For TTGT transformations struct ContractionPlan @@ -569,6 +583,368 @@ namespace mlir std::string bestPermStr_; }; /// struct ContractionPlan + struct ContractionPlanDyn + { + ContractionPlanDyn(OpBuilder& builder, Location loc, IndexVector a_perm, Value a_shape, IndexVector b_perm, + Value b_shape, IndexVector c_perm, Value c_shape) + : a_perm_{a_perm}, b_perm_{b_perm}, c_perm_{c_perm} + { + + /// get M-N-K indices for gemm + std::tie(m_indices_, n_indices_, k_indices_) = + getIndices(a_perm_, b_perm_, c_perm_); + + /// compute size map for each index + for (size_t i = 0; i < a_perm_.size(); i++) + { + auto dim = builder.create(loc, a_shape, i); + size_map_.insert({a_perm_[i], dim}); + } + for (size_t i = 0; i < b_perm_.size(); i++) + { + auto dim = builder.create(loc, b_shape, i); + size_map_.insert({b_perm_[i], dim}); + } + for (size_t i = 0; i < c_perm_.size(); i++) + { + auto dim = builder.create(loc, c_shape, i); + size_map_.insert({c_perm_[i],dim}); + } + + /// compute sizes for M-N-K sizes + m_size_ = builder.create(loc, 1); + for (const auto &idx : m_indices_) + { + auto cur_size = size_map_[idx]; + m_size_ = builder.create(loc, m_size_, cur_size); + } + n_size_ = builder.create(loc, 1); + for (const auto &idx : n_indices_) + { + auto cur_size = size_map_[idx]; + n_size_ = builder.create(loc, n_size_, cur_size); + } + k_size_ = builder.create(loc, 1); + for (const auto &idx : k_indices_) + { + auto cur_size = size_map_[idx]; + k_size_ = builder.create(loc, k_size_, cur_size); + } + } + + std::tuple + getIndices(IndexVector A_perm, IndexVector B_perm, IndexVector C_perm) const + { + IndexVector mIndices, nIndices, kIndices; + + std::set A_perm_set(A_perm.begin(), A_perm.end()), + B_perm_set(B_perm.begin(), B_perm.end()), + C_perm_set(C_perm.begin(), C_perm.end()); + + std::set *A_perm_ptr{&A_perm_set}, *B_perm_ptr{&B_perm_set}; + + std::set A_int_C, B_int_C, A_un_B; + + std::set_intersection(A_perm_ptr->begin(), A_perm_ptr->end(), + C_perm_set.begin(), C_perm_set.end(), + std::inserter(A_int_C, A_int_C.begin())); + std::set_difference(A_int_C.begin(), A_int_C.end(), B_perm_ptr->begin(), + B_perm_ptr->end(), std::back_inserter(mIndices)); + + std::set_intersection(B_perm_ptr->begin(), B_perm_ptr->end(), + C_perm_set.begin(), C_perm_set.end(), + std::inserter(B_int_C, B_int_C.begin())); + std::set_difference(B_int_C.begin(), B_int_C.end(), A_perm_ptr->begin(), + A_perm_ptr->end(), std::back_inserter(nIndices)); + + std::set_union(A_perm_ptr->begin(), A_perm_ptr->end(), B_perm_ptr->begin(), + B_perm_ptr->end(), std::inserter(A_un_B, A_un_B.begin())); + std::set_difference(A_un_B.begin(), A_un_B.end(), C_perm_set.begin(), + C_perm_set.end(), std::back_inserter(kIndices)); + + return std::make_tuple(mIndices, nIndices, kIndices); + } + + Value getTransposeTime(OpBuilder& builder, Location loc, Value mem_size, const IndexVector &perm) const + { + Value memSizeI = builder.create(loc, IntegerType::get(builder.getContext(), 64), mem_size); // ensure we have an integer for division + Value memSizeF = builder.create(loc, Float64Type::get(builder.getContext()), memSizeI); // convert to float for division + Value result; + if (perm[0] != 0) + { + // TODO(gkestor): needs to be adjusted according to our transpose method + result = builder.create(loc, + memSizeF, + builder.create(loc, llvm::APFloat(0.71), Float64Type::get(builder.getContext()))); + // result = mem_size / 0.71; + } + else + { + result = memSizeF; + } + return result; + } + + // double flopCount() const + // { + // double overall_size = m_size_ * n_size_ * k_size_; + // int op_factor = n_indices_.size() == 0 ? 1 : 2; + + // return overall_size * op_factor; + // } + + std::string contractionString(const IndexVector &a_idx, + const IndexVector &b_idx, + const IndexVector &c_idx) const + { + std::string result = "contr_C"; + + for (const auto &idx : c_idx) + { + result += "_" + std::to_string(idx); + } + result += "_A"; + + for (const auto &idx : a_idx) + { + result += "_" + std::to_string(idx); + } + result += "_B"; + + for (const auto &idx : b_idx) + { + result += "_" + std::to_string(idx); + } + + return result; + } + + IndexVector getPermutation(const IndexVector &in_idx, + const IndexVector &out_idx) const + { + IndexVector result; + for (const auto &idx : out_idx) + { + auto it = std::find(in_idx.begin(), in_idx.end(), idx); + assert(it != in_idx.end() && "Wrong permutation"); + result.push_back(std::distance(in_idx.begin(), it)); + } + return result; + } + + // double getTotalTime() + // { + // IndexVector a_perm, b_perm, c_perm; + // double minTime; + // std::tie(a_perm, b_perm, c_perm, minTime) = computeBestPermutations(); + // double result = flopCount() + minTime; + + // return result; + // } + + void computePermutations(OpBuilder& builder, Location loc) + { + IndexVector a_perm, b_perm, c_perm; + computeAllPermutations(builder, loc); + } + + std::tuple findPermutationsAtN(int whichperm) + { + int curper = 1; + IndexVector m_idx{m_indices_}, n_idx{n_indices_}, k_idx{k_indices_}; + std::sort(m_idx.begin(), m_idx.end()); + std::sort(n_idx.begin(), n_idx.end()); + std::sort(k_idx.begin(), k_idx.end()); + + IndexVector a_candidate, b_candidate, c_candidate; + for (size_t i = 0; i < 2; i++) + { + bool doSwap = (i != 0); + do + { + do + { + do + { + if (curper == whichperm) + { + IndexVector a_idx, b_idx, c_idx; + + if (i == 1) + { + a_idx.insert(a_idx.end(), k_idx.begin(), k_idx.end()); + a_idx.insert(a_idx.end(), m_idx.begin(), m_idx.end()); + b_idx.insert(b_idx.end(), n_idx.begin(), n_idx.end()); + b_idx.insert(b_idx.end(), k_idx.begin(), k_idx.end()); + c_idx.insert(c_idx.end(), n_idx.begin(), n_idx.end()); + c_idx.insert(c_idx.end(), m_idx.begin(), m_idx.end()); + } + else + { + a_idx.insert(a_idx.end(), m_idx.begin(), m_idx.end()); + a_idx.insert(a_idx.end(), k_idx.begin(), k_idx.end()); + b_idx.insert(b_idx.end(), k_idx.begin(), k_idx.end()); + b_idx.insert(b_idx.end(), n_idx.begin(), n_idx.end()); + c_idx.insert(c_idx.end(), m_idx.begin(), m_idx.end()); + c_idx.insert(c_idx.end(), n_idx.begin(), n_idx.end()); + } + + a_candidate = a_idx; + b_candidate = b_idx; + c_candidate = c_idx; + swapAB_ = doSwap; + } + curper++; + + } while (std::next_permutation(k_idx.begin(), k_idx.end())); + } while (std::next_permutation(n_idx.begin(), n_idx.end())); + } while (std::next_permutation(m_idx.begin(), m_idx.end())); + } + + assert(whichperm <= curper && "Cannot find the selected permutation"); + IndexVector best_a_perm, best_b_perm, best_c_perm; + bestPermStr_ = contractionString(a_candidate, b_candidate, c_candidate); + best_a_perm = getPermutation(a_perm_, a_candidate); + best_b_perm = getPermutation(b_perm_, b_candidate); + best_c_perm = getPermutation(c_perm_, c_candidate); + + return std::make_tuple(best_a_perm, best_b_perm, best_c_perm); + } + + void computeAllPermutations(OpBuilder &builder, Location loc) + { + IndexVector m_idx{m_indices_}, n_idx{n_indices_}, k_idx{k_indices_}; + std::sort(m_idx.begin(), m_idx.end()); + std::sort(n_idx.begin(), n_idx.end()); + std::sort(k_idx.begin(), k_idx.end()); + + IndexVector a_candidate, b_candidate, c_candidate; + + do + { + do + { + do + { + for (size_t i = 0; i < 2; i++) + { + IndexVector a_idx, b_idx, c_idx; + Value a_size, b_size, c_size; + Value transposeTime = builder.create(loc, llvm::APFloat(0.0), Float64Type::get(builder.getContext())); // initialize to 0.0 + + if (i == 1) + { + a_idx.insert(a_idx.end(), k_idx.begin(), k_idx.end()); + a_idx.insert(a_idx.end(), m_idx.begin(), m_idx.end()); + b_idx.insert(b_idx.end(), n_idx.begin(), n_idx.end()); + b_idx.insert(b_idx.end(), k_idx.begin(), k_idx.end()); + c_idx.insert(c_idx.end(), n_idx.begin(), n_idx.end()); + c_idx.insert(c_idx.end(), m_idx.begin(), m_idx.end()); + } + else + { + a_idx.insert(a_idx.end(), m_idx.begin(), m_idx.end()); + a_idx.insert(a_idx.end(), k_idx.begin(), k_idx.end()); + b_idx.insert(b_idx.end(), k_idx.begin(), k_idx.end()); + b_idx.insert(b_idx.end(), n_idx.begin(), n_idx.end()); + c_idx.insert(c_idx.end(), m_idx.begin(), m_idx.end()); + c_idx.insert(c_idx.end(), n_idx.begin(), n_idx.end()); + } + + // Compute the sizes dynamically + a_size = builder.create(loc, k_size_, m_size_); //k_size_ * m_size_; + b_size = builder.create(loc, n_size_, k_size_); //n_size_ * k_size_; + c_size = builder.create(loc, m_size_, n_size_); //m_size_ * n_size_; + + if (a_perm_ != a_idx) + { + m_transposeA.push_back(true); // mark that A needs transpose for this permutation + transposeTime = builder.create( + loc, + transposeTime, + getTransposeTime(builder, loc, a_size, + getPermutation(a_perm_, a_idx))); + } + else { + m_transposeA.push_back(false); // mark that A needs transpose for this permutation + + } + + if (b_perm_ != b_idx) + { + m_transposeB.push_back(true); // mark that A needs transpose for this permutation + + transposeTime = builder.create( + loc, + transposeTime, + getTransposeTime(builder, loc, b_size, + getPermutation(b_perm_, b_idx))); + } + else { + m_transposeB.push_back(false); // mark that B does not need transpose for this permutation + } + + if (c_perm_ != c_idx) + { + m_transposeC.push_back(true); // mark that C needs transpose for this permutation + transposeTime = builder.create( + loc, + transposeTime, + getTransposeTime(builder, loc, c_size, + getPermutation(c_perm_, c_idx))); + // transposeTime += + // getTransposeTime(c_size, getPermutation(c_perm_, c_idx)); + } + else { + m_transposeC.push_back(false); // mark that C does not need transpose for this permutation + } + + m_contraction_time.push_back(transposeTime); + m_contraction_permutations.push_back({getPermutation(a_perm_, a_idx), getPermutation(b_perm_, b_idx), getPermutation(c_perm_, c_idx)}); // store the permutations for this contraction + // Store the swapAB flag for this contraction ? + m_swapAB.push_back(i == 1); + } + + } while (std::next_permutation(k_idx.begin(), k_idx.end())); + + } while (std::next_permutation(n_idx.begin(), n_idx.end())); + + } while (std::next_permutation(m_idx.begin(), m_idx.end())); + + IndexVector best_a_perm, best_b_perm, best_c_perm; + + // bestPermStr_ = contractionString(a_candidate, b_candidate, c_candidate); + + } + + IndexVector a_perm_; + IndexVector b_perm_; + IndexVector c_perm_; + + IndexVector m_indices_; + IndexVector n_indices_; + IndexVector k_indices_; + + Value m_size_; + Value n_size_; + Value k_size_; + + std::unordered_map size_map_; + std::vector m_transposeA; + std::vector m_transposeB; + std::vector m_transposeC; + std::vector m_swapAB; + std::vector m_contraction_time; + std::vector> m_contraction_permutations; // to store the permutations for each contraction + + + + bool swapAB_; + bool inA_; + + std::string bestPermStr_; + }; /// struct ContractionPlanDyn + } /// namespace tensorAlgebra } /// namespace mlir diff --git a/include/comet/Utils/debug.h b/include/comet/Utils/debug.h index bef84e80..1630c723 100644 --- a/include/comet/Utils/debug.h +++ b/include/comet/Utils/debug.h @@ -1,3 +1,24 @@ +// +// Copyright 2022 Battelle Memorial Institute +// +// Redistribution and use in source and binary forms, with or without modification, +// are permitted provided that the following conditions are met: +// +// 1. Redistributions of source code must retain the above copyright notice, this list of conditions +// and the following disclaimer. +// +// 2. Redistributions in binary form must reproduce the above copyright notice, this list of conditions +// and the following disclaimer in the documentation and/or other materials provided with the distribution. +// +// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED +// WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR +// PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, +// INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE +// GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, +// WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +// OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. +// + #undef comet_debug #undef comet_pdump #undef comet_vdump diff --git a/integration_test/compound_exps/CSR_mult_spTranspose_CSR.ta b/integration_test/compound_exps/CSR_mult_spTranspose_CSR.ta deleted file mode 100644 index 1e5a6e88..00000000 --- a/integration_test/compound_exps/CSR_mult_spTranspose_CSR.ta +++ /dev/null @@ -1,35 +0,0 @@ -# RUN: export SPARSE_FILE_NAME0=%comet_integration_test_data_dir/test_rank2.mtx -# RUN: export SORT_TYPE=SEQ_QSORT -# RUN: comet-opt --opt-comp-workspace --convert-ta-to-it --convert-to-loops --convert-to-llvm %s &> CSR_mult_spTranspose_CSR.llvm -# RUN: mlir-cpu-runner CSR_mult_spTranspose_CSR.llvm -O3 -e main -entry-point-result=void -shared-libs=%comet_utility_library_dir/libcomet_runner_utils%shlibext | FileCheck %s - -def main() { - #IndexLabel Declarations - IndexLabel [i] = [?]; - IndexLabel [j] = [?]; - IndexLabel [k] = [?]; - - #Tensor Declarations - Tensor A([k, j], CSR); - Tensor B([i, j], CSR); - Tensor C([k, i], CSR); - - #Tensor Readfile Operation - A[k, j] = comet_read(0); - B[i, j] = comet_read(0); - - #Tensor Transpose - C[k, i] = A[k, j] * transpose(B[i, j],{j,i}); - print(C); -} -# Print the result for verification. -# CHECK: data = -# CHECK-NEXT: 5, -# CHECK-NEXT: data = -# CHECK-NEXT: 0, -# CHECK-NEXT: data = -# CHECK-NEXT: 0,2,4,5,7,9, -# CHECK-NEXT: data = -# CHECK-NEXT: 0,3,1,4,2,0,3,1,4, -# CHECK-NEXT: data = -# CHECK-NEXT: 2.96,9.7,10.25,22.9,9,9.7,32.81,22.9,52.04, \ No newline at end of file diff --git a/integration_test/compound_exps/Dense_eltwise_sTranspose_CSR.ta b/integration_test/compound_exps/Dense_eltwise_sTranspose_CSR.ta deleted file mode 100644 index 1627ed9c..00000000 --- a/integration_test/compound_exps/Dense_eltwise_sTranspose_CSR.ta +++ /dev/null @@ -1,36 +0,0 @@ -# RUN: export SPARSE_FILE_NAME0=%comet_integration_test_data_dir/test_rank2.mtx -# RUN: export SORT_TYPE=SEQ_QSORT -# RUN: comet-opt --convert-ta-to-it --convert-to-loops --convert-to-llvm %s &> Dense_eltwise_sTranspose_CSR.llvm -# RUN: mlir-cpu-runner Dense_eltwise_sTranspose_CSR.llvm -O3 -e main -entry-point-result=void -shared-libs=%comet_utility_library_dir/libcomet_runner_utils%shlibext | FileCheck %s - - -def main() { - #IndexLabel Declarations - IndexLabel [i] = [?]; - IndexLabel [j] = [?]; - - #Tensor Declarations - Tensor B([i, j], CSR); - Tensor C([j, i], CSR); - Tensor A([j, i], Dense); - - #Tensor Readfile Operation - B[i, j] = comet_read(0); - A[j, i] = 3.2; - - #Tensor Transpose - C[j, i] = A[j, i] .* transpose(B[i, j],{j,i}); - print(C); -} - -# Print the result for verification. -# CHECK: data = -# CHECK-NEXT: 5, -# CHECK-NEXT: data = -# CHECK-NEXT: 0, -# CHECK-NEXT: data = -# CHECK-NEXT: 0,2,4,5,7,9, -# CHECK-NEXT: data = -# CHECK-NEXT: 0,3,1,4,2,0,3,1,4, -# CHECK-NEXT: data = -# CHECK-NEXT: 3.2,13.12,6.4,16.64,9.6,4.48,12.8,8,16, \ No newline at end of file diff --git a/integration_test/compound_exps/Dense_mult_spTranspose_CSR.ta b/integration_test/compound_exps/Dense_mult_spTranspose_CSR.ta deleted file mode 100644 index 0aba55bb..00000000 --- a/integration_test/compound_exps/Dense_mult_spTranspose_CSR.ta +++ /dev/null @@ -1,28 +0,0 @@ -# RUN: export SPARSE_FILE_NAME0=%comet_integration_test_data_dir/test_rank2.mtx -# RUN: export SORT_TYPE=SEQ_QSORT -# RUN: comet-opt --convert-ta-to-it --convert-to-loops --convert-to-llvm %s &> Dense_mult_spTranspose_CSR.llvm -# RUN: mlir-cpu-runner Dense_mult_spTranspose_CSR.llvm -O3 -e main -entry-point-result=void -shared-libs=%comet_utility_library_dir/libcomet_runner_utils%shlibext | FileCheck %s - -def main() { - #IndexLabel Declarations - IndexLabel [i] = [?]; - IndexLabel [j] = [?]; - IndexLabel [k] = [5]; - - #Tensor Declarations - Tensor B([i, j], CSR); - Tensor A([k, j], Dense); - Tensor C([k, i], Dense); - - #Tensor Readfile Operation - B[i, j] = comet_read(0); - A[k, j] = 2.3; - C[k, i] = 0.0; - - #Tensor Transpose - C[k, i] = A[k, j] * transpose(B[i, j],{j,i}); - print(C); -} -# Print the result for verification. -# CHECK: data = -# CHECK-NEXT: 5.52,10.35,6.9,18.63,23.46,5.52,10.35,6.9,18.63,23.46,5.52,10.35,6.9,18.63,23.46,5.52,10.35,6.9,18.63,23.46,5.52,10.35,6.9,18.63,23.46, \ No newline at end of file diff --git a/integration_test/compound_exps/spTranspose_CSR_eltwise_CSR.ta b/integration_test/compound_exps/spTranspose_CSR_eltwise_CSR.ta deleted file mode 100644 index 17e894d0..00000000 --- a/integration_test/compound_exps/spTranspose_CSR_eltwise_CSR.ta +++ /dev/null @@ -1,35 +0,0 @@ -# RUN: export SPARSE_FILE_NAME0=%comet_integration_test_data_dir/test_rank2.mtx -# RUN: export SORT_TYPE=SEQ_QSORT -# RUN: comet-opt --opt-comp-workspace --convert-ta-to-it --convert-to-loops --convert-to-llvm %s &> spTranspose_CSR_eltwise_CSR.llvm -# RUN: mlir-cpu-runner spTranspose_CSR_eltwise_CSR.llvm -O3 -e main -entry-point-result=void -shared-libs=%comet_utility_library_dir/libcomet_runner_utils%shlibext | FileCheck %s - -def main() { - #IndexLabel Declarations - IndexLabel [i] = [?]; - IndexLabel [j] = [?]; - - #Tensor Declarations - Tensor A([i, j], CSR); - Tensor B([j, i], CSR); - Tensor C([j, i], CSR); - - #Tensor Readfile Operation - A[i, j] = comet_read(0); - B[j, i] = comet_read(0); - - #Tensor Transpose - C[j, i] = transpose(A[i, j],{j,i}) .* B[j, i]; - print(C); -} - -# Print the result for verification. -# CHECK: data = -# CHECK-NEXT: 5, -# CHECK-NEXT: data = -# CHECK-NEXT: 0, -# CHECK-NEXT: data = -# CHECK-NEXT: 0,2,4,5,7,9, -# CHECK-NEXT: data = -# CHECK-NEXT: 0,3,1,4,2,0,3,1,4, -# CHECK-NEXT: data = -# CHECK-NEXT: 1,5.74,4,13,9,5.74,16,13,25, \ No newline at end of file diff --git a/integration_test/compound_exps/spTranspose_CSR_eltwise_Dense.ta b/integration_test/compound_exps/spTranspose_CSR_eltwise_Dense.ta deleted file mode 100644 index 243d6eb3..00000000 --- a/integration_test/compound_exps/spTranspose_CSR_eltwise_Dense.ta +++ /dev/null @@ -1,35 +0,0 @@ -# RUN: export SPARSE_FILE_NAME0=%comet_integration_test_data_dir/test_rank2.mtx -# RUN: export SORT_TYPE=SEQ_QSORT -# RUN: comet-opt --convert-ta-to-it --convert-to-loops --convert-to-llvm %s &> spTranspose_CSR_eltwise_Dense.llvm -# RUN: mlir-cpu-runner spTranspose_CSR_eltwise_Dense.llvm -O3 -e main -entry-point-result=void -shared-libs=%comet_utility_library_dir/libcomet_runner_utils%shlibext | FileCheck %s - -def main() { - #IndexLabel Declarations - IndexLabel [i] = [?]; - IndexLabel [j] = [?]; - - #Tensor Declarations - Tensor A([i, j], CSR); - Tensor C([j, i], CSR); - Tensor B([j, i], Dense); - - #Tensor Readfile Operation - A[i, j] = comet_read(0); - B[j, i] = 2.3; - - #Tensor Transpose - C[j, i] = transpose(A[i, j],{j,i}) .* B[j, i]; - print(C); -} - -# Print the result for verification. -# CHECK: data = -# CHECK-NEXT: 5, -# CHECK-NEXT: data = -# CHECK-NEXT: 0, -# CHECK-NEXT: data = -# CHECK-NEXT: 0,2,4,5,7,9, -# CHECK-NEXT: data = -# CHECK-NEXT: 0,3,1,4,2,0,3,1,4, -# CHECK-NEXT: data = -# CHECK-NEXT: 2.3,9.43,4.6,11.96,6.9,3.22,9.2,5.75,11.5, \ No newline at end of file diff --git a/integration_test/compound_exps/spTranspose_CSR_mult_CSR.ta b/integration_test/compound_exps/spTranspose_CSR_mult_CSR.ta deleted file mode 100644 index c627df46..00000000 --- a/integration_test/compound_exps/spTranspose_CSR_mult_CSR.ta +++ /dev/null @@ -1,38 +0,0 @@ -# RUN: export SPARSE_FILE_NAME0=%comet_integration_test_data_dir/test_rank2.mtx -# RUN: export SORT_TYPE=SEQ_QSORT -# RUN: comet-opt --opt-comp-workspace --convert-ta-to-it --convert-to-loops --convert-to-llvm %s &> spTranspose_CSR_mult_CSR.llvm -# RUN: mlir-cpu-runner spTranspose_CSR_mult_CSR.llvm -O3 -e main -entry-point-result=void -shared-libs=%comet_utility_library_dir/libcomet_runner_utils%shlibext | FileCheck %s - -## GOOD TEST (merge request) - -def main() { - #IndexLabel Declarations - IndexLabel [i] = [?]; - IndexLabel [j] = [?]; - IndexLabel [k] = [?]; - - #Tensor Declarations - Tensor A([i, j], CSR); - Tensor B([i, k], CSR); - Tensor C([j, k], CSR); - - #Tensor Readfile Operation - A[i, j] = comet_read(0); - B[i, k] = comet_read(0); - - #Tensor Transpose - C[j, k] = transpose(A[i, j],{j,i}) * B[i, k]; - print(C); -} - -# Print the result for verification. -# CHECK: data = -# CHECK-NEXT: 5, -# CHECK-NEXT: data = -# CHECK-NEXT: 0, -# CHECK-NEXT: data = -# CHECK-NEXT: 0,2,4,5,7,9, -# CHECK-NEXT: data = -# CHECK-NEXT: 0,3,1,4,2,0,3,1,4, -# CHECK-NEXT: data = -# CHECK-NEXT: 17.81,17.8,31.04,31,9,17.8,17.96,31,31.25, \ No newline at end of file diff --git a/integration_test/ops/scalar.ta b/integration_test/ops/scalar.ta deleted file mode 100644 index 7452037b..00000000 --- a/integration_test/ops/scalar.ta +++ /dev/null @@ -1,15 +0,0 @@ -# RUN: comet-opt --convert-to-loops --convert-to-llvm %s &> scalars.llvm -# RUN: mlir-cpu-runner scalars.llvm -O3 -e main -entry-point-result=void -shared-libs=%comet_utility_library_dir/libcomet_runner_utils%shlibext | FileCheck %s - -def main() { - var a = 5 + 1; - var b = a + 5 + 1; - var c = b / 2; - var d = c * 3; - var e = d - 1; - print(e); -} - -# Print the result for verification. -# CHECK: data = -# CHECK-NEXT: 17, \ No newline at end of file diff --git a/integration_test/ops/transpose_COO_matrix.ta b/integration_test/ops/transpose_COO_matrix.ta deleted file mode 100644 index 5e8e67da..00000000 --- a/integration_test/ops/transpose_COO_matrix.ta +++ /dev/null @@ -1,32 +0,0 @@ -# RUN: comet-opt --convert-to-loops --convert-to-llvm %s &> transpose_COO_matrix.llvm -# RUN: export SPARSE_FILE_NAME0=%comet_integration_test_data_dir/test_rank2.mtx -# RUN: export SORT_TYPE=SEQ_QSORT -# RUN: mlir-cpu-runner transpose_COO_matrix.llvm -O3 -e main -entry-point-result=void -shared-libs=%comet_utility_library_dir/libcomet_runner_utils%shlibext | FileCheck %s - -def main() { - #IndexLabel Declarations - IndexLabel [i] = [?]; - IndexLabel [j] = [?]; - - #Tensor Declarations - Tensor A([i, j], COO); - Tensor B([j, i], COO); - - #Tensor Readfile Operation - A[i, j] = comet_read(0); - - #Tensor Transpose - B[j, i] = transpose(A[i, j],{j,i}); - print(B); -} - -# CHECK: data = -# CHECK-NEXT: 0,9, -# CHECK-NEXT: data = -# CHECK-NEXT: 0,0,1,1,2,3,3,4,4, -# CHECK-NEXT: data = -# CHECK-NEXT: -1, -# CHECK-NEXT: data = -# CHECK-NEXT: 0,3,1,4,2,0,3,1,4, -# CHECK-NEXT: data = -# CHECK-NEXT: 1,4.1,2,5.2,3,1.4,4,2.5,5, \ No newline at end of file diff --git a/integration_test/ops/transpose_COO_tensor.ta b/integration_test/ops/transpose_COO_tensor.ta deleted file mode 100644 index 920228ad..00000000 --- a/integration_test/ops/transpose_COO_tensor.ta +++ /dev/null @@ -1,38 +0,0 @@ -# RUN: comet-opt --convert-to-loops --convert-to-llvm %s &> transpose_COO_tensor.llvm -# RUN: export SPARSE_FILE_NAME0=%comet_integration_test_data_dir/test_rank3.tns -# RUN: export SORT_TYPE=SEQ_QSORT -# RUN: mlir-cpu-runner transpose_COO_tensor.llvm -O3 -e main -entry-point-result=void -shared-libs=%comet_utility_library_dir/libcomet_runner_utils%shlibext | FileCheck %s - -def main() { - #IndexLabel Declarations - IndexLabel [i] = [?]; - IndexLabel [j] = [?]; - IndexLabel [k] = [?]; - - #Tensor Declarations - Tensor A([i, j, k], COO); - Tensor B([j, i, k], COO); - - #Tensor Readfile Operation - A[i, j, k] = comet_read(0); - - #Tensor Transpose - B[j, i, k] = transpose(A[i, j, k],{j, i, k}); - print(B); -} - -# Print the result for verification. -# CHECK: data = -# CHECK-NEXT: 0,3, -# CHECK-NEXT: data = -# CHECK-NEXT: 1,3,6, -# CHECK-NEXT: data = -# CHECK-NEXT: 0, -# CHECK-NEXT: data = -# CHECK-NEXT: 2,1,3, -# CHECK-NEXT: data = -# CHECK-NEXT: 0, -# CHECK-NEXT: data = -# CHECK-NEXT: 3,2,5, -# CHECK-NEXT: data = -# CHECK-NEXT: 2.11,1.3,3, \ No newline at end of file diff --git a/integration_test/ops/transpose_CSF_tensor.ta b/integration_test/ops/transpose_CSF_tensor.ta deleted file mode 100644 index 6e20ced3..00000000 --- a/integration_test/ops/transpose_CSF_tensor.ta +++ /dev/null @@ -1,38 +0,0 @@ -# RUN: comet-opt --convert-to-loops --convert-to-llvm %s &> transpose_CSF_tensor.llvm -# RUN: export SPARSE_FILE_NAME0=%comet_integration_test_data_dir/test_rank3.tns -# RUN: export SORT_TYPE=SEQ_QSORT -# RUN: mlir-cpu-runner transpose_CSF_tensor.llvm -O3 -e main -entry-point-result=void -shared-libs=%comet_utility_library_dir/libcomet_runner_utils%shlibext | FileCheck %s - -def main() { - #IndexLabel Declarations - IndexLabel [i] = [?]; - IndexLabel [j] = [?]; - IndexLabel [k] = [?]; - - #Tensor Declarations - Tensor A([i, j, k], CSF); - Tensor B([k, i, j], CSF); - - #Tensor Readfile Operation - A[i, j, k] = comet_read(0); - - #Tensor Transpose - B[k, i, j] = transpose(A[i, j, k],{k, i, j}); - print(B); -} - -# Print the result for verification. -# CHECK: data = -# CHECK-NEXT: 0,3, -# CHECK-NEXT: data = -# CHECK-NEXT: 2,3,5, -# CHECK-NEXT: data = -# CHECK-NEXT: 0,1,2,3, -# CHECK-NEXT: data = -# CHECK-NEXT: 1,2,3, -# CHECK-NEXT: data = -# CHECK-NEXT: 0,1,2,3, -# CHECK-NEXT: data = -# CHECK-NEXT: 3,1,6, -# CHECK-NEXT: data = -# CHECK-NEXT: 1.3,2.11,3, diff --git a/integration_test/ops/transpose_CSR_matrix.ta b/integration_test/ops/transpose_CSR_matrix.ta deleted file mode 100644 index 8e377a14..00000000 --- a/integration_test/ops/transpose_CSR_matrix.ta +++ /dev/null @@ -1,33 +0,0 @@ -# RUN: comet-opt --convert-to-loops --convert-to-llvm %s &> transpose_CSR.llvm -# RUN: export SPARSE_FILE_NAME0=%comet_integration_test_data_dir/test_rank2.mtx -# RUN: export SORT_TYPE=SEQ_QSORT -# RUN: mlir-cpu-runner transpose_CSR.llvm -O3 -e main -entry-point-result=void -shared-libs=%comet_utility_library_dir/libcomet_runner_utils%shlibext | FileCheck %s - -def main() { - #IndexLabel Declarations - IndexLabel [i] = [?]; - IndexLabel [j] = [?]; - - #Tensor Declarations - Tensor A([i, j], CSR); - Tensor B([j, i], CSR); - - #Tensor Readfile Operation - A[i, j] = comet_read(0); - - #Tensor Transpose - B[j, i] = transpose(A[i, j],{j,i}); - print(B); -} - -# Print the result for verification. -# CHECK: data = -# CHECK-NEXT: 5, -# CHECK-NEXT: data = -# CHECK-NEXT: -1, -# CHECK-NEXT: data = -# CHECK-NEXT: 0,2,4,5,7,9, -# CHECK-NEXT: data = -# CHECK-NEXT: 0,3,1,4,2,0,3,1,4, -# CHECK-NEXT: data = -# CHECK-NEXT: 1,4.1,2,5.2,3,1.4,4,2.5,5, diff --git a/integration_test/opts/ccsd_t1_21_ttgt_tiling.ta b/integration_test/opts/ccsd_t1_21_ttgt_tiling.ta deleted file mode 100644 index 4d558506..00000000 --- a/integration_test/opts/ccsd_t1_21_ttgt_tiling.ta +++ /dev/null @@ -1,24 +0,0 @@ -# RUN: comet-opt --opt-matmul-tiling --convert-tc-to-ttgt --convert-to-llvm %s &> ccsd_t1_21_ttgt_tiling.llvm -# RUN: mlir-cpu-runner ccsd_t1_21_ttgt_tiling.llvm -O3 -e main -entry-point-result=void -shared-libs=%comet_utility_library_dir/libcomet_runner_utils%shlibext | FileCheck %s - -def main() { - #IndexLabel Declarations - IndexLabel [i, c] = [2]; - IndexLabel [m, n, a] = [4]; - - Tensor v([i, c, m, n], {Dense}); - Tensor t2([m, n, c, a], {Dense}); - Tensor i0([i, a], {Dense}); - - v[i, c, m, n] = 2.3; - t2[m, n, c, a] = 3.4; - i0[i, a] = 0.0; - - #Tensor contraction - i0[i, a] = v[i, c, m, n] * t2[m, n, c, a]; #ccsd_t1 21st expression - print(i0); -} - -# Print the result for verification. -# CHECK: data = -# CHECK-NEXT: 250.24,250.24,250.24,250.24,250.24,250.24,250.24,250.24, \ No newline at end of file diff --git a/integration_test/opts/spgemm_w_compressed_workspace.ta b/integration_test/opts/spgemm_w_compressed_workspace.ta deleted file mode 100644 index 6c6daba8..00000000 --- a/integration_test/opts/spgemm_w_compressed_workspace.ta +++ /dev/null @@ -1,39 +0,0 @@ -# Sparse matrix sparse matrix multiplication -# Sparse matrix is in CSR format. Currently workspace transformation on the IndexTree dialect works for only CSR format -# RUN: comet-opt --opt-comp-workspace --convert-ta-to-it --convert-to-loops --convert-to-llvm %s &> spgemm_w_compressed_workspace.llvm -# RUN: export SPARSE_FILE_NAME0=%comet_integration_test_data_dir/test_rank2.mtx -# RUN: export SPARSE_FILE_NAME1=%comet_integration_test_data_dir/test_rank2.mtx -# RUN: mlir-cpu-runner spgemm_w_compressed_workspace.llvm -O3 -e main -entry-point-result=void -shared-libs=%comet_utility_library_dir/libcomet_runner_utils%shlibext | FileCheck %s - - -def main() { - #IndexLabel Declarations - IndexLabel [a] = [?]; - IndexLabel [b] = [?]; - IndexLabel [c] = [?]; - - #Tensor Declarations - Tensor A([a, b], {CSR}); - Tensor B([b, c], {CSR}); - Tensor C([a, c], {CSR}); - - #Tensor Readfile Operation - A[a, b] = comet_read(0); - B[b, c] = comet_read(1); - - #Tensor Contraction - C[a, c] = A[a, b] * B[b, c]; - print(C); -} - -# Print the result for verification. -# CHECK: data = -# CHECK-NEXT: 5, -# CHECK-NEXT: data = -# CHECK-NEXT: 0, -# CHECK-NEXT: data = -# CHECK-NEXT: 0,2,4,5,7,9, -# CHECK-NEXT: data = -# CHECK-NEXT: 0,3,1,4,2,0,3,1,4, -# CHECK-NEXT: data = -# CHECK-NEXT: 6.74,7,17,17.5,9,20.5,21.74,36.4,38, \ No newline at end of file diff --git a/lib/Conversion/BlockedGpuToTriton/BlockedGpuToTriton.cpp b/lib/Conversion/BlockedGpuToTriton/BlockedGpuToTriton.cpp new file mode 100644 index 00000000..de19eab8 --- /dev/null +++ b/lib/Conversion/BlockedGpuToTriton/BlockedGpuToTriton.cpp @@ -0,0 +1,1217 @@ +#include "comet/Conversion/BlockedGpuToTriton/BlockedGpuToTriton.h" +#include "comet/Conversion/BlockedGpuToTriton/BlockedGpuToTritonConversion.h" +#include "mlir/Dialect/Arith/IR/Arith.h" +#include "mlir/Dialect/Arith/Utils/Utils.h" +#include "mlir/Dialect/Bufferization/IR/Bufferization.h" +#include "mlir/Dialect/Func/IR/FuncOps.h" +#include "mlir/Dialect/GPU/IR/GPUDialect.h" +#include "mlir/Dialect/Index/IR/IndexOps.h" +#include "mlir/Dialect/Linalg/IR/Linalg.h" +#include "mlir/Dialect/MemRef/IR/MemRef.h" +#include "mlir/IR/Builders.h" +#include "mlir/IR/BuiltinOps.h" +#include "mlir/IR/BuiltinTypeInterfaces.h" +#include "mlir/IR/BuiltinTypes.h" +#include "mlir/IR/MLIRContext.h" +#include "mlir/IR/Operation.h" +#include "mlir/IR/PatternMatch.h" +#include "mlir/IR/TypeRange.h" +#include "mlir/IR/Value.h" +#include "mlir/IR/ValueRange.h" +#include "mlir/Pass/Pass.h" +#include "mlir/Dialect/GPU/IR/GPUDialect.h" +#include "mlir/Pass/Pass.h" +#include "mlir/Dialect/SCF/IR/SCF.h" +#include "mlir/Dialect/Index/IR/IndexDialect.h" +#include "mlir/Dialect/Tensor/IR/Tensor.h" +#include "mlir/Pass/PassManager.h" +#include "mlir/Support/LLVM.h" +#include "mlir/Transforms/DialectConversion.h" +#include "mlir/Transforms/GreedyPatternRewriteDriver.h" +#include "mlir/Transforms/Passes.h" +#include "triton/Dialect/Triton/IR/Dialect.h" +#include "triton/Dialect/Triton/IR/Types.h" +#include "llvm/ADT/MapVector.h" +#include "llvm/ADT/STLExtras.h" +#include "llvm/ADT/SmallVector.h" +#include "llvm/Support/Casting.h" +#include "llvm/Support/LogicalResult.h" +#include "llvm/Support/raw_ostream.h" +#include +#include +#include + +#define GEN_PASS_CLASSES +#include "comet/Conversion/BlockedGpuToTriton/Passes.h.inc" +using namespace mlir; + + +class ConvertGpuFuncToTritonFunc : public OpConversionPattern{ + public: + using mlir::OpConversionPattern::OpConversionPattern; + ConvertGpuFuncToTritonFunc(mlir::MLIRContext* ctx) : mlir::OpConversionPattern(ctx) {} + mlir::LogicalResult + matchAndRewrite(mlir::gpu::GPUFuncOp gpuFunc, OpAdaptor adaptor, + mlir::ConversionPatternRewriter &rewriter) const override { + + auto converter = getTypeConverter(); + llvm::SmallVector materializedValues; + gpu::GPUModuleOp gpuModuleOp = gpuFunc->getParentOfType(); + rewriter.setInsertionPointToStart(gpuModuleOp.getBody()); + + llvm::SmallVector tritonFuncTypes; + llvm::SmallVector sizes; + for(auto argType: gpuFunc.getArgumentTypes()) + { + llvm::SmallVector currentTritonFuncTypes; + if(failed(converter->convertTypes(argType, currentTritonFuncTypes))) + { + return failure(); + } + tritonFuncTypes.insert(tritonFuncTypes.end(), currentTritonFuncTypes.begin(), currentTritonFuncTypes.end()); + sizes.push_back(currentTritonFuncTypes.size()); + } + + auto tritonFuncType = rewriter.getFunctionType(TypeRange(tritonFuncTypes), gpuFunc.getFunctionType().getResults()); + auto tritonFunc = rewriter.create(gpuFunc->getLoc(), "tt_"+ gpuModuleOp.getName().str() +gpuFunc.getName().str(), tritonFuncType); + auto ttFuncBlock = tritonFunc.addEntryBlock(); + rewriter.setInsertionPointToStart(ttFuncBlock); + size_t prev = 0; + for(size_t i = 0; i < gpuFunc.getBody().getArguments().size(); i++) + { + if(!converter->isLegal( gpuFunc.getBody().getArguments()[i].getType())) + { + auto conveted = converter->materializeArgumentConversion(rewriter, gpuFunc.getBody().getArguments()[i].getLoc(), gpuFunc.getBody().getArguments()[i].getType(), ttFuncBlock->getArguments().slice(prev, sizes[i])); + materializedValues.push_back(conveted); + } + else + { + materializedValues.push_back(ttFuncBlock->getArguments()[prev]); + } + rewriter.replaceAllUsesWith(gpuFunc.getBody().getArguments()[i], materializedValues.back()); + prev += sizes[i]; + } + + rewriter.eraseOp(gpuFunc.getBody().front().getTerminator()); + rewriter.mergeBlocks(&gpuFunc.getBody().front(), ttFuncBlock, materializedValues); + rewriter.setInsertionPointToEnd(ttFuncBlock); + rewriter.create(tritonFunc->getLoc()); + rewriter.setInsertionPointToEnd(gpuModuleOp.getBody()); + func::FuncOp newFunc = rewriter.replaceOpWithNewOp(gpuFunc, gpuFunc.getName(), gpuFunc.getFunctionType()); + auto entryBlock = newFunc.addEntryBlock(); + rewriter.setInsertionPointToEnd(entryBlock); + rewriter.create(gpuFunc.getLoc()); + newFunc->setAttr(gpu::GPUDialect::getKernelFuncAttrName(), + rewriter.getUnitAttr()); + newFunc.setArgAttrsAttr(gpuFunc.getArgAttrsAttr()); + + return success(); + } + +}; + +class ConvertForOp : public OpConversionPattern{ + public: + using mlir::OpConversionPattern::OpConversionPattern; + ConvertForOp(mlir::MLIRContext* ctx) : mlir::OpConversionPattern(ctx) {} + mlir::LogicalResult + matchAndRewrite(mlir::scf::ForOp forOp, OpAdaptor adaptor, + mlir::ConversionPatternRewriter &rewriter) const override { + + llvm::SmallVector materializedValues; + + llvm::SmallVector newForValues; + llvm::SmallVector sizes; + for(auto arg: forOp->getOperands()) + { + + if(arg.getType().isIndex()) + { + auto cast_index = rewriter.create(forOp->getLoc(), IntegerType::get(getContext(), 32), arg); + newForValues.push_back(cast_index); + } + else + { + newForValues.push_back(arg); + } + } + + auto newForOp = rewriter.create(forOp->getLoc(), newForValues[0], newForValues[1], newForValues[2], ValueRange(newForValues).drop_front(3)); + rewriter.setInsertionPointToStart(newForOp.getBody()); + llvm::SmallVector newForArgs; + + for(auto arg: newForOp.getBody()->getArguments()) + { + if(arg.getType().isInteger()) + { + auto cast_index = rewriter.create(forOp->getLoc(), rewriter.getIndexType(), arg); + newForArgs.push_back(cast_index); + } + else + { + newForArgs.push_back(arg); + } + } + + rewriter.mergeBlocks(forOp.getBody(), newForOp.getBody(), newForArgs); + rewriter.replaceAllOpUsesWith(forOp, newForOp); + rewriter.eraseOp(forOp); + + return success(); + } + +}; + + +class ConvertBlockId : public OpConversionPattern{ + public: + using mlir::OpConversionPattern::OpConversionPattern; + ConvertBlockId(mlir::MLIRContext* ctx) : mlir::OpConversionPattern(ctx) {} + mlir::LogicalResult + matchAndRewrite(mlir::gpu::BlockIdOp blockIdOp, OpAdaptor adaptor, + mlir::ConversionPatternRewriter &rewriter) const override { + + triton::ProgramIDDim axis; + switch (blockIdOp.getDimension()) + { + case mlir::gpu::Dimension::x: + { + axis = triton::ProgramIDDim::X; + break; + } + case mlir::gpu::Dimension::y: + { + axis = triton::ProgramIDDim::Y; + break; + } + case mlir::gpu::Dimension::z: + { + axis = triton::ProgramIDDim::Z; + break; + } + } + + auto getProgram = rewriter.create(blockIdOp->getLoc(), axis); + auto cast = rewriter.create(blockIdOp->getLoc(), rewriter.getIndexType(), getProgram); + rewriter.replaceOp(blockIdOp, cast); + + return success(); + } +}; + +class ConvertBlockDim : public OpConversionPattern{ + public: + using mlir::OpConversionPattern::OpConversionPattern; + ConvertBlockDim(mlir::MLIRContext* ctx) : mlir::OpConversionPattern(ctx) {} + mlir::LogicalResult + matchAndRewrite(mlir::gpu::BlockDimOp BlockDimOp, OpAdaptor adaptor, + mlir::ConversionPatternRewriter &rewriter) const override { + + triton::ProgramIDDim axis; + switch (BlockDimOp.getDimension()) + { + case mlir::gpu::Dimension::x: + { + axis = triton::ProgramIDDim::X; + break; + } + case mlir::gpu::Dimension::y: + { + axis = triton::ProgramIDDim::Y; + break; + } + case mlir::gpu::Dimension::z: + { + axis = triton::ProgramIDDim::Z; + break; + } + } + + auto getProgram = rewriter.create(BlockDimOp->getLoc(), axis); + auto cast = rewriter.create(BlockDimOp->getLoc(), rewriter.getIndexType(), getProgram); + rewriter.replaceOp(BlockDimOp, cast); + + return success(); + } +}; + +class ConvertTensorSplatOp : public OpConversionPattern{ + public: + using mlir::OpConversionPattern::OpConversionPattern; + ConvertTensorSplatOp(mlir::MLIRContext* ctx) : mlir::OpConversionPattern(ctx) {} + mlir::LogicalResult + matchAndRewrite(mlir::tensor::SplatOp SplatOp, OpAdaptor adaptor, + mlir::ConversionPatternRewriter &rewriter) const override { + + MLIRContext* ctx = rewriter.getContext(); + RankedTensorType outTensorType; + Value input; + if(isa(adaptor.getInput().getType())) + { + IntegerType i32 = IntegerType::get(ctx, 32); + input = rewriter.create(SplatOp->getLoc(), i32, adaptor.getInput()); + outTensorType = RankedTensorType::get(SplatOp.getResult().getType().getShape(), i32); + } + else + { + input = adaptor.getInput(); + outTensorType = SplatOp.getResult().getType(); + } + + Value ttSplatOp = rewriter.create(SplatOp->getLoc(), outTensorType, input); + rewriter.replaceOpWithNewOp(SplatOp, SplatOp.getType(), ttSplatOp); + + return success(); + } +}; + +template +std::pair generate_triton_blocked_bounds_and_offsets(ConversionPatternRewriter& rewriter, T sliceOp, RankedTensorType sliceT, SmallVector collapsed_indices={}) +{ + llvm::SmallVector blockedOffsets; + llvm::SmallVector blockedBounds; + SmallVector tritonOffsets, tritonBounds; + llvm::SmallMapVector dimToValue; + size_t dimIndex = 0; + for(size_t i = 0; i < sliceT.getRank(); i++) + { + if(sliceOp.getStaticSizes()[i] == ShapedType::kDynamic) + { + dimToValue[i] = sliceOp.getSizes()[dimIndex]; + dimIndex ++; + } + } + + + for(size_t i = 0; i < sliceOp.getOffsets().size(); i++) + { + auto offset = sliceOp.getOffsets()[i]; + if(auto castOp = mlir::dyn_cast(offset.getDefiningOp())) + { + Value cast; + if(auto shaped = mlir::dyn_cast(castOp.getInputs().front().getType())) + { + cast = rewriter.create(sliceOp->getLoc(), RankedTensorType::get(shaped.getShape(), rewriter.getIntegerType(32)), castOp.getInputs().front()); + } + else + { + cast = rewriter.create(sliceOp->getLoc(), rewriter.getIntegerType(32), castOp.getInputs().front()); + } + blockedOffsets.push_back(cast); + } + else if(mlir::isa(offset.getType())) + { + auto cast = rewriter.create(sliceOp->getLoc(), rewriter.getIntegerType(32), offset); + blockedOffsets.push_back(cast); + } + + auto blockSize = sliceT.getDimSize(i); + if(sliceOp.getStaticSizes()[i] == ShapedType::kDynamic) + { + + Value bound = dimToValue[i]; + if(auto min = mlir::dyn_cast_if_present(bound.getDefiningOp())) + { + bound = min.getRhs(); + } + Value cast = rewriter.create(bound.getLoc(), rewriter.getIntegerType(32), bound); + + Value blockedBound; + Value blockedBlockSize; + if(blockSize > 1) + { + blockedBound = rewriter.create(bound.getLoc(), RankedTensorType::get({blockSize}, rewriter.getIntegerType(32)), cast); + blockedBlockSize = rewriter.create(bound.getLoc(), RankedTensorType::get({blockSize}, rewriter.getIntegerType(32)), 0, blockSize); + } + else + { + blockedBound = cast; + blockedBlockSize = rewriter.create(bound.getLoc(), 0, 32); + } + auto cmp = rewriter.create(bound.getLoc(), arith::CmpIPredicate::slt, blockedBlockSize, blockedBound); + blockedBounds.push_back(cmp); + tritonBounds.push_back(cmp); + } + else + { + blockedBounds.push_back(rewriter.create(sliceOp.getLoc(), 1 ,1 )); + tritonBounds.push_back(rewriter.create(sliceOp.getLoc(),1 ,1)); + } + } + + + SmallVector resultShape; + + for(size_t i = 0; i < blockedOffsets.size(); i++) + { + auto blockedOffset = blockedOffsets[i]; + Value strideVal = rewriter.create(sliceOp->getLoc(), rewriter.getIntegerType(32), sliceOp.getStrides()[i]); + Value stridesBlocked = strideVal; + auto shapedOffset = mlir::dyn_cast(blockedOffset.getType()); + if(shapedOffset) + { + stridesBlocked = rewriter.create(sliceOp->getLoc(), blockedOffset.getType(), strideVal); + } + + Value offsetBlocked = rewriter.create(sliceOp->getLoc(), blockedOffset, stridesBlocked); + tritonOffsets.push_back(offsetBlocked); + if(auto shaped = mlir::dyn_cast(blockedOffset.getType())) + { + resultShape.push_back(shaped.getDimSize(0)); + } + else + { + resultShape.push_back(1); + } + } + + SmallVector collapsed_offsets; + SmallVector collapsed_bounds; + SmallVector collapsed_shape; + if(collapsed_indices.empty()) + { + collapsed_offsets = tritonOffsets; + for(auto d: resultShape) + { + if (d != 1) + { + collapsed_shape.push_back(d); + } + } + // collapsed_shape = resultShape; + + collapsed_bounds = tritonBounds; + } + + for(auto collapsed: collapsed_indices) + { + size_t shaped_i = 0; + for(size_t i = 0; i < collapsed.size(); i++) + { + auto shapedOffset = mlir::dyn_cast(tritonOffsets[collapsed[i]].getType()); + if(shapedOffset && shapedOffset.getDimSize(0) != 1) + { + shaped_i = i; + } + } + collapsed_offsets.push_back(tritonOffsets[shaped_i]); + for(size_t i = 0; i < collapsed.size(); i++) + { + if(i == shaped_i) + continue; // skip the shaped index + + Value offset = tritonOffsets[collapsed[i]]; + if(tritonOffsets[collapsed[i]].getType() != collapsed_offsets.back().getType()) + { + if(!isa(tritonOffsets[collapsed[i]].getType())) + { + offset = rewriter.create(sliceOp->getLoc(), collapsed_offsets.back().getType(), tritonOffsets[collapsed[i]]); + } + else + { + offset = rewriter.create(sliceOp->getLoc(), collapsed_offsets.back().getType(), tritonOffsets[collapsed[i]]); + } + } + auto res = rewriter.create(sliceOp->getLoc(), collapsed_offsets.back(), offset); + collapsed_offsets.back() = res; + } + if(resultShape[shaped_i] != 1) + { + collapsed_shape.push_back(resultShape[shaped_i]); + } + } + + for(auto collapsed: collapsed_indices) + { + size_t shaped_i = 0; + for(size_t i = 0; i < collapsed.size(); i++) + { + auto shapedOffset = mlir::dyn_cast(tritonBounds[collapsed[i]].getType()); + if(shapedOffset && shapedOffset.getDimSize(0) != 1) + { + shaped_i = i; + } + } + collapsed_bounds.push_back(tritonBounds[shaped_i]); + for(size_t i = 0; i < collapsed.size(); i++) + { + if(i == shaped_i) + continue; // skip the shaped index + + Value offset = tritonBounds[collapsed[i]]; + if(tritonBounds[collapsed[i]].getType() != collapsed_bounds.back().getType()) + { + offset = rewriter.create(sliceOp->getLoc(), collapsed_bounds.back().getType(), tritonBounds[collapsed[i]]); + } + auto res = rewriter.create(sliceOp->getLoc(), collapsed_bounds.back(), offset); + collapsed_bounds.back() = res; + } + } + + for(size_t i = 0; i < collapsed_offsets.size(); i++) + { + for(size_t j = 0; j (collapsed_offsets[i].getType())) + { + collapsed_offsets[i] = rewriter.create(sliceOp->getLoc(), RankedTensorType::get(1, rewriter.getIntegerType(32)), collapsed_offsets[i]); + } + collapsed_offsets[i] = rewriter.create(sliceOp->getLoc(), collapsed_offsets[i], j); + + } + } + + if(!collapsed_shape.empty()) + { + collapsed_offsets[i] = rewriter.create(sliceOp->getLoc(), RankedTensorType::get(collapsed_shape, rewriter.getIntegerType(32)), collapsed_offsets[i]); + } + } + + auto resultType = RankedTensorType::get(collapsed_shape, rewriter.getIntegerType(1)); + size_t k = 0; + for(size_t i = 0; i < collapsed_bounds.size(); i++) + { + for(size_t j = 0; j (collapsed_bounds[i].getType())) + { + collapsed_bounds[i] = rewriter.create(sliceOp->getLoc(), RankedTensorType::get(1, rewriter.getIntegerType(1)), collapsed_bounds[i]); + + } + collapsed_bounds[i] = rewriter.create(sliceOp->getLoc(), collapsed_bounds[i], j); + + } + } + if(collapsed_bounds[i].getType() != resultType) + { + if(!mlir::isa(collapsed_bounds[i].getType())) + { + collapsed_bounds[i] = rewriter.create(sliceOp.getLoc(), RankedTensorType::get(resultType.getShape(), rewriter.getIntegerType(1)), collapsed_bounds[i]); + } + else if(mlir::cast(collapsed_bounds[i].getType()).getRank() != resultType.getRank()) + { + for(size_t j = 0; j < resultType.getRank(); j++) + { + if(resultType.getDimSize(j) != mlir::cast(collapsed_bounds[i].getType()).getDimSize(k)) + { + collapsed_bounds[i] = rewriter.create(sliceOp->getLoc(), collapsed_bounds[i], k++); + } + else + { + k++; + } + } + } + } + if(!collapsed_shape.empty()) + { + collapsed_bounds[i] = rewriter.create(sliceOp->getLoc(), resultType, collapsed_bounds[i]); + } + } + + Value combinedOffsetBlocked, combinedBoundBlocked = nullptr; + combinedOffsetBlocked = collapsed_offsets.front(); + if(!collapsed_bounds.empty()) + { + combinedBoundBlocked = collapsed_bounds.front(); + } + for(size_t i = 1; i < collapsed_offsets.size(); i++) + { + combinedOffsetBlocked = rewriter.create(sliceOp->getLoc(), combinedOffsetBlocked, collapsed_offsets[i]); + } + + for(size_t i = 1; i < collapsed_bounds.size(); i++) + { + combinedBoundBlocked = rewriter.create(sliceOp->getLoc(), combinedBoundBlocked, collapsed_bounds[i]); + } + + return std::make_pair(combinedOffsetBlocked, combinedBoundBlocked); +} + +// class ConvertTensorInsert: public OpConversionPattern{ +// public: +// using mlir::OpConversionPattern::OpConversionPattern; +// ConvertInsertSlice(mlir::MLIRContext* ctx) : mlir::OpConversionPattern(ctx) {} +// mlir::LogicalResult +// matchAndRewrite(mlir::tensor::InsertSliceOp insertSliceOp, OpAdaptor adaptor, +// mlir::ConversionPatternRewriter &rewriter) const override { +// } + +// }; + + + +class ConvertInsertSlice : public OpConversionPattern{ + public: + using mlir::OpConversionPattern::OpConversionPattern; + ConvertInsertSlice(mlir::MLIRContext* ctx) : mlir::OpConversionPattern(ctx) {} + mlir::LogicalResult + matchAndRewrite(mlir::tensor::InsertSliceOp insertSliceOp, OpAdaptor adaptor, + mlir::ConversionPatternRewriter &rewriter) const override { + + Value dest = adaptor.getDest(); + Value ttSource = adaptor.getSource(); + Value ttDest; + bufferization::ToTensorOp toTensor = mlir::dyn_cast(dest.getDefiningOp()); + UnrealizedConversionCastOp castOp; + if(toTensor) + { + if((castOp = mlir::dyn_cast(toTensor.getMemref().getDefiningOp()))) + { + if(triton::PointerType ttPtr = mlir::dyn_cast(castOp->getOperand(0).getType())) + { + ttDest = castOp->getOperand(0); + } + } + } + else if((castOp = mlir::dyn_cast(toTensor.getMemref().getDefiningOp()))) + { + if(triton::PointerType ttPtr = mlir::dyn_cast(castOp->getOperand(0).getType())) + { + ttDest = castOp->getOperand(0); + } + } + + if(!ttDest) + { + return failure(); + } + + RankedTensorType sliceT; + Operation* toReplace = insertSliceOp; + SmallVector expanded_indices; + + if(auto tensorCastOp = mlir::dyn_cast_if_present(insertSliceOp.getSource().getDefiningOp())) + { + sliceT = mlir::cast(tensorCastOp.getSource().getType()); + if(auto expand_op = mlir::dyn_cast_if_present(tensorCastOp.getSource().getDefiningOp())) + { + expanded_indices = expand_op.getReassociationIndices(); + ttSource = expand_op.getSrc(); + } + else if(auto expand_op = mlir::dyn_cast_if_present(insertSliceOp.getSource().getDefiningOp())) + { + expanded_indices = expand_op.getReassociationIndices(); + ttSource = expand_op.getSrc(); + } + } + else + { + sliceT = mlir::cast(insertSliceOp.getSource().getType()); + if(auto expand_op = mlir::dyn_cast(insertSliceOp.getSource().getDefiningOp())) + { + expanded_indices = expand_op.getReassociationIndices(); + ttSource = expand_op.getSrc(); + } + } + + auto [combinedOffsetBlocked, combinedBoundBlocked] = generate_triton_blocked_bounds_and_offsets(rewriter, insertSliceOp, sliceT, expanded_indices); + std::vector resultShape; + if(auto offset_shaped = dyn_cast(combinedOffsetBlocked.getType())) + { + resultShape = offset_shaped.getShape(); + } + + Value blockedPtr = ttDest; + if(!resultShape.empty()) + { + blockedPtr = rewriter.create(toReplace->getLoc(), RankedTensorType::get(resultShape, ttDest.getType()), ttDest); + } + + Value ptr; + if(!resultShape.empty()) + { + ptr = rewriter.create(toReplace->getLoc(), RankedTensorType::get(resultShape, ttDest.getType()), blockedPtr, combinedOffsetBlocked); + } + else + { + ptr = rewriter.create(toReplace->getLoc(), ttDest.getType(), blockedPtr, combinedOffsetBlocked); + } + + if(ShapedType sourceShaped = mlir::dyn_cast(ttSource.getType()); sourceShaped && !sourceShaped.hasStaticShape()) + { + auto castOp = rewriter.create(toReplace->getLoc(), RankedTensorType::get(resultShape, insertSliceOp.getDestType().getElementType()), ttSource); + // if(combinedBoundBlocked) + // { + // rewriter.create( insertSliceOp->getLoc(), ptr, castOp, combinedBoundBlocked, mlir::triton::CacheModifier::NONE, mlir::triton::EvictionPolicy::NORMAL); + // } + // else + { + if(auto ranked_source = dyn_cast(castOp.getType())) + { + if(!isa(ptr.getType())) + { + if(llvm::all_of(ranked_source.getShape(), [](int64_t d){ + return d == 1; + })) + { + ptr = rewriter.create(ptr.getLoc(), RankedTensorType::get(ranked_source.getShape(), ptr.getType()), ptr); + } + } + } + rewriter.create( toReplace->getLoc(), ptr, castOp, combinedBoundBlocked, mlir::triton::CacheModifier::NONE, mlir::triton::EvictionPolicy::NORMAL); + } + auto placeholder = rewriter.create(toReplace->getLoc(), insertSliceOp.getResultType(), ValueRange(ptr)); + rewriter.replaceOp(insertSliceOp, placeholder->getResult(0)); + } + else + { + // if(combinedBoundBlocked) + // { + // rewriter.create( insertSliceOp->getLoc(), ptr, ttSource, combinedBoundBlocked, mlir::triton::CacheModifier::NONE, mlir::triton::EvictionPolicy::NORMAL); + // } + // else + { + if(auto ranked_source = dyn_cast(ttSource.getType())) + { + if(!isa(ptr.getType())) + { + if(llvm::all_of(ranked_source.getShape(), [](int64_t d){ + return d == 1; + })) + { + ptr = rewriter.create(ptr.getLoc(), RankedTensorType::get(ranked_source.getShape(), ptr.getType()), ptr); + } + } + } + rewriter.create( insertSliceOp->getLoc(), ptr, ttSource, combinedBoundBlocked, mlir::triton::CacheModifier::NONE, mlir::triton::EvictionPolicy::NORMAL); + } + + auto placeholder = rewriter.create(insertSliceOp->getLoc(), insertSliceOp.getResultType(), ValueRange(ptr)); + rewriter.replaceOp(insertSliceOp, placeholder->getResult(0)); + } + + if(toTensor->getUsers().empty()) + { + rewriter.eraseOp(toTensor); + } + if(castOp->getUsers().empty()) + { + rewriter.eraseOp(castOp); + } + return success(); + } +}; + +class ConvertExtractSlice : public OpConversionPattern{ + public: + using mlir::OpConversionPattern::OpConversionPattern; + ConvertExtractSlice(mlir::MLIRContext* ctx) : mlir::OpConversionPattern(ctx) {} + mlir::LogicalResult + matchAndRewrite(mlir::tensor::ExtractSliceOp extractSliceOp, OpAdaptor adaptor, + mlir::ConversionPatternRewriter &rewriter) const override { + Value source = adaptor.getSource(); + Value ttSource; + bufferization::ToTensorOp toTensor = mlir::dyn_cast(source.getDefiningOp()); + UnrealizedConversionCastOp castOp; + if(toTensor) + { + if((castOp = mlir::dyn_cast_if_present(toTensor.getMemref().getDefiningOp()))) + { + if(triton::PointerType ttPtr = mlir::dyn_cast(castOp->getOperand(0).getType())) + { + ttSource = castOp->getOperand(0); + } + } + } + else if((castOp = mlir::dyn_cast_if_present(toTensor.getMemref().getDefiningOp()))) + { + if(triton::PointerType ttPtr = mlir::dyn_cast(castOp->getOperand(0).getType())) + { + ttSource = castOp->getOperand(0); + } + } + + if(!ttSource) + { + return failure(); + } + + RankedTensorType sliceT; + SmallVector collapsed_indices; + Operation* toReplace = extractSliceOp; + auto tensorCastOp = mlir::dyn_cast(*extractSliceOp->getUsers().begin()); + Value extractOp = nullptr; + Value collapseOp = nullptr; + if(tensorCastOp) + { + sliceT = mlir::cast(tensorCastOp.getResult().getType()); + if(auto collapse_op = mlir::dyn_cast(*tensorCastOp->getUsers().begin())) + { + collapsed_indices = collapse_op.getReassociationIndices(); + toReplace = collapse_op; + collapseOp = collapse_op; + if(auto extract_op = mlir::dyn_cast(*collapse_op->getUsers().begin())) + { + extractOp = extract_op; + toReplace = extract_op; + } + } + else if(auto extract_op = mlir::dyn_cast(*tensorCastOp->getUsers().begin())) + { + extractOp = extract_op; + toReplace = extract_op; + } + } + else + { + sliceT = mlir::cast(extractSliceOp.getResultType()); + if(auto collapse_op = mlir::dyn_cast(*extractSliceOp->getUsers().begin())) + { + collapsed_indices = collapse_op.getReassociationIndices(); + toReplace = collapse_op; + collapseOp = collapse_op; + + if(auto extract_op = mlir::dyn_cast(*collapse_op->getUsers().begin())) + { + extractOp = extract_op; + toReplace = extract_op; + } + } + else if(auto extract_op = mlir::dyn_cast(*extractSliceOp->getUsers().begin())) + { + extractOp = extract_op; + toReplace = extract_op; + } + } + + + // RankedTensorType sliceT = mlir::cast(tensorCastOp.getResult().getType()); + auto [combinedOffsetBlocked, combinedBoundBlocked] = generate_triton_blocked_bounds_and_offsets(rewriter, extractSliceOp, sliceT, collapsed_indices); + + std::vector resultShape; + if(auto ranked_offset_type = mlir::dyn_cast(combinedOffsetBlocked.getType())) + { + resultShape = ranked_offset_type.getShape().vec(); + } + + Value blockedPtr = ttSource; + + if(!resultShape.empty()) + { + blockedPtr = rewriter.create(extractSliceOp->getLoc(), RankedTensorType::get(resultShape, ttSource.getType()), ttSource); + } + + Value ptr; + if(!resultShape.empty()) + { + ptr = rewriter.create(extractSliceOp->getLoc(), RankedTensorType::get(resultShape, ttSource.getType()), blockedPtr, combinedOffsetBlocked); + } + else + { + ptr = rewriter.create(extractSliceOp->getLoc(), ttSource.getType(), blockedPtr, combinedOffsetBlocked); + } + + + + triton::LoadOp loadOp; + if(combinedBoundBlocked) + { + loadOp = rewriter.create(extractSliceOp->getLoc(), ptr, combinedBoundBlocked, mlir::triton::CacheModifier::NONE, mlir::triton::EvictionPolicy::NORMAL, false); + } + else + { + loadOp = rewriter.create(extractSliceOp->getLoc(), ptr, mlir::triton::CacheModifier::NONE, mlir::triton::EvictionPolicy::NORMAL, false); + + } + if(auto shaped = mlir::dyn_cast(loadOp.getType())) + { + SmallVector rank; + for(int64_t i = 0; i < shaped.getRank(); i ++) + { + rank.push_back(ShapedType::kDynamic); + } + auto cast = rewriter.create(loadOp->getLoc(), RankedTensorType::get(rank, shaped.getElementType()), loadOp); + rewriter.replaceOpUsesWithIf(extractSliceOp, cast->getResults(), + [&](OpOperand& opOperand) { + return opOperand.get().getType() != loadOp.getType(); + }); + rewriter.replaceOpUsesWithIf(extractSliceOp, loadOp->getResults(), + [&](OpOperand& opOperand) { + return opOperand.get().getType() == loadOp.getType(); + }); + + rewriter.eraseOp(extractSliceOp); + } + else + { + rewriter.replaceOp(extractSliceOp, loadOp); + } + + if(toTensor->getUsers().empty()) + { + rewriter.eraseOp(toTensor); + } + if(castOp->getUsers().empty()) + { + rewriter.eraseOp(castOp); + } + if(toReplace != extractSliceOp) + { + rewriter.replaceOp(toReplace, loadOp); + if(tensorCastOp) + { + rewriter.eraseOp(tensorCastOp); + } + if(extractOp && collapseOp) + { + rewriter.eraseOp(collapseOp.getDefiningOp()); + } + } + + return success(); + } +}; + +class ConvertToTensor : public OpConversionPattern{ + public: + using mlir::OpConversionPattern::OpConversionPattern; + ConvertToTensor(mlir::MLIRContext* ctx) : mlir::OpConversionPattern(ctx) {} + mlir::LogicalResult + matchAndRewrite(mlir::bufferization::ToTensorOp toTensor, OpAdaptor adaptor, + mlir::ConversionPatternRewriter &rewriter) const override { + + if(toTensor->getUsers().empty()) + { + rewriter.eraseOp(toTensor); + return success(); + } + else + { + // for(auto user: toTensor->getUsers()) + // { + // user->dump(); + // } + } + + return failure(); + } +}; + +class ConvertUnrealizedCast : public OpConversionPattern{ + public: + using mlir::OpConversionPattern::OpConversionPattern; + ConvertUnrealizedCast(mlir::MLIRContext* ctx) : mlir::OpConversionPattern(ctx) {} + mlir::LogicalResult + matchAndRewrite(mlir::UnrealizedConversionCastOp cast, OpAdaptor adaptor, + mlir::ConversionPatternRewriter &rewriter) const override { + if(cast->getUsers().empty()) + { + rewriter.eraseOp(cast); + return success(); + } + else + { + if(adaptor.getInputs().size() == 1) + { + if(auto index = mlir::dyn_cast_if_present(adaptor.getInputs()[0].getDefiningOp())) + { + ShapedType shape = mlir::cast(cast->getResultTypes()[0]); + auto range_op = rewriter.create(cast->getLoc(), RankedTensorType::get(shape.getShape(), rewriter.getIntegerType(32)), 0, shape.getShape()[0]); + auto indexCast = rewriter.create(cast->getLoc(), shape, range_op); + rewriter.replaceOp(cast, indexCast); + return success(); + } + } + } + + return failure(); + } +}; + +template +class ConvertArithIndex : public OpConversionPattern { + public: + using mlir::OpConversionPattern::OpConversionPattern; + ConvertArithIndex(mlir::MLIRContext* ctx) : mlir::OpConversionPattern(ctx) {} + mlir::LogicalResult + matchAndRewrite(T arith_op, typename T::Adaptor adaptor, + mlir::ConversionPatternRewriter &rewriter) const override { + + Operation* op = arith_op.getOperation(); + auto converter = mlir::OpConversionPattern::getTypeConverter(); + llvm::SmallVector convertedTypes; + llvm::SmallVector newOperands; + if(failed(converter->convertTypes(op->getOperandTypes(), convertedTypes))) + { + return failure(); + } + + for(auto operand: llvm::zip(convertedTypes,op->getOperands())) + { + newOperands.push_back(converter->materializeTargetConversion(rewriter, op->getLoc(), std::get<0>(operand), std::get<1>(operand))); + } + + Value new_arith_op = rewriter.create(op->getLoc(), newOperands); + rewriter.replaceOp(arith_op, converter->materializeSourceConversion(rewriter, op->getLoc(), op->getResultTypes()[0], new_arith_op)); + + return success(); + } +}; + +class ConvertLinalgReduceOp : public OpConversionPattern { + public: + using mlir::OpConversionPattern::OpConversionPattern; + ConvertLinalgReduceOp(mlir::MLIRContext* ctx) : mlir::OpConversionPattern(ctx) {} + mlir::LogicalResult + matchAndRewrite(linalg::ReduceOp reduceOp, OpAdaptor adaptor, + mlir::ConversionPatternRewriter &rewriter) const override { + + assert(adaptor.getDimensions().size() == 1); + triton::ReduceOp ttReduceOp = rewriter.create(reduceOp->getLoc(), adaptor.getInputs(), static_cast(adaptor.getDimensions().front())); + auto arg_types = reduceOp.getBody()->getArgumentTypes(); + std::vector locs = {reduceOp->getOperand(0).getLoc(), reduceOp->getOperand(1).getLoc()}; + // rewriter.setInsertionPointToStart(ttReduceOp.getBody()); + Block* reduceBody = rewriter.createBlock(&ttReduceOp.getBodyRegion(),{}, arg_types, locs); + + auto yieldOp = reduceOp.getBody()->getTerminator(); + rewriter.mergeBlocks(reduceOp.getBody(), reduceBody, reduceBody->getArguments()); + rewriter.setInsertionPointToEnd(ttReduceOp.getBody()); + rewriter.replaceOpWithNewOp(yieldOp, yieldOp->getOperands()); + rewriter.replaceAllUsesWith(reduceOp.getResult(0), ttReduceOp.getResult()); + rewriter.eraseOp(reduceOp); + if(!isa(ttReduceOp->getResultTypes()[0])) + { + + if(auto extractOp = dyn_cast(*ttReduceOp->getUsers().begin())) + { + rewriter.replaceAllUsesWith(extractOp, ttReduceOp->getResult(0)); + rewriter.eraseOp(extractOp); + } + } + + return success(); + } +}; + +class ConvertLinalgBroadcastOp : public OpConversionPattern { + public: + using mlir::OpConversionPattern::OpConversionPattern; + ConvertLinalgBroadcastOp(mlir::MLIRContext* ctx) : mlir::OpConversionPattern(ctx) {} + mlir::LogicalResult + matchAndRewrite(linalg::BroadcastOp bcastOp, OpAdaptor adaptor, + mlir::ConversionPatternRewriter &rewriter) const override { + + assert(adaptor.getDimensions().size() == 1); + auto expandDimsOp = rewriter.create(bcastOp->getLoc(), adaptor.getInput(), static_cast(adaptor.getDimensions().front())); + rewriter.replaceOpWithNewOp(bcastOp, adaptor.getInit().getType(), expandDimsOp); + return success(); + } +}; + +class ConvertLinalgMatmulOp : public OpConversionPattern { + public: + using mlir::OpConversionPattern::OpConversionPattern; + ConvertLinalgMatmulOp(mlir::MLIRContext* ctx) : mlir::OpConversionPattern(ctx) {} + mlir::LogicalResult + matchAndRewrite(linalg::MatmulOp matmulOp, OpAdaptor adaptor, + mlir::ConversionPatternRewriter &rewriter) const override { + + assert(matmulOp->hasOneUse()); + arith::AddFOp addOp = mlir::dyn_cast(*matmulOp->getUsers().begin()); + assert(addOp); + Value c; + if(matmulOp.getResult(0) == addOp.getLhs()) + { + c = addOp.getRhs(); + } + else + { + c = addOp.getLhs(); + } + + + triton::DotOp dot = rewriter.create(matmulOp->getLoc(), matmulOp.getInputs()[0], matmulOp.getInputs()[1], c); + rewriter.eraseOp(matmulOp); + rewriter.replaceOp(*matmulOp->getUsers().begin(), dot); + // rewriter.eraseOp(matmulOp->getUsers()) + return success(); + } +}; + + + +template +class ConvertArithDynamicShape : public OpConversionPattern { + public: + using mlir::OpConversionPattern::OpConversionPattern; + ConvertArithDynamicShape(mlir::MLIRContext* ctx) : mlir::OpConversionPattern(ctx) {} + mlir::LogicalResult + matchAndRewrite(T arith_op, typename T::Adaptor adaptor, + mlir::ConversionPatternRewriter &rewriter) const override { + + Operation* op = arith_op.getOperation(); + llvm::SmallVector newOperands; + for(auto operand: op->getOperands()) + { + if(auto shapedOperand = mlir::dyn_cast(operand.getType())) + { + if(shapedOperand.hasStaticShape()) + { + newOperands.push_back(operand); + } + else + { + if(auto castOp = mlir::dyn_cast_if_present(operand.getDefiningOp())) + { + newOperands.push_back(castOp.getSource()); + } + else + { + return failure(); + } + } + } + else + { + newOperands.push_back(operand); + } + } + + auto new_arith_op = rewriter.create(op->getLoc(), newOperands); + rewriter.replaceOpWithNewOp(op, op->getResultTypes().front(), new_arith_op.getResult()); + return success(); + } +}; + +class ConvertIndexConstant : public OpConversionPattern { + public: + using mlir::OpConversionPattern::OpConversionPattern; + ConvertIndexConstant(mlir::MLIRContext* ctx) : mlir::OpConversionPattern(ctx) {} + mlir::LogicalResult + matchAndRewrite(index::ConstantOp constantOp, OpAdaptor adaptor, + mlir::ConversionPatternRewriter &rewriter) const override { + + Value i32_const = rewriter.create(constantOp->getLoc(), constantOp.getValue().getSExtValue(), 32); + rewriter.replaceOpWithNewOp(constantOp, IndexType::get(getContext()), i32_const); + + return success(); + } +}; + +class ConvertExtractStridedMetadata : public OpConversionPattern { + public: + using mlir::OpConversionPattern::OpConversionPattern; + ConvertExtractStridedMetadata(mlir::MLIRContext* ctx) : mlir::OpConversionPattern(ctx) {} + mlir::LogicalResult + matchAndRewrite(mlir::memref::ExtractStridedMetadataOp metadaOp, OpAdaptor adaptor, + mlir::ConversionPatternRewriter &rewriter) const override { + + if(UnrealizedConversionCastOp castOp = mlir::cast_if_present(metadaOp.getSource().getDefiningOp())) + { + rewriter.replaceAllUsesWith(metadaOp.getResults(), castOp.getOperands()); + rewriter.eraseOp(metadaOp); + } + else + { + return failure(); + } + + return success(); + } +}; + +class ConvertBlockedGpuToTriton: public CometBlockedGpuToTritonBase { + + public: + ConvertBlockedGpuToTriton() = default; + + void runOnOperation() override + { + for(gpu::GPUModuleOp gpuModule : getOperation().getOps()) + { + RewritePatternSet patterns(&getContext()); + + mlir::comet::TritonTypeConverter converter(&getContext()); + mlir::comet::TritonConversionTarget target(getContext(), converter); + patterns.insert( converter, &getContext()); + patterns.insert< + ConvertExtractSlice, + ConvertInsertSlice, + ConvertBlockDim, + ConvertBlockId, + ConvertToTensor, + ConvertTensorSplatOp, + ConvertExtractStridedMetadata, + ConvertLinalgReduceOp, + ConvertLinalgBroadcastOp, + ConvertLinalgMatmulOp + >( &getContext()); + + if (failed(applyPartialConversion(gpuModule, target, std::move(patterns)))) + { + return signalPassFailure(); + } + + auto ttFuncs = getOperation().getOps(); + target.addIllegalOp(); + RewritePatternSet patternsClean(&getContext()); + patternsClean.insert(&getContext()); + + for(auto func: ttFuncs) + { + if (failed(applyPatternsAndFoldGreedily(func, std::move(patternsClean)))) + { + return signalPassFailure(); + } + + } + PassManager pm(&getContext()); + pm.addPass(mlir::createCanonicalizerPass()); + if (failed(pm.run(gpuModule))) { + signalPassFailure(); + return; + } + + RewritePatternSet patternsClean2(&getContext()); + target.addIllegalOp(); + patternsClean2.insert(&getContext()); + + if (failed(applyPartialConversion(gpuModule, target, std::move(patternsClean2)))) + { + return signalPassFailure(); + } + + RewritePatternSet patternsClean3(&getContext()); + target.addDynamicallyLegalOp([&](Operation *op) { + return llvm::all_of(op->getOperandTypes(), + [&](Type type) { return converter.isLegal(type); }) && + llvm::all_of(op->getResultTypes(), + [&](Type type) { return converter.isLegal(type); }); + }); + + patternsClean3.insert, ConvertArithIndex, ConvertArithIndex >(converter, &getContext()); + + if (failed(applyPartialConversion(gpuModule, target, std::move(patternsClean3)))) + { + return signalPassFailure(); + } + + RewritePatternSet patternsClean4(&getContext()); + target.addDynamicallyLegalOp([&](Operation *op) { + return llvm::all_of(op->getOperandTypes(), + [&](Type type) { if(RankedTensorType shapedType = mlir::dyn_cast(type)){return shapedType.hasStaticShape(); } return true; }) && + llvm::all_of(op->getResultTypes(), + [&](Type type) { if(RankedTensorType shapedType = mlir::dyn_cast(type)){return shapedType.hasStaticShape(); } return true; }); + }); + + patternsClean4.insert, ConvertArithDynamicShape, ConvertArithDynamicShape, ConvertArithDynamicShape, ConvertArithDynamicShape, ConvertArithDynamicShape >(&getContext()); + + if (failed(applyPartialConversion(gpuModule, target, std::move(patternsClean4)))) + { + return signalPassFailure(); + } + } + + } +}; + + + +std::unique_ptr> mlir::comet::createConvertBlockedGpuToTritonPass() { + return std::make_unique(); +} \ No newline at end of file diff --git a/lib/Conversion/BlockedGpuToTriton/BlockedGpuToTritonConversion.cpp b/lib/Conversion/BlockedGpuToTriton/BlockedGpuToTritonConversion.cpp new file mode 100644 index 00000000..6eb32f93 --- /dev/null +++ b/lib/Conversion/BlockedGpuToTriton/BlockedGpuToTritonConversion.cpp @@ -0,0 +1,106 @@ +#include "comet//Conversion/BlockedGpuToTriton/BlockedGpuToTritonConversion.h" +#include "mlir/Dialect/Affine/IR/AffineOps.h" +#include "mlir/Dialect/Arith/IR/Arith.h" +#include "mlir/Dialect/Bufferization/IR/Bufferization.h" +#include "mlir/Dialect/Func/IR/FuncOps.h" +#include "mlir/Dialect/GPU/IR/GPUDialect.h" +#include "mlir/Dialect/Linalg/IR/Linalg.h" +#include "mlir/Dialect/SCF/IR/SCF.h" +#include "mlir/Dialect/Tensor/IR/Tensor.h" +#include "mlir/IR/BuiltinOps.h" +#include "mlir/IR/BuiltinTypes.h" +#include "mlir/IR/Types.h" +#include "mlir/IR/Visitors.h" +#include "triton/Dialect/Triton/IR/Dialect.h" +#include "mlir/IR/BuiltinDialect.h" +#include "llvm/ADT/STLExtras.h" + +#include +#include + + +mlir::comet::TritonTypeConverter::TritonTypeConverter(MLIRContext *context) +: context(context) +{ + + addConversion([](mlir::Type t) { + return t; + }); + + addConversion([this] (mlir::MemRefType memrefType, SmallVectorImpl & convertedTypes){ + Type i32 = mlir::IntegerType::get(this->context, 32); + if(mlir::isa(memrefType.getElementType())) + { + convertedTypes.push_back(mlir::triton::PointerType::get(i32, 1)); + } + else + { + convertedTypes.push_back(mlir::triton::PointerType::get(memrefType.getElementType(), 1)); + } + + convertedTypes.push_back(i32); // offset + for(int64_t i = 0; i < memrefType.getRank(); i++) // sizes + { + convertedTypes.push_back(i32); + } + for(int64_t i = 0; i < memrefType.getRank(); i++) // strides + { + convertedTypes.push_back(i32); + } + + return success(); + }); + + addConversion([this] (mlir::IndexType indexType){ + { + return IntegerType::get(this->context, 32); + } + }); + + addConversion([this] (mlir::RankedTensorType rankedType){ + { + if(rankedType.getElementType().isIndex()) + { + return RankedTensorType::get(rankedType.getShape(), IntegerType::get(this->context, 32)); + } + else + { + return rankedType; + } + } + }); + + addSourceMaterialization([](OpBuilder &builder, Type type, ValueRange inputs, Location loc) -> std::optional { + if (inputs.size() != 1) + return std::nullopt; + return builder.create(loc, type, inputs)->getResult(0); + }); + + addTargetMaterialization([](OpBuilder &builder, Type type, ValueRange inputs, Location loc) -> std::optional { + if (inputs.size() != 1) + return std::nullopt; + return builder.create(loc, type, inputs)->getResult(0); + }); + + addArgumentMaterialization([](OpBuilder &builder, MemRefType memrefType, ValueRange values, Location loc)-> std::optional { + assert(values.size() == static_cast(memrefType.getRank()) * 2 + 2); + return builder.create(loc, memrefType, values)->getResult(0); + }); + + addArgumentMaterialization([](OpBuilder &builder, IndexType indexType, ValueRange values, Location loc)-> std::optional { + return builder.create(loc, indexType, values)->getResult(0); + }); +} + +mlir::comet::TritonConversionTarget::TritonConversionTarget(MLIRContext &context, TritonTypeConverter &typeConverter) + : ConversionTarget(context) +{ + addDynamicallyLegalOp( [&](Operation* op) { + // return true; + return llvm::all_of(op->getOperandTypes(), [&](Type type) { return typeConverter.isLegal(type); }); + }); + addLegalDialect(); + addLegalOp(); + addIllegalDialect(); + addIllegalOp(); +} \ No newline at end of file diff --git a/lib/Conversion/BlockedGpuToTriton/CMakeLists.txt b/lib/Conversion/BlockedGpuToTriton/CMakeLists.txt new file mode 100644 index 00000000..baeda955 --- /dev/null +++ b/lib/Conversion/BlockedGpuToTriton/CMakeLists.txt @@ -0,0 +1,7 @@ +add_llvm_library(COMETBlockedGpuToTriton + BlockedGpuToTriton.cpp + BlockedGpuToTritonConversion.cpp + + DEPENDS + BlockedGpuToTritonPassIncGen +) \ No newline at end of file diff --git a/lib/Conversion/CMakeLists.txt b/lib/Conversion/CMakeLists.txt index dec998e4..b853d454 100644 --- a/lib/Conversion/CMakeLists.txt +++ b/lib/Conversion/CMakeLists.txt @@ -1,8 +1,24 @@ add_subdirectory(IndexTreeToSCF) add_subdirectory(TensorAlgebraToIndexTree) add_subdirectory(TensorAlgebraToSCF) +if(ENABLE_FPGA_TARGET) +add_subdirectory(ParallelLoopsToGpuFPGA) +add_subdirectory(GpuToOCLSPIRV) +add_subdirectory(GpuHostToMCLRT) +endif() if(ENABLE_GPU_TARGET) -add_subdirectory(ParallelLoopsToGpu) -add_subdirectory(GpuToTriton) +add_subdirectory(ForallToGpu) +add_subdirectory(GpuToBlockedGpu) +add_subdirectory(BlockedGpuToTriton) +#add_subdirectory(GpuToTriton) +add_subdirectory(PrepareGpuHost) +add_subdirectory(GpuUtils) +endif() + +if(ENABLE_NVIDIA_GPU_BACKEND) add_subdirectory(TritonToCuda) +endif() + +if(ENABLE_AMD_GPU_BACKEND) +add_subdirectory(TritonToHIP) endif() \ No newline at end of file diff --git a/lib/Conversion/ForallToGpu/CMakeLists.txt b/lib/Conversion/ForallToGpu/CMakeLists.txt new file mode 100644 index 00000000..fc689776 --- /dev/null +++ b/lib/Conversion/ForallToGpu/CMakeLists.txt @@ -0,0 +1,9 @@ +add_llvm_library(COMETForallToGpu + ForallToGpu.cpp + + DEPENDS + ForallConversionPassIncGen + + # LINK_LIBS + # MLIRPASS +) \ No newline at end of file diff --git a/lib/Conversion/ForallToGpu/ForallToGpu.cpp b/lib/Conversion/ForallToGpu/ForallToGpu.cpp new file mode 100644 index 00000000..2310e645 --- /dev/null +++ b/lib/Conversion/ForallToGpu/ForallToGpu.cpp @@ -0,0 +1,660 @@ +// +// Copyright 2022 Battelle Memorial Institute +// +// Redistribution and use in source and binary forms, with or without modification, +// are permitted provided that the following conditions are met: +// +// 1. Redistributions of source code must retain the above copyright notice, this list of conditions +// and the following disclaimer. +// +// 2. Redistributions in binary form must reproduce the above copyright notice, this list of conditions +// and the following disclaimer in the documentation and/or other materials provided with the distribution. +// +// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED +// WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR +// PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, +// INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE +// GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, +// WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +// OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. +// + + +#include +#include +#include +#include +#include +#include +#include "mlir/Dialect/Arith/IR/Arith.h" +#include "mlir/IR/AffineExpr.h" +#include "mlir/IR/Attributes.h" +#include "mlir/IR/Block.h" +#include "mlir/IR/Builders.h" +#include "mlir/IR/BuiltinAttributes.h" +#include "mlir/IR/BuiltinTypes.h" +#include "mlir/IR/Location.h" +#include "mlir/IR/Operation.h" +#include "mlir/IR/PatternMatch.h" +#include "mlir/IR/Value.h" +#include "mlir/IR/ValueRange.h" +#include "mlir/Interfaces/SideEffectInterfaces.h" +#include "mlir/Pass/Pass.h" +#include "mlir/Dialect/DLTI/DLTI.h" +#include "mlir/Dialect/MemRef/IR/MemRef.h" +#include "mlir/Dialect/GPU/IR/GPUDialect.h" +#include "mlir/Dialect/SCF/IR/SCF.h" +#include "mlir/Dialect/Affine/IR/AffineOps.h" +#include "mlir/Support/LogicalResult.h" +#include "mlir/Transforms/DialectConversion.h" +#include "mlir/Transforms/GreedyPatternRewriteDriver.h" +#include "mlir/Pass/PassManager.h" +#include "mlir/Transforms/Passes.h" +#include "mlir/Dialect/SCF/Transforms/Transforms.h" + + +#include "comet/Conversion/ForallToGpu/ForallToGpu.h" + +#include "mlir/Dialect/Func/IR/FuncOps.h" +#include "llvm/ADT/ArrayRef.h" +#include "llvm/ADT/MapVector.h" +#include "llvm/ADT/STLExtras.h" +#include "llvm/ADT/SmallVector.h" +#include "llvm/Support/Casting.h" +#include "llvm/Support/raw_ostream.h" + +#define GEN_PASS_CLASSES +#include "comet/Conversion/ForallToGpu/Passes.h.inc" + +// *********** For debug purpose *********// +// #define COMET_DEBUG_MODE +#include "comet/Utils/debug.h" +#undef COMET_DEBUG_MODE +// *********** For debug purpose *********// + +namespace { + +bool contains_arg(mlir::Block& block, mlir::BlockArgument arg) +{ + auto store_ops = block.getOps(); + for(auto store_op: store_ops) + { + for(auto index: store_op.getIndices()) + { + if (auto affine_expr = llvm::dyn_cast_or_null(index.getDefiningOp())) + { + for(auto op: affine_expr.getOperands()) + { + if(op == arg) + { + // llvm::errs() << "operation: " << store_op << " contains: " << op << "\n"; + return true; + } + } + } + else if(auto block_arg = mlir::dyn_cast_if_present(index)) + { + if(block_arg == arg) + { + // llvm::errs() << "operation: " << store_op << " contains: " << op << "\n"; + return true; + } + } + else if(llvm::isa_and_present(index.getDefiningOp())) { + for(auto op: index.getDefiningOp()->getOperands()) + { + if(op == arg) + { + return true; + } + } + } + else { + llvm::errs() << "Load operation without affine expression\n"; + block.dump(); + index.dump(); + store_op->dump(); + exit(1); + } + } + } + + return false; +} + +bool is_reduction_(mlir::Region& region, mlir::BlockArgument arg) +{ + for (mlir::Block &block : region.getBlocks()) + { + if(contains_arg(block, arg)) + { + return false; + } + + for (mlir::Operation &innerOp : block) + { + for (mlir::Region &innerRegion : innerOp.getRegions()) + { + if(!is_reduction_(innerRegion, arg)) + { + return false; + } + } + } + } + + return true; +} + +bool is_reduction(mlir::scf::ForOp forOp) +{ + return is_reduction_(forOp.getBodyRegion(), forOp.getBody()->getArgument(0)); +} +using namespace mlir; + + +std::pair> tileParallelLoop(ConversionPatternRewriter& rewriter, scf::ParallelOp& op, ArrayRef tileSizes) { + rewriter.setInsertionPoint(op); + auto zero = rewriter.create(op.getLoc(), 0); + SmallVector tileSizeConstants; + tileSizeConstants.reserve(op.getUpperBound().size()); + for (size_t i = 0, end = op.getUpperBound().size(); i != end; ++i) { + if (tileSizes[i] != -1) + tileSizeConstants.push_back( + rewriter.create(op.getLoc(), tileSizes[i])); + else + // Just pick 1 for the remaining dimensions. + tileSizeConstants.push_back( + rewriter.create(op.getLoc(), 1)); + } + + // Create the outer loop with adjusted steps. + SmallVector newSteps; + newSteps.reserve(op.getStep().size()); + for (auto step : llvm::zip(op.getStep(), tileSizeConstants)) { + newSteps.push_back(rewriter.create(op.getLoc(), std::get<0>(step), + std::get<1>(step))); + } + auto outerLoop = rewriter.create(op.getLoc(), op.getLowerBound(), + op.getUpperBound(), newSteps); + rewriter.setInsertionPointToStart(outerLoop.getBody()); + + + + // Create the inner loop with adjusted bounds. + SmallVector newBounds; + newBounds.reserve(op.getUpperBound().size()); + bool needInboundCheck = false; + for (auto [lowerBound, upperBound, newStep, iv, step, tileSizeConstant] : + llvm::zip(outerLoop.getLowerBound(), outerLoop.getUpperBound(), + outerLoop.getStep(), outerLoop.getInductionVars(), + op.getStep(), tileSizeConstants)) + { + auto tileSize = + cast(tileSizeConstant.getDefiningOp()).value(); + // Compute min(size, dim - offset) to avoid out-of-bounds accesses. + auto minMap = AffineMap::get( + /*dimCount=*/2, /*symbolCount=*/0, + {getAffineConstantExpr(/*position=*/tileSize, rewriter.getContext()), + getAffineDimExpr(/*position=*/0, rewriter.getContext()) - + getAffineDimExpr(/*position=*/1, rewriter.getContext())}, + rewriter.getContext()); + // Collect the statically known loop bounds + auto lowerBoundConstant = + dyn_cast_or_null(lowerBound.getDefiningOp()); + auto upperBoundConstant = + dyn_cast_or_null(upperBound.getDefiningOp()); + auto stepConstant = + dyn_cast_or_null(step.getDefiningOp()); + + // If the loop bounds and the loop step are constant and if the number of + // loop iterations is an integer multiple of the tile size, we use a static + // bound for the inner loop. + if (lowerBoundConstant && upperBoundConstant && stepConstant) { + auto numIterations = llvm::divideCeil(upperBoundConstant.value() - + lowerBoundConstant.value(), + stepConstant.value()); + if (numIterations % tileSize == 0) { + newBounds.push_back(newStep); + continue; + } + } + + // Otherwise, we dynamically compute the bound for + // each iteration of the outer loop. + newBounds.push_back( + rewriter.create(op.getLoc(), rewriter.getIndexType(), minMap, + ValueRange{upperBound, iv})); + } + + SmallVector newArgs; + SmallVector innerLoops; + SmallVector inductionVars; + for(size_t i = 0; i < tileSizes.size(); i++) + { + if(tileSizes[i] == -1) + { + newArgs.push_back(outerLoop.getBody()->getArgument(i)); + + continue; + } + auto innerLoop = rewriter.create(op.getLoc(), zero, newBounds[i], op.getStep()[i]); + innerLoop->setAttr("blockSize", rewriter.getUI32IntegerAttr(tileSizes[i])); + innerLoops.push_back(innerLoop); + rewriter.setInsertionPointToStart(innerLoop.getBody()); + auto new_arg = rewriter.create(op.getLoc(), innerLoop.getBody()->getArgument(0), outerLoop.getBody()->getArgument(i)); + newArgs.push_back(new_arg); + inductionVars.push_back(innerLoop.getInductionVar()); + } + + rewriter.eraseOp(op.getBody()->getTerminator()); + rewriter.eraseOp(innerLoops.back().getBody()->getTerminator()); + rewriter.setInsertionPointToStart(innerLoops.back().getBody()); + + rewriter.mergeBlocks(op.getBody(), innerLoops.back().getBody(), newArgs); + rewriter.setInsertionPointToEnd(innerLoops.back().getBody()); + rewriter.create(op->getLoc()); + rewriter.eraseOp(op); + + return std::make_pair(outerLoop, innerLoops); +} + +std::pair tileForLoop(OpBuilder& builder, scf::ForOp& op, int64_t tileSize) { + builder.setInsertionPoint(op); + auto zero = builder.create(op.getLoc(), 0); + Value tileSizeConstant = builder.create(op.getLoc(), tileSize); + + // Create the outer loop with adjusted steps. + Value newStep = builder.create(op.getLoc(), op.getStep(), tileSizeConstant); + + auto outerLoop = builder.create(op.getLoc(), op.getLowerBound(), op.getUpperBound(), newStep, op.getInitArgs()); + outerLoop->setAttr("reduceDim", builder.getUnitAttr()); + builder.setInsertionPointToStart(outerLoop.getBody()); + + // Create the inner loop with adjusted bounds. + Value newBound; + auto tileSizeNew = + cast(tileSizeConstant.getDefiningOp()).value(); + // Compute min(size, dim - offset) to avoid out-of-bounds accesses. + auto minMap = AffineMap::get( + /*dimCount=*/2, /*symbolCount=*/0, + {getAffineConstantExpr(/*position=*/tileSizeNew, builder.getContext()), + getAffineDimExpr(/*position=*/0, builder.getContext()) - + getAffineDimExpr(/*position=*/1, builder.getContext())}, + builder.getContext()); + // Collect the statically known loop bounds + auto lowerBoundConstant = + dyn_cast_or_null(outerLoop.getLowerBound().getDefiningOp()); + auto upperBoundConstant = + dyn_cast_or_null(outerLoop.getUpperBound().getDefiningOp()); + auto stepConstant = + dyn_cast_or_null(op.getStep().getDefiningOp()); + + // If the loop bounds and the loop step are constant and if the number of + // loop iterations is an integer multiple of the tile size, we use a static + // bound for the inner loop. + if (lowerBoundConstant && upperBoundConstant && stepConstant) { + auto numIterations = llvm::divideCeil(upperBoundConstant.value() - + lowerBoundConstant.value(), + stepConstant.value()); + if (numIterations % tileSize == 0) { + newBound = outerLoop.getStep(); + } + } + + // Otherwise, we dynamically compute the bound for + // each iteration of the outer loop. + newBound = builder.create(op.getLoc(), builder.getIndexType(), minMap, ValueRange{outerLoop.getUpperBound(), outerLoop.getInductionVar()}); + + SmallVector newArgs; + SmallVector innerLoops; + SmallVector inductionVars; + auto innerLoop = builder.create(op.getLoc(), zero, newBound, op.getStep(), outerLoop.getRegionIterArgs()); + innerLoop->setAttr("blockSize", builder.getUI32IntegerAttr(tileSize)); + innerLoops.push_back(innerLoop); + builder.setInsertionPointToStart(innerLoop.getBody()); + auto new_arg = builder.create(op.getLoc(), innerLoop.getBody()->getArgument(0), outerLoop.getBody()->getArgument(0)); + newArgs.push_back(new_arg); + newArgs.insert(newArgs.end(), innerLoop.getRegionIterArgs().begin(), innerLoop.getRegionIterArgs().end()); + inductionVars.push_back(innerLoop.getInductionVar()); + builder.setInsertionPointToEnd(innerLoops.back().getBody()); + auto temp_yieldOp = builder.create(op->getLoc()); + + builder.setInsertionPointToEnd(outerLoop.getBody()); + builder.create(op->getLoc(), innerLoop.getResults()); + + for(auto& inner_op: llvm::make_early_inc_range(op.getBody()->getOperations())) + { + inner_op.moveBefore(innerLoops.back().getBody()->getTerminator()); + } + temp_yieldOp->erase(); + for(auto [old_arg, new_arg] : llvm::zip(op.getBody()->getArguments(), newArgs) ) + { + old_arg.replaceAllUsesWith(new_arg); + } + op->replaceAllUsesWith(outerLoop); + op->erase(); + + return std::make_pair(outerLoop, innerLoop); +} + +class ForallOpToGpu: public mlir::OpConversionPattern { +private: + int blockX, blockY, blockR; +public: + using mlir::OpConversionPattern::OpConversionPattern; + ForallOpToGpu(mlir::MLIRContext* ctx, int blockX, int blockY, int blockR) : mlir::OpConversionPattern(ctx), blockX(blockX), blockY(blockY), blockR(blockR) {} + mlir::LogicalResult + matchAndRewrite(mlir::scf::ForallOp forAllOp, OpAdaptor adaptor, + mlir::ConversionPatternRewriter &rewriter) const override { + + SmallVector worklist; + scf::ForallOp otherforAllOp = NULL; + for(auto& op: forAllOp.getBody()->without_terminator()) + { + if(auto inner_forAllOp = mlir::dyn_cast(op)) + { + otherforAllOp = inner_forAllOp; + break; + } + else + { + worklist.push_back(&op); + } + } + SmallVector canMove; + if(otherforAllOp != NULL) + { + for(auto op: worklist) + { + if(op->getRegions().size() == 0) + { + canMove.push_back(op); + } + } + } + + scf::ParallelOp parOp = nullptr; + if(worklist.size() == canMove.size() && otherforAllOp) + { + // if(!canMove.empty()) + { + rewriter.setInsertionPointToStart(otherforAllOp.getBody()); + for(auto op: canMove) + { + auto clone = rewriter.clone(*op); + rewriter.replaceAllOpUsesWith(op, clone); + rewriter.eraseOp(op); + } + rewriter.setInsertionPoint(forAllOp); + SmallVector lbs = forAllOp.getLowerBound(rewriter); + SmallVector ubs = forAllOp.getUpperBound(rewriter); + SmallVector steps = forAllOp.getStep(rewriter); + SmallVector otherlbs = otherforAllOp.getLowerBound(rewriter); + SmallVector otherubs = otherforAllOp.getUpperBound(rewriter); + SmallVector othersteps = otherforAllOp.getStep(rewriter); + + std::vector attrs; + if(forAllOp->hasAttr("parallelDim")) + { + attrs.push_back(forAllOp->getAttrOfType("parallelDim")); + } + if(otherforAllOp->hasAttrOfType("parallelDim")) + { + attrs.push_back(otherforAllOp->getAttrOfType("parallelDim")); + } + // auto attrsAttr = ArrayRef(attrs); + auto arrayAttr = rewriter.getNamedAttr("parallelDim", rewriter.getArrayAttr(ArrayRef(attrs))); + + auto combinedParOp = rewriter.create(forAllOp->getLoc(), ValueRange({lbs.front(), otherlbs.front()}), ValueRange({ubs.front(), otherubs.front()}), ValueRange({steps.front(), othersteps.front()})); + if(attrs.size() > 0) + { + combinedParOp->setAttrs(arrayAttr); + } + // for(auto attr: attrs) + // { + // combinedParOp->setAttr("parallelDim", attr); + // } + Operation* terminator = combinedParOp.getBody()->getTerminator(); + rewriter.eraseOp(terminator); + terminator = forAllOp.getBody()->getTerminator(); + rewriter.eraseOp(terminator); + rewriter.mergeBlocks(forAllOp.getBody(), combinedParOp.getBody(), combinedParOp.getInductionVars().front()); + rewriter.eraseOp(forAllOp); + terminator = otherforAllOp.getBody()->getTerminator(); + rewriter.mergeBlocks(otherforAllOp.getBody(), combinedParOp.getBody(), combinedParOp.getInductionVars().back()); + rewriter.eraseOp(terminator); + rewriter.eraseOp(otherforAllOp); + parOp = combinedParOp; + rewriter.setInsertionPointToEnd(parOp.getBody()); + rewriter.create(parOp->getLoc()); + } + } + else + { + SmallVector lbs = forAllOp.getLowerBound(rewriter); + SmallVector ubs = forAllOp.getUpperBound(rewriter); + SmallVector steps = forAllOp.getStep(rewriter); + std::vector attrs; + if(forAllOp->hasAttr("parallelDim")) + { + attrs.push_back(forAllOp->getAttrOfType("parallelDim")); + } + auto arrayAttr = rewriter.getNamedAttr("parallelDim", rewriter.getArrayAttr(ArrayRef(attrs))); + + parOp = rewriter.create(forAllOp->getLoc(), lbs, ubs, steps); + if(attrs.size() > 0) + { + parOp->setAttrs(arrayAttr); + } + Operation* terminator = parOp.getBody()->getTerminator(); + rewriter.eraseOp(terminator); + terminator = forAllOp.getBody()->getTerminator(); + rewriter.mergeBlocks(forAllOp.getBody(), parOp.getBody(), parOp.getInductionVars()); + rewriter.eraseOp(terminator); + rewriter.eraseOp(forAllOp); + rewriter.setInsertionPointToEnd(parOp.getBody()); + rewriter.create(parOp->getLoc()); + } + + llvm::SmallVector allTileSizes = {blockY, blockX}; + llvm::SmallVector tileSizes; + std::copy(allTileSizes.begin(), allTileSizes.begin() + parOp.getInductionVars().size(), std::back_inserter(tileSizes)); + SmallVector allStringAttrs = {rewriter.getAttr("dimY_grid"), rewriter.getAttr("dimX_grid")}; + SmallVector stringAttrs; + std::copy(allStringAttrs.begin(), allStringAttrs.begin() + parOp.getInductionVars().size(), std::back_inserter(stringAttrs)); + + auto dim0 = rewriter.getAffineDimExpr(0); + auto dim1 = rewriter.getAffineDimExpr(0); + auto yMap = mlir::AffineMap::get(1, 0, dim0); + auto xMap = mlir::AffineMap::get(1, 0, dim1); + SmallVector allGpuAttrs = {mlir::gpu::ParallelLoopDimMappingAttr::get(rewriter.getContext(), ::mlir::gpu::Processor::BlockY, yMap, yMap), mlir::gpu::ParallelLoopDimMappingAttr::get(rewriter.getContext(), ::mlir::gpu::Processor::BlockX, xMap, xMap)}; + SmallVector gpuAttrs; + std::copy(allGpuAttrs.begin(), allGpuAttrs.begin() + parOp.getInductionVars().size(), std::back_inserter(gpuAttrs)); + if(!parOp->hasAttr("parallelDim")) + { + + auto tiledLoop = tileParallelLoop(rewriter, parOp, tileSizes); + tiledLoop.first->setAttr("mapping", rewriter.getArrayAttr(gpuAttrs)); + tiledLoop.first->setAttr("parallelDim", rewriter.getArrayAttr(stringAttrs)); + } + else + { + auto parallelDimAttr = mlir::cast(parOp->getAttr("parallelDim")); + if(parallelDimAttr.size() == parOp.getInductionVars().size()) + { + rewriter.modifyOpInPlace(parOp, [&]() { + parOp->setAttr("parallelDim", rewriter.getArrayAttr(stringAttrs)); + parOp->setAttr("mapping", rewriter.getArrayAttr(gpuAttrs)); + }); + } + else + { + if(mlir::cast(parallelDimAttr[0]).getValue() == "dimY_grid") + { + tileSizes[0] = -1; + } + else if(mlir::cast(parallelDimAttr[0]).getValue() == "dimX_grid") + { + tileSizes[1] = -1; + } + + auto tiledLoop = tileParallelLoop(rewriter, parOp, tileSizes); + tiledLoop.first->setAttr("mapping", rewriter.getArrayAttr(gpuAttrs)); + tiledLoop.first->setAttr("parallelDim", rewriter.getArrayAttr(stringAttrs)); + } + } + + return success(); + } +}; +class ConvertForallToGpu: public CometForallToGpuBase { +public: + ConvertForallToGpu() = default; + ConvertForallToGpu(int blockX, int blockY, int blockR) { + this->blockX = blockX; + this->blockY = blockY; + this->blockR = blockR; + } + + void runOnOperation() override { + mlir::MLIRContext *context = &getContext(); + mlir::func::FuncOp funcOp = getOperation(); + if (funcOp.isDeclaration()) + { + return; + } + + mlir::RewritePatternSet patterns(context); + patterns.insert(context, blockX, blockY, blockR); + + mlir::ConversionTarget target(*context); + target.addLegalDialect(); + target.addLegalOp(); + target.addIllegalOp(); + // target.addDynamicallyLegalOp([] (mlir::scf::ForallOp op) { + // return !op->hasAttr("parallelDim"); + // }); + + target.addDynamicallyLegalOp([](mlir::scf::ParallelOp op) -> bool { + return op->hasAttr("mapping") && op->hasAttr("parallelDim"); + }); + + if (mlir::failed(mlir::applyPartialConversion(funcOp, target, std::move(patterns)))) + { + return signalPassFailure(); + } + + + llvm::SmallVector toReduceForOps; + funcOp->walk([&toReduceForOps](mlir::scf::ForOp forOp) { + if(forOp->getParentOfType()) + { + if(is_reduction(forOp)) + { + toReduceForOps.push_back(forOp); + } + } + }); + + OpBuilder builder(funcOp); + + for(scf::ForOp forOp: llvm::make_early_inc_range(toReduceForOps)) + { + llvm::SmallVector inductionVars; + llvm::SmallVector loopInvMemOps; + inductionVars.push_back(forOp.getInductionVar()); + + forOp->walk([&loopInvMemOps](Operation* op){ + if(mlir::isa(op)) + { + loopInvMemOps.push_back(op); + } + }); + + + for(auto inductionVar: inductionVars) + { + for(auto user: llvm::make_early_inc_range(inductionVar.getUsers())) + { + if(mlir::isa(user)) + { + auto it = std::find(loopInvMemOps.begin(), loopInvMemOps.end(), user); + loopInvMemOps.erase(it); + } + for(auto res: user->getResults()) + { + inductionVars.push_back(res); + } + } + } + + llvm::SmallMapVector, 4> loadStorePairs; + for(auto memOp: loopInvMemOps) + { + if(memref::StoreOp storeOp = dyn_cast(memOp)) + { + auto it = loadStorePairs.find(storeOp.getMemRef()); + assert(it!= loadStorePairs.end()); + it->second.push_back(storeOp); + } + else if(memref::LoadOp loadOp = dyn_cast(memOp)) + { + auto it = loadStorePairs.find(loadOp.getMemRef()); + assert(it == loadStorePairs.end()); + loadStorePairs[loadOp.getMemRef()].push_back(loadOp); + // it->second.push_back(loadOp); + } + else + { + assert(false && "UNREACHABLE. Should vectore should only contain store or load operations"); + } + } + + SmallVector iterArgs, yieldOps; + SmallVector storeOps; + for(auto pair: loadStorePairs) + { + if(pair.second.size() == 2) + { + pair.second[0]->moveBefore(forOp); + pair.second[1]->moveAfter(forOp); + iterArgs.push_back(pair.second[0]->getResult(0)); + yieldOps.push_back(pair.second[1]->getOperand(0)); + storeOps.push_back(pair.second[1]); + } + } + builder.setInsertionPoint(forOp); + auto newForOp = builder.create(forOp->getLoc(), forOp.getLowerBound(), forOp.getUpperBound(), forOp.getStep(), ValueRange(iterArgs)); + // newForOp->dump(); + builder.setInsertionPointToEnd(newForOp.getBody()); + builder.create(forOp->getLoc(), yieldOps); + for(auto iterArg: iterArgs) { + iterArg.replaceAllUsesExcept(newForOp.getRegionIterArg(0), newForOp); + } + for(auto storeOp: storeOps) { + storeOp->setOperand(0, newForOp.getResult(0)); + } + forOp.getBody()->getTerminator()->erase(); + + for(auto& op: llvm::make_early_inc_range(forOp.getBody()->getOperations())) + { + op.moveBefore(newForOp.getBody()->getTerminator()); + } + forOp.getInductionVar().replaceAllUsesWith(newForOp.getInductionVar()); + auto [outerLoop, innerLoop] = tileForLoop(builder, newForOp, blockR); + forOp->erase(); + } + } +}; +} + + +std::unique_ptr> mlir::comet::createConvertForallToGpuPass() { + return std::make_unique(); +} + + +std::unique_ptr> mlir::comet::createConvertForallToGpuPass(int blockX, int blockY, int blockR) { + return std::make_unique(blockX, blockY, blockR); +} diff --git a/lib/Conversion/ForallToGpu/try.mlir b/lib/Conversion/ForallToGpu/try.mlir new file mode 100644 index 00000000..a56f8300 --- /dev/null +++ b/lib/Conversion/ForallToGpu/try.mlir @@ -0,0 +1,167 @@ +"func.func"() <{function_type = () -> (), sym_name = "main"}> ({ + %0 = "arith.constant"() <{value = 1.700000e+00 : f64}> : () -> f64 + %1 = "arith.constant"() <{value = 0.000000e+00 : f64}> : () -> f64 + %2 = "arith.constant"() <{value = 0 : i64}> : () -> i64 + %3 = "arith.constant"() <{value = 10 : index}> : () -> index + %4 = "arith.constant"() <{value = 9 : index}> : () -> index + %5 = "arith.constant"() <{value = 8 : index}> : () -> index + %6 = "arith.constant"() <{value = 7 : index}> : () -> index + %7 = "arith.constant"() <{value = 6 : index}> : () -> index + %8 = "arith.constant"() <{value = 5 : index}> : () -> index + %9 = "arith.constant"() <{value = 4 : index}> : () -> index + %10 = "arith.constant"() <{value = 1 : i32}> : () -> i32 + %11 = "arith.constant"() <{value = 0 : i32}> : () -> i32 + %12 = "arith.constant"() <{value = 3 : index}> : () -> index + %13 = "arith.constant"() <{value = 2 : index}> : () -> index + %14 = "arith.constant"() <{value = -1 : index}> : () -> index + %15 = "arith.constant"() <{value = 1 : index}> : () -> index + %16 = "arith.constant"() <{value = 0 : index}> : () -> index + %17 = "memref.alloc"() <{operandSegmentSizes = array}> : () -> memref<13xindex> + %18 = "memref.cast"(%17) : (memref<13xindex>) -> memref<*xindex> + "func.call"(%11, %16, %14, %15, %14, %18, %10) <{callee = @read_input_sizes_2D_f64}> {filename = "SPARSE_FILE_NAME0"} : (i32, index, index, index, index, memref<*xindex>, i32) -> () + %19 = "memref.load"(%17, %16) <{nontemporal = false}> : (memref<13xindex>, index) -> index + %20 = "memref.load"(%17, %15) <{nontemporal = false}> : (memref<13xindex>, index) -> index + %21 = "memref.load"(%17, %13) <{nontemporal = false}> : (memref<13xindex>, index) -> index + %22 = "memref.load"(%17, %12) <{nontemporal = false}> : (memref<13xindex>, index) -> index + %23 = "memref.load"(%17, %9) <{nontemporal = false}> : (memref<13xindex>, index) -> index + %24 = "memref.load"(%17, %8) <{nontemporal = false}> : (memref<13xindex>, index) -> index + %25 = "memref.load"(%17, %7) <{nontemporal = false}> : (memref<13xindex>, index) -> index + %26 = "memref.load"(%17, %6) <{nontemporal = false}> : (memref<13xindex>, index) -> index + %27 = "memref.load"(%17, %5) <{nontemporal = false}> : (memref<13xindex>, index) -> index + %28 = "memref.load"(%17, %4) <{nontemporal = false}> : (memref<13xindex>, index) -> index + %29 = "memref.load"(%17, %3) <{nontemporal = false}> : (memref<13xindex>, index) -> index + %30 = "memref.alloc"(%19) <{operandSegmentSizes = array}> : (index) -> memref + "scf.for"(%16, %19, %15) ({ + ^bb0(%arg19: index): + "memref.store"(%2, %30, %arg19) <{nontemporal = false}> : (i64, memref, index) -> () + "scf.yield"() : () -> () + }) : (index, index, index) -> () + %31 = "memref.cast"(%30) : (memref) -> memref<*xi64> + %32 = "memref.alloc"(%20) <{operandSegmentSizes = array}> : (index) -> memref + "scf.for"(%16, %20, %15) ({ + ^bb0(%arg18: index): + "memref.store"(%2, %32, %arg18) <{nontemporal = false}> : (i64, memref, index) -> () + "scf.yield"() : () -> () + }) : (index, index, index) -> () + %33 = "memref.cast"(%32) : (memref) -> memref<*xi64> + %34 = "memref.alloc"(%21) <{operandSegmentSizes = array}> : (index) -> memref + "scf.for"(%16, %21, %15) ({ + ^bb0(%arg17: index): + "memref.store"(%2, %34, %arg17) <{nontemporal = false}> : (i64, memref, index) -> () + "scf.yield"() : () -> () + }) : (index, index, index) -> () + %35 = "memref.cast"(%34) : (memref) -> memref<*xi64> + %36 = "memref.alloc"(%22) <{operandSegmentSizes = array}> : (index) -> memref + "scf.for"(%16, %22, %15) ({ + ^bb0(%arg16: index): + "memref.store"(%2, %36, %arg16) <{nontemporal = false}> : (i64, memref, index) -> () + "scf.yield"() : () -> () + }) : (index, index, index) -> () + %37 = "memref.cast"(%36) : (memref) -> memref<*xi64> + %38 = "memref.alloc"(%23) <{operandSegmentSizes = array}> : (index) -> memref + "scf.for"(%16, %23, %15) ({ + ^bb0(%arg15: index): + "memref.store"(%2, %38, %arg15) <{nontemporal = false}> : (i64, memref, index) -> () + "scf.yield"() : () -> () + }) : (index, index, index) -> () + %39 = "memref.cast"(%38) : (memref) -> memref<*xi64> + %40 = "memref.alloc"(%24) <{operandSegmentSizes = array}> : (index) -> memref + "scf.for"(%16, %24, %15) ({ + ^bb0(%arg14: index): + "memref.store"(%2, %40, %arg14) <{nontemporal = false}> : (i64, memref, index) -> () + "scf.yield"() : () -> () + }) : (index, index, index) -> () + %41 = "memref.cast"(%40) : (memref) -> memref<*xi64> + %42 = "memref.alloc"(%25) <{operandSegmentSizes = array}> : (index) -> memref + "scf.for"(%16, %25, %15) ({ + ^bb0(%arg13: index): + "memref.store"(%2, %42, %arg13) <{nontemporal = false}> : (i64, memref, index) -> () + "scf.yield"() : () -> () + }) : (index, index, index) -> () + %43 = "memref.cast"(%42) : (memref) -> memref<*xi64> + %44 = "memref.alloc"(%26) <{operandSegmentSizes = array}> : (index) -> memref + "scf.for"(%16, %26, %15) ({ + ^bb0(%arg12: index): + "memref.store"(%2, %44, %arg12) <{nontemporal = false}> : (i64, memref, index) -> () + "scf.yield"() : () -> () + }) : (index, index, index) -> () + %45 = "memref.cast"(%44) : (memref) -> memref<*xi64> + %46 = "memref.alloc"(%27) <{operandSegmentSizes = array}> : (index) -> memref + "scf.for"(%16, %27, %15) ({ + ^bb0(%arg11: index): + "memref.store"(%1, %46, %arg11) <{nontemporal = false}> : (f64, memref, index) -> () + "scf.yield"() : () -> () + }) : (index, index, index) -> () + %47 = "memref.cast"(%46) : (memref) -> memref<*xf64> + "func.call"(%11, %16, %14, %15, %14, %31, %33, %35, %37, %39, %41, %43, %45, %47, %10) <{callee = @read_input_2D_f64_i64}> {filename = "SPARSE_FILE_NAME0"} : (i32, index, index, index, index, memref<*xi64>, memref<*xi64>, memref<*xi64>, memref<*xi64>, memref<*xi64>, memref<*xi64>, memref<*xi64>, memref<*xi64>, memref<*xf64>, i32) -> () + %48 = "memref.alloc"(%29) <{alignment = 32 : i64, operandSegmentSizes = array}> : (index) -> memref + "scf.for"(%16, %29, %15) ({ + ^bb0(%arg9: index): + "scf.for"(%16, %9, %15) ({ + ^bb0(%arg10: index): + "memref.store"(%0, %48, %arg9, %arg10) <{nontemporal = false}> : (f64, memref, index, index) -> () + "scf.yield"() : () -> () + }) : (index, index, index) -> () + "scf.yield"() : () -> () + }) : (index, index, index) -> () + %49 = "memref.alloc"(%28) <{alignment = 32 : i64, operandSegmentSizes = array}> : (index) -> memref + "scf.for"(%16, %28, %15) ({ + ^bb0(%arg7: index): + "scf.for"(%16, %9, %15) ({ + ^bb0(%arg8: index): + "memref.store"(%1, %49, %arg7, %arg8) <{nontemporal = false}> : (f64, memref, index, index) -> () + "scf.yield"() : () -> () + }) : (index, index, index) -> () + "scf.yield"() : () -> () + }) : (index, index, index) -> () + %50 = "arith.constant"() <{value = 0 : index}> : () -> index + %51 = "arith.constant"() <{value = 1 : index}> : () -> index + %52 = "arith.constant"() <{value = 0 : index}> : () -> index + %53 = "arith.constant"() <{value = 4 : index}> : () -> index + %54 = "arith.constant"() <{value = 1 : index}> : () -> index + %55 = "arith.constant"() <{value = 0 : index}> : () -> index + %56 = "arith.constant"() <{value = 1 : index}> : () -> index + %57 = "arith.constant"() <{value = 32 : index}> : () -> index + %58 = "arith.muli"(%51, %56) <{overflowFlags = #arith.overflow}> : (index, index) -> index + %59 = "arith.muli"(%54, %57) <{overflowFlags = #arith.overflow}> : (index, index) -> index + "scf.parallel"(%50, %52, %28, %53, %58, %59) <{operandSegmentSizes = array}> ({ + ^bb0(%arg0: index, %arg1: index): + %61 = "affine.min"(%28, %arg0) <{map = affine_map<(d0, d1) -> (1, d0 - d1)>}> : (index, index) -> index + %62 = "affine.min"(%53, %arg1) <{map = affine_map<(d0, d1) -> (32, d0 - d1)>}> : (index, index) -> index + "scf.for"(%55, %62, %54) ({ + ^bb0(%arg2: index): + %63 = "arith.addi"(%arg2, %arg1) <{overflowFlags = #arith.overflow}> : (index, index) -> index + %64 = "arith.index_cast"(%65) : (i64) -> index + %65 = "memref.load"(%38, %68) <{nontemporal = false}> : (memref, index) -> i64 + %66 = "arith.index_cast"(%67) : (i64) -> index + %67 = "memref.load"(%38, %arg0) <{nontemporal = false}> : (memref, index) -> i64 + %68 = "arith.addi"(%arg0, %15) <{overflowFlags = #arith.overflow}> : (index, index) -> index + %69 = "memref.load"(%49, %arg0, %63) <{nontemporal = false}> : (memref, index, index) -> f64 + %70 = "arith.constant"() <{value = 0 : index}> : () -> index + %71 = "arith.constant"() <{value = 32 : index}> : () -> index + %72 = "arith.muli"(%15, %71) <{overflowFlags = #arith.overflow}> : (index, index) -> index + %73 = "scf.for"(%66, %64, %72, %69) ({ + ^bb0(%arg3: index, %arg4: f64): + %74 = "affine.min"(%64, %arg3) <{map = affine_map<(d0, d1) -> (32, d0 - d1)>}> : (index, index) -> index + %75 = "scf.for"(%70, %74, %15, %arg4) ({ + ^bb0(%arg5: index, %arg6: f64): + %76 = "arith.addi"(%arg5, %arg3) <{overflowFlags = #arith.overflow}> : (index, index) -> index + %77 = "memref.load"(%40, %76) <{nontemporal = false}> : (memref, index) -> i64 + %78 = "arith.index_cast"(%77) : (i64) -> index + %79 = "memref.load"(%46, %76) <{nontemporal = false}> : (memref, index) -> f64 + %80 = "memref.load"(%48, %78, %63) <{nontemporal = false}> : (memref, index, index) -> f64 + %81 = "arith.mulf"(%79, %80) <{fastmath = #arith.fastmath}> : (f64, f64) -> f64 + %82 = "arith.addf"(%arg6, %81) <{fastmath = #arith.fastmath}> : (f64, f64) -> f64 + "scf.yield"(%82) : (f64) -> () + }) {blockSize = 32 : ui32} : (index, index, index, f64) -> f64 + "scf.yield"(%75) : (f64) -> () + }) {reduceDim} : (index, index, index, f64) -> f64 + "memref.store"(%73, %49, %arg0, %63) <{nontemporal = false}> : (f64, memref, index, index) -> () + "scf.yield"() : () -> () + }) {blockSize = 32 : ui32} : (index, index, index) -> () + "scf.reduce"() : () -> () + }) {mapping = [#gpu.loop_dim_map (d0), bound = (d0) -> (d0)>, #gpu.loop_dim_map (d0), bound = (d0) -> (d0)>], parallelDim = ["dimY_grid", "dimX_grid"]} : (index, index, index, index, index, index) -> () + %60 = "memref.cast"(%49) : (memref) -> memref<*xf64> + "func.call"(%60) <{callee = @comet_print_memref_f64}> : (memref<*xf64>) -> () + "func.return"() : () -> () +}) : () -> () \ No newline at end of file diff --git a/lib/Conversion/GpuHostToMCLRT/CMakeLists.txt b/lib/Conversion/GpuHostToMCLRT/CMakeLists.txt new file mode 100644 index 00000000..bc76fafd --- /dev/null +++ b/lib/Conversion/GpuHostToMCLRT/CMakeLists.txt @@ -0,0 +1,6 @@ +add_llvm_library(COMETGpuHostToMCLRT + GpuHostToMCLRTPass.cpp + + DEPENDS + GpuHostToMCLRTConversionPassIncGen +) \ No newline at end of file diff --git a/lib/Conversion/GpuHostToMCLRT/GpuHostToMCLRTPass.cpp b/lib/Conversion/GpuHostToMCLRT/GpuHostToMCLRTPass.cpp new file mode 100644 index 00000000..cfd16344 --- /dev/null +++ b/lib/Conversion/GpuHostToMCLRT/GpuHostToMCLRTPass.cpp @@ -0,0 +1,346 @@ + +#include +#include +#include +#include + +#include "comet/Conversion/GpuHostToMCLRT/GpuHostToMCLRTPass.h" +#include "mlir/Dialect/LLVMIR/LLVMDialect.h" +#include "mlir/Dialect/MemRef/IR/MemRef.h" +#include "mlir/Dialect/SCF/IR/SCF.h" +#include "mlir/Dialect/SPIRV/IR/SPIRVDialect.h" +#include "mlir/Dialect/GPU/IR/GPUDialect.h" +#include "mlir/Dialect/Index/IR/IndexDialect.h" +#include "mlir/Dialect/Index/IR/IndexOps.h" +#include "mlir/Dialect/Func/IR/FuncOps.h" +#include "mlir/IR/BuiltinOps.h" +#include "mlir/Pass/Pass.h" +#include "minos.h" + +#define GEN_PASS_CLASSES +#include "comet/Conversion/GpuHostToMCLRT/Passes.h.inc" + + +static constexpr const char *kMCLInitFname = "mcl_init"; +static constexpr const char *kMCLFinitFname = "mcl_finit"; +static constexpr const char *kMCLPrgLoadFname = "mcl_prg_load"; +static constexpr const char *kMCLCreateTaskFname = "mcl_task_create"; +static constexpr const char *kMCLSetKernelFname = "mcl_task_set_kernel"; +static constexpr const char *kMCLSetArgFname = "mcl_task_set_arg"; +static constexpr const char *kMCLExecFname = "mcl_exec"; +static constexpr const char *kMCLWaitFname = "mcl_wait"; +static constexpr const char *kMCLWaitAllFname = "mcl_wait_all"; + +mlir::func::FuncOp declare_function(mlir::MLIRContext* ctx, mlir::OpBuilder& builder, const char* name, mlir::TypeRange inputs, mlir::TypeRange outputs) { + + mlir::func::FuncOp func = builder.create(builder.getUnknownLoc(), name, builder.getFunctionType(inputs, outputs)); + func.setVisibility(mlir::SymbolTable::Visibility::Private); + + return func; +} + + +void declare_mcl_funcs(mlir::MLIRContext* ctx, mlir::OpBuilder& builder) { + mlir::IntegerType u64 = mlir::IntegerType::get(ctx, 64); + mlir::IntegerType i32 = mlir::IntegerType::get(ctx, 32); + mlir::LLVM::LLVMPointerType ptr = mlir::LLVM::LLVMPointerType::get(ctx); + mlir::IndexType index = mlir::IndexType::get(ctx); + + declare_function(ctx, builder, kMCLInitFname, {u64, u64}, {i32}); + declare_function(ctx, builder, kMCLFinitFname, {}, {i32}); + declare_function(ctx, builder, kMCLCreateTaskFname, {}, {ptr}); + declare_function(ctx, builder, kMCLPrgLoadFname, {ptr, ptr, u64}, {i32}); + declare_function(ctx, builder, kMCLSetKernelFname, {ptr, ptr, u64}, {i32}); + declare_function(ctx, builder, kMCLSetArgFname, {ptr, u64, ptr, index, u64}, {i32}); + declare_function(ctx, builder, kMCLExecFname, {ptr, ptr, ptr, u64}, {i32}); + declare_function(ctx, builder, kMCLWaitFname, {ptr}, {i32}); + declare_function(ctx, builder, kMCLWaitAllFname, {}, {i32}); +} + +void insert_mcl_init(mlir::MLIRContext* ctx, mlir::OpBuilder& builder, mlir::Location loc, mlir::Value num_workers, mlir::Value flags) +{ + mlir::IntegerType i32 = mlir::IntegerType::get(ctx, 32); + builder.create(loc, kMCLInitFname, mlir::TypeRange({i32}), mlir::ValueRange({num_workers, flags})); +} + +void insert_mcl_finit(mlir::MLIRContext* ctx, mlir::OpBuilder& builder, mlir::Location loc) +{ + mlir::IntegerType i32 = mlir::IntegerType::get(ctx, 32); + builder.create(loc, kMCLFinitFname, mlir::TypeRange({i32}), mlir::ValueRange()); +} + +mlir::Value insert_task_create(mlir::MLIRContext* ctx, mlir::OpBuilder& builder, mlir::Location loc) +{ + mlir::LLVM::LLVMPointerType ptr = mlir::LLVM::LLVMPointerType::get(ctx); + return builder.create(loc, kMCLCreateTaskFname, mlir::TypeRange({ptr}), mlir::ValueRange()).getResult(0); +} + +mlir::Value insert_prg_load_call(mlir::MLIRContext* ctx, mlir::OpBuilder& builder, mlir::Location loc, const std::string &prgName, mlir::Value flags) +{ + static int prg_n = 0; + + mlir::IntegerType i32 = mlir::IntegerType::get(ctx, 32); + mlir::IntegerType i8 = mlir::IntegerType::get(ctx, 8); + mlir::LLVM::LLVMArrayType prgNameArrType = mlir::LLVM::LLVMArrayType::get(i8, prgName.size() + 1); + mlir::LLVM::LLVMArrayType copsType = mlir::LLVM::LLVMArrayType::get(i8, 1); + + auto prev = builder.saveInsertionPoint(); + builder.setInsertionPointToStart(builder.getInsertionPoint()->getParentOfType().getBody()); + + auto xclbinpath = builder.create(loc, prgNameArrType, true, mlir::LLVM::Linkage::Internal, "prg_"+std::to_string(prg_n)+"_path", mlir::StringAttr::get(prgName + '\0', i8)); + // mlir::LLVM::createGlobalString(loc, builder, "prg_"+std::to_string(prg_n)+"_path", prgName, mlir::LLVM::linkage::Linkage::Private ); + // mlir::LLVM::createGlobalString(loc, builder, "cops", "", mlir::LLVM::linkage::Linkage::Private ); + auto cops_global = builder.create(loc, copsType, true, mlir::LLVM::Linkage::Internal, "cops", mlir::StringAttr::get(std::string() + '\0', i8)); + builder.restoreInsertionPoint(prev); + mlir::Value prg_path = builder.create(loc, xclbinpath); + mlir::Value cops = builder.create(loc, cops_global); + prg_n++; + return builder.create(loc, kMCLPrgLoadFname, mlir::TypeRange({i32}), mlir::ValueRange({prg_path, cops, flags})).getResult(0); +} + +mlir::Value insert_set_task_kernel_call(mlir::MLIRContext* ctx, mlir::OpBuilder& builder, mlir::Location loc, mlir::gpu::LaunchFuncOp launchOp, mlir::Value mcl_handle) +{ + static int kernel_n = 0; + mlir::IntegerType i8 = mlir::IntegerType::get(ctx, 8); + mlir::IntegerType i32 = mlir::IntegerType::get(ctx, 32); + mlir::IntegerType u64 = mlir::IntegerType::get(ctx, 64); + mlir::LLVM::LLVMArrayType kernelNameArrType = mlir::LLVM::LLVMArrayType::get(i8, launchOp.getKernelName().size() + 1); + auto prev = builder.saveInsertionPoint(); + builder.setInsertionPointToStart(builder.getInsertionPoint()->getParentOfType().getBody()); + auto kernel_global = builder.create(loc, kernelNameArrType, true, mlir::LLVM::Linkage::Internal, "kernel_"+std::to_string(kernel_n)+"_path", mlir::StringAttr::get(launchOp.getKernelName().str() + '\0', i8)); + builder.restoreInsertionPoint(prev); + // mlir::LLVM::createGlobalString(loc, builder, "kernel_"+std::to_string(kernel_n)+"_str", launchOp.getKernelName(), mlir::LLVM::linkage::Linkage::Private ); + mlir::Value kernel_str = builder.create(loc, kernel_global); + mlir::Value num_args = builder.create(loc, mlir::IntegerAttr::get(u64, launchOp.getNumKernelOperands())); + kernel_n++; + return builder.create(loc, kMCLSetKernelFname, mlir::TypeRange({i32}), mlir::ValueRange({mcl_handle, kernel_str,num_args})).getResult(0); +} + +mlir::Value get_memref_num_elements(mlir::MLIRContext* ctx, mlir::OpBuilder& builder, mlir::Location loc, mlir::Value memref) +{ + mlir::Value rank = builder.create(loc, memref); + mlir::Value zero = builder.create(loc, 0); + mlir::Value one = builder.create(loc, 1); + + mlir::scf::ForOp forOp = builder.create(loc, zero, rank, one, mlir::ValueRange({one})); + mlir::Block* body = forOp.getBody(); + mlir::Value inductionvar = forOp.getInductionVar(); + mlir::IRRewriter::InsertPoint ip = builder.saveInsertionPoint(); + builder.setInsertionPointToStart(body); + mlir::Value dim = builder.create(loc, memref, inductionvar); + auto mul = builder.create(loc, forOp.getRegionIterArg(0), dim); + builder.create(loc, mlir::ValueRange({mul})); + builder.restoreInsertionPoint(ip); + + return forOp.getResult(0); +} + +void insert_set_task_arg_calls(mlir::MLIRContext* ctx, mlir::OpBuilder& builder, mlir::Location loc, mlir::Value mcl_handle, mlir::gpu::LaunchFuncOp launchOp, mlir::gpu::GPUFuncOp gpuFuncOp) +{ + mlir::IntegerType i32 = mlir::IntegerType::get(ctx, 32); + mlir::IntegerType u64 = mlir::IntegerType::get(ctx, 64); + mlir::LLVM::LLVMPointerType ptr = mlir::LLVM::LLVMPointerType::get(ctx); + for(size_t i = 0; i < launchOp.getNumKernelOperands(); i++) + { + mlir::Value arg = launchOp.getKernelOperand(i); + mlir::Value arg_id = builder.create(loc, mlir::IntegerAttr::get(u64, i)); + mlir::Type element_type; + mlir::Value size_in_bytes; + mlir::Value element_type_size; + mlir::Value num_elements; + uint64_t flag_val = 0; + mlir::Value flag; + if(auto memref = mlir::dyn_cast(arg.getType())) + { + element_type = memref.getElementType(); + if(memref.hasStaticShape()) + { + num_elements = builder.create(loc, memref.getNumElements()); + } + else + { + num_elements = get_memref_num_elements(ctx, builder, loc, arg); + } + arg = builder.create(loc, arg); + arg = builder.create(loc, u64, arg); + arg = builder.create(loc, ptr, arg); + flag_val |= MCL_ARG_BUFFER; + } + else + { + auto zeroIndex = builder.create(loc, 0); + auto alloca = builder.create(loc, mlir::MemRefType::get({1}, arg.getType())); + builder.create(loc, arg, alloca, mlir::ValueRange({zeroIndex})); + element_type = arg.getType(); + arg = builder.create(loc, alloca); + arg = builder.create(loc, u64, arg); + arg = builder.create(loc, ptr, arg); + num_elements = builder.create(loc, 1); + + flag_val |= MCL_ARG_SCALAR | MCL_ARG_INPUT; + } + if(element_type.isIntOrFloat()) + { + element_type_size = builder.create(loc, element_type.getIntOrFloatBitWidth()/8); + } + else if(element_type.isIndex()) + { + auto bits_in_byte = builder.create(loc, 8); + element_type_size = builder.create(loc, builder.create(loc), bits_in_byte); + } + size_in_bytes = builder.create(loc, element_type_size, num_elements); + for(auto& use: gpuFuncOp.getArgument(i).getUses()) + { + if(auto loadOp = mlir::dyn_cast(use.getOwner()); loadOp && loadOp.getMemRef() == use.get()) + { + flag_val |= MCL_ARG_INPUT; + } + else if(auto storeOp = mlir::dyn_cast(use.getOwner()); storeOp && storeOp.getMemRef() == use.get()) + { + flag_val |= MCL_ARG_OUTPUT; + } + } + flag = builder.create(loc, mlir::IntegerAttr::get(u64, flag_val)); + builder.create(loc, kMCLSetArgFname, mlir::TypeRange({i32}), mlir::ValueRange({mcl_handle, arg_id, arg, size_in_bytes, flag})); + } +} + +void insert_task_exec_calls(mlir::MLIRContext* ctx, mlir::OpBuilder& builder, mlir::Location loc, mlir::Value mcl_handle, mlir::Value flags, mlir::gpu::LaunchFuncOp launchOp) +{ + mlir::IntegerType u64 = mlir::IntegerType::get(ctx, 64); + mlir::LLVM::LLVMPointerType ptr = mlir::LLVM::LLVMPointerType::get(ctx); + + auto grid = launchOp.getGridSizeOperandValues(); + mlir::Value grid_sizes[3] = { builder.create(loc, u64, grid.x), + builder.create(loc, u64, grid.y), + builder.create(loc, u64, grid.z) + }; + + mlir::Value global_wg_mem = builder.create(loc, mlir::MemRefType::get({3}, u64)); + + auto block = launchOp.getBlockSizeOperandValues(); + mlir::Value block_sizes[3] = { + builder.create(loc, u64, block.x), + builder.create(loc, u64, block.y), + builder.create(loc, u64, block.z) + }; + + mlir::Value block_mem = builder.create(loc, mlir::MemRefType::get({3}, u64)); + mlir::Value indices[3] = {builder.create(loc, 0), builder.create(loc, 1), builder.create(loc, 2)}; + + for(size_t i = 0; i < 3; i++) + { + mlir::Value global_index = builder.create(loc, grid_sizes[i], block_sizes[i]); + builder.create(loc, global_index, global_wg_mem, indices[i]); + builder.create(loc, block_sizes[i], block_mem, indices[i]); + } + + mlir::IntegerType i32 = mlir::IntegerType::get(ctx, 32); + mlir::Value global_wg_mem_ptr = builder.create(loc, global_wg_mem); + global_wg_mem_ptr = builder.create(loc, u64, global_wg_mem_ptr); + global_wg_mem_ptr = builder.create(loc, ptr, global_wg_mem_ptr); + mlir::Value block_mem_ptr = builder.create(loc, block_mem); + block_mem_ptr = builder.create(loc, u64, block_mem_ptr); + block_mem_ptr = builder.create(loc, ptr, block_mem_ptr); + builder.create(loc, kMCLExecFname, mlir::TypeRange({i32}), mlir::ValueRange({mcl_handle, global_wg_mem_ptr, block_mem_ptr, flags})); +} + +void insert_task_wait_calls(mlir::MLIRContext* ctx, mlir::OpBuilder& builder, mlir::Location loc, mlir::Value mcl_handle) +{ + mlir::IntegerType i32 = mlir::IntegerType::get(ctx, 32); + // builder.create(loc, kMCLWaitFname, mlir::TypeRange({i32}), mlir::ValueRange({mcl_handle})); + builder.create(loc, kMCLWaitAllFname, mlir::TypeRange({i32}), mlir::ValueRange()); +} + + + +class ConvertGpuHostToMCLRT + : public ConvertGpuHostToMCLRTPassBase { +public: + ConvertGpuHostToMCLRT() = default; + ConvertGpuHostToMCLRT(const char* xclbin_path) + { + this->xclbin_path = xclbin_path; + } + + void runOnOperation() override { + mlir::MLIRContext* ctx = &getContext(); + mlir::ModuleOp module = getOperation(); + llvm::SmallVector launch_ops; + llvm::SmallMapVector gpu_func_ops; + mlir::IntegerType u64 = mlir::IntegerType::get(ctx, 64); + + module->walk([&launch_ops](mlir::gpu::LaunchFuncOp launch_op) { + launch_ops.push_back(launch_op); + }); + auto gpu_modules = module.getOps(); + + for(auto gpu_module: gpu_modules) + { + gpu_module.walk([&gpu_func_ops](mlir::gpu::GPUFuncOp gpu_func) { + gpu_func_ops.insert(std::make_pair(gpu_func.getNameAttr(), gpu_func)); + }); + + } + + + mlir::OpBuilder builder(module); + builder.setInsertionPointToStart(module.getBody()); + declare_mcl_funcs(ctx, builder); + mlir::func::FuncOp parent_func_op = nullptr; + + + for(auto launch_op: launch_ops) + { + if(parent_func_op == nullptr) + { + parent_func_op = launch_op->getParentOfType(); + } + else + { + assert(parent_func_op == launch_op->getParentOfType()); + } + mlir::Location loc = launch_op->getLoc(); + builder.setInsertionPoint(launch_op); + mlir::Value target = builder.create(launch_op->getLoc(), mlir::IntegerAttr::get(u64, MCL_TASK_FPGA)); + insert_prg_load_call(ctx, builder, loc, std::string(this->xclbin_path), target); + mlir::Value mcl_handle = insert_task_create(ctx, builder, loc); + insert_set_task_kernel_call(ctx, builder, loc, launch_op, mcl_handle); + insert_set_task_arg_calls(ctx, builder, loc, mcl_handle, launch_op, gpu_func_ops[launch_op.getKernelName()]); + insert_task_exec_calls(ctx, builder, loc, mcl_handle, target, launch_op); + insert_task_wait_calls(ctx, builder, loc, mcl_handle); + launch_op->erase(); + } + if(parent_func_op) + { + builder.setInsertionPointToStart(&parent_func_op.getBody().front()); + mlir::Value num_workers = builder.create(parent_func_op->getLoc(), mlir::IntegerAttr::get(u64,1)); + mlir::Value flags = builder.create(parent_func_op->getLoc(), mlir::IntegerAttr::get(u64,0)); + insert_mcl_init(ctx, builder, parent_func_op->getLoc(), num_workers, flags); + auto return_op = *parent_func_op.getOps().begin(); + builder.setInsertionPoint(return_op); + insert_mcl_finit(ctx, builder, parent_func_op->getLoc()); + } + + // for(auto func_op: gpu_func_ops) + // { + // func_op.second->erase(); + // } + + // for(auto gpu_module: llvm::make_early_inc_range(gpu_modules)) + // { + // gpu_module->erase(); + // } + } +}; + + +std::unique_ptr> mlir::comet::createConvertGpuHostToMCLRTPass() +{ + return std::make_unique<::ConvertGpuHostToMCLRT>(); +} + +std::unique_ptr> mlir::comet::createConvertGpuHostToMCLRTPass(const char* xclbin_path) +{ + return std::make_unique<::ConvertGpuHostToMCLRT>(xclbin_path); +} \ No newline at end of file diff --git a/lib/Conversion/GpuToBlockedGpu/CMakeLists.txt b/lib/Conversion/GpuToBlockedGpu/CMakeLists.txt new file mode 100644 index 00000000..fef9b8c3 --- /dev/null +++ b/lib/Conversion/GpuToBlockedGpu/CMakeLists.txt @@ -0,0 +1,6 @@ +add_llvm_library(COMETGpuToBlockedGpu + GpuToBlockedGpu.cpp + + DEPENDS + GpuToBlockedGpuPassIncGen +) \ No newline at end of file diff --git a/lib/Conversion/GpuToBlockedGpu/GpuToBlockedGpu.cpp b/lib/Conversion/GpuToBlockedGpu/GpuToBlockedGpu.cpp new file mode 100644 index 00000000..73179199 --- /dev/null +++ b/lib/Conversion/GpuToBlockedGpu/GpuToBlockedGpu.cpp @@ -0,0 +1,1432 @@ +#include "comet/Conversion/GpuToBlockedGpu/GpuToBlockedGpu.h" +#include "comet/Conversion/ForallToGpu/ForallToGpu.h" +#include "mlir/Conversion/AffineToStandard/AffineToStandard.h" +#include "mlir/Dialect/Arith/IR/Arith.h" +#include "mlir/Dialect/Bufferization/IR/Bufferization.h" +#include "mlir/Dialect/Func/IR/FuncOps.h" +#include "mlir/Dialect/GPU/IR/GPUDialect.h" +#include "mlir/Dialect/Index/IR/IndexOps.h" +#include "mlir/Dialect/Linalg/IR/Linalg.h" +#include "mlir/Dialect/MemRef/IR/MemRef.h" +#include "mlir/Dialect/Tensor/IR/Tensor.h" +#include "mlir/IR/Block.h" +#include "mlir/IR/Builders.h" +#include "mlir/IR/BuiltinAttributes.h" +#include "mlir/IR/BuiltinOps.h" +#include "mlir/IR/BuiltinTypeInterfaces.h" +#include "mlir/IR/BuiltinTypes.h" +#include "mlir/IR/OpDefinition.h" +#include "mlir/IR/Operation.h" +#include "mlir/IR/TypeRange.h" +#include "mlir/IR/Value.h" +#include "mlir/IR/ValueRange.h" +#include "mlir/IR/Visitors.h" +#include "mlir/Pass/Pass.h" +#include "mlir/Pass/PassManager.h" +#include "mlir/Support/LLVM.h" +#include "mlir/Transforms/DialectConversion.h" +#include "comet/Conversion/GpuToBlockedGpu/Passes.h" +#include "mlir/Dialect/SCF/IR/SCF.h" +#include "mlir/Dialect/Vector/IR/VectorOps.h" +#include "mlir/Dialect/Affine/IR/AffineOps.h" +#include "llvm/ADT/ArrayRef.h" +#include "llvm/ADT/MapVector.h" +#include "llvm/ADT/STLExtras.h" +#include "llvm/ADT/SmallSet.h" +#include "llvm/ADT/SmallVector.h" +#include "llvm/Support/Casting.h" +#include "llvm/Support/raw_ostream.h" +#include +#include +#include +#include +#include +#include +#include +#include +using namespace::mlir; + +#define GEN_PASS_CLASSES +#include "comet/Conversion/GpuToBlockedGpu/Passes.h.inc" + +struct BlockInfo { + BlockInfo(size_t& argIndex, uint64_t& blockSize, mlir::Value index, mlir::Value max) : argIndex(argIndex), blockSize(blockSize), index(index), max(max) {} + size_t argIndex; + uint64_t blockSize; + mlir::Value index; + mlir::Value max; +}; + +bool checkOperandsSameOrientation(Operation* op, llvm::MapVector>& useToIndices) +{ + if(op->getNumOperands() == 1) + { + return true; + } + + llvm::SmallVector indices; + for(auto operand: op->getOperands()) + { + if(useToIndices.find(operand) != useToIndices.end()) + { + indices = useToIndices[operand]; + } + } + + return llvm::all_of(op->getOperands(), [&useToIndices, &indices](Value val){ + if(useToIndices.find(val) == useToIndices.end()) + { + return true; + } + else + { + if(indices.size() != useToIndices[val].size()) + { + return false; + } + else + { + for(auto [i0, i1]: zip(indices, useToIndices[val])) + { + if(i0 != i1) + { + return false; + } + } + + return true; + } + } + }); +} + +bool checkOperandsSameShape(Operation* op) +{ + if(op->getNumOperands() == 1) + { + return true; + } + + ShapedType shape0 = mlir::dyn_cast(op->getOperandTypes().front()); + return llvm::all_of(op->getOperandTypes(), [&shape0](Type type){ + ShapedType shape1 = mlir::dyn_cast(type); + if(!shape0 && !shape1) + { + return true; + } + else if(shape0 && !shape1) + { + return false; + } + else if(!shape0 && shape1) + { + return false; + } + else if(shape0.getRank() != shape1.getRank()) + { + return false; + } + else + { + return llvm::all_of(zip(shape0.getShape(), shape1.getShape()), [](std::tuple dims) { + auto [first, second] = dims; + return first == second; + }); + } + }); +} + +bool checkOperandsSameType(Operation* op) +{ + + if(op->getNumOperands() == 1) + { + return true; + } + Type type0 = op->getOperandTypes().front(); + return llvm::all_of(op->getOperandTypes(), [&type0](Type type){return type == type0;}); +} + + +void broadcastWithRespectTo(OpBuilder& builder, OpOperand& op, SmallVector indicesRes, llvm::MapVector>& useToIndices) +{ + assert(useToIndices.find(op.get()) != useToIndices.end()); + auto operandIndices = useToIndices[op.get()]; + SmallVector resShape; + for(auto index: operandIndices) + { + RankedTensorType indexShaped = cast(index.getType()); + resShape.push_back(indexShaped.getShape()[0]); + } + + + SmallVector dimensions; + for(size_t i = 0; i < indicesRes.size(); i++) + { + if(std::find(operandIndices.begin(), operandIndices.end(), indicesRes[i]) == operandIndices.end()) + { + dimensions.push_back(i); + RankedTensorType indexShaped = cast(indicesRes[i].getType()); + resShape.insert(resShape.begin()+i, indexShaped.getShape()[0]); + operandIndices.insert(operandIndices.begin()+i, indicesRes[i]); + } + } + Type resElementType = dyn_cast(op.getOwner()->getResultTypes().front()) ? dyn_cast(op.getOwner()->getResultTypes().front()).getElementType() : op.getOwner()->getResultTypes().front(); + + RankedTensorType resType = RankedTensorType::get(resShape, resElementType); + + + Value newOperand = op.get(); + + if(!dimensions.empty()) + { + Value init = builder.create(op.getOwner()->getLoc(), resType.getShape(), resType.getElementType()); + newOperand = builder.create(op.getOwner()->getLoc(), newOperand, init, dimensions)->getResult(0); + bool isDone = useToIndices.insert(std::make_pair(newOperand, operandIndices)).second; + assert(isDone); + } + op.set(newOperand); +} + +class ConvertGpuToBlockedGpu: public CometGpuToBlockedGpuBase { + public: + ConvertGpuToBlockedGpu() = default; + + void runOnOperation() override + { + mlir::gpu::GPUFuncOp funcOp = getOperation(); + mlir::OpBuilder builder(funcOp); + + for(auto memrefArg : funcOp.getArguments()) + { + if(mlir::isa(memrefArg.getType())) + { + std::vector toExamine; + toExamine.insert(toExamine.end(), memrefArg.getUsers().begin(), memrefArg.getUsers().end()); + for(size_t i = 0; i < toExamine.size(); i++) + { + auto user = toExamine[i]; + if(mlir::isa(user)) + { + funcOp.setArgAttr(memrefArg.getArgNumber(), "gpu.read", builder.getUnitAttr()); + } + else if(mlir::isa(user)) + { + funcOp.setArgAttr(memrefArg.getArgNumber(), "gpu.write", builder.getUnitAttr()); + } + + if(user->getNumResults() > 0) + { + for(auto res: user->getResults()) + { + if(isa(res.getType())) + { + // If the result is a memref, we need to check its users as well + // to see if it is used in a store or load operation + toExamine.insert(toExamine.end(), res.getUsers().begin(), res.getUsers().end()); + } + } + } + } + } + } + + for(auto arg: funcOp.getArguments()) + { + if(mlir::isa(arg.getType())) + { + builder.setInsertionPointToStart(&funcOp.getBody().front()); + auto tensor = builder.create(arg.getLoc(), arg); + for(auto user: llvm::make_early_inc_range(arg.getUsers())) + { + if(mlir::memref::StoreOp storeOp = mlir::dyn_cast(user)) + { + + builder.setInsertionPoint(storeOp); + mlir::tensor::InsertOp insertOp = builder.create(storeOp->getLoc(), storeOp.getValueToStore(), tensor, storeOp.getIndices()); + /// TODO: If someone accesses the same memref (load) within the same block but no other store in-between, replace its use with the output of this operation + // storeOp.getMemRef().replaceUsesWithIf(insertOp.getResult(), [&](OpOperand& op) { + + // return op.getOwner()->getBlock() == storeOp->getBlock() && insertOp->isBeforeInBlock(op.getOwner()); + // }); + storeOp->erase(); + // StoreOp has no uses + // storeOp->replaceAllUsesWith(insertOp); + } + else if(mlir::memref::LoadOp loadOp = mlir::dyn_cast(user)) + { + builder.setInsertionPoint(loadOp); + mlir::tensor::ExtractOp extractOp = builder.create(loadOp->getLoc(), tensor, loadOp.getIndices()); + loadOp->replaceAllUsesWith(extractOp); + loadOp->erase(); + } + } + } + } + + llvm::MapVector> extracts; + llvm::MapVector> inserts; + llvm::SmallVector forOps; + std::vector needCheck; + llvm::MapVector> useToIndices; + + funcOp->walk([&](mlir::scf::ForOp forOp) { + // forOps with "blockSize" attribute induction variables need to be expanded to blocked indices + // i.e., for i : 0 -> M must be converted to a vector [0, 1, ..., M-1] that represents the indices taken by i + if(forOp->hasAttr("blockSize")) + { + forOps.push_back(forOp); + uint64_t blockSize = forOp->getAttrOfType("blockSize").getUInt(); + builder.setInsertionPointToStart(forOp.getBody()); + + // Since there is no built-in op that creates a range, we use unrealized casts to represent the expansion of an IV to a vector + auto blockedIndex = builder.create(forOp.getInductionVar().getLoc(), TypeRange({RankedTensorType::get({static_cast(blockSize)}, builder.getIndexType())}), forOp.getInductionVar()); + forOp.getInductionVar().replaceAllUsesExcept(blockedIndex.getResult(0), blockedIndex); + SmallVector indices; + indices.push_back(blockedIndex->getResult(0)); + bool isDone = useToIndices.insert(std::make_pair(blockedIndex.getResult(0), indices)).second; + assert(isDone); + // We have replaced indices with blocked indices, we need to make sure their users + // will be converted to blocks as well + for(auto res: blockedIndex->getResults()) + { + needCheck.push_back(res); + } + llvm::SmallVector inductionVars; + inductionVars.push_back(blockedIndex->getResult(0)); + + // If blocked indices (or their users) are used in a tensor insert/extract op, we cast them to a singl index + // so that the respective operation remains valid + for(size_t i = 0; i < inductionVars.size(); i++) + { + mlir::Value inductionVar = inductionVars[i]; + for(auto user: llvm::make_early_inc_range(inductionVar.getUsers())) + { + if(mlir::tensor::InsertOp insertOp = mlir::dyn_cast(user)) + { + if (inserts.find(insertOp) == inserts.end()) + { + llvm::SmallVector vec; + uint64_t one = 1; + for(size_t i = 0; i < insertOp.getIndices().size(); i++) + { + + vec.push_back(BlockInfo(i, one, insertOp.getIndices()[i], nullptr)); + } + inserts[insertOp] = vec; + } + for(size_t i = 0; i < insertOp.getIndices().size(); i++) + { + if(insertOp.getIndices()[i] == inductionVar) + { + builder.setInsertionPoint(insertOp); + auto cast = builder.create(inductionVar.getLoc(), builder.getIndexType(), inductionVar); + inserts[insertOp][i]= BlockInfo(i, blockSize, cast->getResult(0), forOp.getUpperBound()); + SmallVector indices; + indices.push_back(blockedIndex.getResult(0)); + bool isDone = useToIndices.insert(std::make_pair(cast.getResult(0), indices)).second; + assert(isDone); + } + } + } + else if(mlir::tensor::ExtractOp extractOp = mlir::dyn_cast(user)) + { + if (extracts.find(extractOp) == extracts.end()) + { + llvm::SmallVector vec; + uint64_t one = 1; + for(size_t i = 0; i < extractOp.getIndices().size(); i++) + { + + vec.push_back(BlockInfo(i, one, extractOp.getIndices()[i], nullptr)); + } + extracts[extractOp] = vec; + } + for(size_t i = 0; i < extractOp.getIndices().size(); i++) + { + if(extractOp.getIndices()[i] == inductionVar) + { + builder.setInsertionPoint(extractOp); + auto cast = builder.create(inductionVar.getLoc(), builder.getIndexType(), inductionVar); + extracts[extractOp][i] = BlockInfo(i, blockSize, cast->getResult(0), forOp.getUpperBound()); + SmallVector indices; + indices.push_back(blockedIndex.getResult(0)); + bool isDone = useToIndices.insert(std::make_pair(cast.getResult(0), indices)).second; + assert(isDone); + } + } + + for(auto res: user->getResults()) + { + inductionVars.push_back(res); + } + } + else + { + for(auto res: user->getResults()) + { + inductionVars.push_back(res); + } + } + } + } + } + }); + + // Convert extractOp to extractSliceOp as it is a closer representation of triton's blocked memory loads + // We also keep the information regarding the boolean masks that need to be generated in order to avoid + // reading beyond memory boundaries + funcOp->walk([&](tensor::ExtractOp extractOp){ + auto extract = extracts.find(extractOp); + if(extract == extracts.end()) + { + int64_t rank = extractOp.getTensor().getType().getRank(); + builder.setInsertionPoint(extractOp); + bufferization::ToTensorOp toTensorOp = mlir::cast(extractOp.getTensor().getDefiningOp()); + auto metaData = builder.create(extractOp->getLoc(), toTensorOp.getMemref()); + mlir::SmallVector blockSizes; + mlir::SmallVector blockSizesLiteral; + mlir::SmallVector static_offsets; + mlir::SmallVector offsets; + mlir::SmallVector stridesLiteral; + mlir::SmallVector strides; + SmallVector static_shape; + + for(int64_t i = 0; i < rank; i++) + { + blockSizesLiteral.push_back(1); + static_shape.push_back(1); + offsets.push_back(extractOp.getIndices()[i]); + stridesLiteral.push_back(ShapedType::kDynamic); + static_offsets.push_back(ShapedType::kDynamic); + Value stride = metaData->getResult(2 + rank + i); // tt.ptr + offset + (size) * rank + (stride) * rank + strides.push_back(stride); + } + + auto type = mlir::RankedTensorType::get(blockSizesLiteral, extractOp.getTensor().getType().getElementType()); + auto extractSlice = builder.create(extractOp->getLoc(), type, extractOp.getTensor(), mlir::ValueRange(offsets), mlir::ValueRange(blockSizes), mlir::ValueRange(strides), static_offsets, mlir::ArrayRef(blockSizesLiteral), mlir::ArrayRef(stridesLiteral)); + // extractSlice.dump(); + + // for(auto block: extract->second) + // { + // static_shape.push_back(block.blockSize); + // } + Value castSliceShape = builder.create(extractOp->getLoc(), mlir::RankedTensorType::get(static_shape, extractOp.getTensor().getType().getElementType()), extractSlice); + auto resultShape = mlir::cast(castSliceShape.getType()); + SmallVector,2> all_collapsed_indices; + + SmallVector collapsed_indices; + for(auto [i, d] : llvm::enumerate(resultShape.getShape())) + { + collapsed_indices.push_back(i); + if(d != 1) + { + all_collapsed_indices.push_back(collapsed_indices); + collapsed_indices.clear(); + } + } + if(!collapsed_indices.empty()) + { + if(!all_collapsed_indices.empty()) + { + all_collapsed_indices.back().insert(all_collapsed_indices.back().end(), collapsed_indices.begin(), collapsed_indices.end()); + } + else + { + all_collapsed_indices.push_back(collapsed_indices); + } + collapsed_indices.clear(); + } + + if(all_collapsed_indices.size() != resultShape.getRank()) + { + castSliceShape = builder.create(extractOp->getLoc(), castSliceShape, all_collapsed_indices); + } + if(mlir::cast(castSliceShape.getType()).getRank() == 1 && mlir::cast(castSliceShape.getType()).getDimSize(0) == 1) + { + castSliceShape = builder.create(extractOp->getLoc(), castSliceShape, mlir::ValueRange{builder.create(extractOp->getLoc(), 0)}); + } + + extractOp.replaceAllUsesWith(castSliceShape); + extractOp->erase(); + // needCheck.push_back(castSliceShape); + // if(inserted) + // { + // SmallVector indices; + // for(auto offset: offsets) + // { + // assert(useToIndices.find(offset) != useToIndices.end()); + // indices.insert(indices.end(),useToIndices[offset].begin(), useToIndices[offset].end()); + // // if(auto index = dyn_cast_if_present(offset.getDefiningOp())) + // // { + // // indices.push_back(index->getOperand(0)); + // // } + // // else + // // { + // // indices.push_back(offset); + // // } + // } + // assert(!indices.empty()); + // bool isDone = useToIndices.insert(std::make_pair(castSliceShape,indices)).second; + // assert(isDone); + // isDone = useToIndices.insert(std::make_pair(extractSlice,indices)).second; + // assert(isDone); + return WalkResult::advance(); // Continue walking + } + mlir::SmallVector blockSizes; + mlir::SmallVector blockSizesLiteral; + mlir::SmallVector static_offsets; + mlir::SmallVector offsets; + mlir::SmallVector stridesLiteral; + mlir::SmallVector strides; + std::sort(extract->second.begin(), extract->second.end(), [](BlockInfo& a, BlockInfo& b) {return a.argIndex < b.argIndex;} ); + int64_t rank = extract->first.getTensor().getType().getRank(); + bufferization::ToTensorOp toTensorOp = mlir::cast(extract->first.getTensor().getDefiningOp()); + builder.setInsertionPoint(extract->first); + + // Get information regarding the size and stride of the memref + auto metaData = builder.create(extract->first->getLoc(), toTensorOp.getMemref()); + for(int64_t i = 0; i < rank; i++) + { + if(extract->second[i].max) + { + blockSizes.push_back(extract->second[i].max); + blockSizesLiteral.push_back(ShapedType::kDynamic); + } + else + { + blockSizesLiteral.push_back(1); + } + offsets.push_back(extract->second[i].index); + stridesLiteral.push_back(ShapedType::kDynamic); + static_offsets.push_back(ShapedType::kDynamic); + Value stride = metaData->getResult(2 + rank + i); // tt.ptr + offset + (size) * rank + (stride) * rank + strides.push_back(stride); + } + + + auto type = mlir::RankedTensorType::get(blockSizesLiteral, extract->first.getTensor().getType().getElementType()); + auto extractSlice = builder.create(extract->first->getLoc(), type, extract->first.getTensor(), mlir::ValueRange(offsets), mlir::ValueRange(blockSizes), mlir::ValueRange(strides), static_offsets, mlir::ArrayRef(blockSizesLiteral), mlir::ArrayRef(stridesLiteral)); + + // extractSlice.dump(); + + SmallVector static_shape; + for(auto block: extract->second) + { + static_shape.push_back(block.blockSize); + } + Value castSliceShape = builder.create(extract->first->getLoc(), mlir::RankedTensorType::get(static_shape, extract->first.getTensor().getType().getElementType()), extractSlice); + auto resultShape = mlir::cast(castSliceShape.getType()); + SmallVector,2> all_collapsed_indices; + + SmallVector collapsed_indices; + for(auto [i, d] : llvm::enumerate(resultShape.getShape())) + { + collapsed_indices.push_back(i); + if(d != 1) + { + all_collapsed_indices.push_back(collapsed_indices); + collapsed_indices.clear(); + } + } + if(!collapsed_indices.empty()) + { + if(!all_collapsed_indices.empty()) + { + all_collapsed_indices.back().insert(all_collapsed_indices.back().end(), collapsed_indices.begin(), collapsed_indices.end()); + } + else + { + all_collapsed_indices.push_back(collapsed_indices); + } + collapsed_indices.clear(); + } + + if(all_collapsed_indices.size() != resultShape.getRank()) + { + castSliceShape = builder.create(extractOp->getLoc(), castSliceShape, all_collapsed_indices); + } + if(mlir::cast(castSliceShape.getType()).getRank() == 1 && mlir::cast(castSliceShape.getType()).getDimSize(0) == 1) + { + castSliceShape = builder.create(extractOp->getLoc(), castSliceShape, mlir::ValueRange{builder.create(extractOp->getLoc(), 0)}); + } + extract->first.replaceAllUsesWith(castSliceShape); + extract->first->erase(); + needCheck.push_back(castSliceShape); + // if(inserted) + // { + SmallVector indices; + for(auto [index, offset]: enumerate(offsets)) + { + if(resultShape.getDimSize(index) != 1) + { + indices.insert(indices.end(),useToIndices[offset].begin(), useToIndices[offset].end()); + } + // assert(useToIndices.find(offset) != useToIndices.end()); + // if(auto index = dyn_cast_if_present(offset.getDefiningOp())) + // { + // indices.push_back(index->getOperand(0)); + // } + // else + // { + // indices.push_back(offset); + // } + } + // assert(!indices.empty()); + bool isDone = useToIndices.insert(std::make_pair(castSliceShape,indices)).second; + assert(isDone); + isDone = useToIndices.insert(std::make_pair(extractSlice,indices)).second; + assert(isDone); + // llvm::errs() << "FOUND\n"; + // castSliceShape->dump(); + // llvm::errs() << castSliceShape->getResult(0).getAsOpaquePointer() << "\n"; + + // llvm::errs() << "======= START ===========\n"; + // for(auto index: useToIndices[castSliceShape->getResult(0)]) + // { + // llvm::errs() << index.getAsOpaquePointer() << "\n"; + // } + // llvm::errs() << "======= END ===========\n"; + // } + // else + // { + // assert(false && "Already there"); + // llvm::errs() << "Already there : \n"; + // already_there->dump(); + // llvm::errs() << "Trying to insert : \n"; + // castSliceShape->dump(); + // } + return WalkResult::advance(); // Continue walking + }); + + std::deque worklist; + funcOp.walk([&](Operation *op) { + worklist.push_back(op); + }); + + + + + + SmallVector newOpsToCheck; + while(!worklist.empty()) + { + auto user = worklist.front(); + // user->dump(); + worklist.pop_front(); + if(!user) + { + continue; + } + // user->dump(); + bool isNotPartOfReduction = llvm::any_of(user->getOperands(), [&](Value val) { + if(BlockArgument block_arg = mlir::dyn_cast(val)) + { + // block_arg.dump(); + if(auto owner_for = mlir::dyn_cast(block_arg.getOwner()->getParentOp())) + { + if(auto redTarget = std::find(owner_for.getRegionIterArgs().begin(), owner_for.getRegionIterArgs().end(), block_arg); redTarget != owner_for.getRegionIterArgs().end()) + { + // assert(useToIndices.find(val) != useToIndices.end()); + auto indices = useToIndices[val]; + useToIndices[user->getResult(0)] = indices; + return true; + } + } + } + else if(val.getDefiningOp()->hasAttr("notPartOfReduction")) + { + assert(useToIndices.find(val) != useToIndices.end()); + if(user->getNumResults() > 0) + { + useToIndices[user->getResult(0)] = useToIndices[val]; + } + return true; + } + return false; + }); + if(isNotPartOfReduction && !isa(user) && !isa(user)) + { + user->setAttr("notPartOfReduction", builder.getUnitAttr()); + } + + bool sameOpResType = user->hasTrait(); + bool sameOpResShape = user->hasTrait(); + bool sameOpType = user->hasTrait(); + bool elewise = user->hasTrait() && !isNotPartOfReduction; + // If inputs and outputs should have the same shape + // This includes most of `arith` operations. Instead of handling each operation + // explicitly, we handle them based on whether they have the specific trait + if(elewise) + { + // llvm::errs() << "Elementwise\n"; + // user->dump(); + for(auto operand: user->getOperands()) + { + if(llvm::isa_and_present(operand.getDefiningOp())) + { + + } + else if(llvm::isa(operand) && operand.getParentBlock()->getParentOp()->hasAttr("reduceDim")) + { + + } + else if(llvm::isa(operand) && isa(operand.getParentBlock()->getParentOp())) + { + + } + else + { + // assert(useToIndices.find(operand) != useToIndices.end()); + } + } + + // llvm::errs() << "Has SameOperandsAmdResultTypeType trait\n"; + bool operandsSameShape = checkOperandsSameShape(user); + bool operandsSameOrientation = checkOperandsSameOrientation(user, useToIndices); + // If input operands do not have the same type/shape + if(!operandsSameShape || !operandsSameOrientation) + { + // llvm::errs() << "Not same shape\n"; + // user->dump(); + auto shaped0 = mlir::dyn_cast(user->getOperandTypes()[0]); + auto shaped1 = mlir::dyn_cast(user->getOperandTypes()[1]); + auto indices0 = useToIndices.find(user->getOperand(0)); + auto indices1 = useToIndices.find(user->getOperand(1)); + // If LHS is shaped and RHS is scalar, simply "splat"/expand the value to + // a tensor of same shape as the LHS + + if(shaped0 && !shaped1) + { + assert(indices1 == useToIndices.end()); + builder.setInsertionPoint(user); + auto outputType = RankedTensorType::get(shaped0.getShape(), user->getOperandTypes()[1]); + auto splatOp = builder.create(user->getLoc(), outputType, user->getOperands()[1]); + assert(useToIndices.find(user->getOperand(0)) != useToIndices.end()); + SmallVector indices = useToIndices[user->getOperand(0)]; + useToIndices[splatOp] = indices; + // bool isDone = useToIndices.insert(std::make_pair(splatOp, indices)).second; + // assert(isDone); + user->getOpOperands()[1].set(splatOp); + } + // Same as above with LHS being scalar and RHS being shaped + else if(!shaped0 && shaped1) + { + assert(indices0 == useToIndices.end()); + builder.setInsertionPoint(user); + auto outputType = RankedTensorType::get(shaped1.getShape(), user->getOperandTypes()[0]); + auto splatOp = builder.create(user->getLoc(), outputType, user->getOperands()[0]); + assert(useToIndices.find(user->getOperand(1)) != useToIndices.end()); + SmallVector indices = useToIndices[user->getOperand(1)]; + useToIndices[splatOp] = indices; + // bool isDone = useToIndices.insert(std::make_pair(splatOp, indices)).second; + // assert(isDone); + + user->getOpOperands()[0].set(splatOp); + } + // Both are shaped, but the shapes do not match or they have different indexing orientation + // We need to infer how to reshape them based on their indexing + else + { + // llvm::errs() << "Not same shape but both ranked\n"; + // user->dump(); + assert(shaped0.getElementType() == shaped1.getElementType() && "Inputs need to have same element types"); + + SmallVector& indices0 = useToIndices[user->getOperand(0)]; + SmallVector& indices1 = useToIndices[user->getOperand(1)]; + + assert(!indices0.empty()); + assert(!indices1.empty()); + if(shaped0.getRank() == shaped1.getRank()) + { + // llvm::errs() << "Not same shape but same rank\n"; + // user->dump(); + auto forOp = user->getParentOfType(); + + if(forOp && forOp->hasAttr("blockSize") && forOp.getNumRegionIterArgs() > 0) // if it's a reduction loop + { + auto redForOp = forOp; + + /// TODO: Shapes same rank but different possibly mean matrix multiplication + if(llvm::all_of(user->getUsers(), [](Operation* op){ + return !op->hasAttr("notPartOfReduction"); + })) + { + assert(redForOp); + // assert(redForOp.getRegionIterArgs().size() == 1); + auto userUser = user->getUses().begin(); + while(!isa(userUser->getOwner())) + { + userUser = userUser->getOwner()->getUses().begin(); + } + + + if(mlir::isa(user) && user->hasOneUse() && mlir::isa(*user->getUsers().begin()) && shaped0.getRank() > 1) + { + builder.setInsertionPoint(user); + + auto userUserT = mlir::dyn_cast(redForOp->getOperand(3 + userUser->getOperandNumber()).getType()); + auto empty = builder.create(user->getLoc(), userUserT.getShape(), userUserT.getElementType()); + auto matmul = builder.create(user->getLoc(), user->getOperands(), ValueRange({empty})); + user->replaceAllUsesWith(matmul); + // llvm::errs() << "Erasing: " << user <<"\n"; + user->erase(); + newOpsToCheck.push_back(matmul.getResult(0)); + auto found = useToIndices.find(redForOp->getOperand(3 + userUser->getOperandNumber())); + assert(found != useToIndices.end()); + + bool isDone = useToIndices.insert(std::make_pair(matmul->getResult(0), useToIndices[found->first])).second; + assert(isDone); + continue; + } + } + } + } + + // This could be improved by first trying to make the shapes equal with respect to each other rather than + // the eventual result first + + SmallVector indices0_sorted, indices1_sorted; + for(auto index: indices0) + { + builder.setInsertionPoint(user); + + indices0_sorted.push_back(index.getAsOpaquePointer()); + } + for(auto index: indices1) + { + builder.setInsertionPoint(user); + indices1_sorted.push_back(index.getAsOpaquePointer()); + } + std::sort(indices0_sorted.begin(),indices0_sorted.end()); + std::sort(indices1_sorted.begin(),indices1_sorted.end()); + + SmallVector diff_0_not_in_1; + std::set_difference(indices0_sorted.begin(), indices0_sorted.end(), indices1_sorted.begin(), indices1_sorted.end(), std::back_inserter(diff_0_not_in_1)); + SmallVector diff_1_not_in_0; + std::set_difference(indices1_sorted.begin(), indices1_sorted.end(), indices0_sorted.begin(), indices0_sorted.end(), std::back_inserter(diff_1_not_in_0)); + + if(shaped0.getRank() > shaped1.getRank() && diff_1_not_in_0.empty()) // everything in indices1 is also in indices0 + { + broadcastWithRespectTo(builder, user->getOpOperand(1), useToIndices[user->getOpOperand(0).get()], useToIndices); + } + else if(shaped0.getRank() < shaped1.getRank() && diff_0_not_in_1.empty()) // everything in indices0 is also in indices1 + { + broadcastWithRespectTo(builder, user->getOpOperand(0), useToIndices[user->getOpOperand(1).get()], useToIndices); + } + else + { + auto result = user->getUses().begin(); + while(!isa(result->getOwner()) && !isa(result->getOwner())) + { + result = result->getOwner()->getUses().begin(); + } + + SmallVector indicesRes; + if(isa(result->getOwner())) + { + auto parentOp = result->getOwner()->getParentOp(); + assert(useToIndices.find(parentOp->getResult(result->getOperandNumber())) != useToIndices.end()); + indicesRes = useToIndices[parentOp->getResult(result->getOperandNumber())]; + } + else + { + tensor::InsertOp insertOp = cast(result->getOwner()); + for(auto index: insertOp.getIndices()) + { + assert(useToIndices.find(index) != useToIndices.end()); + assert(useToIndices[index].size() == 1); + indicesRes.push_back(useToIndices[index][0]); + } + } + + builder.setInsertionPoint(user); + + for(auto& operand: user->getOpOperands()) + { + broadcastWithRespectTo(builder, operand, indicesRes, useToIndices); + } + + bool needFurtherBroadcasting = checkOperandsSameOrientation(user, useToIndices); + // Handle the case where dimensions are still missing from one or both tensors + if(needFurtherBroadcasting) + { + auto& operand0 = user->getOpOperand(0); + auto& operand1 = user->getOpOperand(1); + SmallVector indicesRes = useToIndices[operand1.get()]; + broadcastWithRespectTo(builder, operand0, indicesRes, useToIndices); + indicesRes = useToIndices[operand0.get()]; + broadcastWithRespectTo(builder, operand1, indicesRes, useToIndices); + } + } + } + + } + + // For operations we have already "blocked" their input operands + // their outputs would still have remained scalar. + // This can be easily fixed by re-creating the same operation with + // the new "blocked" inputs, which will automatically update their output + // type + Operation* blockedOp = nullptr; + if(user->getOperandTypes()[0] != user->getResultTypes()[0] && sameOpResType) + { + builder.setInsertionPoint(user); + blockedOp = builder.create(user->getLoc(), user->getName().getIdentifier(), user->getOperands(), TypeRange(user->getOperandTypes()[0]), user->getAttrs()); + } + else if(elewise || sameOpResShape) + { + if(RankedTensorType input0T = dyn_cast(user->getOperandTypes()[0]) ) + { + if(!isa(user->getResultTypes()[0])) + { + builder.setInsertionPoint(user); + + blockedOp = builder.create(user->getLoc(), user->getName().getIdentifier(), user->getOperands(), TypeRange(RankedTensorType::get(input0T.getShape(), user->getResultTypes()[0])), user->getAttrs()); + } + + } + } + if(blockedOp) + { + user->replaceAllUsesWith(blockedOp); + newOpsToCheck.insert(newOpsToCheck.end(), blockedOp->getResults().begin(), blockedOp->getResults().end()); + if(useToIndices.find(user->getOperands()[0])!=useToIndices.end()) + { + bool isDone = useToIndices.insert(std::make_pair(blockedOp->getResult(0), useToIndices[user->getOperands()[0]])).second; + assert(isDone); + } + else if(user->getOperands().size() > 1 && useToIndices.find(user->getOperands()[1])!=useToIndices.end()) + { + bool isDone = useToIndices.insert(std::make_pair(blockedOp->getResult(0), useToIndices[user->getOperands()[1]])).second; + assert(isDone); + } + // llvm::errs() << "Erasing: " << user <<"\n"; + user->erase(); + + } + + // llvm::errs() << "FOUND\n"; + // blockedOp->dump(); + // llvm::errs() << blockedOp->getResult(0).getAsOpaquePointer() << "\n"; + // llvm::errs() << "Operand 1\n"; + // llvm::errs() << user->getOperands()[1].getAsOpaquePointer() << "\n"; + // llvm::errs() << "======= START ===========\n"; + // for(auto index: useToIndices[blockedOp->getResult(0)]) + // { + // llvm::errs() << index.getAsOpaquePointer() << "\n"; + // } + // llvm::errs() << "======= END ===========\n"; + // else + // { + // llvm::errs() << "Skipped :\n"; + // user->dump(); + // } + } + else if(scf::ForOp forOp = mlir::dyn_cast(user)) + { + // If one of the changed operations has a use in a ForOp operation's initArgs + // we need to change the forOp accordingly and update its inner uses as well as its results + bool needsFixes = false; + for(size_t i = 0; i < forOp.getInitArgs().size(); i++) + { + if(forOp.getInitArgs()[i].getType() != forOp.getRegionIterArgs()[i].getType()) + { + needsFixes = true; + } + } + if(!needsFixes) + { + continue; + } + // op.dump(); + builder.setInsertionPoint(user); + auto newForOp = builder.create(user->getLoc(), forOp.getLowerBound(), forOp.getUpperBound(), forOp.getStep(), forOp.getInitArgs()); + if(forOp->hasAttr("blockSize")) + { + newForOp->setAttr("blockSize", forOp->getAttr("blockSize")); + } + else if(forOp->hasAttr("reduceDim")) + { + newForOp->setAttr("reduceDim", forOp->getAttr("reduceDim")); + + } + builder.setInsertionPointToEnd(newForOp.getBody()); + auto newYield = builder.create(forOp->getLoc(), forOp.getBody()->getTerminator()->getOperands()); + // forOp.getBody()->getTerminator()->erase(); + for(auto& inner_op: llvm::make_early_inc_range(forOp.getBody()->getOperations())) + { + inner_op.moveBefore(newForOp.getBody()->getTerminator()); + } + // llvm::errs() << "Erasing: " << newYield << newYield.getAsOpaquePointer() <<"\n"; + newYield->erase(); + + // newForOp.getBody()->getTerminator()->erase(); + for(auto inner_arg: zip(forOp.getBody()->getArguments(), newForOp.getBody()->getArguments())) + { + std::get<0>(inner_arg).replaceAllUsesWith(std::get<1>(inner_arg)); + newOpsToCheck.push_back(std::get<1>(inner_arg)); + } + for(auto [old_arg, new_arg]: zip(forOp.getInits(), newForOp.getRegionIterArgs())) + { + assert(useToIndices.find(old_arg) != useToIndices.end()); + bool isDone = useToIndices.insert(std::make_pair(new_arg, useToIndices[old_arg])).second; + assert(isDone); + } + forOp->replaceAllUsesWith(newForOp); + for(auto [res, init_val]: llvm::zip(newForOp.getResults(), newForOp.getInits())) + { + assert(useToIndices.find(init_val) != useToIndices.end()); + useToIndices[res] = useToIndices[init_val]; + } + newOpsToCheck.insert(newOpsToCheck.end() , newForOp->getResults().begin(), newForOp->getResults().end()); + // llvm::errs() << "Erasing: " << forOp << forOp.getAsOpaquePointer() <<"\n"; + forOp->erase(); + + } + } + + funcOp->walk([&](tensor::InsertOp insertOp){ + auto insert = inserts.find(insertOp); + if(insert == inserts.end()) + { + bufferization::ToTensorOp toTensorOp = mlir::cast(insertOp.getDest().getDefiningOp()); + builder.setInsertionPoint(insertOp); + auto metaData = builder.create(insertOp->getLoc(), toTensorOp.getMemref()); + mlir::SmallVector blockSizes; + mlir::SmallVector blockSizesLiteral; + + mlir::SmallVector static_offsets; + mlir::SmallVector offsets; + mlir::SmallVector stridesLiteral; + mlir::SmallVector strides; + + int64_t rank = insertOp.getResult().getType().getRank(); + for(int64_t i = 0; i < rank; i++) + { + blockSizesLiteral.push_back(1); + offsets.push_back(insertOp.getIndices()[i]); + stridesLiteral.push_back(ShapedType::kDynamic); + static_offsets.push_back(ShapedType::kDynamic); + Value stride = metaData.getResult(2 + rank + i); // tt.ptr + offset + (size) * rank + (stride) * rank + strides.push_back(stride); + } + + auto type = mlir::RankedTensorType::get(blockSizesLiteral, insertOp.getResult().getType().getElementType()); + builder.setInsertionPoint(insertOp); + Value insert_val = insertOp.getScalar(); + auto insert_shape = dyn_cast(insert_val.getType()); + + if(insert_shape && insert_shape.getRank() != type.getRank()) + { + SmallVector,2> all_expanded_indices; + + SmallVector expanded_indices; + SmallVector shape; + int64_t static_index = 0; + for(auto [i, d] : llvm::enumerate(type.getShape())) + { + expanded_indices.push_back(i); + if(d != 1) + { + shape.push_back(insert_shape.getDimSize(static_index++)); + all_expanded_indices.push_back(expanded_indices); + expanded_indices.clear(); + } + else + { + shape.push_back(1); + } + } + if(!expanded_indices.empty()) + { + if(!all_expanded_indices.empty()) + { + all_expanded_indices.back().insert(all_expanded_indices.back().end(), expanded_indices.begin(), expanded_indices.end()); + } + else + { + all_expanded_indices.push_back(expanded_indices); + } + expanded_indices.clear(); + } + + insert_val = builder.create(insertOp->getLoc(), mlir::RankedTensorType::get(shape, insert_shape.getElementType()), insert_val, all_expanded_indices); + } + else if(!insert_shape) + { + insert_val = builder.create(insertOp->getLoc(), type, insert_val); + } + + auto castSliceShape = builder.create(insertOp->getLoc(), type, insert_val); + + auto insertSlice = builder.create(insertOp->getLoc(), castSliceShape, insertOp.getDest(), mlir::ValueRange(offsets), mlir::ValueRange(blockSizes), mlir::ValueRange(strides), static_offsets, mlir::ArrayRef(blockSizesLiteral), mlir::ArrayRef(stridesLiteral)); + // llvm::errs() << "Erasing: " << insertOp << insertOp.getAsOpaquePointer(); + insertOp->erase(); + + return WalkResult::advance(); // Continue walking + + } + mlir::SmallVector blockSizes; + mlir::SmallVector blockSizesLiteral; + + mlir::SmallVector static_offsets; + mlir::SmallVector offsets; + mlir::SmallVector stridesLiteral; + mlir::SmallVector strides; + std::sort(insert->second.begin(), insert->second.end(), [](BlockInfo& a, BlockInfo& b) {return a.argIndex < b.argIndex;} ); + bufferization::ToTensorOp toTensorOp = mlir::cast(insert->first.getDest().getDefiningOp()); + builder.setInsertionPoint(insert->first); + auto metaData = builder.create(insert->first->getLoc(), toTensorOp.getMemref()); + + int64_t rank = insert->first.getResult().getType().getRank(); + for(int64_t i = 0; i < rank; i++) + { + if(insert->second[i].max) + { + blockSizes.push_back(insert->second[i].max); + blockSizesLiteral.push_back(ShapedType::kDynamic); + } + else + { + blockSizesLiteral.push_back(1); + } + offsets.push_back(insert->second[i].index); + stridesLiteral.push_back(ShapedType::kDynamic); + static_offsets.push_back(ShapedType::kDynamic); + Value stride = metaData.getResult(2 + rank + i); // tt.ptr + offset + (size) * rank + (stride) * rank + strides.push_back(stride); + } + auto type = mlir::RankedTensorType::get(blockSizesLiteral, insert->first.getResult().getType().getElementType()); + builder.setInsertionPoint(insert->first); + Value insert_val = insert->first.getScalar(); + auto insert_shape = dyn_cast(insert_val.getType()); + if(insert_shape && type.getRank() != insert_shape.getRank()) + { + SmallVector,2> all_expanded_indices; + + SmallVector expanded_indices; + SmallVector shape; + int64_t static_index = 0; + for(auto [i, offset] : llvm::enumerate(offsets)) + { + expanded_indices.push_back(i); + RankedTensorType shaped_offset = nullptr; + if(auto unrealizedCast = dyn_cast_if_present(offset.getDefiningOp())) + { + shaped_offset = dyn_cast(unrealizedCast.getInputs().front().getType()); + } + if(!shaped_offset || (shaped_offset && shaped_offset.getRank() == 1 && shaped_offset.getDimSize(0) == 1)) + { + shape.push_back(1); + } + else + { + shape.push_back(insert_shape.getDimSize(static_index++)); + all_expanded_indices.push_back(expanded_indices); + expanded_indices.clear(); + } + } + if(!expanded_indices.empty()) + { + if(!all_expanded_indices.empty()) + { + all_expanded_indices.back().insert(all_expanded_indices.back().end(), expanded_indices.begin(), expanded_indices.end()); + } + else + { + all_expanded_indices.push_back(expanded_indices); + } + expanded_indices.clear(); + } + + insert_val = builder.create(insert->first->getLoc(), mlir::RankedTensorType::get(shape, insert_shape.getElementType()), insert_val, all_expanded_indices); + } + else if(!insert_shape) + { + auto splatShape = type.getShape().vec(); + for(size_t i = 0; i < splatShape.size(); i++) + { + splatShape[i] = 1; + } + + insert_val = builder.create(insertOp->getLoc(), mlir::RankedTensorType::get(splatShape, insert_val.getType()) , insert_val); + } + + auto castSliceShape = builder.create(insert->first->getLoc(), type, insert_val); + + auto insertSlice = builder.create(insert->first->getLoc(), castSliceShape, insert->first.getDest(), mlir::ValueRange(offsets), mlir::ValueRange(blockSizes), mlir::ValueRange(strides), static_offsets, mlir::ArrayRef(blockSizesLiteral), mlir::ArrayRef(stridesLiteral)); + // llvm::errs() << "Erasing: " << insert->first << insert->first.getAsOpaquePointer(); + insert->first->erase(); + + /// TODO: It might make sense to replaceAll uses at some point later + // insert->first.replaceAllUsesWith(insertSlice.getResult()); + return WalkResult::advance(); // Continue walking + }); + + + // Since we might have replaced forOps if they were using initArgs + // we need to collect forOps with "blockSize" again + forOps.clear(); + funcOp->walk([&](mlir::scf::ForOp forOp) { + if(forOp->hasAttr("blockSize")) + { + forOps.push_back(forOp); + }}); + + for(scf::ForOp forOp: llvm::make_early_inc_range(forOps)) + { + // Since for loops with blockSize attribute are replaced with blocked operations + // we can safely remove them + if(forOp.getInitArgs().empty()) + { + builder.setInsertionPoint(forOp); + auto c0 = builder.create(forOp->getLoc(), 0); + forOp.getInductionVar().replaceAllUsesWith(c0); + auto yieldOp = *forOp.getOps().begin(); + // llvm::errs() << "Erasing: " << yieldOp << yieldOp.getAsOpaquePointer() << "\n"; + yieldOp->erase(); + + forOp->getBlock()->getOperations().splice(forOp->getIterator(), forOp.getBody()->getOperations()); + // llvm::errs() << "Erasing: " << forOp << forOp.getAsOpaquePointer() << "\n"; + forOp->erase(); + } + else // Reduction loops need special handling + { + SmallVector notPartOfReduction; + forOp->walk([¬PartOfReduction](Operation* op) { + if(op->hasAttr("notPartOfReduction")) + { + notPartOfReduction.push_back(op); + } + }); + llvm::SmallSet skip; + // SmallVector opsNotInReduction; + for(auto op: llvm::make_early_inc_range(notPartOfReduction)) + { + // op->dump(); + Value isBlockArgument0 = dyn_cast(op->getOperands()[0]); + Value isBlockArgument1 = dyn_cast(op->getOperands()[1]); + bool is0notPartOfReduction = isBlockArgument0 ? isBlockArgument0.getParentBlock()->getParentOp()->hasAttr("blockSize") : op->getOperands()[0].getDefiningOp()->hasAttr("notPartOfReduction"); + bool is1notPartOfReduction = isBlockArgument1 ? isBlockArgument1.getParentBlock()->getParentOp()->hasAttr("blockSize") : op->getOperands()[1].getDefiningOp()->hasAttr("notPartOfReduction"); + Value v0 = op->getOperand(0); + Value v1 = op->getOperand(1); + SmallVector& indices0 = useToIndices[v0]; + SmallVector& indices1 = useToIndices[v1]; + + // If one of the, possibly, two operands is not operating as part of the reduction, i.e., has the attribute notPartOfReduction + // this operand is either the accumulating target buffer (init) or an operation on it + Value init, toReduced; + int64_t reduce_operand, other_operand; + linalg::ReduceOp reduceOp; + if((is0notPartOfReduction && is1notPartOfReduction) || ((indices0.size() == 0 || indices1.size() == 0) && (indices0.size() != 0 || indices1.size() != 0))) + { + if((reduceOp = dyn_cast_if_present(op->getOperand(0).getDefiningOp()))) + { + reduce_operand = 0; + other_operand = 1; + } + else if((reduceOp = dyn_cast_if_present(op->getOperand(1).getDefiningOp()))) + { + reduce_operand = 1; + other_operand = 0; + } + // else + // { + // assert(false && "Unexpected case where none of the operands are part of a reduce operation"); + // } + + if(reduceOp) + { + + if(!dyn_cast(op->getOperand(other_operand).getType())) + { + builder.setInsertionPoint(reduceOp.getBody()->getTerminator()); + Operation* newRes = builder.create(op->getLoc(), op->getName().getIdentifier(), {op->getOperand(other_operand), reduceOp.getBody()->getTerminator()->getOperands().back()}, op->getResultTypes()); + reduceOp.getBody()->getTerminator()->getOpOperands().back().set(newRes->getResult(0)); + op->replaceAllUsesWith(reduceOp); + // llvm::errs() << "Erasing: " << op <<"\n"; + + op->erase(); + + continue; + } + else + { + assert(false && "Reductions in the for of C = a op C op B ... are not supported. If your reduction can be expressed as C = C op (a op B) please rewrite it as such."); + } + } + } + if(!reduceOp && is0notPartOfReduction) + { + init = op->getOperands()[0]; + toReduced = op->getOperands()[1]; + } + else if(!reduceOp && is1notPartOfReduction) + { + init = op->getOperands()[1]; + toReduced = op->getOperands()[0]; + } + + builder.setInsertionPoint(op); + + + + // v0.dump(); + // v1.dump(); + + + assert(useToIndices.find(v0) != useToIndices.end()); + assert(useToIndices.find(v1) != useToIndices.end()); + + // Collect the induction variable indices used to form these values from load operations + /// TODO: Handle the case where a scalar is used, i.e., A[i] + 1, 1 will not have indices + /// but it is valid to be expanded to any direction + if(indices0.size() < indices1.size()) + { + std::swap(indices0, indices1); + std::swap(v0, v1); + } + if(indices0.size() != indices1.size()) + { + // op->getOpOperand(0).get().dump(); + // op->getOpOperand(1).get().dump(); + // op->dump(); + SmallVector dimensions; + for(size_t i = 0; i < indices0.size(); i++) + { + // auto it = std::find(indices1.begin(), indices1.end(), indices0[i]); + if(std::find(indices1.begin(), indices1.end(), indices0[i]) == indices1.end()) + { + dimensions.push_back(i); + } + } + bool needsExtract = false; + /// TODO: Handle cases where the init and reduced tensors are of the same rank (e.g., matrix multiplication, reducing a row vector against a column vector etc) + if(!isa(init.getType() )) + { + needsExtract = true; + auto prev_indices = useToIndices[init]; + init = builder.create(init.getLoc(), RankedTensorType::get({}, init.getType()), init); + useToIndices[init] = prev_indices; + } + Operation* reduceOp = builder.create(op->getLoc(), toReduced, init, dimensions, [&](OpBuilder builder, Location loc, ValueRange values){ + auto res = builder.create(loc, op->getName().getIdentifier(), values, op->getResultTypes()); + builder.create(loc, res->getResults()); + }); + reduceOp->setAttr("notPartOfReduction", builder.getUnitAttr()); + assert(useToIndices.find(init) != useToIndices.end()); + useToIndices[reduceOp->getResult(0)] = useToIndices[init]; + if(needsExtract) + { + reduceOp = builder.create(reduceOp->getLoc(), reduceOp->getResult(0), ValueRange()); + reduceOp->setAttr("notPartOfReduction", builder.getUnitAttr()); + useToIndices[reduceOp->getResult(0)] = useToIndices[init]; + } + + op->replaceAllUsesWith(reduceOp); + // llvm::errs() << "Erasing: " << op <<"\n"; + // op->dump(); + + op->erase(); + } + else + { + /// TODO: Handle other cases where shapes are same rank + if(mlir::isa(op) && llvm::any_of(op->getOperands(), [](Value val){ + return llvm::isa_and_present(val.getDefiningOp()); + })) + { + + builder.setInsertionPoint(op); + auto addFOp = builder.create(op->getLoc(), op->getOperands()); + op->replaceAllUsesWith(addFOp); + // llvm::errs() << "Erasing: " << op <<"\n"; + op->erase(); + } + else + { + + SmallVector dimensions; + if(is1notPartOfReduction) + { + broadcastWithRespectTo(builder, op->getOpOperand(0) , indices1 , useToIndices); + indices0 = useToIndices[op->getOperand(0)]; + for(size_t i = 0; i < indices0.size(); i++) + { + // auto it = std::find(indices1.begin(), indices1.end(), indices0[i]); + if(std::find(indices1.begin(), indices1.end(), indices0[i]) == indices1.end()) + { + dimensions.push_back(i); + } + } + toReduced = op->getOperand(0); + } + else if(is0notPartOfReduction) + { + broadcastWithRespectTo(builder, op->getOpOperand(1) , indices0 , useToIndices); + indices1 = useToIndices[op->getOperand(1)]; + for(size_t i = 0; i < indices1.size(); i++) + { + // auto it = std::find(indices1.begin(), indices1.end(), indices1[i]); + if(std::find(indices0.begin(), indices0.end(), indices1[i]) == indices0.end()) + { + dimensions.push_back(i); + } + } + toReduced = op->getOperand(1); + } + + + auto reduceOp = builder.create(op->getLoc(), toReduced, init, dimensions, [&](OpBuilder builder, Location loc, ValueRange values){ + auto res = builder.create(loc, op->getName().getIdentifier(), values, op->getResultTypes()); + builder.create(loc, res->getResults()); + }); + reduceOp->setAttr("notPartOfReduction", builder.getUnitAttr()); + assert(useToIndices.find(init) != useToIndices.end()); + useToIndices[reduceOp.getResult(0)] = useToIndices[init]; + + op->replaceAllUsesWith(reduceOp); + // llvm::errs() << "Erasing: " << op <<"\n"; + // op->dump(); + + op->erase(); + } + } + + } + + builder.setInsertionPoint(forOp); + auto c0 = builder.create(forOp->getLoc(), 0); + forOp.getInductionVar().replaceAllUsesWith(c0); + for(auto [index ,iter_arg]: enumerate(forOp.getRegionIterArgs())) + { + iter_arg.replaceAllUsesWith(forOp.getInitArgs()[index]); + } + forOp->replaceAllUsesWith(forOp.getBody()->getTerminator()->getOperands()); + // llvm::errs() << "Erasing: " << forOp.getBody()->getTerminator() << "\n"; + + forOp.getBody()->getTerminator()->erase(); + forOp->getBlock()->getOperations().splice(forOp->getIterator(), forOp.getBody()->getOperations()); + // llvm::errs() << "Erasing: " << forOp << forOp.getAsOpaquePointer() << "\n"; + + forOp->erase(); + + } + } + + + PassManager pm(funcOp.getContext()); + + + pm.addPass(mlir::createLowerAffinePass()); + if (failed(pm.run(funcOp))) { + signalPassFailure(); + return; + } + } +}; + +std::unique_ptr> mlir::comet::createConvertGpuToBlockedGpuPass() { + return std::make_unique(); +} \ No newline at end of file diff --git a/lib/Conversion/GpuToOCLSPIRV/CMakeLists.txt b/lib/Conversion/GpuToOCLSPIRV/CMakeLists.txt new file mode 100644 index 00000000..16f35c8b --- /dev/null +++ b/lib/Conversion/GpuToOCLSPIRV/CMakeLists.txt @@ -0,0 +1,7 @@ +add_llvm_library(COMETGpuToOCLSPIRV + GpuToOCLSPIRVPass.cpp + GPUToSPIRVPass.cpp + + DEPENDS + GpuToOCLSPIRVConversionPassIncGen +) \ No newline at end of file diff --git a/lib/Conversion/GpuToOCLSPIRV/GPUToSPIRVPass.cpp b/lib/Conversion/GpuToOCLSPIRV/GPUToSPIRVPass.cpp new file mode 100644 index 00000000..5081b411 --- /dev/null +++ b/lib/Conversion/GpuToOCLSPIRV/GPUToSPIRVPass.cpp @@ -0,0 +1,167 @@ +//===- GPUToSPIRVPass.cpp - GPU to SPIR-V Passes --------------------------===// +// +// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. +// See https://llvm.org/LICENSE.txt for license information. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception +// +//===----------------------------------------------------------------------===// +// +// This file implements a pass to convert a kernel function in the GPU Dialect +// into a spirv.module operation. +// +//===----------------------------------------------------------------------===// + +#include "comet/Conversion/GpuToOCLSPIRV/GPUToSPIRVPass.h" + +#include "mlir/Conversion/ArithToSPIRV/ArithToSPIRV.h" +#include "mlir/Conversion/FuncToSPIRV/FuncToSPIRV.h" +#include "mlir/Conversion/GPUToSPIRV/GPUToSPIRV.h" +#include "mlir/Conversion/MemRefToSPIRV/MemRefToSPIRV.h" +#include "mlir/Conversion/SCFToSPIRV/SCFToSPIRV.h" +#include "mlir/Conversion/VectorToSPIRV/VectorToSPIRV.h" +#include "mlir/Dialect/Func/IR/FuncOps.h" +#include "mlir/Dialect/GPU/IR/GPUDialect.h" +#include "mlir/Dialect/SPIRV/IR/SPIRVDialect.h" +#include "mlir/Dialect/SPIRV/IR/SPIRVOps.h" +#include "mlir/Dialect/SPIRV/Transforms/SPIRVConversion.h" +#include "mlir/IR/PatternMatch.h" + +namespace mlir { +#define GEN_PASS_DEF_CONVERTGPUTOSPIRV +#include "mlir/Conversion/Passes.h.inc" +} // namespace mlir + +using namespace mlir; + +namespace { +/// Pass to lower GPU Dialect to SPIR-V. The pass only converts the gpu.func ops +/// inside gpu.module ops. i.e., the function that are referenced in +/// gpu.launch_func ops. For each such function +/// +/// 1) Create a spirv::ModuleOp, and clone the function into spirv::ModuleOp +/// (the original function is still needed by the gpu::LaunchKernelOp, so cannot +/// replace it). +/// +/// 2) Lower the body of the spirv::ModuleOp. +struct GPUToSPIRVPass final : impl::ConvertGPUToSPIRVBase { + explicit GPUToSPIRVPass(bool mapMemorySpace, bool use64bitIndex) + : mapMemorySpace(mapMemorySpace), use64bitIndex(use64bitIndex) {} + void runOnOperation() override; + +private: + bool mapMemorySpace, use64bitIndex; +}; + +void GPUToSPIRVPass::runOnOperation() { + MLIRContext *context = &getContext(); + ModuleOp module = getOperation(); + + SmallVector gpuModules; + OpBuilder builder(context); + + auto targetEnvSupportsKernelCapability = [](gpu::GPUModuleOp moduleOp) { + Operation *gpuModule = moduleOp.getOperation(); + auto targetAttr = spirv::lookupTargetEnvOrDefault(gpuModule); + spirv::TargetEnv targetEnv(targetAttr); + return targetEnv.allows(spirv::Capability::Kernel); + }; + + module.walk([&](gpu::GPUModuleOp moduleOp) { + // Clone each GPU kernel module for conversion, given that the GPU + // launch op still needs the original GPU kernel module. + // For Vulkan Shader capabilities, we insert the newly converted SPIR-V + // module right after the original GPU module, as that's the expectation of + // the in-tree Vulkan runner. + // For OpenCL Kernel capabilities, we insert the newly converted SPIR-V + // module inside the original GPU module, as that's the expectaion of the + // normal GPU compilation pipeline. + if (targetEnvSupportsKernelCapability(moduleOp)) { + builder.setInsertionPoint(moduleOp.getBody(), + moduleOp.getBody()->begin()); + } else { + builder.setInsertionPoint(moduleOp.getOperation()); + } + gpuModules.push_back(builder.clone(*moduleOp.getOperation())); + }); + + // Run conversion for each module independently as they can have different + // TargetEnv attributes. + for (Operation *gpuModule : gpuModules) { + spirv::TargetEnvAttr targetAttr = + spirv::lookupTargetEnvOrDefault(gpuModule); + + // Map MemRef memory space to SPIR-V storage class first if requested. + if (mapMemorySpace) { + spirv::MemorySpaceToStorageClassMap memorySpaceMap = + targetEnvSupportsKernelCapability( + dyn_cast(gpuModule)) + ? spirv::mapMemorySpaceToOpenCLStorageClass + : spirv::mapMemorySpaceToVulkanStorageClass; + spirv::MemorySpaceToStorageClassConverter converter(memorySpaceMap); + spirv::convertMemRefTypesAndAttrs(gpuModule, converter); + + // Check if there are any illegal ops remaining. + std::unique_ptr target = + spirv::getMemorySpaceToStorageClassTarget(*context); + gpuModule->walk([&target, this](Operation *childOp) { + if (target->isIllegal(childOp)) { + childOp->emitOpError("failed to legalize memory space"); + signalPassFailure(); + return WalkResult::interrupt(); + } + return WalkResult::advance(); + }); + } + + std::unique_ptr target = + SPIRVConversionTarget::get(targetAttr); + + SPIRVConversionOptions options; + options.use64bitIndex = this->use64bitIndex; + SPIRVTypeConverter typeConverter(targetAttr, options); + populateMMAToSPIRVCoopMatrixTypeConversion(typeConverter); + + RewritePatternSet patterns(context); + populateGPUToSPIRVPatterns(typeConverter, patterns); + populateGpuWMMAToSPIRVCoopMatrixKHRConversionPatterns(typeConverter, + patterns); + + // TODO: Change SPIR-V conversion to be progressive and remove the following + // patterns. + ScfToSPIRVContext scfContext; + populateSCFToSPIRVPatterns(typeConverter, scfContext, patterns); + mlir::arith::populateArithToSPIRVPatterns(typeConverter, patterns); + populateMemRefToSPIRVPatterns(typeConverter, patterns); + populateFuncToSPIRVPatterns(typeConverter, patterns); + populateVectorToSPIRVPatterns(typeConverter, patterns); + + if (failed(applyFullConversion(gpuModule, *target, std::move(patterns)))) + return signalPassFailure(); + } + + // For OpenCL, the gpu.func op in the original gpu.module op needs to be + // replaced with an empty func.func op with the same arguments as the gpu.func + // op. The func.func op needs gpu.kernel attribute set. + module.walk([&](gpu::GPUModuleOp moduleOp) { + if (targetEnvSupportsKernelCapability(moduleOp)) { + moduleOp.walk([&](gpu::GPUFuncOp funcOp) { + builder.setInsertionPoint(funcOp); + auto newFuncOp = builder.create( + funcOp.getLoc(), funcOp.getName(), funcOp.getFunctionType()); + auto entryBlock = newFuncOp.addEntryBlock(); + builder.setInsertionPointToEnd(entryBlock); + builder.create(funcOp.getLoc()); + newFuncOp->setAttr(gpu::GPUDialect::getKernelFuncAttrName(), + builder.getUnitAttr()); + funcOp.erase(); + }); + } + }); +} + +} // namespace + +std::unique_ptr> +mlir::createConvertGPUToSPIRVPass2(bool mapMemorySpace, bool use64bitIndex) { + return std::make_unique(mapMemorySpace, use64bitIndex); +} diff --git a/lib/Conversion/GpuToOCLSPIRV/GpuToOCLSPIRVPass.cpp b/lib/Conversion/GpuToOCLSPIRV/GpuToOCLSPIRVPass.cpp new file mode 100644 index 00000000..151c3391 --- /dev/null +++ b/lib/Conversion/GpuToOCLSPIRV/GpuToOCLSPIRVPass.cpp @@ -0,0 +1,96 @@ +#include "mlir/Conversion/AffineToStandard/AffineToStandard.h" +#include "mlir/Dialect/SPIRV/IR/SPIRVOps.h" +#include "mlir/Dialect/SPIRV/IR/TargetAndABI.h" +#include "mlir/Dialect/SPIRV/Transforms/Passes.h" +#include "mlir/IR/BuiltinOps.h" +#include "mlir/Pass/Pass.h" +#include "mlir/Dialect/Arith/IR/Arith.h" +#include "mlir/Dialect/SCF/IR/SCF.h" +#include "mlir/Dialect/SPIRV/IR/SPIRVDialect.h" +#include "mlir/Dialect/SPIRV/Transforms/SPIRVConversion.h" +#include "mlir/Dialect/GPU/IR/GPUDialect.h" +#include "mlir/Pass/PassManager.h" +#include "mlir/Support/LLVM.h" +#include "mlir/Target/SPIRV/Serialization.h" +#include +#include +#include "comet/Conversion/GpuToOCLSPIRV/GpuToOCLSPIRVPass.h" +#include "comet/Conversion/GpuToOCLSPIRV/GPUToSPIRVPass.h" + +#define GEN_PASS_CLASSES +#include "comet/Conversion/GpuToOCLSPIRV/Passes.h.inc" + +mlir::spirv::TargetEnvAttr getKernelEnv(mlir::Operation *op) { + auto triple = mlir::spirv::VerCapExtAttr::get(mlir::spirv::Version::V_1_0, + {mlir::spirv::Capability::Kernel, mlir::spirv::Capability::Addresses, mlir::spirv::Capability::Int64, mlir::spirv::Capability::Float64, }, + mlir::ArrayRef(), op->getContext()); + return mlir::spirv::TargetEnvAttr::get(triple, mlir::spirv::getDefaultResourceLimits(op->getContext())); +} +class ConvertGpuToOCLSPIRV + : public ConvertGpuKernelToOCLSPIRVPassBase { +public: + ConvertGpuToOCLSPIRV() = default; + ConvertGpuToOCLSPIRV(int block_size_x, int block_size_y, int block_size_z, const char* outpath) { + this->blockX = block_size_x; + this->blockY = block_size_y; + this->blockZ = block_size_z; + this->spirv_bin_path = outpath; + } + void runOnOperation() override { + mlir::MLIRContext *context = &getContext + (); + mlir::ModuleOp module = getOperation(); + module->setAttr(mlir::spirv::getTargetEnvAttrName(), getKernelEnv(module)); + module->walk([context, this](mlir::gpu::GPUFuncOp gpuFunc) { + mlir::StringRef attrName = mlir::spirv::getEntryPointABIAttrName(); + if (!mlir::gpu::GPUDialect::isKernel(gpuFunc) || gpuFunc->getAttr(attrName)) + { + return; + } + mlir::SmallVector workgroupSizeVec = {this->blockX , this->blockY , this->blockZ}; + gpuFunc->setAttr(attrName, + mlir::spirv::getEntryPointABIAttr(context, workgroupSizeVec)); + }); + + mlir::PassManager pm(context); + pm.addPass(mlir::createConvertGPUToSPIRVPass2(true, true)); + + if (failed(pm.run(module))) { + signalPassFailure(); + return; + } + + + mlir::SmallVector spirvModules; + module->walk([&spirvModules](mlir::spirv::ModuleOp module) { + spirvModules.push_back(module); + }); + + for(mlir::spirv::ModuleOp module: spirvModules) + { + mlir::PassManager pm(module.getContext()); + pm.addPass(mlir::spirv::createSPIRVLowerABIAttributesPass()); + pm.addPass(mlir::spirv::createSPIRVUpdateVCEPass()); + + if (failed(pm.run(module))) { + signalPassFailure(); + return; + } + + mlir::SmallVector binary; + if (!failed(mlir::spirv::serialize(module, binary))) + { + std::string pathname(this->spirv_bin_path); + pathname += module.getName()->data(); + pathname += ".bin"; + std::ofstream binout(pathname, std::ios::binary); + binout.write(reinterpret_cast(binary.data()), binary.size() * sizeof(uint32_t)); + } + } + } +}; + +std::unique_ptr> +mlir::comet::createConvertGPUKernelToOCLSPIRVPass(int block_size_x, int block_size_y, int block_size_z, const char* spirv_bin_path) { + return std::make_unique<::ConvertGpuToOCLSPIRV>(block_size_x, block_size_y, block_size_z,spirv_bin_path); +} \ No newline at end of file diff --git a/lib/Conversion/GpuToTriton/CMakeLists.txt b/lib/Conversion/GpuToTriton/CMakeLists.txt index 947c6e4c..d905d327 100644 --- a/lib/Conversion/GpuToTriton/CMakeLists.txt +++ b/lib/Conversion/GpuToTriton/CMakeLists.txt @@ -12,6 +12,6 @@ add_llvm_library(COMETGpuToTriton TritonIR TritonGPUIR TritonGPUTransforms - TritonNvidiaGPUTransforms - NVGPUIR + #TritonNvidiaGPUTransforms + #NVGPUIR ) diff --git a/lib/Conversion/GpuToTriton/GpuToTritonConversion.cpp b/lib/Conversion/GpuToTriton/GpuToTritonConversion.cpp index 90c12368..6a915433 100644 --- a/lib/Conversion/GpuToTriton/GpuToTritonConversion.cpp +++ b/lib/Conversion/GpuToTriton/GpuToTritonConversion.cpp @@ -1,3 +1,24 @@ +// +// Copyright 2022 Battelle Memorial Institute +// +// Redistribution and use in source and binary forms, with or without modification, +// are permitted provided that the following conditions are met: +// +// 1. Redistributions of source code must retain the above copyright notice, this list of conditions +// and the following disclaimer. +// +// 2. Redistributions in binary form must reproduce the above copyright notice, this list of conditions +// and the following disclaimer in the documentation and/or other materials provided with the distribution. +// +// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED +// WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR +// PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, +// INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE +// GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, +// WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +// OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. +// + #include "mlir/Dialect/Affine/IR/AffineOps.h" #include "mlir/Dialect/Arith/IR/Arith.h" #include "mlir/Dialect/GPU/IR/GPUDialect.h" @@ -158,13 +179,15 @@ mlir::comet::GpuConversionTarget2::GpuConversionTarget2( } if(prev != op.getType()) { - if(!prev.isa() && !op.getType().isa()) + auto prevRankedT = dyn_cast(prev); + auto opRankedT = dyn_cast(op.getType()); + if(!prevRankedT && !opRankedT) { return true; } - else if(prev.isa() && op.getType().isa()) + else if(prevRankedT && opRankedT) { - if(prev.cast().getShape() == op.getType().cast().getShape()) + if(prevRankedT.getShape() == opRankedT.getShape()) { return true; } @@ -196,13 +219,15 @@ mlir::comet::GpuConversionTarget2::GpuConversionTarget2( } if(prev != op.getType()) { - if(!prev.isa() && !op.getType().isa()) + auto prevRankedT = dyn_cast(prev); + auto opRankedT = dyn_cast(op.getType()); + if(!prevRankedT && !opRankedT) { return true; } - else if(prev.isa() && op.getType().isa()) + else if(prevRankedT && opRankedT) { - if(prev.cast().getShape() == op.getType().cast().getShape()) + if(prevRankedT.getShape() == opRankedT.getShape()) { return true; } @@ -222,17 +247,19 @@ mlir::comet::GpuConversionTarget2::GpuConversionTarget2( { if(op.getTrueValue().getType() == op.getResult().getType() && op.getFalseValue().getType() == op.getTrueValue().getType()) { - if(op.getCondition().getType().isa() && !op.getFalseValue().getType().isa()) + auto falseT = dyn_cast(op.getFalseValue().getType()); + auto condT = dyn_cast(op.getCondition().getType()); + if(condT && !falseT) { return false; } - else if(!op.getCondition().getType().isa() && op.getFalseValue().getType().isa()) + else if(!condT && falseT) { return false; } - else if(op.getCondition().getType().isa() && op.getFalseValue().getType().isa()) + else if(condT && falseT) { - if(op.getCondition().getType().cast().getShape() == op.getFalseValue().getType().cast().getShape()) + if(condT.getShape() == falseT.getShape()) { return true; } @@ -318,7 +345,7 @@ mlir::comet::GpuConversionTarget2::GpuConversionTarget2( { for(auto opr: op.getArgumentTypes()) { - if(opr.isa()) + if(isa(opr)) { return false; } @@ -342,7 +369,7 @@ mlir::comet::GpuTypeConverter::GpuTypeConverter(MLIRContext *context) addConversion([this](MemRefType memrefType, SmallVectorImpl &results) -> LogicalResult { - if(memrefType.getElementType().isa()) + if(isa(memrefType.getElementType())) { results.push_back(mlir::triton::PointerType::get( mlir::IntegerType::get(this->context, 32), 1)); } diff --git a/lib/Conversion/GpuToTriton/GpuToTritonPass.cpp b/lib/Conversion/GpuToTriton/GpuToTritonPass.cpp index a89872b8..5771783d 100644 --- a/lib/Conversion/GpuToTriton/GpuToTritonPass.cpp +++ b/lib/Conversion/GpuToTriton/GpuToTritonPass.cpp @@ -1,4 +1,25 @@ - +// +// Copyright 2022 Battelle Memorial Institute +// +// Redistribution and use in source and binary forms, with or without modification, +// are permitted provided that the following conditions are met: +// +// 1. Redistributions of source code must retain the above copyright notice, this list of conditions +// and the following disclaimer. +// +// 2. Redistributions in binary form must reproduce the above copyright notice, this list of conditions +// and the following disclaimer in the documentation and/or other materials provided with the distribution. +// +// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED +// WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR +// PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, +// INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE +// GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, +// WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +// OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. +// + +#include #include #include #include "comet/Conversion/GpuToTriton/GpuToTritonPass.h" @@ -12,6 +33,7 @@ #include "mlir/Dialect/GPU/IR/GPUDialect.h" #include "mlir/Dialect/SCF/IR/SCF.h" #include "mlir/Dialect/Affine/IR/AffineOps.h" +#include "mlir/Support/LLVM.h" #include "mlir/Transforms/DialectConversion.h" #include "mlir/Pass/PassManager.h" #include "mlir/Transforms/Passes.h" @@ -94,9 +116,9 @@ auto lhsT = lhs.getType(); auto rhsT = rhs.getType(); if(lhsT != rhsT) { - if(auto lhsTensor = lhsT.dyn_cast()) + if(auto lhsTensor = mlir::dyn_cast(lhsT)) { - if(auto rhsTensor = rhsT.dyn_cast() ) // + if(auto rhsTensor = mlir::dyn_cast(rhsT) ) // { if(lhsTensor.getRank() != rhsTensor.getRank()) @@ -173,7 +195,7 @@ auto lhsT = lhs.getType(); } else { - if(auto rhsTensor = rhsT.dyn_cast()) + if(auto rhsTensor = dyn_cast(rhsT)) { lhs = rewriter.createOrFold(op->getLoc(), RankedTensorType::get(rhsTensor.getShape(), lhs.getType()), lhs); } @@ -403,7 +425,7 @@ LogicalResult convertMemoryOp(Operation* op, ConversionPatternRewriter &rewriter if(map.find(exp.getAsOpaquePointer()) == map.end()) { // COMET_ERRS << "HERE\n"; - if(getSymOrDimOperand(aaffineop, exp).getType().isa()) + if(isa(getSymOrDimOperand(aaffineop, exp).getType())) { // COMET_ERRS << "HERE\n"; // COMET_ERRS << getSymOrDimOperand(aaffineop, exp); @@ -581,14 +603,14 @@ LogicalResult convertMemoryOp(Operation* op, ConversionPatternRewriter &rewriter } mlir::Value ptr; - if(RankedTensorType t = op->getOperand(mem_offset).getType().dyn_cast()) + if(RankedTensorType t = dyn_cast(op->getOperand(mem_offset).getType())) { - if(t.getElementType().isa()) + if(mlir::isa(t.getElementType())) { ptr = rewriter.create(op->getLoc(), RankedTensorType::get(t.getShape(), rewriter.getI32Type()), op->getOperand(mem_offset)); } } - else if(op->getOperand(mem_offset).getType().isa()) + else if(mlir::isa(op->getOperand(mem_offset).getType())) { ptr = rewriter.create(op->getLoc(), rewriter.getI32Type(), op->getOperand(mem_offset)); } @@ -598,9 +620,9 @@ LogicalResult convertMemoryOp(Operation* op, ConversionPatternRewriter &rewriter } // auto ptr_array = rewriter.create(op->getLoc(), RankedTensorType::get(block_sizes, ptr.getType()), ptr); mlir::Value ptr_array; - if(op->getOperand(mem_offset +1).getType().isa()) + if(mlir::isa(op->getOperand(mem_offset +1).getType())) { - ptr_array = rewriter.create(op->getLoc(), RankedTensorType::get(op->getOperand(mem_offset +1).getType().cast().getShape(), ptr.getType()), ptr).getResult(); + ptr_array = rewriter.create(op->getLoc(), RankedTensorType::get(cast(op->getOperand(mem_offset +1).getType()).getShape(), ptr.getType()), ptr).getResult(); } else { @@ -627,10 +649,10 @@ LogicalResult convertMemoryOp(Operation* op, ConversionPatternRewriter &rewriter else { mlir::Value toStore; - if (op->getOperand(0).getType().dyn_cast() ) { + if (dyn_cast(op->getOperand(0).getType()) ) { toStore = op->getOperand(0); } - else if(auto ptr_array_type = final_ptr_array.getResult().getType().dyn_cast(); ptr_array_type && !op->getOperand(0).getType().dyn_cast() ) { + else if(auto ptr_array_type = dyn_cast(final_ptr_array.getResult().getType()); ptr_array_type && !dyn_cast(op->getOperand(0).getType()) ) { toStore = rewriter.create(op->getLoc(), RankedTensorType::get(ptr_array_type.getShape(), op->getOperand(0).getType()), op->getOperand(0)); } mlir::triton::StoreOp storeVal; @@ -661,15 +683,15 @@ LogicalResult ExpandScalarTensorArithOp(T op, ConversionPatternRewriter &rewrite Operation* expandedOp = NULL; if(lhs_type != rhs_type) { - if(auto lhs_tensor_type = lhs_type.dyn_cast_or_null()) + if(auto lhs_tensor_type = mlir::dyn_cast_if_present(lhs_type)) { - if(!rhs_type.isa() && lhs_tensor_type.getElementType() == rhs_type) + if(!mlir::dyn_cast_if_present(rhs_type) && lhs_tensor_type.getElementType() == rhs_type) { auto scalarExpanded = rewriter.create(op->getLoc(), lhs.getType(), rhs); expandedOp = rewriter.create(op->getLoc(), lhs, scalarExpanded); } } - else if(auto rhs_tensor_type = rhs_type.dyn_cast_or_null()) + else if(auto rhs_tensor_type = mlir::dyn_cast_if_present(rhs_type)) { if(rhs_tensor_type.getElementType() == lhs_type) { @@ -715,24 +737,26 @@ LogicalResult ExpandScalarTensorSelectOp(arith::SelectOp& op, ConversionPatternR mlir::Type res_type = res.getType(); Operation* expandedOp = NULL; + auto rankedCondT = mlir::dyn_cast(cond_type); + if(lhs_type != rhs_type) { - if(auto lhs_tensor_type = lhs_type.dyn_cast_or_null()) + if(auto lhs_tensor_type = mlir::dyn_cast(lhs_type)) { - if(!cond_type.isa()) + if(!isa(cond_type)) { auto expandedCond = rewriter.create(op->getLoc(), RankedTensorType::get(lhs_tensor_type.getShape(), cond.getType()), cond); cond = expandedCond.getResult(); } - if(!rhs_type.isa() && lhs_tensor_type.getElementType() == rhs_type) + if(!isa(rhs_type) && lhs_tensor_type.getElementType() == rhs_type) { auto scalarExpanded = rewriter.create(op->getLoc(), lhs.getType(), rhs); expandedOp = rewriter.create(op->getLoc(), cond, lhs, scalarExpanded); } } - else if(auto rhs_tensor_type = rhs_type.dyn_cast_or_null()) + else if(auto rhs_tensor_type = mlir::dyn_cast(rhs_type)) { - if(!cond_type.isa()) + if(!isa(cond_type)) { auto expandedCond = rewriter.create(op->getLoc(), RankedTensorType::get(rhs_tensor_type.getShape(), cond.getType()), cond); cond = expandedCond.getResult(); @@ -746,17 +770,17 @@ LogicalResult ExpandScalarTensorSelectOp(arith::SelectOp& op, ConversionPatternR } else if (lhs_type != res_type) { - auto rhs_tensor_type = rhs_type.dyn_cast_or_null(); - if(rhs_tensor_type && !cond_type.isa()) + auto rhs_tensor_type = mlir::dyn_cast(rhs_type); + if(rhs_tensor_type && !rankedCondT) { auto expandedCond = rewriter.create(op->getLoc(), RankedTensorType::get(rhs_tensor_type.getShape(), cond.getType()), cond); cond = expandedCond.getResult(); } expandedOp = rewriter.create(op->getLoc(), cond, lhs, rhs); } - else if (!cond_type.isa() && rhs_type.isa()) + else if (!rankedCondT && mlir::isa(rhs_type)) { - auto rhs_tensor_type = rhs_type.cast(); + auto rhs_tensor_type = mlir::dyn_cast(rhs_type); auto expandedCond = rewriter.create(op->getLoc(), RankedTensorType::get(rhs_tensor_type.getShape(), cond.getType()), cond); cond = expandedCond.getResult(); expandedOp = rewriter.create(op->getLoc(), cond, lhs, rhs); @@ -797,15 +821,15 @@ LogicalResult ExpandScalarTensorArithCmpOp(T op, ConversionPatternRewriter &rewr Operation* expandedOp = NULL; if(lhs_type != rhs_type) { - if(auto lhs_tensor_type = lhs_type.dyn_cast_or_null()) + if(auto lhs_tensor_type = mlir::dyn_cast_if_present(lhs_type)) { - if(!rhs_type.isa() && lhs_tensor_type.getElementType() == rhs_type) + if(!mlir::isa(rhs_type) && lhs_tensor_type.getElementType() == rhs_type) { auto scalarExpanded = rewriter.create(op->getLoc(), lhs.getType(), rhs); expandedOp = rewriter.create(op->getLoc(), op->template getAttrOfType("predicate"), lhs, scalarExpanded); } } - else if(auto rhs_tensor_type = rhs_type.dyn_cast_or_null()) + else if(auto rhs_tensor_type = mlir::dyn_cast_if_present(rhs_type)) { if(rhs_tensor_type.getElementType() == lhs_type) { @@ -849,7 +873,7 @@ void iterateOperations(Operation *op, PatternRewriter& rewriter) { for (Block &block : region) { for(auto arg: block.getArguments()) { - if(arg.getType().isa()) + if(mlir::isa(arg.getType())) { arg.setType(IntegerType::get(arg.getContext(), 32)); } @@ -887,7 +911,7 @@ void convertBlockArgTypes(mlir::triton::FuncOp func, PatternRewriter& rewriter) for (Block &block : func) { for(auto arg: block.getArguments()) { - if(arg.getType().isa()) + if(mlir::isa(arg.getType())) { arg.setType(IntegerType::get(arg.getContext(), 32)); } @@ -980,7 +1004,7 @@ class RewriteReduction : public OpConversionPattern { continue; } isAncestor(offset.getDefiningOp()->getOperand(0).getDefiningOp(), forOp.getLoopRegions()[0]->getArgument(0), local_ops_chain, ops_chain); - mlir::Value base, step; + mlir::Value base; mlir::Value res; std::vector stepsMul, stepsAdd; if(!local_ops_chain.empty()) @@ -1018,11 +1042,11 @@ class RewriteReduction : public OpConversionPattern { { if(s.getType() != offset.getType()) { - if(auto sTensor = s.getType().dyn_cast()) + if(auto sTensor = mlir::dyn_cast(s.getType())) { - if(sTensor.getRank() != offset.getType().cast().getRank()) + if(sTensor.getRank() != cast(offset.getType()).getRank()) { - if(sTensor.getDimSize(0) == offset.getType().cast().getDimSize(0)) + if(sTensor.getDimSize(0) == cast(offset.getType()).getDimSize(0)) { s = rewriter.create(op->getLoc(), s, 1); } @@ -1042,7 +1066,7 @@ class RewriteReduction : public OpConversionPattern { res = rewriter.create(op->getLoc(), res, temp); } } - else if (s.getType().isa() ) + else if (mlir::isa(s.getType()) ) { auto temp = rewriter.create(op->getLoc(), offset.getType(), s); @@ -1148,7 +1172,7 @@ class RewriteReduction : public OpConversionPattern { if (!forOp.getOps().empty()) { dotOp = *forOp.getOps().begin(); - auto elementType = mlir::isa(dotOp->getResultTypes()[0]) ? dotOp->getResultTypes()[0].cast().getElementType() : dotOp->getResultTypes()[0]; + auto elementType = mlir::isa(dotOp->getResultTypes()[0]) ? cast(dotOp->getResultTypes()[0]).getElementType() : dotOp->getResultTypes()[0]; auto init = rewriter.create(forOp->getLoc(), elementType, rewriter.getZeroAttr(elementType)); auto initSplat = rewriter.create(forOp->getLoc(), dotOp->getResultTypes()[0], init); dotOp.getCMutable().assign(initSplat); @@ -1157,7 +1181,7 @@ class RewriteReduction : public OpConversionPattern { if (!forOp.getOps().empty()) { reduceOp = *forOp.getOps().begin(); - auto elementType = mlir::isa(reduceOp->getUsers().begin()->getResultTypes()[0]) ? reduceOp->getUsers().begin()->getResultTypes()[0].cast().getElementType() : reduceOp->getUsers().begin()->getResultTypes()[0]; + auto elementType = mlir::isa(reduceOp->getUsers().begin()->getResultTypes()[0]) ? cast(reduceOp->getUsers().begin()->getResultTypes()[0]).getElementType() : reduceOp->getUsers().begin()->getResultTypes()[0]; auto init = rewriter.create(forOp->getLoc(), elementType, rewriter.getZeroAttr(elementType)); auto initSplat = rewriter.create(forOp->getLoc(), reduceOp->getUsers().begin()->getResultTypes()[0], init); basePtrs.push_back(initSplat); @@ -1237,6 +1261,7 @@ class FinalizeTritonFuncOp : public OpConversionPattern { LogicalResult matchAndRewrite(mlir::triton::FuncOp op, OpAdaptor adaptor, ConversionPatternRewriter &rewriter) const override { + rewriter.startOpModification(op); bool changed = false; std::vector newArgTypes; @@ -1244,14 +1269,14 @@ class FinalizeTritonFuncOp : public OpConversionPattern { { auto argT = arg.getType(); - if(argT.isa()) + if(mlir::isa(argT)) { changed = true; arg.setType(rewriter.getI32Type()); } - else if(RankedTensorType t= argT.dyn_cast()) + else if(RankedTensorType t= dyn_cast(argT)) { - if(t.getElementType().isa()) + if(mlir::isa(t.getElementType())) { changed = true; arg.setType(RankedTensorType::get(t.getShape(), rewriter.getI32Type())); @@ -1269,11 +1294,11 @@ class FinalizeTritonFuncOp : public OpConversionPattern { { newArgTypes.push_back(arg.getType()); } - rewriter.startRootUpdate(op); + rewriter.startOpModification(op); auto functype= mlir::FunctionType::get(getContext(), newArgTypes, {}); op.setType(functype); convertBlockArgTypes(op, rewriter); - rewriter.finalizeRootUpdate(op); + rewriter.finalizeOpModification(op); return success(); } @@ -1326,8 +1351,8 @@ class ArithOpPattern : public OpConversionPattern auto addOp = cast(*op->getUsers().begin()); auto loadRes = addOp.getLhs() == op ? addOp.getRhs() : addOp.getLhs(); rewriter.setInsertionPointAfter(addOp); - auto lhsTensor = op->getOperand(0).getType().cast(); - auto rhsTensor = op->getOperand(1).getType().cast(); + auto lhsTensor = cast(op->getOperand(0).getType()); + auto rhsTensor = cast(op->getOperand(1).getType()); mlir::Operation* dotOp; if (lhsTensor.getRank() == 2 && (rhsTensor.getRank() == 1 || (rhsTensor.getRank() == 2 && (rhsTensor.getDimSize(0) == 1 || rhsTensor.getDimSize(1) == 1)) )) { @@ -1358,9 +1383,9 @@ class ArithOpPattern : public OpConversionPattern } } else { - dotOp = rewriter.create(op->getLoc(), op->getOperand(0), op->getOperand(1), loadRes, rewriter.getBoolAttr(true), rewriter.getI32IntegerAttr(1)); + // dotOp = rewriter.create(op->getLoc(), op->getOperand(0), op->getOperand(1), loadRes, rewriter.getBoolAttr(true), rewriter.getI32IntegerAttr(1)); + dotOp = rewriter.create(op->getLoc(), op->getOperand(0), op->getOperand(1), loadRes); } - // if(op->getOperand(1).getType().cast().) // addOp->replaceAllUsesWith(dotOp); rewriter.replaceAllUsesWith(addOp, dotOp->getResult(0)); @@ -1514,13 +1539,13 @@ class GpuFuncOpToTritonFuncOp : public OpConversionPattern rewriter.create(rewriter.getUnknownLoc()); - rewriter.startRootUpdate(TTFunc); + rewriter.startOpModification(TTFunc); for(auto arg : llvm::zip(gpuFuncOp.getFunctionBody().back().getArguments(), TTFunc.getFunctionBody().back().getArguments())) { std::get<0>(arg).replaceAllUsesWith(std::get<1>(arg)); } - rewriter.finalizeRootUpdate(TTFunc); + rewriter.finalizeOpModification(TTFunc); rewriter.mergeBlocks(&gpuFuncOp.getFunctionBody().back(), &TTFunc.getFunctionBody().front(), TTFunc.getFunctionBody().front().getArguments()); rewriter.eraseOp(gpuFuncOp); @@ -1686,7 +1711,7 @@ class BlockifyReduction : public OpConversionPattern { { // COMET_ERRS << launchOp.getKernelModuleName() <<" vs " << origin_kernel_module_name << "\n"; // COMET_ERRS << launchOp.getKernelName() <<" vs " << origin_kernel_name << "\n"; - if(launchOp.getKernelModuleName().strref().equals(origin_kernel_module_name) && launchOp.getKernelName().strref().equals(origin_kernel_name)) + if(launchOp.getKernelModuleName().strref().compare(origin_kernel_module_name) == 0 && launchOp.getKernelName().strref().compare(origin_kernel_name) == 0) { // COMET_ERRS << "FOUND LAUNCHOP" << origin_kernel_fullname << "\n"; auto block_size_r = cast(launchOp.getKernelOperand(input_block_size.getArgNumber()-2).getDefiningOp()).value(); @@ -1802,9 +1827,9 @@ class ConvertGpuToTriton auto kernel_mod_name = launchOp.getKernelModuleName(); // COMET_ERRS < bool {return modOp.getName().equals(kernel_mod_name); }); + auto kernel_module = *std::find_if(gpuModules.begin(), gpuModules.end(), [&kernel_mod_name] (mlir::gpu::GPUModuleOp modOp) -> bool {return modOp.getName().compare(kernel_mod_name) == 0; }); auto gpufuncOps = kernel_module.getOps(); - auto gpufuncOp = *std::find_if(gpufuncOps.begin(), gpufuncOps.end(), [&kernel_name](mlir::gpu::GPUFuncOp funcOp) -> bool{ return funcOp->getAttrOfType("sym_name").strref().equals(kernel_name) ;}); + auto gpufuncOp = *std::find_if(gpufuncOps.begin(), gpufuncOps.end(), [&kernel_name](mlir::gpu::GPUFuncOp funcOp) -> bool{ return funcOp->getAttrOfType("sym_name").strref().compare(kernel_name) == 0 ;}); auto b_size_x = dyn_cast(launchOp.getBlockSizeX().getDefiningOp()).value(); auto b_size_y = dyn_cast(launchOp.getBlockSizeY().getDefiningOp()).value(); if (gpufuncOp->hasAttr("simplified")) @@ -1828,19 +1853,19 @@ class ConvertGpuToTriton { if(arith::ConstantIndexOp k = llvm::dyn_cast_if_present(k_op.value().getDefiningOp())) { - // if(k.value() == 0 || k.value() == 1) - // { - // bitvector.set(k_op.index()); - // builder.setInsertionPointToStart(&gpufuncOp.getFunctionBody().getBlocks().front()); - // gpufuncOp.getFunctionBody().getArgument(k_op.index()).replaceAllUsesWith(builder.create(gpufuncOp->getLoc(), k.value())); - // } - // else if(k.value() == b_size_x || k.value() == b_size_y) - // { - // bitvector.set(k_op.index()); - // builder.setInsertionPointToStart(&funcOp.getFunctionBody().getBlocks().front()); - // auto k_val = (k.value() == b_size_x) ? b_size_x : b_size_y; - // funcOp.getFunctionBody().getArgument(k_op.index()).replaceAllUsesWith(builder.create(funcOp->getLoc(), k_val)); - // } + if(k.value() == 0 || k.value() == 1) + { + bitvector.set(k_op.index()); + builder.setInsertionPointToStart(&gpufuncOp.getFunctionBody().getBlocks().front()); + gpufuncOp.getFunctionBody().getArgument(k_op.index()).replaceAllUsesWith(builder.create(gpufuncOp->getLoc(), k.value())); + } + else if(k.value() == b_size_x || k.value() == b_size_y) + { + bitvector.set(k_op.index()); + builder.setInsertionPointToStart(&gpufuncOp.getFunctionBody().getBlocks().front()); + auto k_val = (k.value() == b_size_x) ? b_size_x : b_size_y; + gpufuncOp.getFunctionBody().getArgument(k_op.index()).replaceAllUsesWith(builder.create(funcOp->getLoc(), k_val)); + } // else // { // break; @@ -1935,9 +1960,11 @@ class ConvertGpuToTriton }); patterns3.insert(context); - for(auto ttfuncOp: op.getOps()) - { - if (failed(applyPartialConversion(ttfuncOp, target3, std::move(patterns3)))) + std::vector allttfuncOps; + op->walk([&allttfuncOps](triton::FuncOp gMod) {allttfuncOps.push_back(gMod); }); + // for(auto ttfuncOp: op.getOps()) + { + if (failed(applyPartialConversion(allttfuncOps, target3, std::move(patterns3)))) { // COMET_ERRS << "Failed to Lower STCOutputLowering2\n"; signalPassFailure(); diff --git a/lib/Conversion/GpuUtils/CMakeLists.txt b/lib/Conversion/GpuUtils/CMakeLists.txt new file mode 100644 index 00000000..92ed15c7 --- /dev/null +++ b/lib/Conversion/GpuUtils/CMakeLists.txt @@ -0,0 +1,6 @@ +add_llvm_library(COMETGPUUtils + GpuUtils.cpp + + ADDITIONAL_HEADER_DIRS + ${COMET_MAIN_INCLUDE_DIR}/comet/GpuUtils + ) diff --git a/lib/Conversion/GpuUtils/GpuUtils.cpp b/lib/Conversion/GpuUtils/GpuUtils.cpp new file mode 100644 index 00000000..511c9105 --- /dev/null +++ b/lib/Conversion/GpuUtils/GpuUtils.cpp @@ -0,0 +1,530 @@ +#include "comet/Dialect/Utils/Utils.h" +#include "mlir/Dialect/Arith/IR/Arith.h" +#include "mlir/Dialect/GPU/IR/GPUDialect.h" +#include "mlir/Dialect/GPU/Transforms/Passes.h" +#include "mlir/Dialect/LLVMIR/LLVMDialect.h" +#include "mlir/Dialect/MemRef/IR/MemRef.h" +#include "mlir/Dialect/SCF/IR/SCF.h" +#include "mlir/IR/Attributes.h" +#include "mlir/IR/Builders.h" +#include "mlir/IR/BuiltinAttributes.h" +#include "mlir/IR/BuiltinOps.h" +#include "mlir/IR/ValueRange.h" +#include "triton/Dialect/Triton/IR/Dialect.h" +#include +#include +using namespace mlir; + +mlir::LogicalResult specializeGpuHost(mlir::OpBuilder& builder, mlir::ModuleOp modOp, std::string vendor_prefix) +{ + std::map funcs; + std::vector launchOps; + modOp->walk([&launchOps](mlir::gpu::LaunchFuncOp launchOp) { + launchOps.push_back(launchOp); + }); + + std::set initFuncs; + std::vector toAlloc; + + for (auto launchOp : launchOps) { + + if (initFuncs.find(launchOp->getParentOfType()) == + initFuncs.end()) { + builder.setInsertionPointToStart( + &launchOp->getParentOfType() + .getFunctionBody() + .front()); + Value gpu_code = builder.create( + launchOp->getLoc(), LLVM::LLVMPointerType::get(builder.getContext()), + "gpu_code"); + builder.create(launchOp.getLoc(), + vendor_prefix+"SetModuleImage", TypeRange(), + ValueRange({gpu_code})); + + initFuncs.insert(launchOp->getParentOfType()); + } + } + + std::vector gpuAllocs; + modOp->walk([&gpuAllocs](mlir::gpu::AllocOp gpuAllocOp) { + gpuAllocs.push_back(gpuAllocOp); + }); + + for (auto gpuAlloc : gpuAllocs) { + builder.setInsertionPoint(gpuAlloc->getBlock()->getTerminator()); + builder.create(gpuAlloc->getLoc(), ValueRange(), + gpuAlloc.getResult(0)); + + builder.setInsertionPointAfter(gpuAlloc); + Value allocSize; + if (gpuAlloc.getMemref().getType().hasStaticShape()) { + allocSize = builder + .create( + gpuAlloc->getLoc(), + gpuAlloc.getMemref().getType().getNumElements()) + .getResult(); + } else { + Value numElements = builder.create(gpuAlloc->getLoc(), 1); + for(auto dimSize :gpuAlloc.getDynamicSizes()) + { + numElements = builder.create(gpuAlloc->getLoc(), numElements, dimSize); + } + for(int64_t dim =0; dim < gpuAlloc.getMemref().getType().getRank(); dim++) + { + if(!gpuAlloc.getMemref().getType().isDynamicDim(dim)) + { + Value dimSize = builder.create(gpuAlloc->getLoc(), gpuAlloc.getMemref().getType().getDimSize(dim)); + numElements = builder.create(gpuAlloc->getLoc(), numElements, dimSize); + } + } + allocSize = numElements; + } + + // builder.create(gpuAlloc.getLoc(), gpuAlloc.getMemref(), + // 0); + if (FloatType floatType = mlir::dyn_cast( + gpuAlloc.getMemref().getType().getElementType())) { + int width = floatType.getWidth(); + auto cudaOp = builder.create( + gpuAlloc->getLoc(), vendor_prefix+"MallocF" + std::to_string(width), + TypeRange(builder.getIndexType()), ValueRange(allocSize)); + gpuAlloc->replaceAllUsesWith(cudaOp); + } else if (IntegerType intType = mlir::dyn_cast( + gpuAlloc.getMemref().getType().getElementType())) { + int width = intType.getWidth(); + auto cudaOp = builder.create( + gpuAlloc->getLoc(), vendor_prefix+"MallocI" + std::to_string(width), + TypeRange(builder.getIndexType()), ValueRange(allocSize)); + gpuAlloc->replaceAllUsesWith(cudaOp); + } else if (IndexType intType = mlir::dyn_cast( + gpuAlloc.getMemref().getType().getElementType())) { + auto cudaOp = builder.create( + gpuAlloc->getLoc(), vendor_prefix+"MallocI" + std::to_string(64), + TypeRange(builder.getIndexType()), ValueRange(allocSize)); + gpuAlloc->replaceAllUsesWith(cudaOp); + } + gpuAlloc->erase(); + } + + std::vector gpuCopies; + modOp->walk([&gpuCopies](mlir::gpu::MemcpyOp gpuCopy) { + gpuCopies.push_back(gpuCopy); + }); + + llvm::SmallMapVector, TypedValue, 4> memref_to_collapsed_memref; + for (auto cpy : gpuCopies) { + builder.setInsertionPoint(cpy); + auto hToD = builder.create(cpy->getLoc(), 0); + auto dToH = builder.create(cpy->getLoc(), 1); + + if (cpy.getOperand(0).getDefiningOp() && + (isa(cpy.getOperand(0).getDefiningOp()) && + cast(cpy.getOperand(0).getDefiningOp()) + .getCallee() + .starts_with(vendor_prefix+"Malloc"))) { + if(memref_to_collapsed_memref.find(mlir::cast>(cpy.getSrc())) == memref_to_collapsed_memref.end()) + { + TypedValue collapsedMemref = mlir::tensorAlgebra::collapseMemref(cpy.getSrc(), builder); + auto cast = builder.create( + cpy->getLoc(), + MemRefType::get({ShapedType::kDynamic}, + collapsedMemref.getType().getElementType()), + collapsedMemref); + memref_to_collapsed_memref[mlir::cast>(cpy.getSrc())] = cast.getResult(); + } + mlir::Value cast = memref_to_collapsed_memref[mlir::cast>(cpy.getSrc())]; + + + if (IntegerType intType = mlir::dyn_cast( + mlir::cast(cpy.getOperand(1).getType()) + .getElementType())) { + int width = intType.getWidth(); + builder.create( + cpy->getLoc(), vendor_prefix+"MemcpyI" + std::to_string(width), TypeRange(), + ValueRange({cpy.getOperand(0), cast, hToD})); + } else if (FloatType floatType = mlir::dyn_cast( + mlir::cast(cpy.getOperand(1).getType()) + .getElementType())) { + int width = floatType.getWidth(); + builder.create( + cpy->getLoc(), vendor_prefix+"MemcpyF" + std::to_string(width), TypeRange(), + ValueRange({cpy.getOperand(0), cast, hToD})); + } else if (IndexType indexType = mlir::dyn_cast( + mlir::cast(cpy.getOperand(1).getType()) + .getElementType())) { + builder.create( + cpy->getLoc(), vendor_prefix+"MemcpyIndex", TypeRange(), + ValueRange({cpy.getOperand(0), cast, hToD})); + } + } else if (cpy.getOperand(1).getDefiningOp() && + (isa(cpy.getOperand(1).getDefiningOp()) && + cast(cpy.getOperand(1).getDefiningOp()) + .getCallee() + .starts_with(vendor_prefix+"Malloc"))) { + if(memref_to_collapsed_memref.find(mlir::cast>(cpy.getDst())) == memref_to_collapsed_memref.end()) + { + TypedValue collapsedMemref = mlir::tensorAlgebra::collapseMemref(cpy.getDst(), builder); + auto cast = builder.create( + cpy->getLoc(), + MemRefType::get({ShapedType::kDynamic}, + collapsedMemref.getType().getElementType()), + collapsedMemref); + memref_to_collapsed_memref[mlir::cast>(cpy.getDst())] = cast.getResult(); + } + mlir::Value cast = memref_to_collapsed_memref[mlir::cast>(cpy.getDst())]; + + if (IntegerType intType = mlir::dyn_cast( + mlir::cast(cpy.getOperand(0).getType()) + .getElementType())) { + int width = intType.getWidth(); + builder.create( + cpy->getLoc(), vendor_prefix+"MemcpyI" + std::to_string(width), TypeRange(), + ValueRange({cpy.getOperand(1), cast, dToH})); + } else if (FloatType floatType = mlir::dyn_cast( + mlir::cast(cpy.getOperand(0).getType()) + .getElementType())) { + int width = floatType.getWidth(); + builder.create( + cpy->getLoc(), vendor_prefix+"MemcpyF" + std::to_string(width), TypeRange(), + ValueRange({cpy.getOperand(1), cast, dToH})); + } else if (IndexType indexType = mlir::dyn_cast( + mlir::cast(cpy.getOperand(0).getType()) + .getElementType())) { + builder.create( + cpy->getLoc(), vendor_prefix+"MemcpyIndex", TypeRange(), + ValueRange({cpy.getOperand(1), cast, dToH})); + } + } + + cpy->erase(); + } + + for (auto launchOp : launchOps) { + builder.setInsertionPoint(launchOp); + auto zeroIndex = + builder.create(launchOp->getLoc(), 0); + + int64_t numOps = launchOp.getNumKernelOperands(); + std::vector ptr_ops; + // memref::AllocaOp temp = builder.create( + // launchOp->getLoc(), + // MemRefType::get({1}, launchOp.getGridSizeX().getType())); + // builder.create(launchOp->getLoc(), + // launchOp.getGridSizeY(), temp, + // ValueRange({zeroIndex})); + // ptr_ops.push_back( + // builder.create( + // launchOp->getLoc(), temp)); + // temp = builder.create( + // launchOp->getLoc(), + // MemRefType::get({1}, launchOp.getGridSizeY().getType())); + // builder.create(launchOp->getLoc(), + // launchOp.getGridSizeX(), temp, + // ValueRange({zeroIndex})); + // ptr_ops.push_back( + // builder.create( + // launchOp->getLoc(), temp)); + + for (auto op : launchOp.getKernelOperands()) { + if (mlir::isa(op.getType())) { + ptr_ops.push_back( + builder.create( + launchOp->getLoc(), op)); + } else { + auto temp = builder.create( + launchOp->getLoc(), MemRefType::get({1}, op.getType())); + builder.create(launchOp->getLoc(), op, temp, + ValueRange({zeroIndex})); + ptr_ops.push_back( + builder.create( + launchOp->getLoc(), temp)); + } + } + + auto args = builder.create( + launchOp->getLoc(), + MemRefType::get({numOps}, builder.getIndexType())); + auto args_dynamic = builder.create( + launchOp->getLoc(), + MemRefType::get({ShapedType::kDynamic}, + args.getType().getElementType()), + args); + for (size_t i = 0; i < ptr_ops.size(); i++) { + Value op = ptr_ops[i]; + builder.create( + launchOp->getLoc(), op, args, + builder.create(launchOp->getLoc(), i) + .getResult()); + } + + std::string funcName = "tt_"+ launchOp.getKernelModuleName().str() + launchOp.getKernelName().strref().str(); + if (funcs.find(funcName) == funcs.end()) { + // We need the global string to include the \0 character so that it is + // correctly read by the cuda lib Hence the StringRef(funcName.c_str(), + // funcName.size()+1) + funcs[funcName] = LLVM::createGlobalString( + modOp->getLoc(), builder, funcName + "_str", + StringRef(funcName.c_str(), funcName.size() + 1), + LLVM::linkage::Linkage::Private); + } + builder.create( + launchOp->getLoc(), vendor_prefix+"LaunchKernelMLIR", TypeRange(), + ValueRange({launchOp.getGridSizeX(), launchOp.getGridSizeY(), + launchOp.getGridSizeZ(), launchOp.getBlockSizeX(), + launchOp.getBlockSizeY(), launchOp.getBlockSizeZ(), + args_dynamic, funcs[funcName], + builder.create(launchOp->getLoc(), + funcName.size()), + launchOp.getDynamicSharedMemorySize()})); + launchOp->erase(); + } + + std::vector gpuDeallocs; + modOp->walk([&gpuDeallocs](mlir::gpu::DeallocOp gpuDeallocOp) { + gpuDeallocs.push_back(gpuDeallocOp); + }); + + for (auto dealloc : gpuDeallocs) { + builder.setInsertionPoint(dealloc); + builder.create(dealloc->getLoc(), vendor_prefix+"Free", + TypeRange(), + ValueRange(dealloc->getOperand(0))); + dealloc.erase(); + } + + func::FuncOp funcOp = *initFuncs.begin(); + builder.setInsertionPoint(funcOp.getBody().front().getTerminator()); + + builder.create(funcOp.getBody().front().getTerminator()->getLoc(), vendor_prefix+"Finit", TypeRange(), ValueRange()); + + modOp->walk([](mlir::triton::FuncOp TTFuncOp) { TTFuncOp.erase(); }); + modOp->walk([](mlir::gpu::GPUModuleOp gpuMod) { gpuMod.erase(); }); + + return success(); +} + + +void declare_vendor_funcs(OpBuilder& builder, ModuleOp modOp, std::string vendor_prefix) +{ + auto memrefI32 = + MemRefType::get({ShapedType::kDynamic}, builder.getIntegerType(32)); + auto memrefF32 = + MemRefType::get({ShapedType::kDynamic}, builder.getF32Type()); + auto memrefI64 = + MemRefType::get({ShapedType::kDynamic}, builder.getIntegerType(64)); + auto memrefIndex = + MemRefType::get({ShapedType::kDynamic}, builder.getIndexType()); + auto memrefF64 = + MemRefType::get({ShapedType::kDynamic}, builder.getF64Type()); + builder.clearInsertionPoint(); + auto mallocI32Type = builder.getFunctionType({builder.getIndexType()}, + builder.getIndexType()); + auto mallocI32 = builder.create( + builder.getUnknownLoc(), vendor_prefix+"MallocI32", mallocI32Type); + mallocI32.setVisibility(mlir::SymbolTable::Visibility::Private); + modOp.push_back(mallocI32); + + auto mallocI64Type = builder.getFunctionType({builder.getIndexType()}, + builder.getIndexType()); + auto mallocI64 = builder.create( + builder.getUnknownLoc(), vendor_prefix+"MallocI64", mallocI64Type); + mallocI64.setVisibility(mlir::SymbolTable::Visibility::Private); + modOp.push_back(mallocI64); + + auto mallocF32Type = builder.getFunctionType({builder.getIndexType()}, + builder.getIndexType()); + auto mallocF32 = builder.create( + builder.getUnknownLoc(), vendor_prefix+"MallocF32", mallocF32Type); + mallocF32.setVisibility(mlir::SymbolTable::Visibility::Private); + modOp.push_back(mallocF32); + + auto mallocF64Type = builder.getFunctionType({builder.getIndexType()}, + builder.getIndexType()); + auto mallocF64 = builder.create( + builder.getUnknownLoc(), vendor_prefix+"MallocF64", mallocF64Type); + mallocF64.setVisibility(mlir::SymbolTable::Visibility::Private); + modOp.push_back(mallocF64); + + auto cudaMemcpyI32Type = builder.getFunctionType( + {builder.getIndexType(), memrefI32, builder.getIndexType()}, {}); + auto cudaMemcpyI32 = builder.create( + builder.getUnknownLoc(), vendor_prefix+"MemcpyI32", cudaMemcpyI32Type); + cudaMemcpyI32.setVisibility(mlir::SymbolTable::Visibility::Private); + modOp.push_back(cudaMemcpyI32); + + auto cudaMemcpyI64Type = builder.getFunctionType( + {builder.getIndexType(), memrefI64, builder.getIndexType()}, {}); + auto cudaMemcpyI64 = builder.create( + builder.getUnknownLoc(), vendor_prefix+"MemcpyI64", cudaMemcpyI64Type); + cudaMemcpyI64.setVisibility(mlir::SymbolTable::Visibility::Private); + modOp.push_back(cudaMemcpyI64); + + auto cudaMemcpyIndexType = builder.getFunctionType( + {builder.getIndexType(), memrefIndex, builder.getIndexType()}, {}); + auto cudaMemcpyIndex = builder.create( + builder.getUnknownLoc(), vendor_prefix+"MemcpyIndex", cudaMemcpyIndexType); + cudaMemcpyIndex.setVisibility(mlir::SymbolTable::Visibility::Private); + modOp.push_back(cudaMemcpyIndex); + + auto cudaMemcpyF32Type = builder.getFunctionType( + {builder.getIndexType(), memrefF32, builder.getIndexType()}, {}); + auto cudaMemcpyF32 = builder.create( + builder.getUnknownLoc(), vendor_prefix+"MemcpyF32", cudaMemcpyF32Type); + cudaMemcpyF32.setVisibility(mlir::SymbolTable::Visibility::Private); + modOp.push_back(cudaMemcpyF32); + + auto cudaMemcpyF64Type = builder.getFunctionType( + {builder.getIndexType(), memrefF64, builder.getIndexType()}, {}); + auto cudaMemcpyF64 = builder.create( + builder.getUnknownLoc(), vendor_prefix+"MemcpyF64", cudaMemcpyF64Type); + cudaMemcpyF64.setVisibility(mlir::SymbolTable::Visibility::Private); + modOp.push_back(cudaMemcpyF64); + + auto cudaLaunchKernelT = builder.getFunctionType( + {builder.getIndexType(), builder.getIndexType(), builder.getIndexType(), + builder.getIndexType(), builder.getIndexType(), builder.getIndexType(), + MemRefType::get({ShapedType::kDynamic}, builder.getIndexType()), + LLVM::LLVMPointerType::get(builder.getContext()), + builder.getIndexType(), builder.getIntegerType(32)}, + {}); + auto cudaLaunchKernel = builder.create( + builder.getUnknownLoc(), vendor_prefix+"LaunchKernelMLIR", cudaLaunchKernelT); + cudaLaunchKernel.setVisibility(mlir::SymbolTable::Visibility::Private); + modOp.push_back(cudaLaunchKernel); + + auto cudaSetModuleImageT = builder.getFunctionType( + {LLVM::LLVMPointerType::get(builder.getContext())}, {}); + auto cudaSetModuleImage = builder.create( + builder.getUnknownLoc(), vendor_prefix+"SetModuleImage", cudaSetModuleImageT); + cudaSetModuleImage.setVisibility(mlir::SymbolTable::Visibility::Private); + modOp.push_back(cudaSetModuleImage); + + auto cudaFreeT = builder.getFunctionType({builder.getIndexType()}, {}); + auto cudaFree = builder.create(builder.getUnknownLoc(), + vendor_prefix+"Free", cudaFreeT); + cudaFree.setVisibility(mlir::SymbolTable::Visibility::Private); + modOp.push_back(cudaFree); + auto cudaFinitT = builder.getFunctionType({}, {}); + auto cudaFinit = builder.create(builder.getUnknownLoc(), vendor_prefix+"Finit", cudaFinitT); + cudaFinit.setVisibility(mlir::SymbolTable::Visibility::Private); + modOp.push_back(cudaFinit); +} + +mlir::LogicalResult specializeGpuKernel(mlir::OpBuilder& builder, mlir::ModuleOp modOp, mlir::tensorAlgebra::GPUCompilationFormat codeFormat, Attribute target, std::function add_ttir_passes, std::function add_ttgir_passes, std::function add_llir_passes, std::vector llvm_func_attrs) +{ + std::vector TTFuncs; + modOp->walk([&TTFuncs](mlir::triton::FuncOp op) { TTFuncs.push_back(op); }); + + auto tempMod = builder.create(modOp->getLoc(), "gpu_module", target); + std::vector tempMods; + if (TTFuncs.empty()) { + return failure(); + } + for (auto ttFunc : TTFuncs) { + auto tempMod = ModuleOp::create(modOp.getLoc()); + OpBuilder builder(tempMod.getBodyRegion()); + builder.clone(*ttFunc.getOperation()); + + if (!add_ttir_passes(tempMod)) { + return failure(); + } + + if (!add_ttgir_passes(tempMod)) { + return failure(); + } + + if (!add_llir_passes(tempMod)) { + return failure(); + } + + auto oldAttrs = ttFunc->getAttrs().vec(); + + oldAttrs.insert(oldAttrs.end(), tempMod->getAttrs().begin(), + tempMod->getAttrs().end()); + // ttFunc->setAttrs(tempMod->getAttrs()); + ttFunc->setAttrs(oldAttrs); + tempMods.push_back(tempMod); + } + + llvm::SmallMapVector glob_symbols; + // We have lowered Triton kernels to LLVM, now we need to add them to the + // gpu.module + builder.setInsertionPointToStart(&tempMod.getBodyRegion().front()); + for (auto mod : tempMods) { + for (auto &op : *mod.getBody()) { + // Triton will insert one "global_smem" symbol per module so we have to only keep the first + if(LLVM::GlobalOp glob = mlir::dyn_cast(op)) + { + if(glob_symbols.find(glob.getSymName()) == glob_symbols.end() ) + { + glob_symbols[glob.getSymName()] = true; + builder.clone(op); + } + else + { + assert(glob.getSymName() == "global_smem"); + } + } + else + { + if (mlir::isa(op)) { + auto new_attrs = op.getAttrs().vec(); + new_attrs.insert(new_attrs.end(), llvm_func_attrs.begin(), + llvm_func_attrs.end()); + op.setAttrs(new_attrs); + } + builder.clone(op); + } + } + } + + // GPU dialect verifier expects this + modOp->setAttr("gpu.container_module", builder.getUnitAttr()); + + // The function transformGpuModulesToBinaries expects a module that contains + // gpu.module(s), so we create one with the kernels we want to convert to + // cubin + auto tempOuterMod = ModuleOp::create(modOp.getLoc()); + OpBuilder gbuilder(tempOuterMod.getBodyRegion()); + gbuilder.clone(*tempMod); + gpu::TargetOptions opts; + if (codeFormat == + mlir::tensorAlgebra::GPUCompilationFormat::Assembly) { + opts = gpu::TargetOptions({}, {}, {}, gpu::CompilationTarget::Assembly); + } else if (codeFormat == + mlir::tensorAlgebra::GPUCompilationFormat::Binary) { + opts = gpu::TargetOptions({}, {}, {}, gpu::CompilationTarget::Binary); + } else if (codeFormat == + mlir::tensorAlgebra::GPUCompilationFormat::Fatbin) { + opts = gpu::TargetOptions({}, {}, {}, gpu::CompilationTarget::Fatbin); + } else { + assert(false && "Unexpected gpu compilation code format"); + } + auto res = + mlir::gpu::transformGpuModulesToBinaries(tempOuterMod, nullptr, opts); + if (res.failed()) + { + return failure(); + } + + builder.setInsertionPointToStart(modOp.getBody()); + + // GPU kernels are now converted to cubin, add them to the main module as + // global strings + for (auto &op : *tempOuterMod.getBody()) { + if (auto binOp = dyn_cast(op)) { + auto result = mlir::cast(*binOp.getObjects().begin()) + .getObject() + .str(); + auto type = LLVM::LLVMArrayType::get( + IntegerType::get(builder.getContext(), 8), result.size()); + builder.setInsertionPointToStart(modOp.getBody()); + builder.create( + modOp->getLoc(), type, /*isConstant=*/false, + LLVM::Linkage::Internal, "gpu_code", builder.getStringAttr(result), + /*alignment=*/32); + } + } + + return success(); +} diff --git a/lib/Conversion/IndexTreeToSCF/CMakeLists.txt b/lib/Conversion/IndexTreeToSCF/CMakeLists.txt index 1b56ea5b..788ac02a 100644 --- a/lib/Conversion/IndexTreeToSCF/CMakeLists.txt +++ b/lib/Conversion/IndexTreeToSCF/CMakeLists.txt @@ -1,6 +1,8 @@ add_mlir_conversion_library(COMETIndexTreeToSCF IndexTreeToSCF.cpp AbstractLoopOp.cpp + SymbolicDomainConversion.cpp + IndexTreeConversion.cpp ADDITIONAL_HEADER_DIRS ${COMET_MAIN_INCLUDE_DIR}/comet/Conversion/IndexTreeToSCF @@ -13,6 +15,7 @@ add_mlir_conversion_library(COMETIndexTreeToSCF LINK_LIBS PUBLIC MLIRArithDialect + MLIRIndexDialect MLIRIR MLIRMemRefDialect MLIRSCFDialect diff --git a/lib/Conversion/IndexTreeToSCF/IndexTreeConversion.cpp b/lib/Conversion/IndexTreeToSCF/IndexTreeConversion.cpp new file mode 100644 index 00000000..022c5d31 --- /dev/null +++ b/lib/Conversion/IndexTreeToSCF/IndexTreeConversion.cpp @@ -0,0 +1,219 @@ +// +// Copyright 2022 Battelle Memorial Institute +// +// Redistribution and use in source and binary forms, with or without modification, +// are permitted provided that the following conditions are met: +// +// 1. Redistributions of source code must retain the above copyright notice, this list of conditions +// and the following disclaimer. +// +// 2. Redistributions in binary form must reproduce the above copyright notice, this list of conditions +// and the following disclaimer in the documentation and/or other materials provided with the distribution. +// +// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED +// WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR +// PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, +// INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE +// GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, +// WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +// OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. +// + +#include "comet/Dialect/IndexTree/IR/IndexTreeDialect.h" +#include "comet/Dialect/IndexTree/Passes.h" +#include "comet/Dialect/IndexTree/Patterns.h" + +#include "comet/Dialect/TensorAlgebra/IR/TADialect.h" +#include "mlir/Dialect/SCF/IR/SCF.h" +#include "mlir/Dialect/Index/IR/IndexAttrs.h" +#include "mlir/Dialect/Index/IR/IndexOps.h" +#include "mlir/Dialect/Index/IR/IndexDialect.h" +#include "mlir/Dialect/MemRef/IR/MemRef.h" +#include "mlir/Dialect/Bufferization/IR/Bufferization.h" +#include "mlir/Dialect/Func/IR/FuncOps.h" +#include "mlir/Dialect/Tensor/IR/Tensor.h" +#include "mlir/Transforms/DialectConversion.h" +#include "mlir/Dialect/Func/Transforms/DecomposeCallGraphTypes.h" +#include "mlir/Dialect/Func/Transforms/FuncConversions.h" +#include "mlir/Dialect/SCF/Transforms/Patterns.h" +#include "mlir/Pass/Pass.h" + +using namespace mlir; +using llvm::SmallVector; + +namespace mlir { + namespace comet{ + #define GEN_PASS_DEF_INDEXTREE_INLINING + #include "comet/Conversion/Passes.h.inc" + } +} + + +namespace { +class ConvertIndexTreeYieldOpTypes : public OpConversionPattern { + public: + using OpConversionPattern::OpConversionPattern; + LogicalResult + matchAndRewrite(indexTree::YieldOp op, OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const override { + SmallVector unpacked; + for (Value v : adaptor.getOperands()) { + if (auto cast = + dyn_cast_or_null(v.getDefiningOp())) { + if (cast.getInputs().size() != 1) { + unpacked.append(cast.getInputs().begin(), cast.getInputs().end()); + continue; + } + } + // 1 : 1 type conversion. + unpacked.push_back(v); + } + rewriter.replaceOpWithNewOp(op, unpacked); + return success(); + } +}; + + +class ConvertIndexTreeTypes : public OpConversionPattern{ + public: + using OpConversionPattern::OpConversionPattern; + + mlir::LogicalResult + matchAndRewrite(indexTree::IndexTreeOp op, OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const override { + // Convert the operands + SmallVector unpacked; + for (Value v : adaptor.getOperands()) { + if (auto cast = + dyn_cast_or_null(v.getDefiningOp())) { + if (cast.getInputs().size() != 1) { + unpacked.append(cast.getInputs().begin(), cast.getInputs().end()); + continue; + } + } + // 1 : 1 type conversion. + unpacked.push_back(v); + } + + SmallVector dstTypes; + SmallVector offsets; + offsets.push_back(0); + // Do the type conversion and record the offsets. + for (Type type : op.getResultTypes()) { + if (failed(typeConverter->convertTypes(type, dstTypes))) + return rewriter.notifyMatchFailure(op, "could not convert result type"); + offsets.push_back(dstTypes.size()); + } + + uint32_t input_size = offsets[adaptor.getInputs().size()]; + uint32_t intermediates_size = unpacked.size() - input_size; + ValueRange new_args = ValueRange(unpacked); + + // Calls the actual converter implementation to convert the operation. + auto newOp = rewriter.create( + op.getLoc(), + dstTypes, + new_args.slice(0, input_size), + new_args.slice(input_size, intermediates_size) + ); + rewriter.inlineRegionBefore(op.getRegion(), newOp.getRegion(), newOp.getRegion().end()); + if (failed(rewriter.convertRegionTypes(&newOp.getRegion(), *typeConverter))) + return failure(); + + // Packs the return value. + SmallVector packedRets; + for (unsigned i = 1, e = offsets.size(); i < e; i++) { + unsigned start = offsets[i - 1], end = offsets[i]; + unsigned len = end - start; + ValueRange mappedValue = newOp->getResults().slice(start, len); + if (len != 1) { + // 1 : N type conversion. + Type origType = op.getResultTypes()[i - 1]; + Value mat = typeConverter->materializeSourceConversion( + rewriter, op.getLoc(), origType, mappedValue); + if (!mat) { + return rewriter.notifyMatchFailure( + op, "Failed to materialize 1:N type conversion"); + } + packedRets.push_back(mat); + } else { + // 1 : 1 type conversion. + packedRets.push_back(mappedValue.front()); + } + } + + rewriter.replaceOp(op, packedRets); + return success(); + } +}; + +class InlineIndexTreeOp : public OpConversionPattern{ + public: + using OpConversionPattern::OpConversionPattern; + + mlir::LogicalResult + matchAndRewrite(indexTree::IndexTreeOp op, OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const override { + + Block& body = op.getRegion().front(); + Operation* terminator = body.getTerminator(); + rewriter.inlineBlockBefore(&body, op, op->getOperands()); + rewriter.replaceOp(op, terminator->getOperands()); + return success(); + } +}; + +class InlineIndexTreeYieldOp : public OpConversionPattern{ + public: + using OpConversionPattern::OpConversionPattern; + + mlir::LogicalResult + matchAndRewrite(indexTree::YieldOp op, OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const override { + rewriter.eraseOp(op); + return success(); + } +}; +} //namespace + +void mlir::indexTree::populateIndexTreeInliningPatterns(MLIRContext *context, RewritePatternSet &patterns) { + patterns.add(context); +} + +void mlir::indexTree::populateIndexTreeTypeConversionPatterns(MLIRContext *context, RewritePatternSet &patterns, TypeConverter &typeConverter, ConversionTarget& target) { + target.addDynamicallyLegalOp([&](Operation *op) { + return typeConverter.isLegal(op->getResultTypes()); + }); + target.addDynamicallyLegalOp([&](indexTree::YieldOp op) { + return typeConverter.isLegal(op->getOperandTypes()); + }); + + patterns.add(typeConverter, context); +} + +struct IndexTreeInliningPass + : public PassWrapper> +{ + MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(IndexTreeInliningPass) + + void runOnOperation() override + { + // Convert the rest of the index tree dialect to SCF + TypeConverter typeConverter; + typeConverter.addConversion([](Type type) { return type; }); + mlir::ConversionTarget target(getContext()); + target.addLegalDialect(); + target.addIllegalOp(); + + mlir::RewritePatternSet patterns(&getContext()); + mlir::indexTree::populateIndexTreeInliningPatterns(&getContext(), patterns); + if (mlir::failed(mlir::applyPartialConversion(getOperation(), target, std::move(patterns)))) + signalPassFailure(); + } +}; + +/// Lower sparse tensor algebra operation to loops +std::unique_ptr mlir::comet::createIndexTreeInliningPass() +{ + return std::make_unique(); +} \ No newline at end of file diff --git a/lib/Conversion/IndexTreeToSCF/IndexTreeToSCF.cpp b/lib/Conversion/IndexTreeToSCF/IndexTreeToSCF.cpp index 7eb9a539..34ff2a23 100644 --- a/lib/Conversion/IndexTreeToSCF/IndexTreeToSCF.cpp +++ b/lib/Conversion/IndexTreeToSCF/IndexTreeToSCF.cpp @@ -33,20 +33,32 @@ #include "mlir/Dialect/Arith/IR/Arith.h" #include "mlir/Dialect/Bufferization/IR/Bufferization.h" #include "mlir/Dialect/SCF/IR/SCF.h" +#include "mlir/Dialect/Index/IR/IndexOps.h" #include "mlir/Dialect/Func/IR/FuncOps.h" #include "mlir/Dialect/MemRef/IR/MemRef.h" #include "mlir/Dialect/Math/IR/Math.h" #include "mlir/Dialect/Tensor/IR/Tensor.h" +#include "mlir/IR/BuiltinAttributes.h" +#include "mlir/IR/BuiltinTypeInterfaces.h" #include "mlir/Transforms/DialectConversion.h" #include "mlir/Pass/Pass.h" #include "mlir/IR/Dominance.h" +#include "mlir/Transforms/GreedyPatternRewriteDriver.h" +#include "mlir/Analysis/SliceAnalysis.h" +#include "llvm/ADT/STLExtras.h" #include "llvm/Support/Debug.h" +#include "llvm/Support/ScopedPrinter.h" #include "llvm/ADT/StringSet.h" +#include "llvm/ADT/iterator_range.h" +#include "llvm/ADT/TypeSwitch.h" +#include "llvm/Support/Casting.h" + +#include #include #include +#include #include - #include #include #include @@ -63,13 +75,13 @@ using namespace mlir::tensorAlgebra; using llvm::SmallVector; using llvm::StringRef; +using llvm::SmallDenseMap; #define DEBUG_TYPE "lowering-it-to-scf" // *********** For debug purpose *********// -// #define COMET_DEBUG_MODE +//#define COMET_DEBUG_MODE #include "comet/Utils/debug.h" -#undef COMET_DEBUG_MODE // *********** For debug purpose *********// namespace comet @@ -159,236 +171,10 @@ namespace CSR_DIM2_SIZE = 19 }; - /// MASKING_TYPE to indicate what type of masking is used. - enum MASKING_TYPE - { - NO_MASKING = 0, - PUSH_BASED_MASKING = 1, - PULL_BASED_MASKING = 2 - }; - - /// class MaksingInfo, passed as a parameter to the formSemiringLoopBody() to indicate if using masking or not. - struct MaskingInfo - { - public: - MASKING_TYPE mask_type; - - mlir::Value mask_tensor; - mlir::Value mask_rowptr; - mlir::Value mask_col; - mlir::Value mask_val; - - /// TODO(zhen.peng): Pull-based mask info and auxiliary variables. - - public: - MaskingInfo() : mask_type(NO_MASKING) {} - - /// MaskingInfo(MASKING_TYPE type_, mlir::Value states_) : maskType(type_), states(states_) { } - - void dump() - { - switch (mask_type) - { - case NO_MASKING: - std::cout << "maskType: NO_MASKING\n"; - break; - case PUSH_BASED_MASKING: - std::cout << "maskType: PUSH_BASED_MASKING " - << "mask_tensor: "; - mask_tensor.dump(); - /// std::cout << "maskType: PUSH_BASED_MASKING " << "states: "; - /// states.dump(); - break; - case PULL_BASED_MASKING: - std::cout << "maskType: PULL_BASED_MASKING ... Not supported"; - break; - } - } - }; - - class OpsTree - { - /// private: - public: - // std::vector forOps; /// The (nested) for loops - std::vector forOps; /// The (nested) for loops - std::vector accessIdx; /// The coordinate of accessing that dimension - // std::vector symbolicForOps; /// For-loops in symbolic phase (if necessary) - std::vector symbolicForOps; /// For-loops in symbolic phase (if necessary) - std::vector symbolicAccessIdx; /// The accessing index for that for-loop in symbolic phase (if necessary) - /// std::vector cmptOps; /// The computation ops (no used?) - std::vector children; - OpsTree *parent; - int id; /// the index in the ws_op array. The order is the DFS order. - - // std::vector symbolicForOps_debug; /// (no used?) - // std::vector symbolicAccessIdx_debug; /// (no used?) - - public: - OpsTree() = default; - - // OpsTree(std::vector &forOps, std::vector &accessIdx, - // OpsTree *parent, int id) : forOps(forOps), accessIdx(accessIdx), parent(parent), id(id) - // { - // } - - OpsTree(OpsTree *parent, int id) : parent(parent), id(id) - { - } - - // (no used?) - // OpsTree(std::vector &forOps, std::vector &accessIdx, - // OpsTree *parent) : forOps(forOps), accessIdx(accessIdx), parent(parent) - // { - // } - - ~OpsTree() = default; - - void addChild(OpsTree *tree) - { /// const T& node - this->children.push_back(tree); - } - - // std::vector &getForOps() - // { - // return this->forOps; - // } - - // OpsTree *getParent() - // { - // return this->parent; - // } - - // void setForOps(std::vector &forOps) - // { - // this->forOps = forOps; - // } - - std::vector &getChildren() - { - return this->children; - } - }; - - /// ----------------- /// - /// struct to pass symbolic phase information to the numeric phase - /// ----------------- /// - struct SymbolicInfo - { - bool are_inputs_sparse = false; /// If both inputs are sparse. It is true for SpGEMM and sparse elementwise operations. - /// All other members are only used when are_inputs_sparse is true. - - bool has_symbolic_phase = false; /// If current generated code should have a symbolic phase. - /// Currently, if are_inputs_parse == true; then has_symbolic_phase = true; - - Value mtxC_num_rows = nullptr; - Value mtxC_num_cols = nullptr; - - Value mtxC_rowptr = nullptr; /// Output C's rowptr array when do C = A * B and they are all sparse - /// %alloc_100 = memref.alloc(%43) : memref - Value mtxC_col = nullptr; /// Output C's col array when do C = A * B and they are all sparse - /// %alloc_104 = memref.alloc(%44) : memref - Value mtxC_val = nullptr; /// Output C's val array when do C = A * B and they are all sparse - /// %alloc_108 = memref.alloc(%44) : memref - Value mtxC_val_size = nullptr; /// Output C' correct number of non-zeros or C_val_size (ready after symbolic phase) - - Value mtxC_rowptr_size = nullptr; /// rowptr array's size, which is number of columns plus one (num_col + 1). - - Value row_offset = nullptr; /// In Numeric Phase, row_offset is the insertion location in the C_col and C_val. - - Value mtxC = nullptr; /// The sparse tensor - /// It is %55 below. - }; - - /// ----------------- /// - /// Auxiliary structures for the numeric phase - /// ----------------- /// - struct NumericInfo - { - Value ws_bitmap = nullptr; /// workspace's bitmap to tell if a column ID is visited. - Value ws_bitmap_valueAccessIdx = nullptr; /// value access index for the workspace bitmap. - - Value mask_array = nullptr; /// the intermediate dense vector for a row of the mask. - }; - - /// ----------------- /// - /// Remove an operantion's user who is a memref.store - /// This is very ad-hoc, just to avoid segmentation fault for old very large C.val array and C.col array. - /// ----------------- /// - void removeMemrefStoreUser(Value &opd) - { - { - comet_vdump(opd); - } - std::vector users; - for (Operation *user : opd.getUsers()) - { - if (isa(user)) - { - users.push_back(user); - { - comet_pdump(user); - } - } - } - for (Operation *user : users) - { - user->erase(); - } - } - - /// ----------------- /// - /// Find all users of the old_Value, and replace those users' corresponding operand to new_Value. For example, - /// "ta.print"(%old_Value) => "ta.print"(%new_Value) - /// ----------------- /// - void replaceOldValueToNewValue(Value &old_Value, - Value &new_Value) - { - { - comet_vdump(old_Value); - comet_vdump(new_Value); - } - - /// Traverse each user of new_Value - std::vector users; - for (Operation *user : old_Value.getUsers()) - { - users.push_back(user); - } - DominanceInfo domInfo(new_Value.getDefiningOp()); /// To check dominance - for (Operation *user : users) - { - { - comet_debug() << "before replace operand.\n"; - comet_pdump(user); - } - /// Check if new_Value dominates the user - if (!domInfo.dominates(new_Value, user)) - { - continue; - } - uint64_t op_i = 0; - for (Value op : user->getOperands()) - { - /// Find the mtxC in the user's operands - if (op.getDefiningOp() == old_Value.getDefiningOp()) - { - /// Replace the old sparse tensor to the new one - user->setOperand(op_i, new_Value); - { - comet_debug() << "after replace operand.\n"; - comet_pdump(user); - } - } - ++op_i; - } - } - } - /// ----------------- /// /// Add declaration of the function comet_index_func; /// ----------------- /// - void declareSortFunc(ModuleOp &module, + [[maybe_unused]] void declareSortFunc(ModuleOp &module, MLIRContext *ctx, Location loc) { @@ -409,970 +195,97 @@ namespace } } - /// Get mask_rowptr, mask_col, and mask_val arrays. - /// ----------------- /// - /// mask_tensor = %50 - /// mask_rowptr = %alloc_99 - /// mask_col = %alloc_104 - /// mask_val = %alloc_109 - /// ----------------- /// - /// %45 = bufferization.to_tensor %alloc_99 : memref - /// %46 = bufferization.to_tensor %alloc_104 : memref - /// %49 = bufferization.to_tensor %alloc_109 : memref - /// %50 = ta.sptensor_construct(%41, %42, %43, %44, %45, %46, %47, %48, %49, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38) {tensor_rank = 2 : i32} : (tensor, tensor, tensor, tensor, tensor, tensor, tensor, tensor, tensor, index, index, index, index, index, index, index, index, index, index, index) -> (!ta.sptensor, tensor, tensor, tensor, tensor, tensor, tensor, tensor, tensor, index, index, index, index, index, index, index, index, index, index, index>) - /// ----------------- /// - void getMaskSparseTensorInfo(MaskingInfo &maskingInfo /* contents updated after call*/) + Value getSemiringSecondVal(OpBuilder &builder, Location &loc, + llvm::StringRef &semiringSecond, Value &Input0, Value &Input1) { - Value &mask_tensor = maskingInfo.mask_tensor; - - /// A2pos - Value mask_rowtpr_buff = mask_tensor.getDefiningOp()->getOperand(CSR_A2POS); /// 2 - maskingInfo.mask_rowptr = mask_rowtpr_buff.getDefiningOp()->getOperand(0); - - /// A2crd - Value mask_col_buff = mask_tensor.getDefiningOp()->getOperand(CSR_A2CRD); /// 3 - maskingInfo.mask_col = mask_col_buff.getDefiningOp()->getOperand(0); - - /// Aval - Value mask_val_buff = mask_tensor.getDefiningOp()->getOperand(CSR_AVAL); /// 4 - maskingInfo.mask_val = mask_val_buff.getDefiningOp()->getOperand(0); + Value elementWiseResult; + if (semiringSecond == "times") { - comet_vdump(mask_tensor); - comet_vdump(maskingInfo.mask_rowptr); - comet_vdump(maskingInfo.mask_col); - comet_vdump(maskingInfo.mask_val); + elementWiseResult = builder.create(loc, Input0, Input1); } - } - - unsigned int findIndexInVector_OpsTree(std::vector vec, OpsTree *e) - { - /// Check if element e exists in vector - auto it = std::find(vec.begin(), vec.end(), e); - - /// It accepts a range and an element to search in the given range. If element is found then it returns an iterator to the first element in the given range that’s equal to given element, else it returns an end of the list. - unsigned int ret = 0; - if (it != vec.end()) + else if (semiringSecond == "first") { - /// Get index of element from iterator - ret = std::distance(vec.begin(), it); + elementWiseResult = Input0; } - else + else if (semiringSecond == "second") { - ret = vec.size(); - } - return ret; - } - - // Value findCorrespondingAlloc(Value &iOp) - // { - // comet_debug() << "findCorrespondingAlloc for loop upper bound\n"; - // comet_vdump(iOp); - // auto init_alloc = iOp.getDefiningOp()->getOperand(0); - // comet_vdump(init_alloc); - - // while (true) - // { - // if (isa(init_alloc.getDefiningOp())) - // { - // if (init_alloc.getType().dyn_cast().getDimSize(0) != ShapedType::kDynamic) - // { - // return init_alloc; - // } - // } - // if (init_alloc.getDefiningOp()->getNumOperands() > 0) - // { - // init_alloc = init_alloc.getDefiningOp()->getOperand(0); - // } - // else - // { - // /// Alloc related to another sparse tensor construct such as coming from sparse transpose - // comet_debug() << "Return alloc op - comes from sptensor_construct\n"; - // comet_vdump(init_alloc); - // return init_alloc; - // } - // } - // } - - /// Get allocs for a tensor (sparse or dense) - std::vector getAllocs(Value &tensor) - { - comet_vdump(tensor); - std::vector allocs; - if (tensor.getType().isa()) - { /// Dense tensor - comet_debug() << " getAllocs() - it is dense\n"; - if (isa(tensor.getDefiningOp())) - { - Operation *tensorload = cast(tensor.getDefiningOp()); - auto alloc_op = cast(tensorload->getOperand(0).getDefiningOp()); - comet_vdump(alloc_op); - allocs.push_back(alloc_op); - } - else - { - for (unsigned int i = 0; i < tensor.getDefiningOp()->getNumOperands(); i++) - { - if (isa(tensor.getDefiningOp()->getOperand(i).getDefiningOp())) - { - Operation *tensorload = cast(tensor.getDefiningOp()->getOperand(i).getDefiningOp()); - auto alloc_op = cast(tensorload->getOperand(0).getDefiningOp()); - comet_vdump(alloc_op); - allocs.push_back(alloc_op); - } - } - } + elementWiseResult = Input1; } - else if (tensor.getType().isa()) - { /// nSparse tensor - comet_debug() << " getAllocs() - it is sparse\n"; - auto defop = tensor.getDefiningOp(); - - for (unsigned int n = 0; n < defop.getTotalDimArrayCount(); n++) - { - comet_vdump(defop.getIndices()[n]); - Operation *tensorload = defop.getIndices()[n].getDefiningOp(); - auto alloc_op = cast(tensorload->getOperand(0).getDefiningOp()); - allocs.push_back(alloc_op); - comet_vdump(alloc_op); - } + else if (semiringSecond == "atan2") + { + elementWiseResult = builder.create(loc, Input0, Input1); } - else if (dyn_cast(tensor.getDefiningOp())) - { /// ConstantOp - allocs.push_back(tensor); + else if (semiringSecond == "div") + { + elementWiseResult = builder.create(loc, Input0, Input1); } - return allocs; - } - - std::vector> getAllAllocs(std::vector &tensors) - { - std::vector> allAllocs(tensors.size()); - for (unsigned int i = 0; i < tensors.size(); i++) + else if (semiringSecond == "eq") { - allAllocs[i] = getAllocs(tensors[i]); + elementWiseResult = builder.create(loc, CmpFPredicate::OEQ, Input0, Input1); } - return allAllocs; - } - - /// while until parent == null - void getAncestorsOps(OpsTree *opstree, std::vector &ret) - { - - while (opstree->parent != nullptr) + else if (semiringSecond == "ge") { - ret.push_back(opstree->parent); - opstree = opstree->parent; + elementWiseResult = builder.create(loc, CmpFPredicate::OGE, Input0, Input1); } - } - - /// In genForOps, set Insertion Point for numeric loops. - void setInsertionPointInNumericLoops(OpBuilder &builder, - std::vector &ancestorsOps, - OpsTree *opstree) - { - /// If parent is for loop, insert into the body, How to get end of body? - if (ancestorsOps.size() > 0) + else if (semiringSecond == "gt") { - /// ancestorsOps[0] stores the closest parent - AbstractLoopOp parent_forop; - comet_debug() << "\n"; - std::vector parent_forops = ancestorsOps[0]->forOps; - comet_debug() << " parent_forops.size(): " << parent_forops.size() << " \n"; - - parent_forop = parent_forops.back(); - - comet_debug() << " reset the insertion point\n"; - comet_vdump(parent_forop); - - unsigned int order = findIndexInVector_OpsTree(ancestorsOps[0]->getChildren(), opstree); - comet_debug() << " order: " << order << "\n"; - if (order == ancestorsOps[0]->getChildren().size()) - { - llvm::errs() << __FILE__ << ":" << __LINE__ << "ERROR: Not belong to parent's children\n"; - } - else - { - /// Get the children of the parent_forop - comet_debug() << " number of children: " << parent_forops.size() << "\n"; - if (order == 0) - { - /// builder.setInsertionPointToStart(parent_forop.getBody()); - comet_debug() << "Insertion point order == 0\n"; - builder.setInsertionPoint(parent_forop.getBody()->getTerminator()); - } - else - { - comet_debug() << "\n"; - std::vector brother_forops = ancestorsOps[0]->getChildren()[order - 1]->forOps; - if (brother_forops.size() > 0) - { - comet_debug() << " brother_forops.size(): " << brother_forops.size() << "\n"; - if (opstree->forOps.size() == 0) - { - comet_debug() << "\n"; - comet_vdump(brother_forops[0]); - comet_debug() << "Insertion point (brother_forops.size() > 0 && opstree->forOps.size() == 0)\n"; - builder.setInsertionPointAfter(brother_forops[0]); - } - else - { /// current opstree contains loops, insert in the body of the loops - comet_debug() << " -------- current opstree contain loops --- impossible\n"; - comet_debug() << "Insertion point (brother_forops.size() > 0 && opstree->forOps.size() != 0)\n"; - builder.setInsertionPoint(opstree->forOps.back().getBody()->getTerminator()); - } - } - } - } - comet_debug() << " reset the insertion point\n"; + elementWiseResult = builder.create(loc, CmpFPredicate::OGT, Input0, Input1); } - } - - /// ----------------- /// - /// In genForOps, generate for-loop for a indexOp node if the index is corresponding to Format "D" - /// ----------------- /// - void genForOpFormat_D(OpBuilder &builder, - Location &loc, - Value &tensor, - unsigned int id, - unsigned int i, - std::vector> &allAllocs, - llvm::StringRef &iteratorType, - // scf::ForOp &forLoop /* output */, - AbstractLoopOp &forLoop /* output */, - Value &accessIndex /* output */) - { - /// Value upperBound; - /// Value lowerBound; - /// Check which tensor is sparse, which is dense; - /// Since this function only handles mixed sparse/dense, then "D" only occurs in one tensor - /// Both the dense and sparse tensor contain the dim size; But they are different. Use one. - comet_debug() << " "; - comet_vdump(tensor); - - Value lowerBound = builder.create(loc, 0); - auto step = builder.create(loc, 1); - - if (tensor.getType().isa()) - { /// Dense tensor - Value upperBound; - auto dim = builder.create(loc, tensor, id); - upperBound = dim; - // auto loop = builder.create(loc, lowerBound, upperBound, step); - forLoop.buildLoopOp(iteratorType.str(), - builder, - loc, - lowerBound, - upperBound, - step); - comet_debug() << " D Loop\n"; - comet_vdump(forLoop); - - /// opstree->forOps.push_back(loop); - /// opstree->accessIdx.push_back(loop.getInductionVar()); - // forLoop = loop; - accessIndex = forLoop.getInductionVar(); + else if (semiringSecond == "le") + { + elementWiseResult = builder.create(loc, CmpFPredicate::OLE, Input0, Input1); } - else if (tensor.getType().isa()) + else if (semiringSecond == "lt") { - comet_debug() << " \n"; - comet_pdump(tensor.getDefiningOp()); - if (indexTree::IndexTreeComputeRHSOp rhsop = dyn_cast( - tensor.getDefiningOp())) - { - comet_debug() << " \n"; - } + elementWiseResult = builder.create(loc, CmpFPredicate::OLT, Input0, Input1); } - else if (tensor.getType().cast()) + else if (semiringSecond == "land") { - comet_debug() << "cur_idx is in tensor " << i << "\n"; - - auto index_0 = builder.create(loc, 0); - std::vector upper_indices = {index_0}; - Value upperBound = builder.create(loc, allAllocs[i][4 * id], upper_indices); - comet_vdump(allAllocs[i][4 * id]); - // auto loop = builder.create(loc, lowerBound, upperBound, step); - forLoop.buildLoopOp(iteratorType.str(), - builder, - loc, - lowerBound, - upperBound, - step); - comet_debug() << " D Loop\n"; - comet_vdump(forLoop); - - // forLoop = loop; - accessIndex = forLoop.getInductionVar(); + /// land requires integer type input + llvm::errs() << "Not supported semiring operator (only works for int datatypes): " + << "land" + << "\n"; + /// we should not proceed forward from this point to avoid faulty behavior. + exit(1); } - /// } - } - - /// ----------------- /// - /// In genForOps, generate for-loop for a indexOp node if the index is corresponding to Format "CU" - /// ----------------- /// - void genForOpFormat_CU(OpBuilder &builder, - Location &loc, - OpsTree *opstree, - Value &tensor, - unsigned int id, - unsigned int i, - std::vector> &allAllocs, - // scf::ForOp &parent_forop, - AbstractLoopOp &parent_forop, - Value &parent_accessIdx, - llvm::StringRef &iteratorType, - // scf::ForOp &forLoop /* output */, - AbstractLoopOp &forLoop /* output */, - Value &accessIndex /* output */) - { - /// Generate for(int m = pos[0]; m < pos[1]; m++){int i = crd[m];} - /// if i = 0, index is [0,1] - /// if parent loop and child loop is accessing the same sparse tensor (CSF), index is [m, m+1], m is the nearest loop induction variable - /// Otherwise, the m comes from load operation of the input sparse tensor such as - /// j = crd[i]; - /// for (int m = pos[j]; m < pos[j+1]; m++) - - comet_debug() << " format is CU id: " << id << "\n"; - comet_debug() << " Tensor: \n"; - comet_vdump(tensor); - Value index_lower; - Value index_upper; - if (tensor.getType().cast()) + else if (semiringSecond == "lor") { - comet_debug() << " Tensor type is sparse\n"; - if (id == 0) - { /// The first index in the tensor - index_lower = builder.create(loc, 0); - comet_vdump(index_lower); - } - else - { - if (opstree->parent != nullptr) - { - comet_debug() << " opstree->parent is not NULL\n"; - comet_debug() << " parent forop\n"; - comet_vdump(parent_forop); - - /// TODO (PT) Not sure why this (now commented out) code was needed but it breaks spgemm for cases like C = A * A - /// check if parent's and child's upper bounds come from the same sparse tensor - /// auto parent_UpperBound = parent_forop.getUpperBound(); - /// comet_debug() << " parent upperBound:\n"; - /// comet_vdump(parent_UpperBound); - /// auto alloc_parent_bounds = findCorrespondingAlloc(parent_UpperBound); - /// comet_debug() << " parent upperBound alloc\n"; - /// comet_vdump(alloc_parent_bounds); - - /// comet_debug() << " child upperBound:\n"; - /// comet_vdump(allAllocs[i][4 * id]); - /// auto alloc_child_bounds = findCorrespondingAlloc(allAllocs[i][4 * id]); - /// comet_debug() << " child upperBound alloc\n"; - /// comet_vdump(alloc_child_bounds); - - // if (alloc_child_bounds == alloc_parent_bounds) /// m is the nearest loop induction variable - // { - // comet_debug() << " THESAME: Parent and Child has the same alloc\n"; - // index_lower = parent_forop.getInductionVar(); - // } - // else - { /// m comes from the load - comet_debug() << " DIFFERENT:Parent and Child has the different alloc\n"; - // comet_vdump(alloc_parent_bounds); - // comet_vdump(alloc_child_bounds); - index_lower = parent_accessIdx; - } - } - else - llvm::errs() << "ERROR: Unexpected condition\n"; - } - - comet_debug() << " index_lower:"; - comet_vdump(index_lower); - Value const_index_1 = builder.create(loc, 1); - comet_vdump(const_index_1); - index_upper = builder.create(loc, index_lower, const_index_1); - comet_debug() << " AddIOps (index_upper):"; - comet_vdump(index_upper); - - std::vector lower_indices = {index_lower}; - Value lowerBound = builder.create(loc, allAllocs[i][4 * id], lower_indices); /// 2 * id - - std::vector upper_indices = {index_upper}; - Value upperBound = builder.create(loc, allAllocs[i][4 * id], upper_indices); /// 2 * id - auto step = builder.create(loc, 1); - // auto loop = builder.create(loc, lowerBound, upperBound, step); - forLoop.buildLoopOp(iteratorType.str(), - builder, - loc, - lowerBound, - upperBound, - step); - - comet_debug() << " CU Loop\n"; - comet_vdump(forLoop); - - builder.setInsertionPoint(forLoop.getBody()->getTerminator()); - - std::vector crd_indices = {forLoop.getInductionVar()}; - auto get_index = builder.create(loc, allAllocs[i][4 * id + 1], crd_indices); - - comet_debug() << "CU loop generated\n"; - comet_vdump(forLoop); - // forLoop = loop; - accessIndex = get_index; + /// lor requires integer type input + llvm::errs() << "Not supported semiring operator (only works for int datatypes): " + << "lor" + << "\n"; + /// we should not proceed forward from this point to avoid faulty behavior. + exit(1); } - } - - /// ----------------- /// - /// In genForOps, generate for-loop for a indexOp node if the index is corresponding to Format "CN" - /// ----------------- /// - void genForOpFormat_CN(OpBuilder &builder, - Location &loc, - Value &tensor, - unsigned int id, - unsigned int i, - std::vector> &allAllocs, - llvm::StringRef &iteratorType, - // scf::ForOp &forLoop /* output */, - AbstractLoopOp &forLoop /* output */, - Value &accessIndex /* output */) - { - /// Generate for(int m = pos[0]; m < pos[1]; m++){int i = crd[m];} - if (tensor.getType().cast()) + else if (semiringSecond == "lxor") { - auto index_0 = builder.create(loc, 0); - std::vector lower_indices = {index_0}; - Value lowerBound = builder.create(loc, allAllocs[i][4 * id], lower_indices); - - auto index_1 = builder.create(loc, 1); - std::vector upper_indices = {index_1}; - Value upperBound = builder.create(loc, allAllocs[i][4 * id], upper_indices); - auto step = builder.create(loc, 1); - // auto loop = builder.create(loc, lowerBound, upperBound, step); - forLoop.buildLoopOp(iteratorType.str(), - builder, - loc, - lowerBound, - upperBound, - step); - - comet_debug() << " CN Loop\n"; - comet_vdump(forLoop); - - builder.setInsertionPoint(forLoop.getBody()->getTerminator()); - - std::vector crd_indices = {forLoop.getInductionVar()}; - auto get_index = builder.create(loc, allAllocs[i][4 * id + 1], crd_indices); - - // forLoop = loop; - accessIndex = get_index; + /// lxor requires integer type input + llvm::errs() << "Not supported semiring operator: " + << "lxor" + << "\n"; } - } - - /// ----------------- /// - /// In genForOps, generate for-loop for a indexOp node if the index is corresponding to Format "S" - /// ----------------- /// - void genForOpFormat_S(OpBuilder &builder, - Location &loc, - OpsTree *opstree, - Value &tensor, - unsigned int id, - unsigned int i, - std::vector> &allAllocs, - std::vector &opstree_forops, - AbstractLoopOp &parent_forop, - llvm::StringRef &iteratorType, - AbstractLoopOp &forLoop /* output */, - Value &accessIndex /* output */) - { - /// Currently supported formats, Singleton is not the format of first dimension - /// and it doesn't produce a loop - /// Generate: int j = A2crd[m]; - - if (tensor.getType().cast()) + else if (semiringSecond == "minxy") { - comet_debug() << "cur_idx is in tensor " << i << "\n"; - /// Accesing the last level loop info - AbstractLoopOp last_forop; - if (opstree_forops.size() > 0) - { /// current node contain at least 1 level loop - last_forop = opstree_forops.back(); - } - else - { - if (opstree->parent != nullptr) - last_forop = parent_forop; - } - - std::vector crd_indices = {last_forop.getInductionVar()}; - auto get_index = builder.create(loc, allAllocs[i][4 * id + 1], crd_indices); - - /// Adding one iteration loop to provide consistency with the corresponding index tree. - /// Index tree includes an index node for the dimension but "S" format for this dimension - /// doesn't produce a loop. - Value lowerBound = builder.create(loc, 0); - Value upperBound = builder.create(loc, 1); - auto step = builder.create(loc, 1); - // auto loop = builder.create(loc, lowerBound, upperBound, step); - forLoop.buildLoopOp(iteratorType.str(), - builder, - loc, - lowerBound, - upperBound, - step); - comet_debug() << " S Loop\n"; - comet_vdump(forLoop); - // forLoop = loop; - accessIndex = get_index; + Value cmp = builder.create(loc, CmpFPredicate::OLT, Input0, Input1); + elementWiseResult = builder.create(loc, cmp, Input0, Input1); } - else + else if (semiringSecond == "max") { - llvm::errs() << "Not supported tensor type\n"; + Value cmp = builder.create(loc, CmpFPredicate::OGT, Input0, Input1); + elementWiseResult = builder.create(loc, cmp, Input0, Input1); } - } - - /// In genForOps, set Insertion Point for symbolic loops. - void setInsertionPointInSymbolicLoops(OpBuilder &builder, - std::vector &ancestorsOps, - OpsTree *opstree) - { - /// If parent is for loop, insert into the body, How to get end of body? - if (ancestorsOps.size() > 0) + else if (semiringSecond == "ne") { - /// ancestorsOps[0] stores the closest parent - AbstractLoopOp parent_forop; - comet_debug() << "\n"; - std::vector parent_forops = ancestorsOps[0]->symbolicForOps; - comet_debug() << " parent_forops.size(): " << parent_forops.size() << " \n"; - - parent_forop = parent_forops.back(); - - comet_debug() << "symbolic: reset the insertion point\n"; - comet_vdump(parent_forop); - - unsigned int order = findIndexInVector_OpsTree(ancestorsOps[0]->getChildren(), opstree); - comet_debug() << " order: " << order << "\n"; - if (order == ancestorsOps[0]->getChildren().size()) - { - llvm::errs() << __LINE__ << "Not belong to parent's children\n"; - } - else - { - /// Get the children of the parent_forop - comet_debug() << " number of children: " << parent_forops.size() << "\n"; - if (order == 0) - { - /// builder.setInsertionPointToStart(parent_forop.getBody()); - comet_debug() << "Insertion point order == 0\n"; - builder.setInsertionPoint(parent_forop.getBody()->getTerminator()); - } - else - { - comet_debug() << "\n"; - std::vector brother_forops = ancestorsOps[0]->getChildren()[order - 1]->symbolicForOps; - if (brother_forops.size() > 0) - { - comet_debug() << " brother_forops.size(): " << brother_forops.size() << "\n"; - if (opstree->symbolicForOps.size() == 0) - { - comet_debug() << "\n"; - comet_vdump(brother_forops[0]); - comet_debug() << "Insertion point (brother_forops.size() > 0 && opstree->symbolicForOps.size() == 0)\n"; - builder.setInsertionPointAfter(brother_forops[0]); - } - else - { /// current opstree contains loops, insert in the body of the loops - comet_debug() << " -------- current opstree contain loops --- impossible\n"; - comet_debug() << "Insertion point (brother_forops.size() > 0 && opstree->symbolicForOps.size() != 0)\n"; - builder.setInsertionPoint(opstree->symbolicForOps.back().getBody()->getTerminator()); - } - } - else - { - comet_debug() << "brothers have no for-loops. Insert at the end of parent's for-loop body.\n"; - /// builder.setInsertionPointToEnd(parent_forop.getBody()); /// This doesn't work because it inserts even after the scf.yield, which is wrong. - builder.setInsertionPoint(parent_forop.getBody()->getTerminator()); - } - } - } - comet_debug() << " reset the insertion point\n"; + elementWiseResult = builder.create(loc, CmpFPredicate::ONE, Input0, Input1); } - } - - /// In genCmptOps, generate code for a compute node with workspace transformation. - /// For example, A = 0.0 . A could be scalar or vector. - void genWorkspaceCmptOpInitialAssignment(OpBuilder &builder, - Location &loc, - int lhs_loc, - ConstantOp &cstop, - std::vector &nested_forops, - std::vector> &tensors_lhs_Allocs, - std::vector> &main_tensors_all_Allocs, - bool use_dynamic_init, - SymbolicInfo &symbolicInfo) - { - - /// Generate Store 1.0, A[...] this op - /// this case: allPerms[0] is empty, allFormats[0] is empty - comet_vdump(cstop); - comet_debug() << " cstop.getValue(): " << cstop.getValue() << "\n"; - comet_vdump(main_tensors_all_Allocs[lhs_loc].back()); - comet_debug() << " tensors_lhs_Allocs.size(): " << tensors_lhs_Allocs.size() << "\n"; + else if (semiringSecond == "minus") { - comet_vdump(nested_forops[0]); + elementWiseResult = builder.create(loc, Input0, Input1); } - Value local_accessIdx = nested_forops[0].getInductionVar(); - insertInitialize(loc, - cstop, - main_tensors_all_Allocs[lhs_loc].back(), - local_accessIdx, - builder, - use_dynamic_init, - symbolicInfo.mtxC_rowptr /* dynamic_init */); - } - - /// In genCmptOps, generate code for a compute node that copy a sparse input row into a dense vector. - void genWorkspaceCmptOpScatterInputToWorkspace(OpBuilder &builder, - Location &loc, - int main_tensor_nums, - std::vector> &main_tensors_all_Allocs, - std::vector> &allValueAccessIdx) - { - - std::vector allLoads(main_tensor_nums); - for (auto m = 0; m < main_tensor_nums; m++) + else if (semiringSecond == "plusxy") { - Value s = builder.create(loc, - main_tensors_all_Allocs[m][main_tensors_all_Allocs[m].size() - 1], - allValueAccessIdx[m]); - allLoads[m] = s; - comet_debug() << " "; - comet_vdump(s); - } - comet_debug() << " allLoads.size(): " << allLoads.size() << "\n"; - - builder.create(loc, allLoads[0], - main_tensors_all_Allocs[1][main_tensors_all_Allocs[1].size() - 1], - allValueAccessIdx[1]); - } - - /// Generate scf.for op for indices - /// The index is the "idx"th index of "tensor" - void genForOps(std::vector &tensors, - std::vector &ids, - std::vector &formats, - indexTree::IndexTreeOp rootOp, - OpBuilder &builder, - OpsTree *opstree, - SymbolicInfo &symbolicInfo, - llvm::StringRef &iteratorType) - { - comet_debug() << " genForOps indexTreeOp\n"; - comet_vdump(rootOp); - Location loc = rootOp.getLoc(); - /// The insertion location should be "the end of the body of parent loop" - std::vector ancestorsOps; - getAncestorsOps(opstree, ancestorsOps); - comet_debug() << " genForOps ancestorsOps.size(): " << ancestorsOps.size() << "\n"; - for ([[maybe_unused]] unsigned int i = 0; i < ancestorsOps.size(); i++) - { - comet_debug() << " ancestorsOps[" << i << "]->forOps.size(): " << ancestorsOps[i]->forOps.size() - << ", ancestorsOps->id: " - << ancestorsOps[i]->id << "\n"; - } - comet_debug() << "Tensor size: " << tensors.size() << "\n"; - std::vector> allAllocs = getAllAllocs(tensors); - - comet_debug() << "Tensors:\n"; - for ([[maybe_unused]]unsigned int i = 0; i < tensors.size(); i++) - { - comet_vdump(tensors[i]); - } - - /// ----------------- /// - /// Set insertion point - /// ----------------- /// - setInsertionPointInNumericLoops(builder, - ancestorsOps, - opstree); - - for (unsigned int i = 0; i < tensors.size(); i++) - { - if (i > 0) - { - /// insertion point: the body of the previous i's loop body - comet_debug() << " -------- current opstree contain loops\n"; - builder.setInsertionPoint(opstree->forOps.back().getBody()->getTerminator()); - } - - Value &tensor = tensors[i]; - std::string format = formats[i]; - unsigned int id = ids[i]; - - comet_debug() << " current index format: " << format << "\n"; - if (format.compare(0, 1, "D") == 0) - { - /// Symbolic Phase - if (symbolicInfo.has_symbolic_phase) - { - /// Store the insertion point - auto last_insertion_point = builder.saveInsertionPoint(); - - /// Set the insertions point - setInsertionPointInSymbolicLoops(builder, - ancestorsOps, - opstree); - - AbstractLoopOp forLoop; - Value accessIndex; - genForOpFormat_D(builder, - loc, - tensor, - id, - i, - allAllocs, - iteratorType, - forLoop /* output */, - accessIndex /* output */); - opstree->symbolicForOps.push_back(forLoop); - opstree->symbolicAccessIdx.push_back(accessIndex); - - /// Restore the insertion point - builder.restoreInsertionPoint(last_insertion_point); - } - /// Check which tensor is sparse, which is dense; - /// Since this function only handles mixed sparse/dense, then "D" only occurs in one tensor - /// Both the dense and sparse tensor contain the dim size; But they are different. Use one. - AbstractLoopOp forLoop; - Value accessIndex; - genForOpFormat_D(builder, - loc, - tensor, - id, - i, - allAllocs, - iteratorType, - forLoop /* output */, - accessIndex /* output */); - opstree->forOps.push_back(forLoop); - opstree->accessIdx.push_back(accessIndex); - } - /// mix sparse dense tensor contraction, only one sparse tensor - else if (format.compare(0, 2, "CU") == 0) - { - /// Symbolic Phase - if (symbolicInfo.has_symbolic_phase) - { - /// Store the insertion point - auto last_insertion_point = builder.saveInsertionPoint(); - - /// Set the insertions point - setInsertionPointInSymbolicLoops(builder, - ancestorsOps, - opstree); - - AbstractLoopOp forLoop; - Value accessIndex; - AbstractLoopOp parent_forop; - Value parent_accessIdx; - if (nullptr != opstree->parent) - { - parent_forop = opstree->parent->symbolicForOps.back(); - parent_accessIdx = opstree->parent->symbolicAccessIdx.back(); - } - genForOpFormat_CU(builder, - loc, - opstree, - tensor, - id, - i, - allAllocs, - parent_forop, - parent_accessIdx, - iteratorType, - forLoop /* output */, - accessIndex /* output */); - opstree->symbolicForOps.push_back(forLoop); - opstree->symbolicAccessIdx.push_back(accessIndex); - - /// Restore the insertion point - builder.restoreInsertionPoint(last_insertion_point); - } - /// Generate for(int m = pos[0]; m < pos[1]; m++){int i = crd[m];} - /// if i = 0, index is [0,1] - /// if parent loop and child loop is accessing the same sparse tensor (CSF), index is [m, m+1], m is the nearest loop induction variable - /// Otherwise, the m comes from load operation of the input sparse tensor such as - /// j = crd[i]; - /// for (int m = pos[j]; m < pos[j+1]; m++) - - AbstractLoopOp forLoop; - Value accessIndex; - AbstractLoopOp parent_forop; - Value parent_accessIdx; - if (nullptr != opstree->parent) - { - parent_forop = opstree->parent->forOps.back(); - parent_accessIdx = opstree->parent->accessIdx.back(); - } - genForOpFormat_CU(builder, - loc, - opstree, - tensor, - id, - i, - allAllocs, - parent_forop, - parent_accessIdx, - iteratorType, - forLoop /* output */, - accessIndex /* output */); - opstree->forOps.push_back(forLoop); - opstree->accessIdx.push_back(accessIndex); - } - else if (format.compare(0, 2, "CN") == 0) - { - /// Generate for(int m = pos[0]; m < pos[1]; m++){int i = crd[m];} - AbstractLoopOp forLoop; - Value accessIndex; - genForOpFormat_CN(builder, - loc, - tensor, - id, - i, - allAllocs, - iteratorType, - forLoop /* output */, - accessIndex /* output */); - opstree->forOps.push_back(forLoop); - opstree->accessIdx.push_back(accessIndex); - } - else if (format.compare(0, 1, "S") == 0) - { - /// Currently supported formats, Singleton is not the format of first dimension - /// and it doesn't produce a loop - /// Generate: int j = A2crd[m]; - AbstractLoopOp forLoop; - Value accessIndex; - std::vector &opstree_forops = opstree->forOps; - AbstractLoopOp parent_forop; - if (nullptr != opstree->parent) - { - parent_forop = opstree->parent->forOps.back(); - } - genForOpFormat_S(builder, - loc, - opstree, - tensor, - id, - i, - allAllocs, - opstree_forops, - parent_forop, - iteratorType, - forLoop /* output */, - accessIndex /* output */); - opstree->forOps.push_back(forLoop); - opstree->accessIdx.push_back(accessIndex); - } - else - { - llvm::errs() << "Not supported format: " << format << "\n"; - } - - comet_debug() << "finish generate loops for current index format: " << format << "\n"; - } - } - - Value getSemiringSecondVal(OpBuilder &builder, Location &loc, - llvm::StringRef &semiringSecond, Value &Input0, Value &Input1, - bool compressedWorkspace) - { - - Value elementWiseResult; - if (semiringSecond == "times") - { - elementWiseResult = builder.create(loc, Input0, Input1); - } - else if (semiringSecond == "first") - { - elementWiseResult = Input0; - } - else if (semiringSecond == "second") - { - elementWiseResult = Input1; - } - else if (semiringSecond == "atan2") - { - elementWiseResult = builder.create(loc, Input0, Input1); - } - else if (semiringSecond == "div") - { - elementWiseResult = builder.create(loc, Input0, Input1); - } - else if (semiringSecond == "eq") - { - elementWiseResult = builder.create(loc, CmpFPredicate::OEQ, Input0, Input1); - } - else if (semiringSecond == "ge") - { - elementWiseResult = builder.create(loc, CmpFPredicate::OGE, Input0, Input1); - } - else if (semiringSecond == "gt") - { - elementWiseResult = builder.create(loc, CmpFPredicate::OGT, Input0, Input1); - } - else if (semiringSecond == "le") - { - elementWiseResult = builder.create(loc, CmpFPredicate::OLE, Input0, Input1); - } - else if (semiringSecond == "lt") - { - elementWiseResult = builder.create(loc, CmpFPredicate::OLT, Input0, Input1); - } - else if (semiringSecond == "land") - { - /// land requires integer type input - llvm::errs() << "Not supported semiring operator (only works for int datatypes): " - << "land" - << "\n"; - /// we should not proceed forward from this point to avoid faulty behavior. - exit(1); - } - else if (semiringSecond == "lor") - { - /// lor requires integer type input - llvm::errs() << "Not supported semiring operator (only works for int datatypes): " - << "lor" - << "\n"; - /// we should not proceed forward from this point to avoid faulty behavior. - exit(1); - } - else if (semiringSecond == "lxor") - { - /// lxor requires integer type input - llvm::errs() << "Not supported semiring operator: " - << "lxor" - << "\n"; - } - else if (semiringSecond == "minxy") - { - Value cmp = builder.create(loc, CmpFPredicate::OLT, Input0, Input1); - elementWiseResult = builder.create(loc, cmp, Input0, Input1); - } - else if (semiringSecond == "max") - { - Value cmp = builder.create(loc, CmpFPredicate::OGT, Input0, Input1); - elementWiseResult = builder.create(loc, cmp, Input0, Input1); - } - else if (semiringSecond == "ne") - { - elementWiseResult = builder.create(loc, CmpFPredicate::ONE, Input0, Input1); - } - else if (semiringSecond == "minus") - { - elementWiseResult = builder.create(loc, Input0, Input1); - } - else if (semiringSecond == "plusxy") - { - elementWiseResult = builder.create(loc, Input0, Input1); + elementWiseResult = builder.create(loc, Input0, Input1); } else if (semiringSecond == "pairxy") { @@ -1392,8 +305,7 @@ namespace } Value getSemiringFirstVal(OpBuilder &builder, Location &loc, - llvm::StringRef &semiringFirst, Value &Input0, Value &Input1, - bool compressedWorkspace) + llvm::StringRef &semiringFirst, Value &Input0, Value &Input1) { Value reduceResult; @@ -1407,31 +319,13 @@ namespace } else if (semiringFirst == "minxy") { - if (!compressedWorkspace) - { - llvm::errs() << "Not supported semiring operator " - "(please use compressed workspace optimization or opt-comp-workspace " - "where this operation is known to work): " - << "min" - << "\n"; - /// we should not proceed forward from this point to avoid in-correct results from generated code. - } Value cmp = builder.create(loc, CmpFPredicate::OLT, Input0, Input1); - reduceResult = builder.create(loc, cmp, Input0, Input1); + reduceResult = builder.create(loc, cmp, Input0, Input1); } else if (semiringFirst == "max") { - if (!compressedWorkspace) - { - llvm::errs() << "Not supported semiring operator " - "(please use compressed workspace optimization or opt-comp-workspace " - "where this operation is known to work): " - << "max" - << "\n"; - /// we should not proceed forward from this point to avoid in-correct results from generated code. - } Value cmp = builder.create(loc, CmpFPredicate::OGT, Input0, Input1); - reduceResult = builder.create(loc, cmp, Input0, Input1); + reduceResult = builder.create(loc, cmp, Input0, Input1); } else if (semiringFirst == "land") { @@ -1466,3082 +360,1346 @@ namespace return reduceResult; } - /// Generate numeric semiring kernel if statement condition - void genCmptOpKernelIfStatementCondition(OpBuilder &builder, - Location &loc, - NumericInfo &numericInfo, - MaskingInfo &maskingInfo, - scf::IfOp &if_notAlreadySet /* output */) - { - Value &is_visited_alloc = numericInfo.ws_bitmap; - Value &valueAccessIdx = numericInfo.ws_bitmap_valueAccessIdx; - - Value const_i1_false = builder.create(loc, builder.getI1Type(), builder.getBoolAttr(0)); - Value const_i1_true = builder.create(loc, builder.getI1Type(), builder.getBoolAttr(1)); - if (PUSH_BASED_MASKING == maskingInfo.mask_type) - { - /// if (mask_array[j] == true) { /// C[i,k] is allowed by the mask and has not been seen yet - /// if (ws_bitmap[j] != true) { - Value &mask_array = numericInfo.mask_array; - Value ele_mask_array = builder.create(loc, mask_array, ValueRange{valueAccessIdx}); - Value compare_true = builder.create(loc, CmpIPredicate::eq, ele_mask_array, const_i1_true); - auto if_mask_set = builder.create(loc, compare_true, false /* no else region */); - builder.setInsertionPointToStart(&if_mask_set.getThenRegion().front()); - Value ele_bitmap = builder.create(loc, is_visited_alloc, ValueRange{valueAccessIdx}); - Value compare_false = builder.create(loc, CmpIPredicate::eq, ele_bitmap, const_i1_false); - if_notAlreadySet = builder.create(loc, compare_false, /*WithElseRigion*/ true); - { - comet_vdump(ele_mask_array); - comet_vdump(if_mask_set); - comet_vdump(if_notAlreadySet); - } - } - else if (NO_MASKING == maskingInfo.mask_type) - { - /// if (ws_bitmap[j] != true) { - /// Workspace tensors are on the lhs - Value checkAlreadySet = builder.create(loc, is_visited_alloc, ValueRange{valueAccessIdx}); - Value notAlreadySet = builder.create(loc, CmpIPredicate::eq, checkAlreadySet, const_i1_false); - if_notAlreadySet = builder.create(loc, notAlreadySet, /*WithElseRegion*/ true); - { - comet_vdump(checkAlreadySet); - comet_vdump(notAlreadySet); - comet_vdump(if_notAlreadySet); - } - } - else - { - llvm::errs() << "Error: mask_type " << maskingInfo.mask_type << " is not supported.\n"; - } - } - - /// Generate numeric semiring kernel if statement then region - void genCmptOpKernelIfStatementThenRegion(OpBuilder &builder, - Location &loc, - int lhs_loc, - int main_tensor_nums, - scf::IfOp &if_notAlreadySet, - bool compressedWorkspace, - llvm::StringRef &semiringSecond, - std::vector> &main_tensors_all_Allocs, - std::vector> &tensors_lhs_Allocs, - std::vector> &allValueAccessIdx, - SymbolicInfo &symbolicInfo, - NumericInfo &numericInfo) - { - Value &ws_bitmap = numericInfo.ws_bitmap; - Value &ws_bitmap_valueAccessIdx = numericInfo.ws_bitmap_valueAccessIdx; - Value &W_id_list_size = tensors_lhs_Allocs[3][0]; - Value &mtxC_col = symbolicInfo.mtxC_col; - Value &W_data = main_tensors_all_Allocs[lhs_loc].back(); - Value &W_data_valueAccessIdx = ws_bitmap_valueAccessIdx; - - builder.setInsertionPointToStart(&if_notAlreadySet.getThenRegion().front()); - - /// Wj = Aik * Bkj /// computation wj, outer has k, so +=/= need if/else - /// W_already_set[j] = 1 - /// W_index_list[W_index_list_size] = j - /// W_index_list_size++ - - std::vector allLoadsIf(main_tensor_nums); - for (int m = 0; m < main_tensor_nums; m++) - { - Value s = builder.create(loc, main_tensors_all_Allocs[m][main_tensors_all_Allocs[m].size() - 1], allValueAccessIdx[m]); - allLoadsIf[m] = s; - comet_debug() << " "; - comet_vdump(s); - } - comet_debug() << " allLoadsIf.size(): " << allLoadsIf.size() << "\n"; - - comet_debug() << "calculate elementWise operation only\n"; - /// val = A[j_idx] * B[j_idx]; - /// W_data[j_idx] = val; - Value elementWiseResult = getSemiringSecondVal(builder, loc, semiringSecond, allLoadsIf[0], allLoadsIf[1], compressedWorkspace); -#ifdef DEBUG_MODE_LowerIndexTreeToSCFPass - auto store_sum = builder.create(loc, - elementWiseResult, - W_data, - W_data_valueAccessIdx); - comet_vdump(elementWiseResult); - comet_vdump(store_sum); -#else - builder.create(loc, - elementWiseResult, - W_data, - W_data_valueAccessIdx); -#endif - Value const_index_0 = builder.create(loc, 0); - Value const_index_1 = builder.create(loc, 1); - Value const_i1_true = builder.create(loc, builder.getI1Type(), builder.getBoolAttr(1)); - - /// ws_bitmap[j_idx] = true; - builder.create(loc, const_i1_true, ws_bitmap, ws_bitmap_valueAccessIdx); - - Value W_id_list_size_old = builder.create(loc, W_id_list_size, ValueRange{const_index_0}); - - assert(allValueAccessIdx[lhs_loc].size() == 1 && " more than one access id for auxiliary array\n"); - - /// C.col[W_id_list_size] = j_idx; - builder.create(loc, - ws_bitmap_valueAccessIdx, - mtxC_col, - ValueRange{W_id_list_size_old}); - - /// W_id_list_size += 1 - Value W_id_list_size_new = builder.create(loc, W_id_list_size_old, const_index_1); - comet_debug() << " AddIOps (W_index_list_size_new)"; - comet_vdump(W_id_list_size_new); - - builder.create(loc, W_id_list_size_new, W_id_list_size, ValueRange{const_index_0}); - { - comet_vdump(if_notAlreadySet); - } - } - - /// Generate numeric semiring kernel if statement else region - void genCmptOpKernelIfStatementElseRegion(OpBuilder &builder, - Location &loc, - int lhs_loc, - int main_tensor_nums, - scf::IfOp &if_notAlreadySet, - bool compressedWorkspace, - llvm::StringRef &semiringFirst, - llvm::StringRef &semiringSecond, - std::vector> &main_tensors_all_Allocs, - std::vector> &allValueAccessIdx) - { + struct TensorSubsetInfo { + int64_t dim; + int64_t tiles; + int64_t tile_size; + }; - Value &W_data = main_tensors_all_Allocs[lhs_loc].back(); - Value &W_data_valueAccessIdx = allValueAccessIdx[lhs_loc][0]; - builder.setInsertionPointToStart(&if_notAlreadySet.getElseRegion().front()); + class IndexTreeInferOutputSets { + SmallDenseMap> output_sets; - std::vector allLoadsElse(main_tensor_nums); - for (auto m = 0; m < main_tensor_nums; m++) - { - Value s = builder.create(loc, main_tensors_all_Allocs[m][main_tensors_all_Allocs[m].size() - 1], allValueAccessIdx[m]); - allLoadsElse[m] = s; - comet_vdump(s); - } - comet_debug() << " allLoadsElse.size(): " << allLoadsElse.size() << "\n"; + public: + MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(IndexTreeInferOutputSets) - comet_debug() << "calculate elementWise operation and reduction\n"; - Value elementWiseResult = getSemiringSecondVal(builder, loc, semiringSecond, allLoadsElse[0], allLoadsElse[1], compressedWorkspace); - Value reduceResult = getSemiringFirstVal(builder, loc, semiringFirst, allLoadsElse[lhs_loc], elementWiseResult, compressedWorkspace); - builder.create(loc, reduceResult, W_data, W_data_valueAccessIdx); + IndexTreeInferOutputSets(Operation* op) { - comet_vdump(if_notAlreadySet); - } - } + // TODO: Will not work with workspace op or with sparse tensors! + // Also will not work with indices that don't align to tensor dims! - /// Generate the numeric bitmap - /// It should be deprecated in the future, as the bitmap would be lowered from the Index Tree dialect. - void genNumericBitmap(OpBuilder &builder, - Location &loc, - AbstractLoopOp &symbolic_outermost_forLoop, - SymbolicInfo &symbolicInfo, - Value &bitmap_alloc) - { - /// Store the insertion point - auto last_insertion_point = builder.saveInsertionPoint(); - - /// Jump Insertion Point to the front of the 2nd outermost for-loop - builder.setInsertionPoint(symbolic_outermost_forLoop); - - Value const_index_0 = builder.create(loc, 0); - Value const_index_1 = builder.create(loc, 1); - Value &mtxC_dim2_size = symbolicInfo.mtxC_num_cols; - - MemRefType memTy_dynamic_1i = MemRefType::get({ShapedType::kDynamic}, builder.getI1Type()); - bitmap_alloc = builder.create(loc, - memTy_dynamic_1i, - ValueRange{mtxC_dim2_size}, - builder.getI64IntegerAttr(8) /* alignment bytes */); - Value const_i1_0 = builder.create(loc, builder.getI1Type(), builder.getBoolAttr(0)); - scf::ForOp init_forLoop = builder.create(loc, - const_index_0 /* lowerBound */, - mtxC_dim2_size /* upperBound */, - const_index_1 /* step */); - builder.setInsertionPointToStart(init_forLoop.getBody()); - Value i_idx = init_forLoop.getInductionVar(); - builder.create(loc, - const_i1_0, - bitmap_alloc, - ValueRange{i_idx}); - { - comet_vdump(bitmap_alloc); - comet_vdump(init_forLoop); + IndexTreeOp tree = llvm::cast(op); + for(Value input : tree.getBody()->getArguments()) { + for(auto user : input.getUsers()) { + if(auto lhs = llvm::dyn_cast(user)) { + int64_t dim = 0; + for(Value crd : lhs.getCrds()) { + if(auto access = crd.getDefiningOp()) { + auto node = access.getIndex().getDefiningOp(); + TensorSubsetInfo slice = {dim, -1, -1}; + output_sets[node].insert(std::make_pair(input, slice)); + } + dim += 1; + } + } + } + } } - /// Restore the insertion point - builder.restoreInsertionPoint(last_insertion_point); - } - - /// Generate numeric mask-array before the numeric outermost for-loop. - /// Please don't confuse with mark-array. - void genNumericMaskArray(OpBuilder &builder, - Location &loc, - AbstractLoopOp &numeric_outermost_forLoop, - SymbolicInfo &symbolicInfo, - NumericInfo &numericInfo /* output */) - { - /// Store the insertion point - auto last_insertion_point = builder.saveInsertionPoint(); - - /// Set the insertion Point before the numeric outermost for-loop - builder.setInsertionPoint(numeric_outermost_forLoop); - - /// Generate the mask-array - Value &dim2_size = symbolicInfo.mtxC_num_cols; - MemRefType memTy_dynamic_i1 = MemRefType::get({ShapedType::kDynamic}, builder.getI1Type()); - Value mask_array_alloc = builder.create(loc, - memTy_dynamic_i1, - ValueRange{dim2_size}, - builder.getI64IntegerAttr(8) /* alignment bytes */); - - /// Initialize the mask-array - Value const_index_0 = builder.create(loc, 0); - Value const_index_1 = builder.create(loc, 1); - scf::ForOp mask_array_init_loop = builder.create(loc, - const_index_0 /* lowerBound */, - dim2_size /* upperBound */, - const_index_1 /* step */); - builder.setInsertionPointToStart(mask_array_init_loop.getBody()); - Value j_idx = mask_array_init_loop.getInductionVar(); - Value const_i1_false = builder.create(loc, - builder.getI1Type(), - builder.getBoolAttr(false)); - builder.create(loc, - const_i1_false, - mask_array_alloc, - ValueRange{j_idx}); - - numericInfo.mask_array = mask_array_alloc; - + SmallDenseMap getOutputSets(IndexTreeIndicesOp op) { - comet_vdump(mask_array_alloc); - comet_vdump(mask_array_init_loop); - } - - /// Restore the insertion point - builder.restoreInsertionPoint(last_insertion_point); - } + if(output_sets.contains(op)){ + return output_sets[op]; + } - /// Generate setting the mask-array at the begining of the numeric outermost for-loop, - /// and resetting at the end of the outermost for-loop. - /// ----------------- /// - /// %j_loc_start = memref.load %mask_rowptr[%i_idx] : memref /// alloc_16 = mask.rowptr - /// %j_loc_bound = memref.load %mask_rowptr[%i_idx_plus_one] : memref - /// scf.for %arg1 = %j_loc_start to %j_loc_bound step %c1 { - /// %val = memref.load %mask_val[%arg1] : memref - /// %54 = arith.cmpf une, %val, %cst : f64 - /// scf.if %54 { - /// %j_idx = memref.load %mask_col[%arg1] : memref - /// memref.store %true, %mask_array[%j_idx] : memref - /// } - /// } - /// ----------------- /// - /// Reset mask_array at the end of numeric outermost for-loop - /// ----------------- /// - /// scf.for %arg1 = %j_loc_start to %j_loc_bound step %c1 { - /// %j_idx = memref.load %mask_col[%arg1] : memref - /// memref.store %false, %array_mask[%j_idx] : memref - /// } - void genNumericSetAndResetMaskArray(OpBuilder &builder, - Location &loc, - AbstractLoopOp &numeric_outermost_forLoop, - Value &outermost_forLoop_valueAccessIdx, - NumericInfo &numericInfo, - MaskingInfo &maskingInfo) - { - /// Store the insertion point - auto last_insertion_point = builder.saveInsertionPoint(); - - /// Set the insertion Point before the numeric semiring for-loop - builder.setInsertionPointToStart(numeric_outermost_forLoop.getBody()); - - /// Generate the setting for-loop entry - Value &mask_array = numericInfo.mask_array; - Value &mask_rowptr = maskingInfo.mask_rowptr; - Value &mask_col = maskingInfo.mask_col; - Value &mask_val = maskingInfo.mask_val; - Value const_index_1 = builder.create(loc, 1); - Value &i_idx = outermost_forLoop_valueAccessIdx; - Value i_idx_plus_one = builder.create(loc, i_idx, const_index_1); - Value j_loc_start = builder.create(loc, mask_rowptr, ValueRange{i_idx}); - Value j_loc_bound = builder.create(loc, mask_rowptr, ValueRange{i_idx_plus_one}); - scf::ForOp init_for_loop = builder.create(loc, - j_loc_start /* lower_bound */, - j_loc_bound /* upper_bound*/, - const_index_1 /* step */); - - /// Generate the setting for-loop body - builder.setInsertionPointToStart(init_for_loop.getBody()); - Value const_f64_0 = builder.create(loc, builder.getF64Type(), builder.getF64FloatAttr(0)); - Value j_loc = init_for_loop.getInductionVar(); - Value val = builder.create(loc, mask_val, ValueRange{j_loc}); - Value not_zero = builder.create(loc, CmpFPredicate::UNE, val, const_f64_0); - auto if_not_zero = builder.create(loc, not_zero, false /*NoElseRegion*/); - builder.setInsertionPointToStart(&if_not_zero.getThenRegion().front()); - Value j_idx = builder.create(loc, mask_col, ValueRange{j_loc}); - Value const_i1_1 = builder.create(loc, builder.getI1Type(), builder.getBoolAttr(true)); - builder.create(loc, - const_i1_1, - mask_array, - ValueRange{j_idx}); - { - comet_vdump(val); - comet_vdump(if_not_zero); - comet_vdump(init_for_loop); + // This also will cause errors!!! + return SmallDenseMap(); } - /// Generate the resetting for-loop entry after the semiring for-loop - builder.setInsertionPoint(numeric_outermost_forLoop.getBody()->getTerminator()); - scf::ForOp reset_for_loop = builder.create(loc, - j_loc_start /* lower_bound */, - j_loc_bound /* upper_bound*/, - const_index_1 /* step */); - - /// Generate the resetting for-loop body - builder.setInsertionPointToStart(reset_for_loop.getBody()); - j_loc = reset_for_loop.getInductionVar(); - j_idx = builder.create(loc, mask_col, ValueRange{j_loc}); - Value const_i1_0 = builder.create(loc, builder.getI1Type(), builder.getBoolAttr(false)); - builder.create(loc, - const_i1_0, - mask_array, - ValueRange{j_idx}); - { - comet_vdump(reset_for_loop); - comet_vdump(numeric_outermost_forLoop); - } + }; - /// Restore the insertion point - builder.restoreInsertionPoint(last_insertion_point); - } + class IndexVar { + public: + virtual Value getCrd(IRRewriter& rewriter) = 0; + virtual Value getPos(IRRewriter& rewriter, Value tensor, uint32_t dim) = 0; + virtual ~IndexVar(){}; + }; - void formSemiringLoopBody(indexTree::IndexTreeComputeOp &cur_op, - bool comp_worksp_opt, - llvm::StringRef &semiringFirst, - llvm::StringRef &semiringSecond, - OpBuilder &builder, Location &loc, int lhs_loc, - std::vector> &main_tensors_all_Allocs, - std::vector> &tensors_lhs_Allocs, - std::vector> &tensors_rhs_Allocs, - std::vector> &allValueAccessIdx, - std::vector> &allAccessIdx, - std::vector &forLoops /* numeric for-loop statements, from innermost to outermost*/, - std::vector &numeric_nested_forLoop_AccessIdx, - std::vector &symbolic_nested_forops /* symbolic for-loops from innermost to outermost */, - std::vector> &rhsPerms, - SymbolicInfo &symbolicInfo, - NumericInfo &numericInfo, - MaskingInfo &maskingInfo) - { - std::vector> rhsFormats; - getRHSFormatsOfComputeOp(cur_op.getOperation()->getResult(0), rhsFormats); - std::vector> lhsFormats; - getLHSFormatsOfComputeOp(cur_op.getOperation()->getResult(0), lhsFormats); - bool isMixedMode = checkIsMixedMode(rhsFormats); - bool isElementwise = checkIsElementwise(rhsPerms); - comet_debug() << " isElementwise:" << isElementwise << " isMixedMode: " << isMixedMode << "\n"; - auto ctx = builder.getContext(); - IndexType indexType = IndexType::get(ctx); + class LoopInfo : public IndexVar { + protected: + SmallVector currentInputs; + ValueRange results; - if ((semiringFirst.size() == 0) || (semiringSecond.size() == 0)) - llvm::errs() << "Error during semiring parsing!" - << "\n"; + public: + Operation* loopBody; + IRMapping map; // Maps values from outer scope (i.e. TreeRegion) to loop - if (main_tensors_all_Allocs.size() != allValueAccessIdx.size()) - llvm::errs() << "DEBUG ONLY: issue with main_tensor_nums size" - << "\n"; + LoopInfo(ValueRange inputs, ValueRange outputs, Operation* body, IRMapping& ir_map): + currentInputs(inputs), results(outputs), loopBody(body), map(ir_map) {} - auto f64Type = builder.getF64Type(); - auto const_f64_0 = builder.create(loc, f64Type, builder.getF64FloatAttr(0)); + virtual Value getCrd(IRRewriter& rewriter) override = 0; + virtual Value getPos(IRRewriter& rewriter, Value tensor, uint32_t dim) override = 0; + virtual void updateOutput(IRRewriter& rewriter, uint32_t idx, Value newOutput) = 0; + virtual ~LoopInfo(){}; - int main_tensor_nums = main_tensors_all_Allocs.size(); - bool compressedWorkspace = false; + ValueRange getInputs(){return currentInputs;} + Value getInput(uint32_t idx) {return currentInputs[idx];} + virtual Value getInput(IRRewriter& rewriter, uint32_t idx, SmallVector dims) { + return getInput(idx); + } + virtual ValueRange getResults() {return results;} + }; - if (comp_worksp_opt) /// always lhs is dense after workspace transformations - { - compressedWorkspace = true; + class SentinelLoopInfo : LoopInfo { + private: + indexTree::YieldOp terminator; + public: + SentinelLoopInfo(ValueRange inputs, ValueRange outputs, Operation* body, IRMapping& ir_map, indexTree::YieldOp yield_op): + LoopInfo(inputs, outputs, body, ir_map), terminator(yield_op) {} - /// Generate the numeric bitmap - if (numericInfo.ws_bitmap == nullptr) - { - Value bitmap_alloc; - genNumericBitmap(builder, - loc, - symbolic_nested_forops.back(), - symbolicInfo, - bitmap_alloc); - /// TODO(zpeng): numericInfo.ws_bitmap should be lowered from Index Tree dialect. - numericInfo.ws_bitmap = bitmap_alloc; - numericInfo.ws_bitmap_valueAccessIdx = allValueAccessIdx[lhs_loc][0]; + static LoopInfo* build(ValueRange inputs, ValueRange outputs, Operation* body, IRMapping& ir_map, indexTree::YieldOp yield_op) { + return new SentinelLoopInfo(inputs, outputs, body, ir_map, yield_op); } - /// Generate the mask-array (please not confuse with mark-array) - if (PUSH_BASED_MASKING == maskingInfo.mask_type) - { - /// Generate numeric mask-array before the numeric outermost for-loop. - /// Please don't confuse with mark-array. - genNumericMaskArray(builder, - loc, - forLoops.back() /* numeric_outermost_forLoop= */, - symbolicInfo, - numericInfo /* output */); - - /// Generate setting the mask-array before the numeric semiring for-loop and resetting after the semiring for-loop. - genNumericSetAndResetMaskArray(builder, - loc, - forLoops.back() /* numeric_outermost_forLoop */, - numeric_nested_forLoop_AccessIdx.back() /* outermost_forLoop_valueAccessIdx */, - numericInfo, - maskingInfo); + virtual Value getCrd(IRRewriter& rewriter) override {return nullptr;} + virtual Value getPos(IRRewriter& rewriter, Value tensor, uint32_t dim) override {return nullptr;} + virtual void updateOutput(IRRewriter& rewriter, uint32_t idx, Value newOutput) override { + currentInputs[idx] = newOutput; + rewriter.modifyOpInPlace(terminator, [&](){terminator.setOperand(idx, newOutput);}); } + }; - /// Value &is_visited_alloc = tensors_lhs_Allocs[1][0]; - /// Value &is_visited_alloc_valAccessIdx = allValueAccessIdx[lhs_loc][0]; - scf::IfOp if_notAlreadySet; - genCmptOpKernelIfStatementCondition(builder, - loc, - numericInfo, - maskingInfo, - if_notAlreadySet /* output */); - - /// if-then region corresponding to if_notAlreadySet instruction. - /// if (&if_notAlreadySet. getThenRegion()) - if (!if_notAlreadySet.getThenRegion().empty()) + class DenseLoopInfo : LoopInfo { + private: + Value inductionVar; + scf::YieldOp terminator; + + public: + DenseLoopInfo(ValueRange inputs, + ResultRange outputs, + Operation* body, + IRMapping ir_map, + Value i, + scf::YieldOp yield): + LoopInfo(inputs, outputs, body, ir_map), inductionVar(i), terminator(yield){} + + static LoopInfo* build(Operation* domain_op, IRRewriter& rewriter, ValueRange inputs) { - genCmptOpKernelIfStatementThenRegion(builder, - loc, - lhs_loc, - main_tensor_nums, - if_notAlreadySet, - compressedWorkspace, - semiringSecond, - main_tensors_all_Allocs, - tensors_lhs_Allocs, - allValueAccessIdx, - symbolicInfo, - numericInfo); + auto loc = domain_op->getLoc(); + Value lb = rewriter.create(loc, rewriter.getIndexType(), rewriter.getIndexAttr(0)); + Value ub = domain_op->getOperand(0); + Value step = rewriter.create(loc, rewriter.getIndexType(), rewriter.getIndexAttr(1)); + scf::ForOp for_loop = rewriter.create(loc, lb, ub, step, inputs); + + IRMapping map; + rewriter.setInsertionPointToStart(for_loop.getBody()); + auto yield_op = rewriter.create(loc, for_loop.getRegionIterArgs()); + rewriter.setInsertionPointAfter(for_loop); + return new DenseLoopInfo(for_loop.getRegionIterArgs(), for_loop->getResults(), yield_op, map, for_loop.getInductionVar(), yield_op); + } + + Value getCrd(IRRewriter& rewriter) override {return inductionVar;} + + Value getPos(IRRewriter& rewriter, Value tensor, uint32_t dim) override { + if(dyn_cast(tensor.getType())){ + return inductionVar; + } else if(SparseTensorType tt = dyn_cast(tensor.getType())){ + if((TensorFormatEnum)(tt.getFormat()[2 * dim]) == TensorFormatEnum::D) { + return inductionVar; + } + assert(false && "Invalid type passed to DenseLoopInfo getPos"); + return nullptr; + } + assert(false && "Invalid type passed to DenseLoopInfo getPos"); + return nullptr; } - /// if-else region corresponding to if_notAlreadySet instruction. - /// if (&if_notAlreadySet.getElseRegion()) - if (!if_notAlreadySet.getElseRegion().empty()) - { - genCmptOpKernelIfStatementElseRegion(builder, - loc, - lhs_loc, - main_tensor_nums, - if_notAlreadySet, - compressedWorkspace, - semiringFirst, - semiringSecond, - main_tensors_all_Allocs, - allValueAccessIdx); - } - } - else - { /// general dense or mixed mode computation, no need workspace transformations - std::vector allLoads(main_tensor_nums); - for (auto m = 0; m < main_tensor_nums; m++) - { - Value load_op = builder.create(loc, - main_tensors_all_Allocs[m][main_tensors_all_Allocs[m].size() - 1], allValueAccessIdx[m]); - allLoads[m] = load_op; - comet_debug() << " "; - comet_vdump(load_op); + void updateOutput(IRRewriter& rewriter, uint32_t idx, Value newOutput) override { + currentInputs[idx] = newOutput; + rewriter.modifyOpInPlace(terminator, [&](){terminator.setOperand(idx, newOutput);}); } - comet_debug() << " allLoads.size(): " << allLoads.size() << "\n"; + }; - /// if computeOp is elementwise mixed mode operation, the output is sparse - if (isMixedMode && isElementwise && !checkIsDense(lhsFormats[0])) + class DenseParallelLoopInfo : public LoopInfo { + private: + Value inductionVar; + SmallVector terminator_ops; + SmallDenseMap output_sets; + + public: + DenseParallelLoopInfo(ValueRange inputs, + ResultRange outputs, + Operation* body, + IRMapping ir_map, + Value i, + SmallVector& terminator_ops, + SmallDenseMap output_sets): + LoopInfo(inputs, outputs, body, ir_map), inductionVar(i), terminator_ops(terminator_ops), output_sets(output_sets) {} + + static LoopInfo* build(Operation* domain_op, IRRewriter& rewriter, ValueRange inputs, SmallDenseMap output_sets, StringAttr parallelDim = nullptr) { - - int dense_inputtensor_id = 0; - for (unsigned int i = 0; i < rhsFormats.size(); i++) - { - if (checkIsDense(rhsFormats[i])) - { - dense_inputtensor_id = i; - break; - } - } - - int sparse_inputtensor_id = dense_inputtensor_id ? 0 : 1; - std::string sparse_format = getTensorFormat(rhsFormats, sparse_inputtensor_id); - - auto last_insertionPoint = builder.saveInsertionPoint(); - - /// Need to initialize some memory accesses outside the nested loop - /// Reset the insertion point: the body of the innermost loop - comet_debug() << "LoopSize: " << forLoops.size() << " Loop:\n"; - comet_vdump(forLoops[forLoops.size() - 1]); - builder.setInsertionPoint(forLoops[forLoops.size() - 1]); - - Value const_index_0 = builder.create(loc, 0); - MemRefType memTy_alloc_Cnnz = MemRefType::get({1}, indexType); - Value alloc_Cnnz = builder.create(loc, memTy_alloc_Cnnz); - comet_debug() << " AllocOp for Cnnz: "; - comet_vdump(alloc_Cnnz); - - std::vector alloc_Cnnz_insert_loc = {const_index_0}; -#ifdef DEBUG_MODE_LowerIndexTreeToSCFPass - auto store_Cnnz = builder.create(loc, const_index_0, alloc_Cnnz, alloc_Cnnz_insert_loc); - comet_debug() << " StoreOp: "; - comet_vdump(store_Cnnz); -#else - builder.create(loc, const_index_0, alloc_Cnnz, alloc_Cnnz_insert_loc); -#endif - - /// The following code block is needed to update Update C2pos in the case of output tensor is in DCSR - Value Cnnz_index_old; - Value alloc_Cnnz_row; - if (sparse_format.compare("DCSR") == 0) + auto loc = domain_op->getLoc(); + Value ub = domain_op->getOperand(0); + scf::ForallOp for_loop = rewriter.create( + loc, + ArrayRef(ub), + inputs, + std::nullopt, + nullptr); + if(parallelDim != nullptr) { - alloc_Cnnz_row = builder.create(loc, memTy_alloc_Cnnz); -#ifdef DEBUG_MODE_LowerIndexTreeToSCFPass - auto store_Cnnz_row = builder.create(loc, const_index_0, alloc_Cnnz_row, - alloc_Cnnz_insert_loc); - comet_debug() << " StoreOp DCSR: "; - comet_vdump(store_Cnnz_row); -#else - builder.create(loc, const_index_0, alloc_Cnnz_row, alloc_Cnnz_insert_loc); -#endif - /// Get Cnnz_old - Cnnz_index_old = builder.create(loc, alloc_Cnnz, alloc_Cnnz_insert_loc); + for_loop->setAttr("parallelDim", parallelDim); } - builder.restoreInsertionPoint(last_insertionPoint); + Value inductionVar = for_loop.getInductionVar(0); + rewriter.setInsertionPointToStart(for_loop.getBody()); - comet_debug() << " dense_inputtensor_id: " << dense_inputtensor_id << "\n"; - comet_debug() << " sparse_inputtensor_id: " << sparse_inputtensor_id << "\n"; - Value denseInput_is_nonzero = builder.create(loc, CmpFPredicate::ONE, allLoads[dense_inputtensor_id], - const_f64_0); - auto if_nonzero = builder.create(loc, denseInput_is_nonzero, /*WithElseRegion*/ false); - comet_debug() << " If branch:\n"; - comet_vdump(if_nonzero); - - if (!if_nonzero.getThenRegion().empty()) + SmallVector input_slices; + for(auto& input : for_loop.getOutputsMutable()) { - - builder.setInsertionPointToStart(&if_nonzero.getThenRegion().front()); - - comet_debug() << "calculate product and sum in \n"; - Value elementWiseResult = getSemiringSecondVal(builder, loc, semiringSecond, allLoads[0], allLoads[1], - compressedWorkspace); - - /// Get Cnnz - Value Cnnz_index = builder.create(loc, alloc_Cnnz, alloc_Cnnz_insert_loc); - -/// Store product to Cval -#ifdef DEBUG_MODE_LowerIndexTreeToSCFPass - comet_debug() << "Store product to Cval\n"; - auto store_Cval = builder.create(loc, elementWiseResult, main_tensors_all_Allocs[2][main_tensors_all_Allocs[2].size() - 1], Cnnz_index); - comet_debug() << " StoreOp: "; - comet_vdump(store_Cval); - - /// Update C1crd, C2crd - comet_debug() << "Getting A1crd\n"; - comet_debug() << "allValueAccessIdx[" << sparse_inputtensor_id << "].size(): " - << allAccessIdx[sparse_inputtensor_id].size() << "\n"; - comet_vdump(allAccessIdx[sparse_inputtensor_id][0]); - - for (unsigned int i = 0; i < allAccessIdx.size(); i++) + auto& os = output_sets.at(input.get()); + RankedTensorType tt = llvm::cast(input.get().getType()); + int64_t nDims = tt.getRank(); + SmallVector sizes; + for(int i = 0; i < nDims; i++) { - comet_debug() << "allAccessIdx[" << i << "].size(): " << allAccessIdx[i].size() << "\n"; - for (auto n : allAccessIdx[i]) - { - comet_vdump(n); + if(i != os.dim && tt.getDimSize(i) == ShapedType::kDynamic) { + Value idx = rewriter.create(loc, rewriter.getIndexType(), rewriter.getIndexAttr(i)); + sizes.push_back(rewriter.create(loc, rewriter.getIndexType(), input.get(), idx)); } } -#else - builder.create(loc, elementWiseResult, main_tensors_all_Allocs[2][main_tensors_all_Allocs[2].size() - 1], Cnnz_index); -#endif - comet_debug() << "Store C1crd\n"; - /// Branch out COO... CSR... DCSR... - if (sparse_format.compare("COO") == 0) - { - comet_debug() << "COO format for Elementwise MulOp, update all coordinates\n"; - for (unsigned d = 0; d < rhsPerms[sparse_inputtensor_id].size(); d++) - { - Value crd = allAccessIdx[sparse_inputtensor_id][d]; -#ifdef DEBUG_MODE_LowerIndexTreeToSCFPass - auto store_coo_crd = builder.create(loc, crd, main_tensors_all_Allocs[2][4 * d + 1], - Cnnz_index); - comet_debug() << " COO StoreOp: "; - comet_vdump(store_coo_crd); -#else - builder.create(loc, crd, main_tensors_all_Allocs[2][4 * d + 1], Cnnz_index); -#endif - } - } - else if (sparse_format.compare("CSR") == 0 || sparse_format.compare("DCSR") == 0) + SmallVector static_offsets(nDims, 0); + SmallVector static_sizes(tt.getShape()); + static_offsets[os.dim] = ShapedType::kDynamic; + static_sizes[os.dim] = 1; + + auto slice = rewriter.create( + loc, + RankedTensorType::get(static_sizes, tt.getElementType()), + for_loop.getTiedBlockArgument(&input), + ValueRange(inductionVar), + sizes, + ValueRange(), + rewriter.getDenseI64ArrayAttr(static_offsets), + rewriter.getDenseI64ArrayAttr(static_sizes), + rewriter.getDenseI64ArrayAttr(SmallVector(nDims, 1)) + ); + input_slices.push_back(slice->getResult(0)); + output_sets.insert(std::make_pair(slice->getResult(0), os)); + } + + SmallVector terminator_ops; + auto par_op = for_loop.getTerminator(); + auto slice = input_slices.begin(); + for(auto& input : for_loop.getOutputsMutable()) { + rewriter.setInsertionPoint(par_op); + auto& os = output_sets.at(input.get()); + auto tt = cast(slice->getType()); + int64_t nDims = tt.getRank(); + SmallVector sizes; + for(int i = 0; i < nDims; i++) { - for (unsigned int d = forLoops.size() - 1; d < rhsPerms[sparse_inputtensor_id].size(); d++) - { - Value crd = allAccessIdx[sparse_inputtensor_id][d]; -#ifdef DEBUG_MODE_LowerIndexTreeToSCFPass - auto store_csr_crd = builder.create(loc, crd, main_tensors_all_Allocs[2][4 * d + 1], - Cnnz_index); - comet_debug() << " CSR or DCSR StoreOp: "; - comet_vdump(store_csr_crd); -#else - builder.create(loc, crd, main_tensors_all_Allocs[2][4 * d + 1], Cnnz_index); -#endif + if(tt.getDimSize(i) == ShapedType::kDynamic) { + Value idx = rewriter.create(loc, rewriter.getIndexType(), rewriter.getIndexAttr(i)); + sizes.push_back(rewriter.create(loc, rewriter.getIndexType(), input.get(), idx)); } } - /// Update Cnnz - comet_debug() << "Update Cnnz\n"; - Value const_index_1 = builder.create(loc, 1); - Value new_Cnnz_index = builder.create(loc, Cnnz_index, const_index_1); -#ifdef DEBUG_MODE_LowerIndexTreeToSCFPass - comet_debug() << "AddIOps (new_Cnnz_index): "; - comet_vdump(new_Cnnz_index); - auto store_updated_cnnz = builder.create(loc, new_Cnnz_index, alloc_Cnnz, - alloc_Cnnz_insert_loc); - comet_debug() << " Update Cnnz (store new value) StoreOp: "; - comet_vdump(store_updated_cnnz); -#else - builder.create(loc, new_Cnnz_index, alloc_Cnnz, alloc_Cnnz_insert_loc); -#endif + rewriter.setInsertionPointToEnd(par_op.getBody()); + SmallVector static_offsets(nDims, 0); + static_offsets[os.dim] = ShapedType::kDynamic; + + Operation* insert = rewriter.create( + loc, + *slice, + for_loop.getTiedBlockArgument(&input), + inductionVar, + sizes, + ValueRange(), + rewriter.getDenseI64ArrayAttr(static_offsets), + rewriter.getDenseI64ArrayAttr(tt.getShape()), + rewriter.getDenseI64ArrayAttr(SmallVector(nDims, 1)) + ); + terminator_ops.push_back(insert); + slice++; } - - /// Need to identify dense tensor upperbound to be able to update Cpos and Csize arrays - std::vector denseAllocs = tensors_rhs_Allocs[dense_inputtensor_id]; - assert(denseAllocs.size() == 1); - - comet_debug() << " DenseAllocs: "; - auto inputType = denseAllocs[0].getType(); - std::vector denseDimsSize; - for (unsigned rank = 0; rank < inputType.cast().getRank(); rank++) - { - auto dimSize = inputType.cast().getDimSize(rank); - Value upperBound; - if (dimSize == ShapedType::kDynamic) - { - comet_debug() << " This dimension is a dynamic size:\n"; - unsigned dynamicDimPos = inputType.dyn_cast().getDynamicDimIndex(rank); - comet_debug() << " DynamicDimPos: " << dynamicDimPos << "\n"; - upperBound = denseAllocs[0].getDefiningOp()->getOperand(dynamicDimPos); - comet_vdump(upperBound); + rewriter.setInsertionPointAfter(for_loop); + + IRMapping map; + return new DenseParallelLoopInfo(input_slices, for_loop.getResults(), par_op, map, inductionVar, terminator_ops, output_sets); + } + + Value getCrd(IRRewriter& rewriter) override {return inductionVar;} + + Value getPos(IRRewriter& rewriter, Value tensor, uint32_t dim) override { + if(dyn_cast(tensor.getType())){ + // TODO: Fix me. Right now, getCrd on the output tensor is just zero becuase + // we have already extracted the slice that we want. This will only work + // in very limited scenarios + if(std::find(currentInputs.begin(), currentInputs.end(), tensor) != currentInputs.end()){ + Value zero = rewriter.create(tensor.getLoc(), rewriter.getIndexType(), rewriter.getIndexAttr(0)); + return zero; + } else { + return inductionVar; } - else - { - builder.setInsertionPoint(forLoops[1]); /// [PT] This is needed to make sure we can access the values created outside of the loop - comet_debug() << " This dimension is a static size\n"; - upperBound = builder.create(loc, dimSize); - comet_vdump(upperBound); + + } else if(SparseTensorType tt = dyn_cast(tensor.getType())){ + if((TensorFormatEnum)(tt.getFormat()[2 * dim]) == TensorFormatEnum::D) { + return inductionVar; } + assert(false && "Invalid type passed to DenseLoopInfo getPos"); + return nullptr; + } + assert(false && "Invalid type passed to DenseLoopInfo getPos"); + return nullptr; + } - denseDimsSize.push_back(upperBound); + void updateOutput(IRRewriter& rewriter, uint32_t idx, Value newOutput) override { + currentInputs[idx] = newOutput; + Operation* terminator = terminator_ops[idx]; + rewriter.modifyOpInPlace(terminator, [&](){terminator->setOperand(0, newOutput);}); + } + }; + + class SparseLoopInfo : LoopInfo { + private: + Value inductionVar; + scf::YieldOp terminator; + Value controlTensor; + uint32_t dim; + Value crd; + + public: + SparseLoopInfo(ValueRange inputs, ResultRange outputs, Operation* body, IRMapping ir_map, Value i, scf::YieldOp yield, Value tensor, uint32_t d): + LoopInfo(inputs, outputs, body, ir_map), inductionVar(i), terminator(yield), controlTensor(tensor), dim(d), crd(nullptr) {} + + static LoopInfo* build(Operation* domain_op, IRRewriter& rewriter, ValueRange inputs) + { + auto loc = domain_op->getLoc(); + auto sparse_domain = cast(domain_op); + auto index_type = rewriter.getIndexType(); + + Value start_idx = sparse_domain.getParent(); + if(!start_idx){ + start_idx = rewriter.create(loc, index_type, rewriter.getIndexAttr(0)); } - builder.setInsertionPointAfter(forLoops[0]); - /// To update Cpos - if (sparse_format.compare("CSR") == 0) - { - builder.setInsertionPointAfter(forLoops[0]); - Value const_index_1 = builder.create(loc, 1); - Value arg0_next = builder.create(loc, forLoops[1].getInductionVar(), const_index_1); - comet_debug() << "AddIOp (arg0_next): "; - comet_vdump(arg0_next); - - Value Cnnz_index_final = builder.create(loc, alloc_Cnnz, alloc_Cnnz_insert_loc); - builder.create(loc, Cnnz_index_final, main_tensors_all_Allocs[2][4], arg0_next); /// 2 - - builder.setInsertionPointAfter(forLoops[1]); - /// Update C2pos[0] - comet_debug() << "Update C2pos[0]\n"; - std::vector insert_loc_0 = {const_index_0}; - builder.create(loc, const_index_0, main_tensors_all_Allocs[2][4], insert_loc_0); /// 2 - - /// Update C1pos[0] - comet_debug() << "Update C1pos[0]\n"; - Value dim0_index = denseDimsSize[0]; - builder.create(loc, dim0_index, main_tensors_all_Allocs[2][0], insert_loc_0); + Value inc = rewriter.create(loc, index_type, rewriter.getIndexAttr(1)); + Value end_idx = rewriter.create(loc, index_type, start_idx, inc); + Value lb = rewriter.create(loc, sparse_domain.getPos(), start_idx); + lb = rewriter.createOrFold(loc, rewriter.getIndexType(), lb); + Value ub = rewriter.create(loc, sparse_domain.getPos(), end_idx); + ub = rewriter.createOrFold(loc, rewriter.getIndexType(), ub); + Value step = rewriter.create(loc, rewriter.getIndexType(), + rewriter.getIndexAttr(1)); + scf::ForOp for_loop = rewriter.create(loc, lb, ub, step, inputs); + + IRMapping map; + unsigned init_arg_idx = 0; + for(Value init_arg : inputs){ + map.map(init_arg, for_loop.getRegionIterArg(init_arg_idx)); + init_arg_idx += 1; } - else - { - if (sparse_format.compare("DCSR") == 0) - { - /// Update C2pos - comet_debug() << "Update DCSR C2pos\n"; - builder.setInsertionPointAfter(forLoops[0]); - auto Cnnz_index_new = builder.create(loc, alloc_Cnnz, alloc_Cnnz_insert_loc); - auto has_nnz_row = builder.create(loc, CmpIPredicate::ne, Cnnz_index_new, Cnnz_index_old); - auto has_nnz_row_ifOp = builder.create(loc, has_nnz_row, /*WithElseRegion*/ false); - comet_debug() << " If branch:\n"; - comet_vdump(has_nnz_row_ifOp); - - if (!has_nnz_row_ifOp.getThenRegion().empty()) - { - builder.setInsertionPointToStart(&has_nnz_row_ifOp.getThenRegion().front()); - - Value const_index_1 = builder.create(loc, 1); - Value arg0_next = builder.create(loc, forLoops[1].getInductionVar(), const_index_1); - comet_debug() << "AddIOp (arg0_next): "; - comet_vdump(arg0_next); - - Value Cnnz_index_final = builder.create(loc, alloc_Cnnz, alloc_Cnnz_insert_loc); - builder.create(loc, Cnnz_index_final, main_tensors_all_Allocs[2][4], arg0_next); /// C2pos //2 - Value Cnnz_row_index = builder.create(loc, alloc_Cnnz_row, alloc_Cnnz_insert_loc); - Value idx_i = allAccessIdx[sparse_inputtensor_id][0]; - builder.create(loc, /*i*/ idx_i, main_tensors_all_Allocs[2][1], Cnnz_row_index); /// C1crd - Value Cnnz_row_index_new = builder.create(loc, Cnnz_row_index, const_index_1); - comet_debug() << "AddIOp (Cnnz_row_index_new): "; - comet_vdump(Cnnz_row_index_new); - builder.create(loc, Cnnz_row_index_new, alloc_Cnnz_row, - alloc_Cnnz_insert_loc); /// Update Cnnz_row - } + + rewriter.setInsertionPointToStart(for_loop.getBody()); + auto yield_op = rewriter.create(loc, for_loop.getRegionIterArgs()); + rewriter.setInsertionPointAfter(for_loop); - builder.setInsertionPointAfter(forLoops[1]); - Value const_index_1 = builder.create(loc, 1); - std::vector insert_loc_1 = {const_index_1}; + return new SparseLoopInfo(ValueRange(for_loop.getRegionIterArgs()), for_loop.getResults(), yield_op, map, for_loop.getInductionVar(), yield_op, sparse_domain.getTensor(), sparse_domain.getDim()); + } - /// Update C2pos[0] - std::vector insert_loc_0 = {const_index_0}; - builder.create(loc, const_index_0, main_tensors_all_Allocs[2][4], insert_loc_0); /// 2 + Value getCrd(IRRewriter& rewriter) override { + // if(crd != nullptr) return crd; + auto loc = controlTensor.getLoc(); + // SparseTensorType tt = cast(controlTensor.getType()); + crd = rewriter.create(loc, controlTensor, inductionVar, rewriter.getI32IntegerAttr(dim)); + crd = rewriter.createOrFold(loc, rewriter.getIndexType(), crd); + return crd; + } - /// Update C1pos[0], C1pos[1] - Value Cnnz_row_index = builder.create(loc, alloc_Cnnz_row, alloc_Cnnz_insert_loc); - builder.create(loc, const_index_0, main_tensors_all_Allocs[2][0], insert_loc_0); - builder.create(loc, Cnnz_row_index, main_tensors_all_Allocs[2][0], insert_loc_1); - } - else - { - if (sparse_format.compare("COO") == 0) - { - /// Finally, Update C1pos - comet_debug() << "Update C1pos\n"; - builder.setInsertionPointAfter(forLoops[0]); - Value Cnnz_index_final = builder.create(loc, alloc_Cnnz, alloc_Cnnz_insert_loc); - Value const_index_1 = builder.create(loc, 1); - builder.create(loc, const_index_0, main_tensors_all_Allocs[2][0], const_index_0); - builder.create(loc, Cnnz_index_final, main_tensors_all_Allocs[2][0], const_index_1); - } - else - llvm::errs() << "/// Coordinate values are not updated for output sparse tensor in " << sparse_format - << " format\n"; + Value getPos(IRRewriter& rewriter, Value tensor, uint32_t dim) override { + if(tensor == controlTensor && dim == this->dim){ + return inductionVar; + } + + if(dyn_cast(tensor.getType()) || dyn_cast(tensor.getType())){ + return crd; + } else if(SparseTensorType tt = dyn_cast(tensor.getType())){ + if((TensorFormatEnum)(tt.getFormat()[2 * dim]) == TensorFormatEnum::D) { + return crd; } } + auto loc = controlTensor.getLoc(); + Value crd = rewriter.create(loc, controlTensor, inductionVar, rewriter.getI32IntegerAttr(dim)); + crd = rewriter.createOrFold(loc, rewriter.getIndexType(), crd); + Value pos = rewriter.create(loc, rewriter.getIndexType(), tensor, crd, rewriter.getI32IntegerAttr(dim), rewriter.getBoolAttr(true)); + return pos; + } - } /// end if (isMixedMode && isElementwise) - else - { - /// calculate elementWise operation and reduction for general dense or mix mode computation (which has dense output) - comet_debug() - << "calculate elementWise operation and reduction for general dense or mix mode computation (which has dense output)\n"; - Value elementWiseResult = getSemiringSecondVal(builder, loc, semiringSecond, allLoads[0], allLoads[1], - compressedWorkspace); - Value reduceResult = getSemiringFirstVal(builder, loc, semiringFirst, allLoads[2], elementWiseResult, - compressedWorkspace); - builder.create(loc, reduceResult, - main_tensors_all_Allocs[2][main_tensors_all_Allocs[2].size() - 1], - allValueAccessIdx[2]); + void updateOutput(IRRewriter& rewriter, uint32_t idx, Value newOutput) override { + currentInputs[idx] = newOutput; + rewriter.modifyOpInPlace(terminator, [&](){terminator.setOperand(idx, newOutput);}); } - } - } + }; - /// ----------------- /// - /// Generate Cij = Wj node, gathering the results in the workspace to the sparse output C.val. - /// Called by genCmptOps(). - /// ----------------- /// - /// sort(C.col, C.rowptr[i_idx], C.rowptr[i_idx + 1]); - /// for (int j_loc = C.rowptr[i_idx]; j_loc < C.rowptr[i_idx + 1]; ++j_loc) { - /// int j_idx = C.col[j_loc]; - /// C.val[j_idx] = W_data[j_idx]; - /// is_visited[j_idx] = false; - /// } - /// ----------------- /// - /// %rowptr_bound = memref.load %rowptr[%c0] : memref<1xindex> - /// %C_col_ptr = memref.cast %C_col : memref to memref<*xindex> - /// func.call @comet_sort_index(%C_col_ptr, %rowptr_start, %rowptr_bound) : (memref<*xindex>, index, index) -> () - /// - /// scf.for %ptr = %rowptr_start to %rowptr_bound step %c1 { - /// %c_col_id = memref.load %C_col[%ptr] : memref /// c_col_id = C_col[ptr] - /// %data = memref.load %ws_data[%c_col_id] : memref /// data = ws_data[c_col_id] - /// memref.store %data, %C_val[%ptr] : memref /// C_val[ptr] = data - /// memref.store %false, %ws_bitmap[%c_col_id] : memref /// ws_bitmap[c_col_id] = false - /// } - void genWorkspaceCmptOpGatherFromWorkspaceToOutput(OpBuilder &builder, - Location &loc, - std::vector> &tensors_rhs_Allocs, - std::vector &nested_forops, - std::vector &nested_AccessIdx, - SymbolicInfo &symbolicInfo, - NumericInfo &numericInfo) - { - /// Store the insertion point - auto last_insertion_point = builder.saveInsertionPoint(); - - assert(nested_forops.size() >= 2 && nested_AccessIdx.size() >= 2 && "Error: should be at least 2 levels of for-loop.\n"); - AbstractLoopOp &curr_for_loop = nested_forops[0]; - AbstractLoopOp parent_for_loop = nested_forops[1]; - - /// Set the insertion point before the innermost for-loop - builder.setInsertionPoint(curr_for_loop); - - /// Get and set the boundary of current for-loop - /// %rowptr_start = memref::LoadOp %C_rowptr[%i_idx] : memref - /// %id_idx_plus_one = arith.addi %rowptr_start, %c1 : index - /// %rowptr_bound = memref::LoadOp %C_rowptr[%i_idx_plus_one] : memref - Value const_index_1 = builder.create(loc, 1); - Value i_idx = nested_AccessIdx[1]; - Value i_idx_plus_one = builder.create(loc, i_idx, const_index_1); - Value &mtxC_rowptr = symbolicInfo.mtxC_rowptr; - Value rowptr_start = builder.create(loc, mtxC_rowptr, ValueRange{i_idx}); - Value rowptr_bound = builder.create(loc, mtxC_rowptr, ValueRange{i_idx_plus_one}); - { - comet_vdump(parent_for_loop); - comet_vdump(i_idx); - comet_vdump(rowptr_start); - comet_vdump(rowptr_bound); - } - - /// Generate calling comet_sort_index - /// %C_col_ptr = memref.cast %C_col : memref to memref<*xindex> - /// func.call @comet_sort_index(%C_col_ptr, %rowptr_start, %rowptr_bound) : (memref<*xindex>, index, index) -> () - std::string func_name = "comet_sort_index"; - Value &mtxC_col = symbolicInfo.mtxC_col; - IndexType indexType = IndexType::get(builder.getContext()); - Value C_col_cast = builder.create(loc, - UnrankedMemRefType::get(indexType, 0), - mtxC_col); - builder.create(loc, - func_name, - SmallVector{}, - ValueRange{C_col_cast, rowptr_start, rowptr_bound}); - - /// Change current for-loop boundaries - curr_for_loop.setLowerBound(rowptr_start); - curr_for_loop.setUpperBound(rowptr_bound); - - /// Generate current for-loop body - Value &mtxC_val = symbolicInfo.mtxC_val; - Value &ws_data = tensors_rhs_Allocs[0][0]; - Value &ws_bitmap = numericInfo.ws_bitmap; - Value rowptr = curr_for_loop.getInductionVar(); - builder.setInsertionPointToStart(curr_for_loop.getBody()); - Value c_col_id = builder.create(loc, mtxC_col, ValueRange{rowptr}); - Value data = builder.create(loc, ws_data, ValueRange{c_col_id}); - builder.create(loc, - data, - mtxC_val, - ValueRange{rowptr}); - Value const_i1_0 = builder.create(loc, builder.getI1Type(), builder.getBoolAttr(false)); - builder.create(loc, - const_i1_0, - ws_bitmap, - ValueRange{c_col_id}); - { - comet_vdump(c_col_id); - comet_vdump(data); - comet_vdump(curr_for_loop); - } - - /// Free up ws_data and ws_bitmap after - builder.setInsertionPointAfter(parent_for_loop); - builder.create(loc, ws_data); - builder.create(loc, ws_bitmap); - - /// Restore the insertion point - builder.restoreInsertionPoint(last_insertion_point); - } - - /// In genCmptOps, get current compute node's numeric nested for-loop and access indices. - void getNumericNestedForOpsAndAccessIdx(std::vector &ancestorsWps, - std::vector &ancestorsOps, - std::vector &nested_forops /* output */, - std::vector &nested_AccessIdx /* output */, - std::vector &nested_forops_indices /* output */) - { - - for (unsigned int i = 0; i < ancestorsOps.size(); i++) - { - comet_debug() << " ancestorsOps[" << i << "]->forOps.size(): " << ancestorsOps[i]->forOps.size() - << ", ancestorsOps->id: " - << ancestorsOps[i]->id << "\n"; - if (!ancestorsOps[i]->forOps.empty()) - { /// for loops OpsTree node - for (int j = ancestorsOps[i]->forOps.size() - 1; j >= 0; j--) - { - comet_debug() << " j: " << j << "\n"; - nested_forops.push_back(ancestorsOps[i]->forOps[j]); - comet_debug() << "AccessIdx: " << ancestorsOps[i]->accessIdx[j] << "\n"; - nested_AccessIdx.push_back(ancestorsOps[i]->accessIdx[j]); - } - } - } - comet_debug() << " nested_forops.size(): " << nested_forops.size() << "\n"; - for (unsigned int i = 0; i < ancestorsWps.size(); i++) - { - comet_debug() << " "; - comet_vdump(ancestorsWps[i]); - - if (indexTree::IndexTreeIndicesOp cur_op = dyn_cast( - ancestorsWps[i].getDefiningOp())) + class SingletonLoopInfo : LoopInfo { + private: + LoopInfo* parent; + Value inductionVar; + Value tensor; + int32_t dim; + + public: + SingletonLoopInfo(ValueRange inputs, Operation* body, IRMapping ir_map, LoopInfo* parent_info, Value i, Value sparse_tensor, int32_t dim) : + LoopInfo(inputs, ValueRange(inputs), body, ir_map), parent(parent_info), inductionVar(i), tensor(sparse_tensor), dim(dim) + {} + static LoopInfo* build(Operation* domain_op, IRRewriter& rewriter, ValueRange inputs, LoopInfo* parent_info) { - /// Get indices - ArrayAttr op_indices = cur_op.getIndices(); - - if (op_indices.size() > 0) - { /// for loops OpsTree node - for (int j = op_indices.size() - 1; j >= 0; j--) - { - /// Get the indices; - int64_t idx = op_indices[j].cast().getInt(); - nested_forops_indices.push_back(idx); - } + auto sparse_domain = cast(domain_op); + Value inductionVar = sparse_domain.getParent(); + IRMapping map; + return new SingletonLoopInfo(inputs, &(*(rewriter.saveInsertionPoint().getPoint())), map, parent_info, inductionVar, sparse_domain.getTensor(), sparse_domain.getDim()); + } + + Value getCrd(IRRewriter& rewriter) override { + auto loc = rewriter.getUnknownLoc(); // This is not correct + // auto tt = RankedTensorType::get({ShapedType::kDynamic}, rewriter.getIndexType()); + Value crd_tensor = rewriter.create(loc, tensor, rewriter.getI32IntegerAttr(dim)); + Value crd = rewriter.create(loc, crd_tensor, inductionVar); + crd = rewriter.create(loc, rewriter.getIndexType(), crd); + return crd; + } + + Value getPos(IRRewriter& rewriter, Value tensor, uint32_t dim) override { + // This is needed + if(llvm::isa(tensor.getType())) { + // For dense tensors, positions and crd should be the same. + return getCrd(rewriter); } + return inductionVar; } - } - } - /// In genCmptOps, get current compute node's RHS, LHS, tensors, formats, perms, etc. - void getNumericTensors(indexTree::IndexTreeComputeOp &cur_op, - std::vector &tensors_rhs /* output */, - std::vector> &tensors_lhs_Allocs /* output */, - std::vector> &tensors_rhs_Allocs /* output */, - std::vector> &allFormats /*output*/, - std::vector> &allPerms /* output */, - std::vector> &allPerms_rhs /* output */, - std::vector &main_tensors_all /* output */, - std::vector &main_tensors_rhs /* output */) - { - comet_vdump(cur_op); - for (auto n : cur_op.getRhs()) - { - comet_debug() << " "; - comet_vdump(n); - for (unsigned i = 0; i < n.getDefiningOp()->getNumOperands(); i++) - { - comet_debug() << " "; - comet_vdump(n.getDefiningOp()->getOperand(i)); - tensors_rhs.push_back(n.getDefiningOp()->getOperand(i)); - } - } - - std::vector tensors_lhs; /// inner - for (unsigned i = 0; i < cur_op.getLhs().getDefiningOp()->getNumOperands(); i++) - { - comet_debug() << " "; - comet_vdump(cur_op.getLhs().getDefiningOp()->getOperand(i)); - tensors_lhs.push_back(cur_op.getLhs().getDefiningOp()->getOperand(i)); - } - /// Currently, only one case, the rhs is constant. Wj = 0.0; - tensors_lhs_Allocs = getAllAllocs(tensors_lhs); /// output - comet_debug() << " tensors_lhs_Allocs.size(): " << tensors_lhs_Allocs.size() << "\n"; - tensors_rhs_Allocs = getAllAllocs(tensors_rhs); /// output - comet_debug() << " tensors_rhs_Allocs.size(): " << tensors_rhs_Allocs.size() << "\n"; - -#ifdef DEBUG_MODE_LowerIndexTreeToSCFPass - comet_debug() << " tensors_rhs_Allocs: \n"; - for (auto m : tensors_rhs_Allocs) - { - comet_debug() << " "; - for (auto n : m) - { - comet_vdump(n); + void updateOutput(IRRewriter& rewriter, uint32_t idx, Value newOutput) override { + currentInputs[idx] = newOutput; + parent->updateOutput(rewriter, idx, newOutput); } - comet_debug() << "\n"; - } -#endif - getPermsOfComputeOp(cur_op.getOperation()->getResult(0), allPerms); - - comet_debug() << " allPerms: \n"; - for (auto m : allPerms) - { - comet_debug() << " "; /// print_vector(m); - for (auto n : m) - { - comet_debug() << n << " "; + ValueRange getResults() override { + return ValueRange(currentInputs); } - comet_debug() << "\n"; - } - - getFormatsOfComputeOp(cur_op.getOperation()->getResult(0), allFormats); - comet_debug() << " allFormats: \n"; - for (auto m : allFormats) - { - comet_debug() << " "; - for (auto n : m) - { - comet_debug() << n << " "; - } - comet_debug() << "\n"; - } - - comet_debug() << " "; - comet_vdump(cur_op); - - assert(allPerms.size() == allFormats.size() && "allPerms.size() != allFormats.size()\n"); - for (unsigned int m = 0; m < allPerms.size(); m++) - { - assert(allPerms[m].size() == allFormats[m].size() && "allPerms[m].size() != allFormats[m].size()\n"); - } - comet_debug() << " allPerms.size(): " << allPerms.size() << "\n"; - /// tensor_nums means the actual tensors except the auxiliary tensors - /// Suppose for LHSOp, there are "n" real tensors, then allPerms[m].size() - - getRHSPermsOfComputeOp(cur_op.getOperation()->getResult(0), allPerms_rhs); - comet_debug() << " allPerms_rhs.size(): " << allPerms_rhs.size() << "\n"; - std::vector> allPerms_lhs; /// inner - getLHSPermsOfComputeOp(cur_op.getOperation()->getResult(0), allPerms_lhs); - - comet_debug() << " allPerms_lhs.size(): " << allPerms_lhs.size() << "\n"; - std::vector main_tensors_lhs; /// inner - if (tensors_rhs.size() == allPerms_rhs.size()) - { /// all are "main" tensors - main_tensors_rhs.insert(main_tensors_rhs.end(), tensors_rhs.begin(), tensors_rhs.end()); - } - else - { /// the rhs contains the auxiliary tensors - assert(allPerms_rhs.size() == 1 && - " rhs contains auxiliary tensors and main tensors at the same time, not support currently\n"); /// only 1 main tensor on rhs - main_tensors_rhs.push_back(tensors_rhs[0]); - } - comet_debug() << " main_tensors_rhs.size(): " << main_tensors_rhs.size() << "\n"; - - if (tensors_lhs.size() == allPerms_lhs.size()) - { /// all are "main" tensors - main_tensors_lhs.insert(main_tensors_lhs.end(), tensors_lhs.begin(), tensors_lhs.end()); - } - else - { /// the lhs contains the auxiliary tensors - assert(allPerms_lhs.size() == 1 && - " lhs contains auxiliary tensors and main tensors at the same time, not support currently\n"); /// only 1 main tensor on lhs - main_tensors_lhs.push_back(tensors_lhs[0]); - } - comet_debug() << " main_tensors_lhs.size(): " << main_tensors_lhs.size() << "\n"; + + }; - main_tensors_all = main_tensors_rhs; - main_tensors_all.insert(main_tensors_all.end(), main_tensors_lhs.begin(), main_tensors_lhs.end()); - comet_debug() << " main_tensors_all.size(): " << main_tensors_all.size() << "\n"; - } + class WorkspaceLoopInfo : LoopInfo { + private: + Value inductionVar; + scf::YieldOp terminator; + Value workspaceTensor; + Value crd; - /// In genCmptOps, get for-loops' value access indices. - /// A value access index is not necessarily the for-loop's induction variable. - /// For example, To access sparse matrix C.val, we need to get rowptr = C.col[idx], then rowptr is the access index. - /// This function is used both by Numeric Phase and Symbolic Phase. - void getForLoopsValueAccessIdx(OpBuilder &builder, - Location &loc, - int main_tensor_nums, - std::vector> &allPerms, - std::vector> &allFormats, - std::vector &main_tensors_all, - std::vector &nested_forops, - std::vector &nested_AccessIdx, - std::vector &nested_forops_indices, - std::vector> &main_tensors_all_Allocs, - std::vector> &allAccessIdx /* output */, - std::vector> &allValueAccessIdx /* output */) - { + public: + WorkspaceLoopInfo(ValueRange inputs, ResultRange outputs, Operation* body, IRMapping ir_map, Value i, scf::YieldOp yield, Value workspace) : + LoopInfo(inputs, outputs, body, ir_map), inductionVar(i), terminator(yield), workspaceTensor(workspace) {} - std::vector> allLoopsArg(main_tensor_nums); /// inner - /// std::vector> allAccessIdx(main_tensor_nums); /// output - for (unsigned int i = 0; i < main_tensors_all.size(); i++) - { - for (unsigned int j = 0; j < allPerms[i].size(); j++) + static LoopInfo* build(Operation* domain_op, IRRewriter& rewriter, ValueRange inputs) { - unsigned int index_loc = findIndexInVector(nested_forops_indices, allPerms[i][j]); - comet_debug() << " index_loc " << index_loc << "\n"; - comet_debug() << " Perm: " << allPerms[i][j] << "\n"; - comet_debug() << " Format: " << allFormats[i][j] << "\n"; - assert(index_loc < nested_forops.size() && - "index_loc < nested_forops.size(), i.e. the index not exist in nested for loop\n"); - allLoopsArg[i].push_back(nested_forops[index_loc].getInductionVar()); - allAccessIdx[i].push_back(nested_AccessIdx[index_loc]); - } - /// Consider for the case w_index_list_size - /// if allPerms[i].size() == 0 - } - - /// std::vector> allValueAccessIdx(main_tensor_nums); /// output - for (int i = 0; i < main_tensor_nums; i++) - { /// If constantOp, do not consider it - comet_debug() << " "; - comet_vdump(main_tensors_all[i]); - if (main_tensors_all[i].getType().isa()) - { /// sparse tensor - - /// Find the last sparse index m, then loop_arg * all dense loop args - unsigned lastSparseIndexLoc = allPerms[i].size(); - for (int d = (int)allPerms[i].size() - 1; d >= 0; d--) - { - if (allFormats[i][d].compare(0, 1, "D") != 0 && - allFormats[i][d].compare(0, 1, "S") != 0) - { /// sparse dimension and has a loop, i.e. "CU" or "CN" - lastSparseIndexLoc = d; + auto loc = domain_op->getLoc(); + auto workspace_domain_op = llvm::cast(domain_op); + auto index_type = rewriter.getIndexType(); + + Value workspace = workspace_domain_op.getTensor(); + Type workspace_type = workspace.getType(); + Value sorted_workspace = rewriter.create(loc, workspace_type, workspace); + Value lb = rewriter.create(loc, index_type, rewriter.getIndexAttr(0)); + Value ub = rewriter.create(loc, index_type, sorted_workspace); + Value step = rewriter.create(loc, rewriter.getIndexType(), + rewriter.getIndexAttr(1)); + int32_t workspace_idx = -1; + SmallVector mutable_inputs = SmallVector(inputs.begin(), inputs.end()); + for(unsigned i = 0; i < mutable_inputs.size(); i++) { + if(mutable_inputs[i] == workspace){ + mutable_inputs[i] = sorted_workspace; + workspace_idx = i; break; } } - /// Calculate for ModeGeneric style format: [CN, S, D (, ... ) ] - auto valueAccessIdx_part = allLoopsArg[i][lastSparseIndexLoc]; - if (lastSparseIndexLoc < allPerms[i].size() - 1) - { /// There is dense index after the sparse index - unsigned int last_d = lastSparseIndexLoc + 1; - for (unsigned int d = lastSparseIndexLoc + 1; d < allPerms[i].size(); d++) - { /// i=0 - if (allFormats[i][d].compare(0, 1, "D") == 0) - { - /// Get dense dim size - auto index_0 = builder.create(loc, 0); - std::vector upper_indices = {index_0}; - auto upperBound = builder.create(loc, main_tensors_all_Allocs[i][4 * d], upper_indices); - comet_vdump(upperBound); - valueAccessIdx_part = builder.create(loc, upperBound, valueAccessIdx_part); - last_d = d; - } - } - if (allFormats[i][last_d].compare(0, 1, "D") == 0) - { - comet_debug() << " "; - comet_vdump(allLoopsArg[i][allLoopsArg[i].size() - 1]); - comet_vdump(valueAccessIdx_part); - valueAccessIdx_part = builder.create(loc, allLoopsArg[i][allLoopsArg[i].size() - 1], - valueAccessIdx_part); - comet_debug() << " AddIOps (valueAccessIdx_part): "; - comet_vdump(valueAccessIdx_part); - } - } + scf::ForOp for_loop = rewriter.create(loc, lb, ub, step, ValueRange(mutable_inputs)); + Block* loop_body = for_loop.getBody(); + Value loop_workspace = for_loop.getRegionIterArg(workspace_idx); - allValueAccessIdx[i].push_back(valueAccessIdx_part); + rewriter.setInsertionPointToStart(loop_body); + IRMapping map; + auto yield_op = rewriter.create(loc, for_loop.getRegionIterArgs()); + rewriter.setInsertionPointAfter(for_loop); + return new WorkspaceLoopInfo(ValueRange(for_loop.getRegionIterArgs()), for_loop.getResults(), yield_op, map, for_loop.getInductionVar(), yield_op, loop_workspace); } - else if (main_tensors_all[i].getType().isa()) - { /// dense tensor - allValueAccessIdx[i] = allAccessIdx[i]; - } - } -#ifdef DEBUG_MODE_LowerIndexTreeToSCFPass - for (unsigned int i = 0; i < allValueAccessIdx.size(); i++) - { - comet_debug() << "allValueAccessIdx[" << i << "].size(): " << allValueAccessIdx[i].size() - << ", main_tensors_all_Allocs[" << i << "].size()-1: " << main_tensors_all_Allocs[i].size() - 1 - << "\n"; - } -#endif - } + Value getCrd(IRRewriter& rewriter) override { + if(crd != nullptr) return crd; + auto loc = workspaceTensor.getLoc(); + crd = rewriter.create(loc, workspaceTensor, inductionVar, rewriter.getI32IntegerAttr(0)); + crd = rewriter.createOrFold(loc, rewriter.getIndexType(), crd); - /// In genCmptOps, get current compute node's symbolic nested for-loop and access indices. - void getSymbolicNestedForOpsAndAccessIdx(std::vector &ancestorsWps, - std::vector &ancestorsOps, - std::vector &nested_forops /* output */, - std::vector &nested_AccessIdx /* output */, - std::vector &nested_forops_indices /* output */) - { + return crd; + } - for (unsigned int i = 0; i < ancestorsOps.size(); i++) - { - comet_debug() << " ancestorsOps[" << i << "]->forOps.size(): " << ancestorsOps[i]->symbolicForOps.size() - << ", ancestorsOps->id: " - << ancestorsOps[i]->id << "\n"; - if (!ancestorsOps[i]->symbolicForOps.empty()) - { /// for loops OpsTree node - for (int j = ancestorsOps[i]->symbolicForOps.size() - 1; j >= 0; j--) + Value getPos(IRRewriter& rewriter, Value tensor, uint32_t dim) override { + if(llvm::isa(tensor.getType())) { - comet_debug() << " j: " << j << "\n"; - nested_forops.push_back(ancestorsOps[i]->symbolicForOps[j]); - comet_debug() << "AccessIdx: " << ancestorsOps[i]->symbolicAccessIdx[j] << "\n"; - nested_AccessIdx.push_back(ancestorsOps[i]->symbolicAccessIdx[j]); + auto loc = rewriter.getUnknownLoc(); + return rewriter.create(loc, rewriter.getIndexType(), tensor, nullptr, (int32_t)dim, true); } + return inductionVar; } - } - comet_debug() << " nested_forops.size(): " << nested_forops.size() << "\n"; - /// std::vector nested_forops_indices; - for (unsigned int i = 0; i < ancestorsWps.size(); i++) - { - comet_debug() << " "; - comet_vdump(ancestorsWps[i]); - if (indexTree::IndexTreeIndicesOp cur_op = dyn_cast( - ancestorsWps[i].getDefiningOp())) - { - /// Get indices - ArrayAttr op_indices = cur_op.getIndices(); - if (op_indices.size() > 0) - { /// for loops OpsTree node - for (int j = op_indices.size() - 1; j >= 0; j--) - { - /// Get the indices; - int64_t idx = op_indices[j].cast().getInt(); - nested_forops_indices.push_back(idx); - } - } + void updateOutput(IRRewriter& rewriter, uint32_t idx, Value newOutput) override { + currentInputs[idx] = newOutput; + rewriter.modifyOpInPlace(terminator, [&](){terminator.setOperand(idx, newOutput);}); } - } - } - - /// In genCmptOps, generate code for a compute node that does general A = 0.0 but without workspace transformation. - void genCmptOpGeneralInitialAssignment(OpBuilder &builder, - Location &loc, - int lhs_loc, - ConstantOp &cstop, - std::vector &nested_forops, - std::vector> &main_tensors_all_Allocs, - std::vector> &allValueAccessIdx) - { - /// Generate Store 1.0, A[...] this op - /// this case: allPerms[0] is empty, allFormats[0] is empty - comet_debug() << " cstop.getValue(): " << cstop.getValue() << "\n"; - comet_debug() << " "; - comet_vdump(main_tensors_all_Allocs[lhs_loc][main_tensors_all_Allocs[lhs_loc].size() - 1]); - - if (allValueAccessIdx[lhs_loc].size() > 0) - { - builder.create(loc, cstop, - main_tensors_all_Allocs[lhs_loc][main_tensors_all_Allocs[lhs_loc].size() - - 1], - allValueAccessIdx[lhs_loc]); - } - else - { - Value local_accessIdx = nested_forops[0].getInductionVar(); - insertInitialize(loc, - cstop, - main_tensors_all_Allocs[lhs_loc][main_tensors_all_Allocs[lhs_loc].size() - 1], - local_accessIdx, - builder, - false /* use_dynamic_init */, - nullptr /* dynamic_init */); - } - } - - /// In genCmptOps, get LHS nnz value and data array before gathering results from the workspace. - void getLHSBeforeGatherFromWorkspace(OpBuilder &builder, - Location &loc, - int lhs_loc, - Value lhs, - std::vector> &main_tensors_all_Allocs, - unsigned int &lhs_2crd_size_loc /* output */, - unsigned int &lhs_2pos_size_loc /* output */, - Value &lhs_nnz /* output */, - Value &lhs_nnz_alloc /* output */, - Value &lhs_val /* output */) - { - /// Get tensor ranks - auto sp_op = cast(lhs.getDefiningOp()); - int lhs_ranks = sp_op.getTensorRank(); - - //[0...2d,2d+1...4d+1,4d+2...5d+1] - unsigned int lhs_val_size_loc = 8 * lhs_ranks + 1; /// 17 (2d) /// 15 - lhs_2crd_size_loc = 7 * lhs_ranks; /// 14 (2d) /// 12 /// output - lhs_2pos_size_loc = 7 * lhs_ranks - 1; /// 13 (2d) /// 11 /// output - - /// [0...2d, 2d+1...4d+1, 4d+2...5d+1] - comet_pdump(lhs.getDefiningOp()); - comet_pdump(lhs.getDefiningOp()->getParentOp()); - comet_vdump(lhs.getDefiningOp()->getOperand(lhs_val_size_loc)); - - Value lhs_nnz_operand = lhs.getDefiningOp()->getOperand(lhs_val_size_loc); - Value lhs_nnz_op; - comet_vdump(lhs_nnz_operand); - if (isa(lhs_nnz_operand.getDefiningOp())) - { - lhs_nnz_op = lhs_nnz_operand.getDefiningOp()->getOperand(0); - } - else - { - lhs_nnz_op = lhs_nnz_operand; - } - comet_vdump(lhs_nnz_op); - auto lhs_nnz_load = cast(lhs_nnz_op.getDefiningOp()); /// index - lhs_nnz_alloc = cast(lhs_nnz_load.getMemRef().getDefiningOp()); /// index /// output - - Value cst_0_index = builder.create(loc, 0); - lhs_nnz = builder.create(loc, lhs_nnz_alloc, ValueRange{cst_0_index}); /// output - - lhs_val = main_tensors_all_Allocs[lhs_loc].back(); /// output - comet_vdump(lhs_val); - } - - /// In genCmptOps, generate code for Cij = Wj when they are both dense. - void genCmptOpGatherFromDenseToDense(OpBuilder &builder, - Location &loc, - int rhs_loc, - int lhs_loc, - std::vector> &main_tensors_all_Allocs, - std::vector> &allValueAccessIdx) - { - /// %1 = load b[...] - /// store %1, a[...] - comet_debug() << " main_tensors_all_Allocs[" << rhs_loc << "].size(): " - << main_tensors_all_Allocs[rhs_loc].size() << ", allValueAccessIdx[" << rhs_loc << "].size(): " - << allValueAccessIdx[rhs_loc].size() << "\n"; - - Value rhs_value = builder.create(loc, main_tensors_all_Allocs[rhs_loc].back(), allValueAccessIdx[rhs_loc]); - comet_vdump(rhs_value); - - comet_vdump(main_tensors_all_Allocs[lhs_loc].back()); -#ifdef DEBUG_MODE_LowerIndexTreeToSCFPass - auto s1 = builder.create(loc, rhs_value, main_tensors_all_Allocs[lhs_loc].back(), allValueAccessIdx[lhs_loc]); - comet_vdump(s1); -#else - builder.create(loc, rhs_value, main_tensors_all_Allocs[lhs_loc].back(), allValueAccessIdx[lhs_loc]); -#endif - } - - /// Used by genCmptOps, for Cij = Wj without Workspace Transformation - void genCmptOpGatherFromDenseToOutput(OpBuilder &builder, - Location &loc, - int rhs_loc, - int lhs_loc, - unsigned int lhs_2crd_size_loc, - unsigned int lhs_2pos_size_loc, - Value lhs, - Value lhs_nnz, - Value lhs_nnz_alloc, - Value lhs_val, - std::vector> &allFormats, - std::vector> &main_tensors_all_Allocs, - std::vector> &allAccessIdx, - std::vector> &allValueAccessIdx, - std::vector &nested_forops) - { - - /// %1 = load b[...] - /// if(%1 != 0) { - /// Cnnz = load Cop.operand(4d+1) - /// store %1, cval[Cnnz] - /// store Cnnz+1, Cop.operand(4d+1) - /// } - comet_debug() << " main_tensors_all_Allocs[" << rhs_loc << "].size(): " - << main_tensors_all_Allocs[rhs_loc].size() << ", allValueAccessIdx[" << rhs_loc - << "].size(): " << allValueAccessIdx[rhs_loc].size() << "\n"; - Value rhs_value = builder.create(loc, main_tensors_all_Allocs[rhs_loc][main_tensors_all_Allocs[rhs_loc].size() - 1], allValueAccessIdx[rhs_loc]); - comet_debug() << " "; - comet_vdump(rhs_value); - auto f64Type = builder.getF64Type(); - Value const_f64_0 = builder.create(loc, f64Type, builder.getF64FloatAttr(0)); - Value isNonzero = builder.create(loc, CmpFPredicate::ONE, rhs_value, const_f64_0); - comet_debug() << " "; - comet_vdump(isNonzero); - auto if_nonzero = builder.create(loc, isNonzero, /*WithElseRegion*/ false); - comet_debug() << " If branch:\n"; - comet_vdump(if_nonzero); - - if (!if_nonzero.getThenRegion().empty()) - { - auto last_insertionPoint = builder.saveInsertionPoint(); - builder.setInsertionPointToStart(&if_nonzero.getThenRegion().front()); - - builder.create(loc, rhs_value, lhs_val, ValueRange{lhs_nnz}); - - /// update pos/crd arrays - /// Fill C2crd in CSR format, parent loop's accessIdx - /// Check format j in the output - if (allFormats[lhs_loc][allFormats[lhs_loc].size() - 1].compare(0, 2, "CU") == 0) - { - Value crd_index = allAccessIdx[allAccessIdx.size() - 1][allAccessIdx[allAccessIdx.size() - 1].size() - - 1]; - comet_debug() << " "; - comet_vdump(crd_index); - Value lhs_2crd = main_tensors_all_Allocs[lhs_loc][main_tensors_all_Allocs[lhs_loc].size() - 4]; //-2 - comet_debug() << " "; - comet_vdump(lhs_2crd); - - builder.create(loc, crd_index, lhs_2crd, ValueRange{lhs_nnz}); - } - - comet_debug() << "\n"; - Value cst_1_index = builder.create(loc, 1); - comet_debug() << " "; - comet_vdump(lhs_nnz); - Value lhs_nnz_new = builder.create(loc, lhs_nnz, cst_1_index); - comet_debug() << " AddIOps: (lhs_nnz_new)"; - comet_vdump(lhs_nnz_new); - comet_debug() << " "; - comet_vdump(lhs_nnz_alloc); - - Value cst_0_index = builder.create(loc, 0); - builder.create(loc, lhs_nnz_new, lhs_nnz_alloc, ValueRange{cst_0_index}); - - comet_debug() << "\n"; - Value lhs_2crd = lhs.getDefiningOp()->getOperand(lhs_2crd_size_loc); - Value lhs_2crd_op; - comet_vdump(lhs_2crd); - if (isa(lhs_2crd.getDefiningOp())) - { - lhs_2crd_op = lhs_2crd.getDefiningOp()->getOperand(0); - } - else - { - lhs_2crd_op = lhs_2crd; - } - comet_debug() << " "; - comet_vdump(lhs_2crd_op); - auto c2crd_size_load = cast(lhs_2crd_op.getDefiningOp()); /// index - Value c2crd_size_alloc = cast(c2crd_size_load.getMemRef().getDefiningOp()); /// index - comet_debug() << " "; - comet_vdump(c2crd_size_alloc); - - builder.create(loc, lhs_nnz_new, c2crd_size_alloc, ValueRange{cst_0_index}); - - comet_debug() << " \n"; - builder.restoreInsertionPoint(last_insertionPoint); - } - - comet_debug() << " \n"; - auto prev_forop = nested_forops[nested_forops.size() - 1 - 1]; - builder.setInsertionPointAfter(prev_forop); - - comet_debug() << " "; - comet_vdump(lhs.getDefiningOp()->getOperand(lhs_2pos_size_loc)); - Value lhs_2pos_0 = lhs.getDefiningOp()->getOperand(lhs_2pos_size_loc); - Value lhs_2pos_op; - comet_vdump(lhs_2pos_0); - if (isa(lhs_2pos_0.getDefiningOp())) - { - lhs_2pos_op = lhs_2pos_0.getDefiningOp()->getOperand(0); - } - else - { - lhs_2pos_op = lhs_2pos_0; - } - comet_debug() << " "; - comet_vdump(lhs_2pos_op); - auto c2pos_size_load = cast(lhs_2pos_op.getDefiningOp()); /// index - Value c2pos_size_alloc = cast(c2pos_size_load.getMemRef().getDefiningOp()); /// index - Value cst_0_index = builder.create(loc, 0); - Value c2pos_size_value = builder.create(loc, c2pos_size_alloc, ValueRange{cst_0_index}); - - Value lhs_2crd = lhs.getDefiningOp()->getOperand(lhs_2crd_size_loc); - Value lhs_2crd_op; - comet_vdump(lhs_2crd); - if (isa(lhs_2crd.getDefiningOp())) - { - lhs_2crd_op = lhs_2crd.getDefiningOp()->getOperand(0); - } - else - { - lhs_2crd_op = lhs_2crd; - } - comet_debug() << " "; - comet_vdump(lhs_2crd_op); - auto c2crd_size_load = cast(lhs_2crd_op.getDefiningOp()); /// index - Value c2crd_size_alloc = cast(c2crd_size_load.getMemRef().getDefiningOp()); /// index - Value c2crd_size_nnz = builder.create(loc, c2crd_size_alloc, ValueRange{cst_0_index}); - - /// store crd_size into pos - Value lhs_2pos = main_tensors_all_Allocs[lhs_loc][main_tensors_all_Allocs[lhs_loc].size() - 5]; /// -3 - comet_debug() << " "; - comet_vdump(lhs_2pos); - - builder.create(loc, c2crd_size_nnz, lhs_2pos, ValueRange{c2pos_size_value}); - - Value cst_1_index = builder.create(loc, 1); - comet_debug() << " "; - comet_vdump(c2pos_size_value); - Value c2pos_size_value_new = builder.create(loc, c2pos_size_value, cst_1_index); - comet_debug() << " AddIOps (c2pos_size_value_new): "; - comet_vdump(c2pos_size_value_new); - - builder.create(loc, c2pos_size_value_new, c2pos_size_alloc, ValueRange{cst_0_index}); - } + }; - /// From the W_id_list_size, get the output C and C.rowptr, C.col, and C.val. - /// ----------------- /// - /// %55 = "it.ComputeLHS"(%53) {allFormats = [[]], allPerms = [[]]} : (tensor<1xindex>) -> tensor<*xf64> - /// %56 = "it.Compute"(%54, %55) {MaskType = "none", comp_worksp_opt = true, semiring = "noop_times"} : (tensor<*xindex>, tensor<*xf64>) -> i64 - /// %70 = "it.ComputeRHS"(%50, %51, %52, %53) {allFormats = [["D"]], allPerms = [[1]]} : (tensor, tensor, tensor, tensor<1xindex>) -> tensor<*xf64> - /// %93 = "it.Compute"(%70, %92) {MaskType = "none", comp_worksp_opt = true, semiring = "noop_times"} : (tensor<*xf64>, tensor<*xf64>) -> i64 - /// %92 = "it.ComputeLHS"(%91) {allFormats = [["D", "CU"]], allPerms = [[0, 1]]} : (!ta.sptensor, tensor, tensor, tensor, tensor, tensor, tensor, tensor, tensor, index, index, index, index, index, index, index, index, index, index, index>) -> tensor<*xf64> - /// %91 = ta.sptensor_construct(%73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %11, %12) {tensor_rank = 2 : i32} : (tensor, tensor, tensor, tensor, tensor, tensor, tensor, tensor, tensor, index, index, index, index, index, index, index, index, index, index, index) -> (!ta.sptensor, tensor, tensor, tensor, tensor, tensor, tensor, tensor, tensor, index, index, index, index, index, index, index, index, index, index, index>) - /// %77 = bufferization.to_tensor %alloc_156 : memref - /// %alloc_156 = memref.alloc(%71) : memref - void getOutputMtxCRowptrAndDims(indexTree::IndexTreeComputeOp &cur_op, - Value &W_id_list_size, - SymbolicInfo &symbolicInfo /* output */) - { - Value mtxC = nullptr; - for (Operation *u_rhs : W_id_list_size.getUsers()) - { - if (indexTree::IndexTreeComputeRHSOp rhs_op = dyn_cast(u_rhs)) + class IntersectionLoopInfo : LoopInfo { + private: + Value inductionVar; + Value crd; + scf::YieldOp terminator; + llvm::SmallDenseMap, std::pair> controlVars; + public: + IntersectionLoopInfo(ValueRange inputs, ResultRange outputs, Operation* body, IRMapping ir_map, Value i, Value c, scf::YieldOp yield, SmallDenseMap, std::pair> controls) : + LoopInfo(inputs, outputs, body, ir_map), inductionVar(i), crd(c), terminator(yield), controlVars(controls){} + + static LoopInfo* build(Operation* domain_op, IRRewriter& rewriter, ValueRange inputs) { - /// rhs_op is %70 - for (Operation *u_cmpt : u_rhs->getUsers()) - { - if (indexTree::IndexTreeComputeOp cmpt_op = dyn_cast(u_cmpt)) - { - /// cmpt_op is %93 - /// then %93's Operand[1] is %92 - /// %92's Operand[0] is %91 which is the sparse tensor - Value lhs_op = cmpt_op.getOperand(1); /// lhs_op is %92 - mtxC = lhs_op.getDefiningOp()->getOperand(0); /// mtxC is %91 - break; - } + auto loc = domain_op->getLoc(); + auto intersection_domain_op = llvm::cast(domain_op); + auto index_type = rewriter.getIndexType(); + auto context = rewriter.getContext(); + + SmallVector loop_args = SmallVector(inputs); + + SmallDenseMap, std::pair> tensor_access_map; + + // Intersection between sparse domains + auto domains = intersection_domain_op.getDomains(); + Value inc = rewriter.create(loc, index_type, rewriter.getIndexAttr(1)); + SmallVector loop_conditions; + SmallVector array_crds; + + Block* cond_block = new Block(); + Block* body_block = new Block(); + + // Create loop carried arguments for output tensors and iteration counter + IRMapping map; + for(Value init_arg : loop_args){ + cond_block->addArgument(init_arg.getType(), loc); + BlockArgument body_arg = body_block->addArgument(init_arg.getType(), loc); + map.map(init_arg, body_arg); } - } - } - - assert(mtxC && "Error: cannot find mtxC as the output."); - /// %77 is mtxC.getDefiningOp()->getOperand(A2POS) - /// %alloc_156 is C_rowptr - /// %71 is mtxC_rowptr_size - Value C_rowptr = mtxC.getDefiningOp()->getOperand(CSR_A2POS).getDefiningOp()->getOperand(0); /// A2POS is rowptr's location - Value C_rowptr_size = C_rowptr.getDefiningOp()->getOperand(0); - Value C_num_rows = mtxC.getDefiningOp()->getOperand(CSR_DIM1_SIZE); - Value C_num_cols = mtxC.getDefiningOp()->getOperand(CSR_DIM2_SIZE); - symbolicInfo.mtxC = mtxC; - symbolicInfo.mtxC_rowptr = C_rowptr; - symbolicInfo.mtxC_rowptr_size = C_rowptr_size; - symbolicInfo.mtxC_num_rows = C_num_rows; - symbolicInfo.mtxC_num_cols = C_num_cols; - { - comet_vdump(mtxC); - comet_vdump(C_rowptr); - comet_vdump(C_rowptr_size); - comet_vdump(C_num_rows); - comet_vdump(C_num_cols); - } - } - - /// Generate mark before the outer-most symbolic for-loop, - /// and update mark for every idx at the beginning of the outer-most symbolic for-loop. - void genSymbolicMarkAndUpdate(OpBuilder &builder, - Location &loc, - /// std::vector &symbolic_nested_forops, /* from innermost to outermost */ - AbstractLoopOp &outermost_forLoop, /// the outermost for-loop - Value &mark_alloc /* output */, - Value &mark_new_val /* output */) - { - /// Store the insertion point - auto last_insertion_point = builder.saveInsertionPoint(); - - /// Set the insertion point before the outer-most symbolic for-loop - builder.setInsertionPoint(outermost_forLoop); - - /// Generate the variable mark - /// %mark = memref.alloc() : memref<1xindex> - /// memref.store %c0, %mark[%c0] : memref<1xindex> - MemRefType memTy_1xindex = MemRefType::get({1}, builder.getIndexType()); - mark_alloc = builder.create(loc, memTy_1xindex); - Value const_index_0 = builder.create(loc, 0); - builder.create(loc, - const_index_0, - mark_alloc, - ValueRange{const_index_0}); - { - comet_vdump(mark_alloc); - } - - /// Generate updating mark += 2 - /// %c2 = arith.constant 2 : index - /// %old_val = memref.load %mark[%c0] : memref<1xindex> - /// %new_mark = arith.addi %old_val, %c2 : index - /// memref.store %new_mark, %mark[%c0] : memref<1xindex> - builder.setInsertionPointToStart(outermost_forLoop.getBody()); - Value const_index_2 = builder.create(loc, 2); - Value old_mark_val = builder.create(loc, mark_alloc, ValueRange{const_index_0}); - mark_new_val = builder.create(loc, old_mark_val, const_index_2); - builder.create(loc, - mark_new_val, - mark_alloc, - ValueRange{const_index_0}); - { - comet_vdump(outermost_forLoop); - } - - /// Restore the insertion point - builder.restoreInsertionPoint(last_insertion_point); - } - - /// Generate symbolic if statement condition in the CmptOp - /// -------No masking---------- /// - /// if (mark_array[j_idx] != mark) { - /// mark_array[j_idx] = mark; /// C[i_idx, j_idx] has been visited - /// W_id_list_size += 1; - /// } - /// -------Push masking---------- /// - /// if (mark_array[j_idx] == mark) { - /// mark_array[j_idx] = mark + 1; /// C[i_idx, j_idx] has been visited - /// W_id_list_size += 1; - /// } - void genSymbolicIfStatementCondition(OpBuilder &builder, - Location &loc, - AbstractLoopOp &semiringLoop, /// symbolic_nested_forops[0] - Value &mark_array_alloc, /// tensors_lhs_Allocs[1][0] - Value &valueAccessIdx, /// allValueAccessIdx[lhs_loc][0] - Value &mark_new_val, - scf::IfOp &if_statement /* output */, - MaskingInfo &maskingInfo) - { - /// Store the insertion point - auto last_insertion_point = builder.saveInsertionPoint(); - - /// Set the insertion point at the end of the inner-most symbolic for-loop - builder.setInsertionPoint(semiringLoop.getBody()->getTerminator()); - - { - comet_vdump(semiringLoop); - } - /// Generate If statement condition - Value ele_mark_val = builder.create(loc, mark_array_alloc, ValueRange{valueAccessIdx}); - - if (PUSH_BASED_MASKING == maskingInfo.mask_type) - { - Value equal_mask = builder.create(loc, - CmpIPredicate::eq, - ele_mark_val, - mark_new_val); - if_statement = builder.create(loc, equal_mask, false /* No Else Region */); - } - else if (NO_MASKING == maskingInfo.mask_type) - { - - Value not_equal_mark = builder.create(loc, - CmpIPredicate::ne, - ele_mark_val, - mark_new_val); - if_statement = builder.create(loc, not_equal_mark, false /* No Else Region */); - } - else - { - llvm::errs() << "Error: mask_type " << maskingInfo.mask_type << " is not supported.\n"; - } - { - comet_vdump(ele_mark_val); - comet_vdump(if_statement); - comet_vdump(semiringLoop); - } - /// Restore the insertion point - builder.restoreInsertionPoint(last_insertion_point); - } - - /// Generate symbolic if statement then region in the CmptOp - /// -------No masking---------- /// - /// if (mark_array[j_idx] != mark) { - /// mark_array[j_idx] = mark; /// C[i_idx, j_idx] has been visited - /// W_id_list_size += 1; - /// } - /// -------Push masking---------- /// - /// if (mark_array[j_idx] == mark) { - /// mark_array[j_idx] = mark + 1; /// C[i_idx, j_idx] has been visited - /// W_id_list_size += 1; - /// } - void genSymbolicIfStatementThenRegion(OpBuilder &builder, - Location &loc, - scf::IfOp &if_statement, - Value &mark_array_alloc, /// tensors_lhs_Allocs[1][0] - Value &valueAccessIdx, /// allValueAccessIdx[lhs_loc][0] - Value &W_id_list_size, - Value &mark_new_val, - MaskingInfo &maskingInfo) - { - /// Store the insertion point - auto last_insertion_point = builder.saveInsertionPoint(); - - /// Set the insertion point to the beginning of the if statement then region - builder.setInsertionPointToStart(&if_statement.getThenRegion().front()); - - if (PUSH_BASED_MASKING == maskingInfo.mask_type) - { - /// mark_array[j_idx] = mark + 1; - Value const_index_1 = builder.create(loc, 1); - Value mark_value_plus_one = builder.create(loc, mark_new_val, const_index_1); - builder.create(loc, - mark_value_plus_one, - mark_array_alloc, - ValueRange{valueAccessIdx}); - } - else if (NO_MASKING == maskingInfo.mask_type) - { - /// mark_array[j_idx] = mark - builder.create(loc, - mark_new_val, - mark_array_alloc, - ValueRange{valueAccessIdx}); - } - /// W_id_list_size += 1; - - Value const_index_0 = builder.create(loc, 0); - Value const_index_1 = builder.create(loc, 1); - Value old_val = builder.create(loc, W_id_list_size, ValueRange{const_index_0}); - Value new_val = builder.create(loc, old_val, const_index_1); - builder.create(loc, - new_val, - W_id_list_size, - ValueRange{const_index_0}); - - { - comet_vdump(if_statement); - } - /// Restore the insertion point - builder.restoreInsertionPoint(last_insertion_point); - } - - /// Updating output - /// C.rowptr[idx] = W_id_list_size; - void genSymbolicUpdateCRowptr(OpBuilder &builder, - Location &loc, - AbstractLoopOp &outermost_forLoop, - Value &mtxC_rowptr, - Value &valueAccessIdx, - Value &W_id_list_size) - { - { - comet_vdump(mtxC_rowptr); - comet_vdump(valueAccessIdx); - comet_vdump(W_id_list_size); - } - /// Store the insertion point - auto last_insertion_point = builder.saveInsertionPoint(); - - /// Set the insertion point at the end of the outermost for-loop body - builder.setInsertionPoint(outermost_forLoop.getBody()->getTerminator()); - - Value const_index_0 = builder.create(loc, 0); - Value rowptr_val = builder.create(loc, W_id_list_size, ValueRange{const_index_0}); - builder.create(loc, - rowptr_val, - mtxC_rowptr, - ValueRange{valueAccessIdx}); - - { - comet_vdump(outermost_forLoop); - } - /// Restore the insertion point - builder.restoreInsertionPoint(last_insertion_point); - } - - /// Generate the reduce of the output C.rowptr after the outermost for-loop - /// C.rowptr[M] = 0; - /// int C_val_size = 0; - /// for (int i_idx = 0; i_idx < M + 1; ++i_idx) { - /// int curr = C.rowptr[i_idx]; - /// C.rowptr[i_idx] = C_val_size; - /// C_val_size += curr; - /// } - /// C.col = new int[C_val_size] - /// C.val = new f64[C_val_size] - void genSymbolicReduceOutputCRowptrCColCVal(OpBuilder &builder, - Location &loc, - AbstractLoopOp &outermost_forLoop, - SymbolicInfo &symbolicInfo /* output */) - { - Value const_index_0 = builder.create(loc, 0); - Value const_index_1 = builder.create(loc, 1); - - /// C.rowptr[M] = 0 - Value &mtxC_rowptr = symbolicInfo.mtxC_rowptr; - Value &num_rows = symbolicInfo.mtxC_num_rows; - builder.create(loc, - const_index_0, - mtxC_rowptr, - ValueRange{num_rows}); - - /// C_val_size = 0; - MemRefType memTy_1xindex = MemRefType::get({1}, builder.getIndexType()); - Value C_val_size = builder.create(loc, memTy_1xindex); - builder.create(loc, - const_index_0, - C_val_size, - ValueRange{const_index_0}); - - /// for (int i_idx = 0; i_idx < M + 1; ++i_idx) { - /// int curr = C.rowptr[i_idx]; - /// C.rowptr[i_idx] = C_val_size; - /// C_val_size += curr; - /// } - Value &num_rows_plus_one = symbolicInfo.mtxC_rowptr_size; - scf::ForOp reduce_forLoop = builder.create(loc, - const_index_0 /* lowerBound */, - num_rows_plus_one /* upperBound */, - const_index_1 /* step */); - builder.setInsertionPointToStart(reduce_forLoop.getBody()); - Value i_idx = reduce_forLoop.getInductionVar(); - Value curr = builder.create(loc, mtxC_rowptr, ValueRange{i_idx}); - Value size_val = builder.create(loc, C_val_size, ValueRange{const_index_0}); - builder.create(loc, - size_val, - mtxC_rowptr, - ValueRange{i_idx}); - Value new_val = builder.create(loc, curr, size_val); - builder.create(loc, - new_val, - C_val_size, - ValueRange{const_index_0}); - { - comet_vdump(reduce_forLoop); - } - builder.setInsertionPointAfter(reduce_forLoop); - Value mtxC_val_size = builder.create(loc, C_val_size, ValueRange{const_index_0}); - symbolicInfo.mtxC_val_size = mtxC_val_size; - - /// Allocate new C.col and new C.val - MemRefType memTy_alloc_dynamic_index = MemRefType::get({ShapedType::kDynamic}, builder.getIndexType()); - MemRefType memTy_alloc_dynamic_f64 = MemRefType::get({ShapedType::kDynamic}, builder.getF64Type()); - Value new_mtxC_col = builder.create(loc, - memTy_alloc_dynamic_index, - ValueRange{mtxC_val_size}); - Value new_mtxC_val = builder.create(loc, - memTy_alloc_dynamic_f64, - ValueRange{mtxC_val_size}); - symbolicInfo.mtxC_col = new_mtxC_col; - symbolicInfo.mtxC_val = new_mtxC_val; - { - comet_vdump(mtxC_val_size); - comet_vdump(new_mtxC_col); - comet_vdump(new_mtxC_val); - } - } - - /// ----------------- /// - /// Store new mtxC_val_size to the old mtxC's C_col_size (A2crd_size) and C_val_size (Aval_size). - /// Just in case for safety. - /// ----------------- /// - void storeNewMtxCValeSizeToOldMtxC(OpBuilder &builder, - Location &loc, - SymbolicInfo &symbolicInfo) - { - Value &mtxC = symbolicInfo.mtxC; - Value &mtxC_val_size = symbolicInfo.mtxC_val_size; - Value const_index_0 = builder.create(loc, 0); - - { - comet_vdump(mtxC); - comet_vdump(mtxC_val_size); - } - - /// Find the alloc of C_col_size (Arcrd_size) - /// %66 = memref.load %alloc_153[%c0_128] : memref<1xindex> - Value C_col_size_alloc = mtxC.getDefiningOp()->getOperand(CSR_A2CRD_SIZE).getDefiningOp()->getOperand(0); /// 8 - /// Store the new mtxC_val_size to C_col_size -#ifdef DEBUG_MODE_LowerIndexTreeToSCFPass - auto store_C_col_size_alloc = builder.create(loc, - mtxC_val_size, - C_col_size_alloc, - ValueRange{const_index_0}); - comet_vdump(C_col_size_alloc); - comet_vdump(store_C_col_size_alloc); -#else - builder.create(loc, - mtxC_val_size, - C_col_size_alloc, - ValueRange{const_index_0}); -#endif - - /// Find the alloc of C_val_size (Aval_size) - /// %67 = memref.load %alloc_154[%c0_128] : memref<1xindex> - Value C_val_size_alloc = mtxC.getDefiningOp()->getOperand(CSR_AVAL_SIZE).getDefiningOp()->getOperand(0); /// 9 - /// Store the new mtxC_val_size to C_val_size -#ifdef DEBUG_MODE_LowerIndexTreeToSCFPass - auto store_C_val_size_alloc = builder.create(loc, - mtxC_val_size, - C_val_size_alloc, - ValueRange{const_index_0}); - comet_vdump(C_val_size_alloc); - comet_vdump(store_C_val_size_alloc); -#else - builder.create(loc, - mtxC_val_size, - C_val_size_alloc, - ValueRange{const_index_0}); -#endif - } - - /// Dealloc the old C.val and C.col before the outermost_forLoop. - /// Replace the old C.val and C.col with new ones. - void deallocMtxCColCVal(OpBuilder &builder, - Location &loc, - AbstractLoopOp &outermost_forLoop, - SymbolicInfo &symbolicInfo) - { - /// Find old C.col and C.val - Value &mtxC = symbolicInfo.mtxC; - Value old_C_col = mtxC.getDefiningOp()->getOperand(CSR_A2CRD).getDefiningOp()->getOperand(0); - Value old_C_val = mtxC.getDefiningOp()->getOperand(CSR_AVAL).getDefiningOp()->getOperand(0); - - /// Dealloc old C.col and C.val - /// Store the insertion point - auto last_insertion_point = builder.saveInsertionPoint(); - - /// Set the insertion point before the symbolic outermost_forloop - builder.setInsertionPoint(outermost_forLoop); - - builder.create(loc, old_C_col); - builder.create(loc, old_C_val); - - /// Restore the insertion point - builder.restoreInsertionPoint(last_insertion_point); - - /// -------------- /// - /// Remove mtxC_col's user who is a memref.store operation - /// This is very ad-hoc, just to avoid segmentation fault for old very large C.val array and C.col array. - /// -------------- /// - removeMemrefStoreUser(old_C_col); - removeMemrefStoreUser(old_C_col); - - /// Replace old C.col and C.val - /// Just in case of safety. - replaceOldValueToNewValue(old_C_col, symbolicInfo.mtxC_col); - replaceOldValueToNewValue(old_C_val, symbolicInfo.mtxC_val); - } - - /// Generate a new sparse tensor to replace the old output sparse tensor after the numeric outermost for-loop. - /// (e.g., ta.print(old_tensor) -> ta.print(new_tensor) - void genReplaceOutputSparseTensorToNewSparseTensor(OpBuilder &builder, - Location &loc, - AbstractLoopOp &numeric_outermost_forLoop, - SymbolicInfo &symbolicInfo) - { - /// Set the insertion point after the outermost_forloop - builder.setInsertionPointAfter(numeric_outermost_forLoop); - - Value &mtxC = symbolicInfo.mtxC; - Value &mtxC_col = symbolicInfo.mtxC_col; - Value &mtxC_val = symbolicInfo.mtxC_val; - - /// Generate the new mtxC_col and new mtxC_val bufferization.to_tensor - Value mtxC_col_buffer = builder.create(loc, mtxC_col); - Value mtxC_val_buffer = builder.create(loc, mtxC_val); - - auto sp_op = cast(mtxC.getDefiningOp()); - int tensorRanks = sp_op.getTensorRank(); - - /// Get the operands and their types for the sparse tensor ta.sptensor_construct() (which is mtxC). - SmallVector operands; - operands.insert(operands.end(), - mtxC.getDefiningOp()->getOperands().begin(), - mtxC.getDefiningOp()->getOperands().end()); - operands[CSR_A2CRD] = mtxC_col_buffer; /// 3 (A2crd) - operands[CSR_AVAL] = mtxC_val_buffer; /// 4 (AVal) - SmallVector elementTypes; - for (Value &opd : operands) - { - elementTypes.push_back(opd.getType()); - } - auto ty = tensorAlgebra::SparseTensorType::get(elementTypes); - Value sptensor = builder.create(loc, - ty, - operands, - tensorRanks); - { - comet_vdump(mtxC_col_buffer); - comet_vdump(mtxC_val_buffer); - comet_vdump(sptensor); - } - - /// ----------------- /// - /// Find all users of the old sparse tensor mtxC, and replace those users' corresponding operands - /// to the new sparse tensor (sptensor). For example, - /// "ta.print"(%mtxC) => "ta.print"(%sptensor) - /// ----------------- /// - replaceOldValueToNewValue(mtxC, sptensor); - } - - /// Logistics of memory about old mtxC, mtxC.col, and mtxC.val - /// 1. Dealloc the old C.val and C.col before the outermost_forLoop. - /// 2. Change mtxC's old value in C_col_size (A2crd_size) and C_val_size (Aval_size) to new mtxC_val_size. - /// 3. Generate a new sparse tensor to replace the old output sparse tensor after the numeric outermost for-loop. - void logisticsForMtxCColCVal(OpBuilder &builder, - Location &loc, - AbstractLoopOp &symbolic_outermost_forLoop, - SymbolicInfo &symbolicInfo, - AbstractLoopOp &numeric_outermost_forLoop) - { - - /// Dealloc old C.col and C.val - /// Replace the old C.val and C.col with new ones. - deallocMtxCColCVal(builder, - loc, - symbolic_outermost_forLoop, - symbolicInfo); - - /// Change mtxC's old value in C_col_size (A2crd_size) and C_val_size (Aval_size) to new mtxC_val_size. - /// Just in case for safety. - storeNewMtxCValeSizeToOldMtxC(builder, - loc, - symbolicInfo); - - /// Generate a new sparse tensor to replace the old output sparse tensor after the numeric outermost for-loop. - /// (e.g., ta.print(old_tensor) -> ta.print(new_tensor) - genReplaceOutputSparseTensorToNewSparseTensor(builder, - loc, - numeric_outermost_forLoop, - symbolicInfo); - - /// builder.restoreInsertionPoint(last_insertion_point); - } - - /// Initialize the mark-array according to the mask at the beginning of the symbolic outermost for-loop - /// ----------------- /// - /// %j_loc_start = memref.load %mask_rowptr[%i_idx] : memref - /// %j_loc_bound = memref.load %mask_rowptr[%i_idx_plus_one] : memref - /// scf.for %j_loc = %j_loc_start to %j_loc_bound step %c1 { - /// %val = memref.load %mask_val[%j_loc] : memref - /// %70 = arith.cmpf une, %val, %cst : f64 - /// scf.if %70 { - /// %j_idx = memref.load %mask_col[%arg1] : memref - /// memref.store %mark, %mark_array[%j_idx] : memref - /// } - /// } - void genSymbolicInitMarkArrayByMask(OpBuilder &builder, - Location &loc, - AbstractLoopOp &symbolic_outermost_forLoop, - Value &outermost_forLoop_valueAccessIdx, - Value &mark_array_alloc, - Value &mark_new_val, - MaskingInfo &maskingInfo) - { - /// Store the insertion point - auto last_insertion_point = builder.saveInsertionPoint(); - - /// Set the Insertion Point at the beginning of the symbolic outermost for-loop but AFTER the mark_new_val - builder.setInsertionPointAfter(mark_new_val.getDefiningOp()); - - /// Generate the for-loop entry - Value &mask_rowptr = maskingInfo.mask_rowptr; - Value &mask_col = maskingInfo.mask_col; - Value &mask_val = maskingInfo.mask_val; - Value const_index_1 = builder.create(loc, 1); - Value &i_idx = outermost_forLoop_valueAccessIdx; - Value i_idx_plus_one = builder.create(loc, i_idx, const_index_1); - Value j_loc_start = builder.create(loc, mask_rowptr, ValueRange{i_idx}); - Value j_loc_bound = builder.create(loc, mask_rowptr, ValueRange{i_idx_plus_one}); - auto for_loop = builder.create(loc, - j_loc_start /* lower_bound */, - j_loc_bound /* upper_bound*/, - const_index_1 /* step */); - { - comet_vdump(j_loc_start); - comet_vdump(j_loc_bound); - comet_vdump(for_loop); - } - - /// Generate the for-loop body - builder.setInsertionPointToStart(for_loop.getBody()); - Value const_f64_0 = builder.create(loc, builder.getF64Type(), builder.getF64FloatAttr(0)); - Value j_loc = for_loop.getInductionVar(); - Value val = builder.create(loc, mask_val, ValueRange{j_loc}); - Value not_zero = builder.create(loc, CmpFPredicate::UNE, val, const_f64_0); - auto if_not_zero = builder.create(loc, not_zero, false /*NoElseRegion*/); - builder.setInsertionPointToStart(&if_not_zero.getThenRegion().front()); - Value j_idx = builder.create(loc, mask_col, ValueRange{j_loc}); - builder.create(loc, - mark_new_val, - mark_array_alloc, - ValueRange{j_idx}); - - { - comet_vdump(val); - comet_vdump(if_not_zero); - comet_vdump(for_loop); - comet_vdump(symbolic_outermost_forLoop); - } - - /// Restore the insertion point - builder.restoreInsertionPoint(last_insertion_point); - } - - /// Generate the symbolic phase's kernel to compute the rowptr[i_idx] - void genSymbolicSemiringLoopBody(OpBuilder &builder, - Location &loc, - int lhs_loc, - std::vector> &tensors_lhs_Allocs, - std::vector &symbolic_nested_forops, - std::vector &symbolic_nested_AccessIdx, - std::vector> &symbolic_allValueAccessIdx, - SymbolicInfo &symbolicInfo, - std::vector &numeric_nested_forops, - MaskingInfo &maskingInfo) - { - - AbstractLoopOp &outermost_forLoop = symbolic_nested_forops.back(); - Value &outermost_forLoop_valueAccessIdx = symbolic_nested_AccessIdx.back(); - AbstractLoopOp &semiringLoop = symbolic_nested_forops[0]; - Value &mark_array = tensors_lhs_Allocs[1][0]; - Value &W_id_list_size = tensors_lhs_Allocs[3][0]; - Value &semiringLoop_valueAccessIdx = symbolic_allValueAccessIdx[lhs_loc][0]; - - /// Generate mark before symbolic outer-most for-loop - Value mark_alloc; - Value mark_new_val; - genSymbolicMarkAndUpdate(builder, - loc, - outermost_forLoop, /// the outermost for-loop - mark_alloc /* output */, - mark_new_val /* output */); - - if (PUSH_BASED_MASKING == maskingInfo.mask_type) - { - assert(symbolic_nested_forops.size() >= 2 && symbolic_allValueAccessIdx.size() >= 2 && - "Error: The symbolic for-loops should be at least 2 level.\n"); - - /// Initialize the mark-array according to the mask at the beginning of the symbolic outermost for-loop - genSymbolicInitMarkArrayByMask(builder, - loc, - outermost_forLoop, - outermost_forLoop_valueAccessIdx, - mark_array, - mark_new_val, - maskingInfo); - } - - /// Generate if statement condition - /// if (mark_array[j_idx] != mark) { - /// mark_array[j_idx] = mark; /// C[i_idx, j_idx] has been visited - /// W_id_list_size += 1; - /// } - scf::IfOp if_statement; - genSymbolicIfStatementCondition(builder, - loc, - semiringLoop, /// the inner-most for-loop (SemiringLoop) - mark_array, /// mark-array - semiringLoop_valueAccessIdx, /// value access index j_idx - mark_new_val, - if_statement /* output */, - maskingInfo); - - /// Generate if statement then region - /// if (mark_array[j_idx] != mark) { - /// mark_array[j_idx] = mark; /// C[i_idx, j_idx] has been visited - /// W_id_list_size += 1; - /// } - genSymbolicIfStatementThenRegion(builder, - loc, - if_statement, - mark_array, /// mark-array - semiringLoop_valueAccessIdx, /// value access index j_idx - W_id_list_size, /// W_id_list_size - mark_new_val, - maskingInfo); - - /// Updating output - /// C.rowptr[idx] = W_id_list_size; - Value i_idx = outermost_forLoop.getInductionVar(); - genSymbolicUpdateCRowptr(builder, - loc, - outermost_forLoop, - symbolicInfo.mtxC_rowptr, /// mtxC_rowptr - i_idx, /// value access index i_idx - W_id_list_size /* W_id_list_size */); - - /// Store the insertion point - auto last_insertion_point = builder.saveInsertionPoint(); - /// Set the insertion point after the outermost_forloop - builder.setInsertionPointAfter(outermost_forLoop); - - /// Generate the reduce of output C.rowptr and new C.col and new C.val - /// C.rowptr[M] = 0; - /// int C_val_size = 0; - /// for (int i_idx = 0; i_idx < M + 1; ++i_idx) { - /// int curr = C.rowptr[i_idx]; - /// C.rowptr[i_idx] = C_val_size; - /// C_val_size += curr; - /// } - genSymbolicReduceOutputCRowptrCColCVal(builder, - loc, - outermost_forLoop, - symbolicInfo /* output */); - - /// Logistics of memory about old mtxC, mtxC.col, and mtxC.val - /// 1. Dealloc the old C.val and C.col before the outermost_forLoop. - /// 2. Change mtxC's old value in C_col_size (A2crd_size) and C_val_size (Aval_size) to new mtxC_val_size. - /// 3. Generate a new sparse tensor to replace the old output sparse tensor after the numeric outermost for-loop. - AbstractLoopOp &numeric_outermost_forLoop = numeric_nested_forops.back(); - logisticsForMtxCColCVal(builder, - loc, - outermost_forLoop, /// symbolic_outermost_forLoop - symbolicInfo, - numeric_outermost_forLoop); - - /// Restore the insertion point - builder.restoreInsertionPoint(last_insertion_point); - } - - /// 1. Get the nested loops - /// ---1.1 the nested loops corresponding indices can be infered from ancestors_wp - /// 2. get lhs and rhs. if only 1 rhs, then it's a fill op; otherwise, binary op - /// Note: 1. The auxiliary arrays does not contain the perms/formats information - /// 2. We only apply the compressed workspace on the output of the tensor, then in this case, the workspace tensors will not be in the same side with the main tensors. - /// (main tensors: such as A, B, C, w; auxiliary tensors: such as w_index_list ...) - void genCmptOps(indexTree::IndexTreeComputeOp &cur_op, - indexTree::IndexTreeOp &rootOp, - /// PatternRewriter &rewriter, - OpBuilder &builder, - OpsTree *opstree, - std::vector &ancestorsWps, - std::vector &wp_ops, - SymbolicInfo &symbolicInfo, - NumericInfo &numericInfo) - { - comet_debug() << " calling genCmptOps\n"; - Location loc = rootOp.getLoc(); - comet_debug() << " \n"; - - comet_debug() << " Current IndexTreeComputeOp:"; - comet_vdump(cur_op); - - const bool comp_worksp_opt(cur_op.getCompWorkspOpt()); - comet_debug() << " comp_worksp_opt (bool: true is compressed): " << comp_worksp_opt << "\n"; - - /// Two cases: - /// 1. for the initial workspace, only 1 auxiliary vector w - /// 2. for the compressed workspace, there are 4 auxiliaty vectors, w, w_already_set, w_index_list, w_index_list_size - - /// The insertion location should be "the end of the body of parent loop" - std::vector ancestorsOps; - getAncestorsOps(opstree, ancestorsOps); - comet_debug() << " ancestorsOps.size(): " << ancestorsOps.size() << "\n"; - for (unsigned int i = 0; i < ancestorsOps.size(); i++) - { - comet_debug() << " ancestorsOps[i]->id:" << ancestorsOps[i]->id << "\n"; - } - - /// 1. get the nested loops, from innermost to outermost order - std::vector nested_forops; - std::vector nested_AccessIdx; - std::vector nested_forops_indices; /// Each nested indexOp's index value (e.g., indices=[0]) - getNumericNestedForOpsAndAccessIdx(ancestorsWps, - ancestorsOps, - nested_forops /* output */, - nested_AccessIdx /* output */, - nested_forops_indices /* output */); - - comet_debug() << " nested_forops_indices.size(): " << nested_forops_indices.size() << "\n"; - assert( - nested_forops.size() == nested_forops_indices.size() && "nested_forops.size() != nested_forops_indices.size()"); - - /// Reset the insertion point: the body of the innermost loop - assert(nested_forops.size() > 0 && "No loops\n"); - comet_debug() << " "; - comet_pdump(nested_forops[0].getBody()); - comet_debug() << " "; - comet_pdump(nested_forops[0].getBody()->getTerminator()); - builder.setInsertionPoint(nested_forops[0].getBody()->getTerminator()); - { - comet_vdump(nested_forops[0]); - } - - /// Analyze the leafop, Get the tensors, rhs, lhs, and operator_type - /// --- only one rhs, it will be a fill op; if two, check op_type (+, +=, *=) - /// Check the indices contained in each tensor - /// Generate loadOp, compute ops, StoreOp. - std::vector tensors_rhs; - std::vector> tensors_lhs_Allocs; - std::vector> tensors_rhs_Allocs; - std::vector> allFormats; - std::vector> allPerms; - std::vector> allPerms_rhs; - std::vector main_tensors_all; /// main_tensors_all has first RHS tensors then LHS tensors - std::vector main_tensors_rhs; - getNumericTensors(cur_op, - tensors_rhs /* output */, - tensors_lhs_Allocs /* output */, - tensors_rhs_Allocs /* output */, - allFormats /* output */, - allPerms /* output */, - allPerms_rhs /* output */, - main_tensors_all /* output */, - main_tensors_rhs /* output */); - - /// ----------------- /// - /// Get main_tensors_all_Allocs - /// ----------------- /// - int main_tensor_nums = main_tensors_all.size(); /// output - comet_debug() << " main_tensor_nums: " << main_tensor_nums << "\n"; - /// Check the loop arg in each tensor - std::vector> main_tensors_all_Allocs = getAllAllocs(main_tensors_all); /// output - comet_debug() << " main_tensors_all_Allocs.size(): " << main_tensors_all_Allocs.size() << "\n"; - - /// ----------------- /// - /// Get allValueAccessIdx - /// ----------------- /// - /// For every main_tensors_all[i], allAccessIdx[i] is the for-loop's induction variable. - /// However, allValueAccessIdx[i] is not necessarily the induction variable. - /// For CSR, for example, - /// for (j_loc = A.rowptr[idx]; j_loc < A.rowptr[idx + 1]; ++j_loc) { j_idx = A.col[j_loc]; } - /// j_idx is allValueAccessIdx[i], and j_loc is allAccessIdx[i] - std::vector> allAccessIdx(main_tensor_nums); - std::vector> allValueAccessIdx(main_tensor_nums); - getForLoopsValueAccessIdx(builder, - loc, - main_tensor_nums, - allPerms, - allFormats, - main_tensors_all, - nested_forops, - nested_AccessIdx, - nested_forops_indices, - main_tensors_all_Allocs, - allAccessIdx /* output */, - allValueAccessIdx /* output */); - - /// Symbolic Phase preparation - std::vector symbolic_nested_forops; - std::vector symbolic_nested_AccessIdx; - std::vector symbolic_nested_forops_indices; - std::vector> symbolic_allAccessIdx(main_tensor_nums); - std::vector> symbolic_allValueAccessIdx(main_tensor_nums); - if (symbolicInfo.has_symbolic_phase) - { - /// Store the insertion point - auto last_insertion_point = builder.saveInsertionPoint(); - - getSymbolicNestedForOpsAndAccessIdx(ancestorsWps, - ancestorsOps, - symbolic_nested_forops /* output */, - symbolic_nested_AccessIdx /* output */, - symbolic_nested_forops_indices /* output */); - - /// Set the insertion point - builder.setInsertionPoint(symbolic_nested_forops[0].getBody()->getTerminator()); - - getForLoopsValueAccessIdx(builder, - loc, - main_tensor_nums, - allPerms, - allFormats, - main_tensors_all, - symbolic_nested_forops, - symbolic_nested_AccessIdx, - symbolic_nested_forops_indices, - main_tensors_all_Allocs, - symbolic_allAccessIdx /* output */, - symbolic_allValueAccessIdx /* output */); - - /// Restore the insertion point - builder.restoreInsertionPoint(last_insertion_point); - } - - int rhs_loc = 0; - int lhs_loc = main_tensors_rhs.size(); /// lhs_loc is the location of the first LHS tensor in main_tensors_all - - /// New version - Value lhs = cur_op.getLhs().getDefiningOp()->getOperand(0); - comet_vdump(lhs); -/// lhs is TensorLoadOp -#ifdef DEBUG_MODE_LowerIndexTreeToSCFPass - Value lhs_alloc = (lhs.getDefiningOp())->getOperand(0); - comet_vdump(lhs_alloc); -#endif - if (main_tensors_rhs.size() == 1) - { /// Generate "a = b" - if (ConstantOp cstop = dyn_cast(main_tensors_rhs[0].getDefiningOp())) - { /// "a = 1.0" - comet_vdump(cstop); - if (comp_worksp_opt) /// true attr means compressed workspace + + Value loop_ctr_init = rewriter.create(loc, index_type, rewriter.getIndexAttr(0)); + loop_args.push_back(loop_ctr_init); + cond_block->addArgument(index_type, loc); + body_block->addArgument(index_type, loc); + unsigned loop_carry_args = loop_args.size(); + + // Create control iterators for each of the tensors + OpBuilder::InsertPoint before; + for(Value domain : domains) { - /// Symbolic Phase - if (symbolicInfo.has_symbolic_phase) + IndexTreeSparseDomainOp sparse_domain = llvm::cast(domain.getDefiningOp()); + TensorFormatEnum format = (TensorFormatEnum)sparse_domain.getFormat(); + switch(format) { - /// Store the insertion point - auto last_insertion_point = builder.saveInsertionPoint(); - - /// Set the insertion point - builder.setInsertionPoint(symbolic_nested_forops[0].getBody()->getTerminator()); - - /// Symbolic Phase uses the W_id_list_size in the Index Tree (main_tensors_all_Allocs[lhs_loc].back) - /// to record the current row size. - /// W_id_list_size = 0; - /// However, Numeric Phase should use C.rowptr[i_idx] to initialize W_id_list_size. - /// W_id_list_size = C.rowptr[i_idx]; - genWorkspaceCmptOpInitialAssignment(builder, - loc, - lhs_loc, - cstop, - symbolic_nested_forops, - tensors_lhs_Allocs, - main_tensors_all_Allocs, - false /* use_dynamic_init */, - symbolicInfo); - - /// Prepare C, C.rowptr - if (symbolicInfo.mtxC_rowptr == nullptr) + case TensorFormatEnum::D: + case TensorFormatEnum::UNK: + { + assert(false && "Invalid format for IndexTreeSparseDomainOp"); + break; + } + case TensorFormatEnum::CN: + case TensorFormatEnum::S: { - Value &W_id_list_size = lhs; - { - comet_vdump(W_id_list_size); + // Not yet supported!!! + break; + } + case TensorFormatEnum::CU: + { + Value start_idx = sparse_domain.getParent(); + if(!start_idx){ + start_idx = rewriter.create(loc, index_type, rewriter.getIndexAttr(0)); } - getOutputMtxCRowptrAndDims(cur_op, - W_id_list_size, - symbolicInfo /* output */); + Value end_idx = rewriter.create(loc, index_type, start_idx, inc); + Value start = rewriter.create(loc, sparse_domain.getPos(), start_idx); + start = rewriter.createOrFold(loc, rewriter.getIndexType(), start); + Value end = rewriter.create(loc, sparse_domain.getPos(), end_idx); + end = rewriter.createOrFold(loc, rewriter.getIndexType(), end); + loop_args.push_back(start); + before = rewriter.saveInsertionPoint(); + + Value crd_idx = cond_block->addArgument(start.getType(), loc); + rewriter.setInsertionPointToStart(cond_block); + Value cnd = rewriter.create( + loc, rewriter.getI1Type(), + arith::CmpIPredicateAttr::get(context, arith::CmpIPredicate::ult), + crd_idx, end + ); + loop_conditions.push_back(cnd); + + crd_idx = body_block->addArgument(start.getType(), loc); + rewriter.setInsertionPointToStart(body_block); + auto dim = rewriter.getI32IntegerAttr(sparse_domain.getDim()); + Value array_crd = rewriter.create(loc, sparse_domain.getTensor(), crd_idx, dim); + array_crd = rewriter.createOrFold(loc, rewriter.getIndexType(), array_crd); + array_crds.push_back(array_crd); + + tensor_access_map.insert(std::make_pair( + std::make_pair(sparse_domain.getTensor(), sparse_domain.getDim()), + std::make_pair(crd_idx, array_crd) + )); } + } + rewriter.restoreInsertionPoint(before); + } + + // Create while loop + scf::WhileOp while_loop = rewriter.create(loc, cond_block->getArgumentTypes(), loop_args); + while_loop.getBefore().push_front(cond_block); - /// Restore the insertion point - builder.restoreInsertionPoint(last_insertion_point); - } /// End symbolic phase - if (allFormats[lhs_loc].empty()) + rewriter.setInsertionPointToEnd(cond_block); + Value loop_condition = nullptr; + for(Value cnd : loop_conditions) + { + if(loop_condition == nullptr) { - /// The computeOp node is W_id_list_size = 0, - /// then do W_id_list_size = symbolicInfo.mtxC_rowptr[idx] - genWorkspaceCmptOpInitialAssignment(builder, - loc, - lhs_loc, - cstop, - nested_forops, - tensors_lhs_Allocs, - main_tensors_all_Allocs, - true /* use_dynamic_init */, - symbolicInfo); + loop_condition = cnd; + } else { + loop_condition = rewriter.create(loc, rewriter.getI1Type(), loop_condition, cnd); } - else + } + rewriter.create(loc, loop_condition, cond_block->getArguments()); + + while_loop.getAfter().push_front(body_block); + // Create intersection + rewriter.setInsertionPointToEnd(body_block); + Value crd = nullptr; + for(Value array_crd : array_crds){ + if(crd == nullptr) { - /// The computeOp node is V[j] = 0, - /// then do V[j] = 0.0 - genWorkspaceCmptOpInitialAssignment(builder, - loc, - lhs_loc, - cstop, - nested_forops, - tensors_lhs_Allocs, - main_tensors_all_Allocs, - false /* use_dynamic_init */, - symbolicInfo); + crd = array_crd; + } else { + crd = rewriter.create(loc, index_type, crd, array_crd); } } - else - { /// initial workspace - /// Generate Store 1.0, A[...] this op - /// this case: allPerms[0] is empty, allFormats[0] is empty - - genCmptOpGeneralInitialAssignment(builder, - loc, - lhs_loc, - cstop, - nested_forops, - main_tensors_all_Allocs, - allValueAccessIdx); - } - } - else if (main_tensors_rhs[0].getType().isa()) - { /// Cij = Wj - /// When Cij is dense type - if (lhs.getType().isa()) - { - /// %1 = load b[...] - /// store %1, a[...] - genCmptOpGatherFromDenseToDense(builder, - loc, - rhs_loc, - lhs_loc, - main_tensors_all_Allocs, - allValueAccessIdx); - } - /// Cij = Wj - else if (lhs.getType().isa()) - { - unsigned int lhs_2crd_size_loc; - unsigned int lhs_2pos_size_loc; - Value lhs_nnz; - Value lhs_nnz_alloc; - Value lhs_val; - getLHSBeforeGatherFromWorkspace(builder, - loc, - lhs_loc, - lhs, - main_tensors_all_Allocs, - lhs_2crd_size_loc /* output */, - lhs_2pos_size_loc /* output */, - lhs_nnz /* output */, - lhs_nnz_alloc /* output */, - lhs_val /* output */); - - if (comp_worksp_opt) /// true attr means compressed workspace - { - /// Gather results from Workspace to the sparse output - genWorkspaceCmptOpGatherFromWorkspaceToOutput(builder, - loc, - tensors_rhs_Allocs, - nested_forops, - nested_AccessIdx, - symbolicInfo, - numericInfo); - /// } - } - else + Value intersection_cnd = nullptr; + SmallVector intersections; + for(Value array_crd : array_crds) + { + Value is_intersect = rewriter.create( + loc, rewriter.getI1Type(), + arith::CmpIPredicateAttr::get(context, arith::CmpIPredicate::eq), + crd, array_crd + ); + if(intersection_cnd == nullptr) { - /// %1 = load b[...] - /// if(%1 != 0) { - /// Cnnz = load Cop.operand(4d+1) - /// store %1, cval[Cnnz] - /// store Cnnz+1, Cop.operand(4d+1) - /// } - genCmptOpGatherFromDenseToOutput(builder, - loc, - rhs_loc, - lhs_loc, - lhs_2crd_size_loc, - lhs_2pos_size_loc, - lhs, - lhs_nnz, - lhs_nnz_alloc, - lhs_val, - allFormats, - main_tensors_all_Allocs, - allAccessIdx, - allValueAccessIdx, - nested_forops); + intersection_cnd = is_intersect; + } else { + intersection_cnd = rewriter.create(loc, rewriter.getI1Type(), intersection_cnd, is_intersect); } + intersections.push_back(is_intersect); } - } - /// Vj = Bij - else if (main_tensors_rhs[0].getType().isa()) - { - /// %Bvalue = load %Bval[..] - /// store %Bvalue, %v[%j] - /// Symbolic Phase - if (symbolicInfo.has_symbolic_phase) + SmallVector if_types; + for(unsigned i = 0; i < loop_carry_args; i++) { - /// Store the insertion point - auto last_insertion_point = builder.saveInsertionPoint(); - - /// Set the insertion point - builder.setInsertionPoint(symbolic_nested_forops[0].getBody()->getTerminator()); - - genWorkspaceCmptOpScatterInputToWorkspace(builder, - loc, - main_tensor_nums, - main_tensors_all_Allocs, - symbolic_allValueAccessIdx); - - /// Restore the insertion point - builder.restoreInsertionPoint(last_insertion_point); - } /// End symbolic phase - - genWorkspaceCmptOpScatterInputToWorkspace(builder, - loc, - main_tensor_nums, - main_tensors_all_Allocs, - allValueAccessIdx); - } - } - else if (main_tensors_rhs.size() == 2) - { /// Generate " a = b * c" binary op + if_types.push_back(loop_args[i].getType()); + } - comet_debug() << "No masking codegen...\n"; + scf::IfOp if_op = rewriter.create(loc, if_types, intersection_cnd, true); + rewriter.setInsertionPointToStart(if_op.elseBlock()); + rewriter.create( + loc, + std::vector( + body_block->args_begin(), + body_block->args_begin() + if_op->getNumResults()) + ); + rewriter.setInsertionPointToStart(if_op.thenBlock()); + // OpBuilder::InsertPoint loop_end = rewriter.saveInsertionPoint(); + Value induction_var = body_block->getArgument(loop_carry_args - 1); + auto step = rewriter.create(loc, index_type, rewriter.getIndexAttr(1)); + auto loop_ctr = rewriter.create(loc, index_type, induction_var, step); + auto yield_op = rewriter.create( + loc, + std::vector( + body_block->args_begin(), + body_block->args_begin() + if_op->getNumResults()) + ); + yield_op->setOperand(loop_carry_args - 1, loop_ctr.getResult()); + + // Increment each argument + rewriter.setInsertionPointAfter(if_op); + SmallVector yield_args; + for(auto result : if_op.getResults()) { + yield_args.push_back(result); + } + auto cntrl_arg = body_block->args_begin() + loop_carry_args; + for(Value cnd : intersections) + { + Value inc = rewriter.create(loc, index_type, cnd); + yield_args.push_back(rewriter.create(loc, index_type, *cntrl_arg, inc)); + cntrl_arg += 1; + } - auto semiringParts = cur_op.getSemiring().split('_'); - /// check validity of semiring provided by user. - if (!Semiring_reduceOps.contains(semiringParts.first) || !Semiring_ops.contains(semiringParts.second)) - { - llvm::errs() << "Not supported semiring operator: " - << semiringParts.first << " or " << semiringParts.second << " \n"; - llvm::errs() << "Please report this error to the developers!\n"; - /// we should not proceed forward from this point to avoid faults. - } + // Create YieldOp + rewriter.create(loc, yield_args); + rewriter.setInsertionPointAfter(while_loop); - MaskingInfo maskingInfo; - maskingInfo.mask_type = NO_MASKING; - if (symbolicInfo.has_symbolic_phase) - { - /// Store the insertion point - auto last_insertion_point = builder.saveInsertionPoint(); - - /// Set the insertion point - builder.setInsertionPoint(symbolic_nested_forops[0].getBody()->getTerminator()); - - genSymbolicSemiringLoopBody(builder, - loc, - lhs_loc, - tensors_lhs_Allocs, - symbolic_nested_forops, - symbolic_nested_AccessIdx, - symbolic_allValueAccessIdx, - symbolicInfo, - nested_forops /* numeric_nested_forops= */, - maskingInfo); - - /// Restore the insertion point - builder.restoreInsertionPoint(last_insertion_point); + ValueRange inner_inputs = body_block->getArguments().drop_back(loop_args.size() - loop_carry_args - 1); + ResultRange outputs = ResultRange(while_loop->result_begin(), while_loop->result_begin() + (loop_carry_args - 1)); + return new IntersectionLoopInfo(inner_inputs, outputs, loop_ctr, map, induction_var, crd, yield_op, tensor_access_map); } - formSemiringLoopBody(cur_op, - comp_worksp_opt, - semiringParts.first, semiringParts.second, - builder, loc, lhs_loc, - main_tensors_all_Allocs, - tensors_lhs_Allocs, - tensors_rhs_Allocs, - allValueAccessIdx, - allAccessIdx, - nested_forops, - nested_AccessIdx, - symbolic_nested_forops, - allPerms_rhs, - symbolicInfo, - numericInfo, - maskingInfo); - } - else if (main_tensors_rhs.size() == 3) - { /// Generate " a = b * c" binary op with masking - - { - /// comet_pdump(rootOp.getOperation()->getParentOfType()); - comet_pdump(rootOp->getParentOfType()); - } - auto semiringParts = cur_op.getSemiring().split('_'); - /// check validity of semiring provided by user. - if (!Semiring_reduceOps.contains(semiringParts.first) || !Semiring_ops.contains(semiringParts.second)) - { - llvm::errs() << "Not supported semiring operator: " - << semiringParts.first << " or " << semiringParts.second << " \n"; - llvm::errs() << "Please report this error to the developers!\n"; - /// we should not proceed forward from this point to avoid faults. + Value getCrd(IRRewriter& rewriter) override { + return crd; } - auto maskingAttr = cur_op.getMaskType(); - std::string maskingAttrStr(maskingAttr.data()); - comet_debug() << "mask attr: " << maskingAttrStr << "\n"; - - MASKING_TYPE mask_type; - if (maskingAttrStr == "push") - mask_type = MASKING_TYPE::PUSH_BASED_MASKING; - else if (maskingAttrStr == "pull") - mask_type = MASKING_TYPE::PULL_BASED_MASKING; - else if (maskingAttrStr == "auto") - mask_type = MASKING_TYPE::PUSH_BASED_MASKING; - else /// none - mask_type = MASKING_TYPE::NO_MASKING; - - switch (mask_type) - { - case NO_MASKING: - { /// Use no masking; we should not hit this case because it is handled - /// by the previous if-else branch when main_tensors_rhs.size() == 2 - break; + Value getPos(IRRewriter& rewriter, Value tensor, uint32_t dim) override { + auto control = controlVars.find(std::make_pair(tensor, dim)); + if(control != controlVars.end()) + return control->getSecond().first; + auto loc = tensor.getLoc(); + Value pos = rewriter.create(loc, rewriter.getIndexType(), tensor, crd, rewriter.getI32IntegerAttr(dim), rewriter.getBoolAttr(true)); + return pos; } - case PUSH_BASED_MASKING: - { /// Use push-based masking - /// mask_tensor should be the 3rd operand of ComputeRHS (tensors_rhs[2]). - mlir::Value mask_tensor = tensors_rhs[2]; - { - comet_debug() << "mask_tensor\n"; - comet_vdump(mask_tensor); - } - MaskingInfo maskingInfo; - maskingInfo.mask_type = PUSH_BASED_MASKING; - maskingInfo.mask_tensor = mask_tensor; - /// Get mask_rowptr, mask_col, and mask_val arrays - getMaskSparseTensorInfo(maskingInfo /* contents updated after call*/); - - if (symbolicInfo.has_symbolic_phase) - { - /// Store the insertion point - auto last_insertion_point = builder.saveInsertionPoint(); - - /// Set the insertion point - builder.setInsertionPoint(symbolic_nested_forops[0].getBody()->getTerminator()); - - genSymbolicSemiringLoopBody(builder, - loc, - lhs_loc, - tensors_lhs_Allocs, - symbolic_nested_forops, - symbolic_nested_AccessIdx, - symbolic_allValueAccessIdx, - symbolicInfo, - nested_forops /* numeric_nested_forops= */, - maskingInfo); - - /// Restore the insertion point - builder.restoreInsertionPoint(last_insertion_point); - } - formSemiringLoopBody(cur_op, - comp_worksp_opt, - semiringParts.first, semiringParts.second, - builder, loc, lhs_loc, - main_tensors_all_Allocs, - tensors_lhs_Allocs, - tensors_rhs_Allocs, - allValueAccessIdx, - allAccessIdx, - nested_forops, - nested_AccessIdx, - symbolic_nested_forops, - allPerms_rhs, - symbolicInfo, - numericInfo, - maskingInfo); - break; + void updateOutput(IRRewriter& rewriter, uint32_t idx, Value newOutput) override { + currentInputs[idx] = newOutput; + rewriter.modifyOpInPlace(terminator, [&](){terminator.setOperand(idx, newOutput);}); } - case PULL_BASED_MASKING: /// Use pull-based masking - llvm::errs() << "Error: mask type PULL_BASED_MASKING is not supported, yet.\n"; - } - } - else - { - llvm::errs() << "No support for operation with greater than two operands in workspace transforms!" - << "\n"; - } - } + }; - /// ----------------- /// - /// Get the itree roots - /// ----------------- /// - void getIndexTreeOps(func::FuncOp &function, - std::vector &iTreeRoots /* output */) - { - function.walk([&](indexTree::IndexTreeOp op) - { iTreeRoots.push_back(op); }); - } + class MaskedLoopInfo : LoopInfo { + private: + scf::YieldOp terminator; + LoopInfo* internal_loop; + public: + MaskedLoopInfo(ValueRange inputs, ValueRange outputs, Operation* body, IRMapping ir_map, scf::YieldOp yield, LoopInfo* internal) : + LoopInfo(inputs, outputs, body, ir_map), terminator(yield), internal_loop(internal) {} - /// ----------------- /// - /// Delete every objects in opstree_vec, preventing memory leak. - /// ----------------- /// - void cleanOpstreeVec(std::vector &opstree_vec) - { - for (auto &t : opstree_vec) - { - delete t; - } - } - - /// ----------------- /// - /// Check if the Index Tree inputs are all sparse - /// All inputs are sparse if and only if all computeOp nodes are using workspace transformation. - /// ----------------- /// - void checkIfAllSparse(std::vector &wp_ops, - SymbolicInfo &symbolicInfo /* output */) - { - for (Value &op : wp_ops) - { - if (indexTree::IndexTreeComputeOp cur_op = dyn_cast(op.getDefiningOp())) + static LoopInfo* build(Operation* domain_op, IRRewriter& rewriter, ValueRange inputs) { - bool comp_worksp_opt(cur_op.getCompWorkspOpt()); - if (!comp_worksp_opt) + auto loc = domain_op->getLoc(); + auto masked_domain = llvm::cast(domain_op); + [[maybe_unused]] auto index_type = rewriter.getIndexType(); + [[maybe_unused]] auto context = rewriter.getContext(); + + // Create internal loop + Operation* child = masked_domain.getBase().getDefiningOp(); + LoopInfo* loop_info = llvm::TypeSwitch(child) + .Case([&](IndexTreeDenseDomainOp op) { + return DenseLoopInfo::build(op, rewriter, inputs); + }) + .Case([&](IndexTreeSparseDomainOp op) { + switch((TensorFormatEnum)op.getFormat()){ + case TensorFormatEnum::D: + case TensorFormatEnum::UNK: + { + assert(false && "Invalid format for IndexTreeSparseDomainOp"); + return (LoopInfo*)nullptr; + + break; + } + case TensorFormatEnum::CN: + case TensorFormatEnum::CU: + return SparseLoopInfo::build(op, rewriter, inputs); + case TensorFormatEnum::S: + assert(false && "Singleton loop inside a mask is not supported."); + // return SingletonLoopInfo::build(op, rewriter, inputs, parent_info); + } + }) + .Case([&](IndexTreeWorkspaceDomainOp op) { + return WorkspaceLoopInfo::build(op, rewriter, inputs); + }) + .Case([&](IndexTreeDomainIntersectionOp op) { + return IntersectionLoopInfo::build(op, rewriter, inputs); + }) + .Default([](Operation *op) { + assert(false && "IndexNode not given a valid domain"); + return nullptr; + }); + auto after_internal_loop = rewriter.saveInsertionPoint(); + + // Create if op + rewriter.setInsertionPoint(loop_info->loopBody); + SmallVector if_types; + for(Value input : loop_info->getInputs()) { - symbolicInfo.are_inputs_sparse = false; - return; + if_types.push_back(input.getType()); } + Value cond = rewriter.create(loc, masked_domain.getMask(), loop_info->getCrd(rewriter)); + auto if_op = rewriter.create(loc, if_types, cond, true); + rewriter.setInsertionPointToStart(if_op.elseBlock()); + rewriter.create(loc, loop_info->getInputs()); + rewriter.setInsertionPointToStart(if_op.thenBlock()); + SmallVector new_inputs = SmallVector(loop_info->getInputs().begin(), loop_info->getInputs().end()); + scf::YieldOp terminator = rewriter.create(loc, new_inputs); + uint32_t i = 0; + for(Value output : if_op.getResults()){ + loop_info->updateOutput(rewriter, i, output); + i += 1; + } + rewriter.restoreInsertionPoint(after_internal_loop); + + return new MaskedLoopInfo(new_inputs, loop_info->getResults(), terminator, loop_info->map, terminator, loop_info); } - } - symbolicInfo.are_inputs_sparse = true; - } + Value getCrd(IRRewriter& rewriter) override { + return internal_loop->getCrd(rewriter); + } + + Value getPos(IRRewriter& rewriter, Value tensor, uint32_t dim) override { + return internal_loop->getPos(rewriter, tensor, dim); + } - //===----------------------------------------------------------------------===// - /// LowerIndexTreeIRToSCF PASS - //===----------------------------------------------------------------------===// + void updateOutput(IRRewriter& rewriter, uint32_t idx, Value newOutput) override { + currentInputs[idx] = newOutput; + rewriter.modifyOpInPlace(terminator, [&](){terminator.setOperand(idx, newOutput);}); + } + }; - /// Lower the ta.tc (tensor contraction operation in TA dialect) into scf dialect. struct LowerIndexTreeToSCFPass : public PassWrapper> { MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(LowerIndexTreeToSCFPass) - void runOnOperation() override; - - void doLoweringIndexTreeToSCF(indexTree::IndexTreeOp &rootOp, - OpBuilder &builder); - }; - -} /// end anonymous namespace. - -/** - * @brief : - * Goal: IndexTreeOp(i.e. a tree structure), convert into OpsTree(also tree structure) - * Steps: 1.Iterate over IndexTreeOptree - * 2.pass info to opsGen(), including tensors, current workspacetreeop, parent OpsTree node - * -- the parent of "current workspacetreeop" can get from getUser(). Only one user(tree structure) - * -- DFS traverse the workspacetreeop. How? - * */ -void LowerIndexTreeToSCFPass::doLoweringIndexTreeToSCF(indexTree::IndexTreeOp &rootOp, - OpBuilder &builder) -{ - assert(isa(rootOp)); - comet_debug() << "\ndoLoweringIndexTreeToSCF in LowerIndexTreeIRToSCF\n"; - /// auto module = rootOp->getParentOfType(); - { - /// comet_pdump(rootOp.getOperation()->getParentOfType()); - comet_pdump(rootOp->getParentOfType()); - } - /// comet_pdump(rootOp.getOperation()->getParentOp()); - /// Here, should check the operands, at least one operand should be sparse; - /// Otherwise, if all dense operands, just return. - /// rootOp only contains one workspace child, no indices + SmallDenseMap nodeMap; + SmallDenseMap leafMap; + llvm::ScopedPrinter logger{llvm::dbgs()}; - std::vector wp_ops; - dfsRootOpTree(rootOp.getChildren(), wp_ops); - // #ifdef DEBUG_MODE_LowerIndexTreeToSCFPass - comet_debug() << " wp_ops.size(): " << wp_ops.size() << "\n"; - for ([[maybe_unused]] auto n : wp_ops) - { - comet_debug() << " "; - comet_vdump(n); - /// Declare opsTree - } - // #endif - - /// In ops vector, for each op, the parent of each op can get from getUsers() - /// Since it's a tree structure, only one user ==> which is the parent - /// We can initialize the OpsTree structure with this relationship. - /// Search the location of the parent of current op, if rootOp, return ops.size; - /// Otherwise, return the location index. - std::vector parent_idx; - for (unsigned int i = 0; i < wp_ops.size(); i++) - { - mlir::Value wp_op = wp_ops[i]; - mlir::Value wp_parent; - - for (auto n : wp_op.getDefiningOp()->getUsers()) + Value mapInputIntoLoop(Value input, LoopInfo* loop_info) { - comet_debug() << " " << i << " "; - comet_pdump(n); - wp_parent = n->getResult(0); - - comet_debug() << " parent: " << findIndexInVector_Value(wp_ops, wp_parent) << "\n"; - bool isInTree = false; - if (findIndexInVector_Value(wp_ops, wp_parent) < wp_ops.size()) - { - isInTree = true; + auto input_idx_iter = leafMap.find(input); + if(input_idx_iter != leafMap.end()){ + uint32_t input_idx = input_idx_iter->getSecond(); + return loop_info->getInput(input_idx); } - - if (isInTree || isRealRoot(wp_op.getDefiningOp())) - parent_idx.push_back(findIndexInVector_Value(wp_ops, wp_parent)); + return input; } - } - -#ifdef DEBUG_MODE_LowerIndexTreeToSCFPass - comet_debug() << " parent_idx: " << parent_idx.size() << "\n"; - for (auto n : parent_idx) - { - comet_debug() << " " << n << " \n"; - /// Declare opsTree - } -#endif - - std::vector opstree_vec; - for (unsigned int i = 0; i < wp_ops.size(); i++) - { - // std::vector forOps; - // std::vector accessIdx; - OpsTree *parent = nullptr; - if (i >= 1) - { /// Not rootop - parent = opstree_vec[parent_idx[i]]; - } - comet_debug() << " \n"; - // OpsTree *ops = new OpsTree(forOps, accessIdx, parent, i); - OpsTree *ops = new OpsTree(parent, i); - if (parent != nullptr) - { /// add child to the parent - parent->addChild(ops); + void updateOutput(Value input, Value old_output, Value new_output, LoopInfo* loop_info, IRRewriter& rewriter) + { + uint32_t input_idx = leafMap.find(input)->getSecond(); + leafMap.insert(std::pair(old_output, input_idx)); + loop_info->updateOutput(rewriter, input_idx, new_output); } - opstree_vec.push_back(ops); - } - -#ifdef DEBUG_MODE_LowerIndexTreeToSCFPass - { - int opstree_i = 0; - for (auto n : opstree_vec) + Value convertOperand(LoopInfo* loop_info, IndexTreeLHSOperandOp op, StringRef semiring, IRRewriter &rewriter) { - comet_debug() << " " << n->id << "\n"; - comet_debug() << "opstree_vec[" << opstree_i << "] " - << "forOps.size():" << n->forOps.size() << " " - << "accessIdx.size():" << n->accessIdx.size() << "\n"; - /// << "cmptOps.size():" << n->cmptOps.size() << "\n"; - if (n->parent != nullptr) - { - comet_debug() << "parent: " << n->parent->id << "\n"; - } - else - { - comet_debug() << "parent: null \n"; + Location loc = op->getLoc(); + Value tensor = mapInputIntoLoop(op.getTensor(), loop_info); + auto crds = op.getCrds(); + auto positions = op.getPos(); + + TensorType tensor_type; + if (llvm::isa(tensor.getType()) || llvm::isa(tensor.getType())) { + /// LHS operand is a constant value. + return tensor; + } else if((tensor_type = llvm::dyn_cast(tensor.getType()))){ + return rewriter.create(loc, tensor_type.getElementType(), tensor, positions); + } else { + Type element_type; + if(llvm::isa(tensor.getType())) + { + element_type = llvm::cast(tensor.getType()).getElementType(); + } else if(llvm::isa(tensor.getType())) + { + element_type = llvm::cast(tensor.getType()).getElementType(); + } + Value pos = positions[positions.size() - 1]; + double zero = 0; + if(semiring == "minxy"){ + zero = INFINITY; + } + return rewriter.create(loc, element_type, tensor, pos, crds, rewriter.getF64FloatAttr(zero)); + } + } + + Value convertOperand(LoopInfo* loop_info, IndexTreeOperandOp op, StringRef semiring, IRRewriter &rewriter) + { + Location loc = op->getLoc(); + Value tensor = mapInputIntoLoop(op.getTensor(), loop_info); + auto crds = op.getCrds(); + auto positions = op.getPos(); + + TensorType tensor_type; + if (llvm::isa(tensor.getType()) || llvm::isa(tensor.getType())) { + /// RHS operand is a constant value. + return tensor; + } else if((tensor_type = llvm::dyn_cast(tensor.getType()))){ + return rewriter.create(loc, tensor_type.getElementType(), tensor, crds); + } else { + // LHS may not be constant (i.e. if we are inserting into a tensor that we need to resize), + // so cannot directly lower like we can the RHS + Value pos = positions[positions.size() - 1]; + double zero = 0; + if(semiring == "minxy"){ + zero = INFINITY; + } + + return rewriter.create(loc, rewriter.getF64Type(), tensor, pos, crds, rewriter.getF64FloatAttr(zero)); } - ++opstree_i; } - } -#endif - - SymbolicInfo symbolicInfo; - NumericInfo numericInfo; - checkIfAllSparse(wp_ops, - symbolicInfo /* output */); - if (symbolicInfo.are_inputs_sparse) - { - symbolicInfo.has_symbolic_phase = true; - } - for (unsigned int i = 0; i < wp_ops.size(); i++) - { - comet_debug() << " i: " << i << "\n"; - comet_vdump(wp_ops[i]); - if (indexTree::IndexTreeIndicesOp cur_op = dyn_cast(wp_ops[i].getDefiningOp())) + mlir::LogicalResult convertCompute(IndexTreeComputeOp compute_op, IRRewriter &rewriter) { - /// Get indices - ArrayAttr op_indices = cur_op.getIndices(); - comet_debug() << "curOp is IndexTreeIndicesOp\n"; - comet_vdump(cur_op); - - /// cur_op's index attribute, e.g., "indices = [0]" - std::vector indices; - for (unsigned int j = 0; j < op_indices.size(); j++) + auto loc = compute_op->getLoc(); + LoopInfo* parent_info = nodeMap.find(compute_op.getParent())->getSecond(); + rewriter.setInsertionPoint(parent_info->loopBody); + + auto semiringParts = compute_op.getSemiring().split('_'); + Value elementwise_result; + for(auto rhs = compute_op.getRhs().begin(); rhs != compute_op.getRhs().end(); rhs++) { - /// Get the indices; - int idx = op_indices[j].cast().getInt(); - indices.push_back(idx); + Value rhs_value = convertOperand(parent_info, cast((*rhs).getDefiningOp()), semiringParts.first, rewriter); + if(rhs == compute_op.getRhs().begin()){ + elementwise_result = rhs_value; + } else { + elementwise_result = getSemiringSecondVal(rewriter, loc, semiringParts.second, + elementwise_result, rhs_value); + } } - comet_debug() << " indices.size(): " << indices.size() << "\n"; - /// Leaves are the computeOp nodes and the children of cur_op (an index node) - std::vector leafs; + IndexTreeLHSOperandOp lhs = llvm::cast(compute_op.getLhs().getDefiningOp()); + Value reduce_result = convertOperand(parent_info, lhs, semiringParts.first, rewriter); + reduce_result = getSemiringFirstVal(rewriter, loc, semiringParts.first, + reduce_result, elementwise_result); + + Value old_tensor = mapInputIntoLoop(lhs.getTensor(), parent_info); + Value output_tensor; + if (llvm::isa(old_tensor.getType())) { + LLVM_DEBUG({logger.startLine() << __FILE__ << ":" << __LINE__ << " " << old_tensor << "\n";}); + LLVM_DEBUG({logger.startLine() << __FILE__ << ":" << __LINE__ << " " << reduce_result << "\n";}); + output_tensor = reduce_result; + } else if(llvm::isa(old_tensor.getType())) { + output_tensor = rewriter.create(loc, old_tensor.getType(), reduce_result, old_tensor, lhs.getPos()); + } else { + output_tensor = rewriter.create(loc, old_tensor.getType(), old_tensor, lhs.getPos(), lhs.getCrds(), reduce_result); + } + updateOutput(lhs.getTensor(), compute_op.getResult(), output_tensor, parent_info, rewriter); + return success(); + } + + mlir::LogicalResult convertClearWorkspaceOp(IndexTreeCleanWorkspaceOp op, IRRewriter &rewriter) + { + Location loc = op->getLoc(); + LoopInfo* parent_info = nodeMap.find(op.getParent())->getSecond(); + rewriter.setInsertionPoint(parent_info->loopBody); + Value input = mapInputIntoLoop(op.getWorkspace(), parent_info); + Value workspace = rewriter.create(loc, op->getResultTypes(), input); + updateOutput(op.getWorkspace(), op.getResult(), workspace, parent_info, rewriter); + return success(); + } + + mlir::LogicalResult convertSymbolicDomainOp(ComputeSymbolicDomainOp op, IRRewriter &rewriter) + { + + Location loc = op->getLoc(); + LoopInfo* parent_info = nodeMap.find(op.getParent())->getSecond(); + rewriter.setInsertionPoint(parent_info->loopBody); + Value symbolic_domain = mapInputIntoLoop(op.getDomain(), parent_info); + Value new_domain = rewriter.create(loc, + symbolic_domain.getType(), + symbolic_domain, + parent_info->getCrd(rewriter), + op.getIsUniqueAttr()); + updateOutput(op.getDomain(), op.getResult(), new_domain, parent_info, rewriter); + return success(); + } + + mlir::LogicalResult convertSymbolicDomainEndRowOp(ComputeSymbolicDomainRowOp op, IRRewriter &rewriter) + { + Location loc = op->getLoc(); + LoopInfo* parent_info = nodeMap.find(op.getParent())->getSecond(); + rewriter.setInsertionPoint(parent_info->loopBody); + Value symbolic_domain = mapInputIntoLoop(op.getDomain(), parent_info); + Value new_domain = rewriter.create(loc, + symbolic_domain.getType(), + symbolic_domain, + op.getNeedsMarkAttr()); + updateOutput(op.getDomain(), op.getResult(), new_domain, parent_info, rewriter); + return success(); + } + + mlir::LogicalResult convertFillMaskOp(IndexTreeFillMaskOp fill_mask_op, IRRewriter &rewriter) { + auto loc = fill_mask_op.getLoc(); + LoopInfo* parent_info = nodeMap.find(fill_mask_op.getParent())->getSecond(); + rewriter.setInsertionPoint(parent_info->loopBody); + Value mask_tensor = mapInputIntoLoop(fill_mask_op.getInit(), parent_info); + + // Create loop to fill mask + Operation* mask = fill_mask_op.getDomain().getDefiningOp(); + ValueRange fill_loop_inputs(mask_tensor); + LoopInfo* fill_loop = llvm::TypeSwitch(mask) + .Case([&](IndexTreeDenseDomainOp op) { + return DenseLoopInfo::build(op, rewriter, fill_loop_inputs); + }) + .Case([&](IndexTreeSparseDomainOp op) { + switch((TensorFormatEnum)op.getFormat()){ + case TensorFormatEnum::UNK: + case TensorFormatEnum::D: + assert(false && "Invalid format for IndexTreeSparseDomainOp"); + return (LoopInfo*)nullptr; + break; + case TensorFormatEnum::CN: + case TensorFormatEnum::CU: + return SparseLoopInfo::build(op, rewriter, fill_loop_inputs); + case TensorFormatEnum::S: + assert(false && "Singleton loop inside a mask is not supported."); + return (LoopInfo*)nullptr; + // return SingletonLoopInfo::build(op, rewriter, inputs, parent_info); + break; + } + }) + .Case([&](IndexTreeWorkspaceDomainOp op) { + return WorkspaceLoopInfo::build(op, rewriter, fill_loop_inputs); + }) + .Case([&](IndexTreeDomainIntersectionOp op) { + return IntersectionLoopInfo::build(op, rewriter, fill_loop_inputs); + }) + .Default([](Operation *op) { + assert(false && "IndexNode not given a valid domain"); + return nullptr; + }); + + // Fill in bit tensor + auto after_fill_loop = rewriter.saveInsertionPoint(); + rewriter.setInsertionPoint(fill_loop->loopBody); + Value crd = fill_loop->getCrd(rewriter); + Value t = rewriter.create(loc, rewriter.getI1Type(), rewriter.getBoolAttr(true)); + Value updated_tensor = rewriter.create(loc, fill_loop->getInput(0).getType(), t, fill_loop->getInput(0), crd); + fill_loop->updateOutput(rewriter, 0, updated_tensor); + rewriter.restoreInsertionPoint(after_fill_loop); + + updateOutput( + fill_mask_op.getInit(), + fill_mask_op.getResult(), + fill_loop->getResults()[0], + parent_info, + rewriter + ); + + delete fill_loop; + return success(); + } + + mlir::LogicalResult convertZeroMaskOp(IndexTreeZeroMaskOp zero_mask_op, IRRewriter &rewriter) { + auto loc = zero_mask_op.getLoc(); + LoopInfo* parent_info = nodeMap.find(zero_mask_op.getParent())->getSecond(); + rewriter.setInsertionPoint(parent_info->loopBody); + Value mask_tensor = mapInputIntoLoop(zero_mask_op.getInit(), parent_info); + + // Create loop to fill mask + Operation* mask = zero_mask_op.getDomain().getDefiningOp(); + ValueRange zero_loop_inputs(mask_tensor); + LoopInfo* zero_loop = llvm::TypeSwitch(mask) + .Case([&](IndexTreeDenseDomainOp op) { + return DenseLoopInfo::build(op, rewriter, zero_loop_inputs); + }) + .Case([&](IndexTreeSparseDomainOp op) { + switch((TensorFormatEnum)op.getFormat()){ + case TensorFormatEnum::UNK: + case TensorFormatEnum::D: + assert(false && "Invalid format for IndexTreeSparseDomainOp"); + return (LoopInfo*)nullptr; - /// Find leaves of cur_op in the Index Tree (wp_ops). - /// A leaf is a computeOp node and cur_op is one its ancestors. - findLeafs(cur_op, indices, wp_ops, leafs /* output leaves*/); -#ifdef DEBUG_MODE_LowerIndexTreeToSCFPass - comet_debug() << " leafs.size(): " << leafs.size() << "\n"; - for (auto n : leafs) - { - comet_debug() << " "; - comet_vdump(n); - } -#endif - - /// tensors: the tensors that uses the cur_op (index node) as their iterative index. - /// ids: An id is the location (0, 1, 2, ...) of the cur_op (index node) in the tensor's Perms. - /// formats: The format (e.g., "D", "CU", "CN", etc.) for the id-th dimension of the tensor. - /// tensors[i] uses the cur_op (index node) as its iterative index (e.g., [0], [1], etc.), and - /// ids[i] is the location of the iterative index in tensors[i]'s Perms. - /// For example, - /// allPerms = [[0, 1]]; the tensor[i]'s Perms is [0, 1]. If cur_op (iterative index) is indices = [1], then - /// ids[i] = 1, because [1] is at location 1 in [0, 1], i.e., the 1-st dimension of the tensors[i]. - /// formats[i] is "CU" if tensors[i]'s Formats = ["D", "CU"]. - std::vector tensors; - std::vector ids; - std::vector formats; - - comet_vdump(cur_op); - - getFormatsInfo(cur_op, - indices, - leafs, - tensors /* output */, - ids /* output */, - formats /* output */); - llvm::StringRef iteratorType = cur_op.getIteratorType(); - - comet_debug() << " indices.size(): " << indices.size() << " tensors.size(): " << tensors.size() << "\n"; - for ([[maybe_unused]] unsigned int m = 0; m < tensors.size(); m++) - { - comet_debug() << " Formats:" << formats[m] << " " << ids[m] << " "; - comet_vdump(tensors[m]); + break; + case TensorFormatEnum::CN: + case TensorFormatEnum::CU: + return SparseLoopInfo::build(op, rewriter, zero_loop_inputs); + case TensorFormatEnum::S: + assert(false && "Singleton loop inside a mask is not supported."); + return (LoopInfo*)nullptr; + // return SingletonLoopInfo::build(op, rewriter, inputs, parent_info); + } + }) + .Case([&](IndexTreeWorkspaceDomainOp op) { + return WorkspaceLoopInfo::build(op, rewriter, zero_loop_inputs); + }) + .Case([&](IndexTreeDomainIntersectionOp op) { + return IntersectionLoopInfo::build(op, rewriter, zero_loop_inputs); + }) + .Default([](Operation *op) { + assert(false && "IndexNode not given a valid domain"); + return nullptr; + }); + + // Fill in bit tensor + auto after_zero_loop = rewriter.saveInsertionPoint(); + rewriter.setInsertionPoint(zero_loop->loopBody); + Value crd = zero_loop->getCrd(rewriter); + Value f = rewriter.create(loc, rewriter.getI1Type(), rewriter.getBoolAttr(false)); + Value updated_tensor = rewriter.create(loc, zero_loop->getInput(0).getType(), f, zero_loop->getInput(0), crd); + zero_loop->updateOutput(rewriter, 0, updated_tensor); + rewriter.restoreInsertionPoint(after_zero_loop); + updateOutput( + zero_mask_op.getInit(), + zero_mask_op.getResult(), + zero_loop->getResults()[0], + parent_info, + rewriter + ); + + delete zero_loop; + return success(); + } + + mlir::LogicalResult convertTensorAccessOp(IndexTreeIndexToTensorOp access_op, IRRewriter &rewriter) + { + // Find the position to insert these operations based off the nearest use. + // TODO: figure out order of users? + for(auto user : access_op->getUsers()){ + if(auto node = llvm::dyn_cast(user)) { + Value parent = node.getParentNode(); + auto parent_iterator = nodeMap.find(parent); + if(parent_iterator != nodeMap.end()){ + // Jump out of inner loop if necessary + rewriter.setInsertionPoint(parent_iterator->getSecond()->loopBody); + } + break; + } } + LoopInfo* parent_info = nodeMap.find(access_op.getIndex())->getSecond(); + Value tensor = mapInputIntoLoop(access_op.getTensor(), parent_info); + auto dim = access_op.getDim(); - comet_debug() << " call genForOps, i = " << i << "\n"; - genForOps(tensors, ids, formats, rootOp, builder, opstree_vec[i], symbolicInfo, iteratorType); + Value access_crd = parent_info->getCrd(rewriter); + Value access_pos = parent_info->getPos(rewriter, tensor, dim); + if(auto tensor_type = dyn_cast(tensor.getType())) { - comet_pdump(rootOp->getParentOfType()); - } - comet_debug() << " finished call genForOps, i = " << i << "\n"; - } - else if (indexTree::IndexTreeComputeOp cur_op = dyn_cast(wp_ops[i].getDefiningOp())) - { - /// Generate computation ops. - std::vector ancestors_wp; /// workspace tree ancestor - getAncestorsWp(cur_op, ancestors_wp, wp_ops); -#ifdef DEBUG_MODE_LowerIndexTreeToSCFPass - comet_debug() << " Current Op (IndexTreeComputeOp):"; - comet_vdump(cur_op); - for (auto n : ancestors_wp) + TensorFormatEnum format = (TensorFormatEnum) tensor_type.getFormat()[2 * dim]; + if(format == TensorFormatEnum::D) + { + // TODO: This is incorrect, deal with reordering!!!! + if(access_op.getPrevDim()) { + auto loc = access_op.getLoc(); + auto index_type = rewriter.getIndexType(); + Value dim_size = rewriter.create(loc, index_type, tensor, rewriter.getI32IntegerAttr(dim)); + Value pos_start = rewriter.create(loc, index_type, dim_size, access_op.getPrevDim()); + access_pos = rewriter.create(loc, index_type, pos_start, access_pos); + } + } + } + rewriter.replaceAllUsesWith(access_op.getPos(), access_pos); + rewriter.replaceAllUsesWith(access_op.getCrd(), access_crd); + return success(); + } + + mlir::LogicalResult convertRoot(IndexTreeRootOp root_op, IRRewriter &rewriter) { + IndexTreeOp tree = root_op->getParentOfType(); + indexTree::YieldOp yield = cast(tree.getBody()->getTerminator()); + IRMapping map; + LoopInfo* sentinel_info = SentinelLoopInfo::build(tree.getBody()->getArguments(), tree.getResults(), root_op, map, yield); + nodeMap.insert(std::make_pair(root_op.getResult(), sentinel_info)); + + uint32_t i = 0; + for(Value v: tree.getBody()->getArguments()){ + leafMap.insert(std::make_pair(v, i)); + i++; + } + return success(); + } + + mlir::LogicalResult convertIndexNode(IndexTreeIndicesOp index_node_op, IRRewriter &rewriter) + { + IndexTreeOp tree = index_node_op->getParentOfType(); + Operation* domain_op = index_node_op.getDomain().getDefiningOp(); + Value index_node = index_node_op->getResult(0); + LoopInfo* parent_info = nodeMap.find(index_node_op.getParent())->getSecond(); + rewriter.setInsertionPoint(parent_info->loopBody); + + LoopInfo* loop_info = llvm::TypeSwitch(domain_op) + .Case([&](IndexTreeDenseDomainOp op) { + if(index_node_op.getIsParallel()){ + ValueRange inputs = parent_info->getInputs(); + auto output_analysis = Pass::getChildAnalysis(index_node_op->getParentOp()); + auto output_sets = output_analysis.getOutputSets(index_node_op); + uint32_t i = 0; + for(Value v: inputs){ + output_sets.insert(std::make_pair(v, output_sets[tree.getBody()->getArgument(i)])); + i++; + } + return DenseParallelLoopInfo::build(op, rewriter, parent_info->getInputs(), output_sets, index_node_op.getParallelDimAttr()); + } + return DenseLoopInfo::build(op, rewriter, parent_info->getInputs()); + }) + .Case([&](IndexTreeSparseDomainOp op) { + switch((TensorFormatEnum)op.getFormat()){ + case TensorFormatEnum::D: + case TensorFormatEnum::UNK: + { + assert(false && "Invalid format for IndexTreeSparseDomainOp"); + return (LoopInfo*)nullptr; + break; + } + case TensorFormatEnum::CN: + case TensorFormatEnum::CU: + return SparseLoopInfo::build(op, rewriter, parent_info->getInputs()); + case TensorFormatEnum::S: + return SingletonLoopInfo::build(op, rewriter, parent_info->getInputs(), parent_info); + } + }) + .Case([&](IndexTreeWorkspaceDomainOp op) { + op.getTensorMutable().assign(mapInputIntoLoop(op.getTensor(), parent_info)); + return WorkspaceLoopInfo::build(op, rewriter, parent_info->getInputs()); + }) + .Case([&](IndexTreeDomainIntersectionOp op) { + return IntersectionLoopInfo::build(op, rewriter, parent_info->getInputs()); + }) + .Case([&](IndexTreeMaskedDomainOp op) { + op.getMaskMutable().assign(mapInputIntoLoop(op.getMask(), parent_info)); + return MaskedLoopInfo::build(op, rewriter, parent_info->getInputs()); + }) + .Default([](Operation *op) { + assert(false && "IndexNode not given a valid domain"); + return nullptr; + }); + + ValueRange loop_outputs = loop_info->getResults(); + uint32_t i = 0; + for(Value v : loop_outputs) { - comet_debug() << " "; - comet_vdump(n); - } -#endif + parent_info->updateOutput(rewriter, i, v); + i++; + } + nodeMap.insert(std::make_pair(index_node, loop_info)); + rewriter.setInsertionPoint(loop_info->loopBody); + return success(); + } + + mlir::LogicalResult convertTree(IndexTreeOp treeOp, IRRewriter &rewriter) + { + // Walk the Block of the IndexTreeOp and Convert Operands in order to maintain + // order in generated code + SmallVector toDelete; + LLVM_DEBUG({logger.startLine() << "Current Tree: \n" << treeOp << "\n";}); + for (Operation &op : treeOp.getRegion().front()){ + // A node in the tree should be able to use the loopMaps map to find where it + // should be rewritten to + LogicalResult result = llvm::TypeSwitch(&op) + .Case([&](IndexTreeRootOp op) { + LLVM_DEBUG({ + logger.startLine() << "Converting: " << op << "\n"; + }); + return convertRoot(op, rewriter); + }) + .Case([&](IndexTreeComputeOp op) { + LLVM_DEBUG({ + logger.startLine() << "Converting: " << op << "\n"; + }); + return convertCompute(op, rewriter); + }) + .Case([&](IndexTreeCleanWorkspaceOp op) { + LLVM_DEBUG({ + logger.startLine() << "Converting: " << op << "\n"; + }); + return convertClearWorkspaceOp(op, rewriter); + }) + .Case([&](ComputeSymbolicDomainOp op) { + LLVM_DEBUG({ + logger.startLine() << "Converting: " << op << "\n"; + }); + return convertSymbolicDomainOp(op, rewriter); + }) + .Case([&](ComputeSymbolicDomainRowOp op) { + LLVM_DEBUG({ + logger.startLine() << "Converting: " << op << "\n"; + }); + return convertSymbolicDomainEndRowOp(op, rewriter); + }) + .Case([&](IndexTreeFillMaskOp op) { + LLVM_DEBUG({ + logger.startLine() << "Converting: " << op << "\n"; + }); + return convertFillMaskOp(op, rewriter); + }) + .Case([&](IndexTreeZeroMaskOp op) { + LLVM_DEBUG({ + logger.startLine() << "Converting: " << op << "\n"; + }); + return convertZeroMaskOp(op, rewriter); + }) + .Case([&](IndexTreeIndexToTensorOp op) { + LLVM_DEBUG({ + logger.startLine() << "Converting: " << op << "\n"; + }); + return convertTensorAccessOp(op, rewriter); + }) + .Case([&](IndexTreeIndicesOp op) { + LLVM_DEBUG({ + logger.startLine() << "Converting: " << op << "\n"; + }); + LogicalResult result = convertIndexNode(op, rewriter); + // LLVM_DEBUG({logger.startLine() << "Current Tree: \n" << treeOp << "\n";}); + return result; + + }) + .Default([](Operation *op) { + return success(); + }); + + if(result.failed()){ + LLVM_DEBUG({ + logger.startLine() << "Conversion failed!" << "\n"; + }); + return failure(); + } - comet_debug() << " call genCmptOps, i = " << i << "\n"; - /// ancestors_wp can give all the indices of the nested loops - genCmptOps(cur_op, rootOp, builder, opstree_vec[i], ancestors_wp, - wp_ops, symbolicInfo, numericInfo); - { - comet_pdump(rootOp->getParentOfType()); - } - comet_debug() << " finished call genCmptOps, i = " << i << "\n"; - } - } + if(isa(op)){ + break; + } + toDelete.push_back(&op); + } + + LLVM_DEBUG({logger.startLine() << "Current Tree: \n" << treeOp << "\n";}); + for (auto op = toDelete.rbegin(); op != toDelete.rend(); ++op){ + // Erase all the old ops in the region + LLVM_DEBUG({ + logger.startLine() << "Removing: " << (*op)->getName() << "\n"; + if(!(*op)->use_empty()) { + logger.indent(); + for(auto user : (*op)->getUsers()) { + logger.startLine() << "Op still used by: " << user->getName() << "\n"; + } + } + }); - { - comet_debug() << "End of doLoweringIndexTreeToSCF()\n"; - comet_pdump(rootOp->getParentOfType()); - } + logger.resetIndent(); + rewriter.eraseOp(*op); + } - comet_debug() << "Cleaning up IndexTree Operations\n"; - comet_vdump(rootOp); - std::vector operations_dumpster; - rootOp.erase(); - for (auto itOp : wp_ops) - { - if (indexTree::IndexTreeComputeOp cur_op = dyn_cast(itOp.getDefiningOp())) - { - comet_pdump(itOp.getDefiningOp()->getOperand(0).getDefiningOp()); /// RHS - comet_pdump(itOp.getDefiningOp()->getOperand(1).getDefiningOp()); /// LHS - operations_dumpster.push_back(cur_op.getOperand(0).getDefiningOp()); - operations_dumpster.push_back(cur_op.getOperand(1).getDefiningOp()); + return success(); } - comet_pdump(itOp.getDefiningOp()); - itOp.getDefiningOp()->erase(); - } - for (auto op : operations_dumpster) - { - op->erase(); - } -#ifdef DEBUG_MODE_LowerIndexTreeToSCFPass - { - int opstree_i = 0; - for (auto n : opstree_vec) - { - comet_debug() << " " << n->id << "\n"; - comet_debug() << "opstree_vec[" << opstree_i << "] " - << "forOps.size():" << n->forOps.size() << " " - << "accessIdx.size():" << n->accessIdx.size() << "\n"; - /// << "cmptOps.size():" << n->cmptOps.size() << "\n"; - if (n->parent != nullptr) - { - comet_debug() << "parent: " << n->parent->id << "\n"; - } - else + void runOnOperation() override { + std::vector iTrees; + func::FuncOp funcOp = getOperation(); + funcOp.walk([&](IndexTreeOp op){ iTrees.push_back(op); }); + + for(auto op : iTrees) { - comet_debug() << "parent: null \n"; + OpBuilder builder(op); + IRRewriter rewriter(builder); + if(failed(convertTree(op, rewriter))){ + return signalPassFailure(); + } } - ++opstree_i; - } - } -#endif - /// ----------------- /// - /// Free the memory occupied by each element in opstree_vec. - /// ----------------- /// - cleanOpstreeVec(opstree_vec); -} /// End doLoweringIndexTreeToSCF() + TypeConverter typeConverter; + mlir::ConversionTarget target(getContext()); + target.addLegalDialect(); -void LowerIndexTreeToSCFPass::runOnOperation() -{ - comet_debug() << "LowerIndexTreeToSCFPass\n"; - func::FuncOp function = getOperation(); - auto module = function.getOperation()->getParentOfType(); - auto *ctx = &getContext(); - - /// Declare comet_sort_index() - declareSortFunc(module, - ctx, - function.getLoc()); - - std::vector iTreeRoots; - getIndexTreeOps(function, iTreeRoots /* output */); - for (auto root : iTreeRoots) - { - comet_vdump(root); - OpBuilder builder(root); - doLoweringIndexTreeToSCF(root, builder); - } + mlir::RewritePatternSet patterns(&getContext()); + if (mlir::failed(mlir::applyPartialConversion(getOperation(), target, std::move(patterns)))) + signalPassFailure(); + + } + + + }; } /// Lower sparse tensor algebra operation to loops diff --git a/lib/Conversion/IndexTreeToSCF/SymbolicDomainConversion.cpp b/lib/Conversion/IndexTreeToSCF/SymbolicDomainConversion.cpp new file mode 100644 index 00000000..bfcc379d --- /dev/null +++ b/lib/Conversion/IndexTreeToSCF/SymbolicDomainConversion.cpp @@ -0,0 +1,435 @@ +// +// Copyright 2022 Battelle Memorial Institute +// +// Redistribution and use in source and binary forms, with or without modification, +// are permitted provided that the following conditions are met: +// +// 1. Redistributions of source code must retain the above copyright notice, this list of conditions +// and the following disclaimer. +// +// 2. Redistributions in binary form must reproduce the above copyright notice, this list of conditions +// and the following disclaimer in the documentation and/or other materials provided with the distribution. +// +// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED +// WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR +// PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, +// INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE +// GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, +// WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +// OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. +// + +#include "comet/Dialect/IndexTree/IR/IndexTreeDialect.h" +#include "comet/Dialect/Utils/Utils.h" +#include "comet/Dialect/TensorAlgebra/IR/TADialect.h" +#include "comet/Conversion/IndexTreeToSCF/IndexTreeToSCF.h" +#include "comet/Dialect/IndexTree/Patterns.h" + +#include "mlir/Dialect/Arith/IR/Arith.h" +#include "mlir/Dialect/SCF/IR/SCF.h" +#include "mlir/Dialect/Index/IR/IndexAttrs.h" +#include "mlir/Dialect/Index/IR/IndexOps.h" +#include "mlir/Dialect/Index/IR/IndexDialect.h" +#include "mlir/Dialect/MemRef/IR/MemRef.h" +#include "mlir/Dialect/Bufferization/IR/Bufferization.h" +#include "mlir/Dialect/Func/IR/FuncOps.h" +#include "mlir/Dialect/Tensor/IR/Tensor.h" +#include "mlir/IR/BuiltinTypeInterfaces.h" +#include "mlir/IR/BuiltinTypes.h" +#include "mlir/IR/Value.h" +#include "mlir/IR/ValueRange.h" +#include "mlir/Transforms/DialectConversion.h" +#include "mlir/Dialect/Func/Transforms/DecomposeCallGraphTypes.h" +#include "mlir/Dialect/Func/Transforms/FuncConversions.h" +#include "mlir/Dialect/SCF/Transforms/Patterns.h" +#include "mlir/Pass/Pass.h" + +using namespace mlir; +using llvm::SmallVector; + +namespace mlir { + namespace comet{ + #define GEN_PASS_DEF_CONVERTSYMBOLICDOMAINS + #include "comet/Conversion/Passes.h.inc" + } +} + +struct SymbolicDomain { + Value pos_size; + Value pos_alloc_size; + Value crd_size; + Value dim_size; + Value pos; + Value mark_array; +}; + +static bool unpack_symbolic_domain(Value symbolic_domain, SymbolicDomain& result) +{ + if (auto cast = symbolic_domain.getDefiningOp()) { + result.pos_size = cast->getOperand(0); + result.pos_alloc_size = cast->getOperand(1); + result.crd_size = cast->getOperand(2); + result.dim_size = cast->getOperand(3); + result.pos = cast->getOperand(4); + result.mark_array = cast->getOperand(5); + return true; + } + return false; +} + +namespace { +struct ConvertDomainInsertOp + : public OpConversionPattern { + using OpConversionPattern::OpConversionPattern; + LogicalResult + matchAndRewrite(indexTree::SymbolicDomainInsertOp op, OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const final { + auto loc = op.getLoc(); + Type index_type = rewriter.getIndexType(); + SymbolicDomain domain; + if(!unpack_symbolic_domain(llvm::cast(adaptor).getDomain(), domain)){ + return failure(); + } + + Value one = rewriter.create(loc, index_type, rewriter.getIndexAttr(1)); + if(op.getIsUnique()) + { + // If we know the crd is unique, we can just increment the crd_size value + domain.crd_size = rewriter.create(loc, index_type, domain.crd_size, one); + } else + { + + Value mark = rewriter.create(loc, index_type, domain.pos_size, one); + ShapedType domMarkArrayT = cast(domain.mark_array.getType()); + if(mark.getType() != domMarkArrayT.getElementType()) + { + mark = rewriter.create(loc, domMarkArrayT.getElementType() ,mark); + } + Value mark_val = rewriter.create(loc, domain.mark_array, op.getCrd()); + Value is_marked = rewriter.create(loc, + rewriter.getI1Type(), + arith::CmpIPredicate::eq, + mark, + mark_val); + scf::IfOp if_op = rewriter.create(loc, index_type, is_marked, true); + // We have seen this crd before + rewriter.setInsertionPointToStart(if_op.thenBlock()); + rewriter.create(loc, domain.crd_size); + + // We haven't seen this crd before + rewriter.setInsertionPointToStart(if_op.elseBlock()); + rewriter.create(loc, mark, domain.mark_array, op.getCrd()); + Value new_crd_size = rewriter.create(loc, index_type, domain.crd_size, one); + rewriter.create(loc, new_crd_size); + rewriter.setInsertionPointAfter(if_op); + domain.crd_size = if_op.getResult(0); + } + + Value materialized = getTypeConverter()->materializeArgumentConversion( + rewriter, + op.getLoc(), + op.getDomain().getType(), + { + domain.pos_size, + domain.pos_alloc_size, + domain.crd_size, + domain.dim_size, + domain.pos, + domain.mark_array + } + ); + rewriter.replaceOp(op, {materialized}); + return success(); + } +}; + +struct ConvertDomainEndRowOp + : public OpConversionPattern { + using OpConversionPattern::OpConversionPattern; + LogicalResult + matchAndRewrite(indexTree::SymbolicDomainEndRowOp op, OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const final { + auto loc = op.getLoc(); + Type index_type = rewriter.getIndexType(); + SymbolicDomain domain; + if(!unpack_symbolic_domain(llvm::cast(adaptor).getDomain(), domain)){ + return failure(); + } + + Value inc = rewriter.create(loc, index_type, rewriter.getIndexAttr(1)); + Value new_pos_size = rewriter.create(loc, index_type, domain.pos_size, inc); + Value crd_size_cast = rewriter.createOrFold(loc, rewriter.getIntegerType(op.getResult().getType().getIndicesBitwidth()), domain.crd_size); + // TODO: Dynamically resize array? + rewriter.create(loc, crd_size_cast, domain.pos, new_pos_size); + Value materialized = getTypeConverter()->materializeArgumentConversion( + rewriter, + op.getLoc(), + op.getDomain().getType(), + { + new_pos_size, + domain.pos_alloc_size, + domain.crd_size, + domain.dim_size, + domain.pos, + domain.mark_array + } + ); + rewriter.replaceOp(op, {materialized}); + return success(); + } +}; + +struct ConvertDomainDeclarationOp + : public OpConversionPattern { + using OpConversionPattern::OpConversionPattern; + LogicalResult + matchAndRewrite(indexTree::DeclDomainOp op, OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const final { + auto loc = op.getLoc(); + Type indices_type = rewriter.getIntegerType(*op.getIndicesBitwidth()); + Type index_type = rewriter.getIndexType(); + Type memref_type = MemRefType::get({ShapedType::kDynamic,}, indices_type); + + Value zero = rewriter.create(loc, index_type, rewriter.getIndexAttr(0)); + Value inc = rewriter.create(loc, index_type, rewriter.getIndexAttr(1)); + Value pos_alloc_size = rewriter.create(loc, op.getNumRows(), inc); + TypedValue pos = rewriter.create(loc, memref_type, ValueRange{pos_alloc_size}, ValueRange(), nullptr); + Value zero_cast = zero; + zero_cast = rewriter.createOrFold(loc, pos.getType().getElementType(), zero); + + rewriter.create(loc, zero_cast, pos, zero); + Value mark_array = rewriter.create(loc, memref_type, ValueRange{op.getDimSize()}, ValueRange(), nullptr); + auto new_op = rewriter.create( + loc, + op->getResultTypes(), + ValueRange({zero, pos_alloc_size, zero, op.getDimSize(), pos, mark_array}) + ); + rewriter.replaceOp(op, new_op->getResults()); + return success(); + } +}; + +struct ConvertSparseTensorOp + : public OpConversionPattern { + using OpConversionPattern::OpConversionPattern; + LogicalResult + matchAndRewrite(indexTree::IndexTreeSparseTensorOp op, OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const final { + + + for(Value domain : llvm::cast(adaptor).getDomains()) + { + if(llvm::isa(domain.getType())){ + if(!domain.getDefiningOp()) + return failure(); + } + } + + tensorAlgebra::SparseTensorType spType = mlir::cast(op->getResultTypes()[0]); + auto ctx = op.getContext(); + auto format_unk = tensorAlgebra::TensorFormatEnumAttr::get(ctx, tensorAlgebra::TensorFormatEnum::UNK); + auto format_dense = tensorAlgebra::TensorFormatEnumAttr::get(ctx, tensorAlgebra::TensorFormatEnum::D); + auto format_compressed = tensorAlgebra::TensorFormatEnumAttr::get(ctx, tensorAlgebra::TensorFormatEnum::CU); + SmallVector dim_format; + uint32_t rank = 0; + + auto loc = op.getLoc(); + Type index_type = rewriter.getIndexType(); + Type memref_type = MemRefType::get({ShapedType::kDynamic,}, spType.getIndicesType()); + Value zero = rewriter.create(loc, index_type, rewriter.getIndexAttr(0)); + Value one = rewriter.create(loc, index_type, rewriter.getIndexAttr(1)); + Value nnz = one; + llvm::SmallVector dim_sizes, + pos_indices, + crd_indices, + tile_pos_indices, + tile_crd_indices; + + for(Value domain : llvm::cast(adaptor).getDomains()) + { + rank += 1; + if(llvm::isa(domain.getType())) + { + Operation* domain_op = domain.getDefiningOp(); + if(llvm::isa(domain_op)) + { + auto dense_domain_op = llvm::cast(domain_op); + Value dim_size = dense_domain_op.getDimSize(); // Always index type + Value dim_size_cast; + dim_size_cast = rewriter.createOrFold(loc, spType.getIndicesType(), dim_size); + + + Value pos = rewriter.create(loc, MemRefType::get({ShapedType::kDynamic,}, spType.getIndicesType()), ValueRange({rewriter.create(loc, 1).getResult()})); + rewriter.create(loc, dim_size_cast, pos, zero); + Value crd = rewriter.create(loc, MemRefType::get({0,}, spType.getIndicesType())); + Value pos_tile = rewriter.create(loc, MemRefType::get({0,}, spType.getIndicesType())); + Value crd_tile = rewriter.create(loc, MemRefType::get({0,}, spType.getIndicesType())); + + pos_indices.push_back(rewriter.create(loc, pos, rewriter.getUnitAttr(), rewriter.getUnitAttr())); + crd_indices.push_back(rewriter.create(loc, crd, rewriter.getUnitAttr(), rewriter.getUnitAttr())); + tile_pos_indices.push_back(rewriter.create(loc, pos_tile, rewriter.getUnitAttr(), rewriter.getUnitAttr())); + tile_crd_indices.push_back(rewriter.create(loc, crd_tile, rewriter.getUnitAttr(), rewriter.getUnitAttr())); + + dim_sizes.push_back(dim_size); + nnz = rewriter.create(loc, index_type, nnz, dim_size); + dim_format.push_back(format_dense); + dim_format.push_back(format_unk); + } else if(llvm::isa(domain_op)) + { + auto sparse_domain_op = llvm::cast(domain_op); + Value dim_size = sparse_domain_op.getDimSize(); + Value crd_size = sparse_domain_op.getCrdSize(); + + Value pos = sparse_domain_op.getPos(); + Value crd = sparse_domain_op.getCrd(); + Value pos_tile = rewriter.create(loc, MemRefType::get({0,}, spType.getIndicesType())); + Value crd_tile = rewriter.create(loc, MemRefType::get({0,}, spType.getIndicesType())); + + pos_indices.push_back(pos); + crd_indices.push_back(crd); + + tile_pos_indices.push_back(rewriter.create(loc, pos_tile, rewriter.getUnitAttr(), rewriter.getUnitAttr())); + tile_crd_indices.push_back(rewriter.create(loc, crd_tile, rewriter.getUnitAttr(), rewriter.getUnitAttr())); + + dim_sizes.push_back(dim_size); + nnz = crd_size; + dim_format.push_back(format_compressed); + dim_format.push_back(format_unk); + } else + return failure(); + } else if(llvm::isa(domain.getType())) + { + SymbolicDomain domain_struct; + assert(unpack_symbolic_domain(domain, domain_struct)); + Value crd = rewriter.create(loc, memref_type, ValueRange{domain_struct.crd_size}, ValueRange(), nullptr); + Value pos_tile = rewriter.create(loc, MemRefType::get({0,}, spType.getIndicesType())); + Value crd_tile = rewriter.create(loc, MemRefType::get({0,}, spType.getIndicesType())); + + pos_indices.push_back(rewriter.create(loc, domain_struct.pos, rewriter.getUnitAttr(), rewriter.getUnitAttr())); + crd_indices.push_back(rewriter.create(loc, crd, rewriter.getUnitAttr(), rewriter.getUnitAttr())); + tile_pos_indices.push_back(rewriter.create(loc, pos_tile, rewriter.getUnitAttr(), rewriter.getUnitAttr())); + tile_crd_indices.push_back(rewriter.create(loc, crd_tile, rewriter.getUnitAttr(), rewriter.getUnitAttr())); + + dim_sizes.push_back(domain_struct.dim_size); + nnz = domain_struct.crd_size; + dim_format.push_back(format_compressed); + dim_format.push_back(format_unk); + } + } + + //Allocate values array and initialize + Type float_type = llvm::cast(op.getResult().getType()).getElementType(); + Value val_array = rewriter.create(loc, MemRefType::get({ShapedType::kDynamic,}, float_type), ValueRange{nnz}, ValueRange(), nullptr); + Value float_zero = rewriter.create(loc, float_type, rewriter.getFloatAttr(float_type, 0.0)); + auto for_loop = rewriter.create(loc, zero, nnz, one); + rewriter.setInsertionPointToStart(for_loop.getBody()); + auto induction_var = for_loop.getInductionVar(); + rewriter.create(loc, float_zero, val_array, induction_var); + rewriter.setInsertionPointAfter(for_loop); + val_array = rewriter.create(loc, val_array, rewriter.getUnitAttr(), rewriter.getUnitAttr()); + + std::vector args; + args.push_back(val_array); + args.push_back(nnz); + args.insert(args.end(), dim_sizes.begin(), dim_sizes.end()); + Value dims = rewriter.create(loc, ValueRange(dim_sizes)); + + rewriter.replaceOpWithNewOp(op, op.getResult().getType(), dims, pos_indices, crd_indices, tile_pos_indices, tile_crd_indices, val_array, rank, rewriter.getArrayAttr(dim_format)); + return success(); + } +}; + +class EraseDenseDomainOp : public OpConversionPattern{ + public: + using OpConversionPattern::OpConversionPattern; + + mlir::LogicalResult + matchAndRewrite(indexTree::IndexTreeDenseDomainOp op, OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const override { + rewriter.eraseOp(op); + return success(); + } +}; + +class EraseSparseDomainOp : public OpConversionPattern{ + public: + using OpConversionPattern::OpConversionPattern; + + mlir::LogicalResult + matchAndRewrite(indexTree::IndexTreeSparseDomainOp op, OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const override { + rewriter.eraseOp(op); + return success(); + } +}; + +} //namespace + +struct ConvertSymbolicDomainsPass + : public PassWrapper> +{ + MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(ConvertSymbolicDomainsPass) + + void runOnOperation() override + { + // Convert the rest of the index tree dialect to SCF + TypeConverter typeConverter; + typeConverter.addConversion([](Type type) { return type; }); + typeConverter.addConversion( + [](indexTree::SymbolicDomainType domainType, SmallVectorImpl &types) { + auto context = domainType.getContext(); + IntegerType indicesType = IntegerType::get(context, domainType.getIndicesBitwidth()); + Type index_type = IndexType::get(context); + Type memref_type = MemRefType::get({ShapedType::kDynamic,}, indicesType); + types.push_back(index_type); + types.push_back(index_type); + types.push_back(index_type); + types.push_back(index_type); + types.push_back(memref_type); + types.push_back(memref_type); + return success(); + }); + + typeConverter.addSourceMaterialization( + [](OpBuilder &builder, indexTree::SymbolicDomainType resultType, ValueRange inputs, + Location loc) -> std::optional { + assert(inputs.size() == 6); + Value value = builder.create(loc, resultType, inputs)->getResult(0); + return value; + }); + + typeConverter.addArgumentMaterialization( + [](OpBuilder &builder, indexTree::SymbolicDomainType resultType, ValueRange inputs, + Location loc) -> std::optional { + assert(inputs.size() == 6); + Value value = builder.create(loc, resultType, inputs)->getResult(0); + return value; + }); + + mlir::ConversionTarget target(getContext()); + target.addLegalDialect(); + target.addLegalDialect(); + target.addLegalOp(); + target.addIllegalOp(); + target.addIllegalOp(); + + + + mlir::RewritePatternSet patterns(&getContext()); + populateFunctionOpInterfaceTypeConversionPattern(patterns, typeConverter); + populateCallOpTypeConversionPattern(patterns, typeConverter); + scf::populateSCFStructuralTypeConversionsAndLegality(typeConverter, patterns, target); + indexTree::populateIndexTreeTypeConversionPatterns(&getContext(), patterns, typeConverter, target); + patterns.add(typeConverter, &getContext()); + patterns.add(typeConverter, &getContext()); + + if (mlir::failed(mlir::applyPartialConversion(getOperation(), target, std::move(patterns)))) + signalPassFailure(); + } +}; + +/// Lower sparse tensor algebra operation to loops +std::unique_ptr mlir::comet::createConvertSymbolicDomainsPass() +{ + return std::make_unique(); +} \ No newline at end of file diff --git a/lib/Conversion/ParallelLoopsToGpu/CMakeLists.txt b/lib/Conversion/ParallelLoopsToGpu/CMakeLists.txt deleted file mode 100644 index 121c00ae..00000000 --- a/lib/Conversion/ParallelLoopsToGpu/CMakeLists.txt +++ /dev/null @@ -1,9 +0,0 @@ -add_llvm_library(COMETParallelLoopsToGpu - ParallelLoopsToGpu.cpp - - DEPENDS - ParallelLoopsConversionPassIncGen - - # LINK_LIBS - # MLIRPASS -) \ No newline at end of file diff --git a/lib/Conversion/ParallelLoopsToGpuFPGA/CMakeLists.txt b/lib/Conversion/ParallelLoopsToGpuFPGA/CMakeLists.txt new file mode 100644 index 00000000..5efb4466 --- /dev/null +++ b/lib/Conversion/ParallelLoopsToGpuFPGA/CMakeLists.txt @@ -0,0 +1,9 @@ +add_llvm_library(COMETParallelLoopsToGpuFPGA + ParallelLoopsToGpuFPGA.cpp + + DEPENDS + ParallelLoopsToGpuFPGAConversionPassIncGen + + # LINK_LIBS + # MLIRPASS +) \ No newline at end of file diff --git a/lib/Conversion/ParallelLoopsToGpu/ParallelLoopsToGpu.cpp b/lib/Conversion/ParallelLoopsToGpuFPGA/ParallelLoopsToGpuFPGA.cpp similarity index 64% rename from lib/Conversion/ParallelLoopsToGpu/ParallelLoopsToGpu.cpp rename to lib/Conversion/ParallelLoopsToGpuFPGA/ParallelLoopsToGpuFPGA.cpp index 76392358..0e80786a 100644 --- a/lib/Conversion/ParallelLoopsToGpu/ParallelLoopsToGpu.cpp +++ b/lib/Conversion/ParallelLoopsToGpuFPGA/ParallelLoopsToGpuFPGA.cpp @@ -1,9 +1,30 @@ +// +// Copyright 2022 Battelle Memorial Institute +// +// Redistribution and use in source and binary forms, with or without modification, +// are permitted provided that the following conditions are met: +// +// 1. Redistributions of source code must retain the above copyright notice, this list of conditions +// and the following disclaimer. +// +// 2. Redistributions in binary form must reproduce the above copyright notice, this list of conditions +// and the following disclaimer in the documentation and/or other materials provided with the distribution. +// +// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED +// WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR +// PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, +// INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE +// GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, +// WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +// OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. +// #include #include #include "mlir/Dialect/Arith/IR/Arith.h" #include "mlir/IR/Builders.h" #include "mlir/IR/BuiltinAttributes.h" +#include "mlir/IR/BuiltinTypes.h" #include "mlir/Pass/Pass.h" #include "mlir/Dialect/DLTI/DLTI.h" #include "mlir/Dialect/MemRef/IR/MemRef.h" @@ -16,13 +37,15 @@ #include "mlir/Transforms/Passes.h" -#include "comet/Conversion/ParallelLoopsToGpu/ParallelLoopsToGpu.h" -#include "comet/Conversion/ParallelLoopsToGpu/Passes.h" +#include "comet/Dialect/Utils/Utils.h" +#include "comet/Conversion/ParallelLoopsToGpuFPGA/ParallelLoopsToGpuFPGA.h" +#include "comet/Conversion/ParallelLoopsToGpuFPGA/Passes.h" #include "mlir/Dialect/Func/IR/FuncOps.h" +#include "llvm/Support/Casting.h" #define GEN_PASS_CLASSES -#include "comet/Conversion/ParallelLoopsToGpu/Passes.h.inc" +#include "comet/Conversion/ParallelLoopsToGpuFPGA/Passes.h.inc" // *********** For debug purpose *********// // #define COMET_DEBUG_MODE @@ -68,7 +91,7 @@ mlir::Value reassociateIndices(T op, mlir::OpBuilder& builder) } } - if(op.getIndices()[i].template isa()) + if(mlir::isa(op.getIndices()[i])) { dimVals.push_back(op.getIndices()[i]); expr = expr * mlir::getAffineDimExpr(currDimPos++, op->getContext()); @@ -79,7 +102,7 @@ mlir::Value reassociateIndices(T op, mlir::OpBuilder& builder) expr = expr * mlir::getAffineSymbolExpr(currSymPos++, op->getContext()); } - if(newIndex.isa()) + if(mlir::isa(newIndex)) { dimVals.push_back(newIndex); expr = mlir::getAffineDimExpr(currDimPos++, op->getContext()) + expr; @@ -104,7 +127,7 @@ mlir::Value reassociateIndices(T op, mlir::OpBuilder& builder) void collapseMemrefAndUsers(mlir::Value val, mlir::OpBuilder& builder) { - auto memref = val.getType().cast(); + auto memref = mlir::cast(val.getType()); if (memref.getRank() == 1) { return; @@ -150,7 +173,7 @@ bool contains_arg(mlir::Block& block, mlir::BlockArgument arg) { for(auto index: store_op.getIndices()) { - if (auto affine_expr = llvm::dyn_cast_or_null(index.getDefiningOp())) + if (auto affine_expr = mlir::dyn_cast_if_present(index.getDefiningOp())) { for(auto op: affine_expr.getOperands()) { @@ -161,7 +184,7 @@ bool contains_arg(mlir::Block& block, mlir::BlockArgument arg) } } } - else if(auto block_arg = index.dyn_cast_or_null()) + else if(auto block_arg = mlir::dyn_cast(index)) { if(block_arg == arg) { @@ -180,8 +203,8 @@ bool contains_arg(mlir::Block& block, mlir::BlockArgument arg) } else { llvm::errs() << "Load operation without affine expression\n"; - index.dump(); - store_op->dump(); + // index.dump(); + // store_op->dump(); exit(1); } } @@ -219,12 +242,20 @@ bool is_reduction(mlir::scf::ForOp forOp) return is_reduction_(forOp.getBodyRegion(), forOp.getBody()->getArgument(0)); } -class ParallelOpToGpu: public mlir::OpConversionPattern { +mlir::Operation* CeilDivUIOp(mlir::ConversionPatternRewriter &rewriter, mlir::Location loc, mlir::Value lhs, mlir::Value rhs) +{ + auto c1 = rewriter.create(loc, 1); + auto temp = rewriter.create(loc, lhs, rewriter.create(loc, rhs, c1)); + return rewriter.create(loc, temp, rhs); +} + +class ParallelOpToGpuFPGA: public mlir::OpConversionPattern { private: - int blockX, blockY, blockR; +[[maybe_unused]] int blockX, blockY, blockR; + mlir::tensorAlgebra::TargetDevice target; public: using mlir::OpConversionPattern::OpConversionPattern; - ParallelOpToGpu(mlir::MLIRContext* ctx, int blockX, int blockY, int blockR) : mlir::OpConversionPattern(ctx), blockX(blockX), blockY(blockY), blockR(blockR) {} + ParallelOpToGpuFPGA(mlir::MLIRContext* ctx, int blockX, int blockY, int blockR, mlir::tensorAlgebra::TargetDevice target) : mlir::OpConversionPattern(ctx), blockX(blockX), blockY(blockY), blockR(blockR), target(target) {} mlir::LogicalResult matchAndRewrite(mlir::scf::ParallelOp parOp, OpAdaptor adaptor, mlir::ConversionPatternRewriter &rewriter) const override { @@ -248,8 +279,8 @@ class ParallelOpToGpu: public mlir::OpConversionPattern { auto block_size_y = rewriter.create(parOp->getLoc(), blockY ); auto block_size_x = rewriter.create(parOp->getLoc(), blockX ); auto c1 = rewriter.create(parOp->getLoc(), 1); - auto upperBound0 = rewriter.create(parOp->getLoc(), parOp.getUpperBound().front(), block_size_y); - auto upperBound1 = rewriter.create(parOp->getLoc(), parOp.getUpperBound().back(), block_size_x); + auto upperBound0 = CeilDivUIOp(rewriter, parOp->getLoc(), parOp.getUpperBound().front(), block_size_y); + auto upperBound1 = CeilDivUIOp(rewriter, parOp->getLoc(), parOp.getUpperBound().back(), block_size_x); comet_debug() << upperBound0; auto y_loop_grid = rewriter.create(parOp->getLoc(), parOp.getLowerBound().front(), upperBound0->getResult(0), c1->getResult(0)); @@ -274,13 +305,17 @@ class ParallelOpToGpu: public mlir::OpConversionPattern { auto affineIndex = mlir::AffineMap::get(1, 2, {res}, parOp->getContext()); comet_debug() << affineIndex; std::vector range = { y_loop_grid.getBody()->getArgument(0), block_size_y->getResult(0), y_loop_block.getBody()->getArgument(0)}; - auto newIndexY = rewriter.create(parOp->getLoc(), rewriter.create(parOp->getLoc(), affineIndex, range), parOp.getUpperBound().front()); - newIndexY->setAttr("GuardY", rewriter.getUnitAttr()); - // auto newIndexY = rewriter.create(forOp->getLoc(), affineIndex, range); - // auto newIndexY = rewriter.create(forOp->getLoc(),rewriter.create(forOp->getLoc(), y_loop_grid.getBody()->getArgument(0), block_size_y), y_loop_block.getBody()->getArgument(0)); - // rewriter.setInsertionPoint(newIndexY); - rewriter.replaceAllUsesWith(parOp.getBody()->getArgument(0), newIndexY); - + mlir::Operation* newIndexY = rewriter.create(parOp->getLoc(), affineIndex, range); + if(target == mlir::tensorAlgebra::TargetDevice::GPU) + { + newIndexY = rewriter.create(parOp->getLoc(), newIndexY->getResult(0), parOp.getUpperBound().front()); + newIndexY->setAttr("GuardY", rewriter.getUnitAttr()); + // auto newIndexY = rewriter.create(forOp->getLoc(), affineIndex, range); + // auto newIndexY = rewriter.create(forOp->getLoc(),rewriter.create(forOp->getLoc(), y_loop_grid.getBody()->getArgument(0), block_size_y), y_loop_block.getBody()->getArgument(0)); + // rewriter.setInsertionPoint(newIndexY); + } + + rewriter.replaceAllUsesWith(parOp.getBody()->getArgument(0), newIndexY->getResult(0)); // auto upperBound0 = rewriter.create(parOp->getLoc(), parOp.getUpperBound(), block_size_y); // auto y_loop_grid = rewriter.create(parOp->getLoc(), parOp.getLowerBound(), upperBound0->getResult(0), c1->getResult(0)); @@ -301,25 +336,42 @@ class ParallelOpToGpu: public mlir::OpConversionPattern { affineIndex = mlir::AffineMap::get(1, 2, {res}, parOp->getContext()); range = { x_loop_grid.getBody()->getArgument(0), block_size_x->getResult(0), x_loop_block.getBody()->getArgument(0)}; // auto newIndexX = rewriter.create(parOp->getLoc(), affineIndex, range); - auto newIndexX = rewriter.create(parOp->getLoc(), rewriter.create(parOp->getLoc(), affineIndex, range), parOp.getUpperBound().back()); - newIndexX->setAttr("GuardX", rewriter.getUnitAttr()); - + mlir::Operation* newIndexX = rewriter.create(parOp->getLoc(), affineIndex, range); + if(target == mlir::tensorAlgebra::TargetDevice::GPU) + { + newIndexX = rewriter.create(parOp->getLoc(), newIndexX->getResult(0), parOp.getUpperBound().back()); + newIndexX->setAttr("GuardX", rewriter.getUnitAttr()); + } + + rewriter.replaceAllUsesWith(parOp.getBody()->getArgument(1), newIndexX->getResult(0)); // auto newIndexX = rewriter.create(parOp->getLoc(),rewriter.create(parOp->getLoc(), x_loop_grid.getBody()->getArgument(0), block_size_x), x_loop_block.getBody()->getArgument(0)); - rewriter.setInsertionPoint(newIndexX); - rewriter.replaceAllUsesWith(parOp.getBody()->getArgument(1), newIndexX); - + rewriter.setInsertionPointAfter(newIndexX); rewriter.eraseOp(parOp.getBody()->getTerminator()); - rewriter.inlineBlockBefore(parOp.getBody(), x_loop_block.getBody()->getTerminator(), {newIndexY->getResult(0), newIndexX->getResult(0)}); + if(target == mlir::tensorAlgebra::TargetDevice::GPU) + { + rewriter.inlineBlockBefore(parOp.getBody(), x_loop_block.getBody()->getTerminator(), {newIndexY->getResult(0), newIndexX->getResult(0)}); + } + else + { + auto withinY = rewriter.create(parOp.getLoc(), mlir::arith::CmpIPredicate::slt, newIndexY->getResult(0), parOp.getUpperBound().front()); + auto withinX = rewriter.create(parOp.getLoc(), mlir::arith::CmpIPredicate::slt, newIndexX->getResult(0), parOp.getUpperBound().back()); + auto withinXandY = rewriter.create(parOp.getLoc(), withinY, withinX); + + auto ifOp = rewriter.create(parOp->getLoc(), withinXandY); + rewriter.inlineBlockBefore(parOp.getBody(), ifOp.getBody()->getTerminator(), {newIndexY->getResult(0), newIndexX->getResult(0)}); + } + + rewriter.eraseOp(parOp); return mlir::success(); } - else if (!mlir::isa(parOp->getParentOp())) // Y level loop + else if (!mlir::isa(parOp->getParentOp()) && !(parOp->getParentOp() && parOp->getParentOp()->hasAttrOfType("GuardY"))) // Y level loop { // auto block_size_x = rewriter.create(forOp->getLoc(), rewriter.getIndexType() , rewriter.getIndexAttr(blockX) ); auto block_size_y = rewriter.create(parOp->getLoc(), blockY ); auto c1 = rewriter.create(parOp->getLoc(), 1); - auto upperBound0 = rewriter.create(parOp->getLoc(), parOp.getUpperBound().front(), block_size_y); + auto upperBound0 = CeilDivUIOp(rewriter,parOp->getLoc(), parOp.getUpperBound().front(), block_size_y); comet_debug() << upperBound0; // auto upperBound1 = rewriter.create(forOp->getLoc(), forOp.getUpperBound(), block_size_x); auto y_loop_grid = rewriter.create(parOp->getLoc(), parOp.getLowerBound().front(), upperBound0->getResult(0), c1->getResult(0)); @@ -344,33 +396,61 @@ class ParallelOpToGpu: public mlir::OpConversionPattern { auto affineIndex = mlir::AffineMap::get(1, 2, {res}, parOp->getContext()); comet_debug() << affineIndex; std::vector range = { y_loop_grid.getBody()->getArgument(0), block_size_y->getResult(0), y_loop_block.getBody()->getArgument(0)}; - auto newIndexY = rewriter.create(parOp->getLoc(), rewriter.create(parOp->getLoc(), affineIndex, range), parOp.getUpperBound().front()); - newIndexY->setAttr("GuardY", rewriter.getUnitAttr()); + mlir::Operation* newIndexY = rewriter.create(parOp->getLoc(), affineIndex, range); + if(target == mlir::tensorAlgebra::TargetDevice::GPU) + { + newIndexY = rewriter.create(parOp->getLoc(), newIndexY->getResult(0), parOp.getUpperBound().front()); + newIndexY->setAttr("GuardY", rewriter.getUnitAttr()); + } // auto newIndexY = rewriter.create(forOp->getLoc(), affineIndex, range); // auto newIndexY = rewriter.create(forOp->getLoc(),rewriter.create(forOp->getLoc(), y_loop_grid.getBody()->getArgument(0), block_size_y), y_loop_block.getBody()->getArgument(0)); - rewriter.setInsertionPoint(newIndexY); + rewriter.setInsertionPointAfter(newIndexY); - rewriter.replaceAllUsesWith(parOp.getBody()->getArgument(0), newIndexY); + rewriter.replaceAllUsesWith(parOp.getBody()->getArgument(0), newIndexY->getResult(0)); rewriter.eraseOp(parOp.getBody()->getTerminator()); - rewriter.inlineBlockBefore(parOp.getBody(), y_loop_block.getBody()->getTerminator(), newIndexY->getResult(0)); + if(target == mlir::tensorAlgebra::TargetDevice::GPU) + { + rewriter.inlineBlockBefore(parOp.getBody(), y_loop_block.getBody()->getTerminator(), newIndexY->getResult(0)); + } + else + { + auto withinY = rewriter.create(parOp.getLoc(), mlir::arith::CmpIPredicate::slt, newIndexY->getResult(0), parOp.getUpperBound().front()); + auto ifOp = rewriter.create(parOp->getLoc(), withinY); + ifOp->setAttr("GuardY", rewriter.getUnitAttr()); + rewriter.inlineBlockBefore(parOp.getBody(), ifOp.getBody()->getTerminator(), {newIndexY->getResult(0)}); + } rewriter.eraseOp(parOp); return mlir::success(); } - else if ((mlir::isa(parOp->getParentOp()) && mlir::cast(parOp->getParentOp())->getAttrOfType("parallelDim").getValue().equals("dimY_block"))) // X level loop + else if ((parOp->getParentOp() && parOp->getParentOp()->hasAttrOfType("GuardY")) || (mlir::isa(parOp->getParentOp()) && mlir::cast(parOp->getParentOp())->getAttrOfType("parallelDim").getValue().compare("dimY_block") == 0)) // X level loop { + mlir::Value lower_bound; + mlir::Value upper_bound; + if(parOp->getParentOp()->hasAttrOfType("GuardY")) + { + rewriter.setInsertionPoint(parOp->getParentOp()); + lower_bound = parOp.getLowerBound().front().getDefiningOp() ? rewriter.clone(*parOp.getLowerBound().front().getDefiningOp())->getResult(0) : parOp.getLowerBound().front(); + upper_bound = parOp.getUpperBound().front().getDefiningOp() ? rewriter.clone(*parOp.getUpperBound().front().getDefiningOp())->getResult(0) : parOp.getUpperBound().front(); + } + else + { + lower_bound = parOp.getLowerBound().front(); + upper_bound = parOp.getUpperBound().front(); + } auto block_size_x = rewriter.create(parOp->getLoc(), blockX ); auto c1 = rewriter.create(parOp->getLoc(), 1); // auto upperBound0 = rewriter.create(parOp->getLoc(), parOp.getUpperBound(), block_size_y); - auto upperBound1 = rewriter.create(parOp->getLoc(), parOp.getUpperBound().front(), block_size_x); + + auto upperBound1 = CeilDivUIOp(rewriter, parOp->getLoc(), upper_bound, block_size_x); // auto y_loop_grid = rewriter.create(parOp->getLoc(), parOp.getLowerBound(), upperBound0->getResult(0), c1->getResult(0)); - auto x_loop_grid = rewriter.create(parOp->getLoc(), parOp.getLowerBound().front(), upperBound1->getResult(0), c1->getResult(0)); + auto x_loop_grid = rewriter.create(parOp->getLoc(), lower_bound, upperBound1->getResult(0), c1->getResult(0)); x_loop_grid->setAttr("parallelDim", rewriter.getAttr("dimX_grid")); newAttr = mlir::gpu::ParallelLoopDimMappingAttr::get(rewriter.getContext(), ::mlir::gpu::Processor::BlockX, map, map); x_loop_grid->setAttr("mapping", mlir::ArrayAttr::get(parOp->getContext(), newAttr) ); rewriter.setInsertionPointToStart(x_loop_grid.getBody()); // auto y_loop_block = rewriter.create(parOp->getLoc(), parOp.getLowerBound(), block_size_y->getResult(0), c1->getResult(0)); - auto x_loop_block = rewriter.create(parOp->getLoc(), parOp.getLowerBound().front(), block_size_x->getResult(0), c1->getResult(0)); + auto x_loop_block = rewriter.create(parOp->getLoc(), lower_bound, block_size_x->getResult(0), c1->getResult(0)); x_loop_block->setAttr("parallelDim", rewriter.getAttr("dimX_block")); newAttr = mlir::gpu::ParallelLoopDimMappingAttr::get(rewriter.getContext(), ::mlir::gpu::Processor::ThreadX, map, map); x_loop_block->setAttr("mapping", mlir::ArrayAttr::get(parOp->getContext(), newAttr) ); @@ -380,16 +460,45 @@ class ParallelOpToGpu: public mlir::OpConversionPattern { auto affineIndex = mlir::AffineMap::get(1, 2, {res}, parOp->getContext()); std::vector range = { x_loop_grid.getBody()->getArgument(0), block_size_x->getResult(0), x_loop_block.getBody()->getArgument(0)}; // auto newIndexX = rewriter.create(parOp->getLoc(), affineIndex, range); - auto newIndexX = rewriter.create(parOp->getLoc(), rewriter.create(parOp->getLoc(), affineIndex, range), parOp.getUpperBound().front()); - newIndexX->setAttr("GuardX", rewriter.getUnitAttr()); - + mlir::Operation* newIndexX = rewriter.create(parOp->getLoc(), affineIndex, range); + if(target == mlir::tensorAlgebra::TargetDevice::GPU) + { + newIndexX = rewriter.create(parOp->getLoc(), newIndexX->getResult(0), parOp.getUpperBound().front()); + newIndexX->setAttr("GuardX", rewriter.getUnitAttr()); + } + // auto newIndexX = rewriter.create(parOp->getLoc(),rewriter.create(parOp->getLoc(), x_loop_grid.getBody()->getArgument(0), block_size_x), x_loop_block.getBody()->getArgument(0)); - rewriter.setInsertionPoint(newIndexX); - rewriter.replaceAllUsesWith(parOp.getBody()->getArgument(0), newIndexX); - + // rewriter.setInsertionPoint(newIndexX); + rewriter.replaceAllUsesWith(parOp.getBody()->getArgument(0), newIndexX->getResult(0)); + rewriter.eraseOp(parOp.getBody()->getTerminator()); - rewriter.inlineBlockBefore(parOp.getBody(), x_loop_block.getBody()->getTerminator(), newIndexX->getResult(0)); + // parOp->getParentOfType()->dump(); + if(target == mlir::tensorAlgebra::TargetDevice::GPU) + { + rewriter.inlineBlockBefore(parOp.getBody(), x_loop_block.getBody()->getTerminator(), newIndexX->getResult(0)); + } + else + { + auto withinX = rewriter.create(parOp.getLoc(), mlir::arith::CmpIPredicate::slt, newIndexX->getResult(0), upper_bound); + if(auto par_if = mlir::dyn_cast(parOp->getParentOp())) + { + // auto newBlock = rewriter.splitBlock(par_if.getBody(), parOp->getIterator()); + mlir::Value condition = par_if.getCondition(); + auto new_Y_if = rewriter.create(parOp->getLoc(), condition); + // parOp->getParentOfType()->dump(); + + rewriter.eraseOp(par_if.getBody()->getTerminator()); + + rewriter.inlineBlockBefore(par_if.getBody(), new_Y_if); + rewriter.eraseOp(par_if); + rewriter.setInsertionPoint(new_Y_if.getBody()->getTerminator()); + } + auto ifOp = rewriter.create(parOp->getLoc(), withinX); + rewriter.inlineBlockBefore(parOp.getBody(), ifOp.getBody()->getTerminator(), {newIndexX->getResult(0)}); + // parOp->getParentOfType()->dump(); + + } rewriter.eraseOp(parOp); return mlir::success(); @@ -406,7 +515,7 @@ struct DetectReduction : public mlir::OpConversionPattern { DetectReduction(mlir::MLIRContext* ctx, int blockX, int blockY, int blockR) : mlir::OpConversionPattern(ctx), blockX(blockX), blockY(blockY), blockR(blockR) {} private: - int blockX, blockY, blockR; + [[maybe_unused]] int blockX, blockY, blockR; mlir::LogicalResult matchAndRewrite(mlir::scf::ForOp forOp, OpAdaptor adaptor, mlir::ConversionPatternRewriter &rewriter) const override { @@ -426,7 +535,7 @@ struct DetectReduction auto c1 = rewriter.create(forOp->getLoc(), 1); // auto upperBound0 = rewriter.create(forOp->getLoc(), forOp.getUpperBound(), block_size_y); auto outer_lower_bound = rewriter.create(forOp->getLoc(), 0); - auto outer_upper_bound = rewriter.create(forOp->getLoc(), rewriter.create(forOp->getLoc(), forOp.getUpperBound(), forOp.getLowerBound()), block_size_r); + auto outer_upper_bound = CeilDivUIOp(rewriter,forOp->getLoc(), rewriter.create(forOp->getLoc(), forOp.getUpperBound(), forOp.getLowerBound()), block_size_r); // auto upperBound1 = rewriter.create(forOp->getLoc(), forOp.getUpperBound(), block_size_r); // auto y_loop_grid = rewriter.create(forOp->getLoc(), forOp.getLowerBound(), upperBound0->getResult(0), c1->getResult(0)); @@ -464,13 +573,14 @@ struct DetectReduction } }; -class ConvertParallelLoopsToGpu: public CometParallelLoopsToGpuBase { +class ConvertParallelLoopsToGpuFPGA: public CometParallelLoopsToGpuFPGABase { public: - ConvertParallelLoopsToGpu() = default; - ConvertParallelLoopsToGpu(int blockX, int blockY, int blockR) { + ConvertParallelLoopsToGpuFPGA() = default; + ConvertParallelLoopsToGpuFPGA(int blockX, int blockY, int blockR, mlir::tensorAlgebra::TargetDevice target_device) { this->blockX = blockX; this->blockY = blockY; this->blockR = blockR; + this->target_device = target_device; } void runOnOperation() override { @@ -482,27 +592,49 @@ class ConvertParallelLoopsToGpu: public CometParallelLoopsToGpuBase()) + for(auto arg: funcOp.getArguments()) { - builder.setInsertionPointToStart(&funcOp.getBody().getBlocks().front()); - collapseMemrefAndUsers(arg, builder); + if(mlir::isa(arg.getType())) + { + builder.setInsertionPointToStart(&funcOp.getBody().getBlocks().front()); + collapseMemrefAndUsers(arg, builder); + } + } + + // /// Next, memrefs from allocations + auto memref_allocs = funcOp.getOps(); + for(auto memref: memref_allocs) + { + builder.setInsertionPointAfter(memref); + collapseMemrefAndUsers(memref, builder); } - } + } - /// Next, memrefs from allocations - auto memref_allocs = funcOp.getOps(); - for(auto memref: memref_allocs) + mlir::SmallVector forAllLoops; + funcOp->walk([&forAllLoops](mlir::scf::ForallOp forAllOp){forAllLoops.push_back(forAllOp);}); + + for(auto forAllOp: forAllLoops) { - builder.setInsertionPointAfter(memref); - collapseMemrefAndUsers(memref, builder); + builder.setInsertionPoint(forAllOp); + mlir::SmallVector lbs = forAllOp.getLowerBound(builder); + mlir::SmallVector ubs = forAllOp.getUpperBound(builder); + mlir::SmallVector steps = forAllOp.getStep(builder); + auto parallelOp = builder.create(forAllOp->getLoc(), lbs, ubs, steps); + // parallelOp.getRegion().front().erase(); + parallelOp.getRegion().takeBody(forAllOp.getRegion()); + builder.setInsertionPointToEnd(¶llelOp.getRegion().front()); + parallelOp.getRegion().front().getTerminator()->replaceAllUsesWith(builder.create(parallelOp->getLoc())); + parallelOp.getRegion().front().getTerminator()->erase(); + forAllOp.replaceAllUsesWith(parallelOp); + forAllOp->erase(); } mlir::RewritePatternSet patterns(context); - patterns.insert(context, blockX, blockY, blockR); + patterns.insert(context, blockX, blockY, blockR, this->target_device); mlir::ConversionTarget target(*context); target.addLegalDialect(); @@ -544,40 +676,43 @@ class ConvertParallelLoopsToGpu: public CometParallelLoopsToGpuBase(); + target2.addLegalDialect(); - target2.addLegalOp(); - patterns2.insert(context, blockX, blockY, blockR); - target2.addDynamicallyLegalOp([](mlir::scf::ForOp op) -> bool { - mlir::scf::ParallelOp parent = llvm::dyn_cast_or_null(op->getParentOp()); - if(parent && !op->hasAttr("reduceDim")) - { - return false; - } - else + target2.addLegalOp(); + patterns2.insert(context, blockX, blockY, blockR); + target2.addDynamicallyLegalOp([](mlir::scf::ForOp op) -> bool { + mlir::scf::ParallelOp parent = llvm::dyn_cast_or_null(op->getParentOp()); + if(parent && !op->hasAttr("reduceDim")) + { + return false; + } + else + { + return true; + } + }); + + if (mlir::failed(mlir::applyPartialConversion(funcOp, target2, std::move(patterns2)))) { - return true; + signalPassFailure(); } - }); - - if (mlir::failed(mlir::applyPartialConversion(funcOp, target2, std::move(patterns2)))) - { - signalPassFailure(); } } }; } -std::unique_ptr> mlir::comet::createConvertParallelLoopsToGpuPass() { - return std::make_unique(); +std::unique_ptr> mlir::comet::createConvertParallelLoopsToGpuFPGAPass() { + return std::make_unique(); } -std::unique_ptr> mlir::comet::createConvertParallelLoopsToGpuPass(int blockX, int blockY, int blockR) { - return std::make_unique(blockX, blockY, blockR); +std::unique_ptr> mlir::comet::createConvertParallelLoopsToGpuFPGAPass(int blockX, int blockY, int blockR, mlir::tensorAlgebra::TargetDevice target_device) { + return std::make_unique(blockX, blockY, blockR, target_device); } diff --git a/lib/Conversion/PrepareGpuHost/CMakeLists.txt b/lib/Conversion/PrepareGpuHost/CMakeLists.txt new file mode 100644 index 00000000..bae61b95 --- /dev/null +++ b/lib/Conversion/PrepareGpuHost/CMakeLists.txt @@ -0,0 +1,8 @@ +add_llvm_library(COMETPrepareGpuHost + PrepareGpuHost.cpp + + DEPENDS + PrepareGpuHostPassIncGen + + LINK_LIBS PUBLIC +) diff --git a/lib/Conversion/PrepareGpuHost/PrepareGpuHost.cpp b/lib/Conversion/PrepareGpuHost/PrepareGpuHost.cpp new file mode 100644 index 00000000..62ebabda --- /dev/null +++ b/lib/Conversion/PrepareGpuHost/PrepareGpuHost.cpp @@ -0,0 +1,340 @@ +#include "comet/Conversion/PrepareGpuHost/PrepareGpuHostPass.h" +#include "mlir/Dialect/Func/IR/FuncOps.h" +#include "mlir/IR/Attributes.h" +#include "mlir/IR/Builders.h" +#include "mlir/IR/BuiltinAttributes.h" +#include "mlir/IR/Operation.h" +#include "mlir/Pass/Pass.h" +#include "mlir/Dialect/GPU/IR/GPUDialect.h" +#include "mlir/Dialect/LLVMIR/LLVMDialect.h" +#include "mlir/Dialect/MemRef/IR/MemRef.h" +#include "triton/Dialect/Triton/IR/Dialect.h" +#include "llvm/ADT/STLExtras.h" + +using namespace mlir; + +#define GEN_PASS_CLASSES +#include "comet/Conversion/PrepareGpuHost/Passes.h" + +class PrepareGpuHost + : public mlir::comet::PrepareGpuHostBase { +public: + PrepareGpuHost() = default; + + PrepareGpuHost(bool generateAllocsAndTransfers) { + this->generateAllocsAndTransfers = generateAllocsAndTransfers; + } + + void runOnOperation() override { + mlir::ModuleOp modOp = getOperation(); + OpBuilder builder(modOp); + + std::map funcs; + std::map gpu_to_triton_kernel; + std::map gpu_name_to_funcOp; + std::map triton_name_to_triton_func_op; + auto gpuModules = modOp.getOps(); + for(auto gpuModuleOp: gpuModules) + { + auto funcOps = gpuModuleOp.getOps(); + for(func::FuncOp funcOp: llvm::make_early_inc_range(funcOps)) + { + std::map> argsToSet; + if(!funcOp->hasAttr(gpu::GPUDialect::getKernelFuncAttrName())) + { + continue; + } + builder.setInsertionPoint(funcOp); + SmallVector newTypes; + for(auto arg: funcOp.getArguments()) + { + auto argType = arg.getType(); + newTypes.push_back(argType); + if(MemRefType rankedType = dyn_cast(argType)) + { + argsToSet[newTypes.size() - 1] = {funcOp.getArgAttr(arg.getArgNumber(), "gpu.read"), funcOp.getArgAttr(arg.getArgNumber(), "gpu.write")}; + if(rankedType.hasRank()) + { + newTypes.push_back(builder.getIndexType()); + for(int64_t i = 0; i < rankedType.getRank(); i++) + { + newTypes.push_back(builder.getIndexType()); + } + for(int64_t i = 0; i < rankedType.getRank(); i++) + { + newTypes.push_back(builder.getIndexType()); + } + } + else + { + llvm::errs() << "ERROR! Unranked memrefs are not supported currently\n"; + return signalPassFailure(); + } + } + } + + auto funcType = builder.getFunctionType({newTypes}, funcOp->getResultTypes()); + auto newFunc = builder.create(funcOp->getLoc(), funcOp.getName(), funcType); + auto entryBlock = newFunc.addEntryBlock(); + builder.setInsertionPointToEnd(entryBlock); + builder.create(funcOp.getLoc()); + newFunc->setAttr(gpu::GPUDialect::getKernelFuncAttrName(), + builder.getUnitAttr()); + for(auto [argNumber, attrs]: argsToSet) + { + if(attrs[0]) + { + newFunc.setArgAttr(argNumber, "gpu.read", attrs[0]); + } + if(attrs[1]) + { + newFunc.setArgAttr(argNumber, "gpu.write", attrs[1]); + } + } + gpu_name_to_funcOp[funcOp.getName().str()] = newFunc; + funcOp->erase(); + } + } + + modOp->walk([&gpu_to_triton_kernel, &triton_name_to_triton_func_op](mlir::triton::FuncOp TTFuncOp) { + gpu::GPUModuleOp gpuModuleOp = TTFuncOp->getParentOfType(); + gpu_to_triton_kernel[gpuModuleOp.getName().str() +"::"+ TTFuncOp.getName().substr(3 + gpuModuleOp.getName().size()).str()] = TTFuncOp.getName().str(); + triton_name_to_triton_func_op[TTFuncOp.getName().str()] = TTFuncOp; + }); + + modOp.walk([&](gpu::LaunchFuncOp launchOp){ + builder.setInsertionPoint(launchOp); + SmallVector newValues; + for(auto& operand: llvm::make_early_inc_range(launchOp.getKernelOperandsMutable())) + { + newValues.push_back(operand.get()); + if(MemRefType memref = mlir::dyn_cast(operand.get().getType())) + { + auto metadata = builder.create(launchOp->getLoc(), operand.get()); + newValues.insert(newValues.end(), metadata.getResults().begin()+1, metadata.getResults().end()); // Skip the memref + } + } + builder.create(launchOp->getLoc(), launchOp.getKernel(), launchOp.getGridSizeOperandValues(), launchOp.getBlockSizeOperandValues(), launchOp.getDynamicSharedMemorySize(), newValues); + launchOp->erase(); + }); + + std::vector launchOps; + + modOp->walk([&launchOps](mlir::gpu::LaunchFuncOp launchOp) { + launchOps.push_back(launchOp); + }); + + std::set initFuncs; + std::vector toAlloc; + + for (auto launchOp : launchOps) { + auto name = + gpu_to_triton_kernel[(launchOp.getKernelModuleName().strref() + + "::" + launchOp.getKernelName().strref()) + .str()]; + builder.setInsertionPoint(launchOp); + auto ttFunc = triton_name_to_triton_func_op[name]; + Value sharedMem = builder.create( + launchOp->getLoc(), ttFunc->getAttrOfType("triton_gpu.shared").getInt(), 32); + launchOp.getDynamicSharedMemorySizeMutable().assign(sharedMem); + + Value numWarps = builder.create( + launchOp->getLoc(), + ttFunc->getAttrOfType("triton_gpu.num-warps").getInt()); + Value threadsPerWarp = builder.create( + launchOp->getLoc(), + ttFunc->getAttrOfType("triton_gpu.threads-per-warp") + .getInt()); + Value numThreads = builder.create(launchOp->getLoc(), numWarps, threadsPerWarp); + Value one = builder.create(launchOp->getLoc(), 1); + launchOp.getBlockSizeXMutable().assign(numThreads); + launchOp.getBlockSizeYMutable().assign(one); + launchOp.getBlockSizeZMutable().assign(one); + + for (auto &operand : launchOp->getOpOperands()) { + if (isa(operand.get().getType())) { + toAlloc.push_back(operand.get()); + } + } + } + + auto cmp = [](Value a, Value b) { + return a.getAsOpaquePointer() < b.getAsOpaquePointer(); + }; + std::set uniqueGpuAllocs(toAlloc.begin(), + toAlloc.end(), cmp); + for (Value alloc : uniqueGpuAllocs) { + bool isInDevice = false; + MemRefType allocType = mlir::cast(alloc.getType()); + mlir::gpu::AllocOp gpuAlloc; + if (auto defOp = alloc.getDefiningOp()) { + builder.setInsertionPointAfter(defOp); + } else if (mlir::isa(alloc)) { + if(mlir::cast(mlir::cast(alloc).getOwner()->getParentOp()).getArgAttr(mlir::cast(alloc).getArgNumber(), "gpu.indevice")) + { + isInDevice = true; + } + builder.setInsertionPointToStart(alloc.getParentBlock()); + } else { + assert(false && + "Value has not defining Op and is not a block argument."); + } + + if(!isInDevice) + { + if (allocType.hasStaticShape()) { + gpuAlloc = builder.create( + alloc.getLoc(), allocType, ValueRange(), ValueRange(), + ValueRange()); + } else { + std::vector dynDims; + for (size_t i = 0; i < allocType.getShape().size(); i++) { + if (allocType.isDynamicDim(i)) { + dynDims.push_back( + builder.create(alloc.getLoc(), alloc, i)); + } + } + gpuAlloc = builder.create( + alloc.getLoc(), allocType, ValueRange(), dynDims, ValueRange()); + } + + + auto op = alloc.getDefiningOp() != NULL + ? alloc.getDefiningOp()->getResult(0) + : alloc; + for (auto &use : llvm::make_early_inc_range(op.getUses())) { + if (mlir::gpu::LaunchFuncOp launchOp = + dyn_cast(use.getOwner())) { + builder.setInsertionPoint(launchOp); + int offset = launchOp->getNumOperands() - launchOp.getNumKernelOperands(); + int operNum = use.getOperandNumber() - offset; + if(gpu_name_to_funcOp[launchOp.getKernelName().str()].getArgAttr(operNum, "gpu.read")) + { + auto gpuMemCpy = builder.create(launchOp->getLoc(), TypeRange(), + ValueRange(), + gpuAlloc.getMemref(), alloc); + gpuMemCpy->setAttr("gpu.read", builder.getUnitAttr()); + } + use.set(gpuAlloc.getMemref()); + builder.setInsertionPointAfter(launchOp); + if(gpu_name_to_funcOp[launchOp.getKernelName().str()].getArgAttr(operNum, "gpu.write")) + { + auto gpuMemCpy = builder.create(launchOp->getLoc(), TypeRange(), + ValueRange(), alloc, + gpuAlloc.getMemref()); + gpuMemCpy->setAttr("gpu.write", builder.getUnitAttr()); + } + } + } + } + else + { + for (auto &use : llvm::make_early_inc_range(alloc.getUses())) { + if (mlir::gpu::LaunchFuncOp launchOp = + dyn_cast(use.getOwner())) { + builder.setInsertionPoint(launchOp); + auto ptr = builder.create(launchOp->getLoc(), alloc); + use.set(ptr); + } + } + } + } + + std::vector gpuAllocs; + modOp->walk([&gpuAllocs](mlir::gpu::AllocOp gpuAllocOp) { + gpuAllocs.push_back(gpuAllocOp); + }); + + + std::map> memEffects; + modOp->walk([&uniqueGpuAllocs, &memEffects](Operation* memEffect) { + for(auto op: memEffect->getOperands()) + { + if(uniqueGpuAllocs.find(op) != uniqueGpuAllocs.end()) + { + memEffects[op.getAsOpaquePointer()].push_back(memEffect); + } + } + }); + + for(auto& [memref, effects]: memEffects) + { + bool copyIn = true; + std::vector copyDelete; + for(size_t i = 0; i < effects.size(); i++) + { + if(mlir::gpu::MemcpyOp gpuCopy = mlir::dyn_cast(effects[i])) + { + if(!copyIn & gpuCopy->hasAttr("gpu.read")) + { + gpuCopy->erase(); + } + else if(gpuCopy->hasAttr("gpu.read")) + { + copyIn = false; + } + else if(gpuCopy->hasAttr("gpu.write")) + { + copyIn = false; + copyDelete.push_back(gpuCopy); + } + } + else if(mlir::memref::StoreOp store = mlir::dyn_cast(effects[i])) + { + copyIn = true; + if(!copyDelete.empty()) + { + copyDelete.pop_back(); + } + } + else if(mlir::memref::CopyOp copy = mlir::dyn_cast(effects[i])) + { + if(copy.getTarget().getAsOpaquePointer() == memref) + { + copyIn = true; + } + } + else if(mlir::memref::LoadOp load = mlir::dyn_cast(effects[i])) + { + if(!copyDelete.empty()) + { + copyDelete.pop_back(); + } + } + else if(isa(effects[i])) + { + continue; + } + else // Unknown operation, be conservative + { + copyIn = true; + if(!copyDelete.empty()) + { + copyDelete.pop_back(); + } + } + } + + for(auto toDelete: copyDelete) + { + toDelete->erase(); + } + } + + } +}; + + +std::unique_ptr> +mlir::comet::createPrepareGpuHostPass() { + // std::cout << "Running createPrepareGpuHostPass\n"; + + return std::make_unique<::PrepareGpuHost>(); +} + +std::unique_ptr> +mlir::comet::createPrepareGpuHostPass(bool generateAllocsAndTransfers) { + // std::cout << "Running createPrepareGpuHostPass\n"; + + return std::make_unique<::PrepareGpuHost>(generateAllocsAndTransfers); +} \ No newline at end of file diff --git a/lib/Conversion/TensorAlgebraToIndexTree/TensorAlgebraToIndexTree.cpp b/lib/Conversion/TensorAlgebraToIndexTree/TensorAlgebraToIndexTree.cpp index 7a76fd18..ec1a2347 100644 --- a/lib/Conversion/TensorAlgebraToIndexTree/TensorAlgebraToIndexTree.cpp +++ b/lib/Conversion/TensorAlgebraToIndexTree/TensorAlgebraToIndexTree.cpp @@ -28,27 +28,25 @@ #include "comet/Dialect/IndexTree/Transforms/Tensor.h" #include "comet/Dialect/TensorAlgebra/IR/TADialect.h" #include "comet/Dialect/IndexTree/Transforms/UnitExpression.h" -#include "comet/Dialect/IndexTree/IR/IndexTree.h" #include "comet/Dialect/IndexTree/IR/IndexTreeDialect.h" #include "comet/Dialect/IndexTree/Passes.h" #include "comet/Dialect/Utils/Utils.h" #include "mlir/Dialect/Func/IR/FuncOps.h" +#include "mlir/IR/AffineExpr.h" #include "mlir/IR/Block.h" #include "mlir/IR/Operation.h" +#include "mlir/IR/ValueRange.h" using namespace mlir; using namespace mlir::indexTree; using namespace mlir::tensorAlgebra; // *********** For debug purpose *********// -// #define COMET_DEBUG_MODE +//#define COMET_DEBUG_MODE #include "comet/Utils/debug.h" -#undef COMET_DEBUG_MODE // *********** For debug purpose *********// -using namespace mlir; - namespace { struct LowerTensorAlgebraToIndexTreePass @@ -57,7 +55,7 @@ namespace MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(LowerTensorAlgebraToIndexTreePass) LowerTensorAlgebraToIndexTreePass(TargetDevice device) : device(device){}; void runOnOperation() override; - + TargetDevice device; }; @@ -143,11 +141,10 @@ Value getRealLhs(Operation *op) return setOp.getOperand(1); } -Value getRealRhs(Operation *op) +Value getRealRhs(Value val) { /// this will return set_op for transpose, but messes up getUsers() or subsequent calls to it. - Operation *firstUser = op->getNextNode(); - comet_pdump(firstUser); + /// TODO(gkestor): need to find out why user set_op is not showing up in users of TransposeOp /// from the for loop below. once resolved, remove getNextNode(). /// Operation *firstUser; @@ -156,518 +153,360 @@ Value getRealRhs(Operation *op) /// firstUser = user; /// break; //} - - if (isa(op)) + if(Operation* op = val.getDefiningOp()) { - if (isa(firstUser)) - { - TensorSetOp setOp = cast(firstUser); - return setOp.getOperand(1); - } - else - { - llvm::errs() << "ERROR: Transpose has no set_op after it!\n"; - } - } - else - { - /// do nothing - return op->getResult(0); - } - return op->getResult(0); -} + Operation *firstUser = op->getNextNode(); + comet_pdump(firstUser); -void buildDefUseInfo(UnitExpression *e) -{ - auto lhs = e->getLHS(); - lhs->setDefiningExpr(e); - for (auto operand : e->getOperands()) - { - if (auto def = operand->getDefiningExpr()) + if (isa(op)) { - def->addUser(e); + if (isa(firstUser)) + { + TensorSetOp setOp = cast(firstUser); + return setOp.getOperand(1); + } + else + { + llvm::errs() << "ERROR: Transpose has no set_op after it!\n"; + } + } } -} -IndicesType getUnion(IndicesType indices1, IndicesType indices2) -{ - sort(indices1.begin(), indices1.end()); - sort(indices2.begin(), indices2.end()); - - IndicesType allIndices(indices1.size() * 4); - - IndicesType::iterator it = set_union(indices1.begin(), indices1.end(), indices2.begin(), indices2.end(), allIndices.begin()); - allIndices.resize(it - allIndices.begin()); - return allIndices; + return val; } -IndicesType gpuIndices(IndicesType indices1, IndicesType indices2) -{ - sort(indices1.begin(), indices1.end()); - sort(indices2.begin(), indices2.end()); - - IndicesType interIndices; - IndicesType unIndices; - IndicesType difIndices; - IndicesType allIndices; - - std::set_intersection(indices1.begin(), indices1.end(), indices2.begin(), indices2.end(), std::back_inserter(interIndices)); - set_union(indices1.begin(), indices1.end(), indices2.begin(), indices2.end(), std::back_inserter(unIndices)); - std::set_difference(unIndices.begin(), unIndices.end(), interIndices.begin(), interIndices.end(), std::back_inserter(difIndices)); - allIndices = difIndices; - allIndices.insert(allIndices.end(), interIndices.begin(), interIndices.end()); - - // allIndices.resize(it - allIndices.begin()); - return allIndices; -} - -void doTensorMultOp(TensorMultOp op, unique_ptr &tree, TargetDevice device = CPU) -{ - Value rhs1_tensor = getRealRhs(op.getRhs1().getDefiningOp()); - Value rhs2_tensor = getRealRhs(op.getRhs2().getDefiningOp()); +// void buildDefUseInfo(UnitExpression *e) +// { +// auto lhs = e->getLHS(); +// lhs->setDefiningExpr(e); +// for (auto operand : e->getOperands()) +// { +// if (auto def = operand->getDefiningExpr()) +// { +// def->addUser(e); +// } +// } +// } + +// IndicesType getUnion(IndicesType indices1, IndicesType indices2) +// { +// sort(indices1.begin(), indices1.end()); +// sort(indices2.begin(), indices2.end()); + +// IndicesType allIndices(indices1.size() * 4); + +// IndicesType::iterator it = set_union(indices1.begin(), indices1.end(), indices2.begin(), indices2.end(), allIndices.begin()); +// allIndices.resize(it - allIndices.begin()); +// return allIndices; +// } + +// IndicesType gpuIndices(IndicesType indices1, IndicesType indices2) +// { +// sort(indices1.begin(), indices1.end()); +// sort(indices2.begin(), indices2.end()); + +// IndicesType interIndices; +// IndicesType unIndices; +// IndicesType difIndices; +// IndicesType allIndices; + +// std::set_intersection(indices1.begin(), indices1.end(), indices2.begin(), indices2.end(), std::back_inserter(interIndices)); +// set_union(indices1.begin(), indices1.end(), indices2.begin(), indices2.end(), std::back_inserter(unIndices)); +// std::set_difference(unIndices.begin(), unIndices.end(), interIndices.begin(), interIndices.end(), std::back_inserter(difIndices)); +// allIndices = difIndices; +// allIndices.insert(allIndices.end(), interIndices.begin(), interIndices.end()); + +// // allIndices.resize(it - allIndices.begin()); +// return allIndices; +// } + + +template +mlir::LogicalResult generalIndexOperationRewrite( + mlir::Operation* op, + ArrayRef operands, + mlir::ConversionPatternRewriter &rewriter, + TargetDevice device, + bool compute_missing = false) { + comet_pdump(op); + + auto loc = op->getLoc(); + auto context = rewriter.getContext(); + TATensorOp mult_op = llvm::dyn_cast(op); + + Value rhs1_tensor = getRealRhs(mult_op.getRhs1()); + Value rhs2_tensor = getRealRhs(mult_op.getRhs2()); Value lhs_tensor = getRealLhs(op); - Value mask_tensor = op.getMask(); - // rhs1_tensor.getDefiningOp - comet_debug() << "LowerTensorAlgebraToIndexTreePass: doTensorMultOp\n"; - comet_debug() << "rhs1-tensor\n"; comet_vdump(rhs1_tensor); - comet_debug() << "rhs2-tensor\n"; comet_vdump(rhs2_tensor); - comet_debug() << "lhs-tensor\n"; comet_vdump(lhs_tensor); - comet_debug() << "mask-tensor\n"; - comet_vdump(mask_tensor); - - std::vector rhs1_labels = op.getRhs1IndexLabels(); - std::vector rhs2_labels = op.getRhs2IndexLabels(); - std::vector lhs_labels = op.getResultIndexLabels(); - - auto allPerms = getAllPerms(op.getIndexingMaps()); - assert(allPerms.size() == 3); -#ifdef COMET_DEBUG_MODE - comet_debug() << "\n"; - llvm::errs() << "["; - for (auto &perm : allPerms) + Value mask_tensor = nullptr; + if (llvm::isa(op)) { - llvm::errs() << "["; - for (auto &i : perm) + mask_tensor = llvm::cast(op).getMask(); + if(mask_tensor && (mask_tensor == rhs1_tensor || mask_tensor == rhs2_tensor)) { - llvm::errs() << i << ","; + mask_tensor = rewriter.create(loc, mask_tensor.getType(), mask_tensor); } - llvm::errs() << "],"; - } - llvm::errs() << "]\n"; - comet_debug() << ""; -// comet_debug() << allPerms; -#endif - - auto allFormats = getAllFormats(op.getFormatsAttr(), allPerms); - auto SemiringOp = op.getSemiringAttr(); - auto MaskingTypeAttr = op.getMaskTypeAttr(); - - /// If the operation is one of the chosen operations, then record output indices as parallel interators. - bool is_chosen_operations = check_chosen_operations(allPerms, allFormats); - - auto B = tree->getOrCreateTensor(rhs1_tensor, rhs1_labels, allFormats[0]); - auto C = tree->getOrCreateTensor(rhs2_tensor, rhs2_labels, allFormats[1]); - auto A = tree->getOrCreateTensor(lhs_tensor, lhs_labels, allFormats[2]); - - Tensor *M; - std::unique_ptr e; - std::vector empty; - if (mask_tensor != nullptr) /// mask is an optional input - { - comet_debug() << "mask input provided by user\n"; - M = tree->getOrCreateTensor(mask_tensor, empty, allFormats[2]); /// We don't need indexlabel info for the mask - e = make_unique(A, B, C, M, "*"); - } - else - { - comet_debug() << "no mask input provided by user\n"; - e = make_unique(A, B, C, "*"); } - e->setSemiring(SemiringOp.cast().getValue()); - e->setMaskType(MaskingTypeAttr.cast().getValue()); - - e->setOperation(op); - buildDefUseInfo(e.get()); - - auto inputDomains = e->computeInputIterDomains(); - auto outputDomains = e->computeOutputIterDomains(); - - IndicesType rhs1_indices = tree->getIndices(rhs1_labels); - IndicesType rhs2_indices = tree->getIndices(rhs2_labels); - IndicesType allIndices; - - switch (device) + auto indexing_maps = mult_op.getIndexingMaps(); + auto semiring = cast(mult_op.getSemiringAttr()).getValue(); + + auto tensor_type = op->getResultTypes()[0]; + auto itree_op = rewriter.create(loc, tensor_type, ValueRange(lhs_tensor), ValueRange()); + Region* body = &itree_op.getRegion(); + loc = body->getLoc(); + Block* block = rewriter.createBlock(body, {}, TypeRange(lhs_tensor.getType()), {lhs_tensor.getLoc()}); + lhs_tensor = block->getArgument(0); + comet_vdump(itree_op); + comet_pdump(block); + + indexTree::IndexTreeType tree_type = indexTree::IndexTreeType::get(context); + Value parent = rewriter.create(loc, tree_type); + comet_vdump(parent); + + //Construct each index variable + auto lhsMap = cast(indexing_maps[2]).getValue(); + indexTree::IndexNodeType index_node_type = indexTree::IndexNodeType::get(context); + std::vector index_nodes; + bool is_parallel = true; // Outer-most, non-reduction dimensions are parallel + + // TODO: For now, we do not support outputting sparse tensors from a parallel loop. + if(llvm::isa(tensor_type)) { - case mlir::tensorAlgebra::CPU: - { - allIndices = getUnion(rhs1_indices, rhs2_indices); - } - break; - case mlir::tensorAlgebra::GPU: - { - allIndices = gpuIndices(rhs1_indices, rhs2_indices); - } - break; + is_parallel = false; } - // tree->setSizeOfIteratorTypes(allIndices.size()); // Set the total number of iterators - - auto lhsIndices = A->getIndices(); - TreeNode *parent = tree->getRoot(); - for (unsigned long i = 0; i < allIndices.size(); ++i) - { - int index = allIndices[i]; - auto &idomain = inputDomains.at(index); + std::map exprToIndex; - auto node = tree->addIndexNode(index, parent, idomain); - /// If this index appears on the lhs too, set output domain for the index node - /// and also set the index as a parallel iterator - unique_ptr iteratorType(new IteratorType); - comet_debug() << iteratorType->dump() << "\n"; - if (std::find(lhsIndices.begin(), lhsIndices.end(), index) != lhsIndices.end()) + if(device == TargetDevice::GPU || device == TargetDevice::FPGA) + { + bool isSPMMCSR = getTensorFormatString(rhs1_tensor.getType()) == "CSR" and getTensorFormatString(rhs2_tensor.getType()) == "Dense" && isa(op); + + for (unsigned i = 0; i < lhsMap.getNumResults() ; i++) { - auto &odomain = outputDomains.at(index); - node->setOutputDomain(odomain); - if (is_chosen_operations) + if(device == TargetDevice::GPU && isSPMMCSR && i == 0) { - /// If the operation is one of the chosen ones, and the index appears on the lhs, - /// then the index has "parallel" as its iterator type. - iteratorType->setType("parallel"); + parent = rewriter.create(loc, index_node_type, parent, nullptr, is_parallel, rewriter.getStringAttr("dimY_grid")); } + else + { + parent = rewriter.create(loc, index_node_type, parent, nullptr, is_parallel, nullptr); + } + exprToIndex[cast(lhsMap.getResult(i)).getPosition()] = parent; } - comet_debug() << "index " << index << "\n"; - /// Set the iterator type of the node - tree->setIteratorTypeByIndex(index, std::move(iteratorType)); - node->setIteratorType(tree->getIteratorTypeByIndex(index)); - comet_debug() << "tree: " << tree->getIteratorTypeByIndex(index)->dump() << " ptr: " << tree->getIteratorTypeByIndex(index) << "\n"; - comet_debug() << "node: " << node->getIteratorType()->dump() << " ptr: " << node->getIteratorType() << "\n"; - - parent = node; - } - - tree->addComputeNode(std::move(e), parent); -} - -template -void doElementWiseOp(T op, unique_ptr &tree) -{ - std::vector rhs1_labels = op.getRhs1IndexLabels(); - std::vector rhs2_labels = op.getRhs2IndexLabels(); - std::vector lhs_labels = op.getResultIndexLabels(); - - Value rhs1_tensor = getRealRhs(op.getRhs1().getDefiningOp()); - Value rhs2_tensor = getRealRhs(op.getRhs2().getDefiningOp()); - Value lhs_tensor = getRealLhs(op); - - comet_debug() << "LowerTensorAlgebraToIndexTreePass: doElementWiseMultOp\n"; - comet_debug() << "rhs1-tensor\n"; - comet_vdump(rhs1_tensor); - comet_debug() << "rhs2-tensor\n"; - comet_vdump(rhs2_tensor); - comet_debug() << "lhs-tensor\n"; - comet_vdump(lhs_tensor); - - auto allPerms = getAllPerms(op.getIndexingMaps()); - auto allFormats = getAllFormats(op.getFormatsAttr(), allPerms); - auto SemiringOp = op.getSemiringAttr(); - auto maskAttr = "none"; - - assert(allPerms.size() == 3); - - auto B = tree->getOrCreateTensor(rhs1_tensor, rhs1_labels, allFormats[0]); - auto C = tree->getOrCreateTensor(rhs2_tensor, rhs2_labels, allFormats[1]); - auto A = tree->getOrCreateTensor(lhs_tensor, lhs_labels, allFormats[2]); - - auto e = make_unique(A, B, C, "*"); - - e->setOperation(op); - e->setSemiring(SemiringOp.template cast().getValue()); /// for element-wise multiplication - e->setMaskType(maskAttr); /// for element-wise multiplication - buildDefUseInfo(e.get()); - - auto inputDomains = e->computeInputIterDomains(); - auto outputDomains = e->computeOutputIterDomains(); - - /// RHS and LHS indices must be the same for elementwise multiplication - IndicesType allIndices = tree->getIndices(rhs1_labels); - // tree->setSizeOfIteratorTypes(allIndices.size()); // Set the total number of iterators + - auto lhsIndices = A->getIndices(); - TreeNode *parent = tree->getRoot(); - for (unsigned long i = 0; i < allIndices.size(); i++) - { - int index = allIndices[i]; - auto &idomain = inputDomains.at(index); - - auto node = tree->addIndexNode(index, parent, idomain); - unique_ptr iteratorType(new IteratorType); - - /// If this index appears on the lhs too, set output domain for the index node - if (std::find(lhsIndices.begin(), lhsIndices.end(), index) != lhsIndices.end()) + is_parallel = false; + for (unsigned ii = 0; ii < lhsMap.getNumDims(); ii++) { - auto &odomain = outputDomains.at(index); - node->setOutputDomain(odomain); - iteratorType->setType("parallel"); + if(!lhsMap.isFunctionOfDim(ii)){ + parent = rewriter.create(loc, index_node_type, parent, nullptr, is_parallel, nullptr); + exprToIndex[ii] = parent; + } } - - /// Set iterator type. Currently "default" for all elementwise operations. - tree->setIteratorTypeByIndex(index, std::move(iteratorType)); - node->setIteratorType(tree->getIteratorTypeByIndex(index)); - comet_debug() << "tree: " << tree->getIteratorTypeByIndex(index)->dump() << " ptr: " << tree->getIteratorTypeByIndex(index) << "\n"; - comet_debug() << "node: " << node->getIteratorType()->dump() << " ptr: " << node->getIteratorType() << "\n"; - parent = node; } - tree->addComputeNode(std::move(e), parent); - /// cout << "print tree after tc\n"; - /// tree->print(); -} - -/// helper for treeToDialect() -Operation *getSetOpForTC(Operation *op) -{ - assert(isa(op) || isa(op) || isa(op) || isa(op)); - /// TODO(gkestor): fix the issue with getUsers() after getRealRhs(). - comet_debug() << "The following loop may cause issue!\n"; - Operation *firstUser = nullptr; - for (auto user : op->getResult(0).getUsers()) + else { - firstUser = user; - break; - } - - assert(isa(firstUser)); - return firstUser; -} - -/// helper for treeToDialect() -IndexTreeComputeOp createComputeNodeOp(OpBuilder &builder, TreeNode *node, Location &loc) -{ - auto context = builder.getContext(); - IntegerType i64Type = IntegerType::get(context, 64); - auto expr = node->getExpression(); - SmallVector allIndices_rhs; - - for (auto t : expr->getOperands()) - { - SmallVector indices; - for (auto index : t->getIndices()) + for (unsigned i = 0; i < lhsMap.getNumDims(); i++) { - indices.push_back(index); + if(!lhsMap.isFunctionOfDim(i)){ + is_parallel = false; + } + parent = rewriter.create(loc, index_node_type, parent, nullptr, is_parallel, nullptr); + exprToIndex[i] = parent; } - allIndices_rhs.push_back(builder.getI64ArrayAttr(indices)); } - SmallVector allIndices_lhs; - for (auto t : expr->getResults()) - { - SmallVector indices; - for (auto index : t->getIndices()) - { - comet_debug() << index << " \n"; - indices.push_back(index); - } - allIndices_lhs.push_back(builder.getI64ArrayAttr(indices)); - } - SmallVector allFormats_rhs; - for (auto t : expr->getOperands()) + //Construct LHS Operand + llvm::SmallVector pos; + llvm::SmallVector crds; + Value prev_dim = nullptr; + llvm::SmallVector mask_pos; + llvm::SmallVector mask_crds; + Value mask_prev_dim; + auto access_type = rewriter.getIndexType(); + for (size_t i = 0; i < lhsMap.getNumResults(); i++) { - SmallVector formats; - for (auto &f : t->getFormats()) + auto expr = lhsMap.getResult(i); + IndexTreeIndexToTensorOp access_op = rewriter.create( + loc, + TypeRange({access_type, access_type}), + lhs_tensor, + exprToIndex[cast(expr).getPosition()], + rewriter.getUI32IntegerAttr((unsigned)i), + prev_dim + ); + pos.push_back(access_op.getPos()); + crds.push_back(access_op.getCrd()); + prev_dim = pos[pos.size() - 1]; + comet_vdump(access_op); + + if(mask_tensor != nullptr) { - formats.push_back(f); + IndexTreeIndexToTensorOp access_op = rewriter.create( + loc, + TypeRange({access_type, access_type}), + mask_tensor, + exprToIndex[cast(expr).getPosition()], + rewriter.getUI32IntegerAttr((unsigned)i), + mask_prev_dim + ); + mask_pos.push_back(access_op.getPos()); + mask_crds.push_back(access_op.getCrd()); + mask_prev_dim = mask_pos[mask_pos.size() - 1]; } - allFormats_rhs.push_back(builder.getStrArrayAttr(formats)); } - SmallVector allFormats_lhs; - for (auto t : expr->getResults()) + indexTree::OperandType operand_type = indexTree::OperandType::get(context); + Value lhs_operand = rewriter.create(loc, operand_type, + lhs_tensor, pos, + crds); + comet_vdump(lhs_operand); + Value mask_operand = nullptr; + if(mask_tensor != nullptr) { - SmallVector formats; - for (auto &f : t->getFormats()) - { - formats.push_back(f); - } - allFormats_lhs.push_back(builder.getStrArrayAttr(formats)); + mask_operand = rewriter.create(loc, operand_type, + mask_tensor, mask_pos, + mask_crds); } - std::vector t_rhs; - Value t_lhs = expr->getLHS()->getValue(); - for (auto o : expr->getOperands()) + //Construct RHS operands + std::vector rhs_operands; + pos.clear(); + crds.clear(); + prev_dim = nullptr; + auto affineMap = cast(indexing_maps[0]).getValue(); + for (size_t i = 0; i < affineMap.getNumResults(); i++) { - t_rhs.push_back(o->getValue()); + auto expr = affineMap.getResult(i); + IndexTreeIndexToTensorOp access_op = rewriter.create( + loc, + TypeRange({access_type, access_type}), + rhs1_tensor, + exprToIndex[cast(expr).getPosition()], + rewriter.getUI32IntegerAttr((unsigned)i), + prev_dim + ); + pos.push_back(access_op.getPos()); + crds.push_back(access_op.getCrd()); + prev_dim = pos[pos.size() - 1]; + comet_vdump(access_op); } - - /// check if mask exists and add to t_rhs - if (expr->getMask() != nullptr) + rhs_operands.push_back(rewriter.create( + loc, operand_type, rhs1_tensor, pos, crds)); + + pos.clear(); + crds.clear(); + prev_dim = nullptr; + affineMap = cast(indexing_maps[1]).getValue(); + for (size_t i = 0; i < affineMap.getNumResults(); i++) { - comet_debug() << "user has provided mask input\n"; - t_rhs.push_back(expr->getMask()->getValue()); /// add mask to IndexTreeComputeRHSOp + auto expr = affineMap.getResult(i); + IndexTreeIndexToTensorOp access_op = rewriter.create( + loc, + TypeRange({access_type, access_type}), + rhs2_tensor, + exprToIndex[cast(expr).getPosition()], + rewriter.getUI32IntegerAttr((unsigned)i), + prev_dim + ); + pos.push_back(access_op.getPos()); + crds.push_back(access_op.getCrd()); + prev_dim = pos[pos.size() - 1]; + comet_vdump(access_op); } - - Value leafop_rhs = builder.create(loc, - mlir::UnrankedTensorType::get(builder.getF64Type()), t_rhs, - builder.getArrayAttr(allIndices_rhs), - builder.getArrayAttr(allFormats_rhs)); - comet_vdump(leafop_rhs); - Value leafop_lhs = builder.create(loc, - mlir::UnrankedTensorType::get(builder.getF64Type()), t_lhs, - builder.getArrayAttr(allIndices_lhs), - builder.getArrayAttr(allFormats_lhs)); - comet_vdump(leafop_lhs); - - bool comp_worksp_opt = false; /// non-compressed workspace, this is a place-holder and it is updated in workspace transform pass. - llvm::StringRef semiring = expr->getSemiring(); - llvm::StringRef maskType = expr->getMaskType(); - auto leafop = builder.create(loc, i64Type, leafop_rhs, leafop_lhs, builder.getBoolAttr(comp_worksp_opt), builder.getStringAttr(semiring), builder.getStringAttr(maskType)); - - comet_pdump(leafop); - return leafop; + rhs_operands.push_back(rewriter.create( + loc, operand_type, rhs2_tensor, pos, crds)); + + Value compute_op = rewriter.create( + loc, + tensor_type, + parent, + lhs_operand, + rhs_operands, + mask_operand, + rewriter.getStringAttr(semiring), + rewriter.getBoolAttr(compute_missing) + ); + comet_vdump(compute_op); + + rewriter.create(loc, TypeRange(), compute_op); + rewriter.replaceOp(op, itree_op->getResults()); + return success(); } -/** - * This function performs the actual removal of the ta operations in the tree, - * and add corresponding ta.itree operations.› - * @param tree - */ -void treeToDialect(Index_Tree *tree) -{ - vector TAOps = tree->getContainingTAOps(); - unsigned int TAOpsID = 0; - OpBuilder builder(TAOps[TAOpsID]); - auto loc = TAOps[TAOpsID]->getLoc(); - auto context = builder.getContext(); - - std::map nodeToOp; - - IntegerType i64Type = IntegerType::get(context, 64); +struct TensorMultOpLowering : public mlir::ConversionPattern { + TensorMultOpLowering(mlir::MLIRContext *ctx, TargetDevice device) + : mlir::ConversionPattern(TensorMultOp::getOperationName(), 1, ctx), device(device) {} + TargetDevice device; - for (auto &node : tree->getNodesInReverseTopoOrder()) - { - if (node->isComputeNode()) - { - assert(nodeToOp.count(node) == 0); - builder.setInsertionPoint(TAOps[TAOpsID]); - nodeToOp[node] = createComputeNodeOp(builder, node, loc); - TAOpsID++; - } - else if (node->isRealIndexNode()) - { - if (node->getChildren().empty()) - { - continue; /// to skip nodes that become childless after fusion - } - SmallVector children; - for (auto c : node->getChildren()) - { - assert(nodeToOp.count(c) > 0); - children.push_back(nodeToOp[c]); - } - /// assert(!children.empty()); - SmallVector indices; - indices.push_back(node->getIndex()); - auto indicesAttr = builder.getI64ArrayAttr(indices); - - SmallVector ids; - ids.push_back(node->getId()); - - /// new attribute iterator_type - auto dumb_iterator_type = builder.getStringAttr(node->getIteratorType()->getType()); - Value indexNodeOp = builder.create(loc, - i64Type, - children, - indicesAttr, - dumb_iterator_type); + mlir::LogicalResult + matchAndRewrite(mlir::Operation *op, ArrayRef operands, + mlir::ConversionPatternRewriter &rewriter) const final { + return generalIndexOperationRewrite(op, operands, rewriter, this->device); + } +}; - nodeToOp[node] = indexNodeOp; +struct TensorElewsMultOpLowering : public mlir::ConversionPattern { + TensorElewsMultOpLowering(mlir::MLIRContext *ctx, TargetDevice device) + : mlir::ConversionPattern(TensorElewsMultOp::getOperationName(), 1, ctx), device(device) {} + TargetDevice device; - if (node->getParent() != nullptr && node->getParent()->isFillerIndexNode()) - { -#ifdef DEBUG_MODE_LowerTensorAlgebraToIndexTreePass - Value op = builder.create(loc, i64Type, indexNodeOp); - comet_vdump(op); -#else - builder.create(loc, i64Type, indexNodeOp); -#endif - } - } + mlir::LogicalResult + matchAndRewrite(mlir::Operation *op, ArrayRef operands, + mlir::ConversionPatternRewriter &rewriter) const final { + return generalIndexOperationRewrite(op, operands, rewriter, this->device); } - - for (auto op : tree->getContainingTAOps()) - { - auto setOp = getSetOpForTC(op); - setOp->erase(); - op->erase(); +}; + +struct TensorAddOpLowering : public mlir::ConversionPattern { + TensorAddOpLowering(mlir::MLIRContext *ctx, TargetDevice device) + : mlir::ConversionPattern(TensorAddOp::getOperationName(), 1, ctx), device(device) {} + TargetDevice device; + mlir::LogicalResult + matchAndRewrite(mlir::Operation *op, ArrayRef operands, + mlir::ConversionPatternRewriter &rewriter) const final { + return generalIndexOperationRewrite(op, operands, rewriter, this->device, true); } -} +}; + +struct TensorSubtractOpLowering : public mlir::ConversionPattern { + TensorSubtractOpLowering(mlir::MLIRContext *ctx, TargetDevice device) + : mlir::ConversionPattern(TensorSubtractOp::getOperationName(), 1, ctx), device(device) {} + TargetDevice device; + mlir::LogicalResult + matchAndRewrite(mlir::Operation *op, ArrayRef operands, + mlir::ConversionPatternRewriter &rewriter) const final { + return generalIndexOperationRewrite(op, operands, rewriter, this->device, true); + } +}; void LowerTensorAlgebraToIndexTreePass::runOnOperation() { - unique_ptr tree; - func::FuncOp func = getOperation(); - // #ifdef COMET_DEBUG_MODE - // comet_debug() << "Before LowerTensorAlgebraToIndexTreePass\n"; - // func.dump(); - // #endif - - tree = Index_Tree::createTreeWithRoot(); - bool formIndexTreeDialect = false; - - comet_debug() << "IndexTree pass running on Function\n"; - for (Block &B : func.getBody()) - { - for (Operation &op : B) - { - if (isa(&op)) - { - doTensorMultOp(cast(&op), tree, device); - formIndexTreeDialect = true; - } - else if (isa(&op)) - { -#ifdef COMET_DEBUG_MODE - comet_debug() << "\n !!! doElementWiseOp\n"; -#endif - doElementWiseOp(cast(&op), tree); - formIndexTreeDialect = true; - } - else if (isa(&op) || isa(&op)) - { - /// elementwise addition and subtraction - if (isa(&op)) - { -#ifdef COMET_DEBUG_MODE - comet_debug() << "\n !!! doElementWiseOp\n"; -#endif - doElementWiseOp(cast(&op), tree); - } - - if (isa(&op)) - { -#ifdef COMET_DEBUG_MODE - comet_debug() << "\n !!! doElementWiseOp\n"; -#endif - doElementWiseOp(cast(&op), tree); - } - formIndexTreeDialect = true; - } - } - } - - if (formIndexTreeDialect) - { - comet_debug() << " Dumping Index tree IR\n"; - /// only do this for TensorMultOp or TensorElewsMultOp - treeToDialect(tree.get()); - } + comet_pdump(getOperation()->getParentOfType()); + mlir::ConversionTarget target(getContext()); + + target.addLegalDialect(); + target.addLegalOp(); + target.addIllegalOp(); + + mlir::RewritePatternSet patterns(&getContext()); + patterns.add(&getContext(), this->device); + + if (mlir::failed(mlir::applyPartialConversion(getOperation(), target, std::move(patterns)))) + signalPassFailure(); + comet_pdump(getOperation()->getParentOfType()); } /// create all the passes. diff --git a/lib/Conversion/TensorAlgebraToSCF/CMakeLists.txt b/lib/Conversion/TensorAlgebraToSCF/CMakeLists.txt index 4a563c79..0572b56c 100644 --- a/lib/Conversion/TensorAlgebraToSCF/CMakeLists.txt +++ b/lib/Conversion/TensorAlgebraToSCF/CMakeLists.txt @@ -1,10 +1,9 @@ add_mlir_conversion_library(COMETTensorAlgebraToSCF - EarlyLowering.cpp - LateLowering.cpp LowerFunc.cpp - LowerPCToLoops.cpp TensorAlgebraToSCF.cpp - + EarlyLowering.cpp + LateLowering.cpp + SparseTensorConversionPass.cpp ADDITIONAL_HEADER_DIRS ${COMET_MAIN_INCLUDE_DIR}/comet/Conversion/TensorAlgebraToIndexTree diff --git a/lib/Conversion/TensorAlgebraToSCF/EarlyLowering.cpp b/lib/Conversion/TensorAlgebraToSCF/EarlyLowering.cpp index 150cdf01..28eff319 100644 --- a/lib/Conversion/TensorAlgebraToSCF/EarlyLowering.cpp +++ b/lib/Conversion/TensorAlgebraToSCF/EarlyLowering.cpp @@ -37,11 +37,13 @@ #include "mlir/IR/BuiltinAttributes.h" #include "mlir/Pass/Pass.h" #include "mlir/Dialect/Tensor/IR/Tensor.h" +#include "mlir/Support/LLVM.h" #include #include #include #include +#include "llvm/Support/Casting.h" #include "llvm/Support/Debug.h" using namespace mlir; @@ -82,7 +84,7 @@ namespace auto tensorOperand = operands[0]; ToTensorOp tensorLoadOp; - if (!isa(tensorOperand.getDefiningOp())) + if (!mlir::isa_and_present(tensorOperand.getDefiningOp())) { /// TODO: may need to re-visit when doing reduction support. /// the user declared output to have zeros. @@ -92,8 +94,9 @@ namespace tensorLoadOp = cast(tensorOperand.getDefiningOp()); auto memref = tensorLoadOp.getMemref(); auto valueAttr = tensorFillOp.getValue(); + + rewriter.setInsertionPoint(tensorLoadOp); Value constantOp = rewriter.create(loc, llvm::cast(valueAttr)); - rewriter.create(loc, constantOp, memref); rewriter.eraseOp(op); @@ -196,17 +199,14 @@ namespace assert(isa(op)); comet_debug() << "TensorDimOpLowering in format begin\n"; comet_vdump(op); - auto tensor = op.getTensor(); - if (tensor.getType().isa()) + if (isa(tensor.getType())) { - SparseTensorConstructOp spconstruct = cast(tensor.getDefiningOp()); ::mlir::TypedValue<::mlir::IndexType> idx = op.getIndex(); auto realIndex = getConstantIntValue(idx); - op.replaceAllUsesWith(spconstruct.getIndices()[18 + *realIndex]); - rewriter.eraseOp(op); + rewriter.replaceOpWithNewOp(op, rewriter.getIndexType(), tensor, *realIndex); } - else if (tensor.getType().isa()) + else if (isa(tensor.getType())) { ::mlir::TypedValue<::mlir::IndexType> idx = op.getIndex(); auto dim = rewriter.create(op.getLoc(), tensor, idx); @@ -254,8 +254,16 @@ namespace tensorAlgebra::TensorSetOp, tensorAlgebra::IndexLabelOp, tensorAlgebra::DenseConstantOp, + tensorAlgebra::TensorMultOp, tensorAlgebra::ScalarOp, - tensorAlgebra::SparseTensorConstructOp>(); + tensorAlgebra::SpTensorAliasOp, + tensorAlgebra::SpTensorGetDimSize, + tensorAlgebra::SpTensorGetDimPos, + tensorAlgebra::SpTensorGetDimCrd, + tensorAlgebra::SpTensorGetVals, + tensorAlgebra::SparseTensorConstructOp, + tensorAlgebra::TensorSortOp, + tensorAlgebra::AllocWorkspaceOp>(); if (failed(applyPartialConversion(function, target, std::move(patterns)))) { diff --git a/lib/Conversion/TensorAlgebraToSCF/LateLowering.cpp b/lib/Conversion/TensorAlgebraToSCF/LateLowering.cpp index c28fbe27..69e6a0ea 100644 --- a/lib/Conversion/TensorAlgebraToSCF/LateLowering.cpp +++ b/lib/Conversion/TensorAlgebraToSCF/LateLowering.cpp @@ -32,9 +32,15 @@ #include "mlir/Dialect/Affine/IR/AffineOps.h" #include "mlir/Dialect/SCF/IR/SCF.h" #include "mlir/Dialect/Func/IR/FuncOps.h" +#include "mlir/IR/BuiltinOps.h" +#include "mlir/IR/BuiltinTypeInterfaces.h" +#include "mlir/IR/BuiltinTypes.h" #include "mlir/Pass/Pass.h" #include "mlir/Dialect/MemRef/IR/MemRef.h" #include "mlir/Dialect/Bufferization/IR/Bufferization.h" +#include "mlir/Transforms/DialectConversion.h" +#include "llvm/ADT/SmallVector.h" +#include using namespace mlir; using namespace mlir::arith; @@ -70,104 +76,113 @@ namespace Location loc = op->getLoc(); auto module = op->getParentOfType(); auto *ctx = op->getContext(); - FloatType f64Type = FloatType::getF64(ctx); - IndexType indexType = IndexType::get(ctx); - Type unrankedMemrefType_f64 = UnrankedMemRefType::get(f64Type, 0); - auto printTensorF64Func = FunctionType::get(ctx, {mlir::UnrankedMemRefType::get(f64Type, 0)}, {}); - auto printTensorIndexFunc = FunctionType::get(ctx, {mlir::UnrankedMemRefType::get(indexType, 0)}, {}); - auto printScalarFunc = FunctionType::get(ctx, {FloatType::getF64(ctx)}, {}); - - func::FuncOp print_func; + auto inputType = op->getOperand(0).getType(); - /// If the Input type is scalar (F64) - if (inputType.isa()) + if(ShapedType shaped_type = mlir::dyn_cast(inputType)) { - std::string print_scalar_f64Str = "printF64"; - std::string print_newline_Str = "printNewline"; - if (!hasFuncDeclaration(module, "printF64")) - { - print_func = func::FuncOp::create(loc, print_scalar_f64Str, printScalarFunc, ArrayRef{}); - print_func.setPrivate(); - module.push_back(print_func); + auto unrankedMemrefType = mlir::UnrankedMemRefType::get(shaped_type.getElementType(), 0); + auto printTensor = FunctionType::get(ctx, {unrankedMemrefType}, {}); - if (!hasFuncDeclaration(module, "printNewline")) - { - auto printNewLineFunc = FunctionType::get(ctx, {}, {}); - func::FuncOp print_newline = func::FuncOp::create(loc, print_newline_Str, printNewLineFunc, ArrayRef{}); - print_newline.setPrivate(); - module.push_back(print_newline); - } + std::string comet_print; //_f64Str = "comet_print_memref_f64"; + if(shaped_type.getElementType().isF32()) + { + comet_print = "comet_print_memref_f32"; } - rewriter.create(loc, print_scalar_f64Str, SmallVector{}, ValueRange{op->getOperand(0)}); - rewriter.create(loc, print_newline_Str, SmallVector{}, ValueRange{}); - } - else - { - std::string comet_print_f64Str = "comet_print_memref_f64"; - if (!hasFuncDeclaration(module, comet_print_f64Str)) + else if(shaped_type.getElementType().isF64()) + { + comet_print = "comet_print_memref_f64"; + } + else if(shaped_type.getElementType().isInteger(64)) + { + comet_print = "comet_print_memref_i64"; + } + else if(shaped_type.getElementType().isIndex()) + { + comet_print = "comet_print_memref_index"; + } + else if(shaped_type.getElementType().isInteger(32)) + { + comet_print = "comet_print_memref_i32"; + } + else + { + assert(false && "Unexpected type to print"); + } + + + if (!hasFuncDeclaration(module, comet_print)) { - print_func = func::FuncOp::create(loc, comet_print_f64Str, printTensorF64Func, ArrayRef{}); + func::FuncOp print_func = func::FuncOp::create(loc, comet_print, printTensor, ArrayRef{}); print_func.setPrivate(); module.push_back(print_func); } - if (inputType.isa()) + if (isa(inputType)) { auto alloc_op = cast(op->getOperand(0).getDefiningOp()); comet_vdump(alloc_op); - auto u = rewriter.create(loc, unrankedMemrefType_f64, alloc_op); - rewriter.create(loc, comet_print_f64Str, SmallVector{}, ValueRange{u}); + auto u = rewriter.create(loc, unrankedMemrefType, alloc_op); + rewriter.create(loc, comet_print, SmallVector{}, ValueRange{u}); + } + else if (isa(inputType)) + { + auto rhs = op->getOperand(0); + auto tensor_type = llvm::cast(inputType); + auto memref_type = MemRefType::get(tensor_type.getShape(), tensor_type.getElementType()); + auto buffer = rewriter.create(loc, memref_type, rhs); + auto u = rewriter.create(loc, unrankedMemrefType, buffer); + rewriter.create(loc, comet_print, SmallVector{}, ValueRange{u}); } else { - /// If the Input type is tensor - if (inputType.isa()) - { - auto rhs = op->getOperand(0).getDefiningOp(); - auto alloc_op = cast(rhs->getOperand(0).getDefiningOp()); - comet_vdump(alloc_op); - auto u = rewriter.create(loc, unrankedMemrefType_f64, alloc_op); - rewriter.create(loc, comet_print_f64Str, SmallVector{}, ValueRange{u}); - } - else if (inputType.isa()) - { - std::string comet_print_i64Str = "comet_print_memref_i64"; - if (!hasFuncDeclaration(module, comet_print_i64Str)) - { - print_func = func::FuncOp::create(loc, comet_print_i64Str, printTensorIndexFunc, ArrayRef{}); - print_func.setPrivate(); - module.push_back(print_func); - } + llvm::errs() << __FILE__ << " " << __LINE__ << "Unknown Data type\n"; + } + } + /// If the Input type is scalar (F64) + else if (isa(inputType)) + { + std::string print_scalar; + if(inputType.isF64()) + { + print_scalar = "printF64"; + } + else if (inputType.isF32()) + { + print_scalar = "printF32"; + } + else if (inputType.isIndex()) + { + print_scalar = "printI64"; + } + else + { + assert(false && "Unsupported float type"); + } + FunctionType printScalarFunc = FunctionType::get(ctx, {inputType}, {}); - auto sp_op = cast(op->getOperand(0).getDefiningOp()); - Type unrankedMemref_index = mlir::UnrankedMemRefType::get(indexType, 0); - - auto rhs = op->getOperand(0).getDefiningOp(); - for (int rsize = 0; rsize < sp_op.getDimArrayCount(); rsize += 2) - { - /// accessing xD_pos array and creating cast op for its alloc - auto xD_pos = rhs->getOperand(rsize).getDefiningOp(); - auto alloc_rhs = cast(xD_pos->getOperand(0).getDefiningOp()); - auto u = rewriter.create(loc, unrankedMemref_index, alloc_rhs); - rewriter.create(loc, comet_print_i64Str, SmallVector{}, ValueRange{u}); - - /// accessing xD_crd array and creating cast op for its alloc - auto xD_crd = rhs->getOperand(rsize + 1).getDefiningOp(); - alloc_rhs = cast(xD_crd->getOperand(0).getDefiningOp()); - u = rewriter.create(loc, unrankedMemref_index, alloc_rhs); - rewriter.create(loc, comet_print_i64Str, SmallVector{}, ValueRange{u}); - } + std::string print_newline_Str = "printNewline"; + if (!hasFuncDeclaration(module, print_scalar)) + { + func::FuncOp print_func = func::FuncOp::create(loc, print_scalar, printScalarFunc, ArrayRef{}); + print_func.setPrivate(); + module.push_back(print_func); - auto xD_value = rhs->getOperand(sp_op.getValueArrayPos()).getDefiningOp(); - auto alloc_rhs = cast(xD_value->getOperand(0).getDefiningOp()); - auto u = rewriter.create(loc, unrankedMemrefType_f64, alloc_rhs); - rewriter.create(loc, comet_print_f64Str, SmallVector{}, ValueRange{u}); + if (!hasFuncDeclaration(module, "printNewline")) + { + auto printNewLineFunc = FunctionType::get(ctx, {}, {}); + func::FuncOp print_newline = func::FuncOp::create(loc, print_newline_Str, printNewLineFunc, ArrayRef{}); + print_newline.setPrivate(); + module.push_back(print_newline); } - else - llvm::errs() << __FILE__ << " " << __LINE__ << "Unknown Data type\n"; } + rewriter.create(loc, print_scalar, SmallVector{}, ValueRange{op->getOperand(0)}); + rewriter.create(loc, print_newline_Str, SmallVector{}, ValueRange{}); + } + else + { + assert(false && "Unexpected type to print"); } /// Notify the rewriter that this operation has been removed. @@ -275,7 +290,6 @@ namespace return success(); } }; - } /// end anonymous namespace. /// This is a partial lowering to linear algebra of the tensor algebra operations that are @@ -334,7 +348,7 @@ void LateLoweringPass::runOnOperation() patterns.insert(&getContext()); - + /// With the target and rewrite patterns defined, we can now attempt the /// conversion. The conversion will signal failure if any of our `illegal` /// operations were not converted successfully. @@ -349,3 +363,122 @@ std::unique_ptr mlir::comet::createLateLoweringPass() { return std::make_unique(); } + + +namespace +{ + struct BufferizeFunc + : public PassWrapper> + { + MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(BufferizeFunc) + void runOnOperation() override; + }; +} + + +void BufferizeFunc::runOnOperation() +{ + ModuleOp module = getOperation(); + // func::FuncOp function = getOperation(); + mlir::OpBuilder builder(module.getContext()); + + for(func::FuncOp function: module.getOps()) + { + std::vector newTypes; + std::vector resultTypes; + for(auto arg: function.getArguments()) + { + if(RankedTensorType ttype = mlir::dyn_cast(arg.getType())) + { + auto newType = MemRefType::get(ttype.getShape(), ttype.getElementType()); + newTypes.push_back(newType); + } + else + { + newTypes.push_back(arg.getType()); + } + } + for(auto res: function.getResultTypes()) + { + if(RankedTensorType ttype = mlir::dyn_cast(res)) + { + auto newType = MemRefType::get(ttype.getShape(), ttype.getElementType()); + resultTypes.push_back(newType); + } + else + { + resultTypes.push_back(res); + } + } + + if(newTypes.empty() && resultTypes.empty()) + { + return; + } + assert(function.getArguments().size() == newTypes.size()); + auto newFuncType = builder.getFunctionType(newTypes, resultTypes); + function.setType(newFuncType); + + if(!function.isDeclaration()) + { + builder.setInsertionPointToStart(&function.getFunctionBody().front()); + for(auto arg: function.getFunctionBody().getArguments()) + { + if(RankedTensorType ttype = mlir::dyn_cast(arg.getType())) + { + auto newType = MemRefType::get(ttype.getShape(), ttype.getElementType()); + arg.setType(newType); + ToTensorOp to_tensor = builder.create(function->getLoc(), ttype, arg, true, true); + arg.replaceAllUsesExcept(to_tensor, to_tensor); + } + } + for(auto& arg: function.getFunctionBody().front().getTerminator()->getOpOperands()) + { + builder.setInsertionPoint(function.getFunctionBody().front().getTerminator()); + if(RankedTensorType ttype = mlir::dyn_cast(arg.get().getType())) + { + auto newType = MemRefType::get(ttype.getShape(), ttype.getElementType()); + ToMemrefOp to_memref = builder.create(function->getLoc(), newType, arg.get()); + arg.assign(to_memref); + } + } + } + + module->walk([&](mlir::func::CallOp call) + { + if(call.getCallee() == function.getName()) + { + builder.setInsertionPoint(call); + llvm::SmallVector newOperands; + int argIdx = 0; + for (auto oldOperand : call->getOperands()) { + if (oldOperand.getType() != newTypes[argIdx]) { + auto newOperad = builder.create(function->getLoc(), newTypes[argIdx], oldOperand); + newOperands.push_back(newOperad); + } else { + newOperands.push_back(oldOperand); + } + argIdx++; + } + + auto new_call = builder.create(call.getLoc(), function, newOperands); + builder.setInsertionPointAfter(new_call); + for(auto res: llvm::zip(call.getResults(), new_call->getResults())) + { + if(RankedTensorType ttype = mlir::dyn_cast(std::get<0>(res).getType())) + { + ToTensorOp to_tensor = builder.create(function->getLoc(), ttype, std::get<1>(res), true, true); + std::get<0>(res).replaceAllUsesExcept(to_tensor, to_tensor); + } + } + + call.erase(); + } + }); + } +} + +std::unique_ptr mlir::comet::createTABufferizeFunc() +{ + return std::make_unique(); +} \ No newline at end of file diff --git a/lib/Conversion/TensorAlgebraToSCF/LowerFunc.cpp b/lib/Conversion/TensorAlgebraToSCF/LowerFunc.cpp index 93c9575a..f7b9fedb 100644 --- a/lib/Conversion/TensorAlgebraToSCF/LowerFunc.cpp +++ b/lib/Conversion/TensorAlgebraToSCF/LowerFunc.cpp @@ -1,3 +1,24 @@ +// +// Copyright 2022 Battelle Memorial Institute +// +// Redistribution and use in source and binary forms, with or without modification, +// are permitted provided that the following conditions are met: +// +// 1. Redistributions of source code must retain the above copyright notice, this list of conditions +// and the following disclaimer. +// +// 2. Redistributions in binary form must reproduce the above copyright notice, this list of conditions +// and the following disclaimer in the documentation and/or other materials provided with the distribution. +// +// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED +// WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR +// PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, +// INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE +// GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, +// WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +// OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. +// + #include "mlir/IR/BuiltinDialect.h" #include "comet/Dialect/TensorAlgebra/IR/TADialect.h" #include "comet/Dialect/TensorAlgebra/Passes.h" diff --git a/lib/Conversion/TensorAlgebraToSCF/LowerPCToLoops.cpp b/lib/Conversion/TensorAlgebraToSCF/LowerPCToLoops.cpp deleted file mode 100644 index 0d9be872..00000000 --- a/lib/Conversion/TensorAlgebraToSCF/LowerPCToLoops.cpp +++ /dev/null @@ -1,341 +0,0 @@ -//===- LowerPCToLoops.cpp ------===// -// -// Copyright 2022 Battelle Memorial Institute -// -// Redistribution and use in source and binary forms, with or without modification, -// are permitted provided that the following conditions are met: -// -// 1. Redistributions of source code must retain the above copyright notice, this list of conditions -// and the following disclaimer. -// -// 2. Redistributions in binary form must reproduce the above copyright notice, this list of conditions -// and the following disclaimer in the documentation and/or other materials provided with the distribution. -// -// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED -// WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR -// PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, -// INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE -// GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, -// WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE -// OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. -// -//===----------------------------------------------------------------------===// -// -// This file implements a lowering of some programming constructs such as for-loops, etc. -//===----------------------------------------------------------------------===// - -#include "comet/Dialect/TensorAlgebra/IR/TADialect.h" -#include "comet/Dialect/TensorAlgebra/Passes.h" -#include "comet/Dialect/IndexTree/IR/IndexTreeDialect.h" -#include "comet/Dialect/Utils/Utils.h" - -#include "mlir/Dialect/Func/IR/FuncOps.h" -#include "mlir/Pass/Pass.h" - -#include -#include - -using namespace mlir; -using namespace mlir::arith; -using namespace mlir::tensorAlgebra; -using namespace mlir::indexTree; - -#define DEBUG_TYPE "PC-lowering" - -// *********** For debug purpose *********// -// #define COMET_DEBUG_MODE -#include "comet/Utils/debug.h" -#undef COMET_DEBUG_MODE -// *********** For debug purpose *********// - -namespace -{ - struct PCToLoopsLoweringPass - : public PassWrapper> - { - MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(PCToLoopsLoweringPass) - void runOnOperation() override; - - // lowers ForLoopBeginOp and ForLoopEndOp to scf.for, one loop at a time. - // pre-condition: output from ProcessLoopOps(); - void PCToLoopsLowering(tensorAlgebra::ForLoopBeginOp op, tensorAlgebra::ForLoopEndOp op_end, std::vector &listOps); - - // find all ops to be placed inside the loop body. - std::vector ProcessLoopOps(tensorAlgebra::ForLoopBeginOp op_start, tensorAlgebra::ForLoopEndOp op_end); - - // replicates Ops for loop-body - // pre-condition: some ops require previously generated ops. - void replicateOpsForLoopBody(Location loc, OpBuilder &builder, Operation *op, Value &rhs, Value &lhs, Value &compute, Value &indices, Value &transpose); - }; -} // end anonymous namespace. - -// find all ops to be placed inside the loop body. -std::vector PCToLoopsLoweringPass::ProcessLoopOps(tensorAlgebra::ForLoopBeginOp op_start, tensorAlgebra::ForLoopEndOp op_end) -{ - comet_debug() << "START: Pre-processing to detect ops to be placed inside loop body.\n"; - comet_pdump(op_start); - - std::vector loop_blk; - - auto *B = op_start.getOperation()->getBlock(); - bool match = false; - // collect all ops to for `one` loop body - for (Operation &op : *B) - { - if (isa(&op)) - { - // check for match - if (cast(op).getIterator() == op_start.getIterator()) - { - match = true; - continue; // skip this iteration or op; - } - } - - if (isa(&op)) - { - if (cast(op) == op_end) - { - match = false; - break; // we are done with match of `one` loop-body - } - } - - if (match) - { // add to list of ops to be replicated - loop_blk.push_back(&op); - } - } - - comet_debug() << "END: Pre-processing to detect ops to be placed inside loop body.\n"; - return loop_blk; -} - -// replicates Ops for loop-body -// pre-condition: some ops require previously generated ops. -void PCToLoopsLoweringPass::replicateOpsForLoopBody(Location loc, OpBuilder &builder, Operation *op, - Value &rhs, Value &lhs, Value &compute, Value &indices, Value &transpose) -{ - IntegerType i64Type = IntegerType::get(builder.getContext(), 64); - - if (isa(op)) - { - comet_debug() << "creating TransposeOp\n"; - tensorAlgebra::TransposeOp ta_transpose_op = llvm::dyn_cast(op); - std::vector lhs_lbls_value = {ta_transpose_op.getOperand(1), ta_transpose_op.getOperand(2)}; - transpose = builder.create(loc, ta_transpose_op.getOperand(0).getType(), ta_transpose_op.getOperand(0), - lhs_lbls_value, ta_transpose_op.getIndexingMaps(), ta_transpose_op.getFormats()); - } - - // TensorSetOp goes with TransposeOp at this stage of the lowering. - if (isa(op) && transpose != NULL) - { - comet_debug() << "creating SetOp after TransposeOp\n"; - tensorAlgebra::TensorSetOp ta_set_op = llvm::dyn_cast(op); - builder.create(loc, transpose, ta_set_op.getOperand(1)); - } - - // create IndexTreeComputeRHSOp, no dependency to earlier replications - if (isa(op)) - { - indexTree::IndexTreeComputeRHSOp it_compute_rhs_op = llvm::dyn_cast(op); - ArrayAttr op_formats_ArrayAttr = it_compute_rhs_op.getAllFormats(); - ArrayAttr op_perms_ArrayAttr = it_compute_rhs_op.getAllPerms(); - - rhs = builder.create(loc, mlir::UnrankedTensorType::get(builder.getF64Type()), - it_compute_rhs_op->getOperands(), // tensors - op_perms_ArrayAttr, op_formats_ArrayAttr); - } - - // create IndexTreeComputeLHSOp, no dependency to earlier replications - if (isa(op)) - { - indexTree::IndexTreeComputeLHSOp it_compute_lhs_op = llvm::dyn_cast(op); - ArrayAttr op_formats_ArrayAttr = it_compute_lhs_op.getAllFormats(); - ArrayAttr op_perms_ArrayAttr = it_compute_lhs_op.getAllPerms(); - - lhs = builder.create(loc, mlir::UnrankedTensorType::get(builder.getF64Type()), - it_compute_lhs_op->getOperands(), /// tensors - op_perms_ArrayAttr, op_formats_ArrayAttr); - } - - /// create IndexTreeComputeOp only if rhs and lhs are ready - if (isa(op) && rhs != NULL && lhs != NULL) - { - indexTree::IndexTreeComputeOp it_compute_op = llvm::dyn_cast(op); - compute = builder.create(loc, i64Type, rhs, lhs, - it_compute_op.getCompWorkspOpt(), it_compute_op.getSemiring(), it_compute_op.getMaskType()); - } - - /// create IndexTreeIndicesOp from existing IndexTreeIndicesOp (checks condition: indices != NULL) - if (isa(op) && indices != NULL) - { - indexTree::IndexTreeIndicesOp it_indices_op = llvm::dyn_cast(op); - Value indices_op_new = builder.create(loc, i64Type, indices, it_indices_op.getIndices(), it_indices_op.getIteratorType()); - - indices = indices_op_new; /// for subsequent IndexTreeIndicesOp creation - } - - /// create the first instance of IndexTreeIndicesOp from IndexTreeComputeOp - if (isa(op) && compute != NULL && indices == NULL) - { - indexTree::IndexTreeIndicesOp it_indices_op = llvm::dyn_cast(op); - indices = builder.create(loc, i64Type, compute, it_indices_op.getIndices(), it_indices_op.getIteratorType()); - } - - if (isa(op) && indices != NULL) - { - builder.create(loc, i64Type, indices); - } -} - -/// lowers ForLoopBeginOp and ForLoopEndOp to scf.for, one loop at a time. -void PCToLoopsLoweringPass::PCToLoopsLowering(tensorAlgebra::ForLoopBeginOp op, tensorAlgebra::ForLoopEndOp op_end, std::vector &listOps) -{ - comet_debug() << "PCToLoopsLowering start\n"; - OpBuilder builder(op); - - comet_pdump(op); - comet_pdump(op_end); - - auto loc = op->getLoc(); - auto ForLoopStart = cast(op); - - /// get info of loop - auto upperBound = ForLoopStart.getMax(); - auto lowerBound = ForLoopStart.getMin(); - auto step = ForLoopStart.getStep(); - - auto loop = builder.create(op_end->getLoc(), lowerBound, upperBound, step); - comet_vdump(loop); - - auto insertPt = builder.saveInsertionPoint(); - builder.setInsertionPointToStart(loop.getBody()); - - /// loop-body: listOps contains all ops obtained thru ProcessLoopOps() - /// to be placed inside 'one' loop-body. - /// TODO(gkestor): verify all ops are covered here! - Value new_rhs_op, new_lhs_op, new_compute_op, indices_op; - Value transpose_op; - std::vector loop_bounds; /// carries loop bounds - for (unsigned int i = 0; i < listOps.size(); i++) - { - comet_pdump(listOps[i]); - - replicateOpsForLoopBody(loc, builder, listOps[i], new_rhs_op, new_lhs_op, new_compute_op, indices_op, transpose_op); - - /// nested loop - if (isa(listOps[i])) - { - ConstantIndexOp constant_index = llvm::dyn_cast(listOps[i]); - loop_bounds.push_back(builder.create(loc, constant_index.value())); - } - - /// scf::for op - /// the inner loop bodies have already been created. preserve them. - /// the loop_bounds array should have lowerBound, upperBound and step Values. - if (isa(listOps[i]) && loop_bounds.size() == 3) - { - scf::ForOp scf_for_op = llvm::dyn_cast(listOps[i]); - scf::ForOp nested_scf_for_op = builder.create(loc, loop_bounds[0], loop_bounds[1], loop_bounds[2]); - - auto insertPt_nested = builder.saveInsertionPoint(); - builder.setInsertionPointToStart(nested_scf_for_op.getBody()); - - Block *B = scf_for_op.getBody(); - Value nested_rhs_op, nested_lhs_op, nested_compute_op, nested_indices_op; - Value nested_transpose_op; - comet_debug() << "going to replicate ops inside nested loop...\n"; - for (Operation &op_for : *B) - { - replicateOpsForLoopBody(loc, builder, &op_for, nested_rhs_op, nested_lhs_op, nested_compute_op, nested_indices_op, nested_transpose_op); - } - /// need to restore the insertion point to the previous point - builder.restoreInsertionPoint(insertPt_nested); - builder.setInsertionPoint(op_end); /// TODO(gkestor): need to re-visit this for nested loops. - } - } - - /// remove old ops, since now we are done with the clone inside loop-body. - /// this is done in reverse order. - comet_debug() << "Removing ops that have been cloned\n"; - for (unsigned int i = 0; i < listOps.size(); i++) - { - comet_pdump(listOps[listOps.size() - i - 1]); - listOps[listOps.size() - i - 1]->erase(); - } - - /// need to restore the insertion point to the previous point - builder.restoreInsertionPoint(insertPt); - builder.setInsertionPoint(op_end); - - /// remove ForLoopBeginOp and ForLoopEndOp - op->erase(); - op_end->erase(); - - comet_debug() << "PCToLoopsLowering end\n"; -} - -void PCToLoopsLoweringPass::runOnOperation() -{ - comet_debug() << "start PCToLoopsLoweringPass\n"; - - func::FuncOp function = getOperation(); - - std::vector startOps; - std::vector endOps; - - /// collect all the loops (begin and ends) here in vector data-structure. e.g., - /// for-start1 (): - /// for-start2 (): - /// do_work2(); - /// end2 - /// do_work1(); - /// end1 - - /// vector: for-start1, for-start2 - /// vector: end2, end1 - - for (Block &B : function.getBody()) - { - for (Operation &op : B) - { - if (isa(&op)) - { - startOps.push_back(cast(op)); - } - if (isa(&op)) - { - endOps.push_back(cast(op)); - } - } - } - - /// if there are no for-loops, quit. - if (startOps.empty()) - return; - - /// the size of the two datastructure should be same - assert(startOps.size() == endOps.size() && "the for-begins must match the ends"); - std::vector opList; - for (unsigned int i = 0; i < startOps.size(); i++) - { - /// start with inner most loop and move outwards. - opList = ProcessLoopOps(startOps[startOps.size() - i - 1], endOps[i]); /// for-1, end-1 - PCToLoopsLowering(startOps[startOps.size() - i - 1], endOps[i], opList); - opList.clear(); /// clear for next round. - } - startOps.clear(); - endOps.clear(); - - /// debug - /// auto module = function.getOperation()->getParentOfType(); - - comet_debug() << "end PCToLoopsLoweringPass\n"; -} - -/// Create a pass for lowering programming constructs -std::unique_ptr mlir::comet::createPCToLoopsLoweringPass() -{ - return std::make_unique(); -} \ No newline at end of file diff --git a/lib/Conversion/TensorAlgebraToSCF/SparseTensorConversionPass.cpp b/lib/Conversion/TensorAlgebraToSCF/SparseTensorConversionPass.cpp new file mode 100644 index 00000000..a19d077a --- /dev/null +++ b/lib/Conversion/TensorAlgebraToSCF/SparseTensorConversionPass.cpp @@ -0,0 +1,1459 @@ +// +// Copyright 2022 Battelle Memorial Institute +// +// Redistribution and use in source and binary forms, with or without modification, +// are permitted provided that the following conditions are met: +// +// 1. Redistributions of source code must retain the above copyright notice, this list of conditions +// and the following disclaimer. +// +// 2. Redistributions in binary form must reproduce the above copyright notice, this list of conditions +// and the following disclaimer in the documentation and/or other materials provided with the distribution. +// +// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED +// WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR +// PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, +// INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE +// GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, +// WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +// OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. +// + +#include "comet/Dialect/IndexTree/IR/IndexTreeDialect.h" +#include "comet/Dialect/Utils/Utils.h" +#include "comet/Dialect/TensorAlgebra/IR/TADialect.h" +#include "comet/Conversion/IndexTreeToSCF/IndexTreeToSCF.h" +#include "comet/Conversion/TensorAlgebraToSCF/TensorAlgebraToSCF.h" +#include "comet/Dialect/IndexTree/Patterns.h" + +#include "mlir/Dialect/Arith/IR/Arith.h" +#include "mlir/Dialect/SCF/IR/SCF.h" +#include "mlir/Dialect/Index/IR/IndexAttrs.h" +#include "mlir/Dialect/Index/IR/IndexOps.h" +#include "mlir/Dialect/Index/IR/IndexDialect.h" +#include "mlir/Dialect/MemRef/IR/MemRef.h" +#include "mlir/Dialect/Bufferization/IR/Bufferization.h" +#include "mlir/Dialect/Func/IR/FuncOps.h" +#include "mlir/Dialect/Tensor/IR/Tensor.h" +#include "mlir/IR/Builders.h" +#include "mlir/IR/BuiltinAttributes.h" +#include "mlir/IR/BuiltinDialect.h" +#include "mlir/IR/BuiltinOps.h" +#include "mlir/IR/BuiltinTypeInterfaces.h" +#include "mlir/IR/BuiltinTypes.h" +#include "mlir/IR/TypeRange.h" +#include "mlir/IR/ValueRange.h" +#include "mlir/Support/LLVM.h" +#include "mlir/Support/LogicalResult.h" +#include "mlir/Transforms/DialectConversion.h" +#include "mlir/Dialect/Func/Transforms/DecomposeCallGraphTypes.h" +#include "mlir/Dialect/Func/Transforms/FuncConversions.h" +#include "mlir/Dialect/SCF/Transforms/Patterns.h" +#include "mlir/Pass/Pass.h" + +#include "llvm/ADT/SmallVector.h" +#include "llvm/Support/Debug.h" +#include "llvm/Support/ScopedPrinter.h" +#include + +using namespace mlir; +using namespace mlir::tensorAlgebra; + +#define DEBUG_TYPE "sparse_tensor" + +namespace mlir { + namespace comet{ + #define GEN_PASS_DEF_SPARSETENSORCONVERSIONPASS + #include "comet/Conversion/Passes.h.inc" + } +} + +/** Helper structures to turn sparse tensor into pointers */ +struct Dimension { + Value dim_size; + Value insert_pos; + TensorFormatEnum format; + + Value pos; + Value crd; + + bool has_block; + Value block_pos; + Value block_crd; + +}; + +struct SparseTensor { + Value dim_sizes; + SmallVector dims; + Value vals; + Value val_size; +}; + +struct Workspace { + Value workspace; + Value mark_value; + Value mark_array; + Value num_crds; + Value crds; +}; + +static bool unpack_sparse_tensor(Value sparse_tensor, SparseTensor& result) +{ + /** Helper function to turn arguments from an unrealized cast to sparse tensor */ + if (auto cast = + sparse_tensor.getDefiningOp()) { + SparseTensorType type = llvm::dyn_cast(sparse_tensor.getType()); + if(!type) + return false; + + auto format = type.getFormat(); + auto dim_sizes = type.getDims(); + auto cur_arg = cast.getInputs().begin(); + result.dim_sizes = *cur_arg; + ++cur_arg; + + for(unsigned i = 0; i < dim_sizes.size(); i++){ + Dimension d; + d.insert_pos = *cur_arg; + cur_arg++; + d.format = (TensorFormatEnum)format[2 * i]; + switch(d.format){ + case TensorFormatEnum::D: { + d.pos = *cur_arg; + cur_arg++; + break; + } + case TensorFormatEnum::CU: + case TensorFormatEnum::CN: { + d.pos = *cur_arg; + cur_arg++; + + d.crd = *cur_arg; + cur_arg++; + break; + } + case TensorFormatEnum::S: { + d.crd = *cur_arg; + cur_arg++; + break; + } + default: { + assert(false && "Could not unpack unknown format to sparse tensor."); + } + } + result.dims.push_back(d); + } + result.vals = *cur_arg; + return true; + } + + return false; +} + +static void pack_sparse_tensor(SparseTensorType type, SparseTensor& sparse_tensor, SmallVectorImpl& result) +{ + result.push_back(sparse_tensor.dim_sizes); + for(Dimension d : sparse_tensor.dims) + { + result.push_back(d.insert_pos); + switch(d.format){ + case TensorFormatEnum::D: { + result.push_back(d.pos); + break; + } + case TensorFormatEnum::CU: + case TensorFormatEnum::CN: { + result.push_back(d.pos); + result.push_back(d.crd); + break; + } + case TensorFormatEnum::S: { + result.push_back(d.crd); + break; + } + default: { + assert(false && "Could not unpack unknown format to sparse tensor."); + } + } + } + result.push_back(sparse_tensor.vals); + return; +} + +static bool unpack_workspace(Value workspace_val, Workspace& result) +{ + if (auto cast = + workspace_val.getDefiningOp()) { + if(!llvm::isa(workspace_val.getType())){ + return false; + } + auto cur_arg = cast.getInputs().begin(); + result.workspace = *cur_arg; + cur_arg++; + result.mark_value = *cur_arg; + cur_arg++; + result.mark_array = *cur_arg; + cur_arg++; + result.num_crds = *cur_arg; + cur_arg++; + result.crds = *cur_arg; + return true; + } + return false; + +} +static void pack_workspace(WorkspaceType type, Workspace& workspace, SmallVectorImpl& result) +{ + result.push_back(workspace.workspace); + result.push_back(workspace.mark_value); + result.push_back(workspace.mark_array); + result.push_back(workspace.num_crds); + result.push_back(workspace.crds); +} + +namespace { +class ConvertSpTensorConstructOp + : public OpConversionPattern { + using OpConversionPattern::OpConversionPattern; + LogicalResult + matchAndRewrite(SparseTensorConstructOp op, OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const override { + + SparseTensor sp_tensor; + SparseTensorType sp_tensor_type = llvm::cast(op->getResult(0).getType()); + auto dims = op.getDims(); + auto crd = op.getCrdIndices(); + auto pos = op.getPosIndices(); + auto vals = op.getVals(); + unsigned rank = sp_tensor_type.getDims().size(); + sp_tensor.dim_sizes = dims; + for(unsigned i = 0; i < rank; i++) + { + Dimension d; + d.format = (TensorFormatEnum) sp_tensor_type.getFormat()[2 * i]; + d.insert_pos = rewriter.create(op.getLoc(), rewriter.getIndexType(), rewriter.getIndexAttr(0)); + d.pos = pos[i]; + d.crd = crd[i]; + sp_tensor.dims.push_back(d); + } + + sp_tensor.vals = vals; + + SmallVector cast_args; + pack_sparse_tensor(sp_tensor_type, sp_tensor, cast_args); + rewriter.replaceOpWithNewOp(op, sp_tensor_type, cast_args); + return success(); + } +}; + +class ConvertSpTensorAliasOp + : public OpConversionPattern { + using OpConversionPattern::OpConversionPattern; + LogicalResult + matchAndRewrite(SpTensorAliasOp op, OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const override { + rewriter.replaceOp(op, adaptor.getTensor()); + return success(); + } +}; + +class ConvertSpTensorInsertOp + : public OpConversionPattern { + using OpConversionPattern::OpConversionPattern; + ConvertSpTensorInsertOp(MLIRContext *context) + : OpConversionPattern(context) {} + LogicalResult + matchAndRewrite(tensorAlgebra::TensorInsertOp op, OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const override { + if(!llvm::isa(op.getTensor().getType())){ + return failure(); + } + + SparseTensor sp_tensor; + TensorInsertOpAdaptor insertAdpator = llvm::cast(adaptor); + if(!unpack_sparse_tensor(insertAdpator.getTensor(), sp_tensor)) { + return failure(); + } + + // Match successful! + auto loc = op.getLoc(); + Type index_type = rewriter.getIndexType(); + unsigned i = 0; + for(Dimension& dim : sp_tensor.dims) { + if(dim.format != TensorFormatEnum::D) { + Value crd_idx = insertAdpator.getPos()[i]; + Value crd = insertAdpator.getCrds()[i]; + Value crd_tensor = dim.crd; + RankedTensorType crd_tensorT = mlir::cast(dim.crd.getType()); + if(crd.getType() != crd_tensorT.getElementType()) + { + crd = rewriter.createOrFold(loc, crd_tensorT.getElementType(), crd); + } + + crd_tensor = rewriter.create( + loc, + crd_tensor.getType(), + crd, + crd_tensor, + crd_idx); + dim.crd = crd_tensor; + + // TODO: This is wrong if we insert the same crd multiple times but format is CU? + Value inc = rewriter.create(loc, index_type, rewriter.getIndexAttr(1)); + dim.insert_pos = rewriter.create(loc, index_type, dim.insert_pos, inc); + + /** TODO: Implement tensor resize */ + /** TODO: Insert into CSR only has to be done once per idx? */ + } + + i++; + } + Value vals = sp_tensor.vals; + Value val_idx = insertAdpator.getPos()[insertAdpator.getPos().size() - 1]; + vals = rewriter.create(loc, + vals.getType(), + insertAdpator.getValue(), + vals, + val_idx); + sp_tensor.vals = vals; + SparseTensorType sp_tensor_type = llvm::cast(op.getTensor().getType()); + + SmallVector cast_args; + pack_sparse_tensor(sp_tensor_type, sp_tensor, cast_args); + rewriter.replaceOpWithNewOp(op, sp_tensor_type, cast_args); + return success(); + } +}; + +class ConvertSpTensorExtractOp + : public OpConversionPattern { + using OpConversionPattern::OpConversionPattern; + ConvertSpTensorExtractOp(MLIRContext *context) + : OpConversionPattern(context) {} + LogicalResult + matchAndRewrite(tensorAlgebra::TensorExtractOp op, OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const override { + TensorExtractOpAdaptor extractAdaptor = llvm::cast(adaptor); + if(!llvm::isa(extractAdaptor.getTensor().getType())){ + return failure(); + } + + SparseTensor sp_tensor; + if(!unpack_sparse_tensor(extractAdaptor.getTensor(), sp_tensor)) { + return failure(); + } + + llvm::ScopedPrinter logger{llvm::dbgs()}; + LLVM_DEBUG({ + logger.startLine() << "Unpacked sparse tensor: " << extractAdaptor.getTensor().getDefiningOp() << "\n"; + }); + // Match successful! + auto loc = op.getLoc(); + // Type float_type = llvm::cast(sp_tensor.vals.getType()).getElementType(); + Value result = rewriter.create(loc, sp_tensor.vals, extractAdaptor.getPos()); + rewriter.replaceOp(op, {result}); + return success(); + } +}; + +class ConvertSpTensorGetCrd + : public OpConversionPattern { + using OpConversionPattern::OpConversionPattern; + ConvertSpTensorGetCrd(MLIRContext * context) + : OpConversionPattern(context) {} + + LogicalResult matchAndRewrite(SpTensorGetCrd op, OpAdaptor adaptor, ConversionPatternRewriter &rewriter) const final { + auto opAdaptor = llvm::cast(adaptor); + if(!llvm::isa(opAdaptor.getTensor().getType())){ + return failure(); + } + + SparseTensor sp_tensor; + if(!unpack_sparse_tensor(opAdaptor.getTensor(), sp_tensor)) { + return failure(); + } + rewriter.replaceOpWithNewOp(op, op.getType(), sp_tensor.dims[op.getDim()].crd, op.getIdx()); + return success(); + } +}; + +class ConvertSpTensorInsertCrd + : public OpConversionPattern { + using OpConversionPattern::OpConversionPattern; + ConvertSpTensorInsertCrd(MLIRContext *context) + : OpConversionPattern(context) {} + LogicalResult + matchAndRewrite(SpTensorInsertCrd op, OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const final { + SparseTensor sp_tensor; + auto opAdaptor = llvm::cast(adaptor); + if(!unpack_sparse_tensor(opAdaptor.getTensor(), sp_tensor)) { + return failure(); + } + + // Match successful! + auto loc = op.getLoc(); + Type index_type = rewriter.getIndexType(); + Dimension& dim = sp_tensor.dims[opAdaptor.getDim()]; + if(dim.format != TensorFormatEnum::D) { + Value crd_idx = opAdaptor.getIdx(); + Value crd = opAdaptor.getCrd(); + Value crd_tensor = dim.crd; + crd_tensor = rewriter.create(loc, + crd_tensor.getType(), + crd, + crd_tensor, + crd_idx); + dim.crd = crd_tensor; + + // Update tensor insert state + Value inc = rewriter.create(loc, index_type, rewriter.getIndexAttr(1)); + dim.insert_pos = rewriter.create(loc, index_type, dim.insert_pos, inc); + } + + SparseTensorType sp_tensor_type = llvm::cast(opAdaptor.getTensor().getType()); + SmallVector cast_args; + pack_sparse_tensor(sp_tensor_type, sp_tensor, cast_args); + rewriter.replaceOpWithNewOp(op, sp_tensor_type, cast_args); + return success(); + } +}; + +class ConvertSpTensorGetDimSize + : public OpConversionPattern { + using OpConversionPattern::OpConversionPattern; + ConvertSpTensorGetDimSize(MLIRContext *context) + : OpConversionPattern(context) {} + LogicalResult + matchAndRewrite(SpTensorGetDimSize op, OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const override { + SpTensorGetDimSizeAdaptor tensorAdaptor = llvm::cast(adaptor); + SparseTensor sp_tensor; + if(!unpack_sparse_tensor(tensorAdaptor.getTensor(), sp_tensor)) { + return failure(); + } + Value index = rewriter.create(op->getLoc(), tensorAdaptor.getDim()); + rewriter.replaceOpWithNewOp(op, sp_tensor.dim_sizes, index); + return success(); + } +}; + + +class ConvertSpTensorGetDimCrd + : public OpConversionPattern { + using OpConversionPattern::OpConversionPattern; + ConvertSpTensorGetDimCrd(MLIRContext *context) + : OpConversionPattern(context) {} + LogicalResult + matchAndRewrite(SpTensorGetDimCrd op, OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const override { + SpTensorGetDimCrdAdaptor tensorAdaptor = llvm::cast(adaptor); + SparseTensor sp_tensor; + if(!unpack_sparse_tensor(tensorAdaptor.getTensor(), sp_tensor)) { + return failure(); + } + SparseTensorType spTensorType = mlir::cast(tensorAdaptor.getTensor().getType()); + + + if(sp_tensor.dims[tensorAdaptor.getDim()].crd == nullptr) + { + auto zero = rewriter.create(op->getLoc(), 0); + rewriter.replaceOp(op, rewriter.create(op->getLoc(), RankedTensorType::get({ShapedType::kDynamic,}, spTensorType.getIndicesType()), ValueRange(zero))); + } + else + { + rewriter.replaceOp(op, {sp_tensor.dims[tensorAdaptor.getDim()].crd}); + } + return success(); + } +}; + + +/// TODO: Implement this conversion to actually handle blocks/tiles +class ConvertSpTensorGetDimBlockPos + : public OpConversionPattern { + using OpConversionPattern::OpConversionPattern; + ConvertSpTensorGetDimBlockPos(MLIRContext *context) + : OpConversionPattern(context) {} + LogicalResult + matchAndRewrite(SpTensorGetDimBlockPos op, OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const override { + SpTensorGetDimBlockPosAdaptor tensorAdaptor = llvm::cast(adaptor); + + + // SparseTensor sp_tensor; + // if(!unpack_sparse_tensor(tensorAdaptor.getTensor(), sp_tensor)) { + // return failure(); + // } + SparseTensorType spTensorType = mlir::cast(tensorAdaptor.getTensor().getType()); + + auto zero = rewriter.create(op->getLoc(), 0); + rewriter.replaceOp(op, rewriter.create(op->getLoc(), RankedTensorType::get({ShapedType::kDynamic,}, spTensorType.getIndicesType()), ValueRange(zero))); + + return success(); + } +}; + +/// TODO: Implement this conversion to actually handle blocks/tiles +class ConvertSpTensorGetDimBlockCrd + : public OpConversionPattern { + using OpConversionPattern::OpConversionPattern; + ConvertSpTensorGetDimBlockCrd(MLIRContext *context) + : OpConversionPattern(context) {} + LogicalResult + matchAndRewrite(SpTensorGetDimBlockCrd op, OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const override { + SpTensorGetDimBlockCrdAdaptor tensorAdaptor = llvm::cast(adaptor); + + + // SparseTensor sp_tensor; + // if(!unpack_sparse_tensor(tensorAdaptor.getTensor(), sp_tensor)) { + // return failure(); + // } + SparseTensorType spTensorType = mlir::cast(tensorAdaptor.getTensor().getType()); + + auto zero = rewriter.create(op->getLoc(), 0); + rewriter.replaceOp(op, rewriter.create(op->getLoc(), RankedTensorType::get({ShapedType::kDynamic,}, spTensorType.getIndicesType()), ValueRange(zero))); + + return success(); + } +}; + +class ConvertSpTensorGetVals + : public OpConversionPattern { + using OpConversionPattern::OpConversionPattern; + ConvertSpTensorGetVals(MLIRContext *context) + : OpConversionPattern(context) {} + LogicalResult + matchAndRewrite(SpTensorGetVals op, OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const override { + SpTensorGetValsAdaptor tensorAdaptor = llvm::cast(adaptor); + SparseTensor sp_tensor; + if(!unpack_sparse_tensor(tensorAdaptor.getTensor(), sp_tensor)) { + return failure(); + } + rewriter.replaceOp(op, {sp_tensor.vals}); + return success(); + } +}; + + +class ConvertSpTensorGetDimPos + : public OpConversionPattern { + using OpConversionPattern::OpConversionPattern; + ConvertSpTensorGetDimPos(MLIRContext *context) + : OpConversionPattern(context) {} + LogicalResult + matchAndRewrite(SpTensorGetDimPos op, OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const override { + SpTensorGetDimPosAdaptor tensorAdaptor = llvm::cast(adaptor); + SparseTensorType spTensorType = mlir::cast(tensorAdaptor.getTensor().getType()); + SparseTensor sp_tensor; + if(!unpack_sparse_tensor(tensorAdaptor.getTensor(), sp_tensor)) { + return failure(); + } + if(sp_tensor.dims[tensorAdaptor.getDim()].pos == nullptr) + { + auto zero = rewriter.create(op->getLoc(), 0); + rewriter.replaceOp(op, rewriter.create(op->getLoc(), RankedTensorType::get({ShapedType::kDynamic,}, spTensorType.getIndicesType()), ValueRange(zero))); + } + else + { + rewriter.replaceOp(op, {sp_tensor.dims[tensorAdaptor.getDim()].pos}); + } + + return success(); + } +}; + +class ConvertSpTensorFindPos + : public OpConversionPattern { + using OpConversionPattern::OpConversionPattern; + ConvertSpTensorFindPos(MLIRContext *context) + : OpConversionPattern(context) {} + LogicalResult + matchAndRewrite(TensorFindPos op, OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const override { + TensorFindPosAdaptor tensorAdaptor = llvm::cast(adaptor); + SparseTensor sp_tensor; + if(!unpack_sparse_tensor(tensorAdaptor.getTensor(), sp_tensor)) { + return failure(); + } + + if(tensorAdaptor.getIsLinear()) + { + rewriter.replaceOp(op, {sp_tensor.dims[tensorAdaptor.getDim()].insert_pos}); + } else { + assert(false && "Lowering non-unique inserts is not yet supported, please use workspace transform"); + } + + return success(); + } +}; + +class ConvertAllocWorkspaceOp + : public OpConversionPattern { + using OpConversionPattern::OpConversionPattern; + ConvertAllocWorkspaceOp(MLIRContext *context) + : OpConversionPattern(context) {} + LogicalResult + matchAndRewrite(AllocWorkspaceOp op, OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const final { + auto alloc_adaptor = llvm::cast(adaptor); + auto loc = op.getLoc(); + Type index_type = rewriter.getIndexType(); + + Value sp_tensor = op.getTensor(); + auto sp_tensor_type = llvm::cast(sp_tensor.getType()); + auto dims = alloc_adaptor.getDims(); + SmallVector dim_attrs(dims.size(), ShapedType::kDynamic); + SmallVector sizes; + for(auto dim : dims) + { + Value dim_size = rewriter.create(loc, index_type, sp_tensor, llvm::cast(dim)); + sizes.push_back(dim_size); + } + + WorkspaceType wsType = op.getType(); + Workspace workspace; + auto workspace_tensor_type = RankedTensorType::get(dim_attrs, sp_tensor_type.getElementType()); + workspace.workspace = rewriter.create(loc, workspace_tensor_type, sizes); + workspace.mark_value = rewriter.create(loc, rewriter.getI32Type(), rewriter.getI32IntegerAttr(1)); + auto workspace_mark_type = RankedTensorType::get(dim_attrs, rewriter.getI32Type()); + workspace.mark_array = rewriter.create(loc, workspace_mark_type, sizes); + workspace.num_crds = rewriter.create(loc, index_type, rewriter.getIndexAttr(0)); + auto crds_type = RankedTensorType::get({ShapedType::kDynamic,}, wsType.getIndicesType()); + workspace.crds = rewriter.create(loc, crds_type, sizes); + + auto workspace_type = llvm::cast(op->getResult(0).getType()); + /** TODO: Support higher dimensional workspaces! */ + assert(workspace_type.getDims().size() == 1 && "Workspace dimensions > 1 are currently unsupported."); + + SmallVector cast_args; + pack_workspace(workspace_type, workspace, cast_args); + rewriter.replaceOpWithNewOp(op, workspace_type, cast_args); + return success(); + } +}; + +class ConvertWorkspaceGetNNZ + : public OpConversionPattern { + using OpConversionPattern::OpConversionPattern; + ConvertWorkspaceGetNNZ(MLIRContext *context) + : OpConversionPattern(context) {} + + LogicalResult + matchAndRewrite(SpTensorGetNNZ op, OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const final { + auto opAdaptor = llvm::cast(adaptor); + if(!llvm::isa(opAdaptor.getTensor().getType())){ + return failure(); + } + Workspace workspace; + if(!unpack_workspace(opAdaptor.getTensor(), workspace)){ + return failure(); + } + + rewriter.replaceOp(op, {workspace.num_crds,}); + return success(); + } +}; + + +class ConvertReturnOp + : public OpConversionPattern { + using OpConversionPattern::OpConversionPattern; + ConvertReturnOp(MLIRContext *context) + : OpConversionPattern(context) {} + + LogicalResult + matchAndRewrite(func::ReturnOp op, OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const final { + auto converter = getTypeConverter(); + SmallVector newOperands; + for(auto operand: op.getOperands()) + { + if(!converter->isLegal(operand.getType())) + { + if(UnrealizedConversionCastOp cast_op = mlir::dyn_cast_if_present(operand.getDefiningOp()); cast_op && mlir::isa(operand.getType()) ) + { + newOperands.insert(newOperands.end(), cast_op->getOperands().begin(), cast_op->getOperands().end()); + } + else { + return failure(); + } + } + else + { + newOperands.push_back(operand); + } + + } + + func::ReturnOp new_return = rewriter.create(op->getLoc(), newOperands); + rewriter.replaceOp(op, new_return); + + return success(); + } +}; + +class ConvertFunCallOp + : public OpConversionPattern { + using OpConversionPattern::OpConversionPattern; + ConvertFunCallOp(MLIRContext *context) + : OpConversionPattern(context) {} + + LogicalResult + matchAndRewrite(func::CallOp op, OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const final { + auto opAdaptor = llvm::cast(adaptor); + auto converter = getTypeConverter(); + + SmallVector args; + for(auto operand: opAdaptor.getOperands()) + { + if(mlir::isa(operand.getType())) + { + SparseTensor sp_tensor; + if(!unpack_sparse_tensor(operand, sp_tensor)) + { + return failure(); + } + + args.push_back(sp_tensor.dim_sizes); + for(unsigned i = 0; i < sp_tensor.dims.size(); i++) { + Dimension dim = sp_tensor.dims[i]; + args.push_back(dim.insert_pos); + + switch(dim.format) { + case TensorFormatEnum::D: + { + args.push_back(dim.pos); + break; + } + case TensorFormatEnum::CU: + case TensorFormatEnum::CN: + { + args.push_back(dim.pos); //Pos tensor + args.push_back(dim.crd); //Crd tensor + break; + } + case TensorFormatEnum::S: + { + args.push_back(dim.crd); //Crd tensor + break; + } + default: { + } + } + } + args.push_back(sp_tensor.vals); + } + else + { + args.push_back(operand); + } + } + + SmallVector, 8> res_types; + + for(auto resultType: op.getResultTypes()) + { + SmallVector arg_res_types; + if(SparseTensorType spT = mlir::dyn_cast(resultType)) + { + if(failed(converter->convertType(spT,arg_res_types))) + { + return failure(); + } + } + else + { + arg_res_types.push_back(resultType); + } + + res_types.push_back(arg_res_types); + } + + + SmallVector all_res_types; + for(size_t i = 0; i < res_types.size(); i++) + { + all_res_types.append(res_types[i]); + } + auto newCallOp = rewriter.create(op->getLoc(), op.getCallee(), all_res_types, args); + + SmallVector result_values; + size_t start = 0; + for(size_t i = 0; i < op->getResultTypes().size(); i++) + { + result_values.push_back(typeConverter->materializeSourceConversion(rewriter, op->getLoc(), op->getResultTypes()[i], newCallOp->getResults().slice(start, res_types[i].size()))); + start += res_types[i].size(); + } + + rewriter.replaceOp(op, result_values); + + return success(); + } +}; + +class ConvertWorkspaceGetCrds + : public OpConversionPattern { + using OpConversionPattern::OpConversionPattern; + ConvertWorkspaceGetCrds(MLIRContext *context) + : OpConversionPattern(context) {} + + LogicalResult + matchAndRewrite(SpTensorGetCrd op, OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const final { + auto opAdaptor = llvm::cast(adaptor); + if(!llvm::isa(opAdaptor.getTensor().getType())){ + return failure(); + } + Workspace workspace; + if(!unpack_workspace(opAdaptor.getTensor(), workspace)){ + return failure(); + } + + Value new_op = rewriter.create(op->getLoc(), workspace.crds, opAdaptor.getIdx()); + rewriter.replaceOp(op, new_op); + + return success(); + } +}; + +class ConvertWorkspaceGetDimSize + : public OpConversionPattern { + using OpConversionPattern::OpConversionPattern; + ConvertWorkspaceGetDimSize(MLIRContext *context) + : OpConversionPattern(context) {} + LogicalResult + matchAndRewrite(SpTensorGetDimSize op, OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const override { + auto opAdaptor = llvm::cast(adaptor); + if(!llvm::isa(opAdaptor.getTensor().getType())){ + return failure(); + } + + Workspace workspace; + if(!unpack_workspace(opAdaptor.getTensor(), workspace)) { + return failure(); + } + Value dim = rewriter.create(op->getLoc(), rewriter.getIndexType(), rewriter.getIndexAttr(opAdaptor.getDim())); + rewriter.replaceOpWithNewOp(op, op->getResultTypes(), workspace.workspace, dim); + return success(); + } +}; + +class ConvertWorkspaceTensorInsertOp + : public OpConversionPattern { + using OpConversionPattern::OpConversionPattern; + ConvertWorkspaceTensorInsertOp(MLIRContext *context) + : OpConversionPattern(context) {} + LogicalResult + matchAndRewrite(TensorInsertOp op, OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const final { + auto opAdaptor = llvm::cast(adaptor); + if(!llvm::isa(opAdaptor.getTensor().getType())){ + return failure(); + } + WorkspaceType workspace_type = llvm::cast(opAdaptor.getTensor().getType()); + assert(workspace_type.getDims().size() == 1 && "Workspace dimensions > 1 are currently unsupported."); + + Workspace workspace; + if(!unpack_workspace(opAdaptor.getTensor(), workspace)) { + return failure(); + } + + auto loc = op.getLoc(); + auto context = op.getContext(); + ValueRange crds = opAdaptor.getCrds(); + Value crd = crds[opAdaptor.getCrds().size() - 1]; + Value mark_at_crd = rewriter.create( + loc, + workspace.mark_array, + crd + ); + Value not_seen = rewriter.create( + loc, + rewriter.getI1Type(), + arith::CmpIPredicateAttr::get(context, arith::CmpIPredicate::ne), + mark_at_crd, + workspace.mark_value + ); + + Operation* if_op = rewriter.create( + loc, + not_seen, + [workspace, crd] (OpBuilder& builder, Location loc) { + Type index_type = builder.getIndexType(); + Value new_mark = builder.create(loc, workspace.mark_value, workspace.mark_array, crd); + Value crd_cast = crd; + RankedTensorType crdT = mlir::cast(workspace.crds.getType()); + if(crdT.getElementType() != crd.getType()) + { + crd_cast = builder.create(loc, crdT.getElementType(), crd); + } + Value new_crds = builder.create(loc, crd_cast, workspace.crds, workspace.num_crds); + Value inc = builder.create(loc, index_type, builder.getIndexAttr(1)); + Value new_crd_size = builder.create(loc, index_type, workspace.num_crds, inc); + builder.create(loc, ArrayRef({new_mark, new_crd_size, new_crds})); + }, + [workspace] (OpBuilder& builder, Location loc) { + builder.create(loc, ArrayRef({workspace.mark_array, workspace.num_crds, workspace.crds})); + } + ); + workspace.mark_array = if_op->getResult(0); + workspace.num_crds = if_op->getResult(1); + workspace.crds = if_op->getResult(2); + workspace.workspace = rewriter.create( + loc, + workspace.workspace.getType(), + opAdaptor.getValue(), + workspace.workspace, + crd + ); + + SmallVector cast_args; + pack_workspace(workspace_type, workspace, cast_args); + rewriter.replaceOpWithNewOp(op, workspace_type, cast_args); + return success(); + } +}; + +class ConvertWorkspaceTensorExtractOp + : public OpConversionPattern { + using OpConversionPattern::OpConversionPattern; + ConvertWorkspaceTensorExtractOp(MLIRContext *context) + : OpConversionPattern(context) {} + LogicalResult + matchAndRewrite(TensorExtractOp op, OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const final { + auto opAdaptor = llvm::cast(adaptor); + if(!llvm::isa(opAdaptor.getTensor().getType())){ + return failure(); + } + + Workspace workspace; + if(!unpack_workspace(opAdaptor.getTensor(), workspace)) { + return failure(); + } + + auto loc = op.getLoc(); + auto context = op.getContext(); + Value crd = opAdaptor.getCrds()[0]; + crd = rewriter.createOrFold(loc, rewriter.getIndexType(), crd); + Value mark_at_crd = rewriter.create( + loc, + workspace.mark_array, + crd + ); + Value seen = rewriter.create( + loc, + rewriter.getI1Type(), + arith::CmpIPredicateAttr::get(context, arith::CmpIPredicate::eq), + mark_at_crd, + workspace.mark_value + ); + + + Operation* if_op = rewriter.create( + loc, + seen, + [&] (OpBuilder& builder, Location loc) { + Value extracted = builder.create(loc, op->getResultTypes(), workspace.workspace, crd); + builder.create(loc, ArrayRef({extracted})); + }, + [&] (OpBuilder& builder, Location loc) { + // TODO: Does the zero value depend on the semi-ring? + Type result_type = op->getResult(0).getType(); + FloatAttr zero_attr = mlir::cast(op.getZeroAttr()); + Value zero = builder.create(loc, result_type, zero_attr); + builder.create(loc, ArrayRef({zero})); + } + ); + + rewriter.replaceOp(op, if_op->getResults()); + return success(); + } +}; + +class ConvertWorkspaceClearOp + : public OpConversionPattern { + using OpConversionPattern::OpConversionPattern; + ConvertWorkspaceClearOp(MLIRContext *context) + : OpConversionPattern(context) {} + LogicalResult + matchAndRewrite(WorkspaceClearOp op, OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const final { + auto opAdaptor = llvm::cast(adaptor); + if(!llvm::isa(opAdaptor.getTensor().getType())){ + return failure(); + } + WorkspaceType workspace_type = llvm::cast(opAdaptor.getTensor().getType()); + Workspace workspace; + if(!unpack_workspace(opAdaptor.getTensor(), workspace)) { + return failure(); + } + auto loc = op.getLoc(); + Value inc = rewriter.create(loc, rewriter.getI32Type(), rewriter.getI32IntegerAttr(1)); + workspace.mark_value = rewriter.create(loc, rewriter.getI32Type(), workspace.mark_value, inc); + workspace.num_crds = rewriter.create(loc, rewriter.getIndexType(), rewriter.getIndexAttr(0)); + + SmallVector cast_args; + pack_workspace(workspace_type, workspace, cast_args); + rewriter.replaceOpWithNewOp(op, workspace_type, cast_args); + return success(); + } +}; + +class ConvertWorkspaceTensorFindPos + : public OpConversionPattern { + using OpConversionPattern::OpConversionPattern; + ConvertWorkspaceTensorFindPos(MLIRContext *context) + : OpConversionPattern(context) {} + LogicalResult + matchAndRewrite(TensorFindPos op, OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const override { + Workspace workspace; + if(!unpack_workspace(adaptor.getTensor(), workspace)) { + return failure(); + } + rewriter.replaceOp(op, {workspace.num_crds}); + + return success(); + } +}; + +class ConvertWorkspaceAccumulateOp + : public OpConversionPattern { + using OpConversionPattern::OpConversionPattern; + ConvertWorkspaceAccumulateOp(MLIRContext *context) + : OpConversionPattern(context) {} + LogicalResult + matchAndRewrite(WorkspaceAccumulateOp op, OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const override { + + Workspace workspace; + if(!unpack_workspace(adaptor.getTensor(), workspace)) { + return failure(); + } + + auto loc = op.getLoc(); + auto context = op.getContext(); + Value crd = adaptor.getCrds()[0]; + WorkspaceType workspace_type = llvm::cast(adaptor.getTensor().getType()); + + crd = rewriter.createOrFold(loc, rewriter.getIndexType(), crd); + Value mark_at_crd = rewriter.create( + loc, + workspace.mark_array, + crd + ); + Value seen = rewriter.create( + loc, + rewriter.getI1Type(), + arith::CmpIPredicateAttr::get(context, arith::CmpIPredicate::eq), + mark_at_crd, + workspace.mark_value + ); + + + Operation* if_op = rewriter.create( + loc, + seen, + [&] (OpBuilder& builder, Location loc) { + Value extracted = builder.create(loc, workspace_type.getElementType(), workspace.workspace, crd); + Value new_value = builder.create(loc, extracted.getType(), extracted, adaptor.getValue()); + Value new_workspace = rewriter.create( + loc, + workspace.workspace.getType(), + new_value, + workspace.workspace, + crd + ); + builder.create(loc, ArrayRef({workspace.mark_array, workspace.num_crds, workspace.crds, new_workspace})); + }, + [&] (OpBuilder& builder, Location loc) { + Type index_type = builder.getIndexType(); + Value new_mark = builder.create(loc, workspace.mark_value, workspace.mark_array, crd); + Value crd_cast = crd; + RankedTensorType crdT = mlir::cast(workspace.crds.getType()); + if(crdT.getElementType() != crd.getType()) + { + crd_cast = builder.create(loc, crdT.getElementType(), crd); + } + Value new_crds = builder.create(loc, crd_cast, workspace.crds, workspace.num_crds); + Value inc = builder.create(loc, index_type, builder.getIndexAttr(1)); + Value new_crd_size = builder.create(loc, index_type, workspace.num_crds, inc); + Value new_workspace = rewriter.create( + loc, + workspace.workspace.getType(), + adaptor.getValue(), + workspace.workspace, + crd + ); + builder.create(loc, ArrayRef({new_mark, new_crd_size, new_crds, new_workspace})); + } + ); + workspace.mark_array = if_op->getResult(0); + workspace.num_crds = if_op->getResult(1); + workspace.crds = if_op->getResult(2); + workspace.workspace = if_op->getResult(3); + + SmallVector cast_args; + pack_workspace(workspace_type, workspace, cast_args); + rewriter.replaceOpWithNewOp(op, workspace_type, cast_args); + return success(); + + rewriter.replaceOp(op, if_op->getResults()); + } +}; + +class ConvertWorkspaceRead + : public OpConversionPattern { + using OpConversionPattern::OpConversionPattern; + ConvertWorkspaceRead(MLIRContext *context) + : OpConversionPattern(context) {} + LogicalResult + matchAndRewrite(WorkspaceReadOp op, OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const override { + Workspace workspace; + if(!unpack_workspace(adaptor.getTensor(), workspace)) { + return failure(); + } + + rewriter.replaceOpWithNewOp(op, workspace.workspace, adaptor.getCrd()); + return success(); + } +}; + +class ConvertWorkspaceSrtCrd + : public OpConversionPattern { + + using OpConversionPattern::OpConversionPattern; + ConvertWorkspaceSrtCrd(MLIRContext *context) + : OpConversionPattern(context) {} + + LogicalResult + matchAndRewrite(SortCrdOp op, OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const override { + Workspace workspace; + if(!unpack_workspace(adaptor.getTensor(), workspace)) { + return failure(); + } + + auto loc = op.getLoc(); + Value crd = workspace.crds; + Type crdType = crd.getType(); + Value zero = rewriter.create(loc, rewriter.getIndexType(), rewriter.getIndexAttr(0)); + workspace.crds = rewriter.create(loc, crdType, crd, zero, workspace.num_crds); + + SmallVector cast_args; + WorkspaceType workspace_type = llvm::cast(adaptor.getTensor().getType()); + pack_workspace(workspace_type, workspace, cast_args); + rewriter.replaceOpWithNewOp(op, workspace_type, cast_args); + return success(); + } +}; + +class PrintOpLowering : public OpConversionPattern { + using OpConversionPattern::OpConversionPattern; + PrintOpLowering(MLIRContext *context) : OpConversionPattern(context) {} + + LogicalResult + matchAndRewrite(PrintOp op, OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const override + { + Location loc = op->getLoc(); + auto inputType = adaptor.getInput().getType(); + Type index_type = rewriter.getIndexType(); + SmallVector empty_size(1, 1); + auto empty_type = RankedTensorType::get(empty_size, index_type); + Value empty_tensor = rewriter.create(loc, empty_type, ValueRange(), (Value)nullptr); + Value neg = rewriter.create(loc, index_type, rewriter.getIndexAttr(-1)); + Value zero = rewriter.create(loc, index_type, rewriter.getIndexAttr(0)); + empty_tensor = rewriter.create(loc, empty_type, neg, empty_tensor, zero); + + if (isa(inputType)) + { + SparseTensor sp_tensor; + if(!unpack_sparse_tensor(adaptor.getInput(), sp_tensor)) { + return failure(); + } + for (Dimension& dim : sp_tensor.dims) + { + switch(dim.format){ + case TensorFormatEnum::D: { + rewriter.create(loc, dim.pos); + rewriter.create(loc, empty_tensor); + break; + } + case TensorFormatEnum::CU: + case TensorFormatEnum::CN: { + rewriter.create(loc, dim.pos); + rewriter.create(loc, dim.crd); + break; + } + case TensorFormatEnum::S: { + rewriter.create(loc, empty_tensor); + rewriter.create(loc, dim.crd); + break; + } + default: { + assert(false && "Could not print unknown format to sparse tensor."); + } + } + } + rewriter.create(loc, sp_tensor.vals); + rewriter.eraseOp(op); + return success(); + } + return failure(); + } +}; + +class GetTimeLowering : public OpConversionPattern { + using OpConversionPattern::OpConversionPattern; + GetTimeLowering(MLIRContext *context) : OpConversionPattern(context) {} + + LogicalResult + matchAndRewrite(GetTimeOp op, OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const override + { + auto ctx = rewriter.getContext(); + auto module = op->getParentOfType(); + auto f64Type = rewriter.getF64Type(); + std::string getTimeStr = "getTime"; + + if (!hasFuncDeclaration(module, getTimeStr)) + { + auto getTimeFunc = FunctionType::get(ctx, {}, {FloatType::getF64(ctx)}); + /// func @getTime() -> f64 + func::FuncOp func1 = func::FuncOp::create(op->getLoc(), getTimeStr, + getTimeFunc, ArrayRef{}); + func1.setPrivate(); + module.push_back(func1); + } + + rewriter.replaceOpWithNewOp(op, getTimeStr, SmallVector{f64Type}); + + return success(); + } +}; + +class PrintElapsedTimeLowering : public OpConversionPattern { + using OpConversionPattern::OpConversionPattern; + PrintElapsedTimeLowering(MLIRContext *context) : OpConversionPattern(context) {} + + LogicalResult + matchAndRewrite(PrintElapsedTimeOp op, OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const override + { + auto ctx = rewriter.getContext(); + auto module = op->getParentOfType(); + + auto start = adaptor.getStart(); + auto end = adaptor.getEnd(); + std::string printElapsedTimeStr = "printElapsedTime"; + auto f64Type = rewriter.getF64Type(); + + if (!hasFuncDeclaration(module, printElapsedTimeStr)) + { + auto printElapsedTimeFunc = FunctionType::get(ctx, {f64Type, f64Type}, {}); + /// func @printElapsedTime(f64, f64) -> () + func::FuncOp func1 = func::FuncOp::create(op->getLoc(), printElapsedTimeStr, + printElapsedTimeFunc, ArrayRef{}); + func1.setPrivate(); + module.push_back(func1); + } + + rewriter.replaceOpWithNewOp(op, printElapsedTimeStr, SmallVector{}, ValueRange{start, end}); + + return success(); + } +}; + +} + +void mlir::comet::populateSparseTensorConversionPatterns(MLIRContext *context, RewritePatternSet &patterns, TypeConverter &typeConverter) { + typeConverter.addConversion( + [](tensorAlgebra::SparseTensorType type, SmallVectorImpl &types) { + ArrayRef dim_sizes = type.getDims(); + ArrayRef format = type.getFormat(); + + auto context = type.getContext(); + Type index_type = IndexType::get(context); + [[maybe_unused]] bool is_known_size = true; + [[maybe_unused]] int known_size = 1; + types.push_back(RankedTensorType::get({static_cast(dim_sizes.size())}, IndexType::get(context))); //Dimension sizes + for(unsigned i = 0; i < dim_sizes.size(); i++) { + types.push_back(index_type); //Insert pos + switch(format[2 * i]) + { + case TensorFormatEnum::D: + { + if(dim_sizes[i] != ShapedType::kDynamic) { + known_size *= dim_sizes[i]; + } else { + is_known_size = false; + } + auto pos_type = mlir::RankedTensorType::get({ShapedType::kDynamic,}, type.getIndicesType()); + types.push_back(pos_type); //Pos tensor + break; + } + case TensorFormatEnum::CU: + case TensorFormatEnum::CN: + { + Type pos_type = mlir::RankedTensorType::get({ShapedType::kDynamic,}, + type.getIndicesType()); + Type crd_type = mlir::RankedTensorType::get({ShapedType::kDynamic,}, + type.getIndicesType()); + is_known_size = false; + + types.push_back(pos_type); //Pos tensor + types.push_back(crd_type); //Crd tensor + break; + } + case TensorFormatEnum::S: + { + Type crd_type = mlir::RankedTensorType::get({ShapedType::kDynamic,}, type.getIndicesType()); + types.push_back(crd_type); //Crd tensor + break; + } + default: { + assert(false && "Could not unpack unknown format to sparse tensor."); + } + } + } + Type value_type = mlir::RankedTensorType::get({ShapedType::kDynamic,}, type.getElementType()); + types.push_back(value_type); //Value tensor + return success(); + }); + + typeConverter.addConversion( + [](WorkspaceType type, SmallVectorImpl &types) { + Type element_type = type.getElementType(); + ArrayRef dim_sizes = type.getDims(); + auto context = type.getContext(); + types.push_back(RankedTensorType::get(dim_sizes, element_type)); // Workspace + types.push_back(IntegerType::get(context, 32)); // Mark Value + types.push_back(RankedTensorType::get(dim_sizes, IntegerType::get(context, 32))); // Mark array + types.push_back(IndexType::get(context)); // Crd Size + types.push_back(RankedTensorType::get({ShapedType::kDynamic,}, type.getIndicesType()));// Crd tensors + return success(); + }); + + typeConverter.addArgumentMaterialization( + [](OpBuilder &builder, SparseTensorType resultType, ValueRange inputs, + Location loc) -> std::optional { + auto op = builder.create(loc, resultType, inputs); + return op.getResult(0); + }); + + typeConverter.addSourceMaterialization( + [](OpBuilder &builder, SparseTensorType resultType, ValueRange inputs, + Location loc) -> std::optional { + auto op = builder.create(loc, resultType, inputs); + return op.getResult(0); + }); + + typeConverter.addArgumentMaterialization( + [](OpBuilder &builder, WorkspaceType resultType, ValueRange inputs, + Location loc) -> std::optional { + auto op = builder.create(loc, resultType, inputs); + return op.getResult(0); + }); + + typeConverter.addSourceMaterialization( + [](OpBuilder &builder, WorkspaceType resultType, ValueRange inputs, + Location loc) -> std::optional { + auto op = builder.create(loc, resultType, inputs); + return op.getResult(0); + }); + + patterns.add(typeConverter, context); + patterns.add(typeConverter, context); + patterns.add(typeConverter, context); +} + +struct SparseTensorConversionPass : comet::impl::SparseTensorConversionPassBase { + using SparseTensorConversionPassBase::SparseTensorConversionPassBase; + + SparseTensorConversionPass() = default; + SparseTensorConversionPass(const SparseTensorConversionPass &pass) = default; + + void runOnOperation() override { + auto *ctx = &getContext(); + RewritePatternSet patterns(ctx); + TypeConverter typeConverter; + ConversionTarget target(*ctx); + + // Everything in the TADialect must go + target.addIllegalDialect(); + + // The following operations and dialects may be introduced by the + // rewriting rules, and are therefore marked as legal. + target.addLegalOp(); + target.addLegalDialect< + arith::ArithDialect, bufferization::BufferizationDialect, + tensor::TensorDialect, memref::MemRefDialect, scf::SCFDialect, + func::FuncDialect, index::IndexDialect, BuiltinDialect + >(); + + target.addDynamicallyLegalOp([&](func::FuncOp op){ + return typeConverter.isSignatureLegal(op.getFunctionType()); + }); + + target.addLegalOp(); + target.addDynamicallyLegalOp([&](tensorAlgebra::PrintOp op) { + return typeConverter.isLegal(op->getOperandTypes()); + }); + target.addLegalOp(); + + typeConverter.addConversion([](Type type) { return type; }); + + + // Populate with rules and apply rewriting rules. + populateFunctionOpInterfaceTypeConversionPattern(patterns, + typeConverter); + populateCallOpTypeConversionPattern(patterns, typeConverter); + scf::populateSCFStructuralTypeConversionsAndLegality(typeConverter, patterns, + target); + mlir::indexTree::populateIndexTreeTypeConversionPatterns(ctx, patterns, typeConverter, target); + mlir::comet::populateSparseTensorConversionPatterns(ctx, patterns, typeConverter); + if (failed(applyPartialConversion(getOperation(), target, + std::move(patterns)))) + { + return signalPassFailure(); + } + + RewritePatternSet funcPatterns(ctx); + funcPatterns.add(typeConverter, ctx); + TypeConverter funcTypeConverter; + + target.addDynamicallyLegalOp([&](func::CallOp op){ + return typeConverter.isLegal(op); + }); + + target.addDynamicallyLegalOp([&](func::ReturnOp op){ + return typeConverter.isLegal(op); + }); + + target.addLegalOp(); + + populateFunctionOpInterfaceTypeConversionPattern(funcPatterns, + typeConverter); + if (failed(applyPartialConversion(getOperation(), target, + std::move(funcPatterns)))) + { + return signalPassFailure(); + } + } +}; + +std::unique_ptr mlir::comet::createSparseTensorConversionPass() +{ + return std::make_unique(); +} diff --git a/lib/Conversion/TensorAlgebraToSCF/TensorAlgebraToSCF.cpp b/lib/Conversion/TensorAlgebraToSCF/TensorAlgebraToSCF.cpp index 09dd38fb..a75eea87 100644 --- a/lib/Conversion/TensorAlgebraToSCF/TensorAlgebraToSCF.cpp +++ b/lib/Conversion/TensorAlgebraToSCF/TensorAlgebraToSCF.cpp @@ -28,12 +28,22 @@ #include "comet/Dialect/Utils/Utils.h" #include "comet/Conversion/TensorAlgebraToSCF/TensorAlgebraToSCF.h" #include "comet/Dialect/TensorAlgebra/Passes.h" +#include "comet/Dialect/IndexTree/IR/IndexTreeDialect.h" #include "mlir/Dialect/Arith/IR/Arith.h" #include "mlir/Dialect/Bufferization/IR/Bufferization.h" #include "mlir/Dialect/SCF/IR/SCF.h" #include "mlir/Dialect/MemRef/IR/MemRef.h" #include "mlir/Dialect/Func/IR/FuncOps.h" +#include "mlir/Dialect/Tensor/IR/Tensor.h" +#include "mlir/IR/BuiltinAttributes.h" +#include "mlir/IR/BuiltinTypeInterfaces.h" +#include "mlir/IR/BuiltinTypes.h" +#include "mlir/IR/ValueRange.h" +#include "mlir/Support/LLVM.h" +#include "mlir/Support/LogicalResult.h" +#include "llvm/ADT/SmallVector.h" +#include using namespace mlir; using namespace mlir::arith; @@ -97,10 +107,9 @@ namespace Value alloc; if (user_setOp) { - if (isa(setnewop.getOperand(1).getDefiningOp())) + if (auto toTensor = dyn_cast(setnewop.getOperand(1).getDefiningOp())) { - Operation *tensorload = cast(setnewop.getOperand(1).getDefiningOp()); - auto alloc_op = cast(tensorload->getOperand(0).getDefiningOp()); + auto alloc_op = cast(toTensor->getOperand(0).getDefiningOp()); comet_vdump(alloc_op); alloc = alloc_op; } @@ -117,7 +126,7 @@ namespace /// Create these constants up-front to avoid large amounts of redundant /// operations. auto valueShape = memRefType.getShape(); - auto constTensor = op.getValue().getType().cast(); + auto constTensor = mlir::cast(op.getValue().getType()); if(constTensor.getRank() == 1 && constTensor.getDimSize(0) == 1) { auto float_attr = *constantValue.getValues().begin(); @@ -173,7 +182,7 @@ namespace } /// Replace this operation with the generated alloc. - op->replaceAllUsesWith(rewriter.create(op->getLoc(),alloc)); + op->replaceAllUsesWith(rewriter.create(op->getLoc(), alloc, rewriter.getUnitAttr(), rewriter.getUnitAttr())); rewriter.eraseOp(op); comet_debug() << "ConstantOpLowering ends\n"; return success(); @@ -205,6 +214,7 @@ namespace comet_vdump(op); auto *ctx = op->getContext(); auto inputType = op->getOperand(0).getType(); + auto inputTensor = op->getOperand(0); /// Get tensor contraction expression through analyzing the index map ArrayAttr indexMaps = op.getIndexingMaps(); @@ -215,12 +225,13 @@ namespace tensorAlgebra::TensorSetOp setOp; Value lhs; - if (inputType.isa()) + if (isa(inputType)) { /// for dense comet_debug() << "Dense transpose\n"; - auto inputTensorLoadOp = cast(op->getOperand(0).getDefiningOp()); - auto inputMemref = inputTensorLoadOp.getMemref(); + + // auto inputTensorLoadOp = cast(op->getOperand(0).getDefiningOp()); + // auto inputMemref = inputTensorLoadOp.getMemref(); for (auto u : op.getOperation()->getResult(0).getUsers()) { @@ -233,16 +244,13 @@ namespace } comet_vdump(lhs); - auto outputMemref = lhs.getDefiningOp()->getOperand(0); - rewriter.create(loc, inputMemref, outputMemref, llvm::ArrayRef(allPerms[1])); - Value res_value = rewriter.create(loc, outputMemref); - - op.replaceAllUsesWith(res_value); - rewriter.eraseOp(op); - + // auto outputMemref = lhs.getDefiningOp()->getOperand(0); + auto la_transpose = rewriter.create(loc, inputTensor, lhs, llvm::ArrayRef(allPerms[1])); + // Value res_value = rewriter.create(loc, outputMemref, rewriter.getUnitAttr(), rewriter.getUnitAttr()); + rewriter.replaceOp(op, la_transpose.getResults()); return success(); } - else + else if(auto spType = mlir::dyn_cast(inputType)) { /// for sparse tensors int64_t pnum[2]; /// print allPerms @@ -253,18 +261,17 @@ namespace i++; } - ArrayAttr opFormatsArrayAttr = op.getFormats(); - std::string formats_strIn(opFormatsArrayAttr[0].cast().getValue()); - std::string formats_strOut(opFormatsArrayAttr[1].cast().getValue()); + std::string formats_strIn(getTensorFormatString(op.getOperandTypes()[0])); + std::string formats_strOut(getTensorFormatString(op->getResultTypes()[0])); IntegerType i32Type = IntegerType::get(ctx, 32); IndexType indexType = IndexType::get(ctx); - FloatType f64Type = FloatType::getF64(ctx); Value input_perm_num = rewriter.create(loc, i32Type, rewriter.getI32IntegerAttr(pnum[0])); Value output_perm_num = rewriter.create(loc, i32Type, rewriter.getI32IntegerAttr(pnum[1])); - Type unrankedMemrefType_f64 = UnrankedMemRefType::get(f64Type, 0); + UnrankedMemRefType unrankedMemrefType_float = UnrankedMemRefType::get(spType.getElementType(), 0); Type unrankedMemrefType_index = UnrankedMemRefType::get(indexType, 0); + Type unrankedMemrefType_indices_type = UnrankedMemRefType::get(spType.getIndicesType(), 0); mlir::func::FuncOp transpose_func; /// runtime call @@ -284,50 +291,52 @@ namespace for (unsigned int n = 0; n < tensors_num; n++) { - auto tensor_rank_attr = tensors[n].getDefiningOp()->getAttr("tensor_rank"); - auto tensor_rank_int_attr = cast(tensor_rank_attr); - unsigned int tensor_rank = tensor_rank_int_attr.getValue().getLimitedValue(); + auto tensor_rank = cast(tensors[n].getType()).getRank(); comet_debug() << "ATTR_Val: " << tensor_rank << "\n"; comet_debug() << " tensor_rank: " << tensor_rank << "\n"; comet_debug() << " tensor[n]: " << "\n"; comet_pdump(tensors[n].getDefiningOp()); - - for (unsigned int i = 0; i < 4 * tensor_rank + 1; i++) + auto tensor = tensors[n]; + auto tensor_type = cast(tensors[n].getType()); + for(int i = 0; i < tensor_rank; i++) { - auto tensorload_op = tensors[n].getDefiningOp()->getOperand(i); - comet_debug() << " tensorload_op " - << "\n"; - comet_vdump(tensorload_op); - - auto alloc_op = tensorload_op.getDefiningOp()->getOperand(0); - comet_debug() << " alloc_op " - << "\n"; - comet_vdump(alloc_op); - - if (i < 4 * tensor_rank) - { - /// indexes crd's - mlir::Value v = rewriter.create(loc, unrankedMemrefType_index, alloc_op); - alloc_sizes_cast_vecs[n].push_back(v); - } - else - { - /// NNZ vals - mlir::Value v = rewriter.create(loc, unrankedMemrefType_f64, alloc_op); - alloc_sizes_cast_vecs[n].push_back(v); - } + Value pos = rewriter.create(loc, tensor, rewriter.getI32IntegerAttr(i)); + ShapedType shape = cast(pos.getType()); + Value pos_memref = rewriter.create(loc, MemRefType::get(shape.getShape(), shape.getElementType()), pos); + Value pos_v = rewriter.create(loc, unrankedMemrefType_indices_type, pos_memref); + alloc_sizes_cast_vecs[n].push_back(pos_v); + + mlir::Value crd = rewriter.create(loc, tensor, rewriter.getI32IntegerAttr(i)); + mlir::Value crd_memref = rewriter.create(loc, MemRefType::get(shape.getShape(), shape.getElementType()), crd); + mlir::Value crd_v = rewriter.create(loc, unrankedMemrefType_indices_type, crd_memref); + alloc_sizes_cast_vecs[n].push_back(crd_v); + + Value block_pos = rewriter.create(loc, tensor, rewriter.getI32IntegerAttr(i)); + Value block_pos_memref = rewriter.create(loc, MemRefType::get(shape.getShape(), shape.getElementType()), block_pos); + Value block_pos_v = rewriter.create(loc, unrankedMemrefType_indices_type, block_pos_memref); + alloc_sizes_cast_vecs[n].push_back(block_pos_v); + + mlir::Value block_crd = rewriter.create(loc, tensor, rewriter.getI32IntegerAttr(i)); + mlir::Value block_crd_memref = rewriter.create(loc, MemRefType::get(shape.getShape(), shape.getElementType()), block_crd); + mlir::Value block_crd_v = rewriter.create(loc, unrankedMemrefType_indices_type, block_crd_memref); + alloc_sizes_cast_vecs[n].push_back(block_crd_v); } - - auto memrefload_op = tensors[n].getDefiningOp()->getOperand(tensors[n].getDefiningOp()->getNumOperands() - 1); - allocs_for_sparse_tensors[n].push_back(memrefload_op); + Value vals = rewriter.create(loc, tensor); + Value vals_memref = rewriter.create(loc, MemRefType::get({ShapedType::kDynamic}, tensor_type.getElementType()), vals); + Value vals_v = rewriter.create(loc, unrankedMemrefType_float, vals_memref); + alloc_sizes_cast_vecs[n].push_back(vals_v); + + auto dims_tensor = mlir::cast(tensors[n].getDefiningOp()).getDims(); + auto dims_memref = rewriter.create(loc, MemRefType::get(dims_tensor.getType().getShape(), dims_tensor.getType().getElementType()), dims_tensor); + allocs_for_sparse_tensors[n].push_back(dims_memref); comet_debug() << " memrefload_op " << "\n"; comet_vdump(memrefload_op); } - Value last_dim_size_alloc = allocs_for_sparse_tensors[0][0].getDefiningOp()->getOperand(0); + Value last_dim_size_alloc = allocs_for_sparse_tensors[0][0]; comet_debug() << "Alloc for last dim size:\n"; comet_vdump(last_dim_size_alloc); @@ -339,6 +348,7 @@ namespace auto tensor_rank_int_attr = cast(tensor_rank_attr); unsigned int rank_size = tensor_rank_int_attr.getValue().getLimitedValue(); comet_debug() << "ATTR_Val: Rank_size: " << rank_size << "\n"; + assert(rank_size <= 3 && rank_size >=2 && "Rank size not supported"); /// dim format of input tensor std::vector @@ -347,120 +357,57 @@ namespace /// dim format of output tensor std::vector dim_formatOut = mlir::tensorAlgebra::getFormatsValueInt(formats_strOut, rank_size, rewriter, loc, i32Type); - - if (rank_size == 2) - { /// 2D - auto transpose2DF64Func = FunctionType::get(ctx, - {i32Type, i32Type, i32Type, i32Type, - unrankedMemrefType_index, unrankedMemrefType_index, - unrankedMemrefType_index, unrankedMemrefType_index, - unrankedMemrefType_index, unrankedMemrefType_index, - unrankedMemrefType_index, unrankedMemrefType_index, - unrankedMemrefType_f64, - i32Type, i32Type, i32Type, i32Type, - unrankedMemrefType_index, unrankedMemrefType_index, - unrankedMemrefType_index, unrankedMemrefType_index, - unrankedMemrefType_index, unrankedMemrefType_index, - unrankedMemrefType_index, unrankedMemrefType_index, - unrankedMemrefType_f64, - unrankedMemrefType_index}, - {}); - - std::string func_name = "transpose_2D_f64"; - if (!hasFuncDeclaration(module, func_name)) - { - transpose_func = mlir::func::FuncOp::create(loc, func_name, transpose2DF64Func, ArrayRef{}); - transpose_func.setPrivate(); - module.push_back(transpose_func); - } - comet_debug() << "alloc_sizes_vec: " << alloc_sizes_cast_vecs.size() << "\n"; - comet_debug() << "alloc_sizes_vec[0]: " << alloc_sizes_cast_vecs[0].size() << "\n"; - comet_debug() << "alloc_sizes_vec[1]: " << alloc_sizes_cast_vecs[1].size() << "\n"; - comet_debug() << "dim_formatIn: " << dim_formatIn.size() << "\n"; - comet_debug() << "dim_formatOut: " << dim_formatOut.size() << "\n"; - - rewriter.create(loc, func_name, SmallVector{}, - ValueRange{dim_formatIn[0], dim_formatIn[1], dim_formatIn[2], dim_formatIn[3], - alloc_sizes_cast_vecs[0][0], alloc_sizes_cast_vecs[0][1], - alloc_sizes_cast_vecs[0][2], alloc_sizes_cast_vecs[0][3], - alloc_sizes_cast_vecs[0][4], alloc_sizes_cast_vecs[0][5], - alloc_sizes_cast_vecs[0][6], alloc_sizes_cast_vecs[0][7], - alloc_sizes_cast_vecs[0][8], - dim_formatOut[0], dim_formatOut[1], dim_formatOut[2], dim_formatOut[3], - alloc_sizes_cast_vecs[1][0], alloc_sizes_cast_vecs[1][1], - alloc_sizes_cast_vecs[1][2], alloc_sizes_cast_vecs[1][3], - alloc_sizes_cast_vecs[1][4], alloc_sizes_cast_vecs[1][5], - alloc_sizes_cast_vecs[1][6], alloc_sizes_cast_vecs[1][7], - alloc_sizes_cast_vecs[1][8], - sparse_tensor_desc}); + std::string func_name = "transpose_" + std::to_string(rank_size)+ "D" + "_" + "f" + std::to_string(unrankedMemrefType_float.getElementType().getIntOrFloatBitWidth()) + "_i"+std::to_string(spType.getIndicesType().getWidth()); + llvm::SmallVector funcArgTypes; + if( rank_size == 3) + { + funcArgTypes.push_back(i32Type); + funcArgTypes.push_back(i32Type); } - else if (rank_size == 3) - { /// 3D - auto transpose3DF64Func = FunctionType::get(ctx, - {i32Type, i32Type, - i32Type, i32Type, - i32Type, i32Type, - i32Type, i32Type, - unrankedMemrefType_index, unrankedMemrefType_index, - unrankedMemrefType_index, unrankedMemrefType_index, - unrankedMemrefType_index, unrankedMemrefType_index, - unrankedMemrefType_index, unrankedMemrefType_index, - unrankedMemrefType_index, unrankedMemrefType_index, - unrankedMemrefType_index, unrankedMemrefType_index, - unrankedMemrefType_f64, - i32Type, i32Type, - i32Type, i32Type, - i32Type, i32Type, - unrankedMemrefType_index, unrankedMemrefType_index, - unrankedMemrefType_index, unrankedMemrefType_index, - unrankedMemrefType_index, unrankedMemrefType_index, - unrankedMemrefType_index, unrankedMemrefType_index, - unrankedMemrefType_index, unrankedMemrefType_index, - unrankedMemrefType_index, unrankedMemrefType_index, - unrankedMemrefType_f64, - unrankedMemrefType_index}, - {}); - - std::string func_name = "transpose_3D_f64"; - if (!hasFuncDeclaration(module, func_name)) + for(int k = 0; k < 2; k++) + { + for(unsigned i = 0; i < rank_size * 2; i++) + { + funcArgTypes.push_back(i32Type); + } + for(unsigned i = 0; i < rank_size * 4; i++) { - transpose_func = mlir::func::FuncOp::create(loc, func_name, transpose3DF64Func, ArrayRef{}); - transpose_func.setPrivate(); - module.push_back(transpose_func); + funcArgTypes.push_back(unrankedMemrefType_indices_type); } - - rewriter.create(loc, func_name, SmallVector{}, - ValueRange{input_perm_num, output_perm_num, - dim_formatIn[0], dim_formatIn[1], dim_formatIn[2], - dim_formatIn[3], dim_formatIn[4], dim_formatIn[5], - alloc_sizes_cast_vecs[0][0], alloc_sizes_cast_vecs[0][1], - alloc_sizes_cast_vecs[0][2], alloc_sizes_cast_vecs[0][3], - alloc_sizes_cast_vecs[0][4], alloc_sizes_cast_vecs[0][5], - alloc_sizes_cast_vecs[0][6], alloc_sizes_cast_vecs[0][7], - alloc_sizes_cast_vecs[0][8], alloc_sizes_cast_vecs[0][9], - alloc_sizes_cast_vecs[0][10], alloc_sizes_cast_vecs[0][11], - alloc_sizes_cast_vecs[0][12], - dim_formatOut[0], dim_formatOut[1], dim_formatOut[2], - dim_formatOut[3], dim_formatOut[4], dim_formatOut[5], - alloc_sizes_cast_vecs[1][0], alloc_sizes_cast_vecs[1][1], - alloc_sizes_cast_vecs[1][2], alloc_sizes_cast_vecs[1][3], - alloc_sizes_cast_vecs[1][4], alloc_sizes_cast_vecs[1][5], - alloc_sizes_cast_vecs[1][6], alloc_sizes_cast_vecs[1][7], - alloc_sizes_cast_vecs[1][8], alloc_sizes_cast_vecs[1][9], - alloc_sizes_cast_vecs[1][10], alloc_sizes_cast_vecs[1][11], - alloc_sizes_cast_vecs[1][12], - sparse_tensor_desc}); + funcArgTypes.push_back(unrankedMemrefType_float); } - else + funcArgTypes.push_back(unrankedMemrefType_index); + auto transposeFunc = FunctionType::get(ctx, TypeRange(funcArgTypes), {}); + if (!hasFuncDeclaration(module, func_name)) + { + transpose_func = mlir::func::FuncOp::create(loc, func_name, transposeFunc, ArrayRef{}); + transpose_func.setPrivate(); + module.push_back(transpose_func); + } + + std::vector allInputs; + if(rank_size == 3) { - llvm::errs() << "ERROR: Tensors greater than 3 are not currently supported.\n"; + allInputs.push_back(input_perm_num); + allInputs.push_back(output_perm_num); } + allInputs.insert(allInputs.end(), dim_formatIn.begin(), dim_formatIn.end()); + allInputs.insert(allInputs.end(), alloc_sizes_cast_vecs[0].begin(), alloc_sizes_cast_vecs[0].end()); + allInputs.insert(allInputs.end(), dim_formatOut.begin(), dim_formatOut.end()); + allInputs.insert(allInputs.end(), alloc_sizes_cast_vecs[1].begin(), alloc_sizes_cast_vecs[1].end()); + allInputs.push_back(sparse_tensor_desc); + + rewriter.create(loc, func_name, SmallVector{}, ValueRange(allInputs) ); rewriter.eraseOp(setOp); rewriter.eraseOp(op); return success(); } /// end else sparse tensor + else + { + return failure(); + } } /// Tensor TransposeLowering }; @@ -478,48 +425,49 @@ namespace comet_debug() << "Lowering Reduce operation to SCF\n"; Location loc = op.getLoc(); - auto f64Type = rewriter.getF64Type(); + // auto f64Type = rewriter.getF64Type(); auto inputType = op->getOperand(0).getType(); /// Allocate memory for the result and initialized it auto cst_zero = rewriter.create(loc, 0); /// need to access res alloc - MemRefType memTy_alloc_res = MemRefType::get({1}, f64Type); + ShapedType shapeT = mlir::cast(inputType); + MemRefType memTy_alloc_res = MemRefType::get({1}, shapeT.getElementType()); + Value res = rewriter.create(loc, memTy_alloc_res); - Value const_f64_0 = rewriter.create(loc, f64Type, rewriter.getF64FloatAttr(0)); + FloatAttr zero; + if(shapeT.getElementType().isF32()) + { + zero = rewriter.getF32FloatAttr(0); + } + else if(shapeT.getElementType().isF64()) + { + zero = rewriter.getF64FloatAttr(0); + } + else + { + assert(false && "Unexpected type"); + } + Value const_float_0 = rewriter.create(loc, shapeT.getElementType(), zero); std::vector alloc_zero_loc = {cst_zero}; - rewriter.create(loc, const_f64_0, + rewriter.create(loc, const_float_0, res, alloc_zero_loc); comet_vdump(res); - if (inputType.isa()) + if (auto tensorT = dyn_cast(inputType)) { /// tensor is dense comet_debug() << "Input Tensor is dense\n"; std::vector indices; - auto alloc_op = op->getOperand(0).getDefiningOp()->getOperand(0); comet_vdump(alloc_op); + auto lowerBound = rewriter.create(loc, 0); + auto step = rewriter.create(loc, 1); - for (unsigned rank = 0; rank < inputType.cast().getRank(); rank++) - { - auto dimSize = inputType.cast().getDimSize(rank); - Value upperBound; - if (dimSize == ShapedType::kDynamic) - { - comet_debug() << " This dimension is a dynamic size\n"; - comet_vdump(alloc_op); - auto memRefType = alloc_op.getType().dyn_cast(); - unsigned dynamicDimPos = memRefType.getDynamicDimIndex(rank); - comet_debug() << " dynamicDimPos: " << dynamicDimPos << "\n"; - upperBound = alloc_op.getDefiningOp()->getOperand(dynamicDimPos); - } - else - { - upperBound = rewriter.create(loc, dimSize); - } - auto lowerBound = rewriter.create(loc, 0); - auto step = rewriter.create(loc, 1); + for (unsigned rank = 0; rank < tensorT.getRank(); rank++) + { + Value upperBound = rewriter.create(loc, op->getOperand(0), rank); + /// create for loops auto loop = rewriter.create(loc, lowerBound, upperBound, step); indices.push_back(loop.getInductionVar()); @@ -527,58 +475,24 @@ namespace } /// Build loop body - auto load_rhs = rewriter.create(loc, alloc_op, indices); + auto load_rhs = rewriter.create(loc, op->getOperand(0), indices); auto res_load = rewriter.create(loc, res, alloc_zero_loc); auto reduced = rewriter.create(loc, load_rhs, res_load); rewriter.create(loc, reduced, res, alloc_zero_loc); } - else + else if (auto spTensorT = dyn_cast(inputType)) { /// sparse tensor type - assert(inputType.isa()); comet_debug() << "Input Tensor is sparse\n"; comet_pdump(op); - assert(isa(op->getOperand(0).getDefiningOp())); - tensorAlgebra::SparseTensorConstructOp sp_op = cast(op->getOperand(0).getDefiningOp()); - int tensorRanks = sp_op.getTensorRank(); + int tensorRanks = spTensorT.getRank(); comet_debug() << " tensorRank: " << tensorRanks << " \n"; comet_debug() << "Tensor to reduce:\n"; comet_pdump(op->getOperand(0).getDefiningOp()); + Value sp_tensor_values = rewriter.create(loc, RankedTensorType::get({ShapedType::kDynamic,}, spTensorT.getElementType()), op->getOperand(0)); + Value upperBound = rewriter.create(loc, sp_tensor_values, 0); - /// create the lowerBound, upperbound and step for loop - int indexValueSize = sp_op.getIndexValueSize(); - comet_debug() << "indexValueSize in SparseTensorConstructOp:" << indexValueSize << "\n"; - - auto loadOpForNNZ = op->getOperand(0).getDefiningOp()->getOperand(indexValueSize); - comet_debug() << "Corresponding AllocOp from SparseTensorConstructOp:\n"; - comet_vdump(loadOpForNNZ); - auto memAllocForNNZ = loadOpForNNZ.getDefiningOp()->getOperand(0); - comet_debug() << "Corresponding MemAllocOp for NNZ:\n"; - comet_vdump(memAllocForNNZ); - - MemRefType resultMemTy = memAllocForNNZ.getDefiningOp()->getResult(0).getType().cast(); - auto memRefRank = resultMemTy.getRank(); - comet_debug() << "memRefRank for alloc: " << memRefRank << "\n"; - assert(memRefRank == 1); /// Memref rank should be 1 - - auto memRefDimSize = resultMemTy.getDimSize(memRefRank - 1); - comet_debug() << "memRefDimSize for alloc: " << memRefDimSize << "\n"; - - Value upperBound; - if (memRefDimSize == 1) /// size of value array comes from temporary sparse tensor and Dimsize of alloc is one - { - upperBound = rewriter.create(loc, memAllocForNNZ, alloc_zero_loc); - } - else - { - /// size of value array comes from read_input_sizes_2D_f64, and alloc dimsize can be only expected size - auto expectedMemRefSize = sp_op.getTotalParamCount(); - comet_debug() << "tensorRanks: " << tensorRanks << "\n"; - comet_debug() << "expectedMemRefSize: " << expectedMemRefSize << "\n"; - assert(memRefDimSize == expectedMemRefSize); - upperBound = op->getOperand(0).getDefiningOp()->getOperand(indexValueSize); - } comet_debug() << "Upper Bound:\n"; comet_vdump(upperBound); auto lowerBound = rewriter.create(loc, 0); @@ -591,12 +505,8 @@ namespace rewriter.setInsertionPointToStart(loop.getBody()); /// Build loop body - int indexValuePtr = (tensorRanks * 4); /// 4 corresponding to pos, crd - auto alloc_op = op->getOperand(0).getDefiningOp()->getOperand(indexValuePtr).getDefiningOp()->getOperand(0); - comet_debug() << " ValueAllocOp"; - comet_vdump(alloc_op); std::vector indices = {loop.getInductionVar()}; - auto load_rhs = rewriter.create(loc, alloc_op, indices); + auto load_rhs = rewriter.create(loc, sp_tensor_values, indices); auto res_load = rewriter.create(loc, res, alloc_zero_loc); auto reduce = rewriter.create(loc, load_rhs, res_load); rewriter.create(loc, reduce, res, alloc_zero_loc); @@ -605,10 +515,14 @@ namespace rewriter.restoreInsertionPoint(insertPt); comet_vdump(loop); } - + else + { + return failure(); + } + rewriter.setInsertionPoint(op); /// Important to replace all uses of this operation with the new one, otherwise, the current op won't be lowered. - op.replaceAllUsesWith(res); - rewriter.eraseOp(op); + memref::LoadOp load = rewriter.create(op->getLoc(), res, ValueRange(cst_zero)); + rewriter.replaceOp(op, load); return success(); } @@ -641,26 +555,6 @@ namespace comet_vdump(const_index_0); std::vector alloc_zero_loc = {const_index_0}; - if (auto toTensorOp = llvm::dyn_cast_if_present(rhs.getDefiningOp())) - { - rhs = toTensorOp.getMemref(); - // comet_debug() << "RHS is a tensor\n"; - // rhs = rewriter.create(loc, rhs, alloc_zero_loc); - // comet_vdump(rhs); - } - if (auto toTensorOp = llvm::dyn_cast_if_present(lhs.getDefiningOp())) - { - lhs = toTensorOp.getMemref(); - // comet_debug() << "RHS is a tensor\n"; - // rhs = rewriter.create(loc, rhs, alloc_zero_loc); - // comet_vdump(rhs); - } - // if (lhsType.isa()) - // { - // comet_debug() << "LHS is a tensor\n"; - // lhs = rewriter.create(loc, lhs, alloc_zero_loc); - // } - Value res; bool res_comes_from_setop = false; for (auto u : op.getOperation()->getResult(0).getUsers()) @@ -677,12 +571,6 @@ namespace { res = cast(*(++res.getUsers().begin())).getRhs(); } - if(auto toTensor = mlir::dyn_cast_or_null(res.getDefiningOp())) - { - res = toTensor.getMemref(); - } - comet_debug() << "Result from SetOp:\n"; - comet_vdump(res); res_comes_from_setop = true; break; } @@ -703,41 +591,51 @@ namespace Value res_val; if (op_attr.compare("+") == 0) { - rewriter.create(loc, ValueRange{lhs, rhs}, ValueRange(res)); - // res_val = rewriter.create(loc, lhs, rhs); + res_val = rewriter.create(loc, ValueRange{lhs, rhs}, ValueRange(res)).getResultTensors()[0]; } else if (op_attr.compare("-") == 0) { - rewriter.create(loc, ValueRange{lhs, rhs}, ValueRange(res)); - // res_val = rewriter.create(loc, lhs, rhs); + res_val = rewriter.create(loc, ValueRange{lhs, rhs}, ValueRange(res)).getResultTensors()[0]; } else if (op_attr.compare("*") == 0) { - rewriter.create(loc, ValueRange{lhs, rhs}, ValueRange(res)); - // res_val = rewriter.create(loc, lhs, rhs); + res_val = rewriter.create(loc, ValueRange{lhs, rhs}, ValueRange(res)).getResultTensors()[0]; } else if (op_attr.compare("/") == 0) { - rewriter.create(loc, ValueRange{lhs, rhs}, ValueRange(res)); - // res_val = rewriter.create(loc, lhs, rhs); + res_val = rewriter.create(loc, ValueRange{lhs, rhs}, ValueRange(res)).getResultTensors()[0]; } else { llvm::errs() << "ERROR: Unsuported Operation\n"; } - - comet_vdump(res_val); - /// store res_val to res - // rewriter.create(loc, res_val, res); - // [[maybe_unused]] auto storeOp = rewriter.create(loc, res_val, res, alloc_zero_loc); - comet_vdump(storeOp); - - op.replaceAllUsesWith(res); + + op.replaceAllUsesWith(res_val); rewriter.eraseOp(op); return success(); } }; /// ScalarOpsLowering +class ConvertSetOp : public OpConversionPattern { + using OpConversionPattern::OpConversionPattern; + + LogicalResult + matchAndRewrite(TensorSetOp op, OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const final { + + auto opAdaptor = llvm::cast(adaptor); + Value lhs = opAdaptor.getLhs(); + Value rhs = opAdaptor.getRhs(); + rewriter.replaceUsesWithIf(rhs, lhs, [&](OpOperand& use) { + auto user = use.getOwner(); + auto ancestor = op->getBlock()->findAncestorOpInBlock(*user); + return (ancestor && op->isBeforeInBlock(ancestor)); + }); + rewriter.eraseOp(op); + return success(); + } +}; + } /// end anonymous namespace. /// This is a partial lowering to linear algebra of the tensor algebra operations that are @@ -778,10 +676,15 @@ void LowerTensorAlgebraToSCFPass::runOnOperation() scf::SCFDialect, ArithDialect, memref::MemRefDialect, + func::FuncDialect, bufferization::BufferizationDialect>(); - target.addLegalOp(); - + target.addLegalDialect(); + target.addIllegalOp(); /// Now that the conversion target has been defined, we just need to provide /// the set of patterns that will lower the TA operations. @@ -789,7 +692,8 @@ void LowerTensorAlgebraToSCFPass::runOnOperation() patterns.insert(&getContext()); + ConstantOpLowering, + ConvertSetOp>(&getContext()); /// With the target and rewrite patterns defined, we can now attempt the /// conversion. The conversion will signal failure if any of our `illegal` /// operations were not converted successfully. diff --git a/lib/Conversion/TritonToCuda/CMakeLists.txt b/lib/Conversion/TritonToCuda/CMakeLists.txt index a5851d54..d0cdca75 100644 --- a/lib/Conversion/TritonToCuda/CMakeLists.txt +++ b/lib/Conversion/TritonToCuda/CMakeLists.txt @@ -14,4 +14,5 @@ add_llvm_library(COMETTritonToCuda TritonGPUTransforms TritonNvidiaGPUTransforms NVGPUIR + COMETGPUUtils ) diff --git a/lib/Conversion/TritonToCuda/TritonToCudaPass.cpp b/lib/Conversion/TritonToCuda/TritonToCudaPass.cpp index 8cfbcc77..29fdf62a 100644 --- a/lib/Conversion/TritonToCuda/TritonToCudaPass.cpp +++ b/lib/Conversion/TritonToCuda/TritonToCudaPass.cpp @@ -1,33 +1,65 @@ - +// +// Copyright 2022 Battelle Memorial Institute +// +// Redistribution and use in source and binary forms, with or without modification, +// are permitted provided that the following conditions are met: +// +// 1. Redistributions of source code must retain the above copyright notice, this list of conditions +// and the following disclaimer. +// +// 2. Redistributions in binary form must reproduce the above copyright notice, this list of conditions +// and the following disclaimer in the documentation and/or other materials provided with the distribution. +// +// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED +// WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR +// PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, +// INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE +// GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, +// WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +// OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. +// #include #include #include #include "comet/Conversion/TritonToCuda/TritonToCudaPass.h" - +#include "comet/Conversion/GpuUtils/GpuUtils.h" +#include "comet/Dialect/Utils/Utils.h" +#include "mlir/Conversion/ArithToLLVM/ArithToLLVM.h" +#include "mlir/Conversion/IndexToLLVM/IndexToLLVM.h" +#include "mlir/Conversion/SCFToControlFlow/SCFToControlFlow.h" +#include "mlir/Dialect/Func/IR/FuncOps.h" +#include "mlir/Dialect/GPU/IR/GPUDialect.h" +#include "mlir/Dialect/GPU/Transforms/Passes.h" +#include "mlir/Dialect/LLVMIR/LLVMDialect.h" +#include "mlir/Dialect/LLVMIR/NVVMDialect.h" +#include "mlir/Dialect/MemRef/IR/MemRef.h" +#include "mlir/IR/Builders.h" +#include "mlir/IR/BuiltinOps.h" +#include "mlir/IR/BuiltinTypes.h" #include "mlir/IR/Value.h" #include "mlir/Pass/Pass.h" -#include "mlir/Dialect/MemRef/IR/MemRef.h" -#include "mlir/Dialect/GPU/IR/GPUDialect.h" -#include "mlir/Dialect/SCF/IR/SCF.h" -#include "mlir/Dialect/Affine/IR/AffineOps.h" -#include "mlir/Transforms/DialectConversion.h" #include "mlir/Pass/PassManager.h" +#include "mlir/Support/LLVM.h" +#include "mlir/Target/LLVMIR/Export.h" +#include "mlir/Transforms/DialectConversion.h" #include "mlir/Transforms/Passes.h" #include "mlir/Dialect/GPU/IR/GPUDialect.h" #include "mlir/Dialect/Func/IR/FuncOps.h" - +#include "mlir/Conversion/ArithToLLVM/ArithToLLVM.h" #include "mlir/Dialect/Arith/Transforms/Passes.h" -#include "triton/Conversion/NVGPUToLLVM/NVGPUToLLVMPass.h" -#include "triton/Conversion/TritonGPUToLLVM/TritonGPUToLLVMPass.h" -#include "triton/Dialect/Triton/Transforms/Passes.h" -#include "triton/Dialect/TritonGPU/IR/Dialect.h" -#include "triton/Dialect/TritonNvidiaGPU/IR/Dialect.h" -#include "triton/Dialect/NVGPU/IR/Dialect.h" +// #include "triton/Conversion/NVGPUToLLVM/TritonGPUToLLVMPass.h" +// #include "triton/Conversion/TritonGPUToLLVM/TritonGPUToLLVMPass.h" +#include "third_party/nvidia/include/NVGPUToLLVM/NVGPUToLLVMPass.h" +#include "third_party/nvidia/include/TritonNVIDIAGPUToLLVM/Passes.h" +#include "triton/Conversion/TritonGPUToLLVM/Passes.h" #include "triton/Conversion/TritonToTritonGPU/TritonToTritonGPUPass.h" #include "triton/Dialect/Triton/IR/Dialect.h" +#include "triton/Dialect/Triton/Transforms/Passes.h" #include "triton/Dialect/TritonGPU/Transforms/Passes.h" #include "triton/Dialect/TritonNvidiaGPU/Transforms/Passes.h" +#include "llvm/ADT/MapVector.h" +#include "llvm/ADT/STLExtras.h" #include "mlir/Dialect/LLVMIR/NVVMDialect.h" #include "mlir/Conversion/SCFToControlFlow/SCFToControlFlow.h" @@ -39,173 +71,146 @@ #include "llvm/Target/TargetOptions.h" #include "llvm/IR/LegacyPassManager.h" #include +#include #define GEN_PASS_CLASSES #include "comet/Conversion/TritonToCuda/Passes.h" using namespace mlir; -class LowerTritonDeviceToCuda - : public mlir::comet::LowerTritonDeviceToCudaBase { -public: - LowerTritonDeviceToCuda() = default; +namespace cuda_device { - - LowerTritonDeviceToCuda(int numWarps, - int threadsPerWarp, - int numCTAs, - int numStages, - int computeCapability) { - - this->numWarps = numWarps; - this->threadsPerWarp = threadsPerWarp; - this->numCTAs = numCTAs; - this->numStages = numStages; - this->computeCapability = computeCapability; - } - - void runOnOperation() override { - mlir::ModuleOp modOp = getOperation(); - - std::vector TTFuncs; - modOp->walk([&TTFuncs](mlir::triton::FuncOp op) { - TTFuncs.push_back(op); - }); - - auto tempMod = ModuleOp::create(modOp.getLoc()); - OpBuilder builder(tempMod.getBodyRegion()); - - - for(auto ttFunc: TTFuncs) - { - builder.clone(*ttFunc.getOperation()); - // ttFunc->erase(); - } - - if(TTFuncs.empty()) - { - return signalPassFailure(); - } + bool add_ttir_passes(ModuleOp &mod) + { + PassManager pm(mod.getContext()); - PassManager pm(tempMod.getContext()); - + pm.addPass(createInlinerPass()); + pm.addPass(triton::createRewriteTensorPointerPass()); + pm.addPass(triton::createCombineOpsPass()); pm.addPass(createCanonicalizerPass()); + pm.addPass(triton::createReorderBroadcastPass()); pm.addPass(createCSEPass()); - pm.addPass(createCSEPass()); - pm.addPass(createSymbolDCEPass()); - pm.addPass(createCanonicalizerPass()); pm.addPass(createLoopInvariantCodeMotionPass()); - pm.addPass(mlir::triton::createConvertTritonToTritonGPUPass(numWarps.getValue(), threadsPerWarp.getValue(), numCTAs.getValue(), computeCapability.getValue())); - pm.addPass(triton::gpu::createCoalescePass()); + pm.addPass(triton::createLoopUnrollPass()); + if (failed(pm.run(mod))) { + return false; + } + return true; + } + + bool add_ttgir_passes(ModuleOp &mod, int32_t numWarps, int32_t threadsPerWarp, int32_t numStages, int32_t numCTAs, int32_t computeCapability) + { + PassManager pm(mod.getContext()); + + pm.addPass(mlir::triton::createConvertTritonToTritonGPUPass( + "cuda:" + std::to_string(computeCapability), numWarps, + threadsPerWarp, numCTAs)); + pm.addPass(triton::gpu::createTritonGPUCoalesce()); + if (computeCapability / 10 >= 8) { + pm.addPass(triton::gpu::createTritonGPUF32DotTC()); + } pm.addPass(createTritonNvidiaGPUPlanCTAPass()); - pm.addPass(mlir::triton::createRewriteTensorPointerPass(computeCapability.getValue())); - pm.addPass(triton::gpu::createRemoveLayoutConversionsPass()); - pm.addPass(triton::gpu::createOptimizeThreadLocalityPass()); - pm.addPass(triton::gpu::createAccelerateMatmulPass(computeCapability.getValue())); - pm.addPass(triton::gpu::createAccelerateMatmulPass(computeCapability.getValue())); - pm.addPass(triton::gpu::createRemoveLayoutConversionsPass()); - pm.addPass(triton::gpu::createOptimizeDotOperandsPass()); + pm.addPass(triton::gpu::createTritonGPURemoveLayoutConversions()); + pm.addPass(triton::gpu::createTritonGPUOptimizeThreadLocality()); + pm.addPass(triton::gpu::createTritonGPUAccelerateMatmul()); + pm.addPass(triton::gpu::createTritonGPURemoveLayoutConversions()); + mlir::triton::gpu::TritonGPUOptimizeDotOperandsOptions options; + options.hoistLayoutConversion = computeCapability >= 80; + pm.addPass(triton::gpu::createTritonGPUOptimizeDotOperands(options)); pm.addPass(createCSEPass()); - pm.addPass(triton::gpu::createPipelinePass(numStages.getValue(),numWarps.getValue(),numCTAs.getValue(), computeCapability.getValue())); - pm.addPass(createTritonNvidiaGPUMaterializeLoadStorePass(numWarps.getValue(), computeCapability.getValue())); - pm.addPass(triton::gpu::createPrefetchPass()); - pm.addPass(triton::gpu::createOptimizeDotOperandsPass()); - pm.addPass(triton::gpu::createDecomposeConversionsPass()); - pm.addPass(createTritonNvidiaGPUWSFixupMissingAttrs()); - pm.addPass(triton::gpu::createReorderInstructionsPass()); + if (computeCapability / 10 >= 8) { + pm.addPass(triton::gpu::createTritonGPUOptimizeAccumulatorInit()); + pm.addPass(triton::gpu::createTritonGPUCombineTensorSelectAndIf()); + mlir::triton::gpu::TritonGPUPipelineOptions options; + options.numStages = numStages; + pm.addPass(triton::gpu::createTritonGPUPipeline(options)); + } + pm.addPass(triton::gpu::createTritonGPUPrefetch()); + pm.addPass(triton::gpu::createTritonGPUOptimizeDotOperands(options)); + pm.addPass(triton::gpu::createTritonGPUReduceDataDuplication()); + pm.addPass(triton::gpu::createTritonGPUReorderInstructions()); pm.addPass(createCSEPass()); pm.addPass(createSymbolDCEPass()); - pm.addPass(createTritonNvidiaGPUWSFixupMissingAttrs()); - pm.addPass(createCanonicalizerPass()); - pm.addPass(mlir::createConvertSCFToCFPass()); - pm.addPass(createConvertIndexToLLVMPass()); - pm.addPass(mlir::triton::createConvertTritonGPUToLLVMPass()); - pm.addPass(mlir::triton::createConvertNVGPUToLLVMPass()); - // pm.addPass(createConvertToLLVMPass()); + if (computeCapability / 10 >= 9) { + pm.addPass(createTritonNvidiaGPUFenceInsertionPass()); + pm.addPass(createTritonNvidiaGPUTMALoweringPass()); + } pm.addPass(createCanonicalizerPass()); - pm.addPass(createCSEPass()); - pm.addPass(createSymbolDCEPass()); - - if (failed(pm.run(tempMod))) { - signalPassFailure(); - return; + + if (failed(pm.run(mod))) { + return false; } + return true; + } + + bool add_llir_passes(ModuleOp &mod, int32_t computeCapability) + { + PassManager pm(mod.getContext()); + + pm.addPass(triton::NVIDIA::createDecomposeUnsupportedConversionsPass()); + pm.addPass(triton::gpu::createTritonGPUCombineTensorSelectAndIf()); + pm.addPass(createConvertSCFToCFPass()); + pm.addPass(createConvertIndexToLLVMPass()); + pm.addPass(triton::gpu::createAllocateSharedMemoryPass()); + pm.addPass(mlir::triton::createRewriteTensorPointerPass()); + pm.addPass( + triton::createConvertTritonGPUToLLVMPass(computeCapability)); + pm.addPass(triton::createConvertNVGPUToLLVMPass()); + pm.addPass(createArithToLLVMConversionPass()); + pm.addPass(createCanonicalizerPass()); + pm.addPass(createCSEPass()); + pm.addPass(createSymbolDCEPass()); + + if (failed(pm.run(mod))) { + return false; + } + return true; + } +} +class LowerTritonDeviceToCuda + : public mlir::comet::LowerTritonDeviceToCudaBase { +public: + LowerTritonDeviceToCuda() = default; + + LowerTritonDeviceToCuda( + int numWarps, int threadsPerWarp, int numCTAs, int numStages, + int computeCapability, + mlir::tensorAlgebra::GPUCompilationFormat codeFormat) { + + this->numWarps = numWarps; + this->threadsPerWarp = threadsPerWarp; + this->numCTAs = numCTAs; + this->numStages = numStages; + this->computeCapability = computeCapability; + this->codeFormat = codeFormat; + } - modOp->setAttrs(tempMod->getAttrs()); - modOp->setAttr("gpu.container_module", builder.getUnitAttr()); - builder.setInsertionPointToEnd(&modOp.getRegion().getBlocks().back()); - // tempMod.getRegion(). - // mlir::gpu::SerializeToBlobPass() - llvm::LLVMContext llvmContext; - auto llvmModule = translateModuleToLLVMIR(tempMod.getOperation(), llvmContext); - if (!llvmModule) + void runOnOperation() override { + mlir::ModuleOp modOp = getOperation(); + OpBuilder builder(modOp); + std::string chip = "sm_" + std::to_string(computeCapability.getValue()); + auto target = mlir::NVVM::NVVMTargetAttr::get(modOp->getContext(), 3, + "nvptx64-nvidia-cuda", chip); + auto add_ttgir_passes = [this](ModuleOp& modOp) { return cuda_device::add_ttgir_passes(modOp, numWarps, threadsPerWarp, numStages, numCTAs, computeCapability); }; + auto add_llir_passes = [this](ModuleOp& modOp) { return cuda_device::add_llir_passes(modOp, computeCapability); }; + if(failed(specializeGpuKernel(builder, modOp, this->codeFormat, target, cuda_device::add_ttir_passes, add_ttgir_passes, add_llir_passes))) { return signalPassFailure(); } - std::string TargetTriple = "nvptx64-nvidia-cuda"; - llvmModule->setTargetTriple(TargetTriple); - Location loc = tempMod.getLoc(); - std::string error; - const llvm::Target *target = llvm::TargetRegistry::lookupTarget(llvmModule->getTargetTriple(), error); - if (!target) { - emitError(loc, Twine("failed to lookup target: ") + error); - return signalPassFailure(); - } - llvm::TargetOptions opt; - // if (enable_fp_fusion) - // opt.AllowFPOpFusion = llvm::FPOpFusion::Fast; - opt.UnsafeFPMath = false; - opt.NoInfsFPMath = false; - opt.NoNaNsFPMath = true; - opt.TrapUnreachable = true; - llvm::TargetMachine *machine = - target->createTargetMachine(llvmModule->getTargetTriple(), "sm_"+std::to_string(computeCapability.getValue()), "", opt, llvm::Reloc::PIC_, std::nullopt, llvm::CodeGenOptLevel::Aggressive); - - if (!machine) { - emitError(loc, "failed to create target machine"); - return signalPassFailure(); - } - - llvmModule->setDataLayout(machine->createDataLayout()); - std::string result; - { - llvm::raw_string_ostream stream(result); - llvm::buffer_ostream pstream(stream); - for (llvm::Function &f : llvmModule->functions()) - f.addFnAttr(llvm::Attribute::AlwaysInline); - llvm::legacy::PassManager pass; - // emit - // auto fileType = isObject ? llvm::CodeGenFileType::ObjectFile - // : llvm::CodeGenFileType::AssemblyFile; - auto fileType = llvm::CodeGenFileType::AssemblyFile; - machine->addPassesToEmitFile(pass, pstream, nullptr, fileType); - pass.run(*llvmModule); - } - - auto funcOp = *modOp.getOps().begin(); - builder.setInsertionPointToStart(&funcOp.getFunctionBody().front()); - LLVM::createGlobalString(modOp->getLoc(), builder, "ptx", result, LLVM::linkage::Linkage::Private ); - // modOp->setAttr("gpu.ptx", builder.getStringAttr(result)); - // std::cout << result << std::endl; - - - // builder.clone(*tempMod.getOperation()); } }; - -std::unique_ptr> +std::unique_ptr> mlir::comet::createLowerTritonDeviceToCudaPass() { return std::make_unique<::LowerTritonDeviceToCuda>(); } -std::unique_ptr> -mlir::comet::createLowerTritonDeviceToCudaPass(int numWarps, - int threadsPerWarp, - int numCTAs, - int numStages, - int computeCapability) { - return std::make_unique<::LowerTritonDeviceToCuda>(numWarps, threadsPerWarp, numCTAs, numStages, computeCapability); +std::unique_ptr> +mlir::comet::createLowerTritonDeviceToCudaPass( + int numWarps, int threadsPerWarp, int numCTAs, int numStages, + int computeCapability, mlir::tensorAlgebra::GPUCompilationFormat format) { + return std::make_unique<::LowerTritonDeviceToCuda>( + numWarps, threadsPerWarp, numCTAs, numStages, computeCapability, format); } class LowerGpuHostToCuda @@ -216,311 +221,16 @@ class LowerGpuHostToCuda void runOnOperation() override { mlir::ModuleOp modOp = getOperation(); OpBuilder builder(modOp); - std::map funcs; - std::map gpu_to_triton_kernel; - - modOp->walk([&gpu_to_triton_kernel](mlir::triton::FuncOp TTFuncOp) { - gpu_to_triton_kernel[TTFuncOp->getAttrOfType("origin").str()] = TTFuncOp.getName().str(); - TTFuncOp.erase(); - }); - - auto memrefI32 = MemRefType::get({ShapedType::kDynamic}, builder.getIntegerType(32)); - auto memrefF32 = MemRefType::get({ShapedType::kDynamic}, builder.getF32Type()); - auto memrefI64 = MemRefType::get({ShapedType::kDynamic}, builder.getIntegerType(64)); - auto memrefF64 = MemRefType::get({ShapedType::kDynamic}, builder.getF64Type()); - - auto mallocI32Type = builder.getFunctionType({builder.getIndexType()} , builder.getIndexType()); - auto mallocI32 = builder.create(builder.getUnknownLoc(), "cudaMallocI32", mallocI32Type); - mallocI32.setVisibility(mlir::SymbolTable::Visibility::Private); - modOp.push_back(mallocI32); - - auto mallocI64Type = builder.getFunctionType({builder.getIndexType()} , builder.getIndexType()); - auto mallocI64 = builder.create(builder.getUnknownLoc(), "cudaMallocI64", mallocI64Type); - mallocI64.setVisibility(mlir::SymbolTable::Visibility::Private); - modOp.push_back(mallocI64); - - auto mallocF32Type = builder.getFunctionType({builder.getIndexType()} , builder.getIndexType()); - auto mallocF32 = builder.create(builder.getUnknownLoc(), "cudaMallocF32", mallocF32Type); - mallocF32.setVisibility(mlir::SymbolTable::Visibility::Private); - modOp.push_back(mallocF32); - - auto mallocF64Type = builder.getFunctionType({builder.getIndexType()} , builder.getIndexType()); - auto mallocF64 = builder.create(builder.getUnknownLoc(), "cudaMallocF64", mallocF64Type); - mallocF64.setVisibility(mlir::SymbolTable::Visibility::Private); - modOp.push_back(mallocF64); - - auto cudaMemcpyI32Type = builder.getFunctionType({builder.getIndexType(), memrefI32, builder.getIndexType()}, {}); - auto cudaMemcpyI32 = builder.create(builder.getUnknownLoc(), "cudaMemcpyI32", cudaMemcpyI32Type); - cudaMemcpyI32.setVisibility(mlir::SymbolTable::Visibility::Private); - modOp.push_back(cudaMemcpyI32); - - auto cudaMemcpyI64Type = builder.getFunctionType({builder.getIndexType(), memrefI64, builder.getIndexType()}, {}); - auto cudaMemcpyI64 = builder.create(builder.getUnknownLoc(), "cudaMemcpyI64", cudaMemcpyI64Type); - cudaMemcpyI64.setVisibility(mlir::SymbolTable::Visibility::Private); - modOp.push_back(cudaMemcpyI64); - - auto cudaMemcpyF32Type = builder.getFunctionType({builder.getIndexType(), memrefF32, builder.getIndexType()}, {}); - auto cudaMemcpyF32 = builder.create(builder.getUnknownLoc(), "cudaMemcpyF32", cudaMemcpyF32Type); - cudaMemcpyF32.setVisibility(mlir::SymbolTable::Visibility::Private); - modOp.push_back(cudaMemcpyF32); - - auto cudaMemcpyF64Type = builder.getFunctionType({builder.getIndexType(), memrefF64, builder.getIndexType()}, {}); - auto cudaMemcpyF64 = builder.create(builder.getUnknownLoc(), "cudaMemcpyF64", cudaMemcpyF64Type); - cudaMemcpyF64.setVisibility(mlir::SymbolTable::Visibility::Private); - modOp.push_back(cudaMemcpyF64); - - auto cudaLaunchKernelT = builder.getFunctionType({builder.getIndexType(), builder.getIndexType(), builder.getIndexType(), builder.getIndexType(), builder.getIndexType(), builder.getIndexType(), MemRefType::get({ShapedType::kDynamic}, builder.getIndexType()), LLVM::LLVMPointerType::get(builder.getContext()), builder.getIndexType(), builder.getIntegerType(32), builder.getIndexType(), builder.getIndexType()}, {}); - auto cudaLaunchKernel = builder.create(builder.getUnknownLoc(), "cudaLaunchKernel", cudaLaunchKernelT); - cudaLaunchKernel.setVisibility(mlir::SymbolTable::Visibility::Private); - modOp.push_back(cudaLaunchKernel); - - auto cudaSetModuleImageT = builder.getFunctionType({ LLVM::LLVMPointerType::get(builder.getContext())}, {}); - auto cudaSetModuleImage = builder.create(builder.getUnknownLoc(), "cudaSetModuleImage", cudaSetModuleImageT); - cudaSetModuleImage.setVisibility(mlir::SymbolTable::Visibility::Private); - modOp.push_back(cudaSetModuleImage); - - auto cudaFreeT = builder.getFunctionType({builder.getIndexType()}, {}); - auto cudaFree = builder.create(builder.getUnknownLoc(), "cudaFree", cudaFreeT); - cudaFree.setVisibility(mlir::SymbolTable::Visibility::Private); - modOp.push_back(cudaFree); - - - std::vector launchOps; + declare_vendor_funcs(builder, modOp, "cuda"); - modOp->walk([&launchOps](mlir::gpu::LaunchFuncOp launchOp){ - launchOps.push_back(launchOp); - }); - - std::set initFuncs; - std::vector toAlloc; - for(auto launchOp: launchOps) + if(failed(specializeGpuHost(builder, modOp, std::string("cuda")))) { - if(initFuncs.find(launchOp->getParentOfType()) == initFuncs.end()) - { - builder.setInsertionPointToStart(&launchOp->getParentOfType().getFunctionBody().front()); - Value ptx = builder.create(launchOp->getLoc(), LLVM::LLVMPointerType::get(&getContext()) ,"ptx"); - builder.create(launchOp.getLoc(), "cudaSetModuleImage", TypeRange(), ValueRange({ptx})); - - initFuncs.insert(launchOp->getParentOfType()); - } - builder.setInsertionPoint(launchOp); - Value sharedMem = builder.create(launchOp->getLoc(), modOp->getAttrOfType("triton_gpu.shared").getInt(), 32); - launchOp.getDynamicSharedMemorySizeMutable().assign(sharedMem); - for(auto& operand: launchOp->getOpOperands()) - { - // if(!isa(operand.get().getType()) ) - // { - // // auto i32Operand = builder.create(operand.get().getLoc(), builder.getIntegerType(32), operand.get()); - // // // operand.set(i32Operand); - // // cudaCallArgs.push_back(i32Operand); - // } - // else - if(isa(operand.get().getType())) - { - toAlloc.push_back(operand.get()); - } - } - } - - auto cmp = [](Value a, Value b){return a.getAsOpaquePointer() < b.getAsOpaquePointer();}; - std::set uniqueGpuAllocs(toAlloc.begin(), toAlloc.end(), cmp); - for(Value alloc: uniqueGpuAllocs) - { - MemRefType allocType = alloc.getType().cast(); - mlir::gpu::AllocOp gpuAlloc; - if(auto defOp = alloc.getDefiningOp()) - { - builder.setInsertionPointAfter(defOp); - } - else if(alloc.isa()) - { - builder.setInsertionPointToStart(alloc.getParentBlock()); - } - else { - assert(false && "Value has not defining Op and is not a block argument."); - } - - if(allocType.hasStaticShape()) - { - gpuAlloc = builder.create(alloc.getLoc(), allocType, ValueRange(), ValueRange(), ValueRange()); - } - else { - std::vector dynDims; - for(size_t i = 0; i < allocType.getShape().size(); i++) - { - if(allocType.isDynamicDim(i)) - { - dynDims.push_back(builder.create(alloc.getLoc(), alloc, i)); - } - } - gpuAlloc = builder.create(alloc.getLoc(), allocType, ValueRange(), dynDims, ValueRange()); - } - - auto op = alloc.getDefiningOp() != NULL ? alloc.getDefiningOp()->getResult(0) : alloc; - for(auto& use: llvm::make_early_inc_range(op.getUses())) - { - if(mlir::gpu::LaunchFuncOp launchOp = dyn_cast(use.getOwner())) - { - builder.setInsertionPoint(launchOp); - builder.create(launchOp->getLoc(), TypeRange(), ValueRange(), gpuAlloc.getMemref(), alloc); - use.set(gpuAlloc.getMemref()); - builder.setInsertionPointAfter(launchOp); - builder.create(launchOp->getLoc(), TypeRange(), ValueRange(), alloc, gpuAlloc.getMemref()); - } - } - } - - - std::vector gpuAllocs; - modOp->walk([&gpuAllocs](mlir::gpu::AllocOp gpuAllocOp){ - gpuAllocs.push_back(gpuAllocOp); - }); - - for(auto gpuAlloc: gpuAllocs) - { - builder.setInsertionPoint(gpuAlloc->getBlock()->getTerminator()); - builder.create(gpuAlloc->getLoc(), ValueRange(), gpuAlloc.getResult(0)); - - builder.setInsertionPointAfter(gpuAlloc); - Value allocSize; - if( gpuAlloc.getMemref().getType().hasStaticShape()) - { - allocSize = builder.create(gpuAlloc->getLoc(), gpuAlloc.getMemref().getType().getShape()[0]).getResult(); - } - else - { - allocSize = gpuAlloc->getOperand(0); - } - - // builder.create(gpuAlloc.getLoc(), gpuAlloc.getMemref(), 0); - if(FloatType floatType = gpuAlloc.getMemref().getType().getElementType().dyn_cast()) - { - int width = floatType.getWidth(); - auto cudaOp = builder.create(gpuAlloc->getLoc(), "cudaMallocF"+std::to_string(width), TypeRange(builder.getIndexType()), ValueRange(allocSize)); - gpuAlloc->replaceAllUsesWith(cudaOp); - } - else if(IntegerType intType = gpuAlloc.getMemref().getType().getElementType().dyn_cast()) - { - int width = intType.getWidth(); - auto cudaOp = builder.create(gpuAlloc->getLoc(), "cudaMallocI"+std::to_string(width), TypeRange(builder.getIndexType()), ValueRange(allocSize)); - gpuAlloc->replaceAllUsesWith(cudaOp); - } - gpuAlloc->erase(); - } - - - - std::vector gpuCopies; - modOp->walk([&gpuCopies](mlir::gpu::MemcpyOp gpuCopy){ - gpuCopies.push_back(gpuCopy); - }); - - - for(auto cpy: gpuCopies) - { - builder.setInsertionPoint(cpy); - auto hToD = builder.create(cpy->getLoc(), 0); - auto dToH = builder.create(cpy->getLoc(), 1); - - if(cpy.getOperand(0).getDefiningOp() && (isa(cpy.getOperand(0).getDefiningOp()) && cast(cpy.getOperand(0).getDefiningOp()).getCallee().starts_with("cudaMalloc") )) - { - auto cast = builder.create(cpy->getLoc(), MemRefType::get({ShapedType::kDynamic}, cpy.getSrc().getType().getElementType()), cpy.getSrc()); - if(IntegerType intType = cpy.getOperand(1).getType().cast().getElementType().dyn_cast()) - { - int width = intType.getWidth(); - builder.create(cpy->getLoc(), "cudaMemcpyI"+std::to_string(width), TypeRange(), ValueRange({cpy.getOperand(0), cast, hToD})); - } - else if(FloatType floatType = cpy.getOperand(1).getType().cast().getElementType().dyn_cast()) - { - int width = floatType.getWidth(); - builder.create(cpy->getLoc(), "cudaMemcpyF"+std::to_string(width), TypeRange(), ValueRange({cpy.getOperand(0), cast, hToD})); - } - } - else if(cpy.getOperand(1).getDefiningOp() && (isa(cpy.getOperand(1).getDefiningOp()) && cast(cpy.getOperand(1).getDefiningOp()).getCallee().starts_with("cudaMalloc") )) - { - auto cast = builder.create(cpy->getLoc(), MemRefType::get({ShapedType::kDynamic}, cpy.getDst().getType().getElementType()), cpy.getDst()); - - if(IntegerType intType = cpy.getOperand(0).getType().cast().getElementType().dyn_cast()) - { - int width = intType.getWidth(); - builder.create(cpy->getLoc(), "cudaMemcpyI"+std::to_string(width), TypeRange(), ValueRange({cpy.getOperand(1), cast, dToH})); - } - else if(FloatType floatType = cast.getType().cast().getElementType().dyn_cast()) - { - int width = floatType.getWidth(); - builder.create(cpy->getLoc(), "cudaMemcpyF"+std::to_string(width), TypeRange(), ValueRange({cpy.getOperand(1), cast, dToH})); - } - } - - cpy->erase(); - } - - for(auto launchOp: launchOps) - { - builder.setInsertionPoint(launchOp); - auto zeroIndex = builder.create(launchOp->getLoc(), 0); - - int64_t numOps = launchOp.getNumKernelOperands(); - std::vector ptr_ops; - memref::AllocaOp temp = builder.create(launchOp->getLoc(), MemRefType::get({1}, launchOp.getGridSizeX().getType())); - builder.create(launchOp->getLoc(), launchOp.getGridSizeY(), temp, ValueRange({zeroIndex})); - ptr_ops.push_back(builder.create(launchOp->getLoc(), temp)); - temp = builder.create(launchOp->getLoc(), MemRefType::get({1}, launchOp.getGridSizeY().getType())); - builder.create(launchOp->getLoc(), launchOp.getGridSizeX(), temp, ValueRange({zeroIndex})); - ptr_ops.push_back(builder.create(launchOp->getLoc(), temp)); - - for(auto op: launchOp.getKernelOperands()) - { - if(op.getType().isa()) - { - ptr_ops.push_back(builder.create(launchOp->getLoc(), op)); - } - else - { - auto temp = builder.create(launchOp->getLoc(), MemRefType::get({1}, op.getType())); - builder.create(launchOp->getLoc(), op, temp, ValueRange({zeroIndex})); - ptr_ops.push_back(builder.create(launchOp->getLoc(), temp)); - } - } - - auto args = builder.create(launchOp->getLoc(), MemRefType::get({numOps + 2}, builder.getIndexType())); - auto args_dynamic = builder.create(launchOp->getLoc(), MemRefType::get({ShapedType::kDynamic}, args.getType().getElementType()), args); - for(size_t i = 0; i < ptr_ops.size(); i++) - { - Value op = ptr_ops[i]; - builder.create(launchOp->getLoc(), op, args, builder.create(launchOp->getLoc(), i).getResult()); - } - std::string& funcName = gpu_to_triton_kernel[(launchOp.getKernelModuleName().strref()+"::"+launchOp.getKernelName().strref()).str()]; - if(funcs.find(funcName) == funcs.end()) - { - funcs[funcName] = LLVM::createGlobalString(modOp->getLoc(), builder, funcName+"_str", funcName, LLVM::linkage::Linkage::Private ); - } - Value numWarps = builder.create(launchOp->getLoc(), modOp->getAttrOfType("triton_gpu.num-warps").getInt()); - Value threadsPerWarp = builder.create(launchOp->getLoc(), modOp->getAttrOfType("triton_gpu.threads-per-warp").getInt()); - builder.create(launchOp->getLoc(), "cudaLaunchKernel", TypeRange(), ValueRange({ launchOp.getGridSizeX(), launchOp.getGridSizeY(), launchOp.getGridSizeZ(), launchOp.getBlockSizeX(), launchOp.getBlockSizeY(), launchOp.getBlockSizeZ(), args_dynamic, funcs[funcName], builder.create(launchOp->getLoc(), funcName.size()), launchOp.getDynamicSharedMemorySize(), numWarps, threadsPerWarp})); - launchOp->erase(); - } - - - std::vector gpuDeallocs; - modOp->walk([&gpuDeallocs](mlir::gpu::DeallocOp gpuDeallocOp){ - gpuDeallocs.push_back(gpuDeallocOp); - }); - - for(auto dealloc: gpuDeallocs) - { - builder.setInsertionPoint(dealloc); - builder.create(dealloc->getLoc(), "cudaFree", TypeRange(), ValueRange(dealloc->getOperand(0))); - dealloc.erase(); + return signalPassFailure(); } - - modOp->walk([](mlir::gpu::GPUModuleOp gpuMod) {gpuMod.erase();}); } }; std::unique_ptr> mlir::comet::createLowerGpuHostToCudaPass() { - // std::cout << "Running createLowerGpuHostToCudaPass\n"; - return std::make_unique<::LowerGpuHostToCuda>(); } diff --git a/lib/Conversion/TritonToHIP/CMakeLists.txt b/lib/Conversion/TritonToHIP/CMakeLists.txt new file mode 100644 index 00000000..78f20190 --- /dev/null +++ b/lib/Conversion/TritonToHIP/CMakeLists.txt @@ -0,0 +1,16 @@ +add_llvm_library(COMETTritonToHIP + TritonToHIPPass.cpp + + DEPENDS + TritonDeviceToHIPConversionPassIncGen + + LINK_LIBS PUBLIC + MLIRIR + MLIRPass + MLIRTransforms + TritonIR + TritonGPUIR + TritonGPUTransforms + TritonAMDGPUTransforms + COMETGPUUtils +) diff --git a/lib/Conversion/TritonToHIP/TritonToHIPPass.cpp b/lib/Conversion/TritonToHIP/TritonToHIPPass.cpp new file mode 100644 index 00000000..a80368be --- /dev/null +++ b/lib/Conversion/TritonToHIP/TritonToHIPPass.cpp @@ -0,0 +1,194 @@ + +#include "comet/Conversion/TritonToHIP/TritonToHIPPass.h" +#include "comet/Conversion/GpuUtils/GpuUtils.h" +#include "TritonAMDGPUToLLVM/Passes.h" +#include "TritonAMDGPUTransforms/Passes.h" +#include "comet/Dialect/Utils/Utils.h" +#include "mlir/Conversion/IndexToLLVM/IndexToLLVM.h" +#include "mlir/Conversion/SCFToControlFlow/SCFToControlFlow.h" +#include "mlir/Dialect/Func/IR/FuncOps.h" +#include "mlir/Dialect/GPU/IR/GPUDialect.h" +#include "mlir/Dialect/GPU/Transforms/Passes.h" +#include "mlir/Dialect/LLVMIR/NVVMDialect.h" +#include "mlir/Dialect/LLVMIR/ROCDLDialect.h" +#include "mlir/Dialect/MemRef/IR/MemRef.h" +#include "mlir/IR/Builders.h" +#include "mlir/IR/BuiltinOps.h" +#include "mlir/IR/BuiltinTypes.h" +#include "mlir/IR/Value.h" +#include "mlir/Pass/Pass.h" +#include "mlir/Pass/PassManager.h" +#include "mlir/Target/LLVMIR/Export.h" +#include "mlir/Transforms/DialectConversion.h" +#include "mlir/Transforms/Passes.h" +#include "triton/Conversion/TritonGPUToLLVM/Passes.h" +#include "triton/Conversion/TritonToTritonGPU/TritonToTritonGPUPass.h" +#include "triton/Dialect/Triton/IR/Dialect.h" +#include "triton/Dialect/Triton/Transforms/Passes.h" +#include "triton/Dialect/TritonGPU/Transforms/Passes.h" +#include +#include +#include +#include +#include + +#define GEN_PASS_CLASSES +#include "comet/Conversion/TritonToHIP/Passes.h" + +using namespace mlir; + +namespace hip_device +{ + bool add_ttir_passes(ModuleOp &mod) { + PassManager pm(mod.getContext()); + + pm.addPass(createInlinerPass()); + pm.addPass(triton::createRewriteTensorPointerPass()); + pm.addPass(triton::createCombineOpsPass()); + pm.addPass(createCanonicalizerPass()); + pm.addPass(triton::createReorderBroadcastPass()); + pm.addPass(createCSEPass()); + pm.addPass(createLoopInvariantCodeMotionPass()); + pm.addPass(triton::createLoopUnrollPass()); + if (failed(pm.run(mod))) { + return false; + } + return true; + } + + bool add_ttgir_passes(ModuleOp &mod, int32_t numWarps, int32_t threadsPerWarp, int32_t numStages, int32_t numCTAs, std::string computeCapability) { + { + PassManager pm(mod.getContext()); + pm.addPass(mlir::triton::createConvertTritonToTritonGPUPass( + "hip:" + computeCapability, numWarps, + threadsPerWarp, numCTAs)); + if (failed(pm.run(mod))) { + return false; + } + } + PassManager pm(mod.getContext()); + pm.addPass(triton::gpu::createTritonGPUCoalesce()); + + pm.addPass(triton::gpu::createTritonGPURemoveLayoutConversions()); + pm.addPass(triton::gpu::createTritonGPUOptimizeThreadLocality()); + // TODO: This one takes options + pm.addPass(createTritonAMDGPUAccelerateMatmulPass(computeCapability)); + pm.addPass(triton::gpu::createTritonGPURemoveLayoutConversions()); + pm.addPass(createTritonAMDGPUOptimizeEpiloguePass()); + mlir::triton::gpu::TritonGPUOptimizeDotOperandsOptions options; + options.hoistLayoutConversion = true; + pm.addPass(triton::gpu::createTritonGPUOptimizeDotOperands(options)); + // TODO: Tensor core options here + // addd..... + + pm.addPass(createCanonicalizerPass()); + // pm.addPass(triton::createTritonAMDGPUInsertInstructionSchedHintsPass()); + // TODO: Does not work + pm.addPass(triton::gpu::createTritonGPUOptimizeDotOperands(options)); + pm.addPass(triton::gpu::createTritonGPURemoveLayoutConversions()); + pm.addPass(triton::gpu::createTritonGPUReduceDataDuplication()); + if (numStages != 0) { + pm.addPass(createTritonAMDGPUReorderInstructionsPass()); + } + pm.addPass(createTritonAMDGPUCanonicalizePointersPass()); + pm.addPass(createCanonicalizerPass()); + pm.addPass(createCSEPass()); + pm.addPass(createSymbolDCEPass()); + if (failed(pm.run(mod))) { + return false; + } + + return true; + } + + bool add_llir_passes(ModuleOp &mod, std::string computeCapability) { + PassManager pm(mod.getContext()); + + pm.addPass(triton::AMD::createDecomposeUnsupportedConversionsPass( + computeCapability)); + pm.addPass(createConvertSCFToCFPass()); + pm.addPass(createConvertIndexToLLVMPass()); + pm.addPass(triton::gpu::createAllocateSharedMemoryPass()); + pm.addPass( + triton::createConvertTritonAMDGPUToLLVMPass(computeCapability, true)); + pm.addPass(createCanonicalizerPass()); + pm.addPass(createCSEPass()); + pm.addPass(createSymbolDCEPass()); + pm.addPass(triton::createConvertBuiltinFuncToLLVMPass()); + + if (failed(pm.run(mod))) { + return false; + } + + return true; + } +} + +class LowerTritonDeviceToHIP + : public mlir::comet::LowerTritonDeviceToHIPBase { +public: + + LowerTritonDeviceToHIP() = default; + + LowerTritonDeviceToHIP(int numWarps, int threadsPerWarp, int numCTAs, + int numStages, std::string computeCapability, + mlir::tensorAlgebra::GPUCompilationFormat codeFormat) { + + this->numWarps = numWarps; + this->threadsPerWarp = threadsPerWarp; + this->numCTAs = numCTAs; + this->numStages = numStages; + this->computeCapability = computeCapability; + this->codeFormat = codeFormat; + } + + void runOnOperation() override { + mlir::ModuleOp modOp = getOperation(); + OpBuilder builder(modOp); + auto target = mlir::ROCDL::ROCDLTargetAttr::get(modOp->getContext(), 3, + "amdgcn-amd-amdhsa", computeCapability); + auto add_ttgir_passes = [this](ModuleOp& modOp) { return hip_device::add_ttgir_passes(modOp, numWarps, threadsPerWarp, numStages, numCTAs, computeCapability); }; + auto add_llir_passes = [this](ModuleOp& modOp) { return hip_device::add_llir_passes(modOp, computeCapability); }; + std::vector llvm_func_attrs = {builder.getNamedAttr(mlir::gpu::GPUDialect::getKernelFuncAttrName(), builder.getUnitAttr()), builder.getNamedAttr(mlir::ROCDL::ROCDLDialect::getKernelFuncAttrName(), builder.getUnitAttr()) }; + if(failed(specializeGpuKernel(builder, modOp, this->codeFormat, target, hip_device::add_ttir_passes, add_ttgir_passes, add_llir_passes, llvm_func_attrs))) + { + return signalPassFailure(); + } + } +}; + +std::unique_ptr> +mlir::comet::createLowerTritonDeviceToHIPPass() { + return std::make_unique<::LowerTritonDeviceToHIP>(); +} + +std::unique_ptr> +mlir::comet::createLowerTritonDeviceToHIPPass( + int numWarps, int threadsPerWarp, int numCTAs, int numStages, + std::string computeCapability, + mlir::tensorAlgebra::GPUCompilationFormat format) { + return std::make_unique<::LowerTritonDeviceToHIP>( + numWarps, threadsPerWarp, numCTAs, numStages, computeCapability, format); +} + +class LowerGpuHostToHIP + : public mlir::comet::LowerHostToHIPBase { +public: + LowerGpuHostToHIP() = default; + + void runOnOperation() override { + mlir::ModuleOp modOp = getOperation(); + OpBuilder builder(modOp); + declare_vendor_funcs(builder, modOp, "Hip"); + + if(failed(specializeGpuHost(builder, modOp, std::string("Hip")))) + { + return signalPassFailure(); + } + } +}; + +std::unique_ptr> +mlir::comet::createLowerGpuHostToHIPPass() { + return std::make_unique<::LowerGpuHostToHIP>(); +} diff --git a/lib/Dialect/IndexTree/Analysis/CopiedDomainAnalysis.cpp b/lib/Dialect/IndexTree/Analysis/CopiedDomainAnalysis.cpp new file mode 100644 index 00000000..d7c027a6 --- /dev/null +++ b/lib/Dialect/IndexTree/Analysis/CopiedDomainAnalysis.cpp @@ -0,0 +1,80 @@ +#include "llvm/ADT/SmallSet.h" +#include "llvm/ADT/DenseMap.h" + +#include "comet/Dialect/IndexTree/IR/IndexTreeDialect.h" +#include "comet/Dialect/IndexTree/Analysis/CopiedDomainAnalysis.h" + +using namespace mlir; +using namespace mlir::indexTree; + +CopiedDomainAnalysis::CopiedDomainAnalysis(Operation* op) +{ + op->walk([&](IndexTreeOp tree_op) { + tree_op.getBody()->walk([&](IndexTreeComputeOp compute_op) { + analyzeDomains(compute_op); + }); + }); +} + +void CopiedDomainAnalysis::analyzeDomains(IndexTreeComputeOp compute_op) +{ +llvm::SmallDenseSet output_vars; +llvm::SmallDenseMap candidates; +IndexTreeLHSOperandOp lhs = compute_op.getLhs().getDefiningOp(); +Value tensor = compute_op.getResult(); +for(Value pos : lhs.getPos()){ + auto index_to_tensor = pos.getDefiningOp(); + output_vars.insert(index_to_tensor.getIndex()); + candidates.insert(std::make_pair(index_to_tensor.getIndex(), index_to_tensor.getDim())); +} + +llvm::SmallDenseSet appeared_once; +for(Value rhs : compute_op.getRhs()){ + auto positions = rhs.getDefiningOp().getPos(); + for(auto pos : positions) { + auto index_to_tensor = pos.getDefiningOp(); + Value index = index_to_tensor.getIndex(); + + // Reduction variable appears twice in the same expression, and not in the output + if(!output_vars.contains(index) && appeared_once.contains(index)) { + reductionVars.insert(std::make_pair(compute_op, index)); + } + appeared_once.insert(index); + } +} + +for(Value rhs : compute_op.getRhs()){ + auto positions = rhs.getDefiningOp().getPos(); + uint32_t dim = 0; + for(; dim < positions.size(); dim += 1) { + auto index_to_tensor = positions[dim].getDefiningOp(); + Value index = index_to_tensor.getIndex(); + if(reductionVars.contains(std::make_pair(compute_op, index))) { + candidates.erase(index); + break; + } + } + + // Erase everything that comes after a reduction variable + for(dim += 1; dim < positions.size(); dim += 1) { + auto index_to_tensor = positions[dim].getDefiningOp(); + Value index = index_to_tensor.getIndex(); + candidates.erase(index); + } +} + +for(auto it = candidates.begin(); it != candidates.end(); it++) { + uint32_t dim = it->getSecond(); + copiedDomains.insert(std::make_pair(tensor, dim)); +} +} + +bool CopiedDomainAnalysis::isCopiedDomain(Value tensor, unsigned dim) +{ + return copiedDomains.contains(std::make_pair(tensor, dim)); +} + +bool CopiedDomainAnalysis::isReductionVar(IndexTreeComputeOp op, Value index_var) +{ + return reductionVars.contains(std::make_pair(op, index_var)); +} \ No newline at end of file diff --git a/lib/Dialect/IndexTree/CMakeLists.txt b/lib/Dialect/IndexTree/CMakeLists.txt index 3c0f114c..ec7cf0c0 100644 --- a/lib/Dialect/IndexTree/CMakeLists.txt +++ b/lib/Dialect/IndexTree/CMakeLists.txt @@ -1,24 +1,26 @@ -add_llvm_library(COMETIndexTreeDialect +add_mlir_dialect_library(COMETIndexTreeDialect IR/IndexTreeDialect.cpp IR/IndexTree.cpp - Transforms/IterationDomain.cpp - Transforms/Tensor.cpp - Transforms/UnitExpression.cpp - Transforms/WorkspaceTransforms.cpp - Transforms/Fusion.cpp + Transforms/IterationDomainInference.cpp + Transforms/DomainConcretization.cpp + Transforms/SymbolicCompute.cpp + Transforms/WorkspaceTransforms.cpp + Transforms/KernelFusion.cpp + Transforms/DimensionReductionAfterKernelFusion.cpp + Transforms/MaskDomainPatterns.cpp + + Analysis/CopiedDomainAnalysis.cpp ADDITIONAL_HEADER_DIRS ${COMET_MAIN_INCLUDE_DIR}/comet/Dialect/IndexTree - ) - -add_dependencies( - COMETIndexTreeDialect + DEPENDS COMETIndexTreeOpsIncGen + COMETIndexTreeTypesIncGen COMETIndexTreePassIncGen MLIRSupport - ) - -target_link_libraries(COMETIndexTreeDialect MLIRIR) + LINK_LIBS PUBLIC + MLIRIR +) diff --git a/lib/Dialect/IndexTree/IR/IndexTreeDialect.cpp b/lib/Dialect/IndexTree/IR/IndexTreeDialect.cpp index c8aa4789..c82847ac 100644 --- a/lib/Dialect/IndexTree/IR/IndexTreeDialect.cpp +++ b/lib/Dialect/IndexTree/IR/IndexTreeDialect.cpp @@ -25,13 +25,17 @@ // //===----------------------------------------------------------------------===// #include -#include "comet/Dialect/IndexTree/IR/IndexTreeDialect.h" - #include "mlir/IR/DialectImplementation.h" #include "mlir/Dialect/Arith/IR/Arith.h" +#include "mlir/Dialect/SCF/IR/DeviceMappingInterface.h" +#include "llvm/ADT/TypeSwitch.h" + +#include "comet/Dialect/IndexTree/IR/IndexTreeDialect.h" +#include "comet/Dialect/TensorAlgebra/IR/TADialect.h" using namespace mlir; using namespace mlir::indexTree; +using namespace mlir::tensorAlgebra; //===----------------------------------------------------------------------===// // IndexTreeDialect @@ -39,26 +43,16 @@ using namespace mlir::indexTree; #include "comet/Dialect/IndexTree/IR/IndexTreeDialect.cpp.inc" -Type mlir::indexTree::IndexTreeDialect::parseType(DialectAsmParser &parser) const -{ - /// Parse the main keyword for the type. - StringRef keyword; - /// for "range" and "sptensor" type - if (parser.parseKeyword(&keyword)) - return Type(); +//===----------------------------------------------------------------------===// +// Tablegen Type Definitions +//===----------------------------------------------------------------------===// - parser.emitError(parser.getNameLoc(), - "unknown IndexTree type: " + keyword); - return Type(); -} +#define GET_TYPEDEF_CLASSES +#include "comet/Dialect/IndexTree/IR/IndexTreeTypes.cpp.inc" + +// Include the op interface definitions +#include "comet/Dialect/IndexTree/IR/IndexTreeOpInterfaces.cpp.inc" -/// Print an instance of a type registered to the index tree dialect. -/// No type definition yet -void mlir::indexTree::IndexTreeDialect::printType(mlir::Type type, - mlir::DialectAsmPrinter &printer) const -{ - return; -} #define GET_OP_CLASSES #include "comet/Dialect/IndexTree/IR/IndexTreeOps.cpp.inc" @@ -67,8 +61,16 @@ void mlir::indexTree::IndexTreeDialect::printType(mlir::Type type, /// the point of registration of types and operations for the dialect. void IndexTreeDialect::initialize() { + addTypes< +#define GET_TYPEDEF_LIST +#include "comet/Dialect/IndexTree/IR/IndexTreeTypes.cpp.inc" + >(); + addOperations< #define GET_OP_LIST #include "comet/Dialect/IndexTree/IR/IndexTreeOps.cpp.inc" >(); -} \ No newline at end of file +} + +using namespace mlir; +using namespace mlir::indexTree; diff --git a/lib/Dialect/IndexTree/Transforms/DimensionReductionAfterKernelFusion.cpp b/lib/Dialect/IndexTree/Transforms/DimensionReductionAfterKernelFusion.cpp new file mode 100644 index 00000000..ee618d32 --- /dev/null +++ b/lib/Dialect/IndexTree/Transforms/DimensionReductionAfterKernelFusion.cpp @@ -0,0 +1,816 @@ +// +// Copyright 2022 Battelle Memorial Institute +// +// Redistribution and use in source and binary forms, with or without modification, +// are permitted provided that the following conditions are met: +// +// 1. Redistributions of source code must retain the above copyright notice, this list of conditions +// and the following disclaimer. +// +// 2. Redistributions in binary form must reproduce the above copyright notice, this list of conditions +// and the following disclaimer in the documentation and/or other materials provided with the distribution. +// +// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED +// WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR +// PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, +// INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE +// GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, +// WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +// OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. +// + +#include "mlir/Dialect/Func/IR/FuncOps.h" +#include "mlir/Dialect/Index/IR/IndexOps.h" +#include "mlir/Dialect/MemRef/IR/MemRef.h" +#include "mlir/Dialect/SCF/IR/SCF.h" +#include "mlir/Transforms/GreedyPatternRewriteDriver.h" + +#include "llvm/ADT/StringSet.h" + +#include "comet/Dialect/IndexTree/IR/IndexTreeDialect.h" +#include "comet/Dialect/TensorAlgebra/IR/TADialect.h" +#include "comet/Dialect/IndexTree/Passes.h" + +#include + +using namespace mlir; +using namespace mlir::indexTree; + +// *********** For debug purpose *********// +//#define COMET_DEBUG_MODE +#include "comet/Utils/debug.h" +// *********** For debug purpose *********// + +namespace mlir { + namespace comet{ +#define GEN_PASS_DEF_INDEXTREEDIMENSIONREDUCTION +#include "comet/Dialect/IndexTree/Passes.h.inc" + } +} + +namespace { + +struct ITreeInfo { + IndexTreeOp itree_op; + llvm::SmallVector args_inputs; + llvm::SmallVector args_intermediates; + llvm::SmallVector block_args; + llvm::SmallVector args_types; +}; + +template +uint32_t getIdxInVector(const llvm::SmallVector &array, T value) +{ + auto it = std::find(array.begin(), array.end(), value); + if (it != array.end()) { + return it - array.begin(); + } else { + return (uint32_t) -1; + } +} + +void dfs_it(IndexTreeIndicesOp node, + llvm::SmallVector &path, + llvm::DenseMap> &all_paths) +{ + path.push_back(node); + for (auto child : node->getUsers()) { + if (auto computeOp = llvm::dyn_cast(child)) { + all_paths[computeOp] = path; + } else if (auto indexOp = llvm::dyn_cast(child)) { + dfs_it(indexOp, path, all_paths); + } + } + path.pop_back(); +} + +llvm::DenseMap> getAllPathsToComputeOp(func::FuncOp funcOp) +{ + IndexTreeRootOp rootOp; + funcOp.walk([&](IndexTreeRootOp op) { + rootOp = op; + }); + + uint32_t count_child = 0; + IndexTreeIndicesOp child; + for (auto user : rootOp->getUsers()) { + ++count_child; + child = llvm::cast(user); + } + assert(count_child == 1 && "No common IndexOp nodes are shared by ComputeOp nodes."); + llvm::DenseMap> all_paths; + llvm::SmallVector curr_path; + dfs_it(child, curr_path, all_paths); +// {/// test +// for (const auto &[computeOp, path] : all_paths) { +// comet_debug() << "a path for \n"; +// comet_vdump(computeOp); +// for (const auto index : path) { +// comet_vdump(index); +// } +// } +// } + return all_paths; +} + +llvm::SmallVector getCommonIndices(Value computeOp1, + Value computeOp2, + const llvm::DenseMap> &all_paths) +{ + const llvm::SmallVector &path1 = all_paths.at(computeOp1); + const llvm::SmallVector &path2 = all_paths.at(computeOp2); + llvm::SmallVector common_indices; + uint32_t num = 0; + while (num < path1.size() && num < path2.size()) { + if (path1[num] == path2[num]) { + common_indices.push_back(path1[num]); + ++num; + } else { + break; + } + } + + return common_indices; +} + +std::unordered_map> +getAllComputeOpsCommonIndices(func::FuncOp funcOp, + const llvm::SmallVector &computeOps) +{ + llvm::DenseMap> computeOp_to_index_path = getAllPathsToComputeOp(funcOp); + std::unordered_map> computeOp_to_common_indices; + + for (uint32_t computeOp_i = 0; computeOp_i < computeOps.size() - 1; ++computeOp_i) { + computeOp_to_common_indices[computeOp_i + 1] = getCommonIndices(computeOps[computeOp_i], + computeOps[computeOp_i + 1], + computeOp_to_index_path); + } +// {/// test +// for (const auto [computeOp_i, indices] : computeOp_to_common_indices) { +// comet_debug() << "computeOp_i: " << computeOp_i << " num_common_indices: " << indices.size() << "\n"; +// } +// } + + return computeOp_to_common_indices; +} + +Value getOldLhsTensor(Value computeOp, + ITreeInfo &itreeInfo) +{ + IndexTreeComputeOp cmptOp = llvm::cast(computeOp.getDefiningOp()); + Value lhs_tensor = llvm::cast(cmptOp.getLhs().getDefiningOp()).getTensor(); + + uint32_t idx = getIdxInVector(itreeInfo.block_args, lhs_tensor); + if (idx != (uint32_t) -1) { + if (idx < itreeInfo.args_inputs.size()) { + lhs_tensor = itreeInfo.args_inputs[idx]; + } else { + lhs_tensor = itreeInfo.args_intermediates[idx - itreeInfo.args_inputs.size()]; + } + } else { + assert(false && "Expect to find the lhs tensor in itree's block arguments."); + } + + return lhs_tensor; +} + + +std::unordered_map createNewLhsTensors( + llvm::SmallVector &computeOps, + ITreeInfo &itreeInfo, + std::unordered_map> &computeOp_to_common_indices, + mlir::IRRewriter &rewriter, + mlir::Location &loc) +{ + std::unordered_map computeOp_to_new_tensors; + for (uint32_t computeOp_i = 0; computeOp_i < computeOps.size() - 1; ++computeOp_i) { + Value computeOp = computeOps[computeOp_i]; + Value old_tensor = getOldLhsTensor(computeOp, itreeInfo); + uint32_t num_common_indices = computeOp_to_common_indices[computeOp_i + 1].size(); + assert(llvm::isa(old_tensor.getType()) && + "Expect a mlir::TensorType."); + assert(num_common_indices && "Error: number of common indices is zero"); + auto old_tensor_tt = llvm::cast(old_tensor.getType()); + + /// The start dim idx after those common indices + uint32_t dim_base = num_common_indices; +// dim_base = 1; /// test + if (dim_base < old_tensor_tt.getRank()) { + /// The new tensor is still a tensor, with decreased dimensions. + + uint32_t count_of_dyn_dim = 0; /// How many dynamic dimension is in common indices + for (uint32_t dim_i = 0; dim_i < dim_base; ++dim_i) { + if (old_tensor_tt.isDynamicDim(dim_i)) { + ++count_of_dyn_dim; + } + } + SmallVector operands; /// The remaining dynamic dimensions after fusion. + if (old_tensor.getDefiningOp()->getNumOperands() > count_of_dyn_dim) { + operands.insert(operands.begin(), + old_tensor.getDefiningOp()->getOperands().begin() + count_of_dyn_dim, + old_tensor.getDefiningOp()->getOperands().end()); + } + + llvm::SmallVector shape; + for (int64_t dim_i = dim_base; dim_i < old_tensor_tt.getRank(); ++dim_i) { + if (old_tensor_tt.isDynamicDim(dim_i)) { + shape.push_back(mlir::ShapedType::kDynamic); + } else { + shape.push_back(old_tensor_tt.getDimSize(dim_i)); + } + } + + OpBuilder::InsertionGuard guard(rewriter); + rewriter.setInsertionPoint(old_tensor.getDefiningOp()); + Value new_tensor = rewriter.create( + loc, + mlir::RankedTensorType::get(shape, old_tensor_tt.getElementType()), + operands); + computeOp_to_new_tensors[computeOp_i] = new_tensor; + mlir::TypedAttr zero = rewriter.getZeroAttr(old_tensor_tt.getElementType()); + rewriter.create(loc, + new_tensor, + zero); + comet_vdump(new_tensor); + } else { + /// The new tensor is shrunk to a scalar. + OpBuilder::InsertionGuard guard(rewriter); + rewriter.setInsertionPoint(old_tensor.getDefiningOp()); + mlir::TypedAttr zero = rewriter.getZeroAttr(old_tensor_tt.getElementType()); + Value new_scalar = rewriter.create(loc, + old_tensor_tt.getElementType(), + zero); + comet_vdump(new_scalar); + computeOp_to_new_tensors[computeOp_i] = new_scalar; + } + } + + return computeOp_to_new_tensors; +} + +ITreeInfo createNewItreeOp(uint32_t num_computeOps, + ITreeInfo &oldITreeInfo, + const std::unordered_map &computeOp_to_new_tensors, + mlir::IRRewriter &rewriter, + mlir::Location &loc) +{ + IndexTreeOp old_itree = oldITreeInfo.itree_op; + OpBuilder::InsertionGuard guard(rewriter); + rewriter.setInsertionPoint(old_itree); + llvm::SmallVector args_inputs = oldITreeInfo.args_inputs; + llvm::SmallVector args_intermediates; + for (uint32_t computeOp_i = 0; computeOp_i < num_computeOps - 1; ++computeOp_i) { + args_intermediates.push_back(computeOp_to_new_tensors.at(computeOp_i)); + } + + llvm::SmallVector locs; + llvm::SmallVector args_types; + for (Value arg : args_inputs) { + locs.push_back(arg.getLoc()); + args_types.push_back(arg.getType()); + } + for (Value arg : args_intermediates) { + locs.push_back(arg.getLoc()); + args_types.push_back(arg.getType()); + } + IndexTreeOp new_itree = rewriter.create(loc, + args_types, + args_inputs, + args_intermediates); + Region *body = &new_itree.getRegion(); + Block *block = rewriter.createBlock(body, {}, TypeRange(args_types), locs); + YieldOp dummy_yield = rewriter.create(loc, TypeRange(), args_inputs); /// create a dummy YieldOp to make the following inlineBlockBefore() work. + /// Move the old itree's block to the new itree. + rewriter.inlineBlockBefore(&old_itree.getRegion().front(), + block->getTerminator(), + block->getArguments()); + rewriter.eraseOp(dummy_yield); + assert(old_itree.getNumResults() == new_itree.getNumResults() && + "Expect old itree and new itree have the same number of inputs and outputs."); + for (uint32_t r_i = 0; r_i < old_itree.getNumResults(); ++r_i) { + rewriter.replaceAllUsesWith(old_itree.getResult(r_i), new_itree.getResult(r_i)); + } + rewriter.eraseOp(old_itree); + comet_vdump(new_itree); + comet_vdump(new_itree->getParentOfType()); + + ITreeInfo newITreeInfo; + newITreeInfo.itree_op = new_itree; + newITreeInfo.args_inputs = args_inputs; + newITreeInfo.args_intermediates = args_intermediates; + for (Value arg : block->getArguments()) { + newITreeInfo.block_args.push_back(arg); + } + newITreeInfo.args_types = args_types; + + return newITreeInfo; +} + +/// Create the new LHS operand and remove the old one. +Value createNewLHSOperandOp(Value computeOp, + uint32_t num_common_indices, + Value lhs_tensor, + mlir::IRRewriter &rewriter, + mlir::Location &loc) +{ + /// Find the old LHSOperandOp + IndexTreeLHSOperandOp old_lhs_operand_op = llvm::cast( + llvm::cast(computeOp.getDefiningOp()).getLhs().getDefiningOp()); + OpBuilder::InsertionGuard guard(rewriter); + rewriter.setInsertionPoint(old_lhs_operand_op); + + /// Create the new LHSOperandOp + llvm::SmallVector pos; + llvm::SmallVector crds; + Value prev_dim = nullptr; + auto access_type = rewriter.getIndexType(); + uint32_t dim_i = 0; + for (Value posOp : old_lhs_operand_op.getPos()) { + if (dim_i < num_common_indices) { + ++dim_i; + continue; + } + IndexTreeIndexToTensorOp tensorDimOp = llvm::cast(posOp.getDefiningOp()); + Value index_node = tensorDimOp.getIndex(); + uint32_t dim = tensorDimOp.getDim() - num_common_indices; + auto access_op = rewriter.create( + loc, + TypeRange({access_type, access_type}), + lhs_tensor, + index_node, + rewriter.getUI32IntegerAttr(dim), + prev_dim); + pos.push_back(access_op.getPos()); + crds.push_back(access_op.getCrd()); + prev_dim = pos.back(); + } + + indexTree::OperandType operand_type = indexTree::OperandType::get(rewriter.getContext()); + Value new_lhs_operand = rewriter.create(loc, + operand_type, + lhs_tensor, + pos, + crds); + comet_vdump(new_lhs_operand); + comet_vdump(new_lhs_operand.getDefiningOp()->getParentOfType()); + + /// Replace and erase + rewriter.replaceAllUsesWith(old_lhs_operand_op, new_lhs_operand); + llvm::SmallVector old_index_to_tensor_dim; + for (auto tensor_dim : old_lhs_operand_op.getPos()) { + old_index_to_tensor_dim.push_back(tensor_dim); + } + std::reverse(old_index_to_tensor_dim.begin(), old_index_to_tensor_dim.end()); + rewriter.eraseOp(old_lhs_operand_op); + for (auto tensor_dim : old_index_to_tensor_dim) { + rewriter.eraseOp(tensor_dim.getDefiningOp()); + } + + comet_vdump(new_lhs_operand.getDefiningOp()->getParentOfType()); + return new_lhs_operand; +} + +Value createNewComputeOp(Value computeOp, + const mlir::Type &tensor_type, + Value lhsOperandOp, + ValueRange rhsOperandOps, + mlir::IRRewriter &rewriter, + mlir::Location &loc) +{ + IndexTreeComputeOp old_compute_op = llvm::cast(computeOp.getDefiningOp()); + OpBuilder::InsertionGuard guard(rewriter); + rewriter.setInsertionPoint(computeOp.getDefiningOp()); + Value new_compute_op = rewriter.create( + loc, + tensor_type, + old_compute_op.getParent(), + lhsOperandOp, + rhsOperandOps, + old_compute_op.getMask(), + old_compute_op.getSemiringAttr(), + old_compute_op.getComputeMissingAttr()); + +// rewriter.replaceAllUsesWith(old_compute_op, new_compute_op); +// rewriter.eraseOp(old_compute_op); + + comet_vdump(new_compute_op); + comet_vdump(new_compute_op.getDefiningOp()->getParentOfType()); + return new_compute_op; +} + +/// Create the new rhs oprand for the new intermeidate tensor, then return it along with the other old rhs operand. +llvm::SmallVector createNewRHSOperandOps(Value prev_old_compute_op, + Value curr_old_compute_op, + uint32_t num_common_indices, + Value intermediate_tensor, + IndexTreeOperandOp &new_rhs_operand_op /*out*/, + mlir::IRRewriter &rewriter, + mlir::Location &loc) +{ + /// Find the intermediate idx + llvm::SmallVector old_rhs_operand_ops = + llvm::cast(curr_old_compute_op.getDefiningOp()).getRhs(); + uint32_t intermediate_idx = 0; + IndexTreeOperandOp old_rhs_operand_op = nullptr; + while (intermediate_idx < old_rhs_operand_ops.size()) { + IndexTreeOperandOp operandOp = + llvm::cast(old_rhs_operand_ops[intermediate_idx].getDefiningOp()); + if (operandOp.getTensor() == prev_old_compute_op) { + old_rhs_operand_op = operandOp; + break; /// Found the intermediate tensor + } + ++intermediate_idx; + } + assert(intermediate_idx < old_rhs_operand_ops.size() && "Error: not found the intermediate tensor."); + + /// Create the new rhs operand + OpBuilder::InsertionGuard guard(rewriter); + rewriter.setInsertionPoint(old_rhs_operand_op); + llvm::SmallVector pos; + llvm::SmallVector crds; + Value prev_dim = nullptr; + auto access_type = rewriter.getIndexType(); + uint32_t dim_i = 0; + for (Value posOp : old_rhs_operand_op.getPos()) { + if (dim_i < num_common_indices) { + ++dim_i; + continue; + } + IndexTreeIndexToTensorOp tensorDimOp = llvm::cast(posOp.getDefiningOp()); + Value index_node = tensorDimOp.getIndex(); + uint32_t dim = tensorDimOp.getDim() - num_common_indices; + auto access_op = rewriter.create( + loc, + TypeRange({access_type, access_type}), + /*lhs_tensor=*/intermediate_tensor, + index_node, + rewriter.getUI32IntegerAttr(dim), + prev_dim); + pos.push_back(access_op.getPos()); + crds.push_back(access_op.getCrd()); + prev_dim = pos.back(); + } + indexTree::OperandType operand_type = indexTree::OperandType::get(rewriter.getContext()); + new_rhs_operand_op = rewriter.create(loc, + operand_type, + intermediate_tensor, + pos, + crds); + comet_vdump(new_rhs_operand_op); + comet_vdump(new_rhs_operand_op->getParentOfType()); + + /// Replace and erase + rewriter.replaceAllUsesWith(old_rhs_operand_op, new_rhs_operand_op); + llvm::SmallVector old_index_to_tensor_dim; + for (auto tensor_dim : old_rhs_operand_op.getPos()) { + old_index_to_tensor_dim.push_back(tensor_dim); + } + std::reverse(old_index_to_tensor_dim.begin(), old_index_to_tensor_dim.end()); + rewriter.eraseOp(old_rhs_operand_op); + for (auto tensor_dim : old_index_to_tensor_dim) { + rewriter.eraseOp(tensor_dim.getDefiningOp()); + } + + /// Collect all rhs operands + llvm::SmallVector rhs_operand_ops; + for (uint32_t r_i = 0; r_i < old_rhs_operand_ops.size(); ++r_i) { + if (r_i == intermediate_idx) { + rhs_operand_ops.push_back(new_rhs_operand_op); + } else { + rhs_operand_ops.push_back(old_rhs_operand_ops[r_i]); + } + } + + return rhs_operand_ops; +} + + +llvm::SmallVector createExtraIndicesOpsForReset(Value prev_old_compute_op, + Value parent_index_op, + mlir::IRRewriter &rewriter, + mlir::Location &loc) +{ + /// Get the domain from previous ComputeOP + /// Domain: previous ComputeOp -> LHSOperandOp -> IndexToTensorDim -> IndexOp -> DenseDomainOp + auto prev_compute_op = llvm::cast(prev_old_compute_op.getDefiningOp()); + auto lhs_operand_op = llvm::cast(prev_compute_op.getLhs().getDefiningOp()); + llvm::SmallVector index_to_tensor_dim_ops; + for (Value pos : lhs_operand_op.getPos()) { + index_to_tensor_dim_ops.push_back(llvm::cast(pos.getDefiningOp())); + } + llvm::SmallVector indices_ops; + for (auto index_to_tensor_dim : index_to_tensor_dim_ops) { + indices_ops.push_back(llvm::cast(index_to_tensor_dim.getIndex().getDefiningOp())); + } + llvm::SmallVector domain_ops; + for (auto index : indices_ops) { + domain_ops.push_back(index.getDomain()); + } + + llvm::SmallVector new_indices_ops; + Value parent = parent_index_op; + for (Value domain : domain_ops) { + Value new_index = rewriter.create(loc, + indexTree::IndexNodeType::get(rewriter.getContext()), + parent, + domain, + false, + nullptr); + new_indices_ops.push_back(new_index); + parent = new_index; + comet_vdump(new_index); + } + + return new_indices_ops; +} + +Value createComputeOpForReset(const llvm::SmallVector &common_indices, + Value prev_old_comput_op, + Value intermediate_tensor, + IndexTreeOperandOp new_rhs_operand_op, + mlir::Type element_type, + mlir::IRRewriter &rewriter, + mlir::Location &loc) +{ + uint32_t num_dims_new_tensor = 0; + for ([[maybe_unused]]Value _ : new_rhs_operand_op.getPos()) { + ++num_dims_new_tensor; + } + indexTree::OperandType operand_type = indexTree::OperandType::get(rewriter.getContext()); + Value parent; + Value lhs_operand_op; + if (num_dims_new_tensor) + { + /// If the intermediate tensor is still a tensor, generate extra indices (loops) to reset it. + llvm::SmallVector new_indices_ops = createExtraIndicesOpsForReset(prev_old_comput_op, + /*parent_index_op=*/common_indices.back(), + rewriter, + loc); + assert(new_indices_ops.size() == num_dims_new_tensor && "Expect to reset the whole new tensor"); + parent = new_indices_ops.back(); + + /// Create the new lhs operand to reset + llvm::SmallVector pos; + llvm::SmallVector crds; + Value prev_dim = nullptr; + auto access_type = rewriter.getIndexType(); + uint32_t index_i = 0; + for (Value posOp : new_rhs_operand_op.getPos()) { + IndexTreeIndexToTensorOp tensorDimOp = llvm::cast(posOp.getDefiningOp()); +// Value index_node = tensorDimOp.getIndex(); + uint32_t dim = tensorDimOp.getDim(); + Value index_node = new_indices_ops[index_i++]; + auto access_op = rewriter.create( + loc, + TypeRange({access_type, access_type}), + /*lhs_tensor=*/intermediate_tensor, + index_node, + rewriter.getUI32IntegerAttr(dim), + prev_dim); + pos.push_back(access_op.getPos()); + crds.push_back(access_op.getCrd()); + prev_dim = pos.back(); + } + + lhs_operand_op = rewriter.create( + loc, + operand_type, + intermediate_tensor, + pos, + crds); + } + else + { + /// If the intermediate tensor is a scalar, link to the last common index + parent = common_indices.back(); + // auto access_type = rewriter.getIndexType(); + lhs_operand_op = rewriter.create( + loc, + operand_type, + intermediate_tensor); + } + + /// Create the rhs operand (i.e., constant 0) + mlir::TypedAttr zero = rewriter.getZeroAttr(element_type); + Value cst_0 = rewriter.create(loc, + element_type, + zero); + Value rhs_operand_op = rewriter.create( + loc, + operand_type, + /*rhs_tensor*/cst_0, + /*pos*/ValueRange{}, + /*crds*/ValueRange{}); + + /// Create the Compute Op for resetting + mlir::StringRef semiring("noop_times"); + bool compute_missing = false; + Value compute_op = rewriter.create( + loc, + intermediate_tensor.getType(), + parent, + lhs_operand_op, + ValueRange{rhs_operand_op}, + /*mask_operand*/nullptr, + rewriter.getStringAttr(semiring), + rewriter.getBoolAttr(compute_missing)); + + return compute_op; +} + +} /// anonymous namespace + +struct IndexTreeDimensionReduction : comet::impl::IndexTreeDimensionReductionBase { + using IndexTreeDimensionReductionBase::IndexTreeDimensionReductionBase; + void runOnOperation() override; +}; + +void IndexTreeDimensionReduction::runOnOperation() +{ + comet_debug() << "IndexTreeDimensionReduction::runOnOperation()\n"; + func::FuncOp funcOp = getOperation(); + comet_vdump(funcOp->getParentOfType()); + + /// Get all itree arguments + ITreeInfo oldITreeInfo; + uint32_t count_itrees = 0; + funcOp.walk([&](IndexTreeOp itreeOp) { + ++count_itrees; + oldITreeInfo.itree_op = itreeOp; + for (Value arg : itreeOp.getInputs()) { + oldITreeInfo.args_inputs.push_back(arg); + } + for (Value arg : itreeOp.getIntermediates()) { + oldITreeInfo.args_intermediates.push_back(arg); + } + for (Value arg : itreeOp.getRegion().getBlocks().front().getArguments()) { + oldITreeInfo.block_args.push_back(arg); + } + }); + assert(count_itrees == 1 && "Expected one single fused itree."); + + /// Get all ComputeOp nodes + llvm::SmallVector computeOps; + funcOp.walk([&](IndexTreeComputeOp computeOp) { + computeOps.push_back(computeOp); + comet_vdump(computeOp); + }); + uint32_t num_computeOps = computeOps.size(); + + /// Get number of common indices for each two consecutive ComputeOp nodes. + std::unordered_map> computeOp_to_common_indices = + getAllComputeOpsCommonIndices(funcOp, computeOps); + assert(computeOp_to_common_indices.size() == num_computeOps - 1 && "N ComputeOps expect N-1 intermediates."); + + /// Create new tensors with dimension reduction + mlir::OpBuilder builder(oldITreeInfo.itree_op); + mlir::IRRewriter rewriter(builder); + Location loc = oldITreeInfo.itree_op.getLoc(); + std::unordered_map computeOp_to_new_tensors = createNewLhsTensors(computeOps, + oldITreeInfo, + computeOp_to_common_indices, + rewriter, + loc); + + /// Create a new itree to update the argument types, moving the body from old itree to the new itree + ITreeInfo newITreeInfo = createNewItreeOp(num_computeOps, + oldITreeInfo, + computeOp_to_new_tensors, + rewriter, + loc); + /// Dimension reduction + llvm::SmallVector new_compute_ops; + llvm::SmallVector new_compute_ops_for_reset; + uint32_t intermediates_idx_base = newITreeInfo.args_inputs.size(); + for (uint32_t computeOp_i = 0; computeOp_i < num_computeOps; ++computeOp_i) { + Value computeOp = computeOps[computeOp_i]; + if (computeOp_i == 0) { + /// The first ComputeOp + /// Only LHS needs dimension reduction + Value lhs_operand_op = createNewLHSOperandOp( + computeOp, + /*num_common_indices=*/computeOp_to_common_indices[computeOp_i + 1].size(), + /*lhs_tensor=*/newITreeInfo.block_args[intermediates_idx_base + computeOp_i], + rewriter, + loc); + /// Create new ComputeOp + auto rhs_operand_ops = + llvm::cast(computeOp.getDefiningOp()).getRhs(); + Value new_compute_op = createNewComputeOp( + computeOp, + /*output_tensor_type=*/newITreeInfo.args_types[intermediates_idx_base + computeOp_i], + lhs_operand_op, + rhs_operand_ops, + rewriter, + loc); + new_compute_ops.push_back(new_compute_op); + } else if (computeOp_i == num_computeOps - 1) { + /// The last ComputeOp + /// Only RHS needs dimension reduction + uint32_t num_common_indices = computeOp_to_common_indices[computeOp_i].size(); + Value prev_new_compute_op = new_compute_ops.back(); + IndexTreeOperandOp new_rhs_operand_op; + llvm::SmallVector rhs_operand_ops = createNewRHSOperandOps( + /*prev_old_compute_op=*/computeOps[computeOp_i - 1], + /*curr_old_compute_op=*/computeOps[computeOp_i], + num_common_indices, + /*intermediate_tensor=*/prev_new_compute_op, + new_rhs_operand_op /*out*/, + rewriter, + loc); + /// Create new ComputeOp + auto lhs_operand_op = + llvm::cast(computeOp.getDefiningOp()).getLhs(); + Value new_compute_op = createNewComputeOp( + computeOp, + /*output_tensor_type=*/newITreeInfo.args_types[0]/*the input type*/, + lhs_operand_op, + rhs_operand_ops, + rewriter, + loc); + new_compute_ops.push_back(new_compute_op); + /// Create ComputeOp for resetting + OpBuilder::InsertionGuard guard(rewriter); + rewriter.setInsertionPointAfter(new_compute_op.getDefiningOp()); + mlir::Type element_type = llvm::cast(computeOps[computeOp_i - 1].getType()).getElementType(); + Value compute_op_for_reset = createComputeOpForReset( + /*common_indices=*/computeOp_to_common_indices[computeOp_i], + /*prev_old_comput_op=*/computeOps[computeOp_i - 1], + prev_new_compute_op, + new_rhs_operand_op, + element_type, + rewriter, + loc); + new_compute_ops_for_reset.push_back(compute_op_for_reset); + } else { + /// The middle ComputeOp + /// TODO: need a case to test this middle Computeop + /// LHS dimension reduction + Value lhs_operand_op = createNewLHSOperandOp( + computeOp, + /*num_common_indices=*/computeOp_to_common_indices[computeOp_i + 1].size(), + /*lhs_tensor=*/newITreeInfo.block_args[intermediates_idx_base + computeOp_i], + rewriter, + loc); + /// RHS dimension reduction + uint32_t num_common_indices = computeOp_to_common_indices[computeOp_i].size(); + Value prev_new_compute_op = new_compute_ops.back(); + IndexTreeOperandOp new_rhs_operand_op; + llvm::SmallVector rhs_operand_ops = createNewRHSOperandOps( + /*prev_old_compute_op=*/computeOps[computeOp_i - 1], + /*curr_old_compute_op=*/computeOps[computeOp_i], + num_common_indices, + /*intermediate_tensor=*/prev_new_compute_op, + new_rhs_operand_op /*out*/, + rewriter, + loc); + /// Create new ComputeOp + Value new_compute_op = createNewComputeOp( + computeOp, + /*output_tensor_type=*/newITreeInfo.args_types[intermediates_idx_base + computeOp_i], + lhs_operand_op, + rhs_operand_ops, + rewriter, + loc); + new_compute_ops.push_back(new_compute_op); + /// Create ComputeOp for resetting + OpBuilder::InsertionGuard guard(rewriter); + rewriter.setInsertionPointAfter(new_compute_op.getDefiningOp()); + mlir::Type element_type = llvm::cast(computeOps[computeOp_i - 1].getType()).getElementType(); + Value compute_op_for_reset = createComputeOpForReset( + /*common_indices=*/computeOp_to_common_indices[computeOp_i], + /*prev_old_comput_op=*/computeOps[computeOp_i - 1], + prev_new_compute_op, + new_rhs_operand_op, + element_type, + rewriter, + loc); + new_compute_ops_for_reset.push_back(compute_op_for_reset); + } + } + + /// Update the YieldOp. + Operation *old_yield_op = newITreeInfo.itree_op.getRegion().getBlocks().front().getTerminator(); + OpBuilder::InsertionGuard guard(rewriter); + rewriter.setInsertionPoint(old_yield_op); + llvm::SmallVector yield_op_args; + yield_op_args.push_back(new_compute_ops.back()); + yield_op_args.insert(yield_op_args.end(), new_compute_ops_for_reset.begin(), new_compute_ops_for_reset.end()); + rewriter.create(loc, TypeRange(), yield_op_args); + rewriter.eraseOp(old_yield_op); + + /// Erase old Compute Ops. + for (uint32_t computeOp_i = 0; computeOp_i < num_computeOps; ++computeOp_i) { + rewriter.replaceAllUsesWith(computeOps[computeOp_i], new_compute_ops[computeOp_i]); + rewriter.eraseOp(computeOps[computeOp_i].getDefiningOp()); + } + + comet_vdump(funcOp->getParentOfType()); + +} + +/// Apply the redundancy-aware kernel fusion on index tree dialect for some compound expressions +std::unique_ptr mlir::comet::createIndexTreeDimensionReductionPass() +{ + return std::make_unique(); +} \ No newline at end of file diff --git a/lib/Dialect/IndexTree/Transforms/DomainConcretization.cpp b/lib/Dialect/IndexTree/Transforms/DomainConcretization.cpp new file mode 100644 index 00000000..57c2482e --- /dev/null +++ b/lib/Dialect/IndexTree/Transforms/DomainConcretization.cpp @@ -0,0 +1,639 @@ +// +// Copyright 2022 Battelle Memorial Institute +// +// Redistribution and use in source and binary forms, with or without modification, +// are permitted provided that the following conditions are met: +// +// 1. Redistributions of source code must retain the above copyright notice, this list of conditions +// and the following disclaimer. +// +// 2. Redistributions in binary form must reproduce the above copyright notice, this list of conditions +// and the following disclaimer in the documentation and/or other materials provided with the distribution. +// +// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED +// WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR +// PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, +// INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE +// GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, +// WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +// OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. +// + +#include "mlir/Dialect/Func/IR/FuncOps.h" +#include "mlir/IR/BuiltinOps.h" +#include "mlir/IR/BuiltinTypeInterfaces.h" +#include "mlir/IR/Value.h" +#include "mlir/Support/LLVM.h" +#include "mlir/Transforms/GreedyPatternRewriteDriver.h" +#include "mlir/Dialect/MemRef/IR/MemRef.h" +#include "mlir/Dialect/Math/IR/Math.h" +#include "mlir/Dialect/Tensor/IR/Tensor.h" +#include "mlir/Dialect/Index/IR/IndexOps.h" +#include "mlir/Pass/Pass.h" + +#include "llvm/ADT/StringSet.h" +#include "llvm/ADT/SmallVector.h" +#include "llvm/ADT/IndexedMap.h" +#include "llvm/ADT/TypeSwitch.h" + +#include "comet/Dialect/IndexTree/IR/IndexTreeDialect.h" +#include "comet/Dialect/TensorAlgebra/IR/TADialect.h" +#include "comet/Dialect/IndexTree/Passes.h" +#include "comet/Dialect/IndexTree/Patterns.h" + +using namespace mlir; +using namespace mlir::indexTree; +using namespace mlir::tensorAlgebra; + +namespace mlir { + namespace comet{ + #define GEN_PASS_DEF_INDEXTREEDOMAINCONCRETIZATION + #include "comet/Dialect/IndexTree/Passes.h.inc" + } +} + +struct ConcretizeTensorDomain : public OpRewritePattern { + ConcretizeTensorDomain(MLIRContext *context) + : OpRewritePattern(context, /*benefit=*/3) {} + + mlir::LogicalResult + liftAccessOp(mlir::Operation *dependent_op, + IndexTreeIndexToTensorOp access_op) const { + if(access_op->isBeforeInBlock(dependent_op)) + return success(); + + Value prev_access_value; + if((prev_access_value = access_op.getPrevDim())) + { + if(mlir::failed( + liftAccessOp( + dependent_op, + llvm::cast(prev_access_value.getDefiningOp()) + ) + )) + return failure(); + } + access_op->moveBefore(dependent_op); + return success(); + } + + mlir::LogicalResult + matchAndRewrite(IndexTreeTensorDomainOp domain_op, + mlir::PatternRewriter &rewriter) const override { + auto loc = domain_op->getLoc(); + auto context = rewriter.getContext(); + indexTree::DomainType domain_type = indexTree::DomainType::get(context); + uint32_t dim = domain_op.getDim(); + Value new_domain; + Value tensor = domain_op.getTensor(); + if(SparseTensorType sp_tensor = mlir::dyn_cast(tensor.getType())) + { + mlir::RewriterBase::InsertPoint prev = rewriter.saveInsertionPoint(); + if(tensor.getDefiningOp()) + { + rewriter.setInsertionPointAfter(tensor.getDefiningOp()); + } + else { + rewriter.setInsertionPointToStart(tensor.getParentBlock()); + } + + //Domain comes from a sparse tensor (may still be dense) + ArrayRef tensor_dim_formats = sp_tensor.getFormat(); + TensorFormatEnum format = static_cast(tensor_dim_formats[2*dim]); + + if(format == TensorFormatEnum::D) + { + auto index_type = rewriter.getIndexType(); + Value max = rewriter.create(loc, index_type, tensor, rewriter.getI32IntegerAttr(dim)); + rewriter.restoreInsertionPoint(prev); + new_domain = rewriter.create(loc, domain_type, max, tensor, rewriter.getI32ArrayAttr({static_cast(dim)})); + } + else + { + Value pos = rewriter.create(loc, tensor, rewriter.getI32IntegerAttr(dim)); + Value crd = rewriter.create(loc, tensor, rewriter.getI32IntegerAttr(dim)); + Value pos_size = rewriter.create(loc, pos, 0); + Value crd_size = rewriter.create(loc, crd, 0); + Value dim_size = rewriter.create(loc, rewriter.getIndexType(), tensor, rewriter.getI32IntegerAttr(dim)); + Value parent = domain_op.getParent(); + if(!parent) + { + // Get associated index + IndexTreeIndicesOp index_op; + Operation* use = *(domain_op->user_begin()); + // TODO: Fix danger of infinite loop!!! + while(!(index_op = llvm::dyn_cast(use))) + { + use = *(use->user_begin()); + } + assert(index_op); + + if(dim == 0) + { + parent = nullptr; + } else + { + // Infer parent index variable + for(Operation* use : index_op->getUsers()) + { + IndexTreeIndexToTensorOp access_op = llvm::dyn_cast(use); + if(!access_op || access_op.getTensor() != tensor || access_op.getDim() != dim) + continue; + + parent = access_op.getPrevDim(); + IndexTreeIndexToTensorOp prev_access_op = + llvm::cast(parent.getDefiningOp()); + if(mlir::failed(this->liftAccessOp(domain_op, prev_access_op))) + return failure(); + + break; + } + } + } + rewriter.restoreInsertionPoint(prev); + new_domain = rewriter.create( + loc, domain_type, tensor, domain_op.getDimAttr(), + TensorFormatEnumAttr::get(context, format), + pos, crd, pos_size, crd_size, dim_size, parent); + } + } + else if(llvm::isa(tensor.getType())) { + auto index_type = rewriter.getIndexType(); + Value dim_size = rewriter.create(loc, index_type, tensor, rewriter.getI32IntegerAttr(dim)); + + Value parent = domain_op.getParent(); + if(!parent) + { + // Get associated index + IndexTreeIndicesOp index_op; + Operation* use = *(domain_op->user_begin()); + // TODO: Fix danger of infinite loop!!! + while(!(index_op = llvm::dyn_cast(use))) + { + use = *(use->user_begin()); + } + assert(index_op); + + if(dim == 0) + { + parent = nullptr; + } else + { + // Infer parent index variable + for(Operation* use : index_op->getUsers()) + { + IndexTreeIndexToTensorOp access_op = llvm::dyn_cast(use); + if(!access_op || access_op.getTensor() != tensor || access_op.getDim() != dim) + continue; + + parent = access_op.getPrevDim(); + IndexTreeIndexToTensorOp prev_access_op = + llvm::cast(parent.getDefiningOp()); + if(mlir::failed(this->liftAccessOp(domain_op, prev_access_op))) + return failure(); + + break; + } + } + } + + new_domain = rewriter.create( + loc, + domain_type, + tensor, + dim_size, + rewriter.getUI32IntegerAttr(dim), + parent + ); + } else { + //Domain is dense + auto tensor_type = llvm::cast(tensor.getType()); + auto max = tensor_type.getShape()[dim]; + Value max_val; + if(max < 0) { + auto prev = rewriter.saveInsertionPoint(); + if(tensor.getDefiningOp()) + { + rewriter.setInsertionPointAfter(tensor.getDefiningOp()); + } + else if(auto block_arg = mlir::dyn_cast(tensor)){ // If no definingOp, it is a block argument + rewriter.setInsertionPointToStart(block_arg.getOwner()); + } + else + { + assert(false && "Unhandled condition"); + } + Value dim_val = rewriter.create(loc, rewriter.getIndexType(), rewriter.getIndexAttr(dim)); + max_val = rewriter.create(loc, rewriter.getIndexType(), tensor, dim_val); + rewriter.restoreInsertionPoint(prev); + } else { + max_val = rewriter.create(loc, rewriter.getIndexType(), rewriter.getIndexAttr(max)); + } + new_domain = rewriter.create(loc, domain_type, max_val, tensor, rewriter.getI32ArrayAttr({static_cast(dim)})); + } + rewriter.replaceOp(domain_op, new_domain); + return success(); + } +}; + + +struct SimplifyIntersectionOp : public OpRewritePattern { + SimplifyIntersectionOp(MLIRContext *context) + : OpRewritePattern(context, /*benefit=*/2) {} + + mlir::LogicalResult + matchAndRewrite(IndexTreeDomainIntersectionOp op, + mlir::PatternRewriter &rewriter) const override { + // Look through the input of the current transpose. + Value first_domain = op->getOperand(0); + SmallVector domains; + SmallVector to_remove; + SmallVector tensors; + SmallVector dims; + SmallVector maximums; + IndexTreeDenseDomainOp operand_op; + for(auto operand : op.getDomains()) + { + if(!llvm::isa(operand.getDefiningOp())){ + return failure(); + } + + if((operand_op = operand.getDefiningOp())) + { + to_remove.push_back(operand_op.getOperation()); + auto operand_tensors = operand_op.getTensors(); + tensors.insert(tensors.end(), operand_tensors.begin(), operand_tensors.end()); + auto tensor_dims = operand_op.getDimsAttr(); + dims.insert(dims.end(), tensor_dims.begin(), tensor_dims.end()); + maximums.push_back(operand_op.getDimSize()); + } + else + { + domains.push_back(operand); + } + } + + if(domains.size() == 0) // All domains are dense + { + if(to_remove.size() > 1){ + auto loc = op->getLoc(); + auto context = rewriter.getContext(); + Value max; + // TODO: Do we need this to check if the domains are compatible? + for(Value new_max : maximums){ + if(arith::ConstantOp c = new_max.getDefiningOp()){ + if(llvm::cast(c.getValue()).getValue().isNegative()){ + continue; + } + continue; + } + max = new_max; + break; + } + indexTree::DomainType domain_type = indexTree::DomainType::get(context); + Value new_domain = rewriter.create(loc, domain_type, max, tensors, rewriter.getArrayAttr(dims)); + rewriter.replaceOp(op, {new_domain}); + } else { + rewriter.replaceOp(op, {first_domain}); + } + } else if(domains.size() == 1) + { + // Remove intersection op completely + rewriter.replaceOp(op, {domains[0]}); + } else + { + // Keep only non-dense operands + if(domains.size() == op.getDomains().size() && op.getDimSize() != nullptr){ + return failure(); + } + auto loc = op->getLoc(); + auto context = rewriter.getContext(); + indexTree::DomainType domain_type = indexTree::DomainType::get(context); + Value dim_size = llvm::dyn_cast(domains[0].getDefiningOp()).getDimensionSize(); + Value new_domain = rewriter.create(loc, domain_type, domains, dim_size); + rewriter.replaceOp(op, {new_domain}); + } + + // Delete newly unused values + for(Operation* unused_op : to_remove) + { + if(unused_op->use_empty()) + rewriter.eraseOp(unused_op); + } + + return success(); + } +}; + +struct SimplifyMaskOp : public mlir::OpRewritePattern { + SimplifyMaskOp(mlir::MLIRContext *context) + : OpRewritePattern(context, /*benefit=*/2) {} + + mlir::LogicalResult + matchAndRewrite(IndexTreeMaskedDomainOp op, + mlir::PatternRewriter &rewriter) const override + { + Operation* mask_domain; + if(llvm::isa(op.getMask().getType())){ + mask_domain = op.getMask().getDefiningOp(); + } else { + auto fill_mask = op.getMask().getDefiningOp(); + mask_domain = fill_mask.getDomain().getDefiningOp(); + } + + Operation* base = op.getBase().getDefiningOp(); + if(!llvm::isa(mask_domain) || !llvm::isa(base)){ + return failure(); + } + + if(llvm::isa(mask_domain)) + { + if(llvm::isa(base)) + { + auto masked_domain = llvm::cast(mask_domain); + auto dense_domain = llvm::cast(base); + dense_domain.getTensorsMutable().append(masked_domain.getTensors()); + auto dims = SmallVector(dense_domain.getDims().getAsRange()); + dims.append(masked_domain.getDims().getAsRange().begin(), masked_domain.getDims().getAsRange().end()); + rewriter.modifyOpInPlace(dense_domain, [&](){dense_domain.setDimsAttr(rewriter.getArrayAttr(dims));}); + } + rewriter.replaceOp(op, op.getBase()); + return success(); + } + + if(llvm::isa(base)) + { + rewriter.replaceOp(op, mask_domain); + return success(); + } + + if(op.getDimSize() == nullptr) { + op.getDimSizeMutable().assign(llvm::cast(base).getDimensionSize()); + return success(); + } + return failure(); + } +}; + +struct SimplifyUnionOp : public mlir::OpRewritePattern { + SimplifyUnionOp(mlir::MLIRContext *context) + : OpRewritePattern(context, /*benefit=*/2) {} + + mlir::LogicalResult + matchAndRewrite(IndexTreeDomainUnionOp op, + mlir::PatternRewriter &rewriter) const override + { + bool can_replace = false; + auto context = rewriter.getContext(); + SmallVector domains; + llvm::SmallDenseSet domain_set; + SmallVector tensors; + SmallVector dims; + indexTree::DomainType domain_type = indexTree::DomainType::get(context); + for(auto operand : op.getDomains()) + { + if(!domain_set.contains(operand)){ + domains.push_back(operand); + domain_set.insert(operand); + } + } + + for(auto operand : domains) + { + if(!llvm::isa(operand.getDefiningOp())){ + return failure(); + } + + if(auto operand_op = operand.getDefiningOp()){ + auto operand_tensors = operand_op.getTensors(); + tensors.insert(tensors.end(), operand_tensors.begin(), operand_tensors.end()); + auto tensor_dims = operand_op.getDimsAttr(); + dims.insert(dims.end(), tensor_dims.begin(), tensor_dims.end()); + can_replace = true; + } else if(auto operand_op = operand.getDefiningOp()) { + tensors.push_back(operand_op.getTensor()); + dims.push_back(operand_op.getDimAttr()); + } + } + + if(can_replace) + { + Value dim_size = llvm::dyn_cast(domains[0].getDefiningOp()).getDimensionSize(); + rewriter.replaceOpWithNewOp(op, domain_type, dim_size, tensors, rewriter.getArrayAttr(dims)); + for(auto operand: domains) + { + rewriter.eraseOp(operand.getDefiningOp()); + } + } + else + { + if(op.getDimSize() != nullptr) + { + return failure(); + } + + Value dim_size = llvm::dyn_cast(domains[0].getDefiningOp()).getDimensionSize(); + auto loc = op->getLoc(); + Value new_op = rewriter.create(loc, domain_type, domains, dim_size); + rewriter.replaceOp(op, {new_op}); + } + return success(); + } +}; + +struct InferOutputDomains : public OpRewritePattern { + InferOutputDomains(MLIRContext *context) + : OpRewritePattern(context, /*benefit=*/2) {} + + Value copyDomain(Value domain, + mlir::PatternRewriter &rewriter, + IRMapping& map, + Location loc, + llvm::SmallDenseMap& index_vars) const + { + Value new_domain; + Operation* domain_op = domain.getDefiningOp(); + if(llvm::isa(domain_op)) + { + auto intersection_domain_op = llvm::cast(domain_op); + for(Value subdomain : intersection_domain_op.getDomains()){ + copyDomain(subdomain, rewriter, map, loc, index_vars); + } + } + if(llvm::isa(domain_op)) + { + auto union_domain_op = llvm::cast(domain_op); + for(Value subdomain : union_domain_op.getDomains()){ + copyDomain(subdomain, rewriter, map, loc, index_vars); + } + } + + if(llvm::isa(domain_op)) + { + auto masked_domain_op = llvm::cast(domain_op); + Value mask = masked_domain_op.getMask(); + if(llvm::isa(mask.getType())) + mask = copyDomain(masked_domain_op.getMask(), rewriter, map, loc, index_vars); + else { + assert(llvm::isa(mask.getType())); // Enforced by verifier create by table gen. + auto fill_mask_op = mask.getDefiningOp(); + assert(fill_mask_op && "Currently do not support creating tensors from arbitrary bitmasks."); + mask = copyDomain(fill_mask_op.getDomain(), rewriter, map, loc, index_vars); + } + copyDomain(masked_domain_op.getBase(), rewriter, map, loc, index_vars); + + new_domain = rewriter.clone(*domain_op, map)->getResult(0); + auto new_masked_domain = new_domain.getDefiningOp(); + rewriter.modifyOpInPlace(new_masked_domain, [&](){new_masked_domain.getMaskMutable().assign(mask);}); + + } else if(llvm::isa(domain_op)) + { + auto sparse_domain_op = llvm::cast(domain_op); + + // Ensure parent domain will also be copied. Otherwise create it + Value new_parent_domain = nullptr; + if(sparse_domain_op.getParent()) + { + auto index_to_tensor_op = sparse_domain_op.getParent().getDefiningOp(); + auto index_var = index_to_tensor_op.getIndex().getDefiningOp(); + if(index_vars.find(index_var) == index_vars.end()){ + new_parent_domain = copyDomain(index_var.getDomain(), rewriter, map, loc, index_vars); + index_vars.insert(std::make_pair(index_var.getResult(), index_var.getDomain())); + } + } + + // Clone without parent + new_domain = rewriter.create(loc, + domain_op->getResultTypes(), + sparse_domain_op.getTensor(), + sparse_domain_op.getDimAttr(), + sparse_domain_op.getFormatAttr(), + sparse_domain_op.getPos(), + sparse_domain_op.getCrd(), + sparse_domain_op.getPosSize(), + sparse_domain_op.getCrdSize(), + sparse_domain_op.getDimSize(), + nullptr); + + if(new_parent_domain) + { + // Create or fold so multiple levels of nested domains are foleded into one + new_domain = rewriter.create(loc, + domain_op->getResultTypes(), + llvm::SmallVector{new_parent_domain, new_domain}, + sparse_domain_op.getDimSize()); + } + map.map(sparse_domain_op, new_domain); + + } else { + // Clone + new_domain = rewriter.clone(*domain_op, map)->getResult(0); + } + + auto new_domain_op = new_domain.getDefiningOp(); + for(auto arg : new_domain_op->getOperands()) + { + Operation* origin = arg.getDefiningOp(); + if(origin && (new_domain_op->getBlock() == origin->getBlock()) && new_domain_op->isBeforeInBlock(origin)) + { + rewriter.modifyOpInPlace(new_domain_op, [&]() { new_domain_op->moveAfter(origin); }); + rewriter.setInsertionPointAfter(new_domain_op); + } + } + return new_domain; + } + + mlir::LogicalResult + matchAndRewrite(IndexTreeSparseTensorOp op, + mlir::PatternRewriter &rewriter) const override { + for(auto domain : op.getDomains()) + { + if(!llvm::isa(domain.getDefiningOp())) + { + return failure(); + } + } + + // Get the LHSOperandOp which creates this tensor + Value tensor = op->getResult(0); + IndexTreeOp tree = nullptr; + BlockArgument tensor_arg = nullptr; + for(OpOperand& operand : tensor.getUses()) + { + if((tree = llvm::dyn_cast(operand.getOwner()))){ + tensor_arg = tree.getBody()->getArgument(operand.getOperandNumber()); + break; + } + } + + IndexTreeLHSOperandOp lhs_op = nullptr; + for(Operation* op : tensor_arg.getUsers()) + { + if(llvm::isa(op)) + { + lhs_op = llvm::cast(op); + } + } + + if(lhs_op == nullptr) + return failure(); + + + auto crds = lhs_op.getCrds(); + unsigned dims = (lhs_op.getNumOperands() - 1) / 2; + Value empty_domain = lhs_op.getOperand(0); + llvm::IndexedMap domains(empty_domain); + llvm::SmallDenseMap index_vars; + domains.resize(dims); + for(Value crd : crds){ + auto access_op = llvm::dyn_cast(crd.getDefiningOp()); + if(access_op == nullptr){ + return failure(); + } + auto index_op = llvm::dyn_cast(access_op.getIndex().getDefiningOp()); + if(index_op == nullptr){ + return failure(); + } + + Value domain = index_op.getDomain(); + index_vars.insert(std::make_pair(index_op.getResult(), domain)); + domains[access_op.getDim()] = domain; + } + + // Successfully matched! Cannot fail after this point. + auto loc = op->getLoc(); + SmallVector new_args; + IRMapping map; + for(unsigned dim = 0; dim < dims; dim++){ + Value domain_copy = copyDomain(domains[dim], rewriter, map, loc, index_vars); + new_args.push_back(domain_copy); + } + auto new_tensor = rewriter.create(loc, op->getResult(0).getType(), new_args); + rewriter.replaceOp(op, new_tensor->getResults()); + return success(); + } +}; + +void indexTree::populateDomainConcretizationPatterns( + MLIRContext *context, RewritePatternSet &patterns) { + patterns.add(context); +} + +struct IndexTreeDomainConcretization : comet::impl::IndexTreeDomainConcretizationBase { + using IndexTreeDomainConcretizationBase::IndexTreeDomainConcretizationBase; + + void runOnOperation() override { + mlir::RewritePatternSet domain_concretization_patterns(&getContext()); + indexTree::populateDomainConcretizationPatterns(&getContext(), domain_concretization_patterns); + indexTree::populateMaskDomainTransformationPatterns(&getContext(), domain_concretization_patterns); + if(failed(mlir::applyPatternsAndFoldGreedily(getOperation(), std::move(domain_concretization_patterns)))) { + return signalPassFailure(); + } + } +}; + +/// Apply the compressed workspace transformations on the index tree IR +std::unique_ptr mlir::comet::createIndexTreeDomainConcretizationPass() +{ + return std::make_unique(); +} \ No newline at end of file diff --git a/lib/Dialect/IndexTree/Transforms/Fusion.cpp b/lib/Dialect/IndexTree/Transforms/Fusion.cpp deleted file mode 100644 index e4300223..00000000 --- a/lib/Dialect/IndexTree/Transforms/Fusion.cpp +++ /dev/null @@ -1,1130 +0,0 @@ -//===- Fusion.cpp ------===// -// -// Copyright 2022 Battelle Memorial Institute -// -// Redistribution and use in source and binary forms, with or without modification, -// are permitted provided that the following conditions are met: -// -// 1. Redistributions of source code must retain the above copyright notice, this list of conditions -// and the following disclaimer. -// -// 2. Redistributions in binary form must reproduce the above copyright notice, this list of conditions -// and the following disclaimer in the documentation and/or other materials provided with the distribution. -// -// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED -// WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR -// PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, -// INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE -// GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, -// WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE -// OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. -// -//===----------------------------------------------------------------------===// -// -// This pass performs redundancy-aware kernel Fusion on index tree dialect -// The details of the partial fusion can be found in the following paper. -// ReACT: Redundancy-Aware Code Generation for Tensor Expressions. -// Tong Zhou, Ruiqin Tian, Rizwan A Ashraf, Roberto Gioiosa, Gokcen Kestor, Vivek Sarkar. -// 2022 31st International Conference on Parallel Architectures and Compilation Techniques (PACT). October 2022. -//===----------------------------------------------------------------------===// - -#include "comet/Dialect/IndexTree/IR/IndexTreeDialect.h" -#include "comet/Dialect/IndexTree/Passes.h" -#include "comet/Dialect/TensorAlgebra/IR/TADialect.h" -#include "comet/Dialect/Utils/Utils.h" -#include "comet/Dialect/IndexTree/Transforms/UnitExpression.h" - -#include "mlir/Dialect/Arith/IR/Arith.h" -#include "mlir/Dialect/Bufferization/IR/Bufferization.h" -#include "mlir/Dialect/MemRef/IR/MemRef.h" -#include "mlir/Dialect/Func/IR/FuncOps.h" - -#include "llvm/Support/Debug.h" -#include -#include -#include - -#include -#include -#include -#include -#include -#include -#include - -using namespace mlir; -using namespace mlir::arith; -using namespace mlir::bufferization; -using namespace mlir::indexTree; -using namespace mlir::tensorAlgebra; - -using llvm::SmallVector; -using llvm::StringRef; - -#define DEBUG_TYPE "partial-fusion" - -// *********** For debug purpose *********// -// #define COMET_DEBUG_MODE -#include "comet/Utils/debug.h" -#undef COMET_DEBUG_MODE -// *********** For debug purpose *********// - -//===----------------------------------------------------------------------===// -/// KernelFusion PASS -//===----------------------------------------------------------------------===// - -namespace -{ - class IndexTreeKernelFusionPass - : public mlir::PassWrapper> - { - private: - static void test(mlir::func::FuncOp &funcop); - - static std::vector getAllItrees(mlir::func::FuncOp &funcop); - - static std::vector getAllComputeLHSs(mlir::func::FuncOp &funcop); - - static int getIndicesOpsIndex(mlir::Operation *op); - - static std::vector getPathFromRoot(mlir::Operation *op); - - static std::vector getLongestCommonPrefix(std::vector> &paths); - - static mlir::Value createNewTensorDecl(const mlir::Value &old_dense_tensor_decl, - uint32_t rank_base); - - static void createNewTensor(const mlir::Value &old_tensor_alloc, - const mlir::Value &old_tensor_load, - uint32_t rank_base, - mlir::Value &new_tensor_alloc, - mlir::Value &new_tensor_load); - - static mlir::Value createReducedComputeLHS(mlir::Operation *lhs_op, - mlir::Value &new_tensor_load, - uint32_t rank_base); - - static mlir::Value createReducedComputeRHS(mlir::Operation *rhs_op, - mlir::Value &new_tensor_load, - mlir::Value &old_tensor_load, - uint32_t rank_base); - - static void replaceOldOperandToNew(mlir::Operation *old_operand, mlir::Value &new_val); - - static void replaceOldTensorFillOp(const mlir::Value &old_dense_tensor_decl, - const mlir::Value &new_dense_tensor_decl); - - static void replaceOldLinalgFillOp(mlir::Value &old_tensor_alloc, mlir::Value &old_tensor_load, mlir::Value &new_tensor_load); - - static mlir::Value createResetComputeRHS(const mlir::Value &new_dense_tensor_decl, - mlir::Operation *last_common_prefix); - - static mlir::Value createResetComputeLHS(const mlir::Value &new_tensor_load, - mlir::Operation *last_common_prefix, - int &lcp_index, - int64_t &rank); - - static void insertTensorReset( - const std::vector &lcp, - const mlir::Value &new_dense_tensor_decl); - - static void doKernelFusion(std::vector &itrees, mlir::func::FuncOp &funcop); - - static void reduceTensorDimension(std::vector &LHSs, mlir::func::FuncOp &funcop); - - public: - MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(IndexTreeKernelFusionPass) - void runOnOperation() override; - - void RedundancyAwareFusion(mlir::func::FuncOp &funcop); - }; /// class IndexTreeKernelFusionPass -} /// End anonymous namespace - -[[maybe_unused]] void IndexTreeKernelFusionPass::test(mlir::func::FuncOp &funcop) -{ - int level = 0; - funcop.walk([&](mlir::Operation *op) - { - if (llvm::isa(*op)) { - comet_debug() << "level: " << level++ << "\n"; - level = 0; - comet_pdump(op); - auto rootOp = dyn_cast(*op); - [[maybe_unused]] Value computeOp = rootOp.getOperation()->getOperand(0); - comet_vdump(computeOp); - } else if (llvm::isa(*op)) { - comet_debug() << "level: " << level++ << "\n"; - comet_pdump(op); - - auto indicesOp = dyn_cast(*op); - int numOperands = indicesOp.getOperation()->getNumOperands(); - for (int oper_i = 0; oper_i < numOperands; ++oper_i) { - // Value oper = indicesOp.getOperation()->getOperand(oper_i); - comet_debug() << "Operand " << oper_i << ": "; - comet_vdump(indicesOp.getOperation()->getOperand(oper_i)); - } - ArrayAttr idsArrayAttr = indicesOp.getIndices(); - for (auto ida: idsArrayAttr) { - int id = ida.cast().getInt(); - comet_debug() << "id: " << id; - } - comet_debug() << "\n"; - } else if (llvm::isa(*op)) { - comet_debug() << "level: " << level++ << "\n"; - comet_pdump(op); - auto computeOp = dyn_cast(*op); - int num_operands = computeOp.getOperation()->getNumOperands(); - for (int op_i = 0; op_i < num_operands; ++op_i) { - - comet_debug() << "Operand " << op_i << ": "; - comet_vdump(computeOp.getOperand(op_i)); - } - - bool is_comp_worksp_opt = computeOp.getCompWorkspOpt(); - std::string semiring(computeOp.getSemiring().data()); - comet_debug() << " is_comp_worksp_opt: " << is_comp_worksp_opt << " semiring: " << semiring << "\n"; - - std::vector> opFormats; - std::vector> opPerms; - std::vector > inputOutputMapping; - getFormatsPermsOfComputeOp(computeOp, opFormats, opPerms, inputOutputMapping); - /// opFormats - comet_debug() << "["; - for (auto strings: opFormats) { - comet_debug() << "["; - for (auto fmt: strings) { - comet_debug() << fmt << " "; - } - comet_debug() << "]"; - } - comet_debug() << "]\n"; - - /// opPerms - comet_debug() << "["; - for (auto ints: opPerms) { - comet_debug() << "["; - for (auto perm: ints) { - comet_debug() << perm << " "; - } - comet_debug() << "]"; - } - comet_debug() << "]\n"; - - /// inputOutputMapping - comet_debug() << "["; - for (auto bools: inputOutputMapping) { - comet_debug() << "["; - for (auto mp: bools) { - comet_debug() << mp << " "; - } - comet_debug() << "]"; - } - comet_debug() << "]\n"; - } else if (llvm::isa(*op)) { - comet_debug() << "level: " << level << "\n"; - comet_pdump(op); - } else if (llvm::isa(*op)) { - comet_debug() << "level: " << level << "\n"; - comet_pdump(op); - } }); -} - -std::vector IndexTreeKernelFusionPass::getAllItrees(mlir::func::FuncOp &funcop) -{ - std::vector itrees; - funcop.walk([&](indexTree::IndexTreeOp op) - { itrees.push_back(op.getOperation()); }); - - return itrees; -} - -std::vector IndexTreeKernelFusionPass::getAllComputeLHSs(mlir::func::FuncOp &funcop) -{ - std::vector lhss; - funcop.walk([&](indexTree::IndexTreeComputeLHSOp op) - { lhss.push_back(op.getOperation()); }); - - return lhss; -} - -int IndexTreeKernelFusionPass::getIndicesOpsIndex(mlir::Operation *op) -{ - assert(llvm::isa(op) && "Error: op is not IndexTreeIndicesOp."); - auto indices_op = llvm::dyn_cast(*op); - int index = indices_op.getIndices()[0].cast().getInt(); - return index; -} - -std::vector IndexTreeKernelFusionPass::getPathFromRoot(mlir::Operation *op) -{ - std::vector path; - - while (!llvm::isa(*op)) - { - { /// test - comet_debug() << "op\n"; - comet_pdump(op); - } - path.push_back(op); - op = *(op->getUsers().begin()); - } - - std::reverse(path.begin(), path.end()); - return path; -} - -std::vector IndexTreeKernelFusionPass:: - getLongestCommonPrefix(std::vector> &paths) -{ - std::vector lcp; - assert(paths.size() == 2 && "Error: too many paths to handle with."); - - auto &path0 = paths[0]; - auto &path1 = paths[1]; - - uint32_t index = 0; - while (index < path0.size() && index < path1.size()) - { - if (path0[index] == path1[index]) - { - lcp.push_back(path0[index]); - { /// test - /// comet_debug(); - comet_pdump(path0[index]); - } - } - else - { - break; - } - ++index; - } - - return lcp; -} - -mlir::Value IndexTreeKernelFusionPass::createNewTensorDecl( - const mlir::Value &old_dense_tensor_decl, - uint32_t rank_base) -{ - mlir::Operation *old_tensor_op = old_dense_tensor_decl.getDefiningOp(); - auto loc = old_dense_tensor_decl.getLoc(); - OpBuilder builder(old_tensor_op); - - /// Get operands - std::vector operands; - - TensorType old_tensor = old_dense_tensor_decl.getType().cast(); - - for (int64_t i = rank_base; i < old_tensor.getRank(); i++) - { - if (old_tensor.isDynamicDim(i)) - { - operands.push_back(old_tensor_op->getOperands()[i]); - } - else - { - operands.push_back(builder.create(loc, old_tensor.getDimSize(i))); - } - } - - /// Get format - std::string format; - { - mlir::tensorAlgebra::DenseTensorDeclOp old_tensor_decl_op = llvm::dyn_cast( - old_tensor_op); - format = std::string(old_tensor_decl_op.getFormat()); - assert(format == "Dense" && ("Error: only support for Dense old tensor, not " + format + ".\n").c_str()); - } - - /// Create ta.dense_tensor_decl - auto float_type = mlir::RankedTensorType::get({mlir::ShapedType::kDynamic}, - builder.getF64Type()); - mlir::Value new_dense_tensor_decl = builder.create(loc, - float_type, - operands, - format); - - comet_debug() << "new_dense_tensor_decl\n"; - comet_vdump(new_dense_tensor_decl); - - return new_dense_tensor_decl; -} - -[[maybe_unused]] void IndexTreeKernelFusionPass::createNewTensor( - const mlir::Value &old_tensor_alloc, - const mlir::Value &old_tensor_load, - uint32_t rank_base, - mlir::Value &new_tensor_alloc, - mlir::Value &new_tensor_load) -{ - mlir::Operation *old_tensor_alloc_op = old_tensor_alloc.getDefiningOp(); - if (old_tensor_alloc_op->getNumOperands() == 2 && - rank_base == 1) - { - /// TODO(zpeng): here only considered reduce a 2D tensor to an 1D tensor in SDDMM kernel - /// Get the previous memref.load operand - auto loc = old_tensor_alloc.getLoc(); - OpBuilder builder(old_tensor_alloc.getDefiningOp()); - - /// Create ta.index_label - mlir::Value index_label_op = builder.create(loc); - - comet_debug() << "index_label_op\n"; - comet_vdump(index_label_op); - - /// Create ta.dense_tensor_decl - auto float_type = mlir::RankedTensorType::get({mlir::ShapedType::kDynamic}, - builder.getF64Type()); - std::vector operands = {index_label_op}; - std::string format = "Dense"; - mlir::Value dense_tensor_decl_op = builder.create(loc, - float_type, - operands, - format); - - comet_debug() << "dense_tensor_decl_op\n"; - comet_vdump(dense_tensor_decl_op); - - new_tensor_alloc = index_label_op; - new_tensor_load = dense_tensor_decl_op; - } - else if (old_tensor_alloc_op->getNumOperands() <= 1) - { - mlir::TensorType tensor_ty = old_tensor_load.getType().cast(); - uint32_t rank = tensor_ty.getRank(); - { /// test - /// comet_debug(); - comet_debug() << "rank: " << rank << "\n"; - for (uint32_t r_i = 0; r_i < rank; ++r_i) - { - comet_debug() << r_i << ": " << tensor_ty.getDimSize(r_i) << " "; - } - comet_debug() << "\n"; - } - - mlir::OpBuilder builder(old_tensor_alloc.getDefiningOp()); - auto loc = old_tensor_alloc.getLoc(); - std::vector dims_vec; /// int64_t is required for MemRefType - for (uint32_t r_i = rank_base; r_i < rank; ++r_i) - { - dims_vec.push_back(tensor_ty.getDimSize(r_i)); - } - auto dims_ty = mlir::MemRefType::get( - llvm::ArrayRef(dims_vec), - tensor_ty.getElementType()); - - std::vector operands; - operands.insert(operands.end(), - old_tensor_alloc_op->getOperands().begin() + rank_base, - old_tensor_alloc_op->getOperands().end()); - new_tensor_alloc = builder.create(loc, dims_ty, operands, builder.getI64IntegerAttr(32)); - new_tensor_load = builder.create(loc, new_tensor_alloc); - } - else - { - llvm::errs() << "Error: IndexTreeKernelFusionPass::createNewTensor() does not support tensors whose rank is larger than 2."; - } -} - -mlir::Value IndexTreeKernelFusionPass::createReducedComputeLHS( - mlir::Operation *lhs_op, - mlir::Value &new_tensor_load, - uint32_t rank_base) -{ - assert( - llvm::isa(*lhs_op) && "Error: not a type of indexTree::IndexTreeComputeLHSOp"); - OpBuilder builder(lhs_op); - - /// Get old formats and perms - mlir::indexTree::IndexTreeComputeLHSOp it_compute_lhs_op = llvm::dyn_cast( - lhs_op); - ArrayAttr op_formats_ArrayAttr = it_compute_lhs_op.getAllFormats(); - ArrayAttr op_perms_ArrayAttr = it_compute_lhs_op.getAllPerms(); - std::vector> old_formats_strs = convertArrayAttrStrTo2DVector(op_formats_ArrayAttr); - std::vector> old_perms_ints = convertArrayAttrIntTo2DVector(op_perms_ArrayAttr); - - /// Create the new formats - /// i.g., convert [["D", "D"]] to [["D"]] - SmallVector new_formats; - SmallVector formats; - formats.insert(formats.end(), old_formats_strs[0].begin() + rank_base, old_formats_strs[0].end()); - new_formats.push_back(builder.getStrArrayAttr(formats)); - - /// Create the new perms - /// i.g., convert [[1, 0]] to [[0]] - SmallVector new_perms; - SmallVector perms; - perms.insert(perms.end(), old_perms_ints[0].begin() + rank_base, old_perms_ints[0].end()); - new_perms.push_back(builder.getI64ArrayAttr(perms)); - - /// Create the new ComputeLHS - std::vector tensors; - tensors.push_back(new_tensor_load); - mlir::Value new_lhs_op = builder.create( - lhs_op->getLoc(), - mlir::UnrankedTensorType::get(builder.getF64Type()), - tensors, - builder.getArrayAttr(new_perms), - builder.getArrayAttr(new_formats)); - - return new_lhs_op; -} - -mlir::Value IndexTreeKernelFusionPass::createReducedComputeRHS( - mlir::Operation *rhs_op, - mlir::Value &new_tensor_load, - mlir::Value &old_tensor_load, - uint32_t rank_base) -{ - assert( - llvm::isa(*rhs_op) && "Error: not a type of indexTree::IndexTreeComputeRHSOp"); - OpBuilder builder(rhs_op); - - /// Get old formats and perms - mlir::indexTree::IndexTreeComputeRHSOp it_compute_rhs_op = llvm::dyn_cast( - rhs_op); - ArrayAttr op_formats_ArrayAttr = it_compute_rhs_op.getAllFormats(); - ArrayAttr op_perms_ArrayAttr = it_compute_rhs_op.getAllPerms(); - std::vector> old_formats_strs = convertArrayAttrStrTo2DVector(op_formats_ArrayAttr); - std::vector> old_perms_ints = convertArrayAttrIntTo2DVector(op_perms_ArrayAttr); - - /// Locate the operand to be reduced - uint32_t tensor_id = 0; - for (auto val : rhs_op->getOperands()) - { - if (val.getDefiningOp() == old_tensor_load.getDefiningOp()) - { - break; - } - ++tensor_id; - } - - /// Create the new formats - /// i.g., convert [["D", "D"], ["D", "D"]] to [["D"], ["D", "D"]] - SmallVector new_formats; - for (uint32_t f_i = 0; f_i < old_formats_strs.size(); ++f_i) - { - SmallVector formats; - if (f_i == tensor_id) - { /// for the new reduced tensor - formats.insert(formats.end(), old_formats_strs[f_i].begin() + rank_base, old_formats_strs[f_i].end()); - } - else - { /// for other remaining old operands - formats.insert(formats.end(), old_formats_strs[f_i].begin(), old_formats_strs[f_i].end()); - } - new_formats.push_back(builder.getStrArrayAttr(formats)); - } - - /// Create the new perms - /// i.g., convert [[1, 0], [0, 2]] to [[0], [0, 2]] - SmallVector new_perms; - for (uint32_t p_i = 0; p_i < old_perms_ints.size(); ++p_i) - { - SmallVector perms; - if (p_i == tensor_id) - { /// for the new reduced tensor - perms.insert(perms.end(), old_perms_ints[p_i].begin() + rank_base, old_perms_ints[p_i].end()); - } - else - { /// for other remaining old operands - perms.insert(perms.end(), old_perms_ints[p_i].begin(), old_perms_ints[p_i].end()); - } - new_perms.push_back(builder.getI64ArrayAttr(perms)); - } - - /// Create the new ComputeRHS - std::vector tensors; - uint32_t val_i = 0; - for (auto val : rhs_op->getOperands()) - { - if (val_i == tensor_id) - { - tensors.push_back(new_tensor_load); /// new reduced tensor - } - else - { - tensors.push_back(val); /// other remaining old operands - } - ++val_i; - } - mlir::Value new_rhs_op = builder.create( - rhs_op->getLoc(), - mlir::UnrankedTensorType::get(builder.getF64Type()), - tensors, - builder.getArrayAttr(new_perms), - builder.getArrayAttr(new_formats)); - - return new_rhs_op; -} - -void IndexTreeKernelFusionPass::replaceOldOperandToNew(mlir::Operation *old_operand, mlir::Value &new_val) -{ - for (auto user : old_operand->getUsers()) - { - int op_i = 0; - for (auto op : user->getOperands()) - { - if (op.getDefiningOp() == old_operand) - { - /// Replace the old operand to the new one - user->setOperand(op_i, new_val); - break; - } - ++op_i; - } - } - - /// Erase the old operand - old_operand->erase(); -} - -void IndexTreeKernelFusionPass::replaceOldTensorFillOp( - const mlir::Value &old_dense_tensor_decl, - const mlir::Value &new_dense_tensor_decl) -{ - for (auto user : old_dense_tensor_decl.getUsers()) - { - if (mlir::isa(user)) - { - /// Get value - mlir::tensorAlgebra::TensorFillOp fill_op = llvm::dyn_cast(user); - auto value_attr = fill_op.getValueAttr(); - - OpBuilder builder(user); - [[maybe_unused]] auto new_fill_op = builder.create( - user->getLoc(), - new_dense_tensor_decl, - value_attr); - /// builder.getF64FloatAttr(0)); - - comet_debug() << "new_fill_op\n"; - comet_vdump(new_fill_op); - - user->erase(); - break; - } - else - { - comet_debug() << "Error: user's type is not supported, yet.\n"; - continue; - } - } -} - -[[maybe_unused]] void IndexTreeKernelFusionPass::replaceOldLinalgFillOp( - mlir::Value &old_tensor_alloc, - mlir::Value &old_tensor_load, - mlir::Value &new_tensor_load) -{ - for (auto user : old_tensor_alloc.getUsers()) - { - OpBuilder builder(user); - if (mlir::isa(user)) - { - - comet_debug() << "user\n"; - comet_pdump(user); - - /// Get the ConstantOp - Value constant_op; - for (auto op : user->getOperands()) - { - comet_debug() << "op\n"; - comet_vdump(op); - - if (mlir::isa(op.getDefiningOp())) - { - constant_op = llvm::dyn_cast(*op.getDefiningOp()); - break; - } - } - - comet_debug() << "constant_op\n"; - comet_vdump(constant_op); - - /// Create MemRef - /// mlir::Value memref; - if (llvm::isa(new_tensor_load.getDefiningOp())) - { - [[maybe_unused]] auto new_fill_op = builder.create( - user->getLoc(), - new_tensor_load, - builder.getF64FloatAttr(0)); - { - comet_debug() << "new_fill_op\n"; - comet_vdump(new_fill_op); - } - } - else if (llvm::isa(new_tensor_load.getDefiningOp())) - { - /// For GNN kernel. - mlir::Value memref = llvm::dyn_cast(new_tensor_load.getDefiningOp()).getMemref(); - [[maybe_unused]] auto new_fill_op = builder.create( - user->getLoc(), - constant_op, - memref); - - comet_debug() << "new_fill_op\n"; - comet_pdump(new_fill_op); - } - - /// Erase the user - user->erase(); - break; - } - else - { - comet_debug() << "Error: some unsupported users.\n"; - continue; - } - } -} - -mlir::Value IndexTreeKernelFusionPass::createResetComputeRHS( - const mlir::Value &new_dense_tensor_decl, - mlir::Operation *last_common_prefix) -{ - auto loc = last_common_prefix->getLoc(); - OpBuilder builder(last_common_prefix); - /// Generate itComputeRHS, operand is constant 0, allFormats = [[]], allPerms = [[]] - /// Operand is constant 0 - ConstantOp constant_zero; - for (mlir::Operation *user : new_dense_tensor_decl.getUsers()) - { - if (!mlir::isa(user)) - { - continue; - } - /// Get the ConstantOp - for (mlir::Value op : user->getOperands()) - { - if (mlir::isa(op.getDefiningOp())) - { - constant_zero = llvm::dyn_cast(*op.getDefiningOp()); - break; - } - } - break; - } - if (constant_zero.getOperation() == nullptr) - { - comet_debug() << "create ConstantOp\n"; - constant_zero = builder.create(loc, - builder.getF64Type(), - builder.getF64FloatAttr(0)); - } - - comet_debug() << "constant_zero\n"; - comet_vdump(constant_zero); - - std::vector tensors_rhs; - tensors_rhs.push_back(constant_zero); - - SmallVector indices_rhs; - SmallVector empty_index; - indices_rhs.push_back(builder.getI64ArrayAttr(empty_index)); - - SmallVector formats_rhs; - SmallVector empty_format; - formats_rhs.push_back(builder.getStrArrayAttr(empty_format)); - - /// TODO(zpeng): What if the type is not F64? - mlir::Value compute_rhs = builder.create( - loc, - mlir::UnrankedTensorType::get(builder.getF64Type()), - tensors_rhs, - builder.getArrayAttr(indices_rhs), - builder.getArrayAttr(formats_rhs)); - - return compute_rhs; -} - -mlir::Value IndexTreeKernelFusionPass::createResetComputeLHS( - const mlir::Value &new_dense_tensor_decl, - mlir::Operation *last_common_prefix, - int &lcp_index, - int64_t &rank) -{ - auto loc = last_common_prefix->getLoc(); - OpBuilder builder(last_common_prefix); - /// Generate itComputeLHS, operand is tensor_load, allFormats = [["D"]], allPerms = [[1]] - /// Operand is tensor_load - std::vector tensors_lhs; - tensors_lhs.push_back(new_dense_tensor_decl); - lcp_index = getIndicesOpsIndex(last_common_prefix); - mlir::TensorType tensor_ty = new_dense_tensor_decl.getType().cast(); - rank = tensor_ty.getRank(); - - /// Get indices [[1]] - SmallVector indices_lhs; - SmallVector one_index; - for (uint32_t r_i = 0; r_i < rank; ++r_i) - { - one_index.push_back(lcp_index + 1 + r_i); - } - indices_lhs.push_back(builder.getI64ArrayAttr(one_index)); - - /// Get formats [["D"]] - SmallVector formats_lhs; - SmallVector one_format(rank, "D"); - formats_lhs.push_back(builder.getStrArrayAttr(one_format)); - - mlir::Value compute_lhs = builder.create( - loc, - mlir::UnrankedTensorType::get(builder.getF64Type()), - tensors_lhs, - builder.getArrayAttr(indices_lhs), - builder.getArrayAttr(formats_lhs)); - - return compute_lhs; -} - -mlir::Value createResetIndicesOps( - int lcp_index, - int64_t rank, - mlir::Operation *last_common_prefix, - const mlir::Value &compute_op) -{ - auto loc = last_common_prefix->getLoc(); - OpBuilder builder(last_common_prefix); - - /// Create index nodes - int bound_index = lcp_index + rank; - auto i64_type = builder.getI64Type(); - mlir::Value last_indices_op; - - for (int index = bound_index; index > lcp_index; --index) - { - SmallVector indices = {index}; - auto indices_attr = builder.getI64ArrayAttr(indices); - mlir::Value indices_op; - if (index == bound_index) - { - /// The IndicesOp node closest to the ComputeOp node - /// TODO(zhen.peng): new attribute iterator_type - auto dumb_iterator_type = builder.getStringAttr("default"); - indices_op = builder.create( - loc, - i64_type, - compute_op, - indices_attr, - dumb_iterator_type); - } - else - { - /// TODO(zhen.peng): new attribute iterator_type - auto dumb_iterator_type = builder.getStringAttr("default"); - indices_op = builder.create( - loc, - i64_type, - last_indices_op, - indices_attr, - dumb_iterator_type); - } - last_indices_op = indices_op; - } - - return last_indices_op; -} - -void IndexTreeKernelFusionPass::insertTensorReset( - const std::vector &lcp, - const mlir::Value &new_dense_tensor_decl) -{ - /// Generate itComputeRHS, operand is constant 0, allFormats = [[]], allPerms = [[]] - mlir::Operation *last_common_prefix = lcp.back(); - auto loc = last_common_prefix->getLoc(); - OpBuilder builder(last_common_prefix); - /// Operand is constant 0 - mlir::Value compute_rhs = createResetComputeRHS(new_dense_tensor_decl, - last_common_prefix); - - comet_debug() << "compute_rhs\n"; - comet_vdump(compute_rhs); - - /// Generate itComputeLHS, operand is tensor_load, allFormats = [["D"]], allPerms = [[1]] - int lcp_index; - int64_t rank; - mlir::Value compute_lhs = createResetComputeLHS(new_dense_tensor_decl, - last_common_prefix, - lcp_index, - rank); - - comet_debug() << "compute_lhs\n"; - comet_vdump(compute_lhs); - - auto comp_worksp_opt = builder.getBoolAttr(false); - mlir::StringAttr semiring = builder.getStringAttr("noop_times"); - mlir::StringAttr maskType = builder.getStringAttr("none"); - - IntegerType i64Type = IntegerType::get(builder.getContext(), 64); - mlir::Value compute_op = builder.create( - loc, - i64Type, - compute_rhs, - compute_lhs, - comp_worksp_opt, - semiring, - maskType); - { - comet_debug() << "compute_op\n"; - comet_vdump(compute_op); - } - - /// Create index nodes - mlir::Value indices_root = createResetIndicesOps(lcp_index, - rank, - last_common_prefix, - compute_op); - /// Add to the last common prefix - std::vector operands; - operands.insert(operands.end(), last_common_prefix->getOperands().begin(), last_common_prefix->getOperands().end()); - operands.push_back(indices_root); - last_common_prefix->setOperands(operands); -} - -void IndexTreeKernelFusionPass::doKernelFusion( - std::vector &itrees, mlir::func::FuncOp &funcop) -{ - - std::deque> buffer; - { - /// Collects all itrees' operands and initialize the buffer. - std::vector operands; - for (auto itree : itrees) - { - for (auto operand : itree->getOperands()) - { - if (llvm::isa(operand.getDefiningOp())) - { - comet_pdump(operand.getDefiningOp()); - operands.push_back(operand.getDefiningOp()); - } - } - } - - buffer.push_back(std::move(operands)); - } - comet_debug() << "Buffer size " << buffer.size() << "\n"; - while (!buffer.empty()) - { - std::vector operands = buffer.front(); - buffer.pop_front(); - comet_debug() << "Buffer size " << buffer.size() << "\n"; - if (operands.size() < 2) - { - /// Too few nodes to fuse - continue; - } - - /// Nodes to be fused are clustered and then fused to the host node - std::vector is_clustered(operands.size(), false); - for (int host_i = operands.size() - 1; host_i >= 0; --host_i) - { - if (is_clustered[host_i]) - { - continue; - } - - /// Get the host - /// Note: The host is the latest one, otherwise the fused nodes are not in correct usage order and got Error: - mlir::Operation *host = operands[host_i]; - comet_debug() << host << "host\n"; - comet_debug() << "host\n"; - comet_pdump(host); - - is_clustered[host_i] = true; - int host_index = getIndicesOpsIndex(host); - - /// Cluster other nodes to the host - std::vector cluster; - for (int node_i = 0; node_i < host_i; ++node_i) - { - if (is_clustered[node_i]) - { - continue; - } - mlir::Operation *node = operands[node_i]; - int node_index = getIndicesOpsIndex(node); - comet_debug() << "node\n"; - comet_pdump(node); - /// Check if node_i can be fused with host_i - if (node_index == host_index) - { - cluster.push_back(node); - is_clustered[node_i] = true; - } - } - cluster.push_back(host); - - /// Set the operands of the host to the operands of nodes in the cluster - std::vector sub_operands; - for (mlir::Operation *node : cluster) - { - for (auto operand : node->getOperands()) - { - sub_operands.push_back(operand); - } - } - host->setOperands(sub_operands); - - /// Update the buffer - std::vector tmp_operands; - for (auto &op : sub_operands) - { - tmp_operands.push_back(op.getDefiningOp()); - } - buffer.push_back(std::move(tmp_operands)); - - /// Erase other nodes in the cluster and their users - for (int node_i = cluster.size() - 2; node_i >= 0; --node_i) - { - mlir::Operation *node = cluster[node_i]; - comet_debug() << "Erasing \n"; - for (auto u : node->getUsers()) - { - comet_pdump(u); - u->erase(); - } - comet_pdump(node); - node->erase(); - } - } - } -} - -void IndexTreeKernelFusionPass::reduceTensorDimension(std::vector &LHSs, mlir::func::FuncOp &funcop) -{ - for (mlir::Operation *lhs_op : LHSs) - { - comet_debug() << "lhs_op\n"; - comet_pdump(lhs_op); - - mlir::Value tensor = lhs_op->getOperand(0); - comet_debug() << "tensor\n"; - comet_vdump(tensor); - - auto users = tensor.getUsers(); - - comet_debug() << "users\n"; - for ([[maybe_unused]] auto u : users) - { - comet_pdump(u); - } - - int num_users = 0; - mlir::Operation *rhs_op = nullptr; - for (auto u : users) - { - if (llvm::isa(*u)) - { - ++num_users; - } - else if (llvm::isa(*u)) - { - ++num_users; - rhs_op = u; - } - } - if (num_users < 2) - { - /// No need to reduce its dimension - continue; - } - else if (num_users > 2) - { - comet_debug() << "Error: should not have more than 2 users for the tensor.\n"; - continue; - } - - comet_debug() << "rhs_op\n"; - comet_pdump(rhs_op); - - /// Get paths of lsh_op and rhs_op - std::vector> paths; - paths.push_back(getPathFromRoot(lhs_op)); - paths.push_back(getPathFromRoot(rhs_op)); - - /// Get the longest common prefix - std::vector lcp = getLongestCommonPrefix(paths); - if (lcp.empty()) - { - /// No common prefix for reducing tensor dimension - continue; - } - - uint32_t rank_base = lcp.size(); - mlir::Value new_dense_tensor_decl = createNewTensorDecl(tensor, rank_base); - - comet_debug() << "new_dense_tensor_decl\n"; - comet_vdump(new_dense_tensor_decl); - - mlir::Value new_compute_lhs = createReducedComputeLHS(lhs_op, - new_dense_tensor_decl, - rank_base); - comet_debug() << "new_compute_lhs\n"; - comet_vdump(new_compute_lhs); - - mlir::Value new_compute_rhs = createReducedComputeRHS(rhs_op, - new_dense_tensor_decl, - tensor, - rank_base); - - comet_debug() << "new_compute_rhs\n"; - comet_vdump(new_compute_rhs); - - /// Switch ComputeLHS - comet_debug() << "replace LHS to new LSH\n"; - replaceOldOperandToNew(lhs_op, new_compute_lhs); - - /// Switch ComputeRHS - comet_debug() << "replace RHS to new RHS\n"; - replaceOldOperandToNew(rhs_op, new_compute_rhs); - - comet_debug() << "replace ta.fill\n"; - replaceOldTensorFillOp(tensor, new_dense_tensor_decl); - - /// Erase the old dense_tensor_decl - tensor.getDefiningOp()->erase(); - - /// Generate T = 0 to reset the intermediate tensor. The location is under the last common prefix. - insertTensorReset(lcp, new_dense_tensor_decl); - } -} - -void IndexTreeKernelFusionPass::RedundancyAwareFusion(mlir::func::FuncOp &funcop) -{ - comet_vdump(funcop); - comet_debug() << "ParitalFusionIT pass\n"; - - /// Basic partial fusion - std::vector itrees = getAllItrees(funcop); - if (itrees.size() < 2) - { - /// Only one itree node cannot do fusion. - return; - } - - doKernelFusion(itrees, funcop); - - /// Reduce tensor dimension - comet_vdump(funcop); - std::vector LHSs = getAllComputeLHSs(funcop); - reduceTensorDimension(LHSs, funcop); -} - -void IndexTreeKernelFusionPass::runOnOperation() -{ - LLVM_DEBUG(llvm::dbgs() << "start IndexTreeKernelFusionPass\n"); - comet_debug() << " start KernelFusion pass \n"; - func::FuncOp func = getOperation(); - RedundancyAwareFusion(func); -} - -//// Apply the partial fusion on the index tree dialect -std::unique_ptr mlir::comet::createIndexTreeKernelFusionPass() -{ - return std::make_unique(); -} diff --git a/lib/Dialect/IndexTree/Transforms/IterationDomain.cpp b/lib/Dialect/IndexTree/Transforms/IterationDomain.cpp deleted file mode 100644 index 949ea52b..00000000 --- a/lib/Dialect/IndexTree/Transforms/IterationDomain.cpp +++ /dev/null @@ -1,190 +0,0 @@ -// -// Copyright 2022 Battelle Memorial Institute -// -// Redistribution and use in source and binary forms, with or without modification, -// are permitted provided that the following conditions are met: -// -// 1. Redistributions of source code must retain the above copyright notice, this list of conditions -// and the following disclaimer. -// -// 2. Redistributions in binary form must reproduce the above copyright notice, this list of conditions -// and the following disclaimer in the documentation and/or other materials provided with the distribution. -// -// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED -// WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR -// PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, -// INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE -// GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, -// WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE -// OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. -// - -#include "comet/Dialect/IndexTree/Transforms/IterationDomain.h" -#include "comet/Dialect/IndexTree/Transforms/Tensor.h" - -using namespace std; -// Since multiple threads can lower different functions, -// we need one for each thread lowering. -thread_local std::vector> domains; - -IterDomain *IterDomain::makeDomain(Tensor *tensor, int dim) -{ - auto d = make_unique(tensor, dim); - auto p = d.get(); - domains.push_back(std::move(d)); - return p; -} - -IterDomain *IterDomain::conjunct(IterDomain *a, IterDomain *b) -{ - auto d = make_unique('*', a, b); - auto p = d.get(); - domains.push_back(std::move(d)); - return p; -} - -IterDomain *IterDomain::conjunct(std::vector &domains) -{ - assert(!domains.empty()); - auto d = domains[0]; - for (unsigned long i = 1; i < domains.size(); i++) - { - d = conjunct(d, domains[i]); - } - return d; -} - -bool IterDomain::equals(IterDomain *that) -{ - auto thisSimplified = this->getSimplified(); - auto thatSimplified = that->getSimplified(); - return thisSimplified == thatSimplified; -} - -IterDomain *IterDomain::getSimplified() -{ - if (getOp() == '*') - { - if (getLeft()->isDense()) - { - return getRight(); - } - else if (getRight()->isDense()) - { - return getLeft(); - } - } - - return this; -} - -std::string IterDomain::str() -{ - if (isLeafNode()) - { - string s = "(" + getTensor()->str() + "," + to_string(getDim()) + ")"; - return s; - } - else - { - assert(getLeft() != nullptr && getRight() != nullptr); - string s = getLeft()->str() + string(1, getOp()) + getRight()->str(); - return s; - } -} - -std::string IterDomain::getFormat() -{ - return getTensor()->getFormat(getDim()); -} - -bool IterDomain::isDense() -{ - return getFormat() == "D"; -} - -std::string BoolExpr::str() -{ - if (isTrue()) - { - return "true"; - } - - auto t = value.first; - assert(t != nullptr); - auto dim = value.second; - return "(" + t->str() + ", " + to_string(dim) + ")"; -} - -void BoolExpr::setTrue() -{ - isConstantTrue = true; -} - -void BoolExpr::setFalse() -{ - isConstantFalse = true; -} - -unique_ptr BoolExprManager::trueNode; - -std::vector> BoolExprManager::exprs; - -BoolExpr *BoolExprManager::getTrue() -{ - if (trueNode == nullptr) - { - trueNode = make_unique(); - trueNode->setTrue(); - } - - return trueNode.get(); -} - -BoolExpr *BoolExprManager::makeBoolExpr(BDDElem elem) -{ - auto expr = make_unique(elem); - BoolExpr *ret = expr.get(); - exprs.push_back(std::move(expr)); - return ret; -} -Tensor *IterDomain::getTensor() const -{ - return tensor; -} -void IterDomain::setTensor(Tensor *Tensor) -{ - tensor = Tensor; -} -int IterDomain::getDim() const -{ - return dim; -} -void IterDomain::setDim(int Dim) -{ - dim = Dim; -} -char IterDomain::getOp() const -{ - return op; -} -void IterDomain::setOp(char Op) -{ - op = Op; -} -IterDomain *IterDomain::getLeft() const -{ - return left; -} -void IterDomain::setLeft(IterDomain *Left) -{ - left = Left; -} -IterDomain *IterDomain::getRight() const -{ - return right; -} -void IterDomain::setRight(IterDomain *Right) -{ - right = Right; -} diff --git a/lib/Dialect/IndexTree/Transforms/IterationDomainInference.cpp b/lib/Dialect/IndexTree/Transforms/IterationDomainInference.cpp new file mode 100644 index 00000000..fa0bf592 --- /dev/null +++ b/lib/Dialect/IndexTree/Transforms/IterationDomainInference.cpp @@ -0,0 +1,264 @@ +// +// Copyright 2022 Battelle Memorial Institute +// +// Redistribution and use in source and binary forms, with or without modification, +// are permitted provided that the following conditions are met: +// +// 1. Redistributions of source code must retain the above copyright notice, this list of conditions +// and the following disclaimer. +// +// 2. Redistributions in binary form must reproduce the above copyright notice, this list of conditions +// and the following disclaimer in the documentation and/or other materials provided with the distribution. +// +// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED +// WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR +// PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, +// INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE +// GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, +// WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +// OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. +// + +#include "mlir/Dialect/Func/IR/FuncOps.h" +#include "mlir/Dialect/Index/IR/IndexOps.h" +#include "mlir/Transforms/GreedyPatternRewriteDriver.h" +#include "mlir/Dialect/Arith/IR/Arith.h" +#include "mlir/Dialect/Tensor/IR/Tensor.h" + +#include "llvm/ADT/StringSet.h" +#include "llvm/ADT/SmallSet.h" +#include "llvm/ADT/DenseMap.h" + +#include "comet/Dialect/IndexTree/IR/IndexTreeDialect.h" +#include "comet/Dialect/TensorAlgebra/IR/TADialect.h" +#include "comet/Dialect/IndexTree/Passes.h" +#include "comet/Dialect/IndexTree/Patterns.h" + +using namespace mlir; +using namespace mlir::indexTree; + +// *********** For debug purpose *********// +//#define COMET_DEBUG_MODE +#include "comet/Utils/debug.h" +// *********** For debug purpose *********// + +namespace mlir { + namespace comet{ + #define GEN_PASS_DEF_INDEXTREEDOMAININFERENCE + #include "comet/Dialect/IndexTree/Passes.h.inc" + } +} + +struct IndexTreeDomainInference : comet::impl::IndexTreeDomainInferenceBase { + using IndexTreeDomainInferenceBase::IndexTreeDomainInferenceBase; + void runOnOperation() override; +}; + +struct InferIndexDomain : public OpRewritePattern { + CopiedDomainAnalysis& copiedDomains; + InferIndexDomain(MLIRContext *context, CopiedDomainAnalysis& copiedDomains) + : OpRewritePattern(context, /*benefit=*/1), copiedDomains(copiedDomains) {} + + mlir::LogicalResult + matchAndRewrite(IndexTreeIndicesOp op, mlir::PatternRewriter &builder) const override { + if(op.getDomain()) + return failure(); + + comet_vdump(op); + + Location loc = op.getLoc(); + auto context = builder.getContext(); + indexTree::DomainType domain_type = indexTree::DomainType::get(context); + + // Map operands to domains + // Will break if operand appears multiple times for the same index variable + llvm::SmallDenseMap operands_to_domains; + + // Set of tensors that we need to infer domain + llvm::SmallDenseSet intermediate_tensors; + + // Set of all compute operands + llvm::SmallPtrSet compute_ops; + for(Operation* tensor_access_op : op->getUsers()) + { + if(!llvm::isa(tensor_access_op)) + continue; + + comet_pdump(tensor_access_op); + for(Operation* operand_op : tensor_access_op->getUsers()) + { + if(!llvm::isa(operand_op)) + continue; + + comet_pdump(operand_op); + auto tensor_val = llvm::cast(tensor_access_op).getTensor(); + unsigned dim = llvm::cast(tensor_access_op).getDim(); + comet_vdump(tensor_val); + comet_debug() << "dim: " << dim << "\n"; + Value domain; + if(llvm::isa_and_present(tensor_val.getDefiningOp()) && op->isBeforeInBlock(tensor_val.getDefiningOp())) + { + if(copiedDomains.isCopiedDomain(tensor_val, dim)){ + // The domain is the same on the LHS as it is on the RHS + // We need to find the domain of this index variable on the RHS + intermediate_tensors.insert(tensor_val); + } else { + // We promote this domain to dense because we cannot narrow the domain further + // Use negative one to indicate to use the maximums of the other tensors + Value neg_one = builder.create(loc, builder.getI32Type(), builder.getI32IntegerAttr(-1)); + domain = builder.create( + loc, + domain_type, + neg_one, + ValueRange(), + builder.getI32ArrayAttr(llvm::ArrayRef()) + ); + operands_to_domains.insert(std::pair(operand_op, domain)); + } + } else { + domain = builder.create(loc, + domain_type, + tensor_val, + builder.getUI32IntegerAttr(dim), + tensorAlgebra::TensorFormatEnumAttr::get(context, tensorAlgebra::TensorFormatEnum::UNK), + nullptr + ); + operands_to_domains.insert(std::pair(operand_op, domain)); + } + compute_ops.insert(operand_op->user_begin(), operand_op->user_end()); + break; + } + } + + llvm::SmallVector domains; + llvm::SmallVector backup; + auto tree_op = op->getParentOfType(); + auto yield_op = llvm::cast(tree_op.getBody()->getTerminator()); + llvm::SmallDenseSet outputs(yield_op.getOperands().begin(), yield_op.getOperands().begin() + tree_op.getInputs().size()); + tree_op.getBody()->walk([&](IndexTreeComputeOp compute_op) { + // Check if compute op uses this index variable + if(!compute_ops.contains(compute_op.getOperation())){ + return WalkResult::advance(); + } + + Value temp_domain; + // Check if compute op needs intersection + if(compute_op.getComputeMissing()){ + SmallVector union_domains; + for(auto operand_op_val : compute_op.getRhs()) + { + auto operand_op = operand_op_val.getDefiningOp(); + if(operands_to_domains.find(operand_op) != operands_to_domains.end()) { + union_domains.push_back(operands_to_domains[operand_op]); + } + } + + if(union_domains.size() > 1){ + temp_domain = builder.create( + loc, + domain_type, + union_domains, + nullptr + ); + } else { + temp_domain = union_domains[0]; + } + + + } else { + SmallVector intersection_domains; + for(auto operand_op_val : compute_op.getRhs()) + { + auto operand_op = operand_op_val.getDefiningOp(); + comet_pdump(operand_op); + if(operands_to_domains.find(operand_op) != operands_to_domains.end()) + intersection_domains.push_back(operands_to_domains[operand_op]); + } + if(intersection_domains.size() > 1) + temp_domain = builder.create(loc, domain_type, intersection_domains, nullptr); + else + temp_domain = intersection_domains[0]; + } + + if(compute_op.getMask()) { + auto operand_op = compute_op.getMask().getDefiningOp(); + if(operands_to_domains.find(operand_op) != operands_to_domains.end()) + { + Value mask_domain = operands_to_domains[operand_op]; + temp_domain = builder.create( + loc, + domain_type, + mask_domain, + temp_domain, + nullptr + ); + } + } + + + if(intermediate_tensors.contains(compute_op.getResult())){ + for(auto user : compute_op.getResult().getUsers()) { + if(llvm::isa(user)){ + operands_to_domains.insert(std::make_pair(user, temp_domain)); + } + } + } + + if(copiedDomains.isReductionVar(compute_op, op.getResult()) || outputs.contains(compute_op.getResult())) { + domains.push_back(temp_domain); + } + backup.push_back(temp_domain); + return WalkResult::advance(); + }); + + Value final_domain; + if(domains.size() > 1) { + final_domain = builder.create(loc, + domain_type, domains, nullptr); + } else if(domains.size() == 1) { + final_domain = domains[0]; + } else { + // If the index variable is not used in any domain restricted by the output + // Assume we need the inferred domain + // TODO: Analyze which other index variables refer to the same output domain + // and use that to infer the domain. + final_domain = backup[backup.size() - 1]; + } + comet_vdump(final_domain); + + indexTree::IndexNodeType index_node_type = indexTree::IndexNodeType::get(context); + builder.replaceOpWithNewOp( + op, index_node_type, op.getParent(), final_domain, op.getIsParallelAttr(), op.getParallelDimAttr()); + return success(); + } +}; + +struct CreateZeroMaskOp : public OpRewritePattern { + CreateZeroMaskOp(MLIRContext *context) + : OpRewritePattern(context, /*benefit=*/1) {} + + mlir::LogicalResult + matchAndRewrite(IndexTreeFillMaskOp op, mlir::PatternRewriter &builder) const override { + return failure(); + } +}; + +void mlir::indexTree::populateDomainInferencePatterns( + MLIRContext *context, RewritePatternSet &patterns, CopiedDomainAnalysis &copiedDomains) { + patterns.add(context, copiedDomains); + patterns.add(context); +} + +void IndexTreeDomainInference::runOnOperation() { + mlir::RewritePatternSet domain_inference_patterns(&getContext()); + CopiedDomainAnalysis& copiedDomains = getAnalysis(); + populateDomainInferencePatterns(&getContext(), domain_inference_patterns, copiedDomains); + if (mlir::failed(mlir::applyPatternsAndFoldGreedily(getOperation(), std::move(domain_inference_patterns)))) + signalPassFailure(); +} + +/// Apply the compressed workspace transformations on the index tree IR +std::unique_ptr mlir::comet::createIndexTreeDomainInferencePass() +{ + return std::make_unique(); +} \ No newline at end of file diff --git a/lib/Dialect/IndexTree/Transforms/KernelFusion.cpp b/lib/Dialect/IndexTree/Transforms/KernelFusion.cpp new file mode 100644 index 00000000..d9627067 --- /dev/null +++ b/lib/Dialect/IndexTree/Transforms/KernelFusion.cpp @@ -0,0 +1,1226 @@ +// +// Copyright 2022 Battelle Memorial Institute +// +// Redistribution and use in source and binary forms, with or without modification, +// are permitted provided that the following conditions are met: +// +// 1. Redistributions of source code must retain the above copyright notice, this list of conditions +// and the following disclaimer. +// +// 2. Redistributions in binary form must reproduce the above copyright notice, this list of conditions +// and the following disclaimer in the documentation and/or other materials provided with the distribution. +// +// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED +// WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR +// PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, +// INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE +// GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, +// WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +// OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. +// + +#include "mlir/Dialect/Func/IR/FuncOps.h" +#include "mlir/Dialect/Index/IR/IndexOps.h" +#include "mlir/Dialect/MemRef/IR/MemRef.h" +#include "mlir/Dialect/SCF/IR/SCF.h" +#include "mlir/Transforms/GreedyPatternRewriteDriver.h" + +#include "llvm/ADT/STLExtras.h" +#include "llvm/ADT/StringSet.h" + +#include "comet/Dialect/IndexTree/IR/IndexTreeDialect.h" +#include "comet/Dialect/TensorAlgebra/IR/TADialect.h" +#include "comet/Dialect/IndexTree/Passes.h" + +#include +#include "comet/Dialect/IndexTree/Patterns.h" + +using namespace mlir; +using namespace mlir::indexTree; + +// *********** For debug purpose *********// +//#define COMET_DEBUG_MODE +#include "comet/Utils/debug.h" +// *********** For debug purpose *********// + +namespace mlir { + namespace comet{ +#define GEN_PASS_DEF_INDEXTREEKERNELFUSION +#include "comet/Dialect/IndexTree/Passes.h.inc" + } +} + +namespace { + +using DimCompound = std::pair; /// , i.e., Tensor's i-th dimension. + +template +uint32_t getIdxInVector(const llvm::SmallVector &array, T value) +{ + auto it = std::find(array.begin(), array.end(), value); + if (it != array.end()) { + return it - array.begin(); + } else { + return (uint32_t) -1; + } +} + + +Value getRealLhsTensor(IndexTreeOp itree, + Value lhs_tensor) +{ + /// get the block arguments and itree arguments + llvm::SmallVector block_args; + for (Value arg : itree.getRegion().getBlocks().front().getArguments()) { + block_args.push_back(arg); + } + llvm::SmallVector itree_args; + for (Value arg : itree.getInputs()) { + itree_args.push_back(arg); + } + + /// Find the true operand if it is one of block arguments + uint32_t idx = getIdxInVector(block_args, lhs_tensor); + if (idx != (uint32_t) -1) { + return itree_args[idx]; + } else { + return lhs_tensor; + } +} + + +void collectDimCompound(IndexTreeOp itree_op, + uint32_t &num_indexOps /*out*/, + llvm::SmallVector> &indexOp_to_DimCompound /*out*/) +{ + /// Go through each IndexOp's `it.IndexToTensorDim` + itree_op.getOperation()->walk([&](IndexTreeIndicesOp index_op) { + ++num_indexOps; + comet_vdump(index_op); + indexOp_to_DimCompound.push_back(llvm::DenseSet()); + auto &dims_set = indexOp_to_DimCompound.back(); + for (auto user : index_op->getUsers()) { + if (auto indexToTensorDimOp = llvm::dyn_cast(user)) { + comet_vdump(indexToTensorDimOp); + uint32_t dim = indexToTensorDimOp.getDim(); + /// Find the true operand if it is one of block arguments + Value operand = getRealLhsTensor(itree_op, indexToTensorDimOp.getTensor()); + dims_set.insert(std::make_pair(dim, operand)); + } + } + }); +} + + +bool isSameDimension(const llvm::DenseSet &dim1, const llvm::DenseSet &dim2) +{ + /// Find the intersection + if (std::any_of(dim1.begin(), dim1.end(), [&](const DimCompound &entry) { return dim2.contains(entry); })) { + /// If any dimension in dim1 happens to be in dim2, dim1 and dim2 has intersection, thus they are representing the same dimension + return true; + } else { + return false; + } +} + + +llvm::SmallVector findCommonIndex(const llvm::SmallVector> &indices_list1, + const llvm::SmallVector> &indices_list2) +{ + llvm::SmallVector common_indices; + + uint32_t bound = std::min(indices_list1.size(), indices_list2.size()); + for (uint32_t idx = 0; idx < bound; ++idx) { + if (isSameDimension(indices_list1[idx], indices_list2[idx])) { + common_indices.push_back(idx); + } else { + break; + } + } + + comet_debug() << "common_indices.size(): " << common_indices.size() << "\n"; + return common_indices; +} + + +void collectCommonIndices(const llvm::SmallVector &itree_list, + llvm::SmallVector &itree_to_num_indexOps /*out*/, + std::unordered_map> &itree_to_common_indices /*out*/) +{ + llvm::SmallVector>> itree_to_DimCompound(itree_list.size()); + for (uint32_t tree_i = 0; tree_i < itree_list.size(); ++tree_i) { + collectDimCompound(itree_list[tree_i], + itree_to_num_indexOps[tree_i] /*out*/, + itree_to_DimCompound[tree_i] /*out*/); + } + + SmallVector host_indices; + itree_list[0]->walk([&host_indices] (IndexTreeIndicesOp indice) { + host_indices.push_back(indice); + }); + + + for (uint32_t tree_i = 1; tree_i < itree_list.size(); ++tree_i) { + itree_to_common_indices[tree_i] = findCommonIndex(itree_to_DimCompound[0], itree_to_DimCompound[tree_i]); + SmallVector child_indices; + itree_list[tree_i]->walk([&child_indices] (IndexTreeIndicesOp indice) { + child_indices.push_back(indice); + }); + for(size_t index: itree_to_common_indices[tree_i]) + { + if(host_indices[index].getIsParallel() != child_indices[index].getIsParallel()) + { + host_indices[index].setIsParallel(false); + } + } +// {/// test +// for (uint32_t idx : itree_to_common_indices[tree_i]) { +// comet_debug() << idx << "\n"; +// } +// } + } + +// {/// test +// comet_debug() << "itree_to_common_indices.size(): " << itree_to_common_indices.size() << "\n"; +// for (auto v : itree_to_num_indexOps) { +// comet_debug() << v << "\n"; +// } +// } + +} + + +void collectOperandsDims(IndexTreeOp itree, + Value &lhs_tensor /*out*/, + llvm::SmallVector &lhs_dims /*out*/, + llvm::SmallVector &lhs_index_idx /*out*/, + llvm::SmallVector &rhs_tensors /*out*/, + llvm::SmallVector> &rhs_dims /*out*/, + llvm::SmallVector> &rhs_index_idx /*out*/) +{ + /// Collect all IndexOp + llvm::SmallVector indexOp_list; + itree.walk([&](IndexTreeIndicesOp op) { + indexOp_list.push_back(op); + }); + + itree.walk([&](IndexTreeComputeOp computeOp) { + /// LHS + Value lhs_OperandOp = computeOp.getLhs(); + { + auto lhs_op = llvm::cast(lhs_OperandOp.getDefiningOp()); + /// Find the true operand if it is one of block arguments + lhs_tensor = getRealLhsTensor(itree, lhs_op.getTensor()); + auto positions = lhs_op.getPos(); + for (Value pos: positions) + { + auto indexToTensorDimOp = llvm::cast(pos.getDefiningOp()); + uint32_t dim = indexToTensorDimOp.getDim(); + lhs_dims.push_back(std::make_pair(dim, lhs_tensor)); + + /// Find out which IndexOp this IndexToTensorDim links to + Value indexOp = indexToTensorDimOp.getIndex(); + uint32_t idx = getIdxInVector(indexOp_list, indexOp); + assert(idx != (uint32_t) -1 && "Error: index_op did not exist."); + lhs_index_idx.push_back(idx); + + { + comet_vdump(lhs_op); + comet_vdump(pos); + comet_vdump(lhs_tensor); + comet_debug() << "dim: " << dim << "\n"; + comet_debug() << "idx: " << idx << "\n"; + comet_debug() << "\n"; + } + } + } + + /// RHS + llvm::SmallVector rhs_OperandOp; + for (Value rhs : computeOp.getRhs()) { + rhs_OperandOp.push_back(rhs); + } + + /// The mask, if exists, is another OperandOp in Index Tree dialect. It will be the first generated `it.OperandOp`, + /// but will be the last argument of generated `it.ComputeOp`. + Value mask_tensor = computeOp.getMask(); + if (mask_tensor) { + rhs_OperandOp.push_back(mask_tensor); + } + + for (Value rhs_value : rhs_OperandOp) { + auto rhs_op = llvm::cast(rhs_value.getDefiningOp()); + rhs_tensors.push_back(rhs_op.getTensor()); + rhs_dims.push_back(llvm::SmallVector()); + rhs_index_idx.push_back(llvm::SmallVector()); + + auto positions = rhs_op.getPos(); + for (Value pos : positions) { + auto indexToTensorDimOp = llvm::cast(pos.getDefiningOp()); + uint32_t dim = indexToTensorDimOp.getDim(); + rhs_dims.back().push_back(std::make_pair(dim, rhs_tensors.back())); + + /// Find out which IndexOp this IndexToTensorDim links to + Value indexOp = indexToTensorDimOp.getIndex(); + uint32_t idx = getIdxInVector(indexOp_list, indexOp); + assert(idx != (uint32_t) -1 && "Error: index_op did not exist."); + rhs_index_idx.back().push_back(idx); + + { + comet_vdump(rhs_op); + comet_vdump(pos); + comet_vdump(rhs_tensors.back()); + comet_debug() << "dim: " << dim << "\n"; + comet_debug() << "idx: " << idx << "\n"; + comet_debug() << "\n"; + } + } + } + }); +} + +uint32_t collectIntermediateIdx(Value prev_lhs_tensor, + const llvm::SmallVector &curr_rhs_tensors) +{ + uint32_t idx = getIdxInVector(curr_rhs_tensors, prev_lhs_tensor); + assert(idx != (uint32_t) -1 && "Error: none of current rhs tensors is the previous lhs tensor."); + return idx; +} + +[[maybe_unused]] std::unordered_map createNewLhsTensors( + uint32_t num_itrees, + llvm::SmallVector &itree_to_lhs_tensors, + std::unordered_map> &itree_to_common_indices) +{ + std::unordered_map itree_to_new_tensors; + for (uint32_t tree_i = 0; tree_i < num_itrees - 1; ++tree_i) { + /// Old tensor example + /// %10 = "ta.dense_tensor_decl"(%5) <{format = "Dense"}> : (index) -> tensor + Value old_tensor = itree_to_lhs_tensors[tree_i]; + llvm::SmallVector &common_indices = itree_to_common_indices[tree_i + 1]; + assert(llvm::isa(old_tensor.getType()) && + "Error: old_tensor is not a mlir::TensorType."); + assert(!common_indices.empty() && "Error: no common indices."); + auto old_tensor_tt = llvm::cast(old_tensor.getType()); + + /// The start dim idx after those common indices + uint32_t dim_base = common_indices.size(); +// dim_base = 1; /// test + if (dim_base < old_tensor_tt.getRank()) { + /// The new tensor is still a tensor, with decreased dimensions. + + uint32_t count_of_dyn_dim = 0; /// How many dynamic dimension is in common indices + for (uint32_t dim_i = 0; dim_i < dim_base; ++dim_i) { + if (old_tensor_tt.isDynamicDim(dim_i)) { + ++count_of_dyn_dim; + } + } + SmallVector operands; /// The remaining dynamic dimensions after fusion. + if (old_tensor.getDefiningOp()->getNumOperands() > count_of_dyn_dim) { + operands.insert(operands.begin(), + old_tensor.getDefiningOp()->getOperands().begin() + count_of_dyn_dim, + old_tensor.getDefiningOp()->getOperands().end()); + } + + llvm::SmallVector shape; + for (int64_t dim_i = dim_base; dim_i < old_tensor_tt.getRank(); ++dim_i) { + if (old_tensor_tt.isDynamicDim(dim_i)) { + shape.push_back(mlir::ShapedType::kDynamic); + } else { + shape.push_back(old_tensor_tt.getDimSize(dim_i)); + } + } + + auto loc = old_tensor.getLoc(); + mlir::OpBuilder builder(old_tensor.getDefiningOp()); + Value new_tensor = builder.create( + loc, + mlir::RankedTensorType::get(shape, old_tensor_tt.getElementType()), + operands); + itree_to_new_tensors[tree_i] = new_tensor; + mlir::TypedAttr zero = builder.getZeroAttr(old_tensor_tt.getElementType()); + builder.create(loc, + new_tensor, + zero); + comet_vdump(new_tensor); + } else { + /// The new tensor is shrunk to a scalar. + auto loc = old_tensor.getLoc(); + mlir::OpBuilder builder(old_tensor.getDefiningOp()); + mlir::TypedAttr zero = builder.getZeroAttr(old_tensor_tt.getElementType()); + Value new_scalar = builder.create(loc, + old_tensor_tt.getElementType(), + zero); + comet_vdump(new_scalar); + itree_to_new_tensors[tree_i] = new_scalar; + } + } + + return itree_to_new_tensors; +} + +void collectComputeOpInfo(IndexTreeOp itree, + mlir::StringRef &semiring /*out*/, + bool &compute_missing /*out*/) +{ + itree.walk([&](IndexTreeComputeOp op) { + semiring = op.getSemiring(); + compute_missing = op.getComputeMissing(); + comet_debug() << semiring << "\n"; + comet_debug() << compute_missing << "\n"; + }); +} + +/// 1) Generate new IndexToTensorDim for the new LHS operand. +/// 2) Replace the old LHS operand with the new one. +/// 3) Erase the old LHS operand and its IndexToTensorDim. +[[maybe_unused]] Value replaceOldLHSOperand(IndexTreeOp itree_op, + Value lhs_tensor, + llvm::SmallVector> &itree_to_lhs_dims, + llvm::SmallVector> &itree_to_lhs_index_idx, + const std::unordered_map> &itree_to_common_indices, + MLIRContext *context, + mlir::IRRewriter &rewriter, + mlir::Location &loc) +{ + /// Get all IndexOp + llvm::SmallVector index_ops; + itree_op.walk([&](IndexTreeIndicesOp op) { + index_ops.push_back(op); + }); + /// Find the LHS operand + IndexTreeLHSOperandOp old_lhs_operand; + uint32_t lhs_count = 0; + itree_op.walk([&](IndexTreeLHSOperandOp op) { + old_lhs_operand = op; + ++lhs_count; + }); + assert(lhs_count == 1 && "Error: The kernel should only have one lhs operand."); + OpBuilder::InsertionGuard guard(rewriter); + rewriter.setInsertionPoint(old_lhs_operand); + + llvm::SmallVector pos; + llvm::SmallVector crds; + Value prev_dim = nullptr; + auto access_type = rewriter.getIndexType(); + uint32_t tree_i = 0; /// the host itree + llvm::SmallVector &lhs_dims = itree_to_lhs_dims[tree_i]; + llvm::SmallVector &lhs_index_idx = itree_to_lhs_index_idx[tree_i]; + uint32_t dim_base = itree_to_common_indices.at(tree_i + 1).size(); +// dim_base = 1; /// test + for (uint32_t d_i = dim_base; d_i < lhs_dims.size(); ++d_i) { + Value index_node = index_ops[lhs_index_idx[d_i]]; + uint32_t dim = lhs_dims[d_i].first - dim_base; + auto access_op = rewriter.create( + loc, + TypeRange({access_type, access_type}), + lhs_tensor, + index_node, + rewriter.getUI32IntegerAttr(dim), + prev_dim); + pos.push_back(access_op.getPos()); + crds.push_back(access_op.getCrd()); + prev_dim = pos.back(); + } + + indexTree::OperandType operand_type = indexTree::OperandType::get(context); + Value new_lhs_operand = rewriter.create(loc, + operand_type, + lhs_tensor, + pos, + crds); + comet_vdump(new_lhs_operand); + comet_vdump(new_lhs_operand.getDefiningOp()->getParentOfType()); + + /// Replace and erase + rewriter.replaceAllUsesWith(old_lhs_operand, new_lhs_operand); + llvm::SmallVector old_index_to_tensor_dim; + for (auto tensor_dim : old_lhs_operand.getPos()) { + old_index_to_tensor_dim.push_back(tensor_dim); + } + std::reverse(old_index_to_tensor_dim.begin(), old_index_to_tensor_dim.end()); + rewriter.eraseOp(old_lhs_operand); + for (auto tensor_dim : old_index_to_tensor_dim) { + rewriter.eraseOp(tensor_dim.getDefiningOp()); + } + + comet_vdump(new_lhs_operand.getDefiningOp()->getParentOfType()); + return new_lhs_operand; +} + + +/// Create the new ComputeOp for the new LHS operand. +/// Only the return type needs change to the new LHS operand's type. +[[maybe_unused]] Value replaceOldComputeOp(IndexTreeOp itree_op, + const mlir::Type &tensor_type, + mlir::IRRewriter &rewriter, + mlir::Location &loc) +{ + IndexTreeComputeOp old_compute_op; + uint32_t count_compute_op = 0; + itree_op.walk([&](IndexTreeComputeOp op) { + old_compute_op = op; + ++count_compute_op; + }); + assert(count_compute_op == 1 && "Error: expect only one ComputeOp"); + + OpBuilder::InsertionGuard guard(rewriter); + rewriter.setInsertionPoint(old_compute_op); + Value new_compute_op = rewriter.create( + loc, + tensor_type, + old_compute_op.getParent(), + old_compute_op.getLhs(), + old_compute_op.getRhs(), + old_compute_op.getMask(), + old_compute_op.getSemiringAttr(), + old_compute_op.getComputeMissingAttr()); + + comet_vdump(new_compute_op); + rewriter.replaceAllUsesWith(old_compute_op, new_compute_op); + rewriter.eraseOp(old_compute_op); + comet_vdump(new_compute_op.getDefiningOp()->getParentOfType()); + + return new_compute_op; +} + + +/// 1) Create the new itree based on the host itree (the 0th one). +/// 2) Replace the old LHS operand to the new one with shrunk dimensions. +/// 3) Replace the old ComputeOp with the new one with correct return type (shrunk tensor type). +IndexTreeOp createNewITree(IndexTreeOp host_itree, + const llvm::SmallVector &tree_types, +// const llvm::SmallVector &itree_arguments, + const llvm::SmallVector &itree_arguments_inputs, + const llvm::SmallVector &itree_arguments_intermediates, + llvm::SmallVector> &itree_to_lhs_dims, + llvm::SmallVector> &itree_to_lhs_index_idx, +// const std::unordered_map> &itree_to_common_indices, + llvm::SmallVector &itree_to_new_compute_op /*out*/, + MLIRContext *context, + mlir::IRRewriter &rewriter, + mlir::Location &loc) +{ + OpBuilder::InsertionGuard guard(rewriter); + rewriter.setInsertionPoint(host_itree); + /// Create the new itree + llvm::SmallVector locs; +// for (Value arg : itree_arguments) { +// locs.push_back(arg.getLoc()); +// } + for (Value arg : itree_arguments_inputs) { + locs.push_back(arg.getLoc()); + } + for (Value arg : itree_arguments_intermediates) { + locs.push_back(arg.getLoc()); + } + IndexTreeOp new_itree = rewriter.create(loc, tree_types, itree_arguments_inputs, + itree_arguments_intermediates); + Region *body = &new_itree.getRegion(); + Block *block = rewriter.createBlock(body, {}, TypeRange(tree_types), locs); + YieldOp dummy_yield = rewriter.create(loc, TypeRange(), itree_arguments_inputs); /// create a dummy YieldOp to make the following inlineBlockBefore() work. + comet_vdump(new_itree); + /// Move the host itree's block to the new itree. The host itree's argument happens to be the first `intermediates` argument of new_itree. + /// An itree's arguments are it.itree($inputs, $intermediates) + rewriter.inlineBlockBefore(&host_itree.getRegion().front(), + block->getTerminator(), + {block->getArgument(itree_arguments_inputs.size())} /*argValues to replace host itree (the source)'s block arguments*/); + +// /// Replace the old LHS +// replaceOldLHSOperand(new_itree, +// block->getArgument(tree_i), +// itree_to_lhs_dims, +// itree_to_lhs_index_idx, +// itree_to_common_indices, +// context, +// rewriter, +// loc); +// +// /// Replace the old ComputeOp +// Value new_compute_op = replaceOldComputeOp(new_itree, +// tree_types[tree_i], +// rewriter, +// loc); +// itree_to_new_compute_op.push_back(new_compute_op); + body->walk([&](IndexTreeComputeOp computOp) { + itree_to_new_compute_op.push_back(computOp); + }); + assert(itree_to_new_compute_op.size() == 1 && "Expect only one computeOp in the itree."); + + rewriter.eraseOp(dummy_yield); /// delete the dummy YieldOp, so the block will contain only one YieldOp. + comet_vdump(new_itree); + + return new_itree; +} + + +llvm::SmallVector createIndexOps( + uint32_t tree_i, + llvm::SmallVector &host_index_ops, + const llvm::SmallVector &itree_to_num_indexOps, + std::unordered_map> &itree_to_common_indices, + MLIRContext *context, + mlir::IRRewriter &rewriter, + mlir::Location &loc) +{ + /// 1) Record the common + llvm::SmallVector index_ops; + for (uint32_t common_i = 0; common_i < itree_to_common_indices[tree_i].size(); ++common_i) { + index_ops.push_back(host_index_ops[common_i]); + } + /// 2) Add the new index ops + assert(!index_ops.empty() && "Common indices should not be empty."); + Value parent = index_ops.back(); + indexTree::IndexNodeType index_node_type = indexTree::IndexNodeType::get(context); + uint32_t num_new_index_op = itree_to_num_indexOps[tree_i] - itree_to_common_indices[tree_i].size(); + for (uint32_t op_i = 0; op_i < num_new_index_op; ++op_i) { + parent = rewriter.create(loc, index_node_type, parent); + index_ops.push_back(parent); + } + + return index_ops; +} + + +Value createLHSOperand( + uint32_t tree_i, + llvm::SmallVector &index_ops, + Value lhs_tensor, + llvm::SmallVector> &itree_to_lhs_dims, + llvm::SmallVector> &itree_to_lhs_index_idx, +// const std::unordered_map> &itree_to_common_indices, + MLIRContext *context, + mlir::IRRewriter &rewriter, + mlir::Location &loc) +{ + llvm::SmallVector pos; + llvm::SmallVector crds; + Value prev_dim = nullptr; + auto access_type = rewriter.getIndexType(); + llvm::SmallVector &lhs_dims = itree_to_lhs_dims[tree_i]; + llvm::SmallVector &lhs_index_idx = itree_to_lhs_index_idx[tree_i]; + uint32_t dim_base = 0; +// if (itree_to_common_indices.find(tree_i + 1) != itree_to_common_indices.end()) { +// dim_base = itree_to_common_indices.at(tree_i + 1).size(); +// } + for (uint32_t d_i = dim_base; d_i < lhs_dims.size(); ++d_i) + { + Value index_node = index_ops[lhs_index_idx[d_i]]; + uint32_t dim = lhs_dims[d_i].first - dim_base; + auto access_op = rewriter.create( + loc, + TypeRange({access_type, access_type}), + lhs_tensor, + index_node, + rewriter.getUI32IntegerAttr(dim), + prev_dim); + pos.push_back(access_op.getPos()); + crds.push_back(access_op.getCrd()); + prev_dim = pos.back(); + } + + indexTree::OperandType operand_type = indexTree::OperandType::get(context); + Value lhs_operand = rewriter.create(loc, + operand_type, + lhs_tensor, + pos, + crds); + + return lhs_operand; +} + + +llvm::SmallVector createRHSOperands( + uint32_t tree_i, + uint32_t intermediate_idx, + Value prev_computeOp, + llvm::SmallVector &index_ops, + llvm::SmallVector> &itree_to_rhs_tensors, + llvm::SmallVector>> &itree_to_rhs_dims, + llvm::SmallVector>> &itree_to_rhs_index_idx, +// const std::unordered_map> &itree_to_common_indices, + MLIRContext *context, + mlir::IRRewriter &rewriter, + mlir::Location &loc) +{ + llvm::SmallVector rhs_operands; + llvm::SmallVector &rhs_tensors = itree_to_rhs_tensors[tree_i]; + + indexTree::OperandType operand_type = indexTree::OperandType::get(context); + + for (uint32_t rhs_i = 0; rhs_i < rhs_tensors.size(); ++rhs_i) { + llvm::SmallVector pos; + llvm::SmallVector crds; + Value prev_dim = nullptr; + auto access_type = rewriter.getIndexType(); + llvm::SmallVector &rhs_dims = itree_to_rhs_dims[tree_i][rhs_i]; + llvm::SmallVector &rhs_index_idx = itree_to_rhs_index_idx[tree_i][rhs_i]; + Value rhs_tensor = rhs_tensors[rhs_i]; + uint32_t dim_base = 0; + if (rhs_i == intermediate_idx) { + /// This RHS operand is an intermediate tensor + rhs_tensor = prev_computeOp; +// if (itree_to_common_indices.find(tree_i) != itree_to_common_indices.end()) { +// dim_base = itree_to_common_indices.at(tree_i).size(); +// } + } + for (uint32_t d_i = dim_base; d_i < rhs_dims.size(); ++d_i) { + Value index_node = index_ops[rhs_index_idx[d_i]]; + uint32_t dim = rhs_dims[d_i].first - dim_base; + auto access_op = rewriter.create( + loc, + TypeRange({access_type, access_type}), + rhs_tensor, + index_node, + rewriter.getUI32IntegerAttr(dim), + prev_dim); + pos.push_back(access_op.getPos()); + crds.push_back(access_op.getCrd()); + prev_dim = pos.back(); + } + + rhs_operands.push_back( + rewriter.create(loc, + operand_type, + rhs_tensor, + pos, + crds)); + } + return rhs_operands; +} + + +Value createComputeOp( + uint32_t tree_i, + mlir::Type &tensor_type, + llvm::SmallVector &index_ops, + Value lhs_operand, + llvm::SmallVector &rhs_operands, + const llvm::SmallVector &itree_to_semiring, + const llvm::SmallVector &itree_to_compute_missing, + mlir::IRRewriter &rewriter, + mlir::Location &loc) +{ + Value parent = index_ops.back(); + Value mask_operand = nullptr; + if (rhs_operands.size() == 3) { + mask_operand = rhs_operands.back(); + rhs_operands.pop_back(); + } + mlir::StringRef semiring = itree_to_semiring[tree_i]; + bool compute_missing = itree_to_compute_missing[tree_i]; + Value compute_op = rewriter.create( + loc, + tensor_type, + parent, + lhs_operand, + rhs_operands, + mask_operand, + rewriter.getStringAttr(semiring), + rewriter.getBoolAttr(compute_missing)); + + return compute_op; +} + +[[maybe_unused]] Value createComputeOpReset( + uint32_t tree_i, + uint32_t intermediate_idx, + const mlir::Type &tensor_type, + const mlir::Type &element_type, + const llvm::SmallVector &index_ops, + Value lhs_tensor, + const std::unordered_map> &itree_to_common_indices, + llvm::SmallVector>> &itree_to_rhs_dims, + llvm::SmallVector>> &itree_to_rhs_index_idx, + MLIRContext *context, + mlir::IRRewriter &rewriter, + mlir::Location &loc) +{ + indexTree::OperandType operand_type = indexTree::OperandType::get(context); + /// Create LHS operand + llvm::SmallVector pos; + llvm::SmallVector crds; + Value prev_dim = nullptr; + auto access_type = rewriter.getIndexType(); + llvm::SmallVector &lhs_dims = itree_to_rhs_dims[tree_i][intermediate_idx]; + llvm::SmallVector &lhs_index_idx = itree_to_rhs_index_idx[tree_i][intermediate_idx]; + assert(itree_to_common_indices.find(tree_i) != itree_to_common_indices.end() && "Expect common indices."); + uint32_t dim_base = itree_to_common_indices.at(tree_i).size(); + for (uint32_t d_i = dim_base; d_i < lhs_dims.size(); ++d_i) { + Value index_node = index_ops[lhs_index_idx[d_i]]; + uint32_t dim = lhs_dims[d_i].first - dim_base; + auto access_op = rewriter.create( + loc, + TypeRange({access_type, access_type}), + lhs_tensor, + index_node, + rewriter.getUI32IntegerAttr(dim), + prev_dim); + pos.push_back(access_op.getPos()); + crds.push_back(access_op.getCrd()); + prev_dim = pos.back(); + } + Value lhs_operand = rewriter.create(loc, + operand_type, + lhs_tensor, + pos, + crds); + /// Create RHS operand (constant 0) + mlir::TypedAttr zero = rewriter.getZeroAttr(element_type); + Value cst_0 = rewriter.create(loc, + element_type, + zero); + Value rhs_operand = rewriter.create(loc, + operand_type, + /*rhs_tensor*/cst_0, + /*pos*/ValueRange{}, + /*crds*/ValueRange{}); + + /// Create Compute Op for resetting + Value parent; + if (dim_base < lhs_dims.size()) { + /// If the intermediate tensor is still a tensor, link to the last remaining dimension (index) + /// For example, lhs_index_idx could be {1, 0}, but the innermost index should be 1. Thus we need to sort it at first. + llvm::SmallVector copy(lhs_index_idx); + std::sort(copy.begin(), copy.end()); + parent = index_ops[copy.back()]; + } else { + /// If the intermediate tensor is a scalar, link to the last common index + parent = index_ops[itree_to_common_indices.at(tree_i).back()]; + } + mlir::StringRef semiring("noop_times"); + bool compute_missing = false; + Value compute_op = rewriter.create( + loc, + tensor_type, + parent, + lhs_operand, + ValueRange{rhs_operand}, + /*mask_operand*/nullptr, + rewriter.getStringAttr(semiring), + rewriter.getBoolAttr(compute_missing)); + + return compute_op; +} + + +void fuseITrees(IndexTreeOp new_itree, + uint32_t num_itrees, + const llvm::SmallVector &tree_types, + const llvm::SmallVector &itree_to_num_indexOps, + std::unordered_map> &itree_to_common_indices, + const std::unordered_map &itree_to_intermediate_idx, + llvm::SmallVector> &itree_to_lhs_dims, + llvm::SmallVector> &itree_to_lhs_index_idx, + llvm::SmallVector> &itree_to_rhs_tensors, + llvm::SmallVector>> &itree_to_rhs_dims, + llvm::SmallVector>> &itree_to_rhs_index_idx, + const llvm::SmallVector &itree_to_semiring, + const llvm::SmallVector &itree_to_compute_missing, + llvm::SmallVector &itree_to_new_compute_op, + MLIRContext *context, + mlir::IRRewriter &rewriter, + mlir::Location &loc) +{ + /// Get all IndexOp of the new itree + llvm::SmallVector host_index_ops; + new_itree.walk([&](IndexTreeIndicesOp op) { + host_index_ops.push_back(op); + }); + + + indexTree::YieldOp yield_op = llvm::cast(new_itree.getRegion().getBlocks().front().getTerminator()); + OpBuilder::InsertionGuard guard(rewriter); + rewriter.setInsertionPoint(yield_op); + /// TODO: Parallel execution does not seem to work properly when fusion is in place... + for(auto indexOp : host_index_ops) + { + indexOp.setIsParallel(false); + } + /// Fuse each other itree to the new itree. + for (uint32_t tree_i = 1; tree_i < num_itrees; ++tree_i) { + /// Create Index Ops: 1) Record the common index ops, then 2) add the new index ops. + llvm::SmallVector index_ops = createIndexOps(tree_i, + host_index_ops, + itree_to_num_indexOps, + itree_to_common_indices, + context, + rewriter, + loc); + comet_vdump(new_itree); + + /// Create LHS + /// An itree's arguments are it.itree($inputs, $intermediates) +// Value lhs_tensor = new_itree.getRegion().getBlocks().front().getArgument(tree_i); + Value lhs_tensor; + if (tree_i == num_itrees - 1) { + /// It is the final itree, generating the output + lhs_tensor = new_itree.getRegion().getBlocks().front().getArgument(0); + } else { + /// It is to generate an intermediate + uint32_t num_inputs = llvm::SmallVector(new_itree.getInputs()).size(); + lhs_tensor = new_itree.getRegion().getBlocks().front().getArgument(num_inputs + tree_i); + } + Value lhs_operand = createLHSOperand(tree_i, + index_ops, + lhs_tensor, + itree_to_lhs_dims, + itree_to_lhs_index_idx, +// itree_to_common_indices, + context, + rewriter, + loc); + comet_vdump(new_itree); + + /// Create RHS + /// After fusion, one rhs operand should come from the intermediate variable. + Value prev_computeOp = itree_to_new_compute_op[tree_i - 1]; /// the previous ComputeOp + uint32_t intermediate_idx = itree_to_intermediate_idx.at(tree_i); /// Which rhs operand is from the previous lhs operand + llvm::SmallVector rhs_operands = createRHSOperands(tree_i, + intermediate_idx, + prev_computeOp, + index_ops, + itree_to_rhs_tensors, + itree_to_rhs_dims, + itree_to_rhs_index_idx, +// itree_to_common_indices, + context, + rewriter, + loc); + comet_vdump(new_itree); + + /// Create Compute Ops + auto tensor_type = tree_types[tree_i]; + Value compute_op = createComputeOp(tree_i, + tensor_type, + index_ops, + lhs_operand, + rhs_operands, + itree_to_semiring, + itree_to_compute_missing, + rewriter, + loc); + itree_to_new_compute_op.push_back(compute_op); + comet_vdump(new_itree); + comet_debug() << "\n"; + +// /// Create Compute Op to reset the intermediate tensor to 0 +// /// The element type of the intermediate tensor. +// mlir::Type element_type = +// llvm::cast(itree_to_rhs_tensors[tree_i][intermediate_idx].getType()).getElementType(); +// Value compute_op_reset = createComputeOpReset( +// tree_i, +// intermediate_idx, +// /*tensor_type=*/tree_types[tree_i - 1], +// /*element_type=*/element_type, +// index_ops, +// /*lhs_tensor=*/itree_to_new_compute_op[tree_i - 1], +// itree_to_common_indices, +// itree_to_rhs_dims, +// itree_to_rhs_index_idx, +// context, +// rewriter, +// loc); +// itree_to_new_compute_op[tree_i - 1] = compute_op_reset; /// Will be operands of the YieldOp. + } + + /// Create yield op + assert(itree_to_new_compute_op.size() == num_itrees && "Expect each itree has one ComputeOp."); + llvm::SmallVector yieldOpArgs; + /// Inputs go first + yieldOpArgs.push_back(itree_to_new_compute_op.back()); + /// Then are intemediates + yieldOpArgs.insert(yieldOpArgs.end(), itree_to_new_compute_op.begin(), itree_to_new_compute_op.begin() + num_itrees - 1); + + rewriter.create(loc, TypeRange(), yieldOpArgs); + rewriter.eraseOp(yield_op); + comet_vdump(new_itree); +} + + +void createITree( + uint32_t num_itrees, +// const llvm::SmallVector &itree_arguments, + const llvm::SmallVector &itree_arguments_inputs, + const llvm::SmallVector &itree_arguments_intermediates, + llvm::SmallVector &itree_list, + const llvm::SmallVector &itree_to_num_indexOps, + std::unordered_map> &itree_to_common_indices, + const std::unordered_map &itree_to_intermediate_idx, + llvm::SmallVector> &itree_to_lhs_dims, + llvm::SmallVector> &itree_to_lhs_index_idx, + llvm::SmallVector> &itree_to_rhs_tensors, + llvm::SmallVector>> &itree_to_rhs_dims, + llvm::SmallVector>> &itree_to_rhs_index_idx, + const llvm::SmallVector &itree_to_semiring, + const llvm::SmallVector &itree_to_compute_missing) +{ + /// The 1st itree is the host. + /// 1) Its region will be moved to the new itree. + /// 2) Each other itree will be fused to the new itree. + IndexTreeOp host_itree = itree_list.front(); + mlir::OpBuilder builder(host_itree); + mlir::IRRewriter rewriter(builder); + MLIRContext *context = rewriter.getContext(); + llvm::SmallVector tree_types; +// for (Value arg : itree_arguments) { +// tree_types.push_back(arg.getType()); +// } + for (Value arg : itree_arguments_inputs) { + tree_types.push_back(arg.getType()); + } + for (Value arg : itree_arguments_intermediates) { + tree_types.push_back(arg.getType()); + } + + auto loc = host_itree->getLoc(); + + /// Create the new itree. + llvm::SmallVector itree_to_new_compute_op; + IndexTreeOp new_itree = createNewITree(host_itree, + tree_types, +// itree_arguments, + itree_arguments_inputs, + itree_arguments_intermediates, + itree_to_lhs_dims, + itree_to_lhs_index_idx, +// itree_to_common_indices, + itree_to_new_compute_op /*out*/, + context, + rewriter, + loc); + + /// Fuse other itrees to the new itree + fuseITrees(new_itree, + num_itrees, + tree_types, + itree_to_num_indexOps, + itree_to_common_indices, + itree_to_intermediate_idx, + itree_to_lhs_dims, + itree_to_lhs_index_idx, + itree_to_rhs_tensors, + itree_to_rhs_dims, + itree_to_rhs_index_idx, + itree_to_semiring, + itree_to_compute_missing, + itree_to_new_compute_op, + context, + rewriter, + loc); + + /// Update the usage of results of itree + /// Erase uses of output of kernels other than the last kernel + + for (uint32_t tree_i = 0; tree_i < num_itrees - 1; ++tree_i) { + for (auto result : itree_list[tree_i].getResults()) { + for (auto user : result.getUsers()) { + if (auto set_op = dyn_cast(user)) { + rewriter.eraseOp(set_op); + } + } + } + } + /// Replace the use of output of the last kernel + uint32_t new_r_i = 0; + uint32_t tree_i = num_itrees - 1; + for (auto old_result : itree_list[tree_i].getResults()) { + rewriter.replaceAllUsesWith(old_result, new_itree.getResult(new_r_i++)); + } + assert(new_r_i == 1 && "Expect to output only one tensor"); + + /// Delete the old itrees + for (auto itree : itree_list) { + rewriter.eraseOp(itree); + } +} + +} /// anonymous namespace + + + + +struct IndexTreeKernelFusion : comet::impl::IndexTreeKernelFusionBase +{ + using IndexTreeKernelFusionBase::IndexTreeKernelFusionBase; + void runOnOperation() override; +}; + +void IndexTreeKernelFusion::runOnOperation() +{ + comet_debug() << "IndexTreeKernelFusion::runOnOperation()\n"; + func::FuncOp funcOp = getOperation(); + comet_vdump(funcOp->getParentOfType()); + + /// Collect all itrees. + llvm::SmallVector itree_list; + funcOp.walk([&](IndexTreeOp itree){ + comet_vdump(itree); + itree_list.push_back(itree); + }); + + uint32_t num_itrees = itree_list.size(); + if (num_itrees < 2) { + /// Need at least two itrees to fuse. + return; + } + + /// Find the overlap of IndicesOp; + /// Each IndicesOp corresponds to a set of Dimension Compound . + /// For example, + /// %14 = "it.IndexOp"(%13) : (!it.index_tree) -> !it.index + /// %crd_0, %pos_1 = "it.IndexToTensorDim"(%10, %14, %pos) <{dim = 1 : ui32}> : (tensor, !it.index, index) -> (index, index) + /// %crd_8, %pos_9 = "it.IndexToTensorDim"(%7, %14, %pos_7) <{dim = 1 : ui32}> : (tensor, !it.index, index) -> (index, index) + /// Here %14 corresponds to <%10, 1> and <%7, 1>, meaning that %14 is the tensor %10's dim 1 and tensor %7's dim 1. + /// The set of Dimension Compound will be used to tell if two IndexOps are the common index in two itrees. + llvm::SmallVector itree_to_num_indexOps(num_itrees); + std::unordered_map> itree_to_common_indices; + collectCommonIndices(itree_list, + itree_to_num_indexOps /*out*/, + itree_to_common_indices /*out*/); + bool has_common_indices = false; + for (uint32_t tree_i = 1; tree_i < num_itrees; ++tree_i) { + if (!itree_to_common_indices[tree_i].empty()) { + has_common_indices = true; + break; + } + } + if (!has_common_indices) { + /// No need to fuse if itrees don't have any common indices. + return; + } + +// /// Collect all itrees' arguments +// llvm::SmallVector itree_arguments; +// for (IndexTreeOp &itree : itree_list) { +// for (Value arg : itree.getInputs()) { +// comet_vdump(arg); +// itree_arguments.push_back(arg); +// } +// } + /// Collect all itrees' arguments. The last itree's inputs will be the final `inputs`. Other itrees' inputs will be the + /// final `intermediates`. The new fused itree's arguments will be it.itree($inputs, $intermediates). + llvm::SmallVector itree_arguments_inputs; + llvm::SmallVector itree_arguments_intermediates; + for (uint32_t tree_i = 0; tree_i < num_itrees - 1; ++tree_i) { + IndexTreeOp &itree = itree_list[tree_i]; + for (Value arg : itree.getInputs()) { + itree_arguments_intermediates.push_back(arg); + } + } + for (Value arg : itree_list.back().getInputs()) { + itree_arguments_inputs.push_back(arg); + } + + + /// Collect all LHS and RHS operands + + /// `itree_to_lhs_dims`: which is used as a LHSOperandOp's dimensions. + /// `itree_to_lhs_index_idx`: which IndexOp is linked by a LHSOperandOp's dimensions. IndexOp are ordered starting from 0. + /// For example, + /* + %11 = "it.itree"(%10) ({ + ^bb0(%arg0: tensor): + %13 = "it.RootOp"() : () -> !it.index_tree + %14 = "it.IndexOp"(%13) : (!it.index_tree) -> !it.index /// %14 = h + %15 = "it.IndexOp"(%14) : (!it.index) -> !it.index /// %15 = i + %16 = "it.IndexOp"(%15) : (!it.index) -> !it.index /// %16 = k + %crd, %pos = "it.IndexToTensorDim"(%arg0, %15) <{dim = 0 : ui32}> : (tensor, !it.index) -> (index, index) + %crd_0, %pos_1 = "it.IndexToTensorDim"(%arg0, %14, %pos) <{dim = 1 : ui32}> : (tensor, !it.index, index) -> (index, index) + %17 = "it.LHSOperandOp"(%arg0, %pos, %pos_1, %crd, %crd_0) : (tensor, index, index, index, index) -> !it.operand + %crd_2, %pos_3 = "it.IndexToTensorDim"(%4, %15) <{dim = 0 : ui32}> : (!ta.sparse_tensor, !it.index) -> (index, index) + + // ... + }) : (tensor) -> tensor + */ + /// IndexOps %14, %15, and %16 are referred as 0, 1, and 2. + /// itree_to_lhs_dims[0]: {<0, %10>, <1, %10>}. %10 is the real tensor, not %arg0. + /// itree_to_lhs_index_idx[0]: {1, 0}. because %15 is the 1st IndexOp, and %14 is the 0th IndexOp. + /// the rhs counterparts have the same behavior. + llvm::SmallVector itree_to_lhs_tensors(num_itrees); + llvm::SmallVector> itree_to_lhs_dims(num_itrees); + llvm::SmallVector> itree_to_lhs_index_idx(num_itrees); + llvm::SmallVector> itree_to_rhs_tensors(num_itrees); + llvm::SmallVector>> itree_to_rhs_dims(num_itrees); + llvm::SmallVector>> itree_to_rhs_index_idx(num_itrees); + for (uint32_t tree_i = 0; tree_i < num_itrees; ++tree_i) { + collectOperandsDims(itree_list[tree_i], + itree_to_lhs_tensors[tree_i] /*out*/, + itree_to_lhs_dims[tree_i] /*out*/, + itree_to_lhs_index_idx[tree_i] /*out*/, + itree_to_rhs_tensors[tree_i] /*out*/, + itree_to_rhs_dims[tree_i] /*out*/, + itree_to_rhs_index_idx[tree_i] /*out*/); + } + + /// Record which rhs operand is from the previous lhs operand + /// For example, + /* + T[i, h] += B[i, k] * C[k, h]; // the 0th itree + A[i, j] += T[i, h] * D[h, j]; // the 1st itree + */ + /// itree_to_intermediate_idx[1] = 0, because the 1st itree has T as its 0st rhs operand that is the previous lhs operand. + std::unordered_map itree_to_intermediate_idx; + for (uint32_t tree_i = 1; tree_i < num_itrees; ++tree_i) { + itree_to_intermediate_idx[tree_i] = collectIntermediateIdx(itree_to_lhs_tensors[tree_i - 1] /*previous lhs*/, + itree_to_rhs_tensors[tree_i] /*current rhs*/); + } + +// /// Create new lhs tensors with decreased dimensions +// std::unordered_map itree_to_new_lhs_tensors = +// createNewLhsTensors(num_itrees, +// itree_to_lhs_tensors, +// itree_to_common_indices); +// /// Update itrees' arguments +// for (uint32_t tree_i = 0; tree_i < num_itrees - 1; ++tree_i) { +// if (itree_to_new_lhs_tensors.find(tree_i) != itree_to_new_lhs_tensors.end()) { +// itree_arguments[tree_i] = itree_to_new_lhs_tensors[tree_i]; +// } +// } + + /// Collect ComputeOp's information: semiring, and compute_missing. + llvm::SmallVector itree_to_semiring(num_itrees); + llvm::SmallVector itree_to_compute_missing(num_itrees, false); + for (uint32_t tree_i = 0; tree_i < num_itrees; ++tree_i) { + collectComputeOpInfo(itree_list[tree_i], + itree_to_semiring[tree_i] /*out*/, + itree_to_compute_missing[tree_i] /*out*/); + } + + /// Build the new itree + createITree(num_itrees, +// itree_arguments, + itree_arguments_inputs, + itree_arguments_intermediates, + itree_list, + itree_to_num_indexOps, + itree_to_common_indices, + itree_to_intermediate_idx, + itree_to_lhs_dims, + itree_to_lhs_index_idx, + itree_to_rhs_tensors, + itree_to_rhs_dims, + itree_to_rhs_index_idx, + itree_to_semiring, + itree_to_compute_missing); + +// /// Remove old lhs tensors +// for (uint32_t tree_i = 0; tree_i < num_itrees - 1; ++tree_i) { +// if (itree_to_new_lhs_tensors.find(tree_i) != itree_to_new_lhs_tensors.end()) { +// Value old_tensor = itree_to_lhs_tensors[tree_i]; +// for (auto user : old_tensor.getUsers()) { +// user->erase(); +// } +// old_tensor.getDefiningOp()->erase(); +// } +// } + + comet_vdump(funcOp->getParentOfType()); +} + + +/// Apply the redundancy-aware kernel fusion on index tree dialect for some compound expressions +std::unique_ptr mlir::comet::createIndexTreeKernelFusionPass() +{ + return std::make_unique(); +} \ No newline at end of file diff --git a/lib/Dialect/IndexTree/Transforms/MaskDomainPatterns.cpp b/lib/Dialect/IndexTree/Transforms/MaskDomainPatterns.cpp new file mode 100644 index 00000000..4068c249 --- /dev/null +++ b/lib/Dialect/IndexTree/Transforms/MaskDomainPatterns.cpp @@ -0,0 +1,254 @@ +// +// Copyright 2022 Battelle Memorial Institute +// +// Redistribution and use in source and binary forms, with or without modification, +// are permitted provided that the following conditions are met: +// +// 1. Redistributions of source code must retain the above copyright notice, this list of conditions +// and the following disclaimer. +// +// 2. Redistributions in binary form must reproduce the above copyright notice, this list of conditions +// and the following disclaimer in the documentation and/or other materials provided with the distribution. +// +// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED +// WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR +// PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, +// INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE +// GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, +// WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +// OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. +// + +#include "mlir/Dialect/Func/IR/FuncOps.h" +#include "mlir/IR/BuiltinOps.h" +#include "mlir/IR/BuiltinTypeInterfaces.h" +#include "mlir/IR/Value.h" +#include "mlir/Support/LLVM.h" +#include "mlir/Transforms/GreedyPatternRewriteDriver.h" +#include "mlir/Dialect/MemRef/IR/MemRef.h" +#include "mlir/Dialect/Math/IR/Math.h" +#include "mlir/Dialect/Tensor/IR/Tensor.h" +#include "mlir/Dialect/Index/IR/IndexOps.h" +#include "mlir/Pass/Pass.h" + +#include "llvm/ADT/StringSet.h" +#include "llvm/ADT/SmallVector.h" +#include "llvm/ADT/IndexedMap.h" +#include "llvm/ADT/TypeSwitch.h" +#include "llvm/Support/ScopedPrinter.h" + +#include "comet/Dialect/IndexTree/IR/IndexTreeDialect.h" +#include "comet/Dialect/TensorAlgebra/IR/TADialect.h" +#include "comet/Dialect/IndexTree/Passes.h" +#include "comet/Dialect/IndexTree/Patterns.h" + +#define DEBUG_TYPE "mask-domain" + +using namespace mlir; +using namespace mlir::indexTree; +using namespace mlir::tensorAlgebra; + +struct MoveInvariantMaskOp : public mlir::OpRewritePattern { + MoveInvariantMaskOp(mlir::MLIRContext *context) + : OpRewritePattern(context, /*benefit=*/1) {} + + mlir::LogicalResult + liftAccessOp(mlir::Operation *dependent_op, + IndexTreeIndexToTensorOp access_op, + mlir::PatternRewriter &rewriter ) const { + if(access_op->isBeforeInBlock(dependent_op)) + return success(); + + Value prev_access_value; + if((prev_access_value = access_op.getPrevDim())) + { + if(mlir::failed( + liftAccessOp( + dependent_op, + llvm::cast(prev_access_value.getDefiningOp()), + rewriter + ) + )) + return failure(); + } + rewriter.modifyOpInPlace(access_op, [&]() {access_op->moveBefore(dependent_op);}); + return success(); + } + + mlir::LogicalResult + matchAndRewrite(IndexTreeFillMaskOp op, + mlir::PatternRewriter &rewriter) const override + { + if(!llvm::isa(op.getDomain().getDefiningOp())) { + return failure(); + } + + llvm::ScopedPrinter logger{llvm::dbgs()}; + LLVM_DEBUG({ + Operation* domain_op = op.getDomain().getDefiningOp(); + logger.getOStream() << "\n"; + logger.startLine() << "Looking at : '" << domain_op->getName() << "'(" + << op << ") {\n"; + logger.indent(); + + // If the operation has no regions, just print it here. + logger.startLine() << "Concrete domain : '" << domain_op->getName() << "'("<< domain_op << ")\n"; + + }); + + // We want to find all the indices necessary to compute the mask + llvm::SmallDenseSet used_indices; + llvm::SmallVector work_list; + work_list.push_back(op.getDomain()); + while(!work_list.empty()){ + Value domain = work_list.back(); + work_list.pop_back(); + + llvm::TypeSwitch(domain.getDefiningOp()) + .Case([&](IndexTreeDenseDomainOp op) { + return; + }) + .Case([&](IndexTreeSparseDomainOp op) { + auto access = op.getParent().getDefiningOp(); + used_indices.insert(access.getIndex()); + return; + }) + .Case([&](IndexTreeWorkspaceDomainOp op) { + return; + }) + .Case([&](IndexTreeDomainIntersectionOp op) { + work_list.append(op.getDomains().begin(), op.getDomains().end()); + return; + }) + .Case([&](IndexTreeDomainUnionOp op) { + work_list.append(op.getDomains().begin(), op.getDomains().end()); + return; + }) + .Default([](Operation *op) { + assert(false && "IndexNode not given a valid domain"); + return nullptr; + }); + } + + LLVM_DEBUG({ + if(used_indices.size() == 0) { + op->getParentOp()->print(logger.startLine()); + } + + }); + + // We then move the fill mask op into the highest loop possible + Value parent = op.getParent(); + while(parent != nullptr) { + if(used_indices.contains(parent)){ + break; + } + parent = parent.getDefiningOp().getParent(); + } + assert(parent != nullptr && "Could not find proper nesting for creating mask."); + if(parent == op.getParent()) { + return failure(); + } + + // Match success! + // We then find the zero op and move it to the same nesting so that the logic remains correct. + Value previous_parent = op.getParent(); + LLVM_DEBUG({ + op->getParentOp()->print(logger.startLine()); + logger.startLine() << "\n"; + }); + + rewriter.modifyOpInPlace(op, [&](){ + Value cur_parent = op.getParent(); + while(cur_parent != parent) { + op->moveBefore(cur_parent.getDefiningOp()); + cur_parent = llvm::cast(cur_parent.getDefiningOp()).getParent(); + } + op.getParentMutable().assign(parent); + }); + if(!op.getDomain().getDefiningOp()->isBeforeInBlock(op)){ + Operation* domain = op.getDomain().getDefiningOp(); + rewriter.modifyOpInPlace(domain, [&](){domain->moveBefore(op);}); + IndexTreeSparseDomainOp sparse_domain; + if((sparse_domain = llvm::dyn_cast(domain))){ + if (mlir::failed(liftAccessOp(sparse_domain, llvm::cast(sparse_domain.getParent().getDefiningOp()), rewriter))) { + return failure(); + } + } + } + + IndexTreeZeroMaskOp zero_op = nullptr; + for(auto user : previous_parent.getUsers()) { + if((zero_op = llvm::dyn_cast(user)) && zero_op.getInit() == op.getResult()) { + rewriter.modifyOpInPlace(zero_op, [&](){zero_op.getParentMutable().assign(parent);}); + } + } + return success(); + } +}; + +struct CreateFillMaskOp : public mlir::OpRewritePattern { + CreateFillMaskOp(mlir::MLIRContext *context) + : OpRewritePattern(context, /*benefit=*/0) {} + + mlir::LogicalResult + matchAndRewrite(IndexTreeMaskedDomainOp op, + mlir::PatternRewriter &rewriter) const override + { + if(!llvm::isa(op.getMask().getType())) { + return failure(); + } + + if(!llvm::isa(op.getBase().getDefiningOp())) { + return failure(); + } + + IndexTreeOp tree_op = op->getParentOfType(); + if(!tree_op) { + return failure(); + } + + auto mask_domain = op.getMask().getDefiningOp(); + if(!mask_domain) { + return failure(); + } + + auto loc = op.getLoc(); + + // Create bit tensor outside of index tree + auto cur = rewriter.saveInsertionPoint(); + rewriter.setInsertionPoint(tree_op); + auto bit_tensor_type = RankedTensorType::get({ShapedType::kDynamic}, rewriter.getI1Type()); + Value f = rewriter.create(loc, rewriter.getI1Type(), rewriter.getBoolAttr(false)); + Value init_bit_tensor = rewriter.create(loc, bit_tensor_type, f, mask_domain.getDimensionSize()); + SmallVector tree_temps(tree_op.getIntermediates()); + SmallVector tree_types(tree_op->getResultTypes()); + tree_temps.push_back(init_bit_tensor); + tree_types.push_back(bit_tensor_type); + auto new_op = rewriter.create(loc, tree_types, tree_op.getInputs(), tree_temps); + rewriter.inlineRegionBefore(tree_op.getRegion(), new_op.getRegion(), new_op.getRegion().end()); + + rewriter.restoreInsertionPoint(cur); + rewriter.modifyOpInPlace(new_op, [&](){new_op.getBody()->addArgument(bit_tensor_type, loc);}); + init_bit_tensor = new_op.getBody()->getArgument(new_op.getBody()->getNumArguments() - 1); + + Value domain = op.getMask(); + Value mask_tensor = rewriter.create(loc, bit_tensor_type, op.getParentNode(), domain, init_bit_tensor); + rewriter.modifyOpInPlace(op, [&](){op.getMaskMutable().assign(mask_tensor);}); + indexTree::YieldOp yield = llvm::cast(new_op.getBody()->getTerminator()); + rewriter.setInsertionPoint(yield); + mask_tensor = rewriter.create(loc, bit_tensor_type, op.getParentNode(), domain, mask_tensor); + rewriter.modifyOpInPlace(new_op, [&](){yield.getResultsMutable().append(ValueRange(mask_tensor));}); + + for(unsigned i = 0; i < tree_op.getNumResults(); i++){ + rewriter.replaceAllUsesWith(tree_op.getResult(i), new_op.getResult(i)); + } + rewriter.eraseOp(tree_op); + + return success(); + } +}; + +void indexTree::populateMaskDomainTransformationPatterns(MLIRContext *context, RewritePatternSet &patterns) { + patterns.add(context); +} \ No newline at end of file diff --git a/lib/Dialect/IndexTree/Transforms/SymbolicCompute.cpp b/lib/Dialect/IndexTree/Transforms/SymbolicCompute.cpp new file mode 100644 index 00000000..ebee3cdb --- /dev/null +++ b/lib/Dialect/IndexTree/Transforms/SymbolicCompute.cpp @@ -0,0 +1,361 @@ +// +// Copyright 2022 Battelle Memorial Institute +// +// Redistribution and use in source and binary forms, with or without modification, +// are permitted provided that the following conditions are met: +// +// 1. Redistributions of source code must retain the above copyright notice, this list of conditions +// and the following disclaimer. +// +// 2. Redistributions in binary form must reproduce the above copyright notice, this list of conditions +// and the following disclaimer in the documentation and/or other materials provided with the distribution. +// +// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED +// WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR +// PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, +// INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE +// GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, +// WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +// OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. +// + +#include "mlir/Dialect/Func/IR/FuncOps.h" +#include "mlir/IR/BuiltinAttributes.h" +#include "mlir/Transforms/GreedyPatternRewriteDriver.h" +#include "mlir/Dialect/MemRef/IR/MemRef.h" +#include "mlir/Dialect/Math/IR/Math.h" +#include "mlir/Dialect/Tensor/IR/Tensor.h" +#include "mlir/Dialect/Index/IR/IndexOps.h" +#include "mlir/Pass/Pass.h" + +#include "llvm/ADT/StringSet.h" +#include "llvm/ADT/SmallVector.h" +#include "llvm/ADT/IndexedMap.h" +#include +#include + +#include "comet/Dialect/IndexTree/IR/IndexTreeDialect.h" +#include "comet/Dialect/IndexTree/Patterns.h" +#include "comet/Dialect/TensorAlgebra/IR/TADialect.h" +#include "comet/Dialect/IndexTree/Passes.h" + +using namespace mlir; +using namespace mlir::indexTree; +using namespace mlir::tensorAlgebra; + +namespace mlir { + namespace comet{ + #define GEN_PASS_DEF_INDEXTREESYMBOLICCOMPUTEPASS + #include "comet/Dialect/IndexTree/Passes.h.inc" + } +} + +struct CreateSymbolicTree : public OpRewritePattern { + CreateSymbolicTree(MLIRContext *context) + : OpRewritePattern(context, /*benefit=*/1) {} + + Value copyDomain(Value domain, mlir::PatternRewriter &rewriter, + Location loc, IRMapping& map, + llvm::SmallDenseMap, Value>& tensor_to_node, + Value* parent, + Value predecessor = nullptr) const + { + Value new_domain; + Operation* domain_op = domain.getDefiningOp(); + if(llvm::isa(domain_op)) + { + auto intersection_domain_op = llvm::cast(domain_op); + for(Value subdomain : intersection_domain_op.getDomains()){ + copyDomain(subdomain, rewriter, loc, map, tensor_to_node, parent, predecessor); + } + } + else if(llvm::isa(domain_op)) + { + auto union_domain_op = llvm::cast(domain_op); + for(Value subdomain : union_domain_op.getDomains()){ + copyDomain(subdomain, rewriter, loc, map, tensor_to_node, parent, predecessor); + } + } + else if(llvm::isa(domain_op)) + { + auto masked_domain_op = llvm::cast(domain_op); + copyDomain(masked_domain_op.getMask(), rewriter, loc, map, tensor_to_node, parent, predecessor); + copyDomain(masked_domain_op.getBase(), rewriter, loc, map, tensor_to_node, parent, predecessor); + } + + if(llvm::isa(domain_op)){ + auto nested_domain = llvm::cast(domain_op); + for(auto subdomain_itr = nested_domain.getDomains().begin(); subdomain_itr != nested_domain.getDomains().end();) + { + Value subdomain = *subdomain_itr; + new_domain = copyDomain(subdomain, rewriter, loc, map, tensor_to_node, parent, predecessor); + subdomain_itr++; + if(subdomain_itr != nested_domain.getDomains().end()) + { + auto index_node_type = indexTree::IndexNodeType::get(rewriter.getContext()); + indexTree::IndexTreeIndicesOp index_node_op = rewriter.create(loc, index_node_type, *parent, new_domain, false, nullptr); + createMapping(index_node_op, subdomain, tensor_to_node); + *parent = index_node_op.getOutput(); + predecessor = *parent; + } else { + map.map(nested_domain, new_domain); + } + } + } else if(llvm::isa(domain_op)) + { + auto sparse_domain_op = llvm::cast(domain_op); + auto tensor = sparse_domain_op.getTensor(); + int32_t dim = sparse_domain_op.getDim(); + Value parent = nullptr; + if(dim > 0){ + // TODO: Determine parent of this op + // Will be needed for 3 dimensional sparse tensor outputs + if(!predecessor) + { + assert(tensor_to_node[std::make_pair(tensor, dim-1)] != nullptr); + predecessor = tensor_to_node[std::make_pair(tensor, dim-1)]; + } + + + auto tensor_access_op = rewriter.create( + loc, + TypeRange({rewriter.getIndexType(), rewriter.getIndexType()}), + tensor, + predecessor, + dim-1, + nullptr); + parent = tensor_access_op.getPos(); + } + new_domain = rewriter.create(loc, + domain_op->getResultTypes(), + sparse_domain_op.getTensor(), + sparse_domain_op.getDimAttr(), + sparse_domain_op.getFormatAttr(), + sparse_domain_op.getPos(), + sparse_domain_op.getCrd(), + sparse_domain_op.getPosSize(), + sparse_domain_op.getCrdSize(), + sparse_domain_op.getDimSize(), + parent); + map.map(sparse_domain_op, new_domain); + } else { + // Clone + new_domain = rewriter.clone(*domain_op, map)->getResult(0); + } + return new_domain; + } + + void createMapping(IndexTreeIndicesOp node, Value domain, llvm::SmallDenseMap, Value>& tensor_to_node) const + { + Operation* domain_op = domain.getDefiningOp(); + if(llvm::isa(domain_op)) + { + auto intersection_domain_op = llvm::cast(domain_op); + for(Value subdomain : intersection_domain_op.getDomains()){ + createMapping(node, subdomain, tensor_to_node); + } + } else if(llvm::isa(domain_op)) + { + auto union_domain_op = llvm::cast(domain_op); + for(Value subdomain : union_domain_op.getDomains()){ + createMapping(node, subdomain, tensor_to_node); + } + } else if(llvm::isa(domain_op)) + { + auto masked_domain_op = llvm::cast(domain_op); + createMapping(node, masked_domain_op.getMask(), tensor_to_node); + createMapping(node, masked_domain_op.getBase(), tensor_to_node); + } else if (llvm::isa(domain_op)) + { + auto nested_domain_op = llvm::cast(domain_op); + for(Value subdomain : nested_domain_op.getDomains()){ + createMapping(node, subdomain, tensor_to_node); + } + } + else if(llvm::isa(domain_op)) + { + auto sparse_domain_op = llvm::cast(domain_op); + auto tensor = sparse_domain_op.getTensor(); + int32_t dim = sparse_domain_op.getDim(); + tensor_to_node.insert(std::make_pair( + std::make_pair(tensor, dim), + node.getOutput() + )); + } else if(llvm::isa(domain_op)) + { + auto dense_domain_op = llvm::cast(domain_op); + auto tensors = dense_domain_op.getTensors(); + auto dims = dense_domain_op.getDims(); + unsigned i = 0; + for(auto tensor : tensors) + { + int32_t dim = cast(dims[i]).getValue().getSExtValue(); + tensor_to_node.insert(std::make_pair( + std::make_pair(tensor, dim), + node.getOutput() + )); + i += 1; + } + } + } + + mlir::LogicalResult + match(IndexTreeSparseTensorOp op) const override { + for(auto domain : op.getDomains()) + { + Operation* domain_op = domain.getDefiningOp(); + if(domain_op->hasTrait()) + { + return success(); + } + } + return failure(); + } + + void + rewrite(IndexTreeSparseTensorOp it_tensor_decl_op, + mlir::PatternRewriter &rewriter) const override { + auto loc = it_tensor_decl_op->getLoc(); + auto context = rewriter.getContext(); + + // Declare Sparse Domains and allocate position vectors for each dimension + unsigned indicesBitwidth = cast(it_tensor_decl_op->getResultTypes()[0]).getIndicesType().getWidth(); + llvm::SmallDenseMap symbolic_domains; + + llvm::SmallVector input_domains; + + auto domain_type = SymbolicDomainType::get(context, indicesBitwidth); + auto index_type = rewriter.getIndexType(); + Value cur_pos_size = nullptr; + unsigned dim = 0; + for(Value domain : it_tensor_decl_op.getDomains()){ + Operation* domain_op = domain.getDefiningOp(); + + if(domain_op->hasTrait()) + { + if(dim == 0) + { + // If this is the first dimension, the previous dimension could be expanded as "1" + cur_pos_size = rewriter.create(loc, rewriter.getI32Type(), rewriter.getI32IntegerAttr(1)); + } + + Value num_rows = nullptr; + BoolAttr is_dynamic = rewriter.getBoolAttr(0); + if(cur_pos_size == nullptr){ + is_dynamic = rewriter.getBoolAttr(1); + num_rows = rewriter.create(loc, index_type, rewriter.getIndexAttr(0)); + } else { + num_rows = cur_pos_size; + } + auto concrete_domain = llvm::cast(domain_op); + Value dim_size = concrete_domain.getDimensionSize(); + Value symbolic_domain = rewriter.create(loc, domain_type, dim_size, num_rows, is_dynamic, rewriter.getI32IntegerAttr(indicesBitwidth)); + symbolic_domains.insert(std::make_pair(domain, input_domains.size())); + input_domains.push_back(symbolic_domain); + } + + if(llvm::isa(domain_op)) + { + Value dim_size = llvm::cast(domain_op).getDimSize(); + if(dim != 0) + cur_pos_size = rewriter.create(loc, rewriter.getI32Type(), cur_pos_size, dim_size); + else + cur_pos_size = dim_size; + } else + { + cur_pos_size = nullptr; + } + dim += 1; + } + + auto itree_op = rewriter.create(loc, llvm::SmallVector(symbolic_domains.size(), domain_type), input_domains, ValueRange()); + Region* body = &itree_op.getRegion(); + loc = body->getLoc(); + Block* block = rewriter.createBlock(body, {}, llvm::SmallVector(symbolic_domains.size(), domain_type), llvm::SmallVector(symbolic_domains.size(), loc)); + rewriter.setInsertionPointToStart(block); + + indexTree::IndexTreeType tree_type = indexTree::IndexTreeType::get(context); + Value parent = rewriter.create(loc, tree_type); + indexTree::IndexNodeType index_node_type = indexTree::IndexNodeType::get(context); + IRMapping map; + llvm::SmallDenseMap, Value> tensor_to_node; + SmallVector yield_args; + Value prev_dim = parent; + bool is_unique = true; + for (Value domain : it_tensor_decl_op.getDomains()) + { + Operation* domain_op = domain.getDefiningOp(); + Value prev_parent = parent; + Value new_domain = copyDomain(domain, rewriter, loc, map, tensor_to_node, &parent); + indexTree::IndexTreeIndicesOp index_node_op = rewriter.create(loc, index_node_type, parent, new_domain, false, nullptr); + createMapping(index_node_op, domain, tensor_to_node); + if(prev_parent != parent) + { + is_unique = false; + } + parent = index_node_op.getOutput(); + + if(domain_op->hasTrait()) + { + Value symbolic_domain = block->getArgument(symbolic_domains[domain]); + Value new_symbolic_domain = rewriter.create( + loc, + domain_type, + parent, + symbolic_domain, + rewriter.getBoolAttr(is_unique) + ); + new_symbolic_domain = rewriter.create( + loc, + domain_type, + prev_dim, + new_symbolic_domain, + rewriter.getBoolAttr(!is_unique) + ); + yield_args.push_back(new_symbolic_domain); + } + + prev_dim = parent; + } + rewriter.create(loc, TypeRange(), yield_args); + + rewriter.setInsertionPointAfter(itree_op); + SmallVector args; + unsigned i = 0; + for (Value domain : it_tensor_decl_op.getDomains()) + { + if(domain.getDefiningOp()->hasTrait()) + { + args.push_back(itree_op->getResult(i)); + i += 1; + } else + { + args.push_back(domain); + } + } + auto new_tensor = rewriter.create(loc, it_tensor_decl_op->getResultTypes(), args); + rewriter.replaceOp(it_tensor_decl_op, new_tensor->getResults()); + return; + } +}; + +struct IndexTreeSymbolicComputePass : comet::impl::IndexTreeSymbolicComputePassBase { + using IndexTreeSymbolicComputePassBase::IndexTreeSymbolicComputePassBase; + + void runOnOperation() override { + mlir::RewritePatternSet sp_output_patterns(&getContext()); + sp_output_patterns.add(&getContext()); + indexTree::populateMaskDomainTransformationPatterns(&getContext(), sp_output_patterns); + if(failed(mlir::applyPatternsAndFoldGreedily(getOperation(), std::move(sp_output_patterns)))) + { + signalPassFailure(); + } + + } +}; + +/// Apply the compressed workspace transformations on the index tree IR +std::unique_ptr mlir::comet::createIndexTreeSymbolicComputePass() +{ + return std::make_unique(); +} \ No newline at end of file diff --git a/lib/Dialect/IndexTree/Transforms/WorkspaceTransforms.cpp b/lib/Dialect/IndexTree/Transforms/WorkspaceTransforms.cpp index 38505d29..cd44bf0d 100644 --- a/lib/Dialect/IndexTree/Transforms/WorkspaceTransforms.cpp +++ b/lib/Dialect/IndexTree/Transforms/WorkspaceTransforms.cpp @@ -26,6 +26,7 @@ #include "comet/Dialect/IndexTree/IR/IndexTreeDialect.h" #include "comet/Dialect/IndexTree/Passes.h" +#include "comet/Dialect/IndexTree/Patterns.h" #include "comet/Dialect/TensorAlgebra/IR/TADialect.h" #include "comet/Dialect/Utils/Utils.h" @@ -33,7 +34,8 @@ #include "mlir/Dialect/Bufferization/IR/Bufferization.h" #include "mlir/Dialect/Func/IR/FuncOps.h" #include "mlir/Dialect/MemRef/IR/MemRef.h" -#include "mlir/Dialect/Tensor/IR/Tensor.h" +#include "mlir/Transforms/GreedyPatternRewriteDriver.h" + #include "llvm/Support/Debug.h" #include @@ -47,6 +49,7 @@ #include #include #include +#include using namespace mlir; using namespace mlir::bufferization; @@ -65,14 +68,7 @@ using llvm::StringRef; #undef COMET_DEBUG_MODE // *********** For debug purpose *********// -const bool compressedworkspace = true; -struct dimInTensor -{ - int dim; - int tensorId; - int dimOrder; -}; /// Apply workspace transformation on the lhs /// Consider CSR first @@ -90,1010 +86,296 @@ struct dimInTensor /// Apply workspace transformations on the ta.tc and tc.elews_mul namespace { - struct WorkspaceTransformsPass - : public PassWrapper> - { - MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(WorkspaceTransformsPass) - void runOnOperation() override; - void WorkspaceTransforms(mlir::func::FuncOp function); - }; - struct IndexTreeWorkspaceTransformationsPass : public PassWrapper> { MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(IndexTreeWorkspaceTransformationsPass) void runOnOperation() override; - void CompressedWorkspaceTransforms(mlir::func::FuncOp function); }; - } /// end anonymous namespace. -/// Need a function, dfs traverse the itree -/// get the sparse index that is sparse in the output -std::vector getSparseDimsOutput(std::vector> opFormats, std::vector> opPerms) -{ - std::vector sparseDimsOutput; - assert(opFormats.size() > 0 && "opFormats.size() less than 0\n"); - std::vector outputFormat = opFormats[opFormats.size() - 1]; - std::vector outputPerm = opPerms[opPerms.size() - 1]; - for (unsigned int i = 0; i < outputFormat.size(); i++) - { - if (outputFormat[i].compare("D") != 0) - { /// sparse dim - sparseDimsOutput.push_back(outputPerm[i]); - comet_debug() << "sparse dim in output: " << outputPerm[i] << " with format: " << outputFormat[i] << "\n"; +struct TransformSparseOutput : public OpRewritePattern { + TransformSparseOutput(MLIRContext *context) + : OpRewritePattern(context, /*benefit=*/0) {} + + mlir::LogicalResult + matchAndRewrite(IndexTreeComputeOp compute_op, mlir::PatternRewriter &rewriter) const override { + IndexTreeLHSOperandOp lhs_op = compute_op.getLhs().getDefiningOp(); + Value old_output = lhs_op.getTensor(); + + // Check to see if output is sparse + if(!llvm::isa(old_output.getType())) + return failure(); + + // Check to see if there are "redundant" inserts + llvm::SmallDenseMap index_vars; + for(auto pos : lhs_op.getPos()) { + auto index_to_tensor = pos.getDefiningOp(); + if(index_to_tensor){ + index_vars.insert(std::make_pair( + index_to_tensor.getIndex(), + index_to_tensor + )); + } } - } - return sparseDimsOutput; -} -/// get the sparse index that has sparse format in at least two input tensors -/// which tensor, which dimension. use std::pair represent the information -std::vector getSparseDimsInput(std::vector> opFormats, std::vector> opPerms) -{ - std::vector sparseDimsInput; - - std::vector> inputFormats = {opFormats.begin(), opFormats.end() - 1}; - std::vector> inputPerms = {opPerms.begin(), opPerms.end() - 1}; - - /// Get all dims in input tensors - std::vector allPermsInput = getUnionOf2Dvector(inputPerms); - comet_debug() << " allPermsInput.size(): " << allPermsInput.size() << "\n"; - comet_debug() << "allPermsInput: "; - for (auto n : allPermsInput) - { - comet_debug() << n << " "; - } - comet_debug() << "\n"; - - for (unsigned int i = 0; i < allPermsInput.size(); i++) - { - int cur_index = allPermsInput[i]; - comet_debug() << " cur_index: " << cur_index << "\n"; - /// Get the format of cur_index from each input tensor - std::vector cur_formats; - std::vector tensor_ids; - std::vector dim_orders; - for (unsigned int j = 0; j < inputPerms.size(); j++) - { - unsigned int whichFormat = findIndexInVector(inputPerms[j], cur_index); - if (whichFormat < inputPerms[j].size()) - { /// found - std::string format = inputFormats[j][whichFormat]; - cur_formats.push_back(format); - tensor_ids.push_back(j); - dim_orders.push_back(whichFormat); - } + // Find last output dimension + Value parent = compute_op.getParent(); + auto node = parent.getDefiningOp(); + while(index_vars.find(parent) == index_vars.end()) { + parent = node.getParent(); + node = parent.getDefiningOp(); } - comet_debug() << " cur_formats.size(): " << cur_formats.size() << "\n"; - comet_debug() << "cur_formats: "; - for (auto n : cur_formats) + // node contains output domain + + // Find output dimensions after reduction variable + // to include in workspace + unsigned workspace_rank = 0; + llvm::SmallVector accesses; + llvm::SmallVector dims; + while(index_vars.find(parent) != index_vars.end()){ + auto access_op = index_vars[parent]; + accesses.push_back(access_op); + dims.push_back(access_op.getDim()); + workspace_rank++; + + parent = node.getParent(); + node = parent.getDefiningOp(); + if(!node){ + return failure(); + } + } + //Match success! + // Parent contains reduction variable + + // Declare the workspace outside of the tree + auto loc = compute_op.getLoc(); + auto tree_op = compute_op->getParentOfType(); + rewriter.setInsertionPoint(tree_op); + SparseTensorType spType = llvm::cast(old_output.getType()); + llvm::SmallVector dim_sizes(workspace_rank, ShapedType::kDynamic); + Type workspace_type = WorkspaceType::get(compute_op.getContext(), spType.getElementType(), spType.getIndicesType(), dim_sizes); + std::reverse(dims.begin(), dims.end()); + + // Get the argument that corresponds to the tensor + // This may break if the tensor is part of a longer use def chain of computations + BlockArgument output_arg = llvm::cast(old_output); + Value original_tensor = tree_op->getOperand(output_arg.getArgNumber()); + Value workspace = rewriter.create(loc, workspace_type, original_tensor, rewriter.getI32ArrayAttr(dims)); + + // Modify the itree body to expect the workspace as a modifiable "tensor" + Value workspace_arg = tree_op.getBody()->addArgument(workspace_type, loc); + + // Clean the workspace before use + rewriter.setInsertionPoint(node); + Value clean_workspace = rewriter.create(loc, workspace_type, node.getParent(), workspace_arg); + + // Create new compute op + auto context = getContext(); + rewriter.setInsertionPoint(compute_op); + Type index_type = rewriter.getIndexType(); + llvm::SmallVector pos; + llvm::SmallVector crds; + std::reverse(accesses.begin(), accesses.end()); + int32_t dim = 0; + Value prev_dim = nullptr; + for(auto access_op : accesses) { - comet_debug() << n << " "; + auto new_access_op = rewriter.create( + loc, + TypeRange({index_type, index_type}), + clean_workspace, + access_op.getIndex(), + rewriter.getUI32IntegerAttr(dim), + prev_dim + ); + + pos.push_back(new_access_op.getPos()); + crds.push_back(new_access_op.getCrd()); + prev_dim = new_access_op.getPos(); + dim++; } - comet_debug() << "\n"; - /// check if there is sparse format in cur_formats vector - std::vector cur_sparse_formats; - std::vector sparse_tensor_ids; - std::vector sparse_dim_orders; - for (unsigned int j = 0; j < cur_formats.size(); j++) + Type operand_type = OperandType::get(context); + Value new_lhs = rewriter.create( + loc, + operand_type, + clean_workspace, + pos, + crds + ); + Value new_workspace = rewriter.create( + loc, + workspace_type, + compute_op.getParent(), + new_lhs, + compute_op.getRhs(), + nullptr, + compute_op.getSemiringAttr() + ); + + pos.clear(); + crds.clear(); + dim = 0; + prev_dim = nullptr; + for(auto access_op : accesses) { - comet_debug() << " cur_formats[" << j << "]: " << cur_formats[j] << "\n"; - if (cur_formats[j].compare("D") != 0) - { /// sparse format - cur_sparse_formats.push_back(cur_formats[j]); - sparse_tensor_ids.push_back(tensor_ids[j]); - sparse_dim_orders.push_back(dim_orders[j]); - comet_debug() << " sparse dim in format: " << cur_index << " with format: " << cur_formats[j] << "\n"; - } + auto new_access_op = rewriter.create( + loc, + TypeRange({index_type, index_type}), + new_workspace, + access_op.getIndex(), + rewriter.getUI32IntegerAttr(dim), + prev_dim + ); + + pos.push_back(new_access_op.getPos()); + crds.push_back(new_access_op.getCrd()); + prev_dim = new_access_op.getPos(); + dim++; } - if (cur_sparse_formats.size() > 1) - { /// More than one sparse format - struct dimInTensor dim_in_tensor; - dim_in_tensor.dim = cur_index; - dim_in_tensor.tensorId = sparse_tensor_ids[0]; /// Any sparse tensor is ok - dim_in_tensor.dimOrder = sparse_dim_orders[0]; - sparseDimsInput.push_back(dim_in_tensor); + Value new_rhs = rewriter.create( + loc, + operand_type, + new_workspace, + pos, + crds + ); + + rewriter.replaceOpWithNewOp( + compute_op, + old_output.getType(), + compute_op.getParent(), + compute_op.getLhs(), + ValueRange{new_rhs,}, + nullptr, + "noop_noop" + ); + + + // Update the index tree op + SmallVector tree_temps(tree_op.getIntermediates()); + SmallVector tree_types(tree_op->getResultTypes()); + tree_temps.push_back(workspace); + tree_types.push_back(workspace_type); + rewriter.setInsertionPoint(tree_op); + auto newOp = rewriter.create(loc, tree_types, tree_op.getInputs(), tree_temps); + rewriter.inlineRegionBefore(tree_op.getRegion(), newOp.getRegion(), newOp.getRegion().end()); + indexTree::YieldOp yield = cast(newOp.getRegion().getBlocks().front().getTerminator()); + rewriter.modifyOpInPlace(yield, [&]() { + yield->insertOperands(yield->getNumOperands(), ValueRange{workspace}); + }); + for(unsigned i = 0; i < tree_op.getNumResults(); i++){ + rewriter.replaceAllUsesWith(tree_op.getResult(i), newOp.getResult(i)); } - } + rewriter.eraseOp(tree_op); - comet_debug() << "sparseDimsInput: "; - for (auto n : sparseDimsInput) - { - comet_debug() << "(" << n.dim << ", " << n.tensorId << ", " << n.dimOrder << ") "; + return success(); } - comet_debug() << "\n"; - return sparseDimsInput; -} - -/// Split one indicesOp into several one, i.e. each computeOp has its own parent op -/// i -> j -> V=0;V=A;W=V*B ===> i -> j -> V=0; -/// -> j -> V=A; -/// -> j -> W=V*B -void splitIndicesOp(Operation *needSplitNode, Value denseIndicesOp, OpBuilder &builder, Location loc) -{ - while (isa(needSplitNode)) - { - - comet_pdump(needSplitNode); - /// check how many operands, split into many operands. - indexTree::IndexTreeIndicesOp indicesOp = dyn_cast(needSplitNode); - - comet_vdump(indicesOp); - - Operation *indicesOpFirstUsers = *(indicesOp.getOperation()->getResult(0).getUsers().begin()); - comet_pdump(indicesOpFirstUsers); - - builder.setInsertionPoint(indicesOpFirstUsers); - comet_debug() << "\n"; - - if (needSplitNode != denseIndicesOp.getDefiningOp()) - { - ArrayAttr indices = indicesOp.getIndices(); - - comet_debug() << " indicesOp.getOperation()->getNumOperands(): " << indicesOp.getOperation()->getNumOperands() << "\n"; - std::vector operands; - std::vector newIndicesOp; - for (unsigned int i = 0; i < indicesOp.getOperation()->getNumOperands(); i++) - { - operands.push_back(indicesOp.getOperation()->getOperand(i)); - - comet_vdump(indicesOp.getOperation()->getOperand(i)); - comet_vdump(operands[i]); - auto i64Type = builder.getI64Type(); - /// TODO(zhen.peng): new attribute iterator_type - auto dumb_iterator_type = builder.getStringAttr("default"); - Value t1 = builder.create(loc, i64Type, operands[i], indices, dumb_iterator_type); - - comet_debug() << "New IndexTreeIndicesOp added:\n"; - comet_vdump(t1); - newIndicesOp.push_back(t1); - } - - /// put it here - comet_debug() << " finished calling replacereplaceOperands \n"; - /// This parentIndicesOp is the operation that need to be splitted next time - -#ifdef DEBUG_MODE_WorkspaceTransformsPass - comet_vdump(indicesOp.getOperation()->getResult(0)); - for (auto ppp : indicesOp.getOperation()->getResult(0).getUsers()) - { - - comet_pdump(ppp); - } -#endif - - Operation *parentIndicesOp = *(indicesOp.getOperation()->getResult(0).getUsers().begin()); - - comet_pdump(parentIndicesOp); - - replaceOperands(needSplitNode, newIndicesOp); - needSplitNode = parentIndicesOp; +}; - comet_debug() << " plan to erase the following Op\n"; - comet_debug() << " Indices operations:\n"; - comet_vdump(indicesOp); - comet_debug() << " Split Nodes:\n"; - comet_pdump(needSplitNode); - comet_debug() << " Indices op first users:\n"; - comet_pdump(indicesOpFirstUsers); - indicesOp.erase(); - } - else - { - comet_debug() << "\n"; - break; +struct MoveInvariantComputeOp : public OpRewritePattern { + MoveInvariantComputeOp (MLIRContext *context) + : OpRewritePattern(context, /*benefit=*/2) {} + + mlir::LogicalResult + matchAndRewrite(IndexTreeComputeOp compute_op, mlir::PatternRewriter &rewriter) const override { + // Collect all indices used in this compute expression + llvm::SmallDenseSet used_indices; + IndexTreeLHSOperandOp lhs_op = compute_op.getLhs().getDefiningOp(); + for(auto pos : lhs_op.getPos()) { + auto index_to_tensor = pos.getDefiningOp(); + if(!index_to_tensor) + return failure(); + used_indices.insert(index_to_tensor.getIndex()); } - } - comet_debug() << "\n"; -} - -void removeRedundantIndices(std::vector newComputeOps, - std::map indexValueMap, - int denseDimInOutput, - OpBuilder &builder, - Location loc) -{ - - /// Check whether need to remove redundant indices or not - /// Get the - - /// -------Remove redundant indices------------- - /// For C, the 1st dim i is Dense, the second dim j is sparse. - /// ---- the index including i and before i is not included - mlir::Value denseIndicesOp = indexValueMap[denseDimInOutput]; - /// The indices after denseIndicesOp need to be splitted - /// start from the computeOp, - /// Finished one level - /// Only one User, because it's a tree structure, the leaf only has one parent - assert(newComputeOps[0].getDefiningOp()->getResult(0).hasOneUse() && " the computeOp has more than one users\n"); - /// Get the only one user - Operation *onlyUser = *(newComputeOps[0].getDefiningOp()->getResult(0).getUsers().begin()); - comet_pdump(onlyUser); - - /// needSplitNode is the parent node of the "denseIndicesOp" - Operation *needSplitNode = onlyUser; - /// iterate until the - /// call splitIndicesOp function to split indicesOp until latest "root" - splitIndicesOp(needSplitNode, denseIndicesOp, builder, loc); - comet_debug() << "\n"; - - /// Remove the indices for each tensor - /// iterate over all itComputeOps, get the indices for each tensor - for (auto n : newComputeOps) - { - /// get allPerms, put all indices id into a vector, - /// iterater up until reach the root noe, if the index of indicesOp is not in the vector - /// remove this one: set the operand of the parent of the indicesOp into current op - - comet_debug() << " current computeOp: \n"; - comet_vdump(n); - ArrayAttr allperms_rhs = dyn_cast(n.getDefiningOp()->getOperand(0).getDefiningOp()).getAllPerms(); - std::vector> allpermsInt_rhs = convertArrayAttrIntTo2DVector(allperms_rhs); - std::vector permsInt = getUnionOf2Dvector(allpermsInt_rhs); - comet_debug() << " print permsInt: "; - for (auto p : permsInt) - { - comet_debug() << p << " "; + + auto rhs_operands = compute_op.getRhs(); + for(Value rhs : rhs_operands) { + IndexTreeOperandOp operand_op = rhs.getDefiningOp(); + for(auto pos : operand_op.getPos()) { + auto index_to_tensor = pos.getDefiningOp(); + if(!index_to_tensor) + return failure(); + used_indices.insert(index_to_tensor.getIndex()); + } } - comet_debug() << "\n"; - - mlir::Value computeOp = n; - /// iterate over the IndexTreeIndicesOp; - mlir::Value computeOpParent; /// computeOpParent is IndexTreeIndicesOp - - comet_vdump(n); - assert(n.getDefiningOp()->getResult(0).hasOneUse() && " indicesOp has more than one user\n"); - Operation *computeOpParentPointer = *(n.getDefiningOp()->getResult(0).getUsers().begin()); - computeOpParent = computeOpParentPointer->getResult(0); - - comet_pdump(computeOpParentPointer); - - comet_vdump(computeOpParent); - - while (!isRealRoot(computeOpParent.getDefiningOp())) - { - comet_vdump(computeOpParent); - if (isa(computeOpParent.getDefiningOp())) - { - comet_debug() << " indicesOp's parent can not be computeOp\n"; - } - else if (isa(computeOpParent.getDefiningOp())) - { - comet_debug() << " indicesOp's parent is IndexTreeOp\n"; + // We want to find all of the indices that this compute op is nested under + // and check if they are used in this compute expression. Every time we come + // across an unused index, the index nodes that we have seen so far need to be copied + // to form a new branch of the tree. We also keep track of the parent at the fork + llvm::SmallVector seen_indices; + llvm::SmallVector indices_to_copy; + Value parent = compute_op.getParent(); + Value fork = parent; + IndexTreeIndicesOp node = parent.getDefiningOp(); + while(node) { + if(used_indices.find(parent) != used_indices.end()) { + // Used index variable + seen_indices.push_back(parent); + } else { + // Unused index variable + fork = node.getParent(); + indices_to_copy.insert(indices_to_copy.begin(), seen_indices.rbegin(), seen_indices.rend()); + seen_indices.clear(); } - else if (isa(computeOpParent.getDefiningOp())) - { - /// get the indices integer, to see if it is in permsInt - /// if yes, don't remove - /// if no, remove: - indexTree::IndexTreeIndicesOp curIndicesOp = dyn_cast(computeOpParent.getDefiningOp()); - comet_debug() << " \n"; - ArrayAttr idsArrayAttr = curIndicesOp.getIndices(); /// should be 1D vector - std::vector idsVec; - for (auto n : idsArrayAttr) - { - idsVec.push_back(n.cast().getInt()); - } - comet_debug() << " print idsVec: "; - for (auto p : idsVec) - { - comet_debug() << p << " "; - } - comet_debug() << "\n"; - - assert(idsVec.size() == 1 && " indicesOp contain more than 1 index\n"); - bool isNeedRemove = false; - for (auto n : idsVec) - { /// only 1 index actually, because each indicesOp contain one index - if (std::find(permsInt.begin(), permsInt.end(), n) != permsInt.end()) - { - /// found - isNeedRemove = false; - } - else - { /// the index in curIndicesOp is not found in the computeOp indices - isNeedRemove = true; - } - } - - /// if curIndicesOp is the "real root" of the index tree (has only one user) - /// contain more than 1 index - if (curIndicesOp.getOperation()->getNumOperands() > 1 && curIndicesOp.getOperation()->getResult(0).hasOneUse() && isa(*(curIndicesOp.getOperation()->getResult(0).getUsers().begin()))) - { - isNeedRemove = false; - } - comet_debug() << " isNeedRemove = " << isNeedRemove << "\n"; - - if (isNeedRemove) - { - assert(curIndicesOp.getOperation()->getResult(0).hasOneUse() && " indicesOp has more than one user\n"); - Operation *curIndicesOpParent = *(curIndicesOp.getOperation()->getResult(0).getUsers().begin()); - - comet_vdump(computeOpParent); - comet_pdump(curIndicesOpParent); - - computeOpParent.replaceAllUsesWith(computeOp); /// replace all uses of the indexOp with the new indecesOp - computeOp = computeOpParent; - computeOpParent.getDefiningOp()->erase(); /// erase the previous indecesOp - computeOpParent = curIndicesOpParent->getResult(0); - } - else - { -#ifdef DEBUG_MODE_WorkspaceTransformsPass - comet_vdump(curIndicesOp); - int count = 0; - for (auto p : curIndicesOp.getOperation()->getResult(0).getUsers()) - { - comet_pdump(p); - count++; - } - comet_debug() << " count: " << count << "\n"; -#endif - assert(curIndicesOp.getOperation()->getResult(0).hasOneUse() && " indicesOp has more than one user\n"); - Operation *curIndicesOpParent = *(curIndicesOp.getOperation()->getResult(0).getUsers().begin()); - - comet_pdump(curIndicesOpParent); - computeOp = computeOpParent; - computeOpParent = curIndicesOpParent->getResult(0); - } - } + parent = node.getParent(); + node = parent.getDefiningOp(); } - } /// end for n -} - -std::vector CompressedWorkspaceOutput(std::vector sparseDimsOutput, - indexTree::IndexTreeComputeOp itComputeOp, - std::vector> opFormats, - std::vector> opPerms, - std::map indexValueMap, - OpBuilder &builder, indexTree::IndexTreeOp op) -{ - Location loc = op.getLoc(); - auto comp_worksp_opt = builder.getBoolAttr(compressedworkspace); - int sparseDimOutput = -1; - int sparseDimOrderInOutput = -1; - int denseDimInOutput = -1; - auto i64Type = builder.getI64Type(); - - for (unsigned int j = 0; j < opFormats[opFormats.size() - 1].size(); j++) - { - /// sparse dimension - if (opFormats[opFormats.size() - 1][j].compare("D") != 0) - { - sparseDimOutput = opPerms[opPerms.size() - 1][j]; - sparseDimOrderInOutput = j; + if(fork == compute_op.getParent()) { + return failure(); // Match failed, no indces to move. } - else /// dense dimension - denseDimInOutput = opPerms[opPerms.size() - 1][j]; - } - comet_debug() << " " << sparseDimOutput << "\n"; - - /// 3. Find the ta.itIndices op which represents sparseDimOutput - /// Find its parent ... - Value sparseIndicesOp = indexValueMap[sparseDimOutput]; - - comet_vdump(sparseIndicesOp); - comet_debug() << " sparseDimOrderInOutput: " << sparseDimOrderInOutput << "\n"; - Value sparseDimsPerent = indexValueMap[sparseDimOrderInOutput - 1]; - comet_debug() << " sparseDimsPerent: \n"; - comet_vdump(sparseDimsPerent); - - /// Cij = Aik * Bkj ==> - /// ComputeNode(c1): Wj = 0; - /// ComputeNode(c2): Wj += Aik * Bkj; - /// ComputeNode(c3): Cij = Wj - std::vector tensors; - getTensorsOfComputeOp(itComputeOp.getOperation()->getResult(0), tensors); - - /// 4. create W, j dim size of - /// Value outputItComputeOp = itComputeOp.getOperation()->getOperand(itComputeOp.getOperation()->getNumOperands() - 1).getDefiningOp()->getOperand(0); - /// new version - Value outputItComputeOp = tensors[tensors.size() - 1]; - comet_vdump(outputItComputeOp); - - comet_vdump(outputItComputeOp.getDefiningOp()->getOperand(sparseDimOrderInOutput)); - std::string w_format = "Dense"; /// tensor - auto w_type = RankedTensorType::get({mlir::ShapedType::kDynamic}, builder.getF64Type()); - - Operation *itComputeOpFirstUsers = *(itComputeOp.getOperation()->getUsers().begin()); - builder.setInsertionPoint(itComputeOpFirstUsers); /// Insert before itree Op - std::vector w_lbls_value; - if (outputItComputeOp.getType().cast().isDynamicDim(sparseDimOrderInOutput)) - { - auto opIdx = outputItComputeOp.getType().cast().getDynamicDimIndex(sparseDimOrderInOutput); - w_lbls_value.push_back(outputItComputeOp.getDefiningOp()->getOperand(opIdx)); - } - else - { - w_lbls_value.push_back(builder.create(loc, outputItComputeOp.getType().cast().getDimSize(sparseDimOrderInOutput))); - } - - mlir::Value w = builder.create(loc, w_type, w_lbls_value, w_format); - comet_vdump(w); - auto w_index_list_type = RankedTensorType::get({mlir::ShapedType::kDynamic}, builder.getIndexType()); /// tensor - mlir::Value w_already_set = builder.create(loc, w_index_list_type, w_lbls_value, w_format); - comet_vdump(w_already_set); - mlir::Value w_index_list = builder.create(loc, w_index_list_type, w_lbls_value, w_format); - comet_vdump(w_index_list); - - MemRefType w_index_list_size_type = MemRefType::get({1}, builder.getIndexType()); /// tensor<1xindex> - mlir::Value w_index_list_size_alloc = builder.create(loc, w_index_list_size_type); /// tensor<1xindex> - Value w_index_list_size = builder.create(loc, w_index_list_size_alloc); - - std::vector workspaceTensors = {w, w_already_set, w_index_list, w_index_list_size}; - tensors.push_back(w); /// {A, B, C, W} - - std::vector> formats = {opFormats[0], opFormats[1], opFormats[2], {"D"}}; - std::vector> perms = {opPerms[0], opPerms[1], opPerms[2], {sparseDimOutput}}; - - /// Start building an IndexTreeCompute Operation to represent Wj = 0; - std::vector c1_perms_int_0; - std::vector c1_perms_int_1; - std::vector> c1_perms_int = {c1_perms_int_0, c1_perms_int_1}; - std::vector c1_formats_str_0; - std::vector c1_formats_str_1; - std::vector> c1_formats_str = {c1_formats_str_0, c1_formats_str_1}; - - Value const_index_0 = builder.create(loc, 0); - std::vector c1_rhs = {const_index_0}; - mlir::Value c1_lhs = {w_index_list_size}; - std::string semiringName(itComputeOp.getSemiring().data()); - std::string maskNone = "none"; - std::string maskTypeName(itComputeOp.getMaskType().data()); - auto c1_semiring = builder.getStringAttr(semiringName); - auto c1_maskType = builder.getStringAttr(maskNone); /// masking attribute - - /// for c1_rhs - std::vector> c1_rhsop_perms_str = {c1_perms_int_0}; - ArrayAttr c1_rhsop_perms = convert2DVectorToArrayAttrInt(c1_rhsop_perms_str, builder); - std::vector> c1_rhsop_formats_str = {c1_formats_str_0}; - ArrayAttr c1_rhsop_formats = convert2DVectorToArrayAttrStr(c1_rhsop_formats_str, builder); - mlir::Value c1_rhsop = builder.create(loc, mlir::UnrankedTensorType::get(builder.getIndexType()), c1_rhs, c1_rhsop_perms, c1_rhsop_formats); - comet_debug() << "IndexTreeComputeRHS Operation in Output (c1_rhs):\n"; - comet_vdump(c1_rhsop); - /// for c1_lhs - std::vector> c1_lhsop_perms_str = {c1_perms_int_1}; - ArrayAttr c1_lhsop_perms = convert2DVectorToArrayAttrInt(c1_lhsop_perms_str, builder); - std::vector> c1_lhsop_formats_str = {c1_formats_str_1}; - ArrayAttr c1_lhsop_formats = convert2DVectorToArrayAttrStr(c1_lhsop_formats_str, builder); - mlir::Value c1_lhsop = builder.create(loc, mlir::UnrankedTensorType::get(builder.getF64Type()), c1_lhs, c1_lhsop_perms, c1_lhsop_formats); - comet_debug() << "IndexTreeComputeLHS Operation in Output (c1_lhs):\n"; - comet_vdump(c1_lhsop); - - /// for c1 ==> Wj = 0; - mlir::Value c1 = builder.create(loc, builder.getI64Type(), c1_rhsop, c1_lhsop, comp_worksp_opt, c1_semiring, c1_maskType); - comet_debug() << "IndexTreeCompute Operation in Output (c1):\n"; - comet_vdump(c1); - - /// insert c1 to sparseDimsParent - sparseDimsPerent.getDefiningOp()->insertOperands(0, c1); - - /// Start building an IndexTreeCompute Operation to represent Wj += Aik * Bkj; - std::vector c2_tensors = {tensors[0], tensors[1], w}; - std::vector c2_perms_int_0 = opPerms[0]; - std::vector c2_perms_int_1 = opPerms[1]; - std::vector c2_perms_int_2 = {sparseDimOutput}; - std::vector> c2_perms_int = {c2_perms_int_0, c2_perms_int_1, c2_perms_int_2}; - - /// Convert formats string array into StrAttr - std::vector c2_formats_str_0 = opFormats[0]; - std::vector c2_formats_str_1 = opFormats[1]; - std::vector c2_formats_str_2 = {"D"}; - std::vector> c2_formats_str = {c2_formats_str_0, c2_formats_str_1, c2_formats_str_2}; - std::vector c2_rhs; - std::vector> c2_rhsop_formats_str; - std::vector> c2_rhsop_perms_str; - if (tensors.size() > 4) /// masking input is available: tensors = {%op0, %op1, %mask, %out, %W} - { - c2_rhs = {c2_tensors[0], c2_tensors[1], tensors[2]}; /// tensors[2] val is the mask - c2_rhsop_formats_str = {c2_formats_str_0, c2_formats_str_1, opFormats[2]}; /// mask format is same as the output - c2_rhsop_perms_str = {c2_perms_int_0, c2_perms_int_1, opPerms[2]}; /// perms of mask are same as the output - } - else /// no masking input is provided: tensors = {%op0, %op1, %out, %W} - { - c2_rhs = {c2_tensors[0], c2_tensors[1]}; - c2_rhsop_formats_str = {c2_formats_str_0, c2_formats_str_1}; - c2_rhsop_perms_str = {c2_perms_int_0, c2_perms_int_1}; - } - std::vector c2_lhs = workspaceTensors; - - auto c2_semiring = builder.getStringAttr(semiringName); - auto c2_maskType = builder.getStringAttr(maskTypeName); /// masking attribute - - /// for c2_rhsop - ArrayAttr c2_rhsop_perms = convert2DVectorToArrayAttrInt(c2_rhsop_perms_str, builder); - ArrayAttr c2_rhsop_formats = convert2DVectorToArrayAttrStr(c2_rhsop_formats_str, builder); - mlir::Value c2_rhsop = builder.create(loc, mlir::UnrankedTensorType::get(builder.getF64Type()), c2_rhs, c2_rhsop_perms, c2_rhsop_formats); - comet_debug() << "IndexTreeComputeRHS Operation in Output (c2_rhs):\n"; - comet_vdump(c2_rhsop); - - /// for c2_lhsop - std::vector> c2_lhsop_perms_str = {c2_perms_int_2}; - ArrayAttr c2_lhsop_perms = convert2DVectorToArrayAttrInt(c2_lhsop_perms_str, builder); - std::vector> c2_lhsop_formats_str = {c2_formats_str_2}; - ArrayAttr c2_lhsop_formats = convert2DVectorToArrayAttrStr(c2_lhsop_formats_str, builder); - mlir::Value c2_lhsop = builder.create(loc, mlir::UnrankedTensorType::get(builder.getF64Type()), c2_lhs, c2_lhsop_perms, c2_lhsop_formats); - comet_debug() << "IndexTreeComputeLHS Operation in Output (c2_lhs):\n"; - comet_vdump(c2_lhsop); - - /// for c2 - mlir::Value c2 = builder.create(loc, i64Type, c2_rhsop, c2_lhsop, comp_worksp_opt, c2_semiring, c2_maskType); - comet_debug() << "IndexTreeCompute Operation in Output (c2):\n"; - comet_vdump(c2); - - /// Start building an IndexTreeCompute Operation to represent Cij = Wj; - std::vector c3_tensors; - if (tensors.size() > 4) /// masking input is available: tensors = {%op0, %op1, %mask, %out, %W} - { - c3_tensors = {tensors[3]}; - } - else /// masking input is NOT available: tensors = {%op0, %op1, %out, %W} - { - c3_tensors = {tensors[2]}; - } - std::vector c3_perms_int_0 = {sparseDimOutput}; - std::vector c3_perms_int_1 = opPerms[2]; - std::vector> c3_perms_int = {c3_perms_int_0, c3_perms_int_1}; - - /// Convert formats string array into StrAttr - std::vector c3_formats_str_0 = {"D"}; - std::vector c3_formats_str_1 = opFormats[2]; - std::vector> c3_formats_str = {c3_formats_str_0, c3_formats_str_1}; - - std::vector c3_rhs = workspaceTensors; - mlir::Value c3_lhs = c3_tensors[0]; - auto c3_semiring = builder.getStringAttr(semiringName); - auto c3_maskType = builder.getStringAttr(maskNone); /// masking attribute - - /// for c3_rhs - std::vector> c3_rhsop_perms_str = {c3_perms_int_0}; - ArrayAttr c3_rhsop_perms = convert2DVectorToArrayAttrInt(c3_rhsop_perms_str, builder); - std::vector> c3_rhsop_formats_str = {c3_formats_str_0}; - ArrayAttr c3_rhsop_formats = convert2DVectorToArrayAttrStr(c3_rhsop_formats_str, builder); - mlir::Value c3_rhsop = builder.create(loc, mlir::UnrankedTensorType::get(builder.getF64Type()), c3_rhs, c3_rhsop_perms, c3_rhsop_formats); - comet_debug() << "IndexTreeComputeRHS Operation in Output (c3_rhs):\n"; - comet_vdump(c3_rhsop); - - /// for c3_lhs - std::vector> c3_lhsop_perms_str = {c3_perms_int_1}; - ArrayAttr c3_lhsop_perms = convert2DVectorToArrayAttrInt(c3_lhsop_perms_str, builder); - std::vector> c3_lhsop_formats_str = {c3_formats_str_1}; - ArrayAttr c3_lhsop_formats = convert2DVectorToArrayAttrStr(c3_lhsop_formats_str, builder); - mlir::Value c3_lhsop = builder.create(loc, mlir::UnrankedTensorType::get(builder.getF64Type()), c3_lhs, c3_lhsop_perms, c3_lhsop_formats); - comet_debug() << "IndexTreeComputeLHS Operation in Output (c3_lhs):\n"; - comet_vdump(c3_lhsop); - - /// for c3 ==> Cij = Wj; - mlir::Value c3 = builder.create(loc, i64Type, c3_rhsop, c3_lhsop, comp_worksp_opt, c3_semiring, c3_maskType); - comet_debug() << "IndexTreeCompute Operation in Output (c3):\n"; - comet_vdump(c3); - - std::vector newComputeOps = {c2, c3}; - sparseIndicesOp.getDefiningOp()->setOperands(newComputeOps); - - /// remove redundant indices by calling a function - /// in elementwise: not remove - /// in spgemm: remove - /// check if there is redundant index - bool existRedundantIndex = false; - for (auto n : newComputeOps) - { - std::vector> perms; - getPermsOfComputeOp(n, perms); - std::vector allperms = getUnionOf2Dvector(perms); - comet_debug() << " print allperms \n"; - print_vector(allperms); - - std::vector ancestors; - std::vector dfsOps; - dfsRootOpTree(op.getChildren(), dfsOps); - getAncestorsWp(n, ancestors, dfsOps); - comet_debug() << " print ancestors \n"; - print_vector_value(ancestors); - - /// Iterate over every indicesOp - for (auto ancestor : ancestors) + // Success! + IRMapping map; + auto context = rewriter.getContext(); + auto loc = compute_op.getLoc(); + IndexNodeType index_node_type = IndexNodeType::get(context); + parent = fork; + for(auto index : indices_to_copy) { - /// If indicesOp's index is in allperms, no redundant - /// the indicesOp is real root, no redundant - /// Otherwise, redundant - if (isa(ancestor.getDefiningOp())) - { - indexTree::IndexTreeIndicesOp indicesOp = dyn_cast(ancestor.getDefiningOp()); - - ArrayAttr idsArrayAttr = indicesOp.getIndices(); /// should be 1D vector - /// actually only one index for the indicesOp in our implementation - for (auto m : idsArrayAttr) - { - int perm = m.cast().getInt(); - comet_debug() << " perm: " << perm << "\n"; - - if (findIndexInVector(allperms, perm) == allperms.size()) - { /// not exit - comet_debug() << " perm not exist in allperms\n"; - if (!isRealRoot(indicesOp.getOperation())) - { - existRedundantIndex = true; - comet_debug() << " existRedundantIndex: " << existRedundantIndex << "\n"; - } - } - } - } + Value new_index = rewriter.create(loc, index_node_type, parent); + map.map(index, new_index); + parent = new_index; } - } - if (existRedundantIndex) - { - comet_debug() << "There is loop invariant\n"; - removeRedundantIndices(newComputeOps, indexValueMap, denseDimInOutput, builder, loc); - } - - return newComputeOps; -} /// end CompressedWorkspaceOutput() - -void CompressedWorkspaceInput(std::vector computeOps, OpBuilder &builder, Location loc) -{ - auto comp_worksp_opt = builder.getBoolAttr(compressedworkspace); - for (auto computeOp : computeOps) - { - /// 1. get the opFormats and opPerms of the computeOp - std::vector> opFormats; - std::vector> opPerms; - std::vector> inputOutputMapping; - getFormatsPermsOfComputeOp(computeOp, opFormats, opPerms, inputOutputMapping); - comet_debug() << " \n"; - for (auto n : opFormats) - { - - print_vector(n); + for(auto pos : lhs_op.getPos()) { + Operation* index_to_tensor = pos.getDefiningOp(); + rewriter.clone(*index_to_tensor, map); } - for (auto n : opPerms) - { - - print_vector(n); + rewriter.clone(*lhs_op.getOperation(), map); + + for(Value rhs : rhs_operands) { + IndexTreeOperandOp operand_op = rhs.getDefiningOp(); + for(auto pos : operand_op.getPos()) { + Operation* index_to_tensor = pos.getDefiningOp(); + rewriter.clone(*index_to_tensor, map); + } + rewriter.clone(*operand_op.getOperation(), map); } - std::vector tensors; - getTensorsOfComputeOp(computeOp, tensors); - comet_debug() << " tensors.size(): " << tensors.size() << "\n"; - std::vector tensors_rhs; - getInputTensorsOfComputeOp(computeOp, tensors_rhs); - comet_debug() << " tensors_rhs.size(): " << tensors_rhs.size() << "\n"; - std::vector tensors_lhs; - getOutputTensorsOfComputeOp(computeOp, tensors_lhs); - comet_debug() << " tensors_lhs.size(): " << tensors_lhs.size() << "\n"; - - indexTree::IndexTreeComputeOp itComputeOp = dyn_cast(computeOp.getDefiningOp()); - std::string semiringName(itComputeOp.getSemiring().data()); - - std::vector sparseDimsOutput = getSparseDimsOutput(opFormats, opPerms); - std::vector sparseDimsInput = getSparseDimsInput(opFormats, opPerms); - comet_debug() << " sparseDimsInput.size(): " << sparseDimsInput.size() << "\n"; - - if (sparseDimsInput.size() == 1) - { /// solve only 1 sparseDimsInput - /// No need to apply workspace transformation - comet_debug() << " sparseDimsInput[0]: " << sparseDimsInput[0].dim << ", " << sparseDimsInput[0].tensorId << ", " << sparseDimsInput[0].dimOrder << "\n"; - - /// Wj=Aij*Bij ==> - /// ComputeNode(c1): Vj=0; - /// ComputeNode(c2): Vj=Aij; - /// ComputeNode(c3): Wj=Vj*Bij - - Value sparseInput = tensors_rhs[sparseDimsInput[0].tensorId]; - comet_vdump(sparseInput); - - comet_vdump(sparseInput.getDefiningOp()->getOperand(sparseDimsInput[0].dimOrder)); - std::string v_format = "Dense"; /// tensor - auto v_type = RankedTensorType::get({mlir::ShapedType::kDynamic}, builder.getF64Type()); - - builder.setInsertionPoint(computeOp.getDefiningOp()); - std::vector v_lbls_value = {builder.create(loc, sparseInput, sparseDimsInput[0].dimOrder)}; - mlir::Value v = builder.create(loc, v_type, v_lbls_value, v_format); - comet_vdump(v); - - /// Start building an IndexTreeCompute Operation to represent Vj=0 - std::vector c1_perms_int_0; - std::vector c1_perms_int_1 = {sparseDimsInput[0].dim}; - std::vector> c1_perms_int = {c1_perms_int_0, c1_perms_int_1}; - - std::vector c1_formats_str_0; - std::vector c1_formats_str_1 = {"D"}; - std::vector> c1_formats_str = {c1_formats_str_0, c1_formats_str_1}; - - auto i64Type = builder.getI64Type(); - Value const_f64_0 = builder.create(loc, builder.getF64Type(), builder.getF64FloatAttr(0.0)); - std::vector c1_rhs = {const_f64_0}; - mlir::Value c1_lhs = {v}; - std::string semiringName(itComputeOp.getSemiring().data()); - std::string maskNone = "none"; - auto c1_semiring = builder.getStringAttr(semiringName); - auto c1_maskType = builder.getStringAttr(maskNone); /// masking attribute - - /// for c1_rhs - std::vector> c1_rhsop_perms_str = {c1_perms_int_0}; - ArrayAttr c1_rhsop_perms = convert2DVectorToArrayAttrInt(c1_rhsop_perms_str, builder); - std::vector> c1_rhsop_formats_str = {c1_formats_str_0}; - ArrayAttr c1_rhsop_formats = convert2DVectorToArrayAttrStr(c1_rhsop_formats_str, builder); - mlir::Value c1_rhsop = builder.create(loc, - mlir::UnrankedTensorType::get(builder.getF64Type()), - c1_rhs, - c1_rhsop_perms, - c1_rhsop_formats); - comet_debug() << "IndexTreeComputeRHS Operation in Input (c1_rhs):"; - comet_vdump(c1_rhsop); - - /// for c1_lhs - std::vector> c1_lhsop_perms_str = {c1_perms_int_1}; - ArrayAttr c1_lhsop_perms = convert2DVectorToArrayAttrInt(c1_lhsop_perms_str, builder); - std::vector> c1_lhsop_formats_str = {c1_formats_str_1}; - ArrayAttr c1_lhsop_formats = convert2DVectorToArrayAttrStr(c1_lhsop_formats_str, builder); - mlir::Value c1_lhsop = builder.create(loc, - mlir::UnrankedTensorType::get(builder.getF64Type()), - c1_lhs, - c1_lhsop_perms, - c1_lhsop_formats); - comet_debug() << "IndexTreeComputeLHS Operation in Input (c1_lhs):"; - comet_vdump(c1_lhsop); - - /// for c1 - mlir::Value c1 = builder.create(loc, i64Type, - c1_rhsop, - c1_lhsop, - comp_worksp_opt, - c1_semiring, - c1_maskType); - comet_debug() << "IndexTreeCompute Operation in Input (c1): "; - comet_vdump(c1); - - /// Start building an IndexTreeCompute Operation to represent Vj = Aij - std::vector c2_perms_int_0 = opPerms[sparseDimsInput[0].tensorId]; - std::vector c2_perms_int_1 = {sparseDimsInput[0].dim}; - std::vector> c2_perms_int = {c2_perms_int_0, c2_perms_int_1}; - - std::vector c2_formats_str_0 = opFormats[sparseDimsInput[0].tensorId]; - std::vector c2_formats_str_1 = {"D"}; - std::vector> c2_formats_str = {c2_formats_str_0, c2_formats_str_1}; - - std::vector c2_rhs = {tensors_rhs[sparseDimsInput[0].tensorId]}; - - mlir::Value c2_lhs = {v}; - auto c2_semiring = builder.getStringAttr(semiringName); - auto c2_maskType = builder.getStringAttr(maskNone); /// masking attribute - - /// for c2_rhs - std::vector> c2_rhsop_perms_str = {c2_perms_int_0}; - ArrayAttr c2_rhsop_perms = convert2DVectorToArrayAttrInt(c2_rhsop_perms_str, builder); - std::vector> c2_rhsop_formats_str = {c2_formats_str_0}; - ArrayAttr c2_rhsop_formats = convert2DVectorToArrayAttrStr(c2_rhsop_formats_str, builder); - mlir::Value c2_rhsop = builder.create(loc, - mlir::UnrankedTensorType::get(builder.getF64Type()), - c2_rhs, - c2_rhsop_perms, - c2_rhsop_formats); - comet_debug() << "IndexTreeComputeRHS Operation in Input (c2_rhs):"; - comet_vdump(c2_rhsop); - - /// for c2_lhs - std::vector> c2_lhsop_perms_str = {c2_perms_int_1}; - ArrayAttr c2_lhsop_perms = convert2DVectorToArrayAttrInt(c2_lhsop_perms_str, builder); - std::vector> c2_lhsop_formats_str = {c2_formats_str_1}; - ArrayAttr c2_lhsop_formats = convert2DVectorToArrayAttrStr(c2_lhsop_formats_str, builder); - mlir::Value c2_lhsop = builder.create(loc, - mlir::UnrankedTensorType::get(builder.getF64Type()), - c2_lhs, - c2_lhsop_perms, - c2_lhsop_formats); - comet_debug() << "IndexTreeComputeLHS Operation in Input (c2_lhs):"; - comet_vdump(c2_lhsop); - - /// for c2 - mlir::Value c2 = builder.create(loc, i64Type, - c2_rhsop, - c2_lhsop, - comp_worksp_opt, - c2_semiring, - c2_maskType); - comet_debug() << "IndexTreeCompute Operation in Input (c2): "; - comet_vdump(c2); - - /// Start building an IndexTreeCompute Operation to represent Wj=Vj*Bij - std::vector c3_perms_int_0 = {sparseDimsInput[0].dim}; - std::vector c3_perms_int_1 = opPerms[1]; - std::vector c3_perms_int_2 = opPerms[opPerms.size() - 1]; - std::vector> c3_perms_int = {c3_perms_int_0, c3_perms_int_1, c3_perms_int_2}; - - /// Convert formats string array into StrAttr - std::vector c3_formats_str_0 = {"D"}; - std::vector c3_formats_str_1 = opFormats[1]; - std::vector c3_formats_str_2 = opFormats[opFormats.size() - 1]; - - std::vector> c3_formats_str = {c3_formats_str_0, c3_formats_str_1, c3_formats_str_2}; - std::vector c3_rhs = {v, tensors[1]}; - - comet_debug() << " tensors.size(): " << tensors.size() << "\n"; - std::vector c3_lhs = tensors_lhs; - - auto c3_semiring = builder.getStringAttr(semiringName); - auto c3_maskType = builder.getStringAttr(maskNone); /// masking attribute - - /// for c3_rhs - std::vector> c3_rhsop_perms_str = {c3_perms_int_0, c3_perms_int_1}; - ArrayAttr c3_rhsop_perms = convert2DVectorToArrayAttrInt(c3_rhsop_perms_str, builder); - std::vector> c3_rhsop_formats_str = {c3_formats_str_0, c3_formats_str_1}; - ArrayAttr c3_rhsop_formats = convert2DVectorToArrayAttrStr(c3_rhsop_formats_str, builder); - mlir::Value c3_rhsop = builder.create(loc, - mlir::UnrankedTensorType::get(builder.getF64Type()), - c3_rhs, - c3_rhsop_perms, - c3_rhsop_formats); - comet_debug() << "IndexTreeComputeRHS Operation in Input (c3_rhs):"; - comet_vdump(c3_rhsop); - - /// for c3_lhs - std::vector> c3_lhsop_perms_str = {c3_perms_int_2}; - ArrayAttr c3_lhsop_perms = convert2DVectorToArrayAttrInt(c3_lhsop_perms_str, builder); - std::vector> c3_lhsop_formats_str = {c3_formats_str_2}; - ArrayAttr c3_lhsop_formats = convert2DVectorToArrayAttrStr(c3_lhsop_formats_str, builder); - mlir::Value c3_lhsop = builder.create(loc, - mlir::UnrankedTensorType::get(builder.getF64Type()), - c3_lhs, - c3_lhsop_perms, - c3_lhsop_formats); - comet_debug() << "IndexTreeComputeLHS Operation in Input (c3_lhs):"; - comet_vdump(c3_lhsop); - - /// for c3 - mlir::Value c3 = builder.create(loc, i64Type, - c3_rhsop, - c3_lhsop, - comp_worksp_opt, - c3_semiring, - c3_maskType); - comet_debug() << "IndexTreeCompute Operation in Input (t3): "; - comet_vdump(c3); - - /// old version for new children ops - std::vector newComputeOps = {c1, c2, c3}; - replaceOperands(itComputeOp.getOperation(), newComputeOps); - - /// Step 2: split j into 3. - Operation *needSplitNode = *(newComputeOps[0].getDefiningOp()->getResult(0).getUsers().begin()); - Operation *parentSplitNode = *(needSplitNode->getResult(0).getUsers().begin()); - comet_debug() << " call splitIndicesOp for applying workspace in Input \n"; - comet_pdump(needSplitNode); - splitIndicesOp(needSplitNode, parentSplitNode->getResult(0), builder, loc); - comet_debug() << "\n"; - - } /// end if(sparseDimsInput.size() == 1) + Operation* new_compute_op = rewriter.clone(*compute_op.getOperation(), map); + rewriter.replaceOp(compute_op, new_compute_op->getResults()); + return success(); } -} - -void IndexTreeWorkspaceTransformationsPass::CompressedWorkspaceTransforms(mlir::func::FuncOp funcop) -{ - funcop.walk([](indexTree::IndexTreeOp op) - { - OpBuilder builder(op); - comet_vdump(op); - - Location loc = op.getLoc(); - - /// 1. Find its child, until reach the ta.itCompute op - /// Get first user - Value computeOp = op.getOperation()->getOperand(0); - comet_vdump(computeOp); - - /// Only one child?? - /// Build a map, which index is in which IndexTreeIndicesOp - /// ------ Notice: each index is only in one IndicesOp in original index tree here - /// ------ TODO(gkestor): handle more complicate cases: one index is in more than one IndicesOp - /// For an indexTree, the indices ids are - std::map indexValueMap; - - while (!(isa(computeOp.getDefiningOp()))) - { - if (isa(computeOp.getDefiningOp())) - { - auto indicesop = dyn_cast(computeOp.getDefiningOp()); - ArrayAttr idsArrayAttr = indicesop.getIndices(); - for (auto n : idsArrayAttr) - { - int ids = n.cast().getInt(); - indexValueMap.emplace(ids, computeOp); - } - } - computeOp = computeOp.getDefiningOp()->getOperand(0); /// put here - } - comet_vdump(computeOp); - - /// 2. Check if there is sparse dim in the ta.itCompute op, - std::vector> opFormats; - std::vector> opPerms; - std::vector> inputOutputMapping; - getFormatsPermsOfComputeOp(computeOp, opFormats, opPerms, inputOutputMapping); - -#ifdef DEBUG_MODE_WorkspaceTransformsPass - comet_debug() << "Print opFormats:\n"; - for (auto n : opFormats) - { - - print_vector(n); - } -#endif - - indexTree::IndexTreeComputeOp itComputeOp = dyn_cast(computeOp.getDefiningOp()); - - /// Check the input tensors, and the output tensor, to see if it contains sparse dimensions - /// get the dim ids - std::vector sparseDimsOutput = getSparseDimsOutput(opFormats, opPerms); - -#ifdef DEBUG_MODE_WorkspaceTransformsPass - comet_debug() << " Print sparseDimsOutput: "; - for (auto p : sparseDimsOutput) - { - comet_debug() << p << " "; - } - comet_debug() << "\n"; -#endif - - std::vector sparseDimsInput = getSparseDimsInput(opFormats, opPerms); - - if (sparseDimsOutput.size() == 0 && sparseDimsInput.size() == 0) - { - /// No need to apply workspace transformation - comet_debug() << __FILE__ << __LINE__ << " No need to apply workspace transformation\n"; - return; - } - - assert(sparseDimsOutput.size() == 1 && " More than one sparse index in the output, we are expecting to support it in the future\n"); - - std::vector newComputeOps; - /// create three IndexTreeComputeOp op - /// sparse dim in output tensor - if (sparseDimsOutput.size() == 1) - { - newComputeOps = CompressedWorkspaceOutput(sparseDimsOutput, itComputeOp, opFormats, opPerms, indexValueMap, builder, op); - } - /// initially here workspaceOutput content - -#ifdef DEBUG_MODE_WorkspaceTransformsPass - /// Should notice, the itree has been the new itree already after call workspaceOutput - for (auto n : newComputeOps) - { - - comet_vdump(n); - } -#endif - if (sparseDimsInput.size() == 1) - { - comet_vdump(op); - /// Need the newComputeOps - CompressedWorkspaceInput(newComputeOps, builder, loc); - } - - /// Also remove previous IndexTreeComputeOp's LHS and RHS. - indexTree::IndexTreeComputeRHSOp itComputeOp_rhs = dyn_cast(itComputeOp->getOperand(0).getDefiningOp()); - indexTree::IndexTreeComputeLHSOp itComputeOp_lhs = dyn_cast(itComputeOp->getOperand(1).getDefiningOp()); - - itComputeOp.erase(); - itComputeOp_rhs.erase(); - itComputeOp_lhs.erase(); }); /// end function traverse - - comet_debug() << __FILE__ << " " << __LINE__ << "CompressedWorkspaceTransforms pass is done\n"; -} +}; void IndexTreeWorkspaceTransformationsPass::runOnOperation() { comet_debug() << __FILE__ << " " << __LINE__ << " starting CompressedWorkspaceTransforms pass \n"; - func::FuncOp function = getOperation(); - /// Traverse the function, only handle ta.itree operation - CompressedWorkspaceTransforms(function); + mlir::RewritePatternSet workspace_transformation_patterns(&getContext()); + + workspace_transformation_patterns.add(&getContext()); + CopiedDomainAnalysis& copiedDomains = getAnalysis(); + indexTree::populateDomainInferencePatterns(&getContext(), workspace_transformation_patterns, copiedDomains); //For new index variables + indexTree::populateDomainConcretizationPatterns(&getContext(), workspace_transformation_patterns); + if(failed(mlir::applyPatternsAndFoldGreedily(getOperation(), std::move(workspace_transformation_patterns)))) + { + signalPassFailure(); + } comet_debug() << __FILE__ << " " << __LINE__ << " ending CompressedWorkspaceTransforms pass \n"; } diff --git a/lib/Dialect/TensorAlgebra/CMakeLists.txt b/lib/Dialect/TensorAlgebra/CMakeLists.txt index 012a0c91..209b7767 100644 --- a/lib/Dialect/TensorAlgebra/CMakeLists.txt +++ b/lib/Dialect/TensorAlgebra/CMakeLists.txt @@ -1,14 +1,15 @@ add_llvm_library(COMETTensorAlgebraDialect IR/TADialect.cpp - # IR/TATypes.cpp Transforms/Transforms.cpp Transforms/LinalgTransforms.cpp Transforms/TCtoTTGT.cpp + Transforms/TCtoTTGTDyn.cpp Transforms/Passes.cpp Transforms/CheckImplicitTensorDecls.cpp Transforms/TensorDeclLowering.cpp + Transforms/WorkspaceOptimizations.cpp ADDITIONAL_HEADER_DIRS ${COMET_MAIN_INCLUDE_DIR}/comet/Dialect/TensorAlgebra @@ -16,6 +17,7 @@ add_llvm_library(COMETTensorAlgebraDialect add_dependencies( COMETTensorAlgebraDialect + COMETTensorAlgebraTypesIncGen COMETTensorAlgebraOpsIncGen COMETTensorAlgebraPassIncGen ) diff --git a/lib/Dialect/TensorAlgebra/IR/TADialect.cpp b/lib/Dialect/TensorAlgebra/IR/TADialect.cpp index 7eba8c86..c87c57db 100644 --- a/lib/Dialect/TensorAlgebra/IR/TADialect.cpp +++ b/lib/Dialect/TensorAlgebra/IR/TADialect.cpp @@ -26,12 +26,21 @@ // //===----------------------------------------------------------------------===// #include +#include #include "comet/Dialect/TensorAlgebra/IR/TADialect.h" #include "mlir/Dialect/Arith/IR/Arith.h" +#include "mlir/Dialect/Func/IR/FuncOps.h" +#include "mlir/Dialect/MemRef/IR/MemRef.h" +#include "mlir/IR/BuiltinTypeInterfaces.h" +#include "mlir/IR/BuiltinTypes.h" #include "mlir/IR/DialectImplementation.h" #include "mlir/Interfaces/FunctionImplementation.h" #include "mlir/IR/OpImplementation.h" +#include "mlir/Support/LogicalResult.h" +#include "mlir/Dialect/Bufferization/IR/DstBufferizableOpInterfaceImpl.h" +#include "mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h" +#include "llvm/ADT/TypeSwitch.h" using namespace mlir; using namespace mlir::tensorAlgebra; @@ -42,6 +51,10 @@ using namespace mlir::tensorAlgebra; /// TADialect //===----------------------------------------------------------------------===// + +//===----------------------------------------------------------------------===// +/// ConstantOp + /// Build a constant operation. /// The builder is passed as an argument, so is the state that this method is /// expected to fill in order to build the operation. @@ -87,13 +100,13 @@ mlir::LogicalResult DenseConstantOp::verify() { /// If the return type of the constant is not an unranked tensor, the shape /// must match the shape of the attribute holding the data. - auto resultType = getResult().getType().dyn_cast(); + auto resultType = dyn_cast(getResult().getType()); if (!resultType) return success(); /// Check that the rank of the attribute type matches the rank of the constant /// result type. - auto attrType = getValue().getType().cast(); + auto attrType = mlir::cast(getValue().getType()); if (attrType.getRank() != resultType.getRank()) { if(!(attrType.getRank() == 1 && attrType.getDimSize(0) == 1)) @@ -237,8 +250,8 @@ mlir::LogicalResult TAReturnOp::verify() auto resultType = results.front(); /// Check that the result type of the function matches the operand type. - if (inputType == resultType || inputType.isa() || - resultType.isa()) + if (inputType == resultType || isa(inputType) || + isa(resultType)) return mlir::success(); return emitError() << "type of return operand (" << inputType @@ -259,185 +272,246 @@ void TensorDimOp::build(mlir::OpBuilder &builder, mlir::OperationState &result, build(builder, result, builder.getIndexType(), source, indexValue); } + +/// Sort Op +bool hasFuncDeclaration(ModuleOp &module, std::string funcName) + { + for (auto func : module.getOps()) + { + StringAttr func_name = func.getSymNameAttr(); + if (funcName == func_name.getValue()) + return true; + } + return false; + } + +struct SortOpInterface + : public bufferization::DstBufferizableOpInterfaceExternalModel { + LogicalResult bufferize(Operation *op, RewriterBase &rewriter, + const bufferization::BufferizationOptions &options) const { + auto sortOp = cast(op); + auto loc = op->getLoc(); + auto ctx = rewriter.getContext(); + auto module = op->getParentOfType(); + + FailureOr destMemref = getBuffer(rewriter, sortOp.getTensor(), options); + if (failed(destMemref)) + return failure(); + + auto inputType = op->getOperand(0).getType(); + ShapedType shaped_type = mlir::cast(inputType); + + /// Declare comet_sort_index() + Type elemType = shaped_type.getElementType(); + Type indexType = rewriter.getIndexType(); + auto sort_func = FunctionType::get(ctx, {UnrankedMemRefType::get(elemType, 0), indexType, indexType}, {}); + std::string func_name = "comet_sort"; + if(auto integer = dyn_cast(elemType)) + { + func_name += std::to_string(integer.getIntOrFloatBitWidth()); + } + + if (!hasFuncDeclaration(module, func_name)) + { + func::FuncOp func_declare = func::FuncOp::create(loc, + func_name, + sort_func, + ArrayRef{}); + func_declare.setPrivate(); + module.push_back(func_declare); + } + + + auto unrankedMemrefType = mlir::UnrankedMemRefType::get(shaped_type.getElementType(), 0); + auto unrankedMemref = rewriter.create(loc, unrankedMemrefType, *destMemref); + rewriter.create(loc, func_name, SmallVector{}, ValueRange{unrankedMemref, sortOp.getStart(), sortOp.getEnd()}); + bufferization::replaceOpWithBufferizedValues(rewriter, op, *destMemref); + + return success(); + } +}; + //===----------------------------------------------------------------------===// /// TA Types //===----------------------------------------------------------------------===// -namespace mlir +// Implements the shaped type interface for the workspace type +ShapedType WorkspaceType::cloneWith(std::optional> shape, Type elementType) const { - namespace tensorAlgebra - { - namespace detail - { - /// This class represents the internal storage of the tensorAlgebra `SparseTensorType`. - struct SparseTensorTypeStorage : public mlir::TypeStorage - { - /// The `KeyTy` is a required type that provides an interface for the storage - /// instance. This type will be used when uniquing an instance of the type - /// storage. For our struct type, we will unique each instance structurally on - /// the elements that it contains. - using KeyTy = llvm::ArrayRef; - - /// A constructor for the type storage instance. - SparseTensorTypeStorage(llvm::ArrayRef elementTypes) - : elementTypes(elementTypes) {} - - /// Define the comparison function for the key type with the current storage - /// instance. This is used when constructing a new instance to ensure that we - /// haven't already uniqued an instance of the given key. - bool operator==(const KeyTy &key) const { return key == elementTypes; } - - /// Define a hash function for the key type. This is used when uniquing - /// instances of the storage, see the `SparseTensorType::get` method. - /// Note: This method isn't necessary as both llvm::ArrayRef and mlir::Type - /// have hash functions available, so we could just omit this entirely. - static llvm::hash_code hashKey(const KeyTy &key) - { - return llvm::hash_value(key); - } - - /// Define a construction function for the key type from a set of parameters. - /// These parameters will be provided when constructing the storage instance - /// itself. - /// Note: This method isn't necessary because KeyTy can be directly - /// constructed with the given parameters. - static KeyTy getKey(llvm::ArrayRef elementTypes) - { - return KeyTy(elementTypes); - } - - /// Define a construction method for creating a new instance of this storage. - /// This method takes an instance of a storage allocator, and an instance of a - /// `KeyTy`. The given allocator must be used for *all* necessary dynamic - /// allocations used to create the type storage and its internal. - static SparseTensorTypeStorage *construct(mlir::TypeStorageAllocator &allocator, - const KeyTy &key) - { - /// Copy the elements from the provided `KeyTy` into the allocator. - llvm::ArrayRef elementTypes = allocator.copyInto(key); - - /// Allocate the storage instance and construct it. - return new (allocator.allocate()) - SparseTensorTypeStorage(elementTypes); - } - - /// The following field contains the element types of the struct. - llvm::ArrayRef elementTypes; - }; - - } /// end namespace detail - } /// end namespace tensoralgebra -} /// end namespace mlir - -/// Create an instance of a `SparseTensorType` with the given element types. There -/// *must* be at least one element type. -SparseTensorType SparseTensorType::get(llvm::ArrayRef elementTypes) + // TODO: This may (?) require converting dimensions? Not sure + assert(false && "Workspace tensor cannot not be closed into another type"); + return NULL; +} + +bool WorkspaceType::hasRank() const { - assert(!elementTypes.empty() && "expected at least 1 element type"); - - /// Call into a helper 'get' method in 'TypeBase' to get a uniqued instance - /// of this type. The first two parameters are the context to unique in and the - /// kind of the type. The parameters after the type kind are forwarded to the - /// storage instance. - mlir::MLIRContext *ctx = elementTypes.front().getContext(); - return Base::get(ctx, elementTypes); + return true; } -/// Returns the element types of this sparse tensor type. -llvm::ArrayRef SparseTensorType::getElementTypes() +llvm::ArrayRef WorkspaceType::getShape() const { - /// 'getImpl' returns a pointer to the internal storage instance. - return getImpl()->elementTypes; + return getDims(); } -Type mlir::tensorAlgebra::TADialect::parseType(DialectAsmParser &parser) const +// Implements the shaped type interface for the sparse tensor type +ShapedType SparseTensorType::cloneWith(std::optional> shape, Type elementType) const { - /// Parse the main keyword for the type. - StringRef keyword; - /// for "indexlabel" and "spTensor" type - if (parser.parseKeyword(&keyword)) - return Type(); + // TODO: This may (?) require converting dimensions? Not sure + assert(false && "Sparse tensor cannot not be closed into another type"); + return NULL; +} - MLIRContext *context = getContext(); +bool SparseTensorType::hasRank() const +{ + return true; +} - /// Handle 'range' types. - if (keyword == "indexlabel") - { - return IndexLabelType::get(context); +llvm::ArrayRef SparseTensorType::getShape() const +{ + return getDims(); +} + +::mlir::Type SparseTensorType::parse(::mlir::AsmParser &odsParser) { + ::mlir::FailureOr<::mlir::Type> _result_element_type; + ::mlir::FailureOr<::mlir::IntegerType> indices_type; + SmallVector _result_dims; + SmallVector result_formats; + + // Parse literal '<' + if (odsParser.parseLess()) return {}; + + // Parse variable 'element_type' + _result_element_type = FieldParser::parse(odsParser); + if (failed(_result_element_type)) { + odsParser.emitError(odsParser.getCurrentLocation(), "failed to parse SparseTensor parameter 'element_type' which is to be a `Type`"); + return {}; } - /// Parse the element types of the spTensor. - if (keyword == "spTensor") - { - if (parser.parseLess()) - { - return Type(); - } + // Parse literal ',' + if (odsParser.parseComma()) return {}; - SmallVector elementTypes; - do - { - /// Parse the current element type. - llvm::SMLoc typeLoc = parser.getCurrentLocation(); - mlir::Type elementType; - - if (parser.parseType(elementType)) - return nullptr; - - /// Check that the type is either a TensorType or another SparseTensorType. - if (!elementType.isa()) - { - parser.emitError(typeLoc, "element type for a struct must either " - "be a TensorType or a SparseTensorType, got: ") - << elementType; - return Type(); - } - elementTypes.push_back(elementType); - - /// Parse the optional: `,` - } while (succeeded(parser.parseOptionalComma())); - - /// Parse: `>` - if (parser.parseGreater()) - return Type(); - - return SparseTensorType::get(elementTypes); + // Parse variable 'indices_type' + indices_type = ::mlir::FieldParser<::mlir::IntegerType>::parse(odsParser); + if (::mlir::failed(indices_type)) { + odsParser.emitError(odsParser.getCurrentLocation(), "failed to parse SparseTensor parameter 'indices_type' which is to be a `::mlir::IntegerType`"); + return {}; } + // Parse literal ',' + if (odsParser.parseComma()) return {}; - parser.emitError(parser.getNameLoc(), - "unknown TensorAlgebra type: " + keyword); - return Type(); + // Parse variable 'dims' + if(odsParser.parseDimensionList(_result_dims, true, false)) { + odsParser.emitError(odsParser.getCurrentLocation(), "failed to parse SparseTensor parameter 'dims' which is to be a `ArrayRef`"); + return {}; + } + // Parse literal ',' + if (odsParser.parseComma()) return {}; + + // Parse variable 'format' + do { + auto format = ::mlir::FieldParser::parse(odsParser); + if (mlir::failed(format)) // Parse each integer + return {}; + result_formats.push_back(*format); + } while (odsParser.parseOptionalComma().succeeded()); + + // Parse literal '>' + if (odsParser.parseGreater()) return {}; + + return SparseTensorType::get(odsParser.getContext(), + ::mlir::Type((*_result_element_type)), + ::mlir::IntegerType((*indices_type)), + ::llvm::ArrayRef((_result_dims)), + ::llvm::ArrayRef((result_formats))); } -/// IndexLabelType prints as just "indexlabel". -static void print(IndexLabelType type, DialectAsmPrinter &printer) -{ - printer << "indexlabel"; + +void SparseTensorType::print(::mlir::AsmPrinter &odsPrinter) const { + ::mlir::Builder odsBuilder(getContext()); + odsPrinter << "<"; + odsPrinter.printStrippedAttrOrType(getElementType()); + odsPrinter << ","; + odsPrinter << ' '; + odsPrinter.printStrippedAttrOrType(getIndicesType()); + odsPrinter << ","; + odsPrinter << ' '; + odsPrinter.printDimensionList(getDims()); + odsPrinter << ","; + odsPrinter << ' '; + odsPrinter.printStrippedAttrOrType(getFormat()); + odsPrinter << ">"; } -void mlir::tensorAlgebra::TADialect::printType( - Type type, DialectAsmPrinter &printer) const -{ - if (type.isa()) - { - print(type.cast(), printer); +::mlir::Type WorkspaceType::parse(::mlir::AsmParser &odsParser) { + ::mlir::FailureOr<::mlir::Type> _result_element_type; + ::mlir::FailureOr<::mlir::IntegerType> indices_type; + SmallVector _result_dims; + SmallVector result_formats; + + // Parse literal '<' + if (odsParser.parseLess()) return {}; + + // Parse variable 'element_type' + _result_element_type = FieldParser::parse(odsParser); + if (failed(_result_element_type)) { + odsParser.emitError(odsParser.getCurrentLocation(), "failed to parse SparseTensor parameter 'element_type' which is to be a `Type`"); + return {}; } - else if (type.isa()) - { - /// Currently the only sparse tensor type is a struct type. - SparseTensorType sparseTensorType = type.cast(); - /// Print the struct type according to the parser format. - printer << "spTensor<"; - llvm::interleaveComma(sparseTensorType.getElementTypes(), printer); - printer << '>'; + // Parse literal ',' + if (odsParser.parseComma()) return {}; + + // Parse variable 'indices_type' + indices_type = ::mlir::FieldParser<::mlir::IntegerType>::parse(odsParser); + if (::mlir::failed(indices_type)) { + odsParser.emitError(odsParser.getCurrentLocation(), "failed to parse SparseTensor parameter 'indices_type' which is to be a `::mlir::IntegerType`"); + return {}; } - else - { - llvm_unreachable("Unhandled TensorAlgebra type"); + // Parse literal ',' + if (odsParser.parseComma()) return {}; + + // Parse variable 'dims' + if(odsParser.parseDimensionList(_result_dims, true, false)) { + odsParser.emitError(odsParser.getCurrentLocation(), "failed to parse SparseTensor parameter 'dims' which is to be a `ArrayRef`"); + return {}; } + + // Parse literal '>' + if (odsParser.parseGreater()) return {}; + + return WorkspaceType::get(odsParser.getContext(), + ::mlir::Type((*_result_element_type)), + ::mlir::IntegerType((*indices_type)), + ::llvm::ArrayRef((_result_dims))); +} + + +void WorkspaceType::print(::mlir::AsmPrinter &odsPrinter) const { + ::mlir::Builder odsBuilder(getContext()); + odsPrinter << "<"; + odsPrinter.printStrippedAttrOrType(getElementType()); + odsPrinter << ","; + odsPrinter << ' '; + odsPrinter.printStrippedAttrOrType(getIndicesType()); + odsPrinter << ","; + odsPrinter << ' '; + odsPrinter.printDimensionList(getDims()); + odsPrinter << ">"; } + +//===----------------------------------------------------------------------===// +/// TableGen'd type definitions +//===----------------------------------------------------------------------===// +#define GET_TYPEDEF_CLASSES +#include "comet/Dialect/TensorAlgebra/IR/TATypes.cpp.inc" + +//===----------------------------------------------------------------------===// +/// TableGen'd enum definitions +//===----------------------------------------------------------------------===// +#include "comet/Dialect/TensorAlgebra/IR/TAEnums.cpp.inc" + //===----------------------------------------------------------------------===// /// TableGen'd op method definitions //===----------------------------------------------------------------------===// @@ -453,10 +527,31 @@ void mlir::tensorAlgebra::TADialect::printType( /// the point of registration of types and operations for the dialect. void TADialect::initialize() { + addTypes< +#define GET_TYPEDEF_LIST +#include "comet/Dialect/TensorAlgebra/IR/TATypes.cpp.inc" + >(); + + addAttributes< +#define GET_ATTRDEF_LIST +#include "comet/Dialect/TensorAlgebra/IR/TAAttrs.cpp.inc" + >(); + addOperations< #define GET_OP_LIST #include "comet/Dialect/TensorAlgebra/IR/TAOps.cpp.inc" >(); - // addTypes(); - addTypes(); + + // declarePromisedInterface(); /// PT: I think this is not necessary +} + +namespace mlir { + namespace tensorAlgebra { + void registerBufferizableOpInterfaceExternalModels(DialectRegistry ®istry) + { + registry.addExtension(+[](MLIRContext *ctx, TADialect *dialect) { + TensorSortOp::attachInterface(*ctx); + }); + } + } } diff --git a/lib/Dialect/TensorAlgebra/Transforms/CheckImplicitTensorDecls.cpp b/lib/Dialect/TensorAlgebra/Transforms/CheckImplicitTensorDecls.cpp index e05963e9..9f19069e 100644 --- a/lib/Dialect/TensorAlgebra/Transforms/CheckImplicitTensorDecls.cpp +++ b/lib/Dialect/TensorAlgebra/Transforms/CheckImplicitTensorDecls.cpp @@ -64,8 +64,7 @@ template bool isNeedTensorDecl(t op) { bool isUsedInSetSource = true; - std::string op_str = dump2str(op); - mlir::Value result = op.getOperation()->getResult(0); + mlir::Value result = op->getResult(0); comet_debug() << " "; comet_vdump(result); for (auto u1 : result.getUsers()) @@ -75,20 +74,10 @@ bool isNeedTensorDecl(t op) /// If not used as source tensor of set_op, it is tmp result if (isa(u1)) { - comet_debug() << " used in ta.set_new op\n"; auto p = cast(u1).getOperation(); - for (unsigned int i = 0; i < p->getNumOperands(); i++) + if(result == p->getOperand(0)) { - comet_debug() << " the " << i << "th operand\n"; - std::string n_str = dump2str(p->getOperand(i)); - if (n_str.compare(0, op_str.size(), op_str) == 0) - { - comet_debug() << " FIND IT: " << i << "\n"; - if (i == 0) - { /// used as source tensor - isUsedInSetSource = false; - } - } + isUsedInSetSource = false; } } } @@ -103,7 +92,6 @@ void addTensorDecl(t op) op->getOperands(); std::vector lbls_value; mlir::Value ret_value; - std::string ret_format; ArrayAttr imaps = op.getIndexingMaps(); /// Retrieve the size of the tensor based on the affinity maps @@ -112,12 +100,12 @@ void addTensorDecl(t op) /// 3. The result will be the dimension of the matching indices /// e.g For matmulop RHS1: (d0,d1,d2) -> (d0, d1), RHS2: (d0,d1,d2) -> (d1, d2), LHS : (d0,d1,d2) -> (d0,d2) /// So we get (dim(RHS1, 0), dim(RHS2, 1)) because of d0, d2 respectively - auto res_map = imaps[imaps.size() - 1].cast().getValue(); + auto res_map = cast(imaps[imaps.size() - 1]).getValue(); for (auto v : res_map.getResults()) { for (size_t i = 0; i < imaps.size() - 1; i++) { - auto map = imaps[i].cast().getValue(); + auto map = cast(imaps[i]).getValue(); if (auto pos = map.getResultPosition(v)) { lbls_value.push_back(builder.create(location, op->getOperand(i), *pos)); @@ -126,22 +114,28 @@ void addTensorDecl(t op) } ret_value = op.getOperation()->getResult(0); - mlir::ArrayAttr opFormatsArrayAttr = op.getFormats(); - unsigned int i = opFormatsArrayAttr.size() - 1; - std::string ret_format_local(opFormatsArrayAttr[i].cast().getValue()); - ret_format = ret_format_local; + // mlir::ArrayAttr opFormatsArrayAttr = op.getFormats(); + // unsigned int i = opFormatsArrayAttr.size() - 1; + // std::string ret_format_local(cast(opFormatsArrayAttr[i]).getValue()); + // ret_format = ret_format_local; mlir::Value itensor; - if (ret_format.compare("Dense") == 0) + if (isa(ret_value.getType())) { - itensor = builder.create(location, ret_value.getType(), lbls_value, ret_format); + itensor = builder.create(location, ret_value.getType(), lbls_value); builder.create(location, itensor, builder.getF64FloatAttr(0)); } - else + else if (isa(ret_value.getType())) { /// It is a temporal tensor declaration generated by compound expressions, BoolAttr is true /// to identify SparseTensorDeclOp is for temporaries - itensor = builder.create(location, ret_value.getType(), lbls_value, ret_format, true); + itensor = builder.create(location, ret_value.getType(), lbls_value, true); + } + else + { + comet_debug() << "Error: cannot handle the tensor type for TensorDeclLowering\n"; + op.emitOpError("cannot handle the tensor type for TensorDeclLowering"); + return; } comet_debug() << "PreLowering SparseTensorDeclaration creation\n"; comet_vdump(itensor); @@ -173,41 +167,36 @@ void TensorAlgebraCheckImplicitTensorDeclPass::runOnOperation() comet_debug() << " find a transpose op\n"; /// if the output is not used as a source tensor of a set op /// Need to store use a sparse/dense tensor decl op to store the result - if (isa(cur_op)) + if (auto op = dyn_cast(cur_op)) { - auto op = cast(cur_op); if (isNeedTensorDecl(op)) { addTensorDecl(op); } } - else if (isa(cur_op)) + else if (auto op = dyn_cast(cur_op)) { - auto op = cast(cur_op); if (isNeedTensorDecl(op)) { addTensorDecl(op); } } - else if (isa(cur_op)) + else if (auto op = dyn_cast(cur_op)) { - auto op = cast(cur_op); if (isNeedTensorDecl(op)) { addTensorDecl(op); } } - else if (isa(cur_op)) + else if (auto op = dyn_cast(cur_op)) { - auto op = cast(cur_op); if (isNeedTensorDecl(op)) { addTensorDecl(op); } } - else if (isa(cur_op)) + else if (auto op = dyn_cast(cur_op)) { - auto op = cast(cur_op); if (isNeedTensorDecl(op)) { addTensorDecl(op); diff --git a/lib/Dialect/TensorAlgebra/Transforms/LinalgTransforms.cpp b/lib/Dialect/TensorAlgebra/Transforms/LinalgTransforms.cpp index 6499e1da..ac27b260 100644 --- a/lib/Dialect/TensorAlgebra/Transforms/LinalgTransforms.cpp +++ b/lib/Dialect/TensorAlgebra/Transforms/LinalgTransforms.cpp @@ -30,16 +30,30 @@ #include "comet/Dialect/Utils/Utils.h" #include "mlir/Dialect/Affine/IR/AffineOps.h" +#include "mlir/Dialect/Affine/LoopUtils.h" +#include "mlir/Dialect/Arith/IR/Arith.h" #include "mlir/Dialect/Linalg/IR/Linalg.h" #include "mlir/Dialect/Linalg/Utils/Utils.h" #include "mlir/Dialect/Linalg/Transforms/Transforms.h" #include "mlir/Dialect/Linalg/TransformOps/LinalgTransformOps.h" +#include "mlir/Dialect/MemRef/IR/MemRef.h" +#include "mlir/Dialect/SCF/IR/SCF.h" #include "mlir/Dialect/SCF/Transforms/TileUsingInterface.h" #include "mlir/Dialect/Func/IR/FuncOps.h" #include "mlir/Dialect/LLVMIR/LLVMDialect.h" +#include "mlir/Dialect/Tensor/IR/Tensor.h" +#include "mlir/IR/BuiltinAttributes.h" +#include "mlir/IR/BuiltinTypeInterfaces.h" +#include "mlir/IR/BuiltinTypes.h" +#include "mlir/IR/OpDefinition.h" #include "mlir/IR/PatternMatch.h" +#include "mlir/IR/Value.h" +#include "mlir/IR/ValueRange.h" #include "mlir/Pass/Pass.h" #include "mlir/Transforms/GreedyPatternRewriteDriver.h" +#include "llvm/ADT/SmallVector.h" +#include +#include // suppress all warnings coming from inclusion of blis.h in source tree #ifdef __clang__ @@ -218,7 +232,7 @@ namespace return failure(); FailureOr tilingResult = - scf::tileUsingSCFForOp(rewriter, op, options); + scf::tileUsingSCF(rewriter, op, options); if (failed(tilingResult)) return rewriter.notifyMatchFailure(op, "failed to tile operation"); @@ -263,12 +277,15 @@ static void addPatternForTiling(MLIRContext *context, StringRef filterName, StringRef updatedFilterName, ArrayRef tileSizes, + bool parallel = false, ArrayRef interchange = {}) { scf::SCFTilingOptions tilingOptions; SmallVector tileSizesOfr = getAsIndexOpFoldResult(context, tileSizes); tilingOptions.setTileSizes(tileSizesOfr).setInterchange(interchange); + tilingOptions.setLoopType( + parallel ? scf::SCFTilingOptions::LoopType::ForallOp: scf::SCFTilingOptions::LoopType::ForOp); LinalgTransformationFilter filter(StringAttr::get(context, filterName), StringAttr::get(context, updatedFilterName)); patterns.add(context, tilingOptions, filter); @@ -289,8 +306,8 @@ namespace int mc, kc, nc, mr, nr = 0; get_level3_blocksizes(&mc, &kc, &nc, &mr, &nr, sizeof(double)); - addPatternForTiling(ctx, tilingPatterns, "__with_tiling__", "__L2__with_tiling__", {mc, nc, kc}, {1, 2, 0}); - addPatternForTiling(ctx, tilingPatterns, "__L2__with_tiling__", "__micro_kernel__", {mr, nr, kc}, {1, 0, 2}); + addPatternForTiling(ctx, tilingPatterns, "__with_tiling__", "__L2__with_tiling__", {mc, nc}, false, {1, 0}); + addPatternForTiling(ctx, tilingPatterns, "__L2__with_tiling__", "__micro_kernel__", {mr, nr}, true, {1, 0}); if (failed(applyPatternsAndFoldGreedily(getOperation(), std::move(tilingPatterns)))) @@ -333,8 +350,13 @@ static SmallVector extractOperandTypes(Operation *op) /// The underlying descriptor type (e.g. LLVM) does not have layout /// information. Canonicalizing the type at the level of std when going into /// a library call avoids needing to introduce DialectCastOp. - if (auto memrefType = type.dyn_cast()) - result.push_back(makeStridedLayoutDynamic(memrefType)); + if (auto memrefType = dyn_cast(type)) + { + auto newShape = llvm::to_vector<4>(llvm::map_range(memrefType.getShape(), [](int64_t dimSize) { + return ShapedType::kDynamic; + })); + result.push_back(makeStridedLayoutDynamic(MemRefType::get(newShape, memrefType.getElementType(), memrefType.getLayout(), memrefType.getMemorySpace()))); + } else result.push_back(type); } @@ -350,14 +372,17 @@ createTypeCanonicalizedMemRefOperands(OpBuilder &b, Location loc, res.reserve(operands.size()); for (auto op : operands) { - auto memrefType = op.getType().dyn_cast(); + auto memrefType = dyn_cast(op.getType()); if (!memrefType) { res.push_back(op); continue; } + auto newShape = llvm::to_vector<4>(llvm::map_range(memrefType.getShape(), [](int64_t dimSize) { + return ShapedType::kDynamic; + })); Value cast = - b.create(loc, makeStridedLayoutDynamic(memrefType), op); + b.create(loc, makeStridedLayoutDynamic(MemRefType::get(newShape, memrefType.getElementType(), memrefType.getLayout(), memrefType.getMemorySpace())), op); res.push_back(cast); } return res; @@ -383,6 +408,9 @@ static FailureOr getLibraryCallSymbolRef(Operation *op, Patte /// Add inputTypes for mr and nr inputTypes.push_back(IntegerType::get(rewriter.getContext(), 32)); inputTypes.push_back(IntegerType::get(rewriter.getContext(), 32)); + /// Add inputTypes for alpha and beta + inputTypes.push_back(rewriter.getF64Type()); + inputTypes.push_back(rewriter.getF64Type()); if (op->getNumResults() != 0) { @@ -418,12 +446,32 @@ LogicalResult LinalgMatMulOpToLibraryCallPattern::matchAndRewrite( if (failed(libraryCallName)) return failure(); + const char *arch = bli_arch_string(bli_cpuid_query_id()); + + if (!((strcmp("haswell", arch) == 0) || + (strcmp("zen", arch) == 0) || + (strcmp("zen2", arch) == 0) || + (strcmp("zen3", arch) == 0) || + (strcmp("skx", arch) == 0) || + (strcmp("knl", arch) == 0) || + (strcmp("generic", arch) == 0) || + (strcmp("firestorm", arch) == 0))) + { + assert(false && "Unsupported microkernel architecture for BLIS matmul" + " operation. Supported architectures are: haswell, zen, " + "zen2, zen3, skx, knl, firestorm, generic."); + return failure(); + } + int mc, kc, nc, mr, nr = 0; get_level3_blocksizes(&mc, &kc, &nc, &mr, &nr, sizeof(double)); IntegerType i32Type = IntegerType::get(rewriter.getContext(), 32); Value mrValue = rewriter.create(op->getLoc(), i32Type, rewriter.getIntegerAttr(i32Type, mr)); Value nrValue = rewriter.create(op->getLoc(), i32Type, rewriter.getIntegerAttr(i32Type, nr)); + Value alpha = rewriter.create(op->getLoc(), rewriter.getF64Type(), op->getAttrOfType("__alpha__")); + Value beta = rewriter.create(op->getLoc(), rewriter.getF64Type(), op->getAttrOfType("__beta__")); + comet_vdump(mrValue); comet_vdump(nrValue); @@ -432,6 +480,8 @@ LogicalResult LinalgMatMulOpToLibraryCallPattern::matchAndRewrite( operands.insert(operands.end(), op->getOperands().begin(), op->getOperands().end()); operands.push_back(mrValue); operands.push_back(nrValue); + operands.push_back(alpha); + operands.push_back(beta); rewriter.replaceOpWithNewOp( op, libraryCallName->getValue(), TypeRange(), @@ -469,6 +519,187 @@ struct OptDenseTranspose : public ConversionPattern : ConversionPattern(linalg::TransposeOp::getOperationName(), 1, ctx), tile_size(tile_size), seperate_tiles(seperate_tiles) {} + // Old pass, utilizing memrefs but also supports tile_size and separate_tiles arguments +#if 0 + LogicalResult + matchAndRewrite(Operation *input_op, ArrayRef operands, + ConversionPatternRewriter &rewriter) const final + { + comet_debug() << " OptDenseTranspose : public ConversionPattern\n"; + auto op = dyn_cast(input_op); + comet_debug() << " Lowering dense transpose\n"; + assert(isa(op) && + "this operation is not a linalg.transpose"); + + auto ctx = rewriter.getContext(); + Builder builder(ctx); + Location loc = op.getLoc(); + + //auto module = op->getParentOfType(); + comet_vdump(op); + + auto tensorize_input = op->getOperand(0).getDefiningOp(); + auto inputMemref = tensorize_input->getOperand(0); + auto inputType = inputMemref.getType(); + + auto tensorize_output = op->getOperand(1).getDefiningOp(); + auto outputMemref = tensorize_output->getOperand(0); + + comet_debug() << " Input Type:\n"; + comet_vdump(inputType); + auto inputRank = inputType.cast().getRank(); + + std::vector loops; + std::vector indexIterateOrder; + for (int64_t rank = 0; rank < inputRank; rank++) + { + indexIterateOrder.push_back(rank); + int64_t upperBound = inputType.cast().getDimSize(rank); + if (upperBound == ShapedType::kDynamic) + { + assert(false && "TODO: This dimension is a dynamic size"); + } + /// create for loops + auto loop = rewriter.create(loc, 0, upperBound, 1); + loops.push_back(loop); + comet_vdump(loop); + rewriter.setInsertionPointToStart(loop.getBody()); + } + + AffineMap inputIndexingMap = builder.getMultiDimIdentityMap(inputRank); + auto inputIndices = getReassociationIndices(inputIndexingMap); + auto inputIVs = createInductionVarAffine(loops, indexIterateOrder, inputIndices); + + AffineMap outputIndexingMap = AffineMap::getPermutationMap(llvm::to_vector_of(op.getPermutation()), ctx); + SmallVector outputIndices = + getReassociationIndices(outputIndexingMap); + auto outputIVs = createInductionVarAffine(loops, indexIterateOrder, outputIndices); + + /// Build loop body + auto load_rhs = rewriter.create(loc, inputMemref, inputIVs); + comet_vdump(load_rhs); + #ifdef COMET_DEBUG_MODE + comet_vdump(load_rhs); + auto store_lhs = rewriter.create(loc, load_rhs, outputMemref, outputIVs); + comet_vdump(store_lhs); + #else + rewriter.create(loc, load_rhs, outputMemref, outputIVs); + #endif + + /// TransposeOp index permutation + ArrayRef invresults = inputIndexingMap.getResults(); + std::vector sourceOrder; + for (auto a : invresults) + { + if (a.getKind() == AffineExprKind::DimId) + { + AffineDimExpr *b = (AffineDimExpr *)&a; /// down_casting + sourceOrder.push_back(b->getPosition()); + comet_debug() << "Source order: " << b->getPosition() << "\n"; + } + } + + /// AffineMap outvmap = op.getPermutation().getValue(); + ArrayRef outvresults = outputIndexingMap.getResults(); + /// From outer to inner, the destOrder[size -1] is the most important, + std::vector destOrder; + for (auto a : outvresults) + { + if (a.getKind() == AffineExprKind::DimId) + { + AffineDimExpr *b = (AffineDimExpr *)&a; /// down_casting + destOrder.push_back(b->getPosition()); + comet_debug() << "destination order: " << b->getPosition() << "\n"; + } + } + + if (loops.size() > 0) + { + /* Suppose Given best order: a0, a3, a1, a2 + ** Then first step: a0, a1, a3, a2 (exchange loop index 1 and 2) + ** Then second step: a0, a1, a2, a3 (exchange loop order index 2 and 3) + */ + std::vector optimalOrder = destOrder; + /// Call an getLoopOrder algorithm to get the best order + std::vector> loopOrders; + getLoopOrders(loopOrders, destOrder.size(), sourceOrder, destOrder); + optimalOrder = loopOrders[0]; + + std::vector currentOrder; + for (unsigned i = 0; i < destOrder.size(); i++) + { + currentOrder.push_back(i); + } + + for (unsigned i = 0; i < optimalOrder.size(); i++) + { + comet_debug() << "currentOrder[i]: " << currentOrder[i] << " optimalOrder[i]: " << optimalOrder[i] << "\n"; + /// This loop index is the correct loop index, no loop interchange + if (optimalOrder[i] == currentOrder[i]) + { + continue; + } + else + { /// Get the location of the right loop index + for (unsigned j = i + 1; j < currentOrder.size(); j++) + { + if (optimalOrder[i] == currentOrder[j]) + { /// loop j and i exchange + unsigned k = j; + /// k = (i,j]. k is unsigned, should be >= 0. use k-1, so k>=1 + while (k > 0 && k > i) + { + mlir::affine::interchangeLoops(loops[currentOrder[k - 1]], loops[currentOrder[k]]); + std::swap(currentOrder[k - 1], currentOrder[k]); + k--; + } + break; + } + } + } + } + + std::vector newLoops; + for (unsigned i = 0; i < currentOrder.size(); i++) + { + newLoops.push_back(loops[currentOrder[i]]); + } + loops.clear(); + + /// Possible to assign different tile size based on the dimension + if (tile_size > 1) + { + std::vector tileSizes; + for (unsigned i = 0; i < currentOrder.size(); i++) + { + tileSizes.push_back(tile_size); + } + SmallVector tiledNest; + if (failed(mlir::affine::tilePerfectlyNested(newLoops, tileSizes, &tiledNest))) + return failure(); + + comet_vdump(tiledNest[0]); + + /// Separate full and partial tiles. + if (seperate_tiles) + { + auto intraTileLoops = + MutableArrayRef(tiledNest).drop_front(newLoops.size()); + if (failed(separateFullTiles(intraTileLoops))) + return failure(); + } + } /// end if (tilesize > 1) + + } /// end loops.size() < 0 + + rewriter.replaceAllUsesWith(op->getResult(0), outputMemref); + rewriter.eraseOp(op); + + //module.dump(); + return success(); + } +#endif + LogicalResult matchAndRewrite(Operation *input_op, ArrayRef operands, ConversionPatternRewriter &rewriter) const final @@ -485,51 +716,21 @@ struct OptDenseTranspose : public ConversionPattern comet_vdump(op); - auto inputType = op->getOperand(0).getType(); + auto input = op.getInput(); + auto inputType = input.getType(); + + auto output = op.getInit(); + comet_debug() << " Input Type:\n"; comet_vdump(inputType); - auto inputMemref = op->getOperand(0); - auto inputRank = inputType.cast().getRank(); - auto outputMemref = op->getOperand(1); - - std::vector loops; - std::vector indexIterateOrder; - for (int64_t rank = 0; rank < inputRank; rank++) - { - indexIterateOrder.push_back(rank); - int64_t upperBound = inputType.cast().getDimSize(rank); - if (upperBound == ShapedType::kDynamic) - { - assert(false && "TODO: This dimension is a dynamic size"); - } - /// create for loops - auto loop = rewriter.create(loc, 0, upperBound, 1); - loops.push_back(loop); - comet_vdump(loop); - rewriter.setInsertionPointToStart(loop.getBody()); - } + auto inputRank = inputType.getRank(); AffineMap inputIndexingMap = builder.getMultiDimIdentityMap(inputRank); - auto inputIndices = getReassociationIndices(inputIndexingMap); - auto inputIVs = createInductionVarAffine(loops, indexIterateOrder, inputIndices); - AffineMap outputIndexingMap = AffineMap::getPermutationMap(llvm::to_vector_of(op.getPermutation()), ctx); SmallVector outputIndices = getReassociationIndices(outputIndexingMap); - auto outputIVs = createInductionVarAffine(loops, indexIterateOrder, outputIndices); - - /// Build loop body - auto load_rhs = rewriter.create(loc, inputMemref, inputIVs); - comet_vdump(load_rhs); -#ifdef DEBUG_MODE_LINALGTRANSFORMS - comet_vdump(load_rhs); - auto store_lhs = rewriter.create(loc, load_rhs, outputMemref, outputIVs); - comet_vdump(store_lhs); -#else - rewriter.create(loc, load_rhs, outputMemref, outputIVs); -#endif - /// TransposeOp index permutation + /// TransposeOp index permutation ArrayRef invresults = inputIndexingMap.getResults(); std::vector sourceOrder; for (auto a : invresults) @@ -542,7 +743,6 @@ struct OptDenseTranspose : public ConversionPattern } } - /// AffineMap outvmap = op.getPermutation().getValue(); ArrayRef outvresults = outputIndexingMap.getResults(); /// From outer to inner, the destOrder[size -1] is the most important, std::vector destOrder; @@ -556,94 +756,62 @@ struct OptDenseTranspose : public ConversionPattern } } - if (loops.size() > 0) + /* Suppose Given best order: a0, a3, a1, a2 + ** Then first step: a0, a1, a3, a2 (exchange loop index 1 and 2) + ** Then second step: a0, a1, a2, a3 (exchange loop order index 2 and 3) + */ + std::vector optimalOrder = destOrder; + /// Call an getLoopOrder algorithm to get the best order + std::vector> loopOrders; + getLoopOrders(loopOrders, destOrder.size(), sourceOrder, destOrder); + optimalOrder = loopOrders[0]; + + std::vector currentOrder; + for (unsigned i = 0; i < destOrder.size(); i++) { - /* Suppose Given best order: a0, a3, a1, a2 - ** Then first step: a0, a1, a3, a2 (exchange loop index 1 and 2) - ** Then second step: a0, a1, a2, a3 (exchange loop order index 2 and 3) - */ - std::vector optimalOrder = destOrder; - /// Call an getLoopOrder algorithm to get the best order - std::vector> loopOrders; - getLoopOrders(loopOrders, destOrder.size(), sourceOrder, destOrder); - optimalOrder = loopOrders[0]; - - std::vector currentOrder; - for (unsigned i = 0; i < destOrder.size(); i++) - { - currentOrder.push_back(i); - } - - for (unsigned i = 0; i < optimalOrder.size(); i++) - { - comet_debug() << "currentOrder[i]: " << currentOrder[i] << " optimalOrder[i]: " << optimalOrder[i] << "\n"; - /// This loop index is the correct loop index, no loop interchange - if (optimalOrder[i] == currentOrder[i]) - { - continue; - } - else - { /// Get the location of the right loop index - for (unsigned j = i + 1; j < currentOrder.size(); j++) - { - if (optimalOrder[i] == currentOrder[j]) - { /// loop j and i exchange - unsigned k = j; - /// k = (i,j]. k is unsigned, should be >= 0. use k-1, so k>=1 - while (k > 0 && k > i) - { - mlir::affine::interchangeLoops(loops[currentOrder[k - 1]], loops[currentOrder[k]]); - std::swap(currentOrder[k - 1], currentOrder[k]); - k--; - } - break; - } - } - } - } - - std::vector newLoops; - for (unsigned i = 0; i < currentOrder.size(); i++) - { - newLoops.push_back(loops[currentOrder[i]]); - } - loops.clear(); - - /// Possible to assign different tile size based on the dimension - if (tile_size > 1) - { - std::vector tileSizes; - for (unsigned i = 0; i < currentOrder.size(); i++) - { - tileSizes.push_back(tile_size); - } - SmallVector tiledNest; - if (failed(mlir::affine::tilePerfectlyNested(newLoops, tileSizes, &tiledNest))) - return failure(); + currentOrder.push_back(i); + } - comet_vdump(tiledNest[0]); + SmallVector in_ivs; + SmallVector out_ivs; + in_ivs.resize(optimalOrder.size()); + out_ivs.resize(outputIndices[0].size()); + OpFoldResult one = rewriter.createOrFold(loc, 1); + SmallVector ubs; + for (unsigned i = 0; i < optimalOrder.size(); i++) + { + Value upperBound = rewriter.create(loc, input, optimalOrder[i]); + ubs.push_back(upperBound); + } - /// Separate full and partial tiles. - if (seperate_tiles) - { - auto intraTileLoops = - MutableArrayRef(tiledNest).drop_front(newLoops.size()); - if (failed(separateFullTiles(intraTileLoops))) - return failure(); - } - } /// end if (tilesize > 1) + SmallVector ones(optimalOrder.size(), one); + auto forAll = rewriter.create(loc, ubs, output, std::nullopt); + rewriter.setInsertionPointToStart(forAll.getBody()); + auto ivs = forAll.getLoopInductionVars(); + for(size_t i = 0; i< forAll.getLoopInductionVars()->size(); i++) + { + in_ivs[optimalOrder[i]] = forAll.getLoopInductionVars()->data()[i]; + out_ivs[optimalOrder[outputIndices[0][i]]] = forAll.getLoopInductionVars()->data()[i]; + } + auto read_slice = rewriter.create(loc, input, in_ivs, ones, ones); + auto write_slice = rewriter.create(loc, forAll.getRegionIterArgs().front(), out_ivs, ones, ones); + SmallVector zeros_indices(in_ivs.size(), rewriter.create(loc, 0)); - } /// end loops.size() < 0 + auto extracted = rewriter.create(loc, read_slice, zeros_indices); + auto inserted = rewriter.create(loc, extracted, write_slice, zeros_indices); + rewriter.setInsertionPointToEnd(forAll.getTerminator().getBody()); + rewriter.create(loc, inserted, forAll.getRegionIterArgs().front(), out_ivs, ones, ones); + rewriter.replaceAllUsesWith(op->getResult(0), forAll->getResult(0)); rewriter.eraseOp(op); - /// module.dump(); + //module.dump(); return success(); } private: - uint64_t tile_size; - bool seperate_tiles; + [[maybe_unused]] uint64_t tile_size; + [[maybe_unused]] bool seperate_tiles; }; /// Lower Dense Transpose to loops after optimizations namespace @@ -658,7 +826,7 @@ namespace comet_debug() << "OptDenseTransposePass : public PassWrapper\n"; func::FuncOp func = getOperation(); ConversionTarget target(getContext()); - target.addLegalDialect(); + target.addLegalDialect(); RewritePatternSet patterns(&getContext()); patterns.insert(&getContext(), tile_size, seperate_tiles); @@ -674,6 +842,57 @@ namespace uint64_t tile_size; bool seperate_tiles; }; +} /// end anonym +// ous namespace +namespace +{ + struct MatvecToParallelLoops : public ConversionPattern + { + MatvecToParallelLoops(MLIRContext *ctx) + : ConversionPattern(linalg::MatvecOp::getOperationName(), 1, ctx) + {} + + LogicalResult + matchAndRewrite(Operation *input_op, ArrayRef operands, + ConversionPatternRewriter &rewriter) const final + { + + auto op = dyn_cast(input_op); + + if(failed(mlir::linalg::linalgOpToParallelLoops(rewriter, op))) + { + return mlir::failure(); + } + else{ + rewriter.eraseOp(input_op); + return success(); + } + } + }; + + class MatvecToParallelLoopsPass : public PassWrapper> + { + public: + MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(MatvecToParallelLoopsPass) + MatvecToParallelLoopsPass(){}; + void runOnOperation() override + { + comet_debug() << "MatvecToParallelLoopsPass : public PassWrapper\n"; + func::FuncOp func = getOperation(); + ConversionTarget target(getContext()); + target.addLegalDialect(); + RewritePatternSet patterns(&getContext()); + patterns.insert(&getContext() ); + + if (failed(applyPartialConversion(func, target, std::move(patterns)))) + { + llvm::errs() << "Failed to Lower dense transpose operation\n"; + signalPassFailure(); + } + comet_debug() << "MatvecToParallelLoopsPass done\n"; + } + + }; } /// end anonymous namespace /// Create a pass to optimize LinAlg Matmul Op with tiling @@ -697,3 +916,10 @@ std::unique_ptr mlir::comet::createOptDenseTransposePass(uint64_t ti comet_debug() << "LinAlgTransforms createOptDenseTransposePass\n"; return std::make_unique(tile_size, seperate_tiles); } + + +std::unique_ptr mlir::comet::createMatvecToParallelLoopsPass() +{ + comet_debug() << "LinAlgTransforms createMatvecToParallelLoopsPass\n"; + return std::make_unique(); +} \ No newline at end of file diff --git a/lib/Dialect/TensorAlgebra/Transforms/Passes.cpp b/lib/Dialect/TensorAlgebra/Transforms/Passes.cpp index 52f78848..f7ecdc3c 100644 --- a/lib/Dialect/TensorAlgebra/Transforms/Passes.cpp +++ b/lib/Dialect/TensorAlgebra/Transforms/Passes.cpp @@ -23,6 +23,7 @@ #include "comet/Dialect/TensorAlgebra/Passes.h" #include "comet/Dialect/Utils/Utils.h" +#include "mlir/IR/BuiltinTypes.h" #include "mlir/Pass/Pass.h" #include "mlir/Dialect/Affine/IR/AffineOps.h" #include "mlir/Dialect/Arith/IR/Arith.h" @@ -30,6 +31,8 @@ #include "mlir/Dialect/Linalg/IR/Linalg.h" #include "mlir/Dialect/MemRef/IR/MemRef.h" #include "mlir/Dialect/Bufferization/IR/Bufferization.h" +#include "llvm/ADT/StringRef.h" +#include "llvm/Support/Casting.h" #include #include @@ -136,9 +139,9 @@ std::vector findOutput(const std::vector &rhs1Labels, } std::tuple>, std::vector>> -optimalOrder(ArrayRef inLTOps, Operation *outLTOp, +optimalOrder(ArrayRef inLTOps, Operation *outLTOp, const std::map &lblSizes, - const std::map> &lblMaps) + const std::map> &lblMaps) { IndexVector result; for (size_t i = 0; i < inLTOps.size(); i++) @@ -228,13 +231,13 @@ void FindOptimalTCFactorizationPass::FindOptimalTCFactorization(tensorAlgebra::T std::vector MultOpsToRemove; std::vector LTOpsToRemove; - std::vector inLTOps; - std::map inLTValues; + std::vector inLTOps; + std::map inLTValues; comet_debug() << "Chain Multiplication Factorization begin...\n"; std::map lblSizes; std::map labelValues; - std::map> lblMaps; + std::map> lblMaps; /// collect all operands from series of ta.tc ops if (isa(lhsOp)) @@ -244,9 +247,9 @@ void FindOptimalTCFactorizationPass::FindOptimalTCFactorization(tensorAlgebra::T Value currValue = operands[0]; comet_vdump(currValue); Operation *curr = currValue.getDefiningOp(); - while (isa(curr) || !stack.empty()) + while ((curr && isa(curr)) || !stack.empty()) { - while (isa(curr)) + while (curr && isa(curr)) { auto multop = cast(curr); stack.push(curr); @@ -259,25 +262,36 @@ void FindOptimalTCFactorizationPass::FindOptimalTCFactorization(tensorAlgebra::T auto lblOp = labels[i].getDefiningOp(); if (lblSizes.count(lblOp) == 0) { + auto rhs2T = cast(multop.getRhs2().getType()); + /// If dynamic dimension, we need to retrieve the value of the constantIndexOp that was used to create it /// If static, just get it from the tensor type - if (isa(multop.getRhs2().getDefiningOp())) - { - if (multop.getRhs2().getType().cast().isDynamicDim(i)) + if (rhs2T.isDynamicDim(i)) { - ConstantIndexOp val = cast(multop.getRhs2().getDefiningOp()->getOperand(multop.getRhs2().getType().cast().getDynamicDimIndex(i)).getDefiningOp()); - lblSizes[lblOp] = val.value(); + if (DenseTensorDeclOp dense_decl_op = llvm::dyn_cast_if_present(multop.getRhs2().getDefiningOp())) + { + ConstantIndexOp val = cast(dense_decl_op->getOperand(rhs2T.getDynamicDimIndex(i)).getDefiningOp()); + lblSizes[lblOp] = val.value(); + } + else + { + assert(false && "Factorization optimization can only be applied when tensor size can be inferred statically"); + } + } else { - lblSizes[lblOp] = multop.getRhs2().getType().cast().getDimSize(i); + lblSizes[lblOp] = rhs2T.getDimSize(i); } - } labelValues[lblOp] = labels[i]; } labelVec.push_back(lblOp); } - if (isa(multop.getRhs2().getDefiningOp())) + if(!multop.getRhs2().getDefiningOp()) + { + lblMaps[multop.getRhs2().getAsOpaquePointer()] = labelVec; + } + else if (isa(multop.getRhs2().getDefiningOp())) { lblMaps[multop.getRhs2().getDefiningOp()] = labelVec; } @@ -289,27 +303,37 @@ void FindOptimalTCFactorizationPass::FindOptimalTCFactorization(tensorAlgebra::T auto lblOp = labels[i].getDefiningOp(); if (lblSizes.count(lblOp) == 0) { + auto rhs1T = cast(multop.getRhs1().getType()); + /// If dynamic dimension, we need to retrieve the value of the constantIndexOp that was used to create it /// If static, just get it from the tensor type - if (isa(multop.getRhs1().getDefiningOp())) - { - if (multop.getRhs1().getType().cast().isDynamicDim(i)) + if (rhs1T.isDynamicDim(i)) { - ConstantIndexOp val = cast(multop.getRhs1().getDefiningOp()->getOperand(multop.getRhs1().getType().cast().getDynamicDimIndex(i)).getDefiningOp()); - lblSizes[lblOp] = val.value(); + if (isa(multop.getRhs1().getDefiningOp())) + { + ConstantIndexOp val = cast(multop.getRhs1().getDefiningOp()->getOperand(rhs1T.getDynamicDimIndex(i)).getDefiningOp()); + lblSizes[lblOp] = val.value(); + } + else + { + assert(false && "Factorization optimization can only be applied when tensor size can be inferred statically"); + } } else { - lblSizes[lblOp] = multop.getRhs1().getType().cast().getDimSize(i); + lblSizes[lblOp] = rhs1T.getDimSize(i); } - } labelValues[lblOp] = labels[i]; } labelVec.push_back(lblOp); } - if (isa(multop.getRhs1().getDefiningOp())) + if(!multop.getRhs1().getDefiningOp()) + { + lblMaps[multop.getRhs1().getAsOpaquePointer()] = labelVec; + } + else if (isa(multop.getRhs1().getDefiningOp())) { lblMaps[multop.getRhs1().getDefiningOp()] = labelVec; } @@ -317,17 +341,32 @@ void FindOptimalTCFactorizationPass::FindOptimalTCFactorization(tensorAlgebra::T currValue = cast(curr).getOperation()->getOperand(1); curr = currValue.getDefiningOp(); } - - inLTOps.push_back(curr); - inLTValues[curr] = currValue; + if(curr) + { + inLTOps.push_back(curr); + inLTValues[curr] = currValue; + } + else + { + inLTOps.push_back(currValue.getAsOpaquePointer()); + inLTValues[currValue.getAsOpaquePointer()] = currValue; + } curr = stack.top(); stack.pop(); currValue = cast(curr).getOperation()->getOperand(0); curr = currValue.getDefiningOp(); } - inLTOps.push_back(curr); - inLTValues[curr] = currValue; + if(curr) + { + inLTOps.push_back(curr); + inLTValues[curr] = currValue; + } + else + { + inLTOps.push_back(currValue.getAsOpaquePointer()); + inLTValues[currValue.getAsOpaquePointer()] = currValue; + } } auto outLabels = cast(lhsOp).getResultIndexLabels(); @@ -358,7 +397,7 @@ void FindOptimalTCFactorizationPass::FindOptimalTCFactorization(tensorAlgebra::T for (size_t i = 1; i < order.size(); i++) { newRhs2 = inLTValues[inLTOps[order[i]]]; - auto elType = newRhs1.getType().dyn_cast().getElementType(); + auto elType = dyn_cast(newRhs1.getType()).getElementType(); auto newType = RankedTensorType::get(lhsTensorShapes[i - 1], elType); std::vector newSumLabels; @@ -379,7 +418,7 @@ void FindOptimalTCFactorizationPass::FindOptimalTCFactorization(tensorAlgebra::T { newSumLabels.push_back(labelValues[lbl]); } - if (!isa(newRhs1.getDefiningOp())) + if (!isa_and_present(newRhs1.getDefiningOp())) { /// store the output label values for subsequent ta.tc ops ___newSumLabels = newSumLabels; } @@ -388,7 +427,7 @@ void FindOptimalTCFactorizationPass::FindOptimalTCFactorization(tensorAlgebra::T std::vector new_lhs_lbls_value; std::vector new_rhs_lbls_value; - if (isa(newRhs1.getDefiningOp())) + if (isa_and_present(newRhs1.getDefiningOp())) { /// retrieve the labels from prev iteration. new_rhs_lbls_value = ___newSumLabels; for (auto lbl : new_rhs_lbls_value) @@ -439,30 +478,30 @@ void FindOptimalTCFactorizationPass::FindOptimalTCFactorization(tensorAlgebra::T } /// formats - SmallVector formats; - if (isa(newRhs2.getDefiningOp())) - { - auto lhs_format = dyn_cast(newRhs2.getDefiningOp()).getFormat(); - formats.push_back(lhs_format); - } - if (isa(newRhs1.getDefiningOp())) - { - auto rhs_format = dyn_cast(newRhs1.getDefiningOp()).getFormat(); - formats.push_back(rhs_format); - } - if (isa(newRhs1.getDefiningOp()) && - isa(newRhs2.getDefiningOp())) - { - auto rhs_format = dyn_cast(newRhs1.getDefiningOp()).getFormat(); - formats.push_back(rhs_format); - } - if (isa(newRhs1.getDefiningOp())) - { /// for series of ta.mul case - auto lhs_format = dyn_cast(newRhs2.getDefiningOp()).getFormat(); - formats.push_back(lhs_format); - formats.push_back(lhs_format); - } - auto strAttr = builder.getStrArrayAttr(formats); + SmallVector formats = {"Dense", "Dense", "Dense"}; + // StringRef format = "Dense"; + // if (isa(newRhs2.getDefiningOp())) + // { + // auto lhs_format = dyn_cast(newRhs2.getDefiningOp()).getFormat(); + // formats.push_back(lhs_format); + // } + // if (isa(newRhs1.getDefiningOp())) + // { + // auto rhs_format = dyn_cast(newRhs1.getDefiningOp()).getFormat(); + // formats.push_back(rhs_format); + // } + // if (isa(newRhs1.getDefiningOp()) && + // isa(newRhs2.getDefiningOp())) + // { + // auto rhs_format = dyn_cast(newRhs1.getDefiningOp()).getFormat(); + // formats.push_back(rhs_format); + // } + // if (isa(newRhs1.getDefiningOp())) + // { /// for series of ta.mul case + // auto lhs_format = dyn_cast(newRhs2.getDefiningOp()).getFormat(); + // formats.push_back(lhs_format); + // formats.push_back(lhs_format); + // } std::vector lhs_lbls; std::vector rhs_lbls; @@ -539,7 +578,7 @@ void FindOptimalTCFactorizationPass::FindOptimalTCFactorization(tensorAlgebra::T auto SemiringAttr = builder.getStringAttr("plusxy_times"); auto MaskingAttr = builder.getStringAttr("none"); Value tcop = builder.create(loc, newType, newRhs1, newRhs2, - all_labels, affineMapArrayAttr, strAttr, SemiringAttr, + all_labels, affineMapArrayAttr, SemiringAttr, MaskingAttr, nullptr); tcop.getDefiningOp()->setAttr("__alpha__", builder.getF64FloatAttr(1.0)); tcop.getDefiningOp()->setAttr("__beta__", builder.getF64FloatAttr(0.0)); diff --git a/lib/Dialect/TensorAlgebra/Transforms/TCtoTTGT.cpp b/lib/Dialect/TensorAlgebra/Transforms/TCtoTTGT.cpp index 3f8e4149..fc06d88b 100644 --- a/lib/Dialect/TensorAlgebra/Transforms/TCtoTTGT.cpp +++ b/lib/Dialect/TensorAlgebra/Transforms/TCtoTTGT.cpp @@ -34,12 +34,14 @@ #include "mlir/Dialect/Func/IR/FuncOps.h" #include "mlir/Dialect/Bufferization/IR/Bufferization.h" +#include #include #include #include #include #include "mlir/Dialect/MemRef/IR/MemRef.h" +#include "mlir/IR/BuiltinTypes.h" using namespace mlir; using namespace mlir::linalg; @@ -149,7 +151,7 @@ namespace auto alphaAttr = multop.getOperation()->getAttr("__alpha__"); auto betaAttr = multop.getOperation()->getAttr("__beta__"); - Operation *startTime; + Operation *startTime = nullptr; std::string getTimeStr = "getTime"; auto f64Type = rewriter.getF64Type(); if (printFlops) @@ -164,7 +166,7 @@ namespace /// Find summation indices for (const auto &map : indexMaps) { - auto affineMap = map.cast().getValue(); + auto affineMap = cast(map).getValue(); std::vector perm; for (size_t i = 0; i < affineMap.getNumResults(); i++) { @@ -177,8 +179,8 @@ namespace comet_pdump(op); comet_debug() << "\n"; - auto rhs1Tensor = cast(operands[0].getDefiningOp()); - auto rhs2Tensor = cast(operands[1].getDefiningOp()); + // auto rhs1Tensor = cast(operands[0].getDefiningOp()); + // auto rhs2Tensor = cast(operands[1].getDefiningOp()); comet_debug() << "\n"; Value lhsDef; tensorAlgebra::TensorSetOp setnewop; @@ -194,18 +196,54 @@ namespace comet_vdump(lhsDef); } } - auto lhsTensor = cast(lhsDef.getDefiningOp()); + // auto lhsTensor = cast(lhsDef.getDefiningOp()); comet_vdump(setnewop); comet_debug() << "\n"; - Value rhs1Memref = rhs1Tensor.getMemref(); - Value rhs2Memref = rhs2Tensor.getMemref(); - Value lhsMemref = lhsTensor.getMemref(); + Value rhs1Memref, rhs2Memref, lhsMemref; + if(auto tensor_rhs1 = dyn_cast(operands[0].getType())) + { + rhs1Memref = rewriter.createOrFold(loc, MemRefType::get(tensor_rhs1.getShape(), tensor_rhs1.getElementType()), operands[0]); + } + else if(mlir::isa(operands[0].getType())) + { + rhs1Memref = operands[0]; + } + else + { + assert(false && "Unexpected type"); + } + + if(auto tensor_rhs2 = dyn_cast(operands[1].getType())) + { + rhs2Memref = rewriter.createOrFold(loc, MemRefType::get(tensor_rhs2.getShape(), tensor_rhs2.getElementType()), operands[1]); + } + else if(mlir::isa(operands[1].getType())) + { + rhs2Memref = operands[1]; + } + else + { + assert(false && "Unexpected type"); + } + + if(auto tensor_lhs = dyn_cast(lhsDef.getType())) + { + lhsMemref = rewriter.createOrFold(loc, MemRefType::get(tensor_lhs.getShape(), tensor_lhs.getElementType()), lhsDef); + } + else if(mlir::isa(lhsDef.getType())) + { + lhsMemref = lhsDef; + } + else + { + assert(false && "Unexpected type"); + } - auto rhs1MemrefType = rhs1Memref.getType().cast(); - auto rhs2MemrefType = rhs2Memref.getType().cast(); - auto lhsMemrefType = lhsMemref.getType().cast(); + auto rhs1MemrefType = cast(rhs1Memref.getType()); + auto rhs2MemrefType = cast(rhs2Memref.getType()); + auto lhsMemrefType = cast(lhsMemref.getType()); std::vector allShapes{rhs1MemrefType.getShape(), rhs2MemrefType.getShape(), @@ -260,15 +298,15 @@ namespace /// Do transpose if needed if (!rhs1OutMapAttr.getValue().isIdentity()) { + auto shape = rhs1MemrefType.getShape(); std::vector operands; std::vector rhs1Dims; for (auto idx : rhs1OutPerm) { - auto shape = rhs1MemrefType.getShape(); rhs1Dims.push_back(shape[idx]); if (rhs1MemrefType.isDynamicDim(idx)) { - operands.push_back(rhs1Memref.getDefiningOp()->getOperand(rhs1MemrefType.getDynamicDimIndex(idx))); + operands.push_back(rewriter.create(loc, rhs1Memref, idx)); } } @@ -289,13 +327,13 @@ namespace { std::vector operands; std::vector rhs2Dims; + auto shape = rhs2MemrefType.getShape(); for (auto idx : rhs2OutPerm) { - auto shape = rhs2MemrefType.getShape(); rhs2Dims.push_back(shape[idx]); if (rhs2MemrefType.isDynamicDim(idx)) { - operands.push_back(rhs2Memref.getDefiningOp()->getOperand(rhs2MemrefType.getDynamicDimIndex(idx))); + operands.push_back(rewriter.create(loc, rhs2Memref, idx)); } } @@ -317,13 +355,13 @@ namespace { std::vector operands; std::vector lhsDims; + auto shape = lhsMemrefType.getShape(); for (auto idx : lhsOutPerm) { - auto shape = lhsMemrefType.getShape(); lhsDims.push_back(shape[idx]); if (lhsMemrefType.isDynamicDim(idx)) { - operands.push_back(lhsMemref.getDefiningOp()->getOperand(lhsMemrefType.getDynamicDimIndex(idx))); + operands.push_back(rewriter.create(loc, lhsMemref, idx)); } } @@ -331,7 +369,7 @@ namespace MemRefType::get(lhsDims, lhsMemrefType.getElementType()), operands, loc, rewriter); useLHSTranspose = true; - double beta_val = betaAttr.cast().getValueAsDouble(); + double beta_val = cast(betaAttr).getValueAsDouble(); if (beta_val == 0) { @@ -617,7 +655,7 @@ namespace Value lhsExpand = lhsReshape; if (expandLHS) /// LHS tensor was collapsed and now needs to be re-expanded using the same reassociation indices { - auto expandedTensorType = MemRefType::get(lhsAlloc.getType().cast().getShape(), lhsAlloc.getType().cast().getElementType()); + auto expandedTensorType = MemRefType::get(cast(lhsAlloc.getType()).getShape(), cast(lhsAlloc.getType()).getElementType()); comet_debug() << "\nExpanded:\n"; lhsExpand = rewriter.create( @@ -708,7 +746,7 @@ void TALoweringTTGTPass::runOnOperation() auto printFlopFunc = FunctionType::get(ctx, {FloatType::getF64(ctx)}, {}); /// func @getTime() -> f64 - if (!hasFuncDeclaration(module, "getTime")) + if (this->printFlops && !hasFuncDeclaration(module, "getTime")) { mlir::func::FuncOp func1 = mlir::func::FuncOp::create(function.getLoc(), "getTime", getTimeFunc, ArrayRef{}); @@ -717,7 +755,7 @@ void TALoweringTTGTPass::runOnOperation() } /// func @print_flops(%flops) : (f64) -> () - if (!hasFuncDeclaration(module, "print_flops")) + if (this->printFlops && !hasFuncDeclaration(module, "print_flops")) { mlir::func::FuncOp func1 = mlir::func::FuncOp::create(function.getLoc(), "print_flops", printFlopFunc, ArrayRef{}); @@ -729,7 +767,7 @@ void TALoweringTTGTPass::runOnOperation() patterns.insert(&getContext(), isSelectBestPerm, whatPerm, printFlops); ConversionTarget target(getContext()); - target.addLegalDialect(); + target.addLegalDialect(); if (failed(applyPartialConversion(function, target, std::move(patterns)))) { diff --git a/lib/Dialect/TensorAlgebra/Transforms/TCtoTTGTDyn.cpp b/lib/Dialect/TensorAlgebra/Transforms/TCtoTTGTDyn.cpp new file mode 100644 index 00000000..1c14882b --- /dev/null +++ b/lib/Dialect/TensorAlgebra/Transforms/TCtoTTGTDyn.cpp @@ -0,0 +1,791 @@ +//===- TCtoTTGT.cpp ------===// +// +// Copyright 2022 Battelle Memorial Institute +// +// Redistribution and use in source and binary forms, with or without modification, +// are permitted provided that the following conditions are met: +// +// 1. Redistributions of source code must retain the above copyright notice, this list of conditions +// and the following disclaimer. +// +// 2. Redistributions in binary form must reproduce the above copyright notice, this list of conditions +// and the following disclaimer in the documentation and/or other materials provided with the distribution. +// +// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED +// WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR +// PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, +// INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE +// GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, +// WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +// OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. +// +//===----------------------------------------------------------------------===// +// +// This file implements reformulation of tensor contraction operations as Transpose-Transpose-GEMM-Transpose +//===----------------------------------------------------------------------===// + +#include "comet/Dialect/TensorAlgebra/IR/TADialect.h" +#include "comet/Dialect/TensorAlgebra/Passes.h" +#include "comet/Dialect/Utils/Utils.h" + +#include "mlir/Dialect/Arith/Utils/Utils.h" +#include "mlir/Dialect/Linalg/IR/Linalg.h" +#include "mlir/Dialect/Linalg/Transforms/Transforms.h" +#include "mlir/Dialect/Arith/IR/Arith.h" +#include "mlir/Dialect/Func/IR/FuncOps.h" +#include "mlir/Dialect/Bufferization/IR/Bufferization.h" + +#include +#include +#include +#include +#include +#include + +#include "mlir/Dialect/SCF/IR/SCF.h" +#include "mlir/Dialect/Tensor/IR/Tensor.h" +#include "mlir/Dialect/Utils/ReshapeOpsUtils.h" +#include "mlir/IR/BuiltinAttributes.h" +#include "mlir/IR/BuiltinTypeInterfaces.h" +#include "mlir/IR/BuiltinTypes.h" +#include "mlir/IR/OpDefinition.h" +#include "mlir/IR/Value.h" +#include "mlir/IR/ValueRange.h" +#include "llvm/ADT/APFloat.h" +#include "llvm/ADT/ArrayRef.h" +#include "llvm/ADT/SmallVector.h" + +using namespace mlir; +using namespace mlir::linalg; +using namespace mlir::arith; +using namespace mlir::bufferization; + +using namespace mlir::tensorAlgebra; + +// *********** For debug purpose *********// +// #define COMET_DEBUG_MODE +#include "comet/Utils/debug.h" +#undef COMET_DEBUG_MODE +// *********** For debug purpose *********// + +const StringLiteral kLinalgTransformMarker = "__with_tiling__"; + +template +static bool arePermutations(const std::vector &vec1, + const std::vector &vec2) +{ + if (vec1.size() != vec2.size()) + { + return false; + } + std::vector taken(vec1.size(), false); + for (size_t i = 0; i < vec1.size(); i++) + { + auto it = std::find(vec2.begin(), vec2.end(), vec1[i]); + if (it == vec2.end()) + { + return false; + } + if (taken[std::distance(vec2.begin(), it)] == true) + { + return false; + } + taken[std::distance(vec2.begin(), it)] = true; + } + return true; +} + +/// Detect whether memref dims [dim, dim + extent) can be reshaped without +/// copies. +[[maybe_unused]] static bool isReshapableDimBand(unsigned dim, unsigned extent, + ArrayRef sizes, + ArrayRef strides) +{ + assert(sizes.size() == strides.size() && "mismatched ranks"); + /// off by 1 indexing to avoid out of bounds + for (auto idx = dim, e = dim + extent; idx + 1 < e; ++idx) + { + /// Only bands of static shapes are reshapable. This is due to the fact that + /// there is no relation between dynamic sizes and dynamic strides: we do not + /// have enough information to know whether a "-1" size corresponds to the + /// proper symbol in the AffineExpr of a stride. + if (ShapedType::isDynamic(sizes[dim + 1])) + return false; + /// simplify on the fly and catch more reshapable cases. + if (strides[idx] != strides[idx + 1] * sizes[idx + 1]) + return false; + } + return true; +} + +static IndexVector getIndexRange(unsigned lo, unsigned hi, unsigned step = 1) +{ + IndexVector result; + for (unsigned i = lo; i < hi; i += step) + { + result.push_back(i); + } + return result; +} + +//===----------------------------------------------------------------------===// +/// TAEarlyLoweringTTGTPass +//===----------------------------------------------------------------------===// + +/// This is a partial lowering to linear algebra of the tensor algebra +/// operations that are computationally intensive (like matmul for example...) +/// while keeping the rest of the code in the TA dialect. +namespace +{ + + struct TensorContractionOpLoweringTTGT : public ConversionPattern + { + TensorContractionOpLoweringTTGT(MLIRContext *ctx, int whatPerm, bool printFlops) + : ConversionPattern(tensorAlgebra::TensorMultOp::getOperationName(), 1, ctx), + whatPerm(whatPerm), printFlops{printFlops} {} + + /** + * @brief Latest implementation with following optimizations: + * - if no transpose is required there won't be any copy operations + * - if any operand is 2 dimensional no reshape + * - does not copy C + */ + LogicalResult + matchAndRewrite(Operation *op, ArrayRef operands, + ConversionPatternRewriter &rewriter) const final + { + comet_pdump(op); + assert(isa(op)); + + auto ctx = rewriter.getContext(); + auto loc = op->getLoc(); + auto multop = cast(op); + auto alphaAttr = multop.getOperation()->getAttr("__alpha__"); + auto betaAttr = multop.getOperation()->getAttr("__beta__"); + + Operation *startTime = nullptr; + std::string getTimeStr = "getTime"; + auto f64Type = rewriter.getF64Type(); + if (printFlops) + { + startTime = rewriter.create( + op->getLoc(), getTimeStr, SmallVector{f64Type}); + } + + ArrayAttr indexMaps = multop.getIndexingMaps(); + std::vector> allPerms; + + /// Find summation indices + for (const auto &map : indexMaps) + { + auto affineMap = cast(map).getValue(); + std::vector perm; + for (size_t i = 0; i < affineMap.getNumResults(); i++) + { + auto expr = affineMap.getResult(i); + perm.push_back(llvm::cast(expr).getPosition()); + } + + allPerms.push_back(perm); + } + + Value lhsDef; + tensorAlgebra::TensorSetOp setnewop; + for (auto u : multop.getOperation()->getResult(0).getUsers()) + { + comet_pdump(u); + if (isa(u)) + { + setnewop = cast(u); + Value dstTensor = u->getOperand(1); + + lhsDef = dstTensor; + comet_vdump(lhsDef); + } + } + + comet_vdump(setnewop); + comet_debug() << "\n"; + + Value rhs1Tensor = operands[0], rhs2Tensor = operands[1], lhsTensor = lhsDef; + + auto rhs1TensorType = cast(rhs1Tensor.getType()); + auto rhs2TensorType = cast(rhs2Tensor.getType()); + auto lhsTensorType = cast(lhsTensor.getType()); + + std::vector allShapes{rhs1Tensor, + rhs2Tensor, + lhsTensor}; + + ContractionPlanDyn plan(rewriter, loc, allPerms[0], allShapes[0], allPerms[1], + allShapes[1], allPerms[2], allShapes[2]); + + /// computeBestPermutations identifies the optimal index permutation for TTGT + /// it should enable and disable to heuristic + // IndexVector rhs1OutPerm, rhs2OutPerm, lhsOutPerm; + plan.computePermutations(rewriter, loc); + + comet_debug() << "Best permutation : " << plan.bestPermStr_ << "\n"; + + std::set rhsIndices(allPerms[0].begin(), allPerms[0].end()); + rhsIndices.insert(allPerms[1].begin(), allPerms[1].end()); + std::set lhsIndices(allPerms[2].begin(), allPerms[2].end()); + + std::vector sumIndices; + + std::set_difference(rhsIndices.begin(), rhsIndices.end(), + lhsIndices.begin(), lhsIndices.end(), + std::inserter(sumIndices, sumIndices.begin())); + + std::vector rhs1InPerm = getIdentityPermutation(allPerms[0].size()); + std::vector rhs2InPerm = getIdentityPermutation(allPerms[1].size()); + std::vector lhsInPerm = getIdentityPermutation(allPerms[2].size()); + // auto contractionTimes = rewriter.create(loc, plan.m_contraction_time); + + // Value minContractionTime = rewriter.create(loc, llvm::APFloat(std::numeric_limits::max()), FloatType::getF64(ctx)); + std::vector m_contraction_time_indices; + for (size_t i = 0; i < plan.m_contraction_time.size(); ++i) + { + m_contraction_time_indices.push_back(i); + } + scf::IndexSwitchOp switchOp; + Value permutation; + if(whatPerm == -1) // -1 means select the best permutation based on contraction time + { + Value minTime = plan.m_contraction_time.front(); + Value minIndex = rewriter.create(loc, 0); + + for(size_t i = 1; i < plan.m_contraction_time.size(); i++) + { + Value thisIndex = rewriter.create(loc, i); + auto foundMin = rewriter.create(loc, arith::CmpFPredicate::ULT, plan.m_contraction_time[i], minTime); + minIndex = rewriter.create(loc, foundMin, thisIndex, minIndex); + minTime = rewriter.create(loc, foundMin, plan.m_contraction_time[i], minTime); + } + permutation = minIndex; + } + else + { + permutation = rewriter.create(loc, whatPerm); // use the specified permutation from the command line + } + + switchOp = rewriter.create(loc, TypeRange(lhsTensor.getType()), permutation, ArrayRef(m_contraction_time_indices), plan.m_contraction_time.size()); + + auto& defaultCaseRegion = switchOp.getDefaultRegion(); + auto& defaultBlock = defaultCaseRegion.emplaceBlock(); // ensure the default case has a block + rewriter.setInsertionPointToStart(&defaultBlock); + rewriter.create(loc, ValueRange(lhsTensor)); + + auto caseRegions = switchOp.getCaseRegions(); + for(size_t i = 0; i < plan.m_contraction_time.size(); ++i) + { + bool useLHSTranspose = false; + Value rhs1Final = rhs1Tensor; // default to the original tensor + Value rhs2Final = rhs2Tensor; // default to the original tensor + Value lhsFinal = lhsTensor; // default to the original tensor + auto& caseBlock = caseRegions[i].emplaceBlock(); + rewriter.setInsertionPointToStart(&caseBlock); + if(plan.m_transposeA[i]) + { + auto shape = rhs1TensorType.getShape(); + std::vector operands; + std::vector rhs1Dims; + for (auto idx :plan.m_contraction_permutations[i][0]) // rhs1OutPerm + { + rhs1Dims.push_back(shape[idx]); + if (rhs1TensorType.isDynamicDim(idx)) + { + operands.push_back(rewriter.create(loc, rhs1Tensor, idx)); + } + } + + auto outputShape = rewriter.create(loc, rhs1Dims, rhs1TensorType.getElementType(), operands); + + std::vector rhs1OutPerm_int64(plan.m_contraction_permutations[i][0].begin(), plan.m_contraction_permutations[i][0].end()); + rhs1Final = rewriter.create(loc, rhs1Tensor, outputShape, llvm::ArrayRef(rhs1OutPerm_int64)).getResults()[0]; + comet_debug() << "\n"; + comet_vdump(rhs1Final); + } + if(plan.m_transposeB[i]) + { + std::vector operands; + std::vector rhs2Dims; + auto shape = rhs2TensorType.getShape(); + for (auto idx : plan.m_contraction_permutations[i][1]) // use the perm for rhs2 + { + rhs2Dims.push_back(shape[idx]); + if (rhs2TensorType.isDynamicDim(idx)) + { + operands.push_back(rewriter.create(loc, rhs2Tensor, idx)); + } + } + + std::vector rhs2OutPerm_int64(plan.m_contraction_permutations[i][1].begin(), plan.m_contraction_permutations[i][1].end()); + + auto outputShape = rewriter.create(loc, rhs2Dims, rhs2TensorType.getElementType(), operands); + rhs2Final = rewriter.create(loc, rhs2Tensor, outputShape, llvm::ArrayRef(rhs2OutPerm_int64)).getResults()[0]; + comet_debug() << " rhs2Transpose op: " << __LINE__ << "\n"; + comet_vdump(rhs2Final); + } + + if(plan.m_transposeC[i]) + { + std::vector operands; + std::vector lhsDims; + auto shape = lhsTensorType.getShape(); + for (auto idx : plan.m_contraction_permutations[i][2]) // lhsOutPerm + { + lhsDims.push_back(shape[idx]); + if (lhsTensorType.isDynamicDim(idx)) + { + operands.push_back(rewriter.create(loc, lhsTensor, idx)); + } + } + + useLHSTranspose = true; + double beta_val = cast(betaAttr).getValueAsDouble(); + auto outputShape = rewriter.create(loc, lhsDims, lhsTensorType.getElementType(), operands); + + if (beta_val == 0) + { + lhsFinal = outputShape; + } + else + { + std::vector lhsOutPerm_int64(plan.m_contraction_permutations[i][2].begin(), plan.m_contraction_permutations[i][2].end()); + lhsFinal = rewriter.create(loc, lhsTensor, outputShape, llvm::ArrayRef(lhsOutPerm_int64)).getResults()[0]; + } + } + + Value rhs1Reshape = rhs1Final; + Value rhs2Reshape = rhs2Final; + Value lhsReshape = lhsFinal; + + unsigned mIdxSize = plan.m_indices_.size(); + unsigned nIdxSize = plan.n_indices_.size(); + unsigned kIdxSize = plan.k_indices_.size(); + + bool isRHS1SumPermutation = arePermutations(allPerms[0], sumIndices); + bool isRHS2SumPermutation = arePermutations(allPerms[1], sumIndices); + + comet_debug() << __LINE__ << "mIdxSize, nIdxSize, kIdxSize: " + << mIdxSize << ", " + << nIdxSize << ", " + << kIdxSize << " isRHS1SumPermutation, isRHS2SumPermutation: " + << isRHS1SumPermutation << ", " + << isRHS2SumPermutation << "\n"; + + /// Do reshape if needed + if (isRHS1SumPermutation) + { + auto resultShape = rhs1TensorType.getShape(); + + auto rhs1AffineMap = AffineMap::getPermutationMap( + getIdentityPermutation(resultShape.size()), ctx); + + SmallVector rhs1IndexingMap{rhs1AffineMap}; + + SmallVector reassociationIndices = + getReassociationIndices(rhs1IndexingMap); + + comet_debug() << "\n"; + rhs1Reshape = rewriter.create( + loc, rhs1Final, reassociationIndices); + comet_vdump(rhs1Reshape); + } + else if (rhs1TensorType.getShape().size() != 2) + { + auto resultShape = rhs1TensorType.getShape(); + /// Construct combined shape of 2D memrefc + std::vector rhs1_0, rhs1_1; + + if (plan.m_swapAB[i]) + { + rhs1_0 = getIndexRange(0, kIdxSize); + rhs1_1 = getIndexRange(kIdxSize, kIdxSize + mIdxSize); + } + else + { + rhs1_0 = getIndexRange(0, mIdxSize); + rhs1_1 = getIndexRange(mIdxSize, mIdxSize + kIdxSize); + } + + auto rhs1AffineMap = AffineMap::getPermutationMap( + getIdentityPermutation(resultShape.size()), ctx); + auto rhs1Subset0 = rhs1AffineMap.getSubMap(rhs1_0); + auto rhs1Subset1 = rhs1AffineMap.getSubMap(rhs1_1); + + SmallVector rhs1IndexingMap; + + rhs1IndexingMap.push_back(rhs1Subset0); + rhs1IndexingMap.push_back(rhs1Subset1); + + SmallVector reassociationIndices = + getReassociationIndices(rhs1IndexingMap); + + comet_debug() << "\n"; + comet_debug() << " rhs1Alloc: \n"; + comet_vdump(rhs1Alloc); + comet_vdump(rhs1MemrefType); + + rhs1Reshape = rewriter.create( + loc, rhs1Final, reassociationIndices); + comet_debug() << " Before rhs1Reshape: \n"; + comet_vdump(rhs1Reshape); + comet_debug() << " After rhs1Reshape: \n"; + } + + if (isRHS2SumPermutation && rhs2TensorType.getShape().size() != 1) + { + auto resultShape = rhs2TensorType.getShape(); + + auto rhs2AffineMap = AffineMap::getPermutationMap( + getIdentityPermutation(resultShape.size()), ctx); + + SmallVector rhs2IndexingMap{rhs2AffineMap}; + + SmallVector reassociationIndices = + getReassociationIndices(rhs2IndexingMap); + + rhs2Reshape = rewriter.create( + loc, rhs2Final, reassociationIndices); + + comet_debug() << "\n"; + comet_vdump(rhs2Reshape); + } + else if (rhs2TensorType.getShape().size() != 2 && rhs2TensorType.getShape().size() != 1) + { + auto resultShape = rhs2TensorType.getShape(); + + /// Construct combined shape of 2D memref + std::vector rhs2_0, rhs2_1; + + if (plan.m_swapAB[i]) + { + rhs2_0 = getIndexRange(0, nIdxSize); + rhs2_1 = getIndexRange(nIdxSize, nIdxSize + kIdxSize); + } + else + { + rhs2_0 = getIndexRange(0, kIdxSize); + rhs2_1 = getIndexRange(kIdxSize, kIdxSize + nIdxSize); + } + + auto rhs2AffineMap = AffineMap::getPermutationMap( + getIdentityPermutation(resultShape.size()), ctx); + auto rhs2Subset0 = rhs2AffineMap.getSubMap(rhs2_0); + auto rhs2Subset1 = rhs2AffineMap.getSubMap(rhs2_1); + + SmallVector rhs2IndexingMap; + + rhs2IndexingMap.push_back(rhs2Subset0); + rhs2IndexingMap.push_back(rhs2Subset1); + + SmallVector reassociationIndices = + getReassociationIndices(rhs2IndexingMap); + + rhs2Reshape = rewriter.create( + loc, rhs2Final, reassociationIndices); + + comet_debug() << "\n"; + comet_vdump(rhs2Reshape); + } + + bool expandLHS = false; + /// Keep the reassociation indices that will be used for collapsing the LHS tensor + /// The exact same indices can be used to re-expand it back to its original rank (after the potential transpose operation) + SmallVector lhsReassociationIndices; + + comet_debug() << "\n"; + if (isRHS1SumPermutation || (isRHS2SumPermutation && rhs2TensorType.getShape().size() != 1)) + { + comet_debug() << "\n"; + auto resultShape = lhsTensorType.getShape(); + + auto lhsAffineMap = AffineMap::getPermutationMap( + getIdentityPermutation(resultShape.size()), ctx); + + SmallVector lhsIndexingMap{lhsAffineMap}; + + SmallVector reassociationIndices = + getReassociationIndices(lhsIndexingMap); + + /// TODO(gkestor): should it be expandop? + lhsReshape = rewriter.create( + loc, lhsFinal, reassociationIndices); + + comet_debug() << "\n"; + comet_vdump(lhsReshape); + expandLHS = true; + lhsReassociationIndices = reassociationIndices; + } + else if (lhsTensorType.getShape().size() != 2 && lhsTensorType.getShape().size() != 1) + { + comet_debug() << "\n"; + auto resultShape = lhsTensorType.getShape(); + /// Construct combined shape of 2D memref + std::vector lhs_0, lhs_1; + if (plan.m_swapAB[i]) // swap A and B in the matmul + { + lhs_0 = getIndexRange(0, nIdxSize); + lhs_1 = getIndexRange(nIdxSize, nIdxSize + mIdxSize); + } + else + { + lhs_0 = getIndexRange(0, mIdxSize); + lhs_1 = getIndexRange(mIdxSize, mIdxSize + nIdxSize); + } + + auto lhsAffineMap = AffineMap::getPermutationMap( + getIdentityPermutation(resultShape.size()), ctx); + auto lhsSubset0 = lhsAffineMap.getSubMap(lhs_0); + auto lhsSubset1 = lhsAffineMap.getSubMap(lhs_1); + + SmallVector lhsIndexingMap; + + lhsIndexingMap.push_back(lhsSubset0); + lhsIndexingMap.push_back(lhsSubset1); + + SmallVector reassociationIndices = + getReassociationIndices(lhsIndexingMap); + + lhsReshape = rewriter.create( + loc, lhsFinal, reassociationIndices); + comet_debug() << "\n"; + comet_vdump(lhsReshape); + + expandLHS = true; + lhsReassociationIndices = reassociationIndices; + } + + comet_debug() << "\n"; + /// Create linalg matmul op + linalg::MatmulOp matmulop; + linalg::MatvecOp matvecop; + + if (isRHS1SumPermutation) + { + comet_debug() << "\n"; + matvecop = rewriter.create( + loc, ValueRange{rhs2Reshape, rhs1Reshape}, + ValueRange{lhsReshape}); + comet_debug() << "\n"; + comet_vdump(matvecop); + + matvecop.getOperation()->setAttr("__alpha__", alphaAttr); + matvecop.getOperation()->setAttr("__beta__", betaAttr); + lhsReshape = matvecop.getResults()[0]; + /// TODO(gkestor): Add attribute to the linalg.matvec operations + /// matvecop.setAttr(kLinalgTransformMarker, rewriter.getStringAttr(kLinalgTransformMarker)); + } + else if (isRHS2SumPermutation) + { + comet_debug() << "\n"; + matvecop = rewriter.create( + loc, ValueRange{rhs1Reshape, rhs2Reshape}, + ValueRange{lhsReshape}); + comet_debug() << "\n"; + comet_vdump(rhs1Reshape); + comet_vdump(rhs2Reshape); + comet_vdump(lhsReshape); + comet_vdump(matvecop); + + matvecop.getOperation()->setAttr("__alpha__", alphaAttr); + matvecop.getOperation()->setAttr("__beta__", betaAttr); + lhsReshape = matvecop.getResults()[0]; + + /// TODO(gkestor): Add attribute to the linalg.matvec operations + /// matvecop.setAttr(kLinalgTransformMarker, rewriter.getStringAttr(kLinalgTransformMarker)); + } + else + { + comet_debug() << "\n"; + + if (plan.m_swapAB[i]) + { + std::swap(rhs1Reshape, rhs2Reshape); // swap the operands for matmul when A and B are swapped in the contraction + comet_debug() << "Swapping rhs1 and rhs2 for matmul due to contraction swap\n"; + } + + matmulop = rewriter.create( + loc, ValueRange{rhs1Reshape, rhs2Reshape}, + ValueRange{lhsReshape}); + comet_debug() << "\n"; + comet_vdump(rhs1Reshape); + comet_vdump(rhs2Reshape); + comet_vdump(lhsReshape); + comet_vdump(matmulop); + comet_debug() << "\n"; + /// Add attribute to the linalg.matmul operations + auto iterator_types = matmulop.getIteratorTypesArray(); + + matmulop.getOperation()->setAttr(kLinalgTransformMarker, + rewriter.getStringAttr(kLinalgTransformMarker)); + matmulop.getOperation()->setAttr("__alpha__", alphaAttr); + matmulop.getOperation()->setAttr("__beta__", betaAttr); + lhsReshape = matmulop.getResults()[0]; + } + + Value lhsExpand = lhsReshape; + if (expandLHS) /// LHS tensor was collapsed and now needs to be re-expanded using the same reassociation indices + { + auto expandedTensorType = RankedTensorType::get(cast(lhsFinal.getType()).getShape(), cast(lhsFinal.getType()).getElementType()); + SmallVector dims; + for(int64_t i = 0; i < expandedTensorType.getRank(); i++) + { + if(expandedTensorType.isDynamicDim(i)) + { + Value dim = rewriter.create(loc, i); + dims.push_back(rewriter.create(loc, lhsFinal, dim).getResult()); + } + else + { + Value dim = rewriter.create(loc, expandedTensorType.getDimSize(i)); + dims.push_back(dim); + } + } + + + + comet_debug() << "\nExpanded:\n"; + // lhsExpand = rewriter.create( + // loc, expandedTensorType, lhsReshape, getReassociationIndicesAttribute(rewriter, lhsReassociationIndices), dims, expandedTensorType.getShape()); + lhsExpand = rewriter.create( + loc, expandedTensorType, lhsReshape, lhsReassociationIndices, dims); + comet_debug() << "\n"; + comet_vdump(lhsExpand); + } + + /// Copy back the result if needed + if (lhsFinal != lhsTensor && useLHSTranspose) + { + std::vector revLhsOutPerm(plan.m_contraction_permutations[i][2].size()); + for (size_t j = 0; j < revLhsOutPerm.size(); j++) + revLhsOutPerm[plan.m_contraction_permutations[i][2][j]] = j; + + lhsExpand = rewriter.create(loc, lhsExpand, lhsTensor, llvm::ArrayRef(revLhsOutPerm)).getResults()[0]; + comet_vdump(lhsExpand); + comet_debug() << "\n"; + + } + rewriter.create(loc, ValueRange({lhsExpand})); + } + rewriter.setInsertionPointAfter(switchOp); + + // auto argInt = rewriter.create(loc, rewriter.getIntegerType(64), switchOp.getArg()); + // auto argFloat = + // rewriter.create(loc, f64Type, argInt); + // rewriter.create(loc, forOp.getResult(0)); // for debugging purposes, print the minimum time found + // rewriter.create(loc, argFloat); // for debugging purposes, print the minimum permutation index found + + if (printFlops) + { + auto endTime = rewriter.create( + loc, getTimeStr, SmallVector{f64Type}); + + auto start = startTime->getResult(0); + auto end = endTime.getResult(0); + + Value totalTimeValue = + rewriter.create(loc, f64Type, end, start); + + Value opNums = rewriter.create(loc, 2); + opNums = rewriter.create(loc, opNums, + plan.m_size_); // m size + opNums = rewriter.create(loc, opNums, + plan.n_size_); // n size + opNums = rewriter.create(loc, opNums, + plan.k_size_); // k size + opNums = rewriter.create(loc, rewriter.getIntegerType(64), opNums); + opNums = rewriter.create(loc, f64Type, opNums); // convert to f64 for division + + Value flopsOp = + rewriter.create(loc, f64Type, opNums, totalTimeValue); + + /// call @print_flops(%flops) : (f64) -> () + std::string printFlopsStr = "print_flops"; + /// auto printFlopsCall = + rewriter.create( + loc, printFlopsStr, SmallVector{}, ValueRange{flopsOp}); + } + + rewriter.replaceAllUsesWith( + op->getResults(), switchOp.getResults()); // Replace the original op with the final result of the matmul or matvec + // rewriter.replaceUsesWithIf(setnewop->getOperand(1), switchOp.getResult(0), [&](OpOperand& use) { + // auto user = use.getOwner(); + // auto ancestor = switchOp->getBlock()->findAncestorOpInBlock(*user); + // return (ancestor && switchOp->isBeforeInBlock(ancestor)); + // }); + // op->replaceAllUsesWith(switchOp); + // rewriter.eraseOp(setnewop); + rewriter.eraseOp(op); + return success(); + } + + private: + int whatPerm; + bool printFlops; + }; /// namespace + + struct TALoweringTTGTDynPass + : public PassWrapper> + { + MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TALoweringTTGTDynPass) + TALoweringTTGTDynPass(int whatPerm, bool printFlops) : whatPerm(whatPerm), printFlops{printFlops} {}; + void runOnOperation() override; + + private: + int whatPerm; + bool printFlops; + }; + +} /// end anonymous namespace. + +void TALoweringTTGTDynPass::runOnOperation() +{ + func::FuncOp function = getOperation(); + auto module = function.getOperation()->getParentOfType(); + auto *ctx = &getContext(); + + auto getTimeFunc = FunctionType::get(ctx, {}, {FloatType::getF64(ctx)}); + auto printFlopFunc = FunctionType::get(ctx, {FloatType::getF64(ctx)}, {}); + + /// func @getTime() -> f64 + if (this->printFlops && !hasFuncDeclaration(module, "getTime")) + { + mlir::func::FuncOp func1 = mlir::func::FuncOp::create(function.getLoc(), "getTime", getTimeFunc, + ArrayRef{}); + func1.setPrivate(); + module.push_back(func1); + } + + /// func @print_flops(%flops) : (f64) -> () + if (this->printFlops && !hasFuncDeclaration(module, "print_flops")) + { + mlir::func::FuncOp func1 = mlir::func::FuncOp::create(function.getLoc(), "print_flops", + printFlopFunc, ArrayRef{}); + func1.setPrivate(); + module.push_back(func1); + } + + RewritePatternSet patterns(&getContext()); + patterns.insert(&getContext(), whatPerm, printFlops); + + ConversionTarget target(getContext()); + target.addLegalDialect(); + target.addIllegalOp(); + + if (failed(applyPartialConversion(function, target, std::move(patterns)))) + { + llvm::errs() << "Failed to applyPartialConversion in TALoweringTTGTDynPass\n"; + signalPassFailure(); + } +} + +/// Create a pass for lowering operations in the `LinAlg` and `Std` dialects, +/// for a subset of the TA IR (e.g. matmul). +/// ordering of permutation starts with one +std::unique_ptr mlir::comet::createLoweringTTGTDynPass(int whatPerm, bool printFlops) +{ + return std::make_unique(whatPerm, printFlops); +} diff --git a/lib/Dialect/TensorAlgebra/Transforms/TensorDeclLowering.cpp b/lib/Dialect/TensorAlgebra/Transforms/TensorDeclLowering.cpp index df5cc2c5..49ad700c 100644 --- a/lib/Dialect/TensorAlgebra/Transforms/TensorDeclLowering.cpp +++ b/lib/Dialect/TensorAlgebra/Transforms/TensorDeclLowering.cpp @@ -31,14 +31,25 @@ #include "mlir/Dialect/Arith/IR/Arith.h" #include "mlir/Dialect/Func/IR/FuncOps.h" +#include "mlir/Dialect/Linalg/IR/Linalg.h" #include "mlir/Dialect/MemRef/IR/MemRef.h" +#include "mlir/Dialect/Tensor/IR/Tensor.h" #include "mlir/Dialect/Bufferization/IR/Bufferization.h" +#include "mlir/IR/Attributes.h" +#include "mlir/IR/BuiltinAttributes.h" +#include "mlir/IR/BuiltinTypeInterfaces.h" +#include "mlir/IR/BuiltinTypes.h" +#include "mlir/IR/TypeRange.h" +#include "mlir/IR/ValueRange.h" +#include "mlir/Support/LLVM.h" #include #include #include +#include #include +#include "llvm/ADT/SmallVector.h" #include "llvm/Support/Debug.h" using namespace mlir; @@ -51,7 +62,7 @@ using namespace mlir::indexTree; #define DEBUG_TYPE "tensor-decl-lowering" // *********** For debug purpose *********// -//#define COMET_DEBUG_MODE +// #define COMET_DEBUG_MODE #include "comet/Utils/debug.h" #undef COMET_DEBUG_MODE // *********** For debug purpose *********// @@ -61,407 +72,188 @@ using namespace mlir::indexTree; //===----------------------------------------------------------------------===// namespace { - void mixModeEltWiseMultSparseTensorOutputLowering(Value computeOp, Location loc, - std::vector> rshPerms, - std::vector &dimSizes, - std::vector &tensorload_sizes_vec, - std::vector &array_sizes_vec, - PatternRewriter &rewriter) + void insertReadFileLibCall(int rank_size, Type floatEleType, Type indicesType, MLIRContext *ctx, ModuleOp &module, func::FuncOp function) { + comet_debug() << "Inserting insertReadFileLibCall\n"; + IndexType indexType = IndexType::get(function.getContext()); + IntegerType i32Type = IntegerType::get(ctx, 32); + auto unrankedMemref_index = mlir::UnrankedMemRefType::get(indexType, 0); + auto unrankedMemref_element_type = mlir::UnrankedMemRefType::get(floatEleType, 0); + auto unrankedMemref_indices_type = mlir::UnrankedMemRefType::get(indicesType, 0); + llvm::SmallVector inputFuncArgTypes; + llvm::SmallVector inputSizeFuncArgTypes; + inputFuncArgTypes.push_back(i32Type); + inputSizeFuncArgTypes.push_back(i32Type); + for(int i = 0; i < rank_size * 2; i++) + { + inputFuncArgTypes.push_back(indexType); + inputSizeFuncArgTypes.push_back(indexType); + } + for(int i = 0; i < rank_size * 4; i++) + { + inputFuncArgTypes.push_back(unrankedMemref_indices_type); + } + inputSizeFuncArgTypes.push_back(unrankedMemref_index); + inputFuncArgTypes.push_back(unrankedMemref_element_type); + inputSizeFuncArgTypes.push_back(i32Type); + inputFuncArgTypes.push_back(i32Type); - IndexType indexType = IndexType::get(computeOp.getContext()); - FloatType f64Type = FloatType::getF64(computeOp.getContext()); - auto dynamicmemTy_1d_index = MemRefType::get({ShapedType::kDynamic}, indexType); /// memref - auto dynamicmemTy_1d_f64 = MemRefType::get({ShapedType::kDynamic}, f64Type); /// memref - - comet_debug() << "mixModeEltWiseMultSparseTensorOutputLowering computeOp\n"; - comet_vdump(computeOp); - - /// elementwise mul op in mix sparse dense case - /// If elementwise, copy sparse input arrays for elementwise mul - int sparse_inputtensor_id = -1; - auto rhsComputeOp = computeOp.getDefiningOp()->getOperand(0).getDefiningOp(); - - auto first_operand = rhsComputeOp->getOperand(0).getDefiningOp(); - auto second_operand = rhsComputeOp->getOperand(1).getDefiningOp(); - comet_debug() << "EltWiseMult Operands:\n"; - comet_pdump(first_operand); - comet_pdump(second_operand); + auto readInpuFunc = FunctionType::get(ctx, TypeRange(inputFuncArgTypes), {}); + assert(rank_size <=3 && rank_size >=2); - if (isa(first_operand)) + std::string func_name = "read_input_"+std::to_string(rank_size)+"D"; + if (floatEleType.isF32()) { - sparse_inputtensor_id = 0; + func_name += "_f32"; } - else if (isa(second_operand)) + else if (floatEleType.isF64()) { - sparse_inputtensor_id = 1; + func_name += "_f64"; } - else + else { - llvm::errs() << "ERROR: SparseTensorConstructOp was not found as one of the operands for itCompute\n"; + assert(false && "Unexpected type"); } - comet_debug() << " SparseTensorConstructOp for computeOp: \n"; - comet_pdump(rhsComputeOp->getOperand(sparse_inputtensor_id).getDefiningOp()); - auto sptensor_construct_op = cast(rhsComputeOp->getOperand(sparse_inputtensor_id).getDefiningOp()); - - for (unsigned int i = 0; i < 4 * (rshPerms[sparse_inputtensor_id].size()) + 1; i++) + if(indicesType.isIndex()) { - comet_debug() << " in for loop\n"; - Value intput_tensorload_op = cast(sptensor_construct_op.getOperand(i).getDefiningOp()); - Value input_alloc_op = cast(intput_tensorload_op.getDefiningOp()->getOperand(0).getDefiningOp()); - comet_debug() << " AllocOp: "; - comet_vdump(input_alloc_op); - - comet_debug() << " "; - Value input_alloc_op_param = input_alloc_op.getDefiningOp()->getOperand(0); - comet_debug() << " "; - - Value output_alloc_op; - if (i < 4 * (rshPerms[sparse_inputtensor_id].size())) - { - /// Memory allocation for position and coordinate arrays in sparse tensor contractions - output_alloc_op = insertAllocAndInitialize(loc, dynamicmemTy_1d_index, ValueRange{input_alloc_op_param}, rewriter); - } - else - { - /// Memory allocation for value array in sparse tensor contractions - output_alloc_op = insertAllocAndInitialize(loc, dynamicmemTy_1d_f64, ValueRange{input_alloc_op_param}, rewriter); /// Cval array - comet_debug() << " AllocOp: "; - comet_vdump(output_alloc_op); - } - - Value output_tensorload_op = rewriter.create(loc, output_alloc_op); - tensorload_sizes_vec.push_back(output_tensorload_op); + func_name += "_64"; } - comet_debug() << " "; - - /// [0...2d, 2d+1...4d+1, 4d+2...5d+1] - for (unsigned int i = 0; i < 4 * (rshPerms[sparse_inputtensor_id].size()) + 1; i++) + else if(indicesType.isInteger(32)) { - int sizes_i = i + 4 * (rshPerms[sparse_inputtensor_id].size()) + 1; - comet_debug() << " "; - comet_pdump(sptensor_construct_op.getOperand(sizes_i).getDefiningOp()); - - Value input_load_op = sptensor_construct_op.getOperand(sizes_i); - comet_debug() << "Ops push_back for Sparse Tensor Construct Op for MixedMode elementwise multiplication (array_sizes_vec):\n"; - comet_vdump(input_load_op); - array_sizes_vec.push_back(input_load_op); + func_name += "_i32"; } - - for (unsigned int i = 0; i < rshPerms[sparse_inputtensor_id].size(); i++) + else if(indicesType.isInteger(64)) { - int sizes_i = i + 2 * (2 * (rshPerms[sparse_inputtensor_id].size()) + 1); - comet_debug() << " "; - comet_pdump(sptensor_construct_op.getOperand(sizes_i).getDefiningOp()); - - Value input_load_op = sptensor_construct_op.getOperand(sizes_i); - comet_debug() << "Ops push_back for Sparse Tensor Construct Op for MixedMode elementwise multiplication (dimSizes):\n"; - comet_vdump(input_load_op); - dimSizes.push_back(input_load_op); + func_name += "_i64"; + } + else + { + assert(false && "Unexpected type"); } - } - - template - void pureSparseMultSparseTensorOutputLowering(T op, - Location loc, - std::string sparseOutputFormat, - std::vector &dimSizes, - std::vector &tensorload_sizes_vec, - std::vector &array_sizes_vec, - PatternRewriter &rewriter) - { - comet_debug() << " sparse output is used in itComputeOp op\n"; - comet_debug() << " sparseOutputFormat: " << sparseOutputFormat << "\n"; - - comet_vdump(op); - - IndexType indexType = IndexType::get(op.getContext()); - FloatType f64Type = FloatType::getF64(op.getContext()); - auto dynamicmemTy_1d_index = MemRefType::get({ShapedType::kDynamic}, indexType); /// memref - auto dynamicmemTy_1d_f64 = MemRefType::get({ShapedType::kDynamic}, f64Type); /// memref - - Value cst_index_0 = rewriter.create(loc, IndexType::get(op.getContext()), rewriter.getIndexAttr(0)); - comet_vdump(cst_index_0); - Value cst_index_1 = rewriter.create(loc, IndexType::get(op.getContext()), rewriter.getIndexAttr(1)); - comet_vdump(cst_index_1); - comet_vdump(op); - comet_pdump(op.getOperation()); - - unsigned int tensor_rank = op.getResult().getType().getRank(); - - std::vector array_sizes; - std::vector array_sizes_alloc_vec; - std::vector initial_array_sizes; - - if (sparseOutputFormat.compare("CSR") == 0) - { /// CSR format - comet_debug() << " 2D CSR format in sparse output decl op\n"; - /// AllocOp, storeOp, LoadOp - initial_array_sizes.push_back(cst_index_1); - initial_array_sizes.push_back(cst_index_1); - - /// A1tile - initial_array_sizes.push_back(cst_index_0); - initial_array_sizes.push_back(cst_index_0); - - /// The other three size information size.. - /// get the dimension size from operand or type - - TensorType t = op.getResult().getType(); - for(int i = 0; i < t.getRank(); i++) - { - if(t.isDynamicDim(i)) - { - dimSizes.push_back(op.getOperation()->getOperand(i)); - } - else - { - dimSizes.push_back(rewriter.create(loc, t.getShape()[i] )); - } - } - /// The dim size is the second parameter of the - Value dim2_posSize = rewriter.create(loc, dimSizes[0], cst_index_1); - comet_debug() << "AddIOp generated for dim2_posSize:\n"; - comet_vdump(dim2_posSize); - initial_array_sizes.push_back(dim2_posSize); + if (!hasFuncDeclaration(module, func_name)) + { + comet_debug() << "Adding " << func_name <<" to the module\n"; + func::FuncOp func1 = func::FuncOp::create(function.getLoc(), func_name, + readInpuFunc, ArrayRef{}); + func1.setPrivate(); + module.push_back(func1); + } - Value dim2_crdSize = rewriter.create(loc, dimSizes[0], dimSizes[1]); - initial_array_sizes.push_back(dim2_crdSize); - /// A2tile - initial_array_sizes.push_back(cst_index_0); - initial_array_sizes.push_back(cst_index_0); + auto readInputSizesFunc = FunctionType::get(ctx, TypeRange(inputSizeFuncArgTypes), {}); /// last arg (i32Type): readMode - /// Aval - initial_array_sizes.push_back(dim2_crdSize); - comet_debug() << " "; - comet_vdump(dim2_crdSize); + std::string input_size_func_name = "read_input_sizes_"+std::to_string(rank_size)+"D"; + if (floatEleType.isF32()) + { + input_size_func_name+="_f32"; } - else + else if (floatEleType.isF64()) { - llvm::errs() << __FILE__ << ":" << __LINE__ << "Not supported format\n"; + input_size_func_name+="_f64"; } - - /// same with transpose case - comet_debug() << " initial_array_sizes.size(): " << initial_array_sizes.size() << "\n"; - comet_debug() << " tensor_rank: " << tensor_rank << "\n"; - std::vector array_alloc_vec; - for (unsigned int i = 0; i < 4 * tensor_rank + 1; i++) + else { - Value alloc_sizes; - if (i < 4 * tensor_rank) - { - comet_debug() << " Inserting AllocOp: "; - alloc_sizes = insertAllocAndInitialize(loc, dynamicmemTy_1d_index, ValueRange{initial_array_sizes[i]}, rewriter); - comet_debug() << " AllocOp: "; - comet_vdump(alloc_sizes); - } - else - { - alloc_sizes = insertAllocAndInitialize(loc, dynamicmemTy_1d_f64, ValueRange{initial_array_sizes[i]}, rewriter); - comet_debug() << " AllocOp: "; - comet_vdump(alloc_sizes); - } - Value tensorload_sizes = rewriter.create(loc, alloc_sizes); - tensorload_sizes_vec.push_back(tensorload_sizes); - array_alloc_vec.push_back(alloc_sizes); + assert(false && "Unsupported type"); } - /// Initialize the sizes of pos/crd/val arrays - array_sizes.push_back(cst_index_1); /// A1pos_size - array_sizes.push_back(cst_index_1); /// A1crd_size - array_sizes.push_back(cst_index_1); /// A1tile_pos_size - array_sizes.push_back(cst_index_0); /// A1tile_crd_size - array_sizes.push_back(cst_index_1); /// A2pos_size - array_sizes.push_back(cst_index_0); /// A2crd_size - array_sizes.push_back(cst_index_0); /// A2tile_pos_size - array_sizes.push_back(cst_index_1); /// A2tile_crd_size - array_sizes.push_back(cst_index_0); /// Aval_size - /// put the array sizes into alloc/store/loadOp - for (auto size : array_sizes) + + if (!hasFuncDeclaration(module, input_size_func_name)) { - MemRefType memTy_alloc_sizes = MemRefType::get({1}, IndexType::get(op.getContext())); - Value allocop = rewriter.create(loc, memTy_alloc_sizes); - rewriter.create(loc, size, allocop, ValueRange{cst_index_0}); - Value loadop = rewriter.create(loc, allocop, ValueRange{cst_index_0}); - array_sizes_vec.push_back(loadop); - array_sizes_alloc_vec.push_back(allocop); + comet_debug() << "Adding read_input_sizes_2D_f32 to the module\n"; + func::FuncOp func1 = func::FuncOp::create(function.getLoc(), input_size_func_name, + readInputSizesFunc, ArrayRef{}); + func1.setPrivate(); + module.push_back(func1); } - - rewriter.create(loc, dimSizes[0], array_alloc_vec[0], ValueRange{cst_index_0}); } - void insertReadFileLibCall(int rank_size, MLIRContext *ctx, ModuleOp &module, func::FuncOp function) +Value insertSparseTensorDeclOp(PatternRewriter & rewriter, + MLIRContext* ctx, + Location loc, + unsigned rank_size, + std::vector& tensorload_sizes_vec, + std::vector& array_sizes_vec, + std::vector>& allPerms, + std::vector& dimSizes, + Type ty) { - comet_debug() << "Inserting insertReadFileLibCall\n"; - FloatType f32Type, f64Type; - if (VALUETYPE.compare("f32") == 0) - { - f32Type = FloatType::getF32(ctx); - } - else + comet_debug() << " Get users after "; + /// create sparse tensor construct after lowering each sparse tensor output users + comet_debug() << " tensorload_sizes_vec.size(): " << tensorload_sizes_vec.size() << ", rank_size: " << rank_size << "\n"; + /// create sptensor_construct + + ArrayRef dim_formats = mlir::cast(ty).getFormat(); + llvm::SmallVector dim_formats_attr; + for(TensorFormatEnum format: dim_formats) { - f64Type = FloatType::getF64(ctx); + dim_formats_attr.push_back(TensorFormatEnumAttr::get(ctx,format)); } - IndexType indexType = IndexType::get(function.getContext()); - IntegerType i32Type = IntegerType::get(ctx, 32); - auto unrankedMemref_f64 = mlir::UnrankedMemRefType::get(f64Type, 0); - /// TODO(gkestor): there is an issue with F32 UnrankedMemRefType - auto unrankedMemref_f32 = mlir::UnrankedMemRefType::get(f64Type, 0); - auto unrankedMemref_index = mlir::UnrankedMemRefType::get(indexType, 0); - + Value sptensor; if (rank_size == 2) { - comet_debug() << " Rank Size is 2\n"; - auto readInput2DF32Func = FunctionType::get(ctx, {i32Type, indexType, indexType, /// A1_format, A1_tile_format - indexType, indexType, /// A2_format, A2_tile_format - unrankedMemref_index, unrankedMemref_index, /// A1_pos, A1_crd - unrankedMemref_index, unrankedMemref_index, /// A1_tile_pos, A1_tile_crd - unrankedMemref_index, unrankedMemref_index, /// A2_pos, A2_crd - unrankedMemref_index, unrankedMemref_index, /// A2_tile_pos, A2_tile_crd - unrankedMemref_f32, i32Type}, - {}); /// last arg (i32Type): readMode - auto readInput2DF64Func = FunctionType::get(ctx, {i32Type, indexType, indexType, /// A1_format, A1_tile_format - indexType, indexType, /// A2_format, A2_tile_format - unrankedMemref_index, unrankedMemref_index, /// A1_pos, A1_crd - unrankedMemref_index, unrankedMemref_index, /// A1_tile_pos, A1_tile_crd - unrankedMemref_index, unrankedMemref_index, /// A2_pos, A2_crd - unrankedMemref_index, unrankedMemref_index, /// A2_tile_pos, A2_tile_crd - unrankedMemref_f64, i32Type}, - {}); - - if (VALUETYPE.compare("f32") == 0) - { - std::string func_name = "read_input_2D_f32"; - if (!hasFuncDeclaration(module, func_name)) - { - comet_debug() << "Adding read_input_2D_f32 to the module\n"; - func::FuncOp func1 = func::FuncOp::create(function.getLoc(), func_name, - readInput2DF32Func, ArrayRef{}); - func1.setPrivate(); - module.push_back(func1); - } - } - else /// f64 - { - std::string func_name = "read_input_2D_f64"; - if (!hasFuncDeclaration(module, func_name)) - { - comet_debug() << "Adding read_input_2D_f64 to the module\n"; - func::FuncOp func1 = func::FuncOp::create(function.getLoc(), func_name, - readInput2DF64Func, ArrayRef{}); - func1.setPrivate(); - module.push_back(func1); - } - } - - auto readInputSizes2DF64Func = FunctionType::get(ctx, {i32Type, indexType, indexType, indexType, indexType, unrankedMemref_index, i32Type}, {}); /// last arg (i32Type): readMode - - if (VALUETYPE.compare("f32") == 0) - { - std::string func_name = "read_input_sizes_2D_f32"; - if (!hasFuncDeclaration(module, func_name)) - { - comet_debug() << "Adding read_input_sizes_2D_f32 to the module\n"; - func::FuncOp func1 = func::FuncOp::create(function.getLoc(), func_name, - readInputSizes2DF64Func, ArrayRef{}); - func1.setPrivate(); - module.push_back(func1); - } - } - else - { - std::string func_name = "read_input_sizes_2D_f64"; - if (!hasFuncDeclaration(module, func_name)) - { - comet_debug() << "Adding read_input_sizes_2D_f64 to the module\n"; - func::FuncOp func1 = func::FuncOp::create(function.getLoc(), func_name, - readInputSizes2DF64Func, ArrayRef{}); - func1.setPrivate(); - module.push_back(func1); - } - } + Value dims = rewriter.create(loc, ValueRange{dimSizes[0], dimSizes[1]}); + sptensor = rewriter.create(loc, ty, + dims, + ValueRange{ + tensorload_sizes_vec[0], /// A1pos (each dimension consists of pos and crd arrays) + tensorload_sizes_vec[4], /// A2pos + }, + ValueRange{ + tensorload_sizes_vec[1], /// A1crd + tensorload_sizes_vec[5], /// A2crd + }, + ValueRange { + tensorload_sizes_vec[2], /// A1tile_pos + tensorload_sizes_vec[6], /// A2tile_pos + }, + ValueRange { + tensorload_sizes_vec[3], /// A1tile_crd + tensorload_sizes_vec[7], /// A2tile_crd + }, + tensorload_sizes_vec[8], /// Aval + 2, ArrayAttr::get(rewriter.getContext(), dim_formats_attr)); } - /// 3D tensor else if (rank_size == 3) { - auto readInput3DF32Func = FunctionType::get(ctx, {i32Type, indexType, indexType, indexType, indexType, indexType, indexType, /// Dimensions - unrankedMemref_index, unrankedMemref_index, /// A1 - unrankedMemref_index, unrankedMemref_index, /// A1_tile - unrankedMemref_index, unrankedMemref_index, /// A2 - unrankedMemref_index, unrankedMemref_index, /// A2_tile - unrankedMemref_index, unrankedMemref_index, /// A3 - unrankedMemref_index, unrankedMemref_index, /// A3_tile - unrankedMemref_f32, i32Type}, - {}); /// last arg (i32Type): readMode - auto readInput3DF64Func = FunctionType::get(ctx, {i32Type, indexType, indexType, indexType, indexType, indexType, indexType, /// Dimensions - unrankedMemref_index, unrankedMemref_index, /// A1 - unrankedMemref_index, unrankedMemref_index, /// A1_tile - unrankedMemref_index, unrankedMemref_index, /// A2 - unrankedMemref_index, unrankedMemref_index, /// A2_tile - unrankedMemref_index, unrankedMemref_index, /// A3 - unrankedMemref_index, unrankedMemref_index, /// A3_tile - unrankedMemref_f64, i32Type}, - {}); - - if (VALUETYPE.compare("f32") == 0) - { - std::string func_name = "read_input_3D_f32"; - if (!hasFuncDeclaration(module, func_name)) - { - func::FuncOp func1 = func::FuncOp::create(function.getLoc(), func_name, - readInput3DF32Func, ArrayRef{}); - func1.setPrivate(); - module.push_back(func1); - } - } - else - { - std::string func_name = "read_input_3D_f64"; - if (!hasFuncDeclaration(module, func_name)) - { - comet_debug() << " Insert read_input_3D_f64 decl\n"; - func::FuncOp func1 = func::FuncOp::create(function.getLoc(), func_name, - readInput3DF64Func, ArrayRef{}); - func1.setPrivate(); - module.push_back(func1); - } - } - - auto readInputSizes3DF64Func = FunctionType::get(ctx, {i32Type, indexType, indexType, indexType, indexType, indexType, indexType, unrankedMemref_index, i32Type}, {}); /// last arg (i32Type): readMode - - if (VALUETYPE.compare("f32") == 0) - { - - std::string func_name = "read_input_sizes_3D_f32"; - if (!hasFuncDeclaration(module, func_name)) - { - func::FuncOp func1 = func::FuncOp::create(function.getLoc(), func_name, - readInputSizes3DF64Func, ArrayRef{}); - func1.setPrivate(); - module.push_back(func1); - } - } - else - { - std::string func_name = "read_input_sizes_3D_f64"; - if (!hasFuncDeclaration(module, func_name)) - { - comet_debug() << " Insert read_input_sizes_3D_f64 decl\n"; - func::FuncOp func1 = func::FuncOp::create(function.getLoc(), func_name, - readInputSizes3DF64Func, ArrayRef{}); - func1.setPrivate(); - module.push_back(func1); - } - } + Value dims = rewriter.create(loc, ValueRange{dimSizes[0], dimSizes[1], dimSizes[2]}); + + sptensor = rewriter.create(loc, ty, dims, + ValueRange{ + tensorload_sizes_vec[0], /// A1pos (each dimension consists of pos and crd arrays) + tensorload_sizes_vec[4], /// A2pos + tensorload_sizes_vec[8], /// A3pos + }, + ValueRange{ + tensorload_sizes_vec[1], /// A1crd + tensorload_sizes_vec[5], /// A2crd + tensorload_sizes_vec[9], /// A3crd + }, + ValueRange{ + tensorload_sizes_vec[2], /// A1tile_pos + tensorload_sizes_vec[6], /// A2tile_pos + tensorload_sizes_vec[10], /// A3tile_pos + }, + ValueRange{ + tensorload_sizes_vec[3], /// A1tile_crd + tensorload_sizes_vec[7], /// A2tile_crd + tensorload_sizes_vec[11], /// A3tile_crd + }, + // ValueRange{ + tensorload_sizes_vec[12], /// Aval + 3, ArrayAttr::get(rewriter.getContext(), dim_formats_attr)); } else { - llvm::errs() << __LINE__ << "Not supported dims\n"; + llvm::errs() << __FILE__ << ":" << __LINE__ << "ERROR: Not supported format (Tensors of dimensions greater than 3 are currently not supported).\n"; } + + comet_debug() << "SparseTensorConstructOp generated for sparse output tensor:\n"; + comet_vdump(sptensor); + + return sptensor; } /// This a common lowering function used to lower SparseOutputTensorDeclOp and TempSparseOutputTensorDeclOp @@ -476,37 +268,37 @@ namespace { comet_debug() << "lowerSparseOutputTensorDec::TempSparseOutputTensorDeclOp lowering\n"; } + else + { + assert(false && "Op should be either SparseOutputTensorDeclOp or TempSparseOutputTensorDeclOp"); + } + + SparseTensorType spType = mlir::cast(op->getResultTypes()[0]); - assert(isa(op) || (isa(op) && - "Op should be either SparseOutputTensorDeclOp or TempSparseOutputTensorDeclOp")); comet_vdump(op); auto loc = op.getLoc(); - StringRef formatsAttr = op.getFormat(); - std::string formats_str(formatsAttr.data()); - comet_debug() << " --- " << formats_str << "\n"; comet_debug() << " " << op.getNumOperands() << "\n"; - auto rank_size = op.getResult().getType().getRank(); + auto rank_size = mlir::cast(op.getResult().getType()).getRank(); - IndexType indexType = IndexType::get(op.getContext()); - FloatType f64Type = FloatType::getF64(op.getContext()); - if (VALUETYPE.compare(0, 3, "f32") == 0) - f64Type = FloatType::getF32(op.getContext()); + // IndexType indexType = IndexType::get(op.getContext()); + Type valsType = spType.getElementType(); + Type indicesType = spType.getIndicesType(); /// A1_pos ... A_value - auto dynamicmemTy_1d_index = MemRefType::get({ShapedType::kDynamic}, indexType); /// memref - auto dynamicmemTy_1d_f64 = MemRefType::get({ShapedType::kDynamic}, f64Type); /// memref + auto dynamicmemTy_1d_vals_type = MemRefType::get({ShapedType::kDynamic}, valsType); /// memref + auto dynamicmemTy_1d_indices_type = MemRefType::get({ShapedType::kDynamic}, indicesType); + - comet_debug() << " " << formats_str << " isDense: " << isDense(formats_str, ", ") << "\n"; + Value new_tensor; /// sparse output - if (isDense(formats_str, ", ") == false) + if (!isa(op.getResult().getType())) { /// search read_from_file function call to get the input file name /// Currently, has no filename - std::vector tensorload_sizes_vec; std::vector array_sizes_vec; /// Store the size of C1pos, C1crd,..., Cval,C_dim1_size, C_dim2_size.... /// No need to read from file @@ -582,15 +374,6 @@ namespace dstIndexLocInSrcVec.push_back(dstIndexLocInSrc); } - ArrayAttr allFormats = transpose_op.getFormats(); - std::vector allFormatsStr; - for (unsigned int i = 0; i < allFormats.size(); i++) - { - std::string formats_str(allFormats[i].cast().getValue()); - allFormatsStr.push_back(formats_str); - } - std::string src_format = allFormatsStr[0]; - std::string dst_format = allFormatsStr[1]; /// If in COO format, then the sizes are the same as the input /// for A and B: 2x+1 + 2x+1 + x = 5x+2 @@ -600,22 +383,14 @@ namespace comet_vdump(dst_input); comet_debug() << " "; comet_pdump(dst_input.getDefiningOp()); - mlir::TensorType type; - auto res = dst_input.getDefiningOp()->getResult(0); - if (res.getType().isa()) - { - type = res.getType().cast(); - } - else - { - assert(false && "Expected TensorType"); - } + mlir::tensorAlgebra::SparseTensorType type = cast(dst_input.getType()); + // unsigned int dst_rank = dst_input.getDefiningOp()->getNumOperands(); unsigned int dst_rank = type.getRank(); for (unsigned int i = 0; i < dst_rank; i++) { /// 4*rank+2 + i - dimSizes.push_back(src_input.getDefiningOp()->getOperand(8 * dst_rank + 2 + dstIndexLocInSrcVec[i])); + dimSizes.push_back(rewriter.create(loc, src_input, rewriter.getI32IntegerAttr(i))); } Value cst_index_0 = rewriter.create(loc, IndexType::get(op.getContext()), rewriter.getIndexAttr(0)); @@ -627,26 +402,19 @@ namespace /// For COO format, 2D and 3D are the same /// if src format is in COO format, - if (src_format.compare("COO") == 0) + auto srcType = cast(transpose_op.getOperand(0).getType()); + auto src_format = getTensorFormatString(srcType); + if (src_format.compare("COO")) { for (unsigned int i = 0; i < dst_rank; i++) { - /// 2*dst_rank+1 - unsigned int dstIndexLocInSrc = dstIndexLocInSrcVec[i]; - /// src_rank = dst_rank - unsigned int posLocInSrc = (4 * dst_rank + 1) + 4 * dstIndexLocInSrc; - unsigned int crdLocInSrc = posLocInSrc + 1; - - unsigned int posLocInSrc2 = posLocInSrc + 2; - unsigned int crdLocInSrc2 = crdLocInSrc + 2; - - array_sizes_vec.push_back(src_input.getDefiningOp()->getOperand(posLocInSrc)); - array_sizes_vec.push_back(src_input.getDefiningOp()->getOperand(crdLocInSrc)); - array_sizes_vec.push_back(src_input.getDefiningOp()->getOperand(posLocInSrc2)); - array_sizes_vec.push_back(src_input.getDefiningOp()->getOperand(crdLocInSrc2)); + array_sizes_vec.push_back(rewriter.create(loc, rewriter.create(loc, src_input, rewriter.getI32IntegerAttr(0)), 0)); + array_sizes_vec.push_back(rewriter.create(loc, rewriter.create(loc, src_input, rewriter.getI32IntegerAttr(0)), 0)); + array_sizes_vec.push_back(rewriter.create(loc, rewriter.create(loc, src_input, rewriter.getI32IntegerAttr(1)), 0)); + array_sizes_vec.push_back(rewriter.create(loc, rewriter.create(loc, src_input, rewriter.getI32IntegerAttr(1)), 0)); } /// val array size - array_sizes_vec.push_back(src_input.getDefiningOp()->getOperand(8 * dst_rank + 1)); + array_sizes_vec.push_back(rewriter.create(loc, rewriter.create(loc, RankedTensorType::get({ShapedType::kDynamic}, cast(src_input.getType()).getElementType()), src_input), 0)); /// set the pos array size, 1st dim as 2, all others as 1. for (unsigned int i = 0; i < dst_rank * 2; i++) @@ -680,16 +448,15 @@ namespace comet_vdump(crd_size); array_sizes_vec.push_back(crd_size); /// B2pos, Bval are the same size with A2pos, Aval - /// TODO(gkestor): Do not hardcode - array_sizes_vec.push_back(src_input.getDefiningOp()->getOperand(17)); + mlir::Value vals_size = rewriter.create(loc, rewriter.create(loc, RankedTensorType::get({ShapedType::kDynamic}, cast(src_input.getType()).getElementType()), src_input), 0); + array_sizes_vec.push_back(vals_size); /// A2tile array_sizes_vec.push_back(cst_index_0); array_sizes_vec.push_back(cst_index_0); /// Aval - /// TODO(gkestor): Do not hardcode - array_sizes_vec.push_back(src_input.getDefiningOp()->getOperand(17)); + array_sizes_vec.push_back(vals_size); } else if (src_format.compare("ELL") == 0) { @@ -705,16 +472,15 @@ namespace /// A2 array_sizes_vec.push_back(cst_index_1); - /// TODO(gkestor): Do not hardcode - array_sizes_vec.push_back(src_input.getDefiningOp()->getOperand(14)); + /// TODO(PT): Verify this + array_sizes_vec.push_back(rewriter.create(loc, rewriter.create(loc, src_input, rewriter.getI32IntegerAttr(1)), 0)); /// A2tile array_sizes_vec.push_back(cst_index_0); array_sizes_vec.push_back(cst_index_0); /// Aval - /// TODO(gkestor): Do not hardcode - array_sizes_vec.push_back(src_input.getDefiningOp()->getOperand(17)); + rewriter.create(loc, rewriter.create(loc, src_input), 0); } } /// For 3D, consider CSF @@ -724,7 +490,9 @@ namespace { comet_debug() << " 3D CSF transpose to 3D CSF\n"; array_sizes_vec.push_back(cst_index_2); - mlir::Value src_nnz = src_input.getDefiningOp()->getOperand(25); + mlir::Value vals_size = rewriter.create(loc, rewriter.create(loc, src_input), 0); + + mlir::Value src_nnz = vals_size; mlir::Value src_nnz_add1 = rewriter.create(loc, src_nnz, cst_index_1); comet_debug() << "AddIOp generated for nnz for CSF:\n"; comet_vdump(src_nnz_add1); @@ -747,283 +515,64 @@ namespace comet_debug() << " array_sizes_vec.size(): " << array_sizes_vec.size() << "\n"; comet_debug() << " dst_rank: " << dst_rank << "\n"; + std::vector tensorload_sizes_vec; for (unsigned int i = 0; i < 4 * dst_rank + 1; i++) { Value alloc_sizes; if (i < 4 * dst_rank) { - alloc_sizes = insertAllocAndInitialize(loc, dynamicmemTy_1d_index, ValueRange{array_sizes_vec[i]}, rewriter); + alloc_sizes = insertAllocAndInitialize(loc, dynamicmemTy_1d_indices_type, ValueRange{array_sizes_vec[i]}, rewriter); comet_debug() << " AllocOp: "; comet_vdump(alloc_sizes); } else { - alloc_sizes = insertAllocAndInitialize(loc, dynamicmemTy_1d_f64, ValueRange{array_sizes_vec[i]}, rewriter); + alloc_sizes = insertAllocAndInitialize(loc, dynamicmemTy_1d_vals_type, ValueRange{array_sizes_vec[i]}, rewriter); comet_debug() << " AllocOp: "; comet_vdump(alloc_sizes); } - Value tensorload_sizes = rewriter.create(loc, alloc_sizes); + Value tensorload_sizes = rewriter.create(loc, alloc_sizes, rewriter.getUnitAttr(), rewriter.getUnitAttr()); tensorload_sizes_vec.push_back(tensorload_sizes); } + new_tensor = insertSparseTensorDeclOp(rewriter, op.getContext(), loc, rank_size, tensorload_sizes_vec, array_sizes_vec, allPerms, dimSizes, op.getResult().getType()); + break; } - else if (isa(u)) - { - comet_debug() << " sparse output is used in itComputeOp op\n"; - - /// Set the insertion point before its user - rewriter.setInsertionPoint(u); - - indexTree::IndexTreeComputeLHSOp lhsOp = cast(u); - comet_debug() << " formats_str: " << formats_str << "\n"; - comet_debug() << " current Op: "; - comet_vdump(lhsOp); - - for (auto uLHS : lhsOp.getOperation()->getUsers()) - { - assert(isa(uLHS) && "User of IndexTreeComputeLHSOp can only be IndexTreeComputeOp"); - - comet_debug() << " lhsOp user: "; - comet_pdump(uLHS); - - auto computeOp = cast(uLHS); - comet_debug() << " Get RHS op: "; - comet_vdump(computeOp); - - std::vector> rhsPerms; - getRHSPermsOfComputeOp(computeOp, rhsPerms); - - std::vector> rhsFormats; - getRHSFormatsOfComputeOp(computeOp, rhsFormats); - - comet_debug() << " rhsPerms: \n"; - for (auto m : rhsPerms) - { - comet_debug() << " \n"; - for (auto n : m) - { - comet_debug() << n << " \n"; - } - comet_debug() << "\n"; - } - - comet_debug() << " rhsFormats: \n"; - for (auto m : rhsFormats) - { - comet_debug() << " \n"; - for (auto n : m) - { - comet_debug() << n << " \n"; - } - comet_debug() << "\n"; - } - - bool isElementwise = checkIsElementwise(rhsPerms); - - comet_debug() << "Checking if it is mixed mode\n"; - bool isMixedMode = checkIsMixedMode(rhsFormats); - - comet_debug() << "IsElementWise: " << isElementwise << " isMixedMode: " << isMixedMode << "\n"; - if (isElementwise && isMixedMode) - { - comet_debug() << "It is an elementwise multiplication in mixed Mode sparse = sparse * dense\n"; - if (isMixedMode) - { - comet_debug() << "It is an mix-mode elementwise multiplication in Mix Mode\n"; - mixModeEltWiseMultSparseTensorOutputLowering(computeOp, - loc, - rhsPerms, - dimSizes, - tensorload_sizes_vec, - array_sizes_vec, rewriter); - } - else - { - comet_debug() << "It is an pure-sparse elementwise multiplication\n"; - pureSparseMultSparseTensorOutputLowering<>(op, - loc, - formats_str, - dimSizes, - tensorload_sizes_vec, - array_sizes_vec, - rewriter); - } - } - else - { - if (!isMixedMode) - { - comet_debug() << "It is an pure-sparse multiplication or assigment from dense to sparse (produced after workspace transformations)\n"; - pureSparseMultSparseTensorOutputLowering(op, - loc, - formats_str, - dimSizes, - tensorload_sizes_vec, - array_sizes_vec, - rewriter); - } - else - { - /// TODO(gkestor) Mix-mode sparse computation with sparse output not yet supported such as TTM (tensor times matrix) - /// TODO(gkestor): if the sparsity patterns is known - comet_debug() << "It is an mix mode element-wise multiplication\n"; - mixModeEltWiseMultSparseTensorOutputLowering(computeOp, - loc, - rhsPerms, - dimSizes, - tensorload_sizes_vec, - array_sizes_vec, rewriter); - } - } - } - } - else if (isa(u)) - { - comet_debug() << " Sparse output is used in TensorFillFromFileOp\n"; - auto fillfromfileop = cast(u); - /// Can get filename, from "filename" attribute of fillfromfileop - rewriter.eraseOp(fillfromfileop); - } - else if (isa(u)) - { - comet_debug() << "The tensor is in IndexTreeComputeRHSOp, no action taken\n"; - continue; - } - else if (isa(u)) - { - comet_debug() << "The tensor is in print op, no action taken\n"; - continue; - } - else if (isa(u)) - { - comet_debug() << "The tensor is in sum op, no action taken\n"; - continue; - } - else if (isa(u)) - { - comet_debug() << "The tensor is in dim op, no action taken\n"; - continue; - } - else - { - comet_pdump(u); - llvm::errs() << __FILE__ << __LINE__ << " tensor is used in the following unsupported op\n"; - } - - comet_debug() << " Get users after "; - /// create sparse tensor construct after lowering each sparse tensor output users - comet_debug() << " tensorload_sizes_vec.size(): " << tensorload_sizes_vec.size() << ", rank_size: " << rank_size << "\n"; - /// create sptensor_construct - SmallVector elementTypes; - for (unsigned int i = 0; i < 4 * rank_size + 1; i++) - { - assert(tensorload_sizes_vec.size() > 0 && "ERROR: Please report this error to the developers!"); - comet_debug() << " " << i << " "; - comet_vdump(tensorload_sizes_vec[i]); - elementTypes.push_back(tensorload_sizes_vec[i].getType()); - } - comet_debug() << "\n "; - /// [0 ... 2*rank_size, 2*rank_size+1 ... 4*rank_size+1, 4*rank_size+2 ... 5*rank_size + 1] - /// 2d+1 + 2d+1 + d => 5d+2 - for (unsigned int i = 0; i < 4 * rank_size + 1; i++) - { - assert(array_sizes_vec.size() > 0 && "ERROR: Please report this error to the developers!"); - comet_debug() << " " << i << " "; - comet_vdump(array_sizes_vec[i]); - elementTypes.push_back(array_sizes_vec[i].getType()); - } - comet_debug() << "\n "; - for (unsigned int i = 0; i < rank_size; i++) - { - assert(dimSizes.size() > 0 && "ERROR: Please report this error to the developers!"); - elementTypes.push_back(dimSizes[i].getType()); - } - comet_debug() << "\n "; - - auto ty = tensorAlgebra::SparseTensorType::get(elementTypes); - - Value sptensor; - if (rank_size == 2) - { - sptensor = rewriter.create(loc, ty, - ValueRange{ - tensorload_sizes_vec[0], /// A1pos (each dimension consists of pos and crd arrays) - tensorload_sizes_vec[1], /// A1crd - tensorload_sizes_vec[2], /// A1tile_pos - tensorload_sizes_vec[3], /// A1tile_crd - tensorload_sizes_vec[4], /// A2pos - tensorload_sizes_vec[5], /// A2crd - tensorload_sizes_vec[6], /// A2tile_pos - tensorload_sizes_vec[7], /// A2tile_crd - tensorload_sizes_vec[8], /// Aval - array_sizes_vec[0], /// A1pos_size (size of each pos and crd arrays) - array_sizes_vec[1], /// A1crd_size - array_sizes_vec[2], /// A1tile_pos_size - array_sizes_vec[3], /// A1tile_crd_size - array_sizes_vec[4], /// A2pos_size - array_sizes_vec[5], /// A2crd_size - array_sizes_vec[6], /// A2tile_pos_size - array_sizes_vec[7], /// A2tile_crd_size - array_sizes_vec[8], /// Aval_size (size of value array) - dimSizes[0], /// dim1_size(size of each dimension in sparse tensor) - dimSizes[1] /// dim2_size (size of each dimension in sparse tensor) - }, - 2); - } - else if (rank_size == 3) + else if (isa(u)) { - sptensor = rewriter.create(loc, ty, - ValueRange{ - tensorload_sizes_vec[0], /// A1pos (each dimension consists of pos and crd arrays) - tensorload_sizes_vec[1], /// A1crd - tensorload_sizes_vec[2], /// A1tile_pos - tensorload_sizes_vec[3], /// A1tile_crd - tensorload_sizes_vec[4], /// A2pos - tensorload_sizes_vec[5], /// A2crd - tensorload_sizes_vec[6], /// A2tile_pos - tensorload_sizes_vec[7], /// A2tile_crd - tensorload_sizes_vec[8], /// A3pos - tensorload_sizes_vec[9], /// A3crd - tensorload_sizes_vec[10], /// A3tile_pos - tensorload_sizes_vec[11], /// A3tile_crd - tensorload_sizes_vec[12], /// Aval - array_sizes_vec[0], /// A1pos_size (size of each pos and crd arrays) - array_sizes_vec[1], /// A1crd_size - array_sizes_vec[2], /// A1tile_pos_size - array_sizes_vec[3], /// A1tile_crd_size - array_sizes_vec[4], /// A2pos_size - array_sizes_vec[5], /// A2crd_size - array_sizes_vec[6], /// A2tile_pos_size - array_sizes_vec[7], /// A2tile_crd_size - array_sizes_vec[8], /// A3pos_size - array_sizes_vec[9], /// A3crd_size - array_sizes_vec[10], /// A3tile_pos_size - array_sizes_vec[11], /// A3tile_crd_size - array_sizes_vec[12], /// Aval_size (size of value array) - dimSizes[0], /// dim1_size (size of each dimension in sparse tensor) - dimSizes[1], /// dim2_size (size of each dimension in sparse tensor) - dimSizes[2] /// dim3_size - }, - 3); + comet_debug() << " Sparse output is modified in an index tree\n"; + // Tensor is created as the output of a sparse tensor operation + // For now we defer to the index tree dialect by inserting a tensor decl + // that just contains empty domains. + rank_size = spType.getRank(); + indexTree::DomainType domain_type = indexTree::DomainType::get(op.getContext()); + rewriter.setInsertionPoint(op); + Value empty_domain = rewriter.create(loc, domain_type); + llvm::SmallVector args = llvm::SmallVector(rank_size, empty_domain); + new_tensor = rewriter.create(loc, spType, args); + + // Eventually, there are 2 cases: + // Case 1: We can determine apriori the dimension of the sparse tensor + // This is the case if none of the index variables in the output + // tensor are used in a union or a insersect op. In this case we use + // the sparse tensor decleration of the input in order to determine + // the output tensor. We allocate arrays of the same size and then + // insert a ta.SpTensorDeclOp. + // Case 2: We can't determine the dimension of the sparse tensor. + // This happens in all other cases. Here we insert a tensor + // that is defined with an (at least one) empty domain. In + // the lowering process we can either use the symbolic phase + // to determine the allocations needed, or we can perform the + // allocations during the computational phase + break; } - else - { - llvm::errs() << __FILE__ << ":" << __LINE__ << "ERROR: Not supported format (Tensors of dimensions greater than 3 are currently not supported).\n"; - } - - comet_debug() << "SparseTensorConstructOp generated for sparse output tensor:\n"; - comet_vdump(sptensor); - - /// create ta.index_label operation. - comet_vdump(op); - - op.replaceAllUsesWith(sptensor); - rewriter.replaceOp(op, sptensor); - } /// for (auto u : op.getOperation()->getUsers()) + } + op.replaceAllUsesWith(new_tensor); + rewriter.replaceOp(op, {new_tensor}); } else { /// format == "Dense" - auto resultTensorType = op.getResult().getType().template cast(); + auto resultTensorType = cast(op.getResult().getType()); std::vector cur_indices; std::vector cur_memref; auto resultMemTy = convertTensorToMemRef(resultTensorType); @@ -1040,12 +589,12 @@ namespace } llvm::ArrayRef cur_memref_arrayref = llvm::ArrayRef(cur_memref); - MemRefType memrefType2 = MemRefType::get(cur_memref_arrayref, f64Type); + MemRefType memrefType2 = MemRefType::get(cur_memref_arrayref, valsType); Value alloc_sizes1 = insertAllocAndInitialize(loc, memrefType2, ValueRange(cur_indices), rewriter); comet_debug() << " AllocOp: "; comet_vdump(alloc_sizes1); - Value tensorLoad = rewriter.create(loc, alloc_sizes1); + Value tensorLoad = rewriter.create(loc, alloc_sizes1, rewriter.getUnitAttr(), rewriter.getUnitAttr()); comet_vdump(tensorLoad); op.replaceAllUsesWith(tensorLoad); @@ -1078,7 +627,7 @@ namespace auto resultTensorType = op.getResult().getType(); std::vector cur_indices; std::vector cur_memref; - auto resultMemTy = convertTensorToMemRef(resultTensorType.cast()); + auto resultMemTy = convertTensorToMemRef(cast(resultTensorType)); int j = 0; for (int i = 0; i < resultMemTy.getRank(); i++) @@ -1115,7 +664,7 @@ namespace cast(init_alloc.getDefiningOp()).setAlignmentAttr(rewriter.getI64IntegerAttr(32)); - Value tensorLoad = rewriter.create(loc, init_alloc); + Value tensorLoad = rewriter.create(loc, init_alloc, rewriter.getUnitAttr(), rewriter.getUnitAttr()); comet_debug() << " TensorLoad:\n"; comet_vdump(tensorLoad); @@ -1151,120 +700,56 @@ namespace auto function = cast(op->getParentOp()); auto module = function.getOperation()->getParentOfType(); - std::string op_str = dump2str(op); bool isOutputTensor = false; auto loc = op.getLoc(); - StringRef formatsAttr = op.getFormat(); - std::string formats_str(formatsAttr.data()); + // StringRef formatsAttr = op.getFormat(); + std::string formats_str(getTensorFormatString(op.getType())); comet_debug() << " --- " << formats_str << "\n"; comet_debug() << " " << op.getNumOperands() << "\n"; - auto res = op.getResult(); - mlir::TensorType type; - if(res.getType().isa()) - { - type = res.getType().cast(); - } - else - { - assert(false && "Expected TensorType"); - } + mlir::tensorAlgebra::SparseTensorType type = cast(op.getResult().getType()); auto rank_size = type.getRank(); // auto rank_size = op.getResult().getType().cast().getRank(); // auto rank_size = op.getNumOperands(); IndexType indexType = IndexType::get(op.getContext()); - FloatType f64Type = FloatType::getF64(op.getContext()); - if (VALUETYPE.compare(0, 3, "f32") == 0) - f64Type = FloatType::getF32(op.getContext()); + Type floatEleType = type.getElementType(); + IntegerType indicesType = type.getIndicesType(); for (auto u1 : op.getOperation()->getUsers()) { comet_debug() << "\nCheck the tensor is input or output\n"; comet_pdump(u1); - if (isa(u1)) + if (isa(u1)) { comet_debug() << " used in ta.tc op\n"; - auto p = cast(u1).getOperation(); - for (unsigned int i = 0; i < p->getNumOperands(); i++) + auto p = u1->getOperand(2); + if(p == op) { - /// comet_vdump(n); - std::string n_str = dump2str(p->getOperand(i)); - comet_debug() << "the operands: " << n_str << "\n"; - if (n_str.compare(0, op_str.size(), op_str) == 0) - { - comet_debug() << " FIND IT: " << i << "\n"; - if (i == 2) - { - isOutputTensor = true; - } - } - } - } - else if (isa(u1)) - { - comet_debug() << " used in ta.elews_mul op\n"; - auto p = cast(u1).getOperation(); - for (unsigned int i = 0; i < p->getNumOperands(); i++) - { - std::string n_str = dump2str(p->getOperand(i)); - if (n_str.compare(0, op_str.size(), op_str) == 0) - { - comet_debug() << " FIND IT: " << i << "\n"; - if (i == 2) - { - isOutputTensor = true; - } - } + isOutputTensor = true; } } else if (isa(u1)) { comet_debug() << " used in ta.set op\n"; - auto p = cast(u1).getOperation(); - for (unsigned int i = 0; i < p->getNumOperands(); i++) + auto p = u1->getOperand(1); + if(p == op) { - comet_debug() << " the " << i << "th operand\n"; - std::string n_str = dump2str(p->getOperand(i)); - if (n_str.compare(0, op_str.size(), op_str) == 0) - { - comet_debug() << " FIND IT: " << i << "\n"; - if (i == 1) - { - /// The source tensor of the set op - isOutputTensor = true; - } - } + isOutputTensor = true; } } - else if (isa(u1)) + else if (isa(u1)) { - comet_debug() << " used in transpose op\n"; - auto p = cast(u1).getOperation(); - for (unsigned int i = 0; i < p->getNumOperands(); i++) - { - std::string n_str = dump2str(p->getOperand(i)); - if (n_str.compare(0, op_str.size(), op_str) == 0) - { - comet_debug() << " FIND IT: " << i << "\n"; - if (i == 2) - { - /// output of ta.elews_mul - isOutputTensor = true; - } - } - } + comet_debug() << " used in it.Operand op\n"; } - else if (isa(u1)) + else if (isa(u1)) { - comet_debug() << " used in ta.itComputeRHS op\n"; - isOutputTensor = false; + comet_debug() << " used in it.TensorAccess op\n"; } - else if (isa(u1)) + else if (isa(u1)) { - comet_debug() << " used in ta.itComputeLHS op\n"; - isOutputTensor = true; + comet_debug() << " used in it.Domain op\n"; } else if (isa(u1)) { @@ -1292,26 +777,63 @@ namespace { comet_debug() << " the tensor has use in TensorDimOp and this use will be ignored!\n"; } + else if (isa(u1)) + { + /// do nothing! + comet_debug() << " the tensor has use in AllocWorkspaceOp\n"; + } + else if(isa(u1)) + { + /// do nothing! + comet_debug() << " the tensor has use in func::CallOp\n"; + } + else if(isa(u1)) + { + comet_debug() << " the tensor has use in func::ReturnOp\n"; + } + else if (isa(u1)) + { + /// Tensors that are modified are passed as arguments to the index tree + comet_debug() << " the tensor has use in a Index Tree\n"; + isOutputTensor = true; + } + else if (isa(u1)) + { + /// do nothing! + comet_debug() << " the tensor has use in alias op\n"; + } + else if (isa(u1) + || isa(u1)) + { + /// do nothing! + comet_debug() << " the tensor has use in Domain op\n"; + } + else if (isa(u1) + || isa(u1) + || isa(u1)) + { + /// do nothing! + comet_debug() << " the tensor has use in Sparse Tensor Get Dim op\n"; + } else { u1->dump(); - llvm::errs() << __FILE__ << ":" << __LINE__ << "The tensor is in not supported operation\n"; + llvm::errs() << __FILE__ << ":" << __LINE__ << " The tensor is in not supported operation\n"; } } comet_debug() << " isOutputTensor: " << isOutputTensor << "\n"; /// A1_pos ... A_value - auto dynamicmemTy_1d_index = MemRefType::get({ShapedType::kDynamic}, indexType); /// memref - auto dynamicmemTy_1d_f64 = MemRefType::get({ShapedType::kDynamic}, f64Type); /// memref + auto dynamicmemTy_1d_float = MemRefType::get({ShapedType::kDynamic}, floatEleType); /// memref Type unrankedMemTy_index = UnrankedMemRefType::get(indexType, 0); - Type unrankedMemTy_f64 = UnrankedMemRefType::get(f64Type, 0); + Type unrankedMemTy_float = UnrankedMemRefType::get(floatEleType, 0); comet_debug() << " " << formats_str << " isDense: " << isDense(formats_str, ", ") << "\n"; /// tensor is sparse and input. - if (isDense(formats_str, ", ") == false && isOutputTensor == false) + if (!isa(op.getType()) && isOutputTensor == false) { comet_debug() << " Sparse input tensor \n"; @@ -1327,8 +849,8 @@ namespace { auto fillfromfileop = cast(u); /// Can get filename, from "filename" attribute of fillfromfileop - StringAttr filename = fillfromfileop.getFilename().cast(); - IntegerAttr readModeAttr = fillfromfileop.getReadMode().cast(); + StringAttr filename = cast(fillfromfileop.getFilename()); + IntegerAttr readModeAttr = cast(fillfromfileop.getReadMode()); rewriter.eraseOp(fillfromfileop); comet_debug() << " filename: " << filename.getValue() << "\n"; @@ -1351,6 +873,13 @@ namespace Value alloc_sizes_cast = rewriter.create(loc, unrankedMemTy_index, alloc_sizes); std::vector dim_format = mlir::tensorAlgebra::getFormatsValue(formats_str, rank_size, rewriter, loc, indexType); + std::vector dim_format_int = mlir::tensorAlgebra::getFormats(formats_str, rank_size, ctx); + llvm::SmallVector dim_format_attr; + for(TensorFormatEnum& format: dim_format_int) + { + dim_format_attr.push_back(TensorFormatEnumAttr::get(ctx, format)); + } + auto dim_format_attrs = ArrayAttr::get(ctx, ArrayRef(dim_format_attr)); comet_debug() << " Get the dim_format\n"; /// inform the runtime of what env var to use for parsing input file @@ -1390,17 +919,21 @@ namespace { /// 2D comet_debug() << " 2D\n"; /// Add function definition to the module - insertReadFileLibCall(rank_size, ctx, module, function); + insertReadFileLibCall(rank_size, floatEleType, indicesType, ctx, module, function); std::string read_input_sizes_str; - if (VALUETYPE.compare(0, 3, "f32") == 0) + if (floatEleType.isF32()) { read_input_sizes_str = "read_input_sizes_2D_f32"; } - else + else if(floatEleType.isF64()) { read_input_sizes_str = "read_input_sizes_2D_f64"; } + else + { + assert(false && "Unexpected data type"); + } auto read_input_sizes_Call = rewriter.create(loc, read_input_sizes_str, SmallVector{}, ValueRange{sparseFileID, dim_format[0], dim_format[1], dim_format[2], dim_format[3], @@ -1411,17 +944,22 @@ namespace { /// 3D comet_debug() << " 3D\n"; /// Add function definition to the module - insertReadFileLibCall(rank_size, ctx, module, function); + insertReadFileLibCall(rank_size, floatEleType, indicesType, ctx, module, function); + std::string read_input_sizes_str; - if (VALUETYPE.compare(0, 3, "f32") == 0) + if (floatEleType.isF32()) { read_input_sizes_str = "read_input_sizes_3D_f32"; } - else - { /// default f64 + else if(floatEleType.isF64()) + { read_input_sizes_str = "read_input_sizes_3D_f64"; } + else + { + assert(false && "Unexpected data type"); + } auto read_input_sizes_3D_Call = rewriter.create(loc, read_input_sizes_str, SmallVector{}, ValueRange{sparseFileID, dim_format[0], dim_format[1], /// A1, A1_tile @@ -1448,17 +986,22 @@ namespace std::vector alloc_sizes_cast_vec; std::vector alloc_sizes_vec; + + /// A1_pos ... A_value + auto dynamicmemTy_1d_indices_type = MemRefType::get({ShapedType::kDynamic}, type.getIndicesType()); /// memref + Type unrankedMemTy_indices_type = UnrankedMemRefType::get(type.getIndicesType(), 0); + for (unsigned int i = 0; i < sp_decl.getDimArrayCount(); i++) { std::vector idxes; idxes.push_back(array_sizes[i]); comet_vdump(array_sizes[i]); - Value alloc_size = insertAllocAndInitialize(loc, dynamicmemTy_1d_index, ValueRange{idxes}, rewriter); + Value alloc_size = insertAllocAndInitialize(loc, dynamicmemTy_1d_indices_type, ValueRange{idxes}, rewriter); comet_debug() << " "; comet_vdump(alloc_size); alloc_sizes_vec.push_back(alloc_size); - Value alloc_size_cast = rewriter.create(loc, unrankedMemTy_index, alloc_size); + Value alloc_size_cast = rewriter.create(loc, unrankedMemTy_indices_type, alloc_size); alloc_sizes_cast_vec.push_back(alloc_size_cast); } @@ -1466,11 +1009,11 @@ namespace { std::vector idxes; idxes.push_back(array_sizes[i]); - Value alloc_size = insertAllocAndInitialize(loc, dynamicmemTy_1d_f64, ValueRange{idxes}, rewriter); + Value alloc_size = insertAllocAndInitialize(loc, dynamicmemTy_1d_float, ValueRange{idxes}, rewriter); comet_debug() << " "; comet_vdump(alloc_size); alloc_sizes_vec.push_back(alloc_size); - Value alloc_size_cast = rewriter.create(loc, unrankedMemTy_f64, alloc_size); + Value alloc_size_cast = rewriter.create(loc, unrankedMemTy_float, alloc_size); alloc_sizes_cast_vec.push_back(alloc_size_cast); } @@ -1478,14 +1021,22 @@ namespace if (rank_size == 2) { /// 2D std::string read_input_str; - if (VALUETYPE.compare(0, 3, "f32") == 0) + + if (floatEleType.isF32()) { read_input_str = "read_input_2D_f32"; } - else + else if (floatEleType.isF64()) { read_input_str = "read_input_2D_f64"; } + else + { + assert(false && "Unexpected type"); + } + + read_input_str += "_i"+std::to_string(indicesType.getWidth()); + auto read_input_f64Call = rewriter.create(loc, read_input_str, SmallVector{}, ValueRange{sparseFileID, dim_format[0], dim_format[1], /// A1_format, A1_tile_format @@ -1504,14 +1055,17 @@ namespace else if (rank_size == 3) { /// 3D std::string read_input_str; - if (VALUETYPE.compare(0, 3, "f32") == 0) + if (floatEleType.isF32()) { read_input_str = "read_input_3D_f32"; } - else + else if (floatEleType.isF64()) { read_input_str = "read_input_3D_f64"; } + + read_input_str += "_i"+std::to_string(indicesType.getWidth()); + auto read_input_f64Call = rewriter.create(loc, read_input_str, SmallVector{}, ValueRange{sparseFileID, dim_format[0], dim_format[1], /// A1, A1_tile @@ -1535,43 +1089,65 @@ namespace std::vector alloc_tensor_vec; for (unsigned int i = 0; i < sp_decl.getTotalArrayCount(); i++) { - Value tensorLoad = rewriter.create(loc, alloc_sizes_vec[i]); + Value tensorLoad = rewriter.create(loc, alloc_sizes_vec[i], rewriter.getUnitAttr(), rewriter.getUnitAttr()); alloc_tensor_vec.push_back(tensorLoad); } - /// create sptensor_construct - SmallVector elementTypes; - for (unsigned int i = 0; i < sp_decl.getTotalArrayCount(); i++) - { - elementTypes.push_back(alloc_tensor_vec[i].getType()); - } - for (unsigned int i = 0; i < 5 * rank_size + 1; i++) - { - elementTypes.push_back(array_sizes[i].getType()); - } - - auto ty = tensorAlgebra::SparseTensorType::get(elementTypes); + llvm::SmallVector dim_sizes(rank_size, ShapedType::kDynamic); // TODO: Determine sizes!!!! + auto ty = op.getResult().getType(); Value sptensor; if (rank_size == 2) { - sptensor = rewriter.create(loc, ty, ValueRange{alloc_tensor_vec[0], alloc_tensor_vec[1], /// A1 - alloc_tensor_vec[2], alloc_tensor_vec[3], /// A1_tile - alloc_tensor_vec[4], alloc_tensor_vec[5], /// A2 - alloc_tensor_vec[6], alloc_tensor_vec[7], /// A2_tile - alloc_tensor_vec[8], array_sizes[0], array_sizes[1], array_sizes[2], array_sizes[3], array_sizes[4], array_sizes[5], array_sizes[6], array_sizes[7], array_sizes[8], array_sizes[9], array_sizes[10]}, - 2); + Value dims = rewriter.create(loc, ValueRange{array_sizes[9], array_sizes[10]}); /// I, J + sptensor = rewriter.create(loc, ty, + dims, /// Dim sizes + ValueRange{ + alloc_tensor_vec[0], // A1_pos + alloc_tensor_vec[4], /// A2_pos + }, + ValueRange{ + alloc_tensor_vec[1], /// A1_crd + alloc_tensor_vec[5], /// A2_crd + }, + ValueRange{ + alloc_tensor_vec[2], /// A1_tile_pos + alloc_tensor_vec[6], /// A2_tile_pos + }, + ValueRange{ + alloc_tensor_vec[3], /// A1_tile_crd + alloc_tensor_vec[7], /// A2_tile_crd + }, + alloc_tensor_vec[8], /// Avals + 2, dim_format_attrs); } else if (rank_size == 3) { - sptensor = rewriter.create(loc, ty, ValueRange{alloc_tensor_vec[0], alloc_tensor_vec[1], /// A1 - alloc_tensor_vec[2], alloc_tensor_vec[3], /// A1_tile - alloc_tensor_vec[4], alloc_tensor_vec[5], /// A2 - alloc_tensor_vec[6], alloc_tensor_vec[7], /// A2_tile - alloc_tensor_vec[8], alloc_tensor_vec[9], /// A3 - alloc_tensor_vec[10], alloc_tensor_vec[11], /// A3_tile - alloc_tensor_vec[12], array_sizes[0], array_sizes[1], array_sizes[2], array_sizes[3], array_sizes[4], array_sizes[5], array_sizes[6], array_sizes[7], array_sizes[8], array_sizes[9], array_sizes[10], array_sizes[11], array_sizes[12], array_sizes[13], array_sizes[14], array_sizes[15], array_sizes[16], array_sizes[17], array_sizes[18]}, - 3); + Value dims = rewriter.create(loc, ValueRange{array_sizes[13], array_sizes[14], array_sizes[15]}); /// I, J, K + + sptensor = rewriter.create(loc, ty, dims, + ValueRange { + alloc_tensor_vec[0], /// A1_pos + alloc_tensor_vec[4], /// A2_pos + alloc_tensor_vec[8], /// A3_pos + }, + ValueRange { + alloc_tensor_vec[1], /// A1_crd + alloc_tensor_vec[5], /// A2_crd + alloc_tensor_vec[9], /// A3_crd + }, + ValueRange { + alloc_tensor_vec[2], /// A1_tile_pos + alloc_tensor_vec[6], /// A2_tile_pos + alloc_tensor_vec[10], /// A3_tile_pos + }, + ValueRange { + alloc_tensor_vec[3], /// A1_tile_crd + alloc_tensor_vec[7], /// A2_tile_crd + alloc_tensor_vec[11], /// A3_tile_crd + }, + alloc_tensor_vec[12], /// Avals + 3, dim_format_attrs); } else { @@ -1591,7 +1167,6 @@ namespace /// Is sparse output ,lower to ta.output_tensor_decl auto tensor_decl_value = cast(op); auto labels = tensor_decl_value.getLabels(); - auto tensor_format = tensor_decl_value.getFormat(); auto tensor_type = tensor_decl_value.getType(); auto is_temporal_tensor = tensor_decl_value.getTemporalTensor(); @@ -1600,13 +1175,13 @@ namespace { /// TempSparseOutputTensorDeclOp should be lowered before SparseOutputTensorDeclOp outputtensordecl = rewriter.create(loc, - tensor_type, labels, tensor_format); + tensor_type, labels); comet_debug() << "Gokcen\n"; comet_vdump(outputtensordecl); } else outputtensordecl = rewriter.create(loc, - tensor_type, labels, tensor_format); + tensor_type, labels); comet_debug() << "SparseOutputTensorDecl or TempSparseOutputTensorDeclOp Operation is generated\n"; comet_vdump(outputtensordecl); op.replaceAllUsesWith(outputtensordecl); @@ -1688,6 +1263,7 @@ namespace memref::MemRefDialect, scf::SCFDialect, bufferization::BufferizationDialect, + linalg::LinalgDialect, IndexTreeDialect>(); target.addLegalOp(); + tensorAlgebra::SparseTensorConstructOp, + tensorAlgebra::SpTensorAliasOp, + tensorAlgebra::SpTensorGetDimPos, + tensorAlgebra::SpTensorGetDimCrd, + tensorAlgebra::SpTensorGetVals, + tensorAlgebra::TensorSortOp, + tensorAlgebra::SpTensorGetDimSize>(); if (failed(applyPartialConversion(function, target, std::move(patterns)))) { @@ -1731,6 +1314,7 @@ namespace scf::SCFDialect, mlir::memref::MemRefDialect, IndexTreeDialect, + tensor::TensorDialect, bufferization::BufferizationDialect>(); target.addLegalOp(); @@ -1778,6 +1363,7 @@ namespace scf::SCFDialect, mlir::memref::MemRefDialect, IndexTreeDialect, + tensor::TensorDialect, bufferization::BufferizationDialect>(); target.addIllegalDialect(); @@ -1794,7 +1380,15 @@ namespace tensorAlgebra::IndexLabelOp, tensorAlgebra::DenseConstantOp, tensorAlgebra::TensorDimOp, + tensorAlgebra::SpTensorAliasOp, + tensorAlgebra::SpTensorGetDimCrd, + tensorAlgebra::SpTensorGetDimPos, + tensorAlgebra::SpTensorGetDimSize, + tensorAlgebra::SpTensorGetVals, tensorAlgebra::ScalarOp, + tensorAlgebra::AllocWorkspaceOp, + tensorAlgebra::TensorMultOp, // Should this be dynamically legal to only work with dense tensors? + tensorAlgebra::TensorSortOp, func::CallOp>(); if (failed(applyPartialConversion(function, target, std::move(patterns)))) @@ -1829,6 +1423,7 @@ namespace scf::SCFDialect, mlir::memref::MemRefDialect, IndexTreeDialect, + tensor::TensorDialect, bufferization::BufferizationDialect>(); target.addIllegalDialect(); @@ -1845,6 +1440,14 @@ namespace tensorAlgebra::DenseConstantOp, tensorAlgebra::TensorDimOp, tensorAlgebra::ScalarOp, + tensorAlgebra::AllocWorkspaceOp, + tensorAlgebra::SpTensorAliasOp, + tensorAlgebra::SpTensorGetDimCrd, + tensorAlgebra::SpTensorGetDimPos, + tensorAlgebra::SpTensorGetDimSize, + tensorAlgebra::SpTensorGetVals, + tensorAlgebra::TensorMultOp, // Should this be dynamically legal to only work with dense tensors? + tensorAlgebra::TensorSortOp, func::CallOp>(); if (failed(applyPartialConversion(function, target, std::move(patterns)))) diff --git a/lib/Dialect/TensorAlgebra/Transforms/WorkspaceOptimizations.cpp b/lib/Dialect/TensorAlgebra/Transforms/WorkspaceOptimizations.cpp new file mode 100644 index 00000000..ffc73183 --- /dev/null +++ b/lib/Dialect/TensorAlgebra/Transforms/WorkspaceOptimizations.cpp @@ -0,0 +1,147 @@ +// +// Copyright 2022 Battelle Memorial Institute +// +// Redistribution and use in source and binary forms, with or without modification, +// are permitted provided that the following conditions are met: +// +// 1. Redistributions of source code must retain the above copyright notice, this list of conditions +// and the following disclaimer. +// +// 2. Redistributions in binary form must reproduce the above copyright notice, this list of conditions +// and the following disclaimer in the documentation and/or other materials provided with the distribution. +// +// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED +// WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR +// PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, +// INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE +// GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, +// WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +// OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. +// + +#include "mlir/Dialect/Func/IR/FuncOps.h" +#include "mlir/IR/BuiltinOps.h" +#include "mlir/IR/BuiltinTypeInterfaces.h" +#include "mlir/IR/Value.h" +#include "mlir/Support/LLVM.h" +#include "mlir/Transforms/GreedyPatternRewriteDriver.h" +#include "mlir/Dialect/MemRef/IR/MemRef.h" +#include "mlir/Dialect/Math/IR/Math.h" +#include "mlir/Dialect/Tensor/IR/Tensor.h" +#include "mlir/Dialect/Index/IR/IndexOps.h" +#include "mlir/Pass/Pass.h" + +#include "llvm/ADT/StringSet.h" +#include "llvm/ADT/SmallVector.h" +#include "llvm/ADT/IndexedMap.h" + +#include "comet/Dialect/TensorAlgebra/IR/TADialect.h" +#include "comet/Dialect/TensorAlgebra/Patterns.h" +#include "comet/Dialect/TensorAlgebra/Passes.h" + +using namespace mlir; +using namespace mlir::tensorAlgebra; + +namespace mlir { + namespace comet{ + #define GEN_PASS_DEF_TENSORALGEBRAWORKSPACEOPTIMIZATIONS + #include "comet/Dialect/TensorAlgebra/Passes.h.inc" + } +} + +struct WorkspaceMergeExtract : public OpRewritePattern { + WorkspaceMergeExtract(MLIRContext *context) + : OpRewritePattern(context, /*benefit=*/1) {} + + mlir::LogicalResult + matchAndRewrite(TensorInsertOp op, + mlir::PatternRewriter &rewriter) const override { + Value workspace = op.getTensor(); + WorkspaceType workspace_type = llvm::dyn_cast(workspace.getType()); + if(!workspace_type) { + return failure(); + } + + Value to_insert = op.getValue(); + + //TODO: Right now we only support merging add operations + auto accumulate_op = to_insert.getDefiningOp(); + if(!accumulate_op || !accumulate_op->hasOneUse()) { + return failure(); + } + + auto extract_op = accumulate_op.getOperand(0).getDefiningOp(); + if(!extract_op || !extract_op->hasOneUse()) { + return failure(); + } + + if(extract_op.getTensor() != op.getTensor()){ + return failure(); + } + + rewriter.replaceOpWithNewOp(op, workspace_type, workspace, extract_op.getPos(), extract_op.getCrds(), accumulate_op.getOperand(1)); + rewriter.eraseOp(accumulate_op); + rewriter.eraseOp(extract_op); + return success(); + } +}; + +struct WorkspaceExtractNoCheck : public OpRewritePattern { + WorkspaceExtractNoCheck(MLIRContext *context) + : OpRewritePattern(context, /*benefit=*/1) {} + + mlir::LogicalResult + matchAndRewrite(TensorExtractOp op, + mlir::PatternRewriter &rewriter) const override { + Value workspace = op.getTensor(); + WorkspaceType workspace_type = llvm::dyn_cast(workspace.getType()); + if(!workspace_type) { + return failure(); + } + + Value crd = op.getCrds()[0]; // TODO: Deal with greater than 1 dimensional workspace + + Operation* src = crd.getDefiningOp(); + // Deal with casting + if(llvm::isa(src)){ + src = src->getOperand(0).getDefiningOp(); + } + + SpTensorGetCrd getCrdOp = llvm::dyn_cast(src); + if(!getCrdOp) { + return failure(); + } + + if(getCrdOp.getTensor() != workspace){ + return failure(); + } + + rewriter.replaceOpWithNewOp(op, workspace_type.getElementType(), workspace, getCrdOp.getIdx(), crd); + return success(); + } +}; + + +void tensorAlgebra::populateWorkspaceOptimizationPatterns( + MLIRContext *context, RewritePatternSet &patterns) { + patterns.add(context); +} + +struct TensorAlgebraWorkspaceOptimizations : comet::impl::TensorAlgebraWorkspaceOptimizationsBase { + using TensorAlgebraWorkspaceOptimizationsBase::TensorAlgebraWorkspaceOptimizationsBase; + + void runOnOperation() override { + mlir::RewritePatternSet workspace_optimization_patterns(&getContext()); + tensorAlgebra::populateWorkspaceOptimizationPatterns(&getContext(), workspace_optimization_patterns); + + if(failed(mlir::applyPatternsAndFoldGreedily(getOperation(), std::move(workspace_optimization_patterns)))) { + return signalPassFailure(); + } + } +}; + +/// Apply the compressed workspace transformations on the index tree IR +std::unique_ptr mlir::comet::createWorkspaceOptimizationsPass() +{ + return std::make_unique(); +} \ No newline at end of file diff --git a/lib/Dialect/Utils/Utils.cpp b/lib/Dialect/Utils/Utils.cpp index 7c001081..f2a27110 100644 --- a/lib/Dialect/Utils/Utils.cpp +++ b/lib/Dialect/Utils/Utils.cpp @@ -28,13 +28,16 @@ #include "comet/Dialect/IndexTree/IR/IndexTreeDialect.h" #include "comet/Dialect/Utils/Utils.h" +#include "mlir/Dialect/Linalg/IR/Linalg.h" #include "mlir/Dialect/MemRef/IR/MemRef.h" #include "mlir/Dialect/Affine/IR/AffineOps.h" #include "mlir/Dialect/Func/IR/FuncOps.h" +#include "mlir/Support/LLVM.h" #include "llvm/Support/Debug.h" #include +#include #define DEBUG_TYPE "ta-utils" @@ -45,7 +48,7 @@ // *********** For debug purpose *********// /// TODO(gkestor): supports only f64 - need generalization -std::string VALUETYPE = "f64"; +// std::string VALUETYPE = "f64"; using namespace mlir::arith; using namespace mlir::affine; @@ -136,48 +139,36 @@ namespace mlir auto elementType = memtype.getElementType(); Value cst_init; - if (elementType.isF64()) - { - comet_debug() << "Element type F64\n"; - cst_init = rewriter.create(loc, rewriter.getF64FloatAttr(0.0)); - } - else if (elementType.isF32()) - { - comet_debug() << "Element type F32\n"; - cst_init = rewriter.create(loc, rewriter.getF32FloatAttr(0.0)); - } - else if (elementType.isIndex()) + // if (elementType.isF64()) + // { + // comet_debug() << "Element type F64\n"; + // cst_init = rewriter.create(loc, rewriter.getZeroAttr(elementType)); + // } + // else if (elementType.isF32()) + // { + // comet_debug() << "Element type F32\n"; + // cst_init = rewriter.create(loc, rewriter.getF32FloatAttr(0.0)); + // } + if (elementType.isIndex()) { comet_debug() << "Element type Index\n"; cst_init = rewriter.create(loc, 0); } - else if (elementType.isInteger(1)) - { - comet_debug() << "Element type I1 - boolean\n"; - cst_init = rewriter.create(loc, rewriter.getI1Type(), rewriter.getBoolAttr(0)); - } - else + else { - llvm::errs() << __FILE__ << ":" << __LINE__ << "Not supported memory reference type. Supported element Types are F32, F64, Index \n"; + cst_init = rewriter.create(loc, rewriter.getZeroAttr(elementType)); } + // else if (elementType.isInteger(1)) + // { + // comet_debug() << "Element type I1 - boolean\n"; + // cst_init = rewriter.create(loc, rewriter.getI1Type(), rewriter.getBoolAttr(0)); + // } + // else + // { + // llvm::errs() << __FILE__ << ":" << __LINE__ << "Not supported memory reference type. Supported element Types are F32, F64, Index \n"; + // } - /// TODO(gkestor): add better initialization method based on the dimension, leverage existing operations for initialization - auto lowerBound = rewriter.create(loc, 0); - auto upperBound = alloc_op.getDefiningOp()->getOperand(0); - auto step = rewriter.create(loc, 1); - auto loop = rewriter.create(loc, lowerBound, upperBound, step); - auto insertPt = rewriter.saveInsertionPoint(); - rewriter.setInsertionPointToStart(loop.getBody()); - - /// Build loop body - std::vector indices = {loop.getInductionVar()}; - rewriter.create(loc, cst_init, alloc_op, ValueRange{indices}); - - /// need to restore the insertion point to the previous point - rewriter.restoreInsertionPoint(insertPt); - comet_debug() << " insertAllocAndInitialize loop " - << "\n"; - comet_vdump(loop); + rewriter.create(loc, ValueRange(cst_init), ValueRange(alloc_op)); return alloc_op; } @@ -192,7 +183,7 @@ namespace mlir Value dynamic_init) { [[maybe_unused]] auto lowerBound = builder.create(loc, 0); - MemRefType resultMemTy = alloc_op.getDefiningOp()->getResult(0).getType().cast(); + MemRefType resultMemTy = cast(alloc_op.getDefiningOp()->getResult(0).getType()); std::vector cur_indices; std::vector cur_memref; @@ -330,12 +321,12 @@ namespace mlir /// Find summation indices for (const auto &map : indexMaps) { - auto affineMap = map.cast().getValue(); + auto affineMap = cast(map).getValue(); std::vector perm; for (size_t i = 0; i < affineMap.getNumResults(); i++) { auto expr = affineMap.getResult(i); - perm.push_back(llvm::cast(expr).getPosition()); + perm.push_back(cast(expr).getPosition()); } allPerms.push_back(perm); @@ -354,13 +345,13 @@ namespace mlir comet_vdump(map); std::vector perm; - if (auto arrayattr = map.dyn_cast()) + if (auto arrayattr = dyn_cast(map)) { comet_debug() << " "; for (auto n : arrayattr) { comet_debug() << " "; - if (IntegerAttr i = n.dyn_cast()) + if (IntegerAttr i = dyn_cast(n)) { comet_debug() << " " << i.getInt() << "\n"; perm.push_back(i.getInt()); @@ -500,14 +491,14 @@ namespace mlir } /// for string delimiter - std::vector stringSplit(std::string s, std::string delimiter) + std::vector stringSplit(llvm::StringRef s, llvm::StringRef delimiter) { /// comet_debug() << "split formats string: " << s << ", deli: "<< delimiter << ".\n"; - std::vector res; + std::vector res; std::string format = ""; - for (unsigned int i = 0; i < s.length(); i++) + for (unsigned int i = 0; i < s.size(); i++) { comet_debug() << "s[" << i << "]: " << s[i] << "\n"; if (s[i] != delimiter[0] && s[i] != delimiter[1]) @@ -526,17 +517,17 @@ namespace mlir res.push_back(format); comet_debug() << "The final format: "; - print_vector(res); + print_vector(res); return res; } - std::vector> getAllFormats(ArrayAttr opFormatsArrayAttr, std::vector> allPerms) + std::vector> getAllFormats(ArrayAttr opFormatsArrayAttr, std::vector> allPerms) { - std::vector> allFormats(allPerms.size()); + std::vector> allFormats(allPerms.size()); /// format with each input matrix: ["CSR", "D", "D"] SpMM for (unsigned int i = 0; i < opFormatsArrayAttr.size(); i++) { - std::string formats_str(opFormatsArrayAttr[i].cast().getValue()); + std::string formats_str(cast(opFormatsArrayAttr[i]).getValue()); unsigned int tensorDims = allPerms[i].size(); comet_debug() << "format_str: " << formats_str << ", tensorDims: " << tensorDims << "\n"; @@ -635,16 +626,16 @@ namespace mlir { comet_debug() << " "; /// std::string formats_str; - if (opFormatsArrayAttr[i].dyn_cast()) + if (dyn_cast(opFormatsArrayAttr[i])) { comet_debug() << " yes "; } - else if (mlir::ArrayAttr formatArrayAttr = opFormatsArrayAttr[i].dyn_cast()) + else if (mlir::ArrayAttr formatArrayAttr = dyn_cast(opFormatsArrayAttr[i])) { comet_debug() << " yes " << formatArrayAttr.size() << " "; for (unsigned long j = 0; j < formatArrayAttr.size(); j++) { - if (mlir::StringAttr format = formatArrayAttr[j].dyn_cast()) + if (mlir::StringAttr format = dyn_cast(formatArrayAttr[j])) { std::string formats_str(format.getValue()); comet_debug() << " " << formats_str << " "; @@ -831,7 +822,7 @@ namespace mlir return false; } - std::vector getFormatsValue(std::string formats_str, int rank_size, PatternRewriter &rewriter, Location loc, IndexType indexType) + std::vector getFormatsValue(llvm::StringRef formats_str, int rank_size, PatternRewriter &rewriter, Location loc, IndexType indexType) { Value format_unk = rewriter.create(loc, indexType, rewriter.getIndexAttr(-1)); Value format_dense = rewriter.create(loc, indexType, rewriter.getIndexAttr(0)); @@ -846,42 +837,42 @@ namespace mlir { /// 2D comet_debug() << " 2D\n"; /// Value dim0_format, dim1_format; - if (formats_str.compare(0, 3, "CSR") == 0) + if (formats_str.compare("CSR") == 0) { dim_format.push_back(format_dense); dim_format.push_back(format_unk); dim_format.push_back(format_compressed); dim_format.push_back(format_unk); } - else if (formats_str.compare(0, 4, "DCSR") == 0) + else if (formats_str.compare("DCSR") == 0) { dim_format.push_back(format_compressed); dim_format.push_back(format_unk); dim_format.push_back(format_compressed); dim_format.push_back(format_unk); } - else if (formats_str.compare(0, 3, "COO") == 0) + else if (formats_str.compare("COO") == 0) { /// COO dim_format.push_back(format_compressednonunique); dim_format.push_back(format_unk); dim_format.push_back(format_singleton); dim_format.push_back(format_unk); } - else if (formats_str.compare(0, 3, "ELL") == 0) + else if (formats_str.compare("ELL") == 0) { /// ELL dim_format.push_back(format_dense); dim_format.push_back(format_dense); dim_format.push_back(format_singleton); dim_format.push_back(format_unk); } - else if (formats_str.compare(0, 4, "BCSR") == 0) + else if (formats_str.compare("BCSR") == 0) { /// BCSR dim_format.push_back(format_dense); dim_format.push_back(format_compressednonunique); dim_format.push_back(format_dense); dim_format.push_back(format_dense); } - else if (formats_str.compare(0, 3, "CSB") == 0) + else if (formats_str.compare("CSB") == 0) { /// CSB dim_format.push_back(format_dense); dim_format.push_back(format_dense); @@ -890,22 +881,22 @@ namespace mlir } else if (formats_str.find("D") != std::string::npos || formats_str.find("CU") != std::string::npos || formats_str.find("CN") != std::string::npos || formats_str.find("S") != std::string::npos) { - std::vector format_vec = stringSplit(formats_str, ", "); + std::vector format_vec = stringSplit(formats_str, ", "); for (auto n : format_vec) { - if (n.compare(0, 1, "D") == 0) + if (n.compare("D") == 0) { dim_format.push_back(format_dense); } - else if (n.compare(0, 2, "CU") == 0) + else if (n.compare("CU") == 0) { dim_format.push_back(format_compressed); } - else if (n.compare(0, 2, "CN") == 0) + else if (n.compare("CN") == 0) { dim_format.push_back(format_compressednonunique); } - else if (n.compare(0, 1, "S") == 0) + else if (n.compare("S") == 0) { dim_format.push_back(format_singleton); } @@ -920,7 +911,7 @@ namespace mlir { /// 3D comet_debug() << " 3D\n"; /// Value dim0_format, dim1_format, dim2_format; - if (formats_str.compare(0, 3, "CSF") == 0) + if (formats_str.compare("CSF") == 0) { dim_format.push_back(format_compressed); dim_format.push_back(format_unk); @@ -929,7 +920,7 @@ namespace mlir dim_format.push_back(format_compressed); dim_format.push_back(format_unk); } - else if (formats_str.compare(0, 11, "ModeGeneric") == 0) + else if (formats_str.compare("ModeGeneric") == 0) { dim_format.push_back(format_compressednonunique); dim_format.push_back(format_unk); @@ -938,7 +929,7 @@ namespace mlir dim_format.push_back(format_dense); dim_format.push_back(format_unk); } - else if (formats_str.compare(0, 3, "COO") == 0) + else if (formats_str.compare("COO") == 0) { /// COO dim_format.push_back(format_compressednonunique); dim_format.push_back(format_unk); @@ -949,25 +940,25 @@ namespace mlir } else if (formats_str.find("D") != std::string::npos || formats_str.find("CU") != std::string::npos || formats_str.find("CN") != std::string::npos || formats_str.find("S") != std::string::npos) { - std::vector format_vec = stringSplit(formats_str, ", "); + std::vector format_vec = stringSplit(formats_str, ", "); comet_debug() << " format_vec.size(): " << format_vec.size() << " \n"; for (auto n : format_vec) { comet_debug() << "Current format attribute: " << n << "---\n"; - if (n.compare(0, 1, "D") == 0) + if (n.compare("D") == 0) { dim_format.push_back(format_dense); } - else if (n.compare(0, 2, "CU") == 0) + else if (n.compare("CU") == 0) { dim_format.push_back(format_compressed); } - else if (n.compare(0, 2, "CN") == 0) + else if (n.compare("CN") == 0) { dim_format.push_back(format_compressednonunique); } - else if (n.compare(0, 1, "S") == 0) + else if (n.compare("S") == 0) { dim_format.push_back(format_singleton); } @@ -994,7 +985,78 @@ namespace mlir return dim_format; } - std::vector getFormatsValueInt(std::string formats_str, int rank_size, PatternRewriter &rewriter, Location loc, IntegerType intType) + std::string getTensorFormatString(Type tensorT) + { + if(isa(tensorT)) + { + return "Dense"; + } + + auto sparseTensorT = cast(tensorT); + auto formats = sparseTensorT.getFormat(); + if(sparseTensorT.getRank() == 2) + { + if(formats[0] == TensorFormatEnum::CN && formats[1] == TensorFormatEnum::UNK && formats[2] == TensorFormatEnum::S && formats[3] == TensorFormatEnum::UNK) + { + return "COO"; + } + else if(formats[0] == TensorFormatEnum::D && formats[1] == TensorFormatEnum::UNK && formats[2] == TensorFormatEnum::CU && formats[3] == TensorFormatEnum::UNK) + { + return "CSR"; /// CSR format + } + else if(formats[0] == TensorFormatEnum::CU && formats[1] == TensorFormatEnum::UNK && formats[2] == TensorFormatEnum::CU && formats[3] == TensorFormatEnum::UNK) + { + return "DCSR"; /// DCSR format + } + else if(formats[0] == TensorFormatEnum::D && formats[1] == TensorFormatEnum::D && formats[2] == TensorFormatEnum::S && formats[3] == TensorFormatEnum::UNK) + { + return "ELL"; /// ELL format + } + else if(formats[0] == TensorFormatEnum::D && formats[1] == TensorFormatEnum::CN && formats[2] == TensorFormatEnum::D && formats[3] == TensorFormatEnum::D) + { + return "BCSR"; /// BCSR format + } + else if(formats[0] == TensorFormatEnum::D && formats[1] == TensorFormatEnum::D && formats[2] == TensorFormatEnum::CU && formats[3] == TensorFormatEnum::S) + { + return "CSB"; /// CSB format + } + else + { + //Handle errors for unsupported formats + llvm::errs() << "Unsupported 2D sparse tensor format detected: " + << formats[0] << ", " + << formats[1] << ", " + << formats[2] << ", " + << formats[3] << "\n"; + return ""; // Return empty string to indicate unsupported format + } + } + else if(sparseTensorT.getRank() == 3) + { + if(formats[0] == TensorFormatEnum::CU && formats[1] == TensorFormatEnum::UNK && formats[2] == TensorFormatEnum::CU && formats[3] == TensorFormatEnum::UNK && formats[4] == TensorFormatEnum::CU && formats[5] == TensorFormatEnum::UNK) + { + return "CSF"; + } + else if(formats[0] == TensorFormatEnum::CN && formats[1] == TensorFormatEnum::UNK && formats[2] == TensorFormatEnum::S && formats[3] == TensorFormatEnum::UNK && formats[4] == TensorFormatEnum::D && formats[5] == TensorFormatEnum::UNK) + { + return "ModeGeneric"; + } + else if(formats[0] == TensorFormatEnum::CN && formats[1] == TensorFormatEnum::UNK && formats[2] == TensorFormatEnum::S && formats[3] == TensorFormatEnum::UNK && formats[4] == TensorFormatEnum::S && formats[5] == TensorFormatEnum::UNK) + { + return "COO"; /// COO format for 3D + } + else + { + llvm::errs() << "Unsupported 3D sparse tensor format detected: " << formats[0] << ", " << formats[1] << ", " << formats[2] << "\n"; + } + } + + // Handle cases for unsupported tensor ranks/formats + llvm::errs() << "Unsupported sparse tensor format detected for rank: " << sparseTensorT.getRank() << "\n"; + return ""; // Return empty string to indicate unsupported format + } + + std::vector getFormatsValueInt(llvm::StringRef formats_str, int rank_size, PatternRewriter &rewriter, Location loc, IntegerType intType) { Value format_unk = rewriter.create(loc, intType, rewriter.getIntegerAttr(intType, -1)); Value format_dense = rewriter.create(loc, intType, rewriter.getIntegerAttr(intType, 0)); @@ -1009,42 +1071,42 @@ namespace mlir { /// 2D comet_debug() << " 2D\n"; /// Value dim0_format, dim1_format; - if (formats_str.compare(0, 3, "CSR") == 0) + if (formats_str.compare("CSR") == 0) { dim_format.push_back(format_dense); dim_format.push_back(format_unk); dim_format.push_back(format_compressed); dim_format.push_back(format_unk); } - else if (formats_str.compare(0, 4, "DCSR") == 0) + else if (formats_str.compare("DCSR") == 0) { dim_format.push_back(format_compressed); dim_format.push_back(format_unk); dim_format.push_back(format_compressed); dim_format.push_back(format_unk); } - else if (formats_str.compare(0, 3, "COO") == 0) + else if (formats_str.compare("COO") == 0) { /// COO dim_format.push_back(format_compressednonunique); dim_format.push_back(format_unk); dim_format.push_back(format_singleton); dim_format.push_back(format_unk); } - else if (formats_str.compare(0, 3, "ELL") == 0) + else if (formats_str.compare("ELL") == 0) { /// ELL dim_format.push_back(format_dense); dim_format.push_back(format_dense); dim_format.push_back(format_singleton); dim_format.push_back(format_unk); } - else if (formats_str.compare(0, 4, "BCSR") == 0) + else if (formats_str.compare("BCSR") == 0) { /// BCSR dim_format.push_back(format_dense); dim_format.push_back(format_compressed); dim_format.push_back(format_dense); dim_format.push_back(format_dense); } - else if (formats_str.compare(0, 3, "CSB") == 0) + else if (formats_str.compare("CSB") == 0) { /// CSB dim_format.push_back(format_dense); dim_format.push_back(format_dense); @@ -1053,22 +1115,22 @@ namespace mlir } else if (formats_str.find("D") != std::string::npos || formats_str.find("CU") != std::string::npos || formats_str.find("CN") != std::string::npos || formats_str.find("S") != std::string::npos) { - std::vector format_vec = stringSplit(formats_str, ", "); + std::vector format_vec = stringSplit(formats_str, ", "); for (auto n : format_vec) { - if (n.compare(0, 1, "D") == 0) + if (n.compare("D") == 0) { dim_format.push_back(format_dense); } - else if (n.compare(0, 2, "CU") == 0) + else if (n.compare("CU") == 0) { dim_format.push_back(format_compressed); } - else if (n.compare(0, 2, "CN") == 0) + else if (n.compare("CN") == 0) { dim_format.push_back(format_compressednonunique); } - else if (n.compare(0, 1, "S") == 0) + else if (n.compare("S") == 0) { dim_format.push_back(format_singleton); } @@ -1083,7 +1145,7 @@ namespace mlir { /// 3D comet_debug() << " 3D\n"; /// Value dim0_format, dim1_format, dim2_format; - if (formats_str.compare(0, 3, "CSF") == 0) + if (formats_str.compare("CSF") == 0) { dim_format.push_back(format_compressed); dim_format.push_back(format_unk); @@ -1092,7 +1154,7 @@ namespace mlir dim_format.push_back(format_compressed); dim_format.push_back(format_unk); } - else if (formats_str.compare(0, 11, "ModeGeneric") == 0) + else if (formats_str.compare("ModeGeneric") == 0) { dim_format.push_back(format_compressednonunique); dim_format.push_back(format_unk); @@ -1101,7 +1163,7 @@ namespace mlir dim_format.push_back(format_dense); dim_format.push_back(format_unk); } - else if (formats_str.compare(0, 3, "COO") == 0) + else if (formats_str.compare("COO") == 0) { /// COO dim_format.push_back(format_compressednonunique); dim_format.push_back(format_unk); @@ -1112,25 +1174,190 @@ namespace mlir } else if (formats_str.find("D") != std::string::npos || formats_str.find("CU") != std::string::npos || formats_str.find("CN") != std::string::npos || formats_str.find("S") != std::string::npos) { - std::vector format_vec = stringSplit(formats_str, ", "); + std::vector format_vec = stringSplit(formats_str, ", "); comet_debug() << " format_vec.size(): " << format_vec.size() << " \n"; /// print_vector(format_vec); for (auto n : format_vec) { comet_debug() << "Current format attribute: " << n << "---\n"; - if (n.compare(0, 1, "D") == 0) + if (n.compare("D") == 0) { dim_format.push_back(format_dense); } - else if (n.compare(0, 2, "CU") == 0) + else if (n.compare("CU") == 0) { dim_format.push_back(format_compressed); } - else if (n.compare(0, 2, "CN") == 0) + else if (n.compare("CN") == 0) { dim_format.push_back(format_compressednonunique); } - else if (n.compare(0, 1, "S") == 0) + else if (n.compare("S") == 0) + { + dim_format.push_back(format_singleton); + } + else + { + llvm::errs() << "Uncorrect format attribute: " << n << "---\n"; + } + comet_debug() << " dim_format.size(): " << dim_format.size() << " \n"; + } + comet_debug() << " formats_str: " << formats_str << ", dim_format.size(): " << dim_format.size() << " \n"; + } + } + else + { + llvm::errs() << "Unsupported formats: " << formats_str << " (tensor dimes: " << rank_size << ") \n"; + } + + comet_debug() << " print dim_format: "; + for (auto n : dim_format) + { + comet_debug() << n << " "; + } + comet_debug() << "\n"; + return dim_format; + } + + // TODO (alokvk2): Not good to have this replicated 3 times. Ideally this is only used for "special" formats (i.e. CSR, COO etc.) + // And this converts it to a vector of TAFormatAttrs. + std::vector getFormats(llvm::StringRef formats_str, int rank_size, MLIRContext* ctx) + { + auto format_unk = TensorFormatEnum::UNK; + auto format_dense = TensorFormatEnum::D; + auto format_compressed = TensorFormatEnum::CU; + auto format_compressednonunique = TensorFormatEnum::CN; + auto format_singleton = TensorFormatEnum::S; + /// read_input_sizes_2D_f64 or read_input_sizes_3D_f64 + comet_debug() << "\n"; + std::vector dim_format; + + if (rank_size == 2) + { /// 2D + comet_debug() << " 2D\n"; + /// Value dim0_format, dim1_format; + if (formats_str.compare("CSR") == 0) + { + dim_format.push_back(format_dense); + dim_format.push_back(format_unk); + dim_format.push_back(format_compressed); + dim_format.push_back(format_unk); + } + else if (formats_str.compare("DCSR") == 0) + { + dim_format.push_back(format_compressed); + dim_format.push_back(format_unk); + dim_format.push_back(format_compressed); + dim_format.push_back(format_unk); + } + else if (formats_str.compare("COO") == 0) + { /// COO + dim_format.push_back(format_compressednonunique); + dim_format.push_back(format_unk); + dim_format.push_back(format_singleton); + dim_format.push_back(format_unk); + } + else if (formats_str.compare("ELL") == 0) + { /// ELL + dim_format.push_back(format_dense); + dim_format.push_back(format_dense); + dim_format.push_back(format_singleton); + dim_format.push_back(format_unk); + } + else if (formats_str.compare("BCSR") == 0) + { /// BCSR + dim_format.push_back(format_dense); + dim_format.push_back(format_compressed); + dim_format.push_back(format_dense); + dim_format.push_back(format_dense); + } + else if (formats_str.compare("CSB") == 0) + { /// CSB + dim_format.push_back(format_dense); + dim_format.push_back(format_dense); + dim_format.push_back(format_compressed); + dim_format.push_back(format_singleton); + } + else if (formats_str.find("D") != std::string::npos || formats_str.find("CU") != std::string::npos || formats_str.find("CN") != std::string::npos || formats_str.find("S") != std::string::npos) + { + std::vector format_vec = stringSplit(formats_str, ", "); + for (auto n : format_vec) + { + if (n.compare("D") == 0) + { + dim_format.push_back(format_dense); + } + else if (n.compare("CU") == 0) + { + dim_format.push_back(format_compressed); + } + else if (n.compare("CN") == 0) + { + dim_format.push_back(format_compressednonunique); + } + else if (n.compare("S") == 0) + { + dim_format.push_back(format_singleton); + } + } + } + else + { + llvm::errs() << "Unsupported formats: " << formats_str << " (tensor dimes: " << rank_size << ") \n"; + } + } + else if (rank_size == 3) + { /// 3D + comet_debug() << " 3D\n"; + /// Value dim0_format, dim1_format, dim2_format; + if (formats_str.compare("CSF") == 0) + { + dim_format.push_back(format_compressed); + dim_format.push_back(format_unk); + dim_format.push_back(format_compressed); + dim_format.push_back(format_unk); + dim_format.push_back(format_compressed); + dim_format.push_back(format_unk); + } + else if (formats_str.compare("ModeGeneric") == 0) + { + dim_format.push_back(format_compressednonunique); + dim_format.push_back(format_unk); + dim_format.push_back(format_singleton); + dim_format.push_back(format_unk); + dim_format.push_back(format_dense); + dim_format.push_back(format_unk); + } + else if (formats_str.compare("COO") == 0) + { /// COO + dim_format.push_back(format_compressednonunique); + dim_format.push_back(format_unk); + dim_format.push_back(format_singleton); + dim_format.push_back(format_unk); + dim_format.push_back(format_singleton); + dim_format.push_back(format_unk); + } + else if (formats_str.find("D") != std::string::npos || formats_str.find("CU") != std::string::npos || formats_str.find("CN") != std::string::npos || formats_str.find("S") != std::string::npos) + { + std::vector format_vec = stringSplit(formats_str, ", "); + comet_debug() << " format_vec.size(): " << format_vec.size() << " \n"; + /// print_vector(format_vec); + for (auto n : format_vec) + { + comet_debug() << "Current format attribute: " << n << "---\n"; + if (n.compare("D") == 0) + { + dim_format.push_back(format_dense); + } + else if (n.compare("CU") == 0) + { + dim_format.push_back(format_compressed); + } + else if (n.compare("CN") == 0) + { + dim_format.push_back(format_compressednonunique); + } + else if (n.compare("S") == 0) { dim_format.push_back(format_singleton); } @@ -1193,33 +1420,34 @@ namespace mlir /// parent node can get from getUser() function, only one user since tree structure void dfsRootOpTree(Value tcRootOp, std::vector &ret) { - if (isa(tcRootOp.getDefiningOp())) - { - IndexTreeIndicesOp workspaceop = dyn_cast(tcRootOp.getDefiningOp()); - - comet_debug() << " dfsRootOpTree\n"; - comet_vdump(workspaceop); - - unsigned int sz = workspaceop.getChildren().size(); - - comet_debug() << " " << sz << " "; - ret.push_back(workspaceop); - comet_debug() << " "; - comet_vdump(workspaceop); - - for (unsigned int i = 0; i < sz; i++) - { - Value t = workspaceop.getChildren()[i]; - dfsRootOpTree(t, ret); - } - } - else if (isa(tcRootOp.getDefiningOp())) - { - indexTree::IndexTreeComputeOp leafop = dyn_cast(tcRootOp.getDefiningOp()); - /// comet_debug() << " dfsRootOpTree\n"; - comet_vdump(leafop); - ret.push_back(leafop); - } + return; + // if (isa(tcRootOp.getDefiningOp())) + // { + // IndexTreeIndicesOp workspaceop = dyn_cast(tcRootOp.getDefiningOp()); + + // comet_debug() << " dfsRootOpTree\n"; + // comet_vdump(workspaceop); + + // unsigned int sz = workspaceop.getChildren().size(); + + // comet_debug() << " " << sz << " "; + // ret.push_back(workspaceop); + // comet_debug() << " "; + // comet_vdump(workspaceop); + + // for (unsigned int i = 0; i < sz; i++) + // { + // Value t = workspaceop.getChildren()[i]; + // dfsRootOpTree(t, ret); + // } + // } + // else if (isa(tcRootOp.getDefiningOp())) + // { + // indexTree::IndexTreeComputeOp leafop = dyn_cast(tcRootOp.getDefiningOp()); + // /// comet_debug() << " dfsRootOpTree\n"; + // comet_vdump(leafop); + // ret.push_back(leafop); + // } } void getAncestorsWp(Value op, std::vector &ret /* output ancestors*/, std::vector &dfsOps) @@ -1301,7 +1529,7 @@ namespace mlir std::vector> mapping; for (unsigned int m = 0; m < perms.size(); m++) { - ArrayAttr aa = perms[m].dyn_cast(); + ArrayAttr aa = dyn_cast(perms[m]); std::vector p; for (unsigned int n = 0; n < aa.size(); n++) { @@ -1318,12 +1546,12 @@ namespace mlir std::vector> perms_int; for (unsigned int m = 0; m < perms.size(); m++) { - ArrayAttr aa = perms[m].dyn_cast(); + ArrayAttr aa = dyn_cast(perms[m]); std::vector p; for (unsigned int n = 0; n < aa.size(); n++) { - p.push_back(aa[n].cast().getInt()); - comet_debug() << " convertArrayAttrIntTo2DVector:" << aa[n].cast().getInt() << "\n"; + p.push_back(cast(aa[n]).getInt()); + comet_debug() << " convertArrayAttrIntTo2DVector:" << cast(aa[n]).getInt() << "\n"; } perms_int.push_back(p); } @@ -1367,11 +1595,11 @@ namespace mlir std::vector> formats_str; for (unsigned int m = 0; m < formats.size(); m++) { - ArrayAttr aa = formats[m].dyn_cast(); + ArrayAttr aa = dyn_cast(formats[m]); std::vector p; for (unsigned int n = 0; n < aa.size(); n++) { - std::string format_str(aa[n].cast().getValue()); + std::string format_str(cast(aa[n]).getValue()); p.push_back(format_str); comet_debug() << " convertArrayAttrStrTo2DVector:" << format_str << "\n"; } @@ -1385,99 +1613,105 @@ namespace mlir std::vector> &opPerms, std::vector> &inputOutputMapping) { - indexTree::IndexTreeComputeRHSOp itComputeOp_rhs = dyn_cast(computeOp.getDefiningOp()->getOperand(0).getDefiningOp()); - ArrayAttr opFormatsArrayAttr_rhs = itComputeOp_rhs.getAllFormats(); - ArrayAttr opPermsArrayAttr_rhs = itComputeOp_rhs.getAllPerms(); - indexTree::IndexTreeComputeLHSOp itComputeOp_lhs = dyn_cast(computeOp.getDefiningOp()->getOperand(1).getDefiningOp()); - ArrayAttr opFormatsArrayAttr_lhs = itComputeOp_lhs.getAllFormats(); - ArrayAttr opPermsArrayAttr_lhs = itComputeOp_lhs.getAllPerms(); - assert(opFormatsArrayAttr_rhs.size() == opPermsArrayAttr_rhs.size() && "not equal RHS formats size with perms size\n"); - assert(opFormatsArrayAttr_lhs.size() == opPermsArrayAttr_lhs.size() && "not equal LHS formats size with perms size\n"); - - /// Get output format, vector of vector - /// Convert ArrayAttr into - comet_debug() << "Start printing opFormats_rhs\n"; - std::vector> opFormats_rhs = convertArrayAttrStrTo2DVector(opFormatsArrayAttr_rhs); - comet_debug() << "End printing opFormats_rhs\n"; - std::vector> opPerms_rhs = convertArrayAttrIntTo2DVector(opPermsArrayAttr_rhs); - std::vector> inputMapping = createInputOutputMapping(opPermsArrayAttr_rhs, true); - - comet_debug() << "Start printing opFormats_lhs\n"; - std::vector> opFormats_lhs = convertArrayAttrStrTo2DVector(opFormatsArrayAttr_lhs); - comet_debug() << "End printing opFormats_lhs\n"; - std::vector> opPerms_lhs = convertArrayAttrIntTo2DVector(opPermsArrayAttr_lhs); - std::vector> outputMapping = createInputOutputMapping(opPermsArrayAttr_lhs, false); - - opFormats = opFormats_rhs; - opFormats.insert(opFormats.end(), opFormats_lhs.begin(), opFormats_lhs.end()); - opPerms = opPerms_rhs; - opPerms.insert(opPerms.end(), opPerms_lhs.begin(), opPerms_lhs.end()); - inputOutputMapping = inputMapping; - inputOutputMapping.insert(inputOutputMapping.end(), outputMapping.begin(), outputMapping.end()); + return; + // indexTree::IndexTreeComputeRHSOp itComputeOp_rhs = dyn_cast(computeOp.getDefiningOp()->getOperand(0).getDefiningOp()); + // ArrayAttr opFormatsArrayAttr_rhs = itComputeOp_rhs.getAllFormats(); + // ArrayAttr opPermsArrayAttr_rhs = itComputeOp_rhs.getAllPerms(); + // indexTree::IndexTreeComputeLHSOp itComputeOp_lhs = dyn_cast(computeOp.getDefiningOp()->getOperand(1).getDefiningOp()); + // ArrayAttr opFormatsArrayAttr_lhs = itComputeOp_lhs.getAllFormats(); + // ArrayAttr opPermsArrayAttr_lhs = itComputeOp_lhs.getAllPerms(); + // assert(opFormatsArrayAttr_rhs.size() == opPermsArrayAttr_rhs.size() && "not equal RHS formats size with perms size\n"); + // assert(opFormatsArrayAttr_lhs.size() == opPermsArrayAttr_lhs.size() && "not equal LHS formats size with perms size\n"); + + // /// Get output format, vector of vector + // /// Convert ArrayAttr into + // comet_debug() << "Start printing opFormats_rhs\n"; + // std::vector> opFormats_rhs = convertArrayAttrStrTo2DVector(opFormatsArrayAttr_rhs); + // comet_debug() << "End printing opFormats_rhs\n"; + // std::vector> opPerms_rhs = convertArrayAttrIntTo2DVector(opPermsArrayAttr_rhs); + // std::vector> inputMapping = createInputOutputMapping(opPermsArrayAttr_rhs, true); + + // comet_debug() << "Start printing opFormats_lhs\n"; + // std::vector> opFormats_lhs = convertArrayAttrStrTo2DVector(opFormatsArrayAttr_lhs); + // comet_debug() << "End printing opFormats_lhs\n"; + // std::vector> opPerms_lhs = convertArrayAttrIntTo2DVector(opPermsArrayAttr_lhs); + // std::vector> outputMapping = createInputOutputMapping(opPermsArrayAttr_lhs, false); + + // opFormats = opFormats_rhs; + // opFormats.insert(opFormats.end(), opFormats_lhs.begin(), opFormats_lhs.end()); + // opPerms = opPerms_rhs; + // opPerms.insert(opPerms.end(), opPerms_lhs.begin(), opPerms_lhs.end()); + // inputOutputMapping = inputMapping; + // inputOutputMapping.insert(inputOutputMapping.end(), outputMapping.begin(), outputMapping.end()); } /// Get the formats of the itCompute op void getFormatsOfComputeOp(Value computeOp, std::vector> &opFormats) { - indexTree::IndexTreeComputeRHSOp itComputeOp_rhs = dyn_cast(computeOp.getDefiningOp()->getOperand(0).getDefiningOp()); - ArrayAttr opFormatsArrayAttr_rhs = itComputeOp_rhs.getAllFormats(); - indexTree::IndexTreeComputeLHSOp itComputeOp_lhs = dyn_cast(computeOp.getDefiningOp()->getOperand(1).getDefiningOp()); - ArrayAttr opFormatsArrayAttr_lhs = itComputeOp_lhs.getAllFormats(); - - /// Get output format, vector of vector - /// Convert ArrayAttr into - std::vector> opFormats_rhs = convertArrayAttrStrTo2DVector(opFormatsArrayAttr_rhs); - std::vector> opFormats_lhs = convertArrayAttrStrTo2DVector(opFormatsArrayAttr_lhs); - - opFormats = opFormats_rhs; - opFormats.insert(opFormats.end(), opFormats_lhs.begin(), opFormats_lhs.end()); + return; + // indexTree::IndexTreeComputeRHSOp itComputeOp_rhs = dyn_cast(computeOp.getDefiningOp()->getOperand(0).getDefiningOp()); + // ArrayAttr opFormatsArrayAttr_rhs = itComputeOp_rhs.getAllFormats(); + // indexTree::IndexTreeComputeLHSOp itComputeOp_lhs = dyn_cast(computeOp.getDefiningOp()->getOperand(1).getDefiningOp()); + // ArrayAttr opFormatsArrayAttr_lhs = itComputeOp_lhs.getAllFormats(); + + // /// Get output format, vector of vector + // /// Convert ArrayAttr into + // std::vector> opFormats_rhs = convertArrayAttrStrTo2DVector(opFormatsArrayAttr_rhs); + // std::vector> opFormats_lhs = convertArrayAttrStrTo2DVector(opFormatsArrayAttr_lhs); + + // opFormats = opFormats_rhs; + // opFormats.insert(opFormats.end(), opFormats_lhs.begin(), opFormats_lhs.end()); } /// Get the rhs formats of the itCompute op void getRHSFormatsOfComputeOp(Value computeOp, std::vector> &opFormats) { - indexTree::IndexTreeComputeRHSOp itComputeOp_rhs = dyn_cast(computeOp.getDefiningOp()->getOperand(0).getDefiningOp()); - ArrayAttr opFormatsArrayAttr_rhs = itComputeOp_rhs.getAllFormats(); + return; + // indexTree::IndexTreeComputeRHSOp itComputeOp_rhs = dyn_cast(computeOp.getDefiningOp()->getOperand(0).getDefiningOp()); + // ArrayAttr opFormatsArrayAttr_rhs = itComputeOp_rhs.getAllFormats(); - /// Get output format, vector of vector - /// Convert ArrayAttr into - std::vector> opFormats_rhs = convertArrayAttrStrTo2DVector(opFormatsArrayAttr_rhs); + // /// Get output format, vector of vector + // /// Convert ArrayAttr into + // std::vector> opFormats_rhs = convertArrayAttrStrTo2DVector(opFormatsArrayAttr_rhs); - opFormats = opFormats_rhs; + // opFormats = opFormats_rhs; } /// Get the LHS formats of the itCompute op void getLHSFormatsOfComputeOp(Value computeOp, std::vector> &opFormats) { - indexTree::IndexTreeComputeLHSOp itComputeOp_lhs = dyn_cast(computeOp.getDefiningOp()->getOperand(1).getDefiningOp()); - ArrayAttr opFormatsArrayAttr_lhs = itComputeOp_lhs.getAllFormats(); - std::vector> opFormats_lhs = convertArrayAttrStrTo2DVector(opFormatsArrayAttr_lhs); - opFormats = opFormats_lhs; + return; + // indexTree::IndexTreeComputeLHSOp itComputeOp_lhs = dyn_cast(computeOp.getDefiningOp()->getOperand(1).getDefiningOp()); + // ArrayAttr opFormatsArrayAttr_lhs = itComputeOp_lhs.getAllFormats(); + // std::vector> opFormats_lhs = convertArrayAttrStrTo2DVector(opFormatsArrayAttr_lhs); + // opFormats = opFormats_lhs; } /// Get the input tensors of the itCompute op void getInputTensorsOfComputeOp(Value computeOp, std::vector &inputTensors /* output */) { + return; /// indexTree::IndexTreeComputeOp itComputeOp = dyn_cast(computeOp.getDefiningOp()); - indexTree::IndexTreeComputeRHSOp itComputeOp_rhs = dyn_cast(computeOp.getDefiningOp()->getOperand(0).getDefiningOp()); - comet_debug() << " "; - comet_vdump(itComputeOp_rhs); - for (unsigned int i = 0; i < itComputeOp_rhs.getOperation()->getNumOperands(); i++) - { - comet_debug() << " "; - comet_vdump(itComputeOp_rhs.getOperation()->getOperand(i)); - inputTensors.push_back(itComputeOp_rhs.getOperation()->getOperand(i)); - } + // indexTree::IndexTreeComputeRHSOp itComputeOp_rhs = dyn_cast(computeOp.getDefiningOp()->getOperand(0).getDefiningOp()); + // comet_debug() << " "; + // comet_vdump(itComputeOp_rhs); + // for (unsigned int i = 0; i < itComputeOp_rhs.getOperation()->getNumOperands(); i++) + // { + // comet_debug() << " "; + // comet_vdump(itComputeOp_rhs.getOperation()->getOperand(i)); + // inputTensors.push_back(itComputeOp_rhs.getOperation()->getOperand(i)); + // } } /// Get the output tensors of the itCompute op void getOutputTensorsOfComputeOp(Value computeOp, std::vector &outputTensors /* output */) { - indexTree::IndexTreeComputeLHSOp itComputeOp_lhs = dyn_cast(computeOp.getDefiningOp()->getOperand(1).getDefiningOp()); - for (unsigned int i = 0; i < itComputeOp_lhs.getOperation()->getNumOperands(); i++) - { - outputTensors.push_back(itComputeOp_lhs.getOperation()->getOperand(i)); - } + return; + // indexTree::IndexTreeComputeLHSOp itComputeOp_lhs = dyn_cast(computeOp.getDefiningOp()->getOperand(1).getDefiningOp()); + // for (unsigned int i = 0; i < itComputeOp_lhs.getOperation()->getNumOperands(); i++) + // { + // outputTensors.push_back(itComputeOp_lhs.getOperation()->getOperand(i)); + // } } /// Get indices in current WorkspaceOp cur_op @@ -1488,186 +1722,187 @@ namespace mlir std::vector &ids /* output */, std::vector &formats /* output */) { - /// For each indices, find in each leaf, which tensor, the corresponding format - /// If in all tensors, the formats of the index are D, then D - /// If only one Sparse, then sparse - comet_debug() << " getFormatsInfo:Start Current op\n"; - comet_vdump(cur_op); - comet_debug() << " getFormatsInfo:indices.size(): " << indices.size() << "\n"; - for (unsigned long i = 0; i < indices.size(); i++) - { - comet_debug() << " getFormatsInfo:indices[" << i << "]: " << indices[i] << "\n"; - /// Info for each index - std::string format; - Value tensor; - unsigned int id; - bool isSet = false; - - std::vector formats_leafs; - std::vector tensors_leafs; - std::vector ids_leafs; - - for (unsigned long j = 0; j < leafs.size(); j++) - { - /// Info for each index in leaf[j] - comet_debug() << " getFormatsInfo:LeafOp: "; - comet_vdump(leafs[j]); - std::string format_in_leaf; - Value tensor_in_leaf; - unsigned int id_in_leaf; - bool isSetInLeaf = false; - - /// get All perms and formats info - if (indexTree::IndexTreeComputeOp leafop = dyn_cast(leafs[j].getDefiningOp())) - { - comet_debug() << " getFormatsInfo:leafs[" << j << "] is computeOp\n"; - std::vector> allFormats; - std::vector> allPerms; - std::vector> inputOutputMapping; - OpBuilder builder(leafop); - getFormatsPermsOfComputeOp(leafop, allFormats, allPerms, inputOutputMapping); - - comet_debug() << " getFormatsInfo:Allformats allFormats.size(): " << allFormats.size() << "\n"; - for (auto m : allFormats) - { - comet_debug() << " "; - for (auto n : m) - { - comet_debug() << n << " "; - } - comet_debug() << "\n"; - } - - std::vector leafop_inputTensors; - getInputTensorsOfComputeOp(leafop, leafop_inputTensors /* output */); - comet_debug() << " getFormatsInfo:leafop_inputTensors.size(): " << leafop_inputTensors.size() << "\n"; - - std::vector leafop_outputTensors; - getOutputTensorsOfComputeOp(leafop, leafop_outputTensors /* output */); - comet_debug() << " getFormatsInfo:leafop_outputTensors.size(): " << leafop_outputTensors.size() << "\n"; - - std::vector leafop_tensors = leafop_inputTensors; - leafop_tensors.insert(leafop_tensors.end(), leafop_outputTensors.begin(), leafop_outputTensors.end()); -#ifdef DEBUG_MODE_UTILS - comet_debug() << " getFormatsInfo:leafop_tensors.size(): " << leafop_tensors.size() << "\n"; - for (auto n : leafop_tensors) - { - comet_debug() << " "; - comet_vdump(n); - } -#endif - /// Check if this index is in this leaf's perms - - std::vector formats_local; - std::vector tensors_local; - std::vector ids_local; - std::vector rhs_vs_lhs; - /// This leafOp contain multiple tensors. - comet_debug() << " getFormatsInfo:allPerms.size()" << allPerms.size() << "\n"; - for (unsigned long k = 0; k < allPerms.size(); k++) - { - comet_debug() << " getFormatsInfo:allPerms[" << k << "].size(): " << allPerms[k].size() << ", print allPerms[" << k << "]: "; - print_vector(allPerms[k]); - comet_debug() << " getFormatsInfo:indices[" << i << "]: " << indices[i] << "\n"; - unsigned int idx = findIndexInVector(allPerms[k], indices[i]); - comet_debug() << " getFormatsInfo:idx: " << idx << ", allPerms[" << k << "].size(): " << allPerms[k].size() << "\n"; - if (idx < allPerms[k].size()) - { /// In tensor k - comet_debug() << " getFormatsInfo:AddingLocalFormat[" << k << "][" << idx << "]: " << allFormats[k][idx] << " "; - comet_vdump(leafop_tensors[k]); - formats_local.push_back(allFormats[k][idx]); - tensors_local.push_back(leafop_tensors[k]); - ids_local.push_back(idx); - rhs_vs_lhs.push_back(inputOutputMapping[k][idx]); - } - } - - comet_debug() << " getFormatsInfo:formats_local.size(): " << formats_local.size() << " \n"; - for (unsigned long k = 0; k < formats_local.size(); k++) - { - comet_debug() << " getFormatsInfo:formats_local[k]:" << formats_local[k] << " " << ids_local[k] << " "; - comet_vdump(tensors_local[k]); - } - - /// analyze _local arrays, to get final formats, tensors, idx - if (formats_local.size() > 0) - { - isSetInLeaf = true; - format_in_leaf = formats_local[0]; - tensor_in_leaf = tensors_local[0]; - id_in_leaf = ids_local[0]; - - for (unsigned long k = 1; k < formats_local.size(); k++) - { - if (format_in_leaf.compare(0, 1, "D") == 0 && formats_local[k].compare(0, 1, "D") != 0 && rhs_vs_lhs[k]) - /// if the next format in the local format is not dense and not output - /// rhs_vs_lhs determines if the format comes from input (lhs) or output (rhs) - /// C[i,j] = A[i,k] * B[k, j] -> i is in both input A and output C - /// -> j is in both input B and output C - /// -> k is in both inputs A and B - /// index format information stores in formats_local - { - format_in_leaf = formats_local[k]; - tensor_in_leaf = tensors_local[k]; - id_in_leaf = ids_local[k]; - break; /// Get the first sparse case - } - } - } - - } /// if(indexTree::IndexTreeComputeOp leafop - - if (isSetInLeaf) - { - comet_debug() << " getFormatsInfo:isSetInLeaf: " << isSetInLeaf << ", format_in_leaf: " << format_in_leaf << ", id_in_leaf: " << id_in_leaf << ", tensor: "; - comet_vdump(tensor_in_leaf); - formats_leafs.push_back(format_in_leaf); - tensors_leafs.push_back(tensor_in_leaf); - ids_leafs.push_back(id_in_leaf); - } - - } /// for(auto j = 0; j < leafs.size(); j++){ - - comet_debug() << " getFormatsInfo:formats_leafs.size(): " << formats_leafs.size() << "\n"; - for ([[maybe_unused]] unsigned long k = 0; k < formats_leafs.size(); k++) - { - comet_debug() << " getFormatsInfo:formats_leafs[k]:" << formats_leafs[k] << "\n"; - } - - /// analyze the _leafs info to get the current index format, tensor, id information - for (unsigned long j = 0; j < formats_leafs.size(); j++) - { - if (j == 0) - { - format = formats_leafs[j]; - tensor = tensors_leafs[j]; - id = ids_leafs[j]; - isSet = true; - } - else - { - if (formats_leafs[j].compare(0, 1, "D") != 0) - { /// not D - format = formats_leafs[j]; - tensor = tensors_leafs[j]; - id = ids_leafs[j]; - isSet = true; - break; /// Get the first sparse case - } - } - } - - if (isSet) - { - comet_debug() << " getFormatsInfo:EndFormat: " << format << ", id: " << id << ", tensor: "; - comet_vdump(tensor); - - formats.push_back(format); - tensors.push_back(tensor); - ids.push_back(id); - } - - } /// for(auto i = 0; i < indices.size(); i++){ + return; +// /// For each indices, find in each leaf, which tensor, the corresponding format +// /// If in all tensors, the formats of the index are D, then D +// /// If only one Sparse, then sparse +// comet_debug() << " getFormatsInfo:Start Current op\n"; +// comet_vdump(cur_op); +// comet_debug() << " getFormatsInfo:indices.size(): " << indices.size() << "\n"; +// for (unsigned long i = 0; i < indices.size(); i++) +// { +// comet_debug() << " getFormatsInfo:indices[" << i << "]: " << indices[i] << "\n"; +// /// Info for each index +// std::string format; +// Value tensor; +// unsigned int id; +// bool isSet = false; + +// std::vector formats_leafs; +// std::vector tensors_leafs; +// std::vector ids_leafs; + +// for (unsigned long j = 0; j < leafs.size(); j++) +// { +// /// Info for each index in leaf[j] +// comet_debug() << " getFormatsInfo:LeafOp: "; +// comet_vdump(leafs[j]); +// std::string format_in_leaf; +// Value tensor_in_leaf; +// unsigned int id_in_leaf; +// bool isSetInLeaf = false; + +// /// get All perms and formats info +// if (indexTree::IndexTreeComputeOp leafop = dyn_cast(leafs[j].getDefiningOp())) +// { +// comet_debug() << " getFormatsInfo:leafs[" << j << "] is computeOp\n"; +// std::vector> allFormats; +// std::vector> allPerms; +// std::vector> inputOutputMapping; +// OpBuilder builder(leafop); +// getFormatsPermsOfComputeOp(leafop, allFormats, allPerms, inputOutputMapping); + +// comet_debug() << " getFormatsInfo:Allformats allFormats.size(): " << allFormats.size() << "\n"; +// for (auto m : allFormats) +// { +// comet_debug() << " "; +// for (auto n : m) +// { +// comet_debug() << n << " "; +// } +// comet_debug() << "\n"; +// } + +// std::vector leafop_inputTensors; +// getInputTensorsOfComputeOp(leafop, leafop_inputTensors); +// comet_debug() << " getFormatsInfo:leafop_inputTensors.size(): " << leafop_inputTensors.size() << "\n"; + +// std::vector leafop_outputTensors; +// getOutputTensorsOfComputeOp(leafop, leafop_outputTensors); +// comet_debug() << " getFormatsInfo:leafop_outputTensors.size(): " << leafop_outputTensors.size() << "\n"; + +// std::vector leafop_tensors = leafop_inputTensors; +// leafop_tensors.insert(leafop_tensors.end(), leafop_outputTensors.begin(), leafop_outputTensors.end()); +// #ifdef DEBUG_MODE_UTILS +// comet_debug() << " getFormatsInfo:leafop_tensors.size(): " << leafop_tensors.size() << "\n"; +// for (auto n : leafop_tensors) +// { +// comet_debug() << " "; +// comet_vdump(n); +// } +// #endif +// /// Check if this index is in this leaf's perms + +// std::vector formats_local; +// std::vector tensors_local; +// std::vector ids_local; +// std::vector rhs_vs_lhs; +// /// This leafOp contain multiple tensors. +// comet_debug() << " getFormatsInfo:allPerms.size()" << allPerms.size() << "\n"; +// for (unsigned long k = 0; k < allPerms.size(); k++) +// { +// comet_debug() << " getFormatsInfo:allPerms[" << k << "].size(): " << allPerms[k].size() << ", print allPerms[" << k << "]: "; +// print_vector(allPerms[k]); +// comet_debug() << " getFormatsInfo:indices[" << i << "]: " << indices[i] << "\n"; +// unsigned int idx = findIndexInVector(allPerms[k], indices[i]); +// comet_debug() << " getFormatsInfo:idx: " << idx << ", allPerms[" << k << "].size(): " << allPerms[k].size() << "\n"; +// if (idx < allPerms[k].size()) +// { /// In tensor k +// comet_debug() << " getFormatsInfo:AddingLocalFormat[" << k << "][" << idx << "]: " << allFormats[k][idx] << " "; +// comet_vdump(leafop_tensors[k]); +// formats_local.push_back(allFormats[k][idx]); +// tensors_local.push_back(leafop_tensors[k]); +// ids_local.push_back(idx); +// rhs_vs_lhs.push_back(inputOutputMapping[k][idx]); +// } +// } + +// comet_debug() << " getFormatsInfo:formats_local.size(): " << formats_local.size() << " \n"; +// for (unsigned long k = 0; k < formats_local.size(); k++) +// { +// comet_debug() << " getFormatsInfo:formats_local[k]:" << formats_local[k] << " " << ids_local[k] << " "; +// comet_vdump(tensors_local[k]); +// } + +// /// analyze _local arrays, to get final formats, tensors, idx +// if (formats_local.size() > 0) +// { +// isSetInLeaf = true; +// format_in_leaf = formats_local[0]; +// tensor_in_leaf = tensors_local[0]; +// id_in_leaf = ids_local[0]; + +// for (unsigned long k = 1; k < formats_local.size(); k++) +// { +// if (format_in_leaf.compare(0, 1, "D") == 0 && formats_local[k].compare(0, 1, "D") != 0 && rhs_vs_lhs[k]) +// /// if the next format in the local format is not dense and not output +// /// rhs_vs_lhs determines if the format comes from input (lhs) or output (rhs) +// /// C[i,j] = A[i,k] * B[k, j] -> i is in both input A and output C +// /// -> j is in both input B and output C +// /// -> k is in both inputs A and B +// /// index format information stores in formats_local +// { +// format_in_leaf = formats_local[k]; +// tensor_in_leaf = tensors_local[k]; +// id_in_leaf = ids_local[k]; +// break; /// Get the first sparse case +// } +// } +// } + +// } /// if(indexTree::IndexTreeComputeOp leafop + +// if (isSetInLeaf) +// { +// comet_debug() << " getFormatsInfo:isSetInLeaf: " << isSetInLeaf << ", format_in_leaf: " << format_in_leaf << ", id_in_leaf: " << id_in_leaf << ", tensor: "; +// comet_vdump(tensor_in_leaf); +// formats_leafs.push_back(format_in_leaf); +// tensors_leafs.push_back(tensor_in_leaf); +// ids_leafs.push_back(id_in_leaf); +// } + +// } /// for(auto j = 0; j < leafs.size(); j++){ + +// comet_debug() << " getFormatsInfo:formats_leafs.size(): " << formats_leafs.size() << "\n"; +// for ([[maybe_unused]] unsigned long k = 0; k < formats_leafs.size(); k++) +// { +// comet_debug() << " getFormatsInfo:formats_leafs[k]:" << formats_leafs[k] << "\n"; +// } + +// /// analyze the _leafs info to get the current index format, tensor, id information +// for (unsigned long j = 0; j < formats_leafs.size(); j++) +// { +// if (j == 0) +// { +// format = formats_leafs[j]; +// tensor = tensors_leafs[j]; +// id = ids_leafs[j]; +// isSet = true; +// } +// else +// { +// if (formats_leafs[j].compare(0, 1, "D") != 0) +// { /// not D +// format = formats_leafs[j]; +// tensor = tensors_leafs[j]; +// id = ids_leafs[j]; +// isSet = true; +// break; /// Get the first sparse case +// } +// } +// } + +// if (isSet) +// { +// comet_debug() << " getFormatsInfo:EndFormat: " << format << ", id: " << id << ", tensor: "; +// comet_vdump(tensor); + +// formats.push_back(format); +// tensors.push_back(tensor); +// ids.push_back(id); +// } + +// } /// for(auto i = 0; i < indices.size(); i++){ } /// Find leaves of tcRootOp in the Index Tree (dfsOp). @@ -1681,36 +1916,37 @@ namespace mlir std::vector &dfsOps, std::vector &ret /* output leaves */) { - std::vector> allAncestors(dfsOps.size()); - for (unsigned int i = 0; i < dfsOps.size(); i++) - { - if (IndexTreeComputeOp cur_op = dyn_cast(dfsOps[i].getDefiningOp())) - { - getAncestorsWp(dfsOps[i], allAncestors[i] /* output ancestors */, dfsOps); - comet_debug() << " print allAncestors[" << i << "]: "; - print_vector_value(allAncestors[i]); - } - } - - /// Each wp op in which tensors - if (IndexTreeIndicesOp cur_op = dyn_cast(tcRootOp.getDefiningOp())) - { - comet_debug() << " "; - comet_vdump(tcRootOp); - for (unsigned int j = 0; j < dfsOps.size(); j++) - { - auto idx = findIndexInVector_Value(allAncestors[j], tcRootOp); - if (idx < allAncestors[j].size()) - { - if (indexTree::IndexTreeComputeOp cur_op = dyn_cast(dfsOps[j].getDefiningOp())) - { - ret.push_back(dfsOps[j]); - comet_debug() << " "; - comet_vdump(dfsOps[j]); - } - } - } - } + return; + // std::vector> allAncestors(dfsOps.size()); + // for (unsigned int i = 0; i < dfsOps.size(); i++) + // { + // if (IndexTreeComputeOp cur_op = dyn_cast(dfsOps[i].getDefiningOp())) + // { + // getAncestorsWp(dfsOps[i], allAncestors[i] /* output ancestors */, dfsOps); + // comet_debug() << " print allAncestors[" << i << "]: "; + // print_vector_value(allAncestors[i]); + // } + // } + + // /// Each wp op in which tensors + // if (IndexTreeIndicesOp cur_op = dyn_cast(tcRootOp.getDefiningOp())) + // { + // comet_debug() << " "; + // comet_vdump(tcRootOp); + // for (unsigned int j = 0; j < dfsOps.size(); j++) + // { + // auto idx = findIndexInVector_Value(allAncestors[j], tcRootOp); + // if (idx < allAncestors[j].size()) + // { + // if (indexTree::IndexTreeComputeOp cur_op = dyn_cast(dfsOps[j].getDefiningOp())) + // { + // ret.push_back(dfsOps[j]); + // comet_debug() << " "; + // comet_vdump(dfsOps[j]); + // } + // } + // } + // } } /// new version for new children ops @@ -1753,61 +1989,65 @@ namespace mlir /// Get the output tensors of the itCompute op void getTensorsOfComputeOp(Value computeOp, std::vector &tensors) { - indexTree::IndexTreeComputeRHSOp itComputeOp_rhs = dyn_cast(computeOp.getDefiningOp()->getOperand(0).getDefiningOp()); - comet_debug() << " "; - comet_vdump(itComputeOp_rhs); - for (unsigned int i = 0; i < itComputeOp_rhs.getOperation()->getNumOperands(); i++) - { - comet_debug() << " "; - comet_vdump(itComputeOp_rhs.getOperation()->getOperand(i)); - tensors.push_back(itComputeOp_rhs.getOperation()->getOperand(i)); - } - indexTree::IndexTreeComputeLHSOp itComputeOp_lhs = dyn_cast(computeOp.getDefiningOp()->getOperand(1).getDefiningOp()); - for (unsigned int i = 0; i < itComputeOp_lhs.getOperation()->getNumOperands(); i++) - { - tensors.push_back(itComputeOp_lhs.getOperation()->getOperand(i)); - } + return; + // indexTree::IndexTreeComputeRHSOp itComputeOp_rhs = dyn_cast(computeOp.getDefiningOp()->getOperand(0).getDefiningOp()); + // comet_debug() << " "; + // comet_vdump(itComputeOp_rhs); + // for (unsigned int i = 0; i < itComputeOp_rhs.getOperation()->getNumOperands(); i++) + // { + // comet_debug() << " "; + // comet_vdump(itComputeOp_rhs.getOperation()->getOperand(i)); + // tensors.push_back(itComputeOp_rhs.getOperation()->getOperand(i)); + // } + // indexTree::IndexTreeComputeLHSOp itComputeOp_lhs = dyn_cast(computeOp.getDefiningOp()->getOperand(1).getDefiningOp()); + // for (unsigned int i = 0; i < itComputeOp_lhs.getOperation()->getNumOperands(); i++) + // { + // tensors.push_back(itComputeOp_lhs.getOperation()->getOperand(i)); + // } } /// Get the perms and formats of the itCompute op void getRHSPermsOfComputeOp(Value computeOp, std::vector> &opPerms) { - indexTree::IndexTreeComputeRHSOp itComputeOp_rhs = dyn_cast(computeOp.getDefiningOp()->getOperand(0).getDefiningOp()); - ArrayAttr opPermsArrayAttr_rhs = itComputeOp_rhs.getAllPerms(); - /// Get output format, vector of vector - /// Convert ArrayAttr into - std::vector> opPerms_rhs = convertArrayAttrIntTo2DVector(opPermsArrayAttr_rhs); - opPerms = opPerms_rhs; + return; + // indexTree::IndexTreeComputeRHSOp itComputeOp_rhs = dyn_cast(computeOp.getDefiningOp()->getOperand(0).getDefiningOp()); + // ArrayAttr opPermsArrayAttr_rhs = itComputeOp_rhs.getAllPerms(); + // /// Get output format, vector of vector + // /// Convert ArrayAttr into + // std::vector> opPerms_rhs = convertArrayAttrIntTo2DVector(opPermsArrayAttr_rhs); + // opPerms = opPerms_rhs; } /// Get the perms and formats of the itCompute op void getLHSPermsOfComputeOp(Value computeOp, std::vector> &opPerms) { - indexTree::IndexTreeComputeLHSOp itComputeOp_lhs = dyn_cast(computeOp.getDefiningOp()->getOperand(1).getDefiningOp()); - ArrayAttr opPermsArrayAttr_lhs = itComputeOp_lhs.getAllPerms(); + return; + // indexTree::IndexTreeComputeLHSOp itComputeOp_lhs = dyn_cast(computeOp.getDefiningOp()->getOperand(1).getDefiningOp()); + // ArrayAttr opPermsArrayAttr_lhs = itComputeOp_lhs.getAllPerms(); - /// Get output format, vector of vector - /// Convert ArrayAttr into - std::vector> opPerms_lhs = convertArrayAttrIntTo2DVector(opPermsArrayAttr_lhs); + // /// Get output format, vector of vector + // /// Convert ArrayAttr into + // std::vector> opPerms_lhs = convertArrayAttrIntTo2DVector(opPermsArrayAttr_lhs); - opPerms = opPerms_lhs; + // opPerms = opPerms_lhs; } /// Get the perms and formats of the itCompute op void getPermsOfComputeOp(Value computeOp, std::vector> &opPerms) { - indexTree::IndexTreeComputeRHSOp itComputeOp_rhs = dyn_cast(computeOp.getDefiningOp()->getOperand(0).getDefiningOp()); - ArrayAttr opPermsArrayAttr_rhs = itComputeOp_rhs.getAllPerms(); - indexTree::IndexTreeComputeLHSOp itComputeOp_lhs = dyn_cast(computeOp.getDefiningOp()->getOperand(1).getDefiningOp()); - ArrayAttr opPermsArrayAttr_lhs = itComputeOp_lhs.getAllPerms(); - - /// Get output format, vector of vector - /// Convert ArrayAttr into - std::vector> opPerms_rhs = convertArrayAttrIntTo2DVector(opPermsArrayAttr_rhs); - std::vector> opPerms_lhs = convertArrayAttrIntTo2DVector(opPermsArrayAttr_lhs); - - opPerms = opPerms_rhs; - opPerms.insert(opPerms.end(), opPerms_lhs.begin(), opPerms_lhs.end()); + return; + // indexTree::IndexTreeComputeRHSOp itComputeOp_rhs = dyn_cast(computeOp.getDefiningOp()->getOperand(0).getDefiningOp()); + // ArrayAttr opPermsArrayAttr_rhs = itComputeOp_rhs.getAllPerms(); + // indexTree::IndexTreeComputeLHSOp itComputeOp_lhs = dyn_cast(computeOp.getDefiningOp()->getOperand(1).getDefiningOp()); + // ArrayAttr opPermsArrayAttr_lhs = itComputeOp_lhs.getAllPerms(); + + // /// Get output format, vector of vector + // /// Convert ArrayAttr into + // std::vector> opPerms_rhs = convertArrayAttrIntTo2DVector(opPermsArrayAttr_rhs); + // std::vector> opPerms_lhs = convertArrayAttrIntTo2DVector(opPermsArrayAttr_lhs); + + // opPerms = opPerms_rhs; + // opPerms.insert(opPerms.end(), opPerms_lhs.begin(), opPerms_lhs.end()); } double loopCostHeuristic(const std::vector &loopOrder, size_t dim_, @@ -1819,7 +2059,7 @@ namespace mlir { /// row major: last one has no penalty /// const int idx = loopOrder[dim_-1-i]; /// const int posB = findPos(idx, perm_); - int idx; /// position in sourceOrder + int idx = -1; /// position in sourceOrder int posB = 0; /// position in destOrder for (unsigned ii = 0; ii < sourceOrder.size(); ii++) { @@ -1842,6 +2082,7 @@ namespace mlir /// int importanceB = (1<<(dim_ - posB));/// subsequent indices are half as important /// int penalty = 10 * (1<<(i-1)); /// smaller i has smaller penalty /*row major: */ + assert(idx != -1); int importanceA = (1 << idx); /// stride-1 has the most importance. Larger pos, more important int importanceB = (1 << posB); /// subsequent indices are half as important int penalty = 10 * (1 << (dim_ - (i + 2))); /// smaller i has larger penalty @@ -1932,8 +2173,8 @@ namespace mlir auto rhs1AlphaAttr = rhs1Tensor.getDefiningOp()->getAttr("__alpha__"); auto rhs2AlphaAttr = rhs2Tensor.getDefiningOp()->getAttr("__alpha__"); - alpha *= rhs1AlphaAttr.cast().getValueAsDouble(); - alpha *= rhs2AlphaAttr.cast().getValueAsDouble(); + alpha *= cast(rhs1AlphaAttr).getValueAsDouble(); + alpha *= cast(rhs2AlphaAttr).getValueAsDouble(); unsigned idx = 0; for (auto lbl : rhs1Labels) @@ -1990,45 +2231,14 @@ namespace mlir auto affineMapArrayAttr = rewriter.getAffineMapArrayAttr(affineMaps); comet_debug() << "\n"; - SmallVector formats; - std::vector defops{rhs1Tensor.getDefiningOp(), rhs2Tensor.getDefiningOp(), lhsTensor.getDefiningOp()}; - for (auto defop : defops) - { - comet_debug() << " "; - comet_pdump(defop); - if (isa(defop)) - { - comet_debug() << " is TensorDeclOp\n"; - /// infer the format - auto lhs_format = dyn_cast(defop).getFormat(); - comet_debug() << " lhs_format: " << lhs_format << "\n"; - formats.push_back(lhs_format); - } - else if (isa(defop)) - { - comet_debug() << " is TensorDeclOp\n"; - - /// infer the format - auto lhs_format = dyn_cast(defop).getFormat(); - comet_debug() << " lhs_format: " << lhs_format << "\n"; - formats.push_back(lhs_format); - } - else - { - comet_debug() << " not TensorDeclOp\n"; - } - } - - auto formatAttr = rewriter.getStrArrayAttr(formats); - comet_debug() << " formatAttr: " << formatAttr << "\n"; auto SemiringAttr = rewriter.getStringAttr("none"); auto MaskingAttr = rewriter.getStringAttr("none"); auto tc = rewriter.create(loc, lhsTensor.getType(), rhs1Tensor, rhs2Tensor, lhsLabels, affineMapArrayAttr, - formatAttr, SemiringAttr, MaskingAttr, + SemiringAttr, MaskingAttr, nullptr); /// TODO: masking is an optional operand tc.getOperation()->setAttr("__alpha__", rewriter.getF64FloatAttr(alpha)); tc.getOperation()->setAttr("__beta__", rewriter.getF64FloatAttr(beta)); @@ -2097,6 +2307,45 @@ namespace mlir return reassociation; } + TypedValue collapseMemref(TypedValue val, mlir::OpBuilder& builder) + { + + auto memref = mlir::cast(val.getType()); + if (memref.getRank() == 1) + { + return val; + } + + llvm::SmallVector,1> indices; + indices.push_back(llvm::SmallVector()); + for(int64_t i = 0; i < memref.getRank(); i++) + { + indices[0].push_back(i); + } + + /// Collapse memref to 1D + auto collapsedMemref = builder.create(val.getLoc(), val, mlir::ArrayRef(indices)).getResult(); + return collapsedMemref; + } + + mlir::Value get_memref_num_elements(mlir::MLIRContext* ctx, mlir::OpBuilder& builder, mlir::Location loc, mlir::Value memref) + { + mlir::Value rank = builder.create(loc, memref); + mlir::Value zero = builder.create(loc, 0); + mlir::Value one = builder.create(loc, 1); + + mlir::scf::ForOp forOp = builder.create(loc, zero, rank, one, mlir::ValueRange({one})); + mlir::Block* body = forOp.getBody(); + mlir::Value inductionvar = forOp.getInductionVar(); + mlir::IRRewriter::InsertPoint ip = builder.saveInsertionPoint(); + builder.setInsertionPointToStart(body); + mlir::Value dim = builder.create(loc, memref, inductionvar); + auto mul = builder.create(loc, forOp.getRegionIterArg(0), dim); + builder.create(loc, mlir::ValueRange({mul})); + builder.restoreInsertionPoint(ip); + + return forOp.getResult(0); + } } //// namespace tensorAlgebra } //// namespace mlir diff --git a/lib/ExecutionEngine/CMakeLists.txt b/lib/ExecutionEngine/CMakeLists.txt index 7e59d5af..cc207a6a 100644 --- a/lib/ExecutionEngine/CMakeLists.txt +++ b/lib/ExecutionEngine/CMakeLists.txt @@ -1,5 +1,15 @@ -if(ENABLE_GPU_TARGET) -find_package(CUDAToolkit REQUIRED) +if(ENABLE_NVIDIA_GPU_BACKEND) + find_package(CUDAToolkit) + if(NOT ${CUDAToolkit_FOUND}) + message(WARNING "CUDAToolkit not found. You can still lower to GPU Assembly code (use --gpu-code-format=Assembly), but you will not be able to run it using COMET's execution engine") + endif() +endif() + +if(ENABLE_AMD_GPU_BACKEND) + find_package(hip) + if(NOT ${hip_FOUND}) + message(WARNING "Hip not found. Unable to use GPU target") + endif() endif() set(SOURCES @@ -10,11 +20,16 @@ set(SOURCES ) -if(ENABLE_GPU_TARGET) - set(SOURCES - ${SOURCES} - GpuUtils.cpp - ) +if(ENABLE_NVIDIA_GPU_BACKEND) + if(${CUDAToolkit_FOUND}) + set(SOURCES ${SOURCES} CuGpuUtils.cpp) + endif() +endif() + +if(ENABLE_AMD_GPU_BACKEND) + if(${hip_FOUND}) + set(SOURCES ${SOURCES} HipGpuUtils.cpp) + endif() endif() add_llvm_library(comet_runner_utils @@ -29,10 +44,17 @@ set(LIBS LLVMSupport ) -if(ENABLE_GPU_TARGET) -set(LIBS - ${LIBS} - CUDA::cuda_driver - ) +if(ENABLE_NVIDIA_GPU_BACKEND) + if(${CUDAToolkit_FOUND}) + set(LIBS ${LIBS} CUDA::cuda_driver) + set(LIBS ${LIBS} CUDA::cudart) + endif() +endif() + +if(ENABLE_AMD_GPU_BACKEND) + if(${hip_FOUND}) + set(LIBS ${LIBS} hip::host) + endif() endif() + target_link_libraries(comet_runner_utils ${LIBS}) \ No newline at end of file diff --git a/lib/ExecutionEngine/GpuUtils.cpp b/lib/ExecutionEngine/CuGpuUtils.cpp similarity index 54% rename from lib/ExecutionEngine/GpuUtils.cpp rename to lib/ExecutionEngine/CuGpuUtils.cpp index 8436d073..88320832 100644 --- a/lib/ExecutionEngine/GpuUtils.cpp +++ b/lib/ExecutionEngine/CuGpuUtils.cpp @@ -1,9 +1,32 @@ +// +// Copyright 2022 Battelle Memorial Institute +// +// Redistribution and use in source and binary forms, with or without modification, +// are permitted provided that the following conditions are met: +// +// 1. Redistributions of source code must retain the above copyright notice, this list of conditions +// and the following disclaimer. +// +// 2. Redistributions in binary form must reproduce the above copyright notice, this list of conditions +// and the following disclaimer in the documentation and/or other materials provided with the distribution. +// +// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED +// WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR +// PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, +// INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE +// GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, +// WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +// OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. +// + +#include #include #include #include #include #include #include +#include #include #define CU_CHECK(call) \ @@ -15,22 +38,25 @@ do { \ exit(1); \ } \ } while (0) +#define CUDA_CHECK(call) \ +do { \ + cudaError_t res = call; \ + if (res != cudaSuccess) { \ + fprintf(stderr, "CUDA Error: %s:%d, ", __FILE__, __LINE__); \ + fprintf(stderr, "code: %d\n", res); \ + exit(1); \ + } \ +} while (0) -CUcontext cuContext = NULL; CUmodule cuModule = NULL; -char* moduleImg = NULL; +// char* moduleImg = NULL; -void initCudaCtx() +void initCudaCtx(char* moduleImg) { - if(!cuContext) - { - CU_CHECK(cuInit(0)); - CU_CHECK(cuCtxCreate(&cuContext, 0, 0)); - } - - if(!cuModule && moduleImg) + + if(moduleImg) { - // printf("PTX: \n %s\n", moduleImg); + cudaFree(0); // To initialize the context when needed CU_CHECK(cuModuleLoadData(&cuModule, moduleImg)); } } @@ -38,22 +64,19 @@ void initCudaCtx() template int64_t cudaMalloc(int64_t size) { - initCudaCtx(); - CUdeviceptr device_ptr; - CU_CHECK(cuMemAlloc(&device_ptr, size * sizeof(T))); + void* device_ptr; + cudaMalloc(&device_ptr, size * sizeof(T)); // return (int64_t)malloc(size * sizeof(T)); - return device_ptr; + return (int64_t)(device_ptr); } extern "C" __attribute__((visibility("default"))) void cudaSetModuleImage(char* ptx) { - // printf("Called cudaSetModuleImage"); - - if (moduleImg == NULL) + // if (moduleImg == NULL) { // printf("Setting PTX"); - moduleImg = ptx; + initCudaCtx(ptx); } } @@ -77,33 +100,26 @@ extern "C" __attribute__((visibility("default"))) int64_t cudaMallocI32(int64_t return cudaMalloc(size); } -extern "C" __attribute__((visibility("default"))) void cudaFree(int64_t ptr) { - initCudaCtx(); - - // printf("Freeing memory\n"); - CU_CHECK(cuMemFree(ptr)); - // free((void*)ptr); -} - template void cudaMemcpy(int64_t device, void* ptr, void* aligned_ptr, int64_t offset, int64_t size, int64_t stride, int64_t direction) { - initCudaCtx(); - - // printf("Memcpy memory of size: %ld\n", size); if(direction == 0) // Host to Device { - CU_CHECK(cuMemcpyHtoD(device, aligned_ptr, size * sizeof(T))); + CUDA_CHECK(cudaMemcpy((void*)device, aligned_ptr, size * sizeof(T), cudaMemcpyHostToDevice)); // memcpy((void*)device, aligned_ptr, size * sizeof(T)); } else // Device to Host { - CU_CHECK(cuMemcpyDtoH(aligned_ptr, device, size * sizeof(T))); + CUDA_CHECK(cudaMemcpy(aligned_ptr, (void*)device, size * sizeof(T), cudaMemcpyDeviceToHost)); // memcpy(aligned_ptr, (void*)device, size * sizeof(T)); } } +extern "C" __attribute__((visibility("default"))) void cudaMemcpyIndex(int64_t device, void* ptr, void* aligned_ptr, int64_t offset, int64_t size, int64_t stride, int64_t direction) { + cudaMemcpy(device, ptr, aligned_ptr, offset, size, stride, direction); +} + extern "C" __attribute__((visibility("default"))) void cudaMemcpyI64(int64_t device, void* ptr, void* aligned_ptr, int64_t offset, int64_t size, int64_t stride, int64_t direction) { cudaMemcpy(device, ptr, aligned_ptr, offset, size, stride, direction); } @@ -133,25 +149,20 @@ const int64_t MAX_NUM_BLOCKS_X = 2147483647; const int64_t MAX_NUM_BLOCKS_Y = 65535; const int64_t MAX_NUM_BLOCKS_Z = 65535; -extern "C" __attribute__((visibility("default"))) void cudaLaunchKernel(int64_t realblocksX, int64_t realblocksY, int64_t realblocksZ, int64_t tritonBlockX, int64_t tritonBlockY, int64_t tritonBlockZ, void* ptr, void* aligned_ptr, int64_t offset, int64_t size, int64_t stride, char* kernel, int64_t kernel_name_size, int64_t sharedMem, int64_t numWraps, int64_t threadsPerWarp) +extern "C" __attribute__((visibility("default"))) void cudaLaunchKernelMLIR(int64_t realblocksX, int64_t realblocksY, int64_t realblocksZ, int64_t tritonBlockX, int64_t tritonBlockY, int64_t tritonBlockZ, void* ptr, void* aligned_ptr, int64_t offset, int64_t size, int64_t stride, char* kernel, int64_t kernel_name_size, int64_t sharedMem) { - initCudaCtx(); CUfunction cuFunction = NULL; unsigned blocksPerGridX = std::min(realblocksX, MAX_NUM_BLOCKS_X); unsigned blocksPerGridY = std::min(realblocksY, MAX_NUM_BLOCKS_Y); - unsigned blocksPerGridZ = std::min(realblocksZ, MAX_NUM_BLOCKS_Y); - char* name_with_suffix = (char*)malloc(kernel_name_size+1); - memcpy(name_with_suffix, kernel, kernel_name_size); - name_with_suffix[kernel_name_size] = '\0'; - CU_CHECK(cuModuleGetFunction(&cuFunction, cuModule, name_with_suffix)); - // printf("Kernel name: %s\n", name_with_suffix); + unsigned blocksPerGridZ = std::min(realblocksZ, MAX_NUM_BLOCKS_Z); + CU_CHECK(cuModuleGetFunction(&cuFunction, cuModule, kernel)); void** cast_args = (void**)aligned_ptr; - // for(int i = 0; i < size; i++) - // { - // printf("%p\n", cast_args[i]); - // } - free(name_with_suffix); + CU_CHECK(cuLaunchKernel(cuFunction, blocksPerGridX, blocksPerGridY, blocksPerGridZ, tritonBlockX, tritonBlockY, tritonBlockZ, sharedMem, 0, cast_args, 0)); +} - CU_CHECK(cuLaunchKernel(cuFunction, blocksPerGridX, blocksPerGridY, blocksPerGridZ, numWraps* threadsPerWarp, 1, 1, sharedMem, 0, cast_args, 0)); +extern "C" __attribute__((visibility("default"))) void cudaFinit() +{ + cuModuleUnload(cuModule); + cuModule = NULL; } \ No newline at end of file diff --git a/lib/ExecutionEngine/HipGpuUtils.cpp b/lib/ExecutionEngine/HipGpuUtils.cpp new file mode 100644 index 00000000..2f1c25a1 --- /dev/null +++ b/lib/ExecutionEngine/HipGpuUtils.cpp @@ -0,0 +1,160 @@ +#include "hip/hip_runtime.h" +#include +#include +#include +#include +#include +#include +#include + +#define HIP_CHECK(call) \ + do { \ + hipError_t res = call; \ + if (res != hipSuccess) { \ + fprintf(stderr, "hip Error: %s:%d, ", __FILE__, __LINE__); \ + fprintf(stderr, "code: %d\n", res); \ + fprintf(stderr, "error: %s\n", hipGetErrorString(res)); \ + exit(1); \ + } \ + } while (0) + +hipModule_t hipModule = NULL; +// char* moduleImg = NULL; + +void initHipCtx(char *moduleImg) { + if (moduleImg) { + HIP_CHECK(hipFree(NULL)); + HIP_CHECK(hipModuleLoadData(&hipModule, moduleImg)); + } +} + +template int64_t HipMalloc(int64_t size) { + hipDeviceptr_t device_ptr; + HIP_CHECK(hipMalloc(&device_ptr, size * sizeof(T))); + + // return (int64_t)malloc(size * sizeof(T)); + return (int64_t)device_ptr; +} + +extern "C" __attribute__((visibility("default"))) void +HipSetModuleImage(char *ptx) { + // if (moduleImg == NULL) + { + // printf("Setting PTX"); + initHipCtx(ptx); + } +} + +extern "C" __attribute__((visibility("default"))) int64_t +HipMallocF64(int64_t size) { + // printf("Allocating memory of size: %ld\n", size); + return HipMalloc(size); +} + +extern "C" __attribute__((visibility("default"))) int64_t +HipMallocF32(int64_t size) { + // printf("Allocating memory of size: %ld\n", size); + return HipMalloc(size); +} + +extern "C" __attribute__((visibility("default"))) int64_t +HipMallocI64(int64_t size) { + // printf("Allocating memory of size: %ld\n", size); + return HipMalloc(size); +} + +extern "C" __attribute__((visibility("default"))) int64_t +HipMallocI32(int64_t size) { + // printf("Allocating memory of size: %ld\n", size); + return HipMalloc(size); +} + +extern "C" __attribute__((visibility("default"))) void HipFree(int64_t ptr) { + + // printf("Freeing memory\n"); + HIP_CHECK(hipFree((void *)ptr)); + // free((void*)ptr); +} + +template +void HipMemcpy(int64_t device, void *ptr, void *aligned_ptr, int64_t offset, + int64_t size, int64_t stride, int64_t direction) { + + // printf("Memcpy memory of size: %ld\n", size); + + if (direction == 0) // Host to Device + { + HIP_CHECK(hipMemcpyHtoD((void *)device, aligned_ptr, size * sizeof(T))); + // memcpy((void*)device, aligned_ptr, size * sizeof(T)); + } else // Device to Host + { + HIP_CHECK(hipMemcpyDtoH(aligned_ptr, (void *)device, size * sizeof(T))); + // memcpy(aligned_ptr, (void*)device, size * sizeof(T)); + } +} + +extern "C" __attribute__((visibility("default"))) void +HipMemcpyIndex(int64_t device, void *ptr, void *aligned_ptr, int64_t offset, + int64_t size, int64_t stride, int64_t direction) { + HipMemcpy(device, ptr, aligned_ptr, offset, size, stride, direction); +} + +extern "C" __attribute__((visibility("default"))) void +HipMemcpyI64(int64_t device, void *ptr, void *aligned_ptr, int64_t offset, + int64_t size, int64_t stride, int64_t direction) { + HipMemcpy(device, ptr, aligned_ptr, offset, size, stride, direction); +} + +extern "C" __attribute__((visibility("default"))) void +HipMemcpyF64(int64_t device, void *ptr, void *aligned_ptr, int64_t offset, + int64_t size, int64_t stride, int64_t direction) { + HipMemcpy(device, ptr, aligned_ptr, offset, size, stride, direction); +} + +extern "C" __attribute__((visibility("default"))) void +HipMemcpyI32(int64_t device, void *ptr, void *aligned_ptr, int64_t offset, + int64_t size, int64_t stride, int64_t direction) { + HipMemcpy(device, ptr, aligned_ptr, offset, size, stride, direction); +} + +extern "C" __attribute__((visibility("default"))) void +HipMemcpyF32(int64_t device, void *ptr, void *aligned_ptr, int64_t offset, + int64_t size, int64_t stride, int64_t direction) { + HipMemcpy(device, ptr, aligned_ptr, offset, size, stride, direction); +} + +struct UnrankedMemRef { + + void *ptr; + void *aligned_ptr; + int64_t offset; + int64_t sizes[2]; + int64_t strides[2]; +}; + +const int64_t MAX_NUM_BLOCKS_X = 2147483647; +const int64_t MAX_NUM_BLOCKS_Y = 65535; +const int64_t MAX_NUM_BLOCKS_Z = 65535; + +extern "C" __attribute__((visibility("default"))) void +HipLaunchKernelMLIR(int64_t realblocksX, int64_t realblocksY, int64_t realblocksZ, + int64_t tritonBlockX, int64_t tritonBlockY, + int64_t tritonBlockZ, void *ptr, void *aligned_ptr, + int64_t offset, int64_t size, int64_t stride, char *kernel, + int64_t kernel_name_size, int64_t sharedMem) { + hipFunction_t hipFunction = NULL; + unsigned blocksPerGridX = std::min(realblocksX, MAX_NUM_BLOCKS_X); + unsigned blocksPerGridY = std::min(realblocksY, MAX_NUM_BLOCKS_Y); + unsigned blocksPerGridZ = std::min(realblocksZ, MAX_NUM_BLOCKS_Z); + HIP_CHECK(hipModuleGetFunction(&hipFunction, hipModule, kernel)); + void **cast_args = (void **)aligned_ptr; + HIP_CHECK(hipModuleLaunchKernel(hipFunction, blocksPerGridX, blocksPerGridY, + blocksPerGridZ, tritonBlockX, tritonBlockY, + tritonBlockZ, sharedMem, NULL, cast_args, NULL)); +} + +extern "C" __attribute__((visibility("default"))) void HipFinit() +{ + HIP_CHECK(hipModuleUnload(hipModule)); + hipModule = NULL; +} \ No newline at end of file diff --git a/lib/ExecutionEngine/SparseUtils.cpp b/lib/ExecutionEngine/SparseUtils.cpp index e33bc433..9c40a225 100644 --- a/lib/ExecutionEngine/SparseUtils.cpp +++ b/lib/ExecutionEngine/SparseUtils.cpp @@ -26,6 +26,7 @@ #include "llvm/Support/raw_ostream.h" #include +#include #include #include #include @@ -1442,7 +1443,7 @@ struct FileReaderWrapper bool readFileNameStr(int32_t fileID) { - char *pSparseInput; + char *pSparseInput = NULL; std::string envString; if (fileID >= 0 && fileID < 9999) { @@ -1809,7 +1810,7 @@ void read_input_sizes_2D(int32_t fileID, } } -template +template void read_input_2D(int32_t fileID, int32_t A1format, int32_t A1_tile_format, int32_t A2format, int32_t A2_tile_format, @@ -1825,14 +1826,14 @@ void read_input_2D(int32_t fileID, int32_t readMode) { - auto *desc_A1pos = static_cast *>(A1pos_ptr); - auto *desc_A1crd = static_cast *>(A1crd_ptr); - auto *desc_A2pos = static_cast *>(A2pos_ptr); - auto *desc_A2crd = static_cast *>(A2crd_ptr); - auto *desc_A1tile_pos = static_cast *>(A1tile_pos_ptr); - auto *desc_A1tile_crd = static_cast *>(A1tile_crd_ptr); - auto *desc_A2tile_pos = static_cast *>(A2tile_pos_ptr); - auto *desc_A2tile_crd = static_cast *>(A2tile_crd_ptr); + auto *desc_A1pos = static_cast *>(A1pos_ptr); + auto *desc_A1crd = static_cast *>(A1crd_ptr); + auto *desc_A2pos = static_cast *>(A2pos_ptr); + auto *desc_A2crd = static_cast *>(A2crd_ptr); + auto *desc_A1tile_pos = static_cast *>(A1tile_pos_ptr); + auto *desc_A1tile_crd = static_cast *>(A1tile_crd_ptr); + auto *desc_A2tile_pos = static_cast *>(A2tile_pos_ptr); + auto *desc_A2tile_crd = static_cast *>(A2tile_crd_ptr); auto *desc_Aval = static_cast *>(Aval_ptr); /// For example, A2pos is not used for COO, but initialized with -1 to speficify that it is not used @@ -2117,22 +2118,22 @@ void read_input_sizes_3D(int32_t fileID, /// std::cout << "CSF format\n"; Csf3DTensor csf_3dtensor(FileReader.coo_3dtensor); - desc_sizes->data[0] = csf_3dtensor.A1pos_size; - desc_sizes->data[1] = csf_3dtensor.A1crd_size; - desc_sizes->data[2] = 0; - desc_sizes->data[3] = 0; - desc_sizes->data[4] = csf_3dtensor.A2pos_size; - desc_sizes->data[5] = csf_3dtensor.A2crd_size; - desc_sizes->data[6] = 0; - desc_sizes->data[7] = 0; - desc_sizes->data[8] = csf_3dtensor.A3pos_size; - desc_sizes->data[9] = csf_3dtensor.A3crd_size; - desc_sizes->data[10] = 0; - desc_sizes->data[11] = 0; - desc_sizes->data[12] = csf_3dtensor.Aval_size; - desc_sizes->data[13] = csf_3dtensor.num_index_i; - desc_sizes->data[14] = csf_3dtensor.num_index_j; - desc_sizes->data[15] = csf_3dtensor.num_index_k; + desc_sizes->data[0] = csf_3dtensor.A1pos_size; /// A1pos + desc_sizes->data[1] = csf_3dtensor.A1crd_size; /// A1crd + desc_sizes->data[2] = 0; /// A1_tile_pos + desc_sizes->data[3] = 0; /// A1_tile_crd + desc_sizes->data[4] = csf_3dtensor.A2pos_size; /// A2pos + desc_sizes->data[5] = csf_3dtensor.A2crd_size; /// A2crd + desc_sizes->data[6] = 0; /// A2_tile_pos + desc_sizes->data[7] = 0; /// A2_tile_crd + desc_sizes->data[8] = csf_3dtensor.A3pos_size; /// A3pos + desc_sizes->data[9] = csf_3dtensor.A3crd_size; /// A3crd + desc_sizes->data[10] = 0; /// A3_tile_pos + desc_sizes->data[11] = 0; /// A3_tile_crd + desc_sizes->data[12] = csf_3dtensor.Aval_size; /// Aval + desc_sizes->data[13] = csf_3dtensor.num_index_i; /// I + desc_sizes->data[14] = csf_3dtensor.num_index_j; /// J + desc_sizes->data[15] = csf_3dtensor.num_index_k; /// K } /// Mode-Generic else if (A1format == Compressed_nonunique && A2format == singleton && A3format == Dense) @@ -2140,22 +2141,22 @@ void read_input_sizes_3D(int32_t fileID, /// std::cout << "Mode-Generic format\n"; Mg3DTensor mg_3dtensor(FileReader.coo_3dtensor); - desc_sizes->data[0] = mg_3dtensor.A1pos_size; - desc_sizes->data[1] = mg_3dtensor.A1crd_size; - desc_sizes->data[2] = 0; - desc_sizes->data[3] = 0; - desc_sizes->data[4] = mg_3dtensor.A2pos_size; - desc_sizes->data[5] = mg_3dtensor.A2crd_size; - desc_sizes->data[6] = 0; - desc_sizes->data[7] = 0; - desc_sizes->data[8] = mg_3dtensor.A3pos_size; - desc_sizes->data[9] = mg_3dtensor.A3crd_size; - desc_sizes->data[10] = 0; - desc_sizes->data[11] = 0; - desc_sizes->data[12] = mg_3dtensor.Aval_size; - desc_sizes->data[13] = mg_3dtensor.num_index_i; - desc_sizes->data[14] = mg_3dtensor.num_index_j; - desc_sizes->data[15] = mg_3dtensor.num_index_k; + desc_sizes->data[0] = mg_3dtensor.A1pos_size; /// A1pos + desc_sizes->data[1] = mg_3dtensor.A1crd_size; /// A1crd + desc_sizes->data[2] = 0; /// A1_tile_pos + desc_sizes->data[3] = 0; /// A1_tile_crd + desc_sizes->data[4] = mg_3dtensor.A2pos_size; /// A2pos + desc_sizes->data[5] = mg_3dtensor.A2crd_size; /// A2crd + desc_sizes->data[6] = 0; /// A2_tile_pos + desc_sizes->data[7] = 0; /// A2_tile_crd + desc_sizes->data[8] = mg_3dtensor.A3pos_size; /// A3pos + desc_sizes->data[9] = mg_3dtensor.A3crd_size; /// A3crd + desc_sizes->data[10] = 0; /// A3_tile_pos + desc_sizes->data[11] = 0; /// A3_tile_crd + desc_sizes->data[12] = mg_3dtensor.Aval_size; /// Aval + desc_sizes->data[13] = mg_3dtensor.num_index_i; /// I + desc_sizes->data[14] = mg_3dtensor.num_index_j; /// J + desc_sizes->data[15] = mg_3dtensor.num_index_k; /// K } else { @@ -2163,7 +2164,7 @@ void read_input_sizes_3D(int32_t fileID, } } -template +template void read_input_3D(int32_t fileID, int32_t A1format, int32_t A1_tile_format, int32_t A2format, int32_t A2_tile_format, @@ -2178,12 +2179,12 @@ void read_input_3D(int32_t fileID, { /// TODO(gkestor): readMode is for future use. - auto *desc_A1pos = static_cast *>(A1pos_ptr); - auto *desc_A1crd = static_cast *>(A1crd_ptr); - auto *desc_A2pos = static_cast *>(A2pos_ptr); - auto *desc_A2crd = static_cast *>(A2crd_ptr); - auto *desc_A3pos = static_cast *>(A3pos_ptr); - auto *desc_A3crd = static_cast *>(A3crd_ptr); + auto *desc_A1pos = static_cast *>(A1pos_ptr); + auto *desc_A1crd = static_cast *>(A1crd_ptr); + auto *desc_A2pos = static_cast *>(A2pos_ptr); + auto *desc_A2crd = static_cast *>(A2crd_ptr); + auto *desc_A3pos = static_cast *>(A3pos_ptr); + auto *desc_A3crd = static_cast *>(A3crd_ptr); auto *desc_Aval = static_cast *>(Aval_ptr); FileReaderWrapper FileReader(fileID, true); /// init of COO_3d_tensor @@ -2350,7 +2351,7 @@ void read_input_3D(int32_t fileID, } /// Utility functions to read sparse matrices and fill in the pos and crd arrays per dimension -extern "C" void read_input_2D_f32(int32_t fileID, +extern "C" void read_input_2D_f32_i32(int32_t fileID, int32_t A1format, int32_t A1_tile_format, int32_t A2format, int32_t A2_tile_format, int A1pos_rank, void *A1pos_ptr, @@ -2364,7 +2365,7 @@ extern "C" void read_input_2D_f32(int32_t fileID, int Aval_rank, void *Aval_ptr, int32_t readMode) { - read_input_2D(fileID, + read_input_2D(fileID, A1format, A1_tile_format, A2format, A2_tile_format, A1pos_rank, A1pos_ptr, A1crd_rank, A1crd_ptr, @@ -2374,7 +2375,8 @@ extern "C" void read_input_2D_f32(int32_t fileID, Aval_rank, Aval_ptr, readMode); } -extern "C" void read_input_2D_f64(int32_t fileID, +/// Utility functions to read sparse matrices and fill in the pos and crd arrays per dimension +extern "C" void read_input_2D_f32_i64(int32_t fileID, int32_t A1format, int32_t A1_tile_format, int32_t A2format, int32_t A2_tile_format, int A1pos_rank, void *A1pos_ptr, @@ -2388,7 +2390,31 @@ extern "C" void read_input_2D_f64(int32_t fileID, int Aval_rank, void *Aval_ptr, int32_t readMode) { - read_input_2D(fileID, + read_input_2D(fileID, + A1format, A1_tile_format, + A2format, A2_tile_format, + A1pos_rank, A1pos_ptr, A1crd_rank, A1crd_ptr, + A1tile_pos_rank, A1tile_pos_ptr, A1tile_crd_rank, A1tile_crd_ptr, + A2pos_rank, A2pos_ptr, A2crd_rank, A2crd_ptr, + A2tile_pos_rank, A2tile_pos_ptr, A2tile_crd_rank, A2tile_crd_ptr, + Aval_rank, Aval_ptr, readMode); +} + +extern "C" void read_input_2D_f64_i32(int32_t fileID, + int32_t A1format, int32_t A1_tile_format, + int32_t A2format, int32_t A2_tile_format, + int A1pos_rank, void *A1pos_ptr, + int A1crd_rank, void *A1crd_ptr, + int A1tile_pos_rank, void *A1tile_pos_ptr, + int A1tile_crd_rank, void *A1tile_crd_ptr, + int A2pos_rank, void *A2pos_ptr, + int A2crd_rank, void *A2crd_ptr, + int A2tile_pos_rank, void *A2tile_pos_ptr, + int A2tile_crd_rank, void *A2tile_crd_ptr, + int Aval_rank, void *Aval_ptr, + int32_t readMode) +{ + read_input_2D(fileID, A1format, A1_tile_format, A2format, A2_tile_format, A1pos_rank, A1pos_ptr, A1crd_rank, A1crd_ptr, @@ -2398,7 +2424,31 @@ extern "C" void read_input_2D_f64(int32_t fileID, Aval_rank, Aval_ptr, readMode); } -extern "C" void read_input_3D_f32(int32_t fileID, +extern "C" void read_input_2D_f64_i64(int32_t fileID, + int32_t A1format, int32_t A1_tile_format, + int32_t A2format, int32_t A2_tile_format, + int A1pos_rank, void *A1pos_ptr, + int A1crd_rank, void *A1crd_ptr, + int A1tile_pos_rank, void *A1tile_pos_ptr, + int A1tile_crd_rank, void *A1tile_crd_ptr, + int A2pos_rank, void *A2pos_ptr, + int A2crd_rank, void *A2crd_ptr, + int A2tile_pos_rank, void *A2tile_pos_ptr, + int A2tile_crd_rank, void *A2tile_crd_ptr, + int Aval_rank, void *Aval_ptr, + int32_t readMode) +{ + read_input_2D(fileID, + A1format, A1_tile_format, + A2format, A2_tile_format, + A1pos_rank, A1pos_ptr, A1crd_rank, A1crd_ptr, + A1tile_pos_rank, A1tile_pos_ptr, A1tile_crd_rank, A1tile_crd_ptr, + A2pos_rank, A2pos_ptr, A2crd_rank, A2crd_ptr, + A2tile_pos_rank, A2tile_pos_ptr, A2tile_crd_rank, A2tile_crd_ptr, + Aval_rank, Aval_ptr, readMode); +} + +extern "C" void read_input_3D_f32_i32(int32_t fileID, int32_t A1format, int32_t A1_tile_format, int32_t A2format, int32_t A2_tile_format, int32_t A3format, int32_t A3_tile_format, @@ -2411,7 +2461,7 @@ extern "C" void read_input_3D_f32(int32_t fileID, int Aval_rank, void *Aval_ptr, int32_t readMode) { - read_input_3D(fileID, + read_input_3D(fileID, A1format, A1_tile_format, A2format, A2_tile_format, A3format, A3_tile_format, @@ -2424,7 +2474,7 @@ extern "C" void read_input_3D_f32(int32_t fileID, Aval_rank, Aval_ptr, readMode); } -extern "C" void read_input_3D_f64(int32_t fileID, +extern "C" void read_input_3D_f32_i64(int32_t fileID, int32_t A1format, int32_t A1_tile_format, int32_t A2format, int32_t A2_tile_format, int32_t A3format, int32_t A3_tile_format, @@ -2436,7 +2486,58 @@ extern "C" void read_input_3D_f64(int32_t fileID, int A3tile_pos_rank, void *A3tile_pos_ptr, int A3tile_crd_rank, void *A3tile_crd_ptr, int Aval_rank, void *Aval_ptr, int32_t readMode) { - read_input_3D(fileID, + + read_input_3D(fileID, + A1format, A1_tile_format, + A2format, A2_tile_format, + A3format, A3_tile_format, + A1pos_rank, A1pos_ptr, A1crd_rank, A1crd_ptr, + A1tile_pos_rank, A1tile_pos_ptr, A1tile_crd_rank, A1tile_crd_ptr, + A2pos_rank, A2pos_ptr, A2crd_rank, A2crd_ptr, + A2tile_pos_rank, A2tile_pos_ptr, A2tile_crd_rank, A2tile_crd_ptr, + A3pos_rank, A3pos_ptr, A3crd_rank, A3crd_ptr, + A3tile_pos_rank, A3tile_pos_ptr, A3tile_crd_rank, A3tile_crd_ptr, + Aval_rank, Aval_ptr, readMode); +} + +extern "C" void read_input_3D_f64_i32(int32_t fileID, + int32_t A1format, int32_t A1_tile_format, + int32_t A2format, int32_t A2_tile_format, + int32_t A3format, int32_t A3_tile_format, + int A1pos_rank, void *A1pos_ptr, int A1crd_rank, void *A1crd_ptr, + int A1tile_pos_rank, void *A1tile_pos_ptr, int A1tile_crd_rank, void *A1tile_crd_ptr, + int A2pos_rank, void *A2pos_ptr, int A2crd_rank, void *A2crd_ptr, + int A2tile_pos_rank, void *A2tile_pos_ptr, int A2tile_crd_rank, void *A2tile_crd_ptr, + int A3pos_rank, void *A3pos_ptr, int A3crd_rank, void *A3crd_ptr, + int A3tile_pos_rank, void *A3tile_pos_ptr, int A3tile_crd_rank, void *A3tile_crd_ptr, + int Aval_rank, void *Aval_ptr, int32_t readMode) +{ + read_input_3D(fileID, + A1format, A1_tile_format, + A2format, A2_tile_format, + A3format, A3_tile_format, + A1pos_rank, A1pos_ptr, A1crd_rank, A1crd_ptr, + A1tile_pos_rank, A1tile_pos_ptr, A1tile_crd_rank, A1tile_crd_ptr, + A2pos_rank, A2pos_ptr, A2crd_rank, A2crd_ptr, + A2tile_pos_rank, A2tile_pos_ptr, A2tile_crd_rank, A2tile_crd_ptr, + A3pos_rank, A3pos_ptr, A3crd_rank, A3crd_ptr, + A3tile_pos_rank, A3tile_pos_ptr, A3tile_crd_rank, A3tile_crd_ptr, + Aval_rank, Aval_ptr, readMode); +} + +extern "C" void read_input_3D_f64_i64(int32_t fileID, + int32_t A1format, int32_t A1_tile_format, + int32_t A2format, int32_t A2_tile_format, + int32_t A3format, int32_t A3_tile_format, + int A1pos_rank, void *A1pos_ptr, int A1crd_rank, void *A1crd_ptr, + int A1tile_pos_rank, void *A1tile_pos_ptr, int A1tile_crd_rank, void *A1tile_crd_ptr, + int A2pos_rank, void *A2pos_ptr, int A2crd_rank, void *A2crd_ptr, + int A2tile_pos_rank, void *A2tile_pos_ptr, int A2tile_crd_rank, void *A2tile_crd_ptr, + int A3pos_rank, void *A3pos_ptr, int A3crd_rank, void *A3crd_ptr, + int A3tile_pos_rank, void *A3tile_pos_ptr, int A3tile_crd_rank, void *A3tile_crd_ptr, + int Aval_rank, void *Aval_ptr, int32_t readMode) +{ + read_input_3D(fileID, A1format, A1_tile_format, A2format, A2_tile_format, A3format, A3_tile_format, @@ -2496,7 +2597,12 @@ extern "C" void read_input_sizes_3D_f64(int32_t fileID, //===----------------------------------------------------------------------===// /// Sort a vector within a range [first, last). //===----------------------------------------------------------------------===// -extern "C" void _milr_ciface_comet_sort(UnrankedMemRefType *M, int64_t index_first, int64_t index_last) +extern "C" void _milr_ciface_comet_sort32(UnrankedMemRefType *M, int64_t index_first, int64_t index_last) +{ + cometSortIndex(*M, index_first, index_last); +} + +extern "C" void _milr_ciface_comet_sort64(UnrankedMemRefType *M, int64_t index_first, int64_t index_last) { cometSortIndex(*M, index_first, index_last); } @@ -2504,5 +2610,17 @@ extern "C" void _milr_ciface_comet_sort(UnrankedMemRefType *M, int64_t extern "C" void comet_sort_index(int64_t rank, void *ptr, int64_t index_first, int64_t index_last) { UnrankedMemRefType descriptor = {rank, ptr}; - _milr_ciface_comet_sort(&descriptor, index_first, index_last); + _milr_ciface_comet_sort64(&descriptor, index_first, index_last); +} + +extern "C" void comet_sort32(int64_t rank, void *ptr, int64_t index_first, int64_t index_last) +{ + UnrankedMemRefType descriptor = {rank, ptr}; + _milr_ciface_comet_sort32(&descriptor, index_first, index_last); +} + +extern "C" void comet_sort64(int64_t rank, void *ptr, int64_t index_first, int64_t index_last) +{ + UnrankedMemRefType descriptor = {rank, ptr}; + _milr_ciface_comet_sort64(&descriptor, index_first, index_last); } \ No newline at end of file diff --git a/lib/ExecutionEngine/StatUtils.cpp b/lib/ExecutionEngine/StatUtils.cpp index 33d5bc3f..3823e688 100644 --- a/lib/ExecutionEngine/StatUtils.cpp +++ b/lib/ExecutionEngine/StatUtils.cpp @@ -91,23 +91,51 @@ extern "C" void _mlir_ciface_comet_print_memref_f64(UnrankedMemRefType * cometPrintMemRef(*M); } +extern "C" void _mlir_ciface_comet_print_memref_f32(UnrankedMemRefType *M) +{ + cometPrintMemRef(*M); +} + extern "C" void _mlir_ciface_comet_print_memref_i64(UnrankedMemRefType *M) { cometPrintMemRef(*M); } +extern "C" void _mlir_ciface_comet_print_memref_i32(UnrankedMemRefType *M) +{ + cometPrintMemRef(*M); +} + extern "C" void comet_print_memref_f64(int64_t rank, void *ptr) { UnrankedMemRefType descriptor = {rank, ptr}; _mlir_ciface_comet_print_memref_f64(&descriptor); } +extern "C" void comet_print_memref_f32(int64_t rank, void *ptr) +{ + UnrankedMemRefType descriptor = {rank, ptr}; + _mlir_ciface_comet_print_memref_f32(&descriptor); +} + extern "C" void comet_print_memref_i64(int64_t rank, void *ptr) { UnrankedMemRefType descriptor = {rank, ptr}; _mlir_ciface_comet_print_memref_i64(&descriptor); } +extern "C" void comet_print_memref_i32(int64_t rank, void *ptr) +{ + UnrankedMemRefType descriptor = {rank, ptr}; + _mlir_ciface_comet_print_memref_i32(&descriptor); +} + +extern "C" void comet_print_memref_index(int64_t rank, void *ptr) +{ + UnrankedMemRefType descriptor = {rank, ptr}; + _mlir_ciface_comet_print_memref_i64(&descriptor); +} + //===----------------------------------------------------------------------===// /// Small runtime support library for memset //===----------------------------------------------------------------------===// diff --git a/lib/ExecutionEngine/TransposeUtils.cpp b/lib/ExecutionEngine/TransposeUtils.cpp index a6e4281a..8889a8e2 100644 --- a/lib/ExecutionEngine/TransposeUtils.cpp +++ b/lib/ExecutionEngine/TransposeUtils.cpp @@ -30,6 +30,7 @@ #include "comet/ExecutionEngine/RunnerUtils.h" #include "llvm/Support/raw_ostream.h" +#include #include #include #include @@ -478,7 +479,7 @@ void transpose_sort(int sort_type, vector &coo_ts, int sz, * @param sizes_rank ? * @param sizes_ptr array of length of those arrays */ -template +template void transpose_2D(int32_t A1format, int32_t A1tile_format, int32_t A2format, int32_t A2tile_format, int A1pos_rank, void *A1pos_ptr, int A1crd_rank, void *A1crd_ptr, int A1tile_pos_rank, void *A1tile_pos_ptr, int A1tile_crd_rank, void *A1tile_crd_ptr, @@ -502,28 +503,28 @@ void transpose_2D(int32_t A1format, int32_t A1tile_format, int32_t A2format, int std::string Aspformat; std::string Bspformat; - auto *desc_A1pos = static_cast *>(A1pos_ptr); - auto *desc_A1crd = static_cast *>(A1crd_ptr); - auto *desc_A1tile_pos = static_cast *>(A1tile_pos_ptr); - auto *desc_A1tile_crd = static_cast *>(A1tile_crd_ptr); - auto *desc_A2pos = static_cast *>(A2pos_ptr); - auto *desc_A2crd = static_cast *>(A2crd_ptr); - [[maybe_unused]] auto *desc_A2tile_pos = static_cast *>(A2tile_pos_ptr); - [[maybe_unused]] auto *desc_A2tile_crd = static_cast *>(A2tile_crd_ptr); + auto *desc_A1pos = static_cast *>(A1pos_ptr); + auto *desc_A1crd = static_cast *>(A1crd_ptr); + auto *desc_A1tile_pos = static_cast *>(A1tile_pos_ptr); + auto *desc_A1tile_crd = static_cast *>(A1tile_crd_ptr); + auto *desc_A2pos = static_cast *>(A2pos_ptr); + auto *desc_A2crd = static_cast *>(A2crd_ptr); + [[maybe_unused]] auto *desc_A2tile_pos = static_cast *>(A2tile_pos_ptr); + [[maybe_unused]] auto *desc_A2tile_crd = static_cast *>(A2tile_crd_ptr); auto *desc_Aval = static_cast *>(Aval_ptr); - auto *desc_B1pos = static_cast *>(B1pos_ptr); - auto *desc_B1crd = static_cast *>(B1crd_ptr); - auto *desc_B2pos = static_cast *>(B2pos_ptr); - auto *desc_B2crd = static_cast *>(B2crd_ptr); + auto *desc_B1pos = static_cast *>(B1pos_ptr); + auto *desc_B1crd = static_cast *>(B1crd_ptr); + auto *desc_B2pos = static_cast *>(B2pos_ptr); + auto *desc_B2crd = static_cast *>(B2crd_ptr); auto *desc_Bval = static_cast *>(Bval_ptr); - auto *desc_sizes = static_cast *>(sizes_ptr); + auto *dim_sizes = static_cast *>(sizes_ptr); - auto *desc_B1tile_pos = static_cast *>(B1tile_pos_ptr); - auto *desc_B1tile_crd = static_cast *>(B1tile_crd_ptr); - auto *desc_B2tile_pos = static_cast *>(B2tile_pos_ptr); - auto *desc_B2tile_crd = static_cast *>(B2tile_crd_ptr); + auto *desc_B1tile_pos = static_cast *>(B1tile_pos_ptr); + auto *desc_B1tile_crd = static_cast *>(B1tile_crd_ptr); + auto *desc_B2tile_pos = static_cast *>(B2tile_pos_ptr); + auto *desc_B2tile_crd = static_cast *>(B2tile_crd_ptr); desc_B1tile_pos->data[0] = -1; desc_B1tile_crd->data[0] = -1; desc_B2tile_pos->data[0] = -1; @@ -549,8 +550,8 @@ void transpose_2D(int32_t A1format, int32_t A1tile_format, int32_t A2format, int << "desc_sizes->data[10]: " << desc_sizes->data[10] << "\n"; */ - int rowSize = desc_sizes->data[9]; - int colSize = desc_sizes->data[10]; + int rowSize = dim_sizes->data[0]; + int colSize = dim_sizes->data[1]; if (A1format == Dense && A1tile_format == Dense && A2format == singleton) { @@ -691,8 +692,8 @@ void transpose_2D(int32_t A1format, int32_t A1tile_format, int32_t A2format, int desc_B2pos->data[0] = desc_A2pos->data[0]; /// switch row and col size - desc_sizes->data[9] = colSize; - desc_sizes->data[10] = rowSize; + dim_sizes->data[0] = colSize; + dim_sizes->data[1] = rowSize; } if (Aspformat.compare("CSR") == 0 && Bspformat.compare("CSR") == 0) @@ -702,8 +703,8 @@ void transpose_2D(int32_t A1format, int32_t A1tile_format, int32_t A2format, int /// 1) not by sorting: only works for CSR/matrices /// Atomic-based Transposition: retraverse the matrix from the transposed direction /// B's row size == input's col size - int BRowSize = desc_sizes->data[9]; - int BColSize = desc_sizes->data[10]; + int BRowSize = dim_sizes->data[0]; + int BColSize = dim_sizes->data[1]; /// B's col pos size == B's #rows + 1 desc_B2pos->sizes[0] = BRowSize + 1; @@ -732,8 +733,8 @@ void transpose_2D(int32_t A1format, int32_t A1tile_format, int32_t A2format, int desc_B1crd->data[0] = -1; /// switch row and col size - desc_sizes->data[9] = colSize; - desc_sizes->data[10] = rowSize; + dim_sizes->data[0] = colSize; + dim_sizes->data[1] = rowSize; } else { @@ -807,8 +808,8 @@ void transpose_2D(int32_t A1format, int32_t A1tile_format, int32_t A2format, int } /// switch row and col size - desc_sizes->data[9] = colSize; - desc_sizes->data[10] = rowSize; + dim_sizes->data[0] = colSize; + dim_sizes->data[1] = rowSize; } } @@ -823,7 +824,7 @@ void transpose_2D(int32_t A1format, int32_t A1tile_format, int32_t A2format, int } } -template +template void transpose_3D(int32_t input_permutation, int32_t output_permutation, int32_t A1format, int32_t A1tile_format, int32_t A2format, int32_t A2tile_format, @@ -889,29 +890,29 @@ void transpose_3D(int32_t input_permutation, int32_t output_permutation, } /// auto *desc_A1pos = static_cast *>(A1pos_ptr); - auto *desc_A1crd = static_cast *>(A1crd_ptr); - auto *desc_A2pos = static_cast *>(A2pos_ptr); - auto *desc_A2crd = static_cast *>(A2crd_ptr); - auto *desc_A3pos = static_cast *>(A3pos_ptr); - auto *desc_A3crd = static_cast *>(A3crd_ptr); + auto *desc_A1crd = static_cast *>(A1crd_ptr); + auto *desc_A2pos = static_cast *>(A2pos_ptr); + auto *desc_A2crd = static_cast *>(A2crd_ptr); + auto *desc_A3pos = static_cast *>(A3pos_ptr); + auto *desc_A3crd = static_cast *>(A3crd_ptr); auto *desc_Aval = static_cast *>(Aval_ptr); - auto *desc_B1pos = static_cast *>(B1pos_ptr); - auto *desc_B1crd = static_cast *>(B1crd_ptr); - auto *desc_B2pos = static_cast *>(B2pos_ptr); - auto *desc_B2crd = static_cast *>(B2crd_ptr); - auto *desc_B3pos = static_cast *>(B3pos_ptr); - auto *desc_B3crd = static_cast *>(B3crd_ptr); + auto *desc_B1pos = static_cast *>(B1pos_ptr); + auto *desc_B1crd = static_cast *>(B1crd_ptr); + auto *desc_B2pos = static_cast *>(B2pos_ptr); + auto *desc_B2crd = static_cast *>(B2crd_ptr); + auto *desc_B3pos = static_cast *>(B3pos_ptr); + auto *desc_B3crd = static_cast *>(B3crd_ptr); auto *desc_Bval = static_cast *>(Bval_ptr); - auto *desc_sizes = static_cast *>(sizes_ptr); + auto *desc_sizes = static_cast *>(sizes_ptr); - auto *desc_B1tile_pos = static_cast *>(B1tile_pos_ptr); - auto *desc_B1tile_crd = static_cast *>(B1tile_crd_ptr); - auto *desc_B2tile_pos = static_cast *>(B2tile_pos_ptr); - auto *desc_B2tile_crd = static_cast *>(B2tile_crd_ptr); - auto *desc_B3tile_pos = static_cast *>(B3tile_pos_ptr); - auto *desc_B3tile_crd = static_cast *>(B3tile_crd_ptr); + auto *desc_B1tile_pos = static_cast *>(B1tile_pos_ptr); + auto *desc_B1tile_crd = static_cast *>(B1tile_crd_ptr); + auto *desc_B2tile_pos = static_cast *>(B2tile_pos_ptr); + auto *desc_B2tile_crd = static_cast *>(B2tile_crd_ptr); + auto *desc_B3tile_pos = static_cast *>(B3tile_pos_ptr); + auto *desc_B3tile_crd = static_cast *>(B3tile_crd_ptr); desc_B1tile_pos->data[0] = -1; desc_B1tile_crd->data[0] = -1; desc_B2tile_pos->data[0] = -1; @@ -1347,7 +1348,7 @@ void transpose_3D(int32_t input_permutation, int32_t output_permutation, } /// 2D tensors -extern "C" void transpose_2D_f32(int32_t A1format, int32_t A1tile_format, int32_t A2format, int32_t A2tile_format, +extern "C" void transpose_2D_f32_i32(int32_t A1format, int32_t A1tile_format, int32_t A2format, int32_t A2tile_format, int A1pos_rank, void *A1pos_ptr, int A1crd_rank, void *A1crd_ptr, int A1tile_pos_rank, void *A1tile_pos_ptr, int A1tile_crd_rank, void *A1tile_crd_ptr, int A2pos_rank, void *A2pos_ptr, int A2crd_rank, void *A2crd_ptr, @@ -1361,7 +1362,7 @@ extern "C" void transpose_2D_f32(int32_t A1format, int32_t A1tile_format, int32_ int Bval_rank, void *Bval_ptr, int sizes_rank, void *sizes_ptr) { - transpose_2D(A1format, A1tile_format, A2format, A2tile_format, + transpose_2D(A1format, A1tile_format, A2format, A2tile_format, A1pos_rank, A1pos_ptr, A1crd_rank, A1crd_ptr, A1tile_pos_rank, A1tile_pos_ptr, A1tile_crd_rank, A1tile_crd_ptr, A2pos_rank, A2pos_ptr, A2crd_rank, A2crd_ptr, @@ -1376,7 +1377,66 @@ extern "C" void transpose_2D_f32(int32_t A1format, int32_t A1tile_format, int32_ sizes_rank, sizes_ptr); } -extern "C" void transpose_2D_f64(int32_t A1format, int32_t A1tile_format, int32_t A2format, int32_t A2tile_format, +/// 2D tensors +extern "C" void transpose_2D_f32_i64(int32_t A1format, int32_t A1tile_format, int32_t A2format, int32_t A2tile_format, + int A1pos_rank, void *A1pos_ptr, int A1crd_rank, void *A1crd_ptr, + int A1tile_pos_rank, void *A1tile_pos_ptr, int A1tile_crd_rank, void *A1tile_crd_ptr, + int A2pos_rank, void *A2pos_ptr, int A2crd_rank, void *A2crd_ptr, + int A2tile_pos_rank, void *A2tile_pos_ptr, int A2tile_crd_rank, void *A2tile_crd_ptr, + int Aval_rank, void *Aval_ptr, + int32_t B1format, int32_t B1tile_format, int32_t B2format, int32_t B2tile_format, + int B1pos_rank, void *B1pos_ptr, int B1crd_rank, void *B1crd_ptr, + int B1tile_pos_rank, void *B1tile_pos_ptr, int B1tile_crd_rank, void *B1tile_crd_ptr, + int B2pos_rank, void *B2pos_ptr, int B2crd_rank, void *B2crd_ptr, + int B2tile_pos_rank, void *B2tile_pos_ptr, int B2tile_crd_rank, void *B2tile_crd_ptr, + int Bval_rank, void *Bval_ptr, + int sizes_rank, void *sizes_ptr) +{ + transpose_2D(A1format, A1tile_format, A2format, A2tile_format, + A1pos_rank, A1pos_ptr, A1crd_rank, A1crd_ptr, + A1tile_pos_rank, A1tile_pos_ptr, A1tile_crd_rank, A1tile_crd_ptr, + A2pos_rank, A2pos_ptr, A2crd_rank, A2crd_ptr, + A2tile_pos_rank, A2tile_pos_ptr, A2tile_crd_rank, A2tile_crd_ptr, + Aval_rank, Aval_ptr, + B1format, B1tile_format, B2format, B2tile_format, + B1pos_rank, B1pos_ptr, B1crd_rank, B1crd_ptr, + B1tile_pos_rank, B1tile_pos_ptr, B1tile_crd_rank, B1tile_crd_ptr, + B2pos_rank, B2pos_ptr, B2crd_rank, B2crd_ptr, + B2tile_pos_rank, B2tile_pos_ptr, B2tile_crd_rank, B2tile_crd_ptr, + Bval_rank, Bval_ptr, + sizes_rank, sizes_ptr); +} + +extern "C" void transpose_2D_f64_i32(int32_t A1format, int32_t A1tile_format, int32_t A2format, int32_t A2tile_format, + int A1pos_rank, void *A1pos_ptr, int A1crd_rank, void *A1crd_ptr, + int A1tile_pos_rank, void *A1tile_pos_ptr, int A1tile_crd_rank, void *A1tile_crd_ptr, + int A2pos_rank, void *A2pos_ptr, int A2crd_rank, void *A2crd_ptr, + int A2tile_pos_rank, void *A2tile_pos_ptr, int A2tile_crd_rank, void *A2tile_crd_ptr, + int Aval_rank, void *Aval_ptr, + int32_t B1format, int32_t B1tile_format, int32_t B2format, int32_t B2tile_format, + int B1pos_rank, void *B1pos_ptr, int B1crd_rank, void *B1crd_ptr, + int B1tile_pos_rank, void *B1tile_pos_ptr, int B1tile_crd_rank, void *B1tile_crd_ptr, + int B2pos_rank, void *B2pos_ptr, int B2crd_rank, void *B2crd_ptr, + int B2tile_pos_rank, void *B2tile_pos_ptr, int B2tile_crd_rank, void *B2tile_crd_ptr, + int Bval_rank, void *Bval_ptr, + int sizes_rank, void *sizes_ptr) +{ + transpose_2D(A1format, A1tile_format, A2format, A2tile_format, + A1pos_rank, A1pos_ptr, A1crd_rank, A1crd_ptr, + A1tile_pos_rank, A1tile_pos_ptr, A1tile_crd_rank, A1tile_crd_ptr, + A2pos_rank, A2pos_ptr, A2crd_rank, A2crd_ptr, + A2tile_pos_rank, A2tile_pos_ptr, A2tile_crd_rank, A2tile_crd_ptr, + Aval_rank, Aval_ptr, + B1format, B1tile_format, B2format, B2tile_format, + B1pos_rank, B1pos_ptr, B1crd_rank, B1crd_ptr, + B1tile_pos_rank, B1tile_pos_ptr, B1tile_crd_rank, B1tile_crd_ptr, + B2pos_rank, B2pos_ptr, B2crd_rank, B2crd_ptr, + B2tile_pos_rank, B2tile_pos_ptr, B2tile_crd_rank, B2tile_crd_ptr, + Bval_rank, Bval_ptr, + sizes_rank, sizes_ptr); +} + +extern "C" void transpose_2D_f64_i64(int32_t A1format, int32_t A1tile_format, int32_t A2format, int32_t A2tile_format, int A1pos_rank, void *A1pos_ptr, int A1crd_rank, void *A1crd_ptr, int A1tile_pos_rank, void *A1tile_pos_ptr, int A1tile_crd_rank, void *A1tile_crd_ptr, int A2pos_rank, void *A2pos_ptr, int A2crd_rank, void *A2crd_ptr, @@ -1390,7 +1450,7 @@ extern "C" void transpose_2D_f64(int32_t A1format, int32_t A1tile_format, int32_ int Bval_rank, void *Bval_ptr, int sizes_rank, void *sizes_ptr) { - transpose_2D(A1format, A1tile_format, A2format, A2tile_format, + transpose_2D(A1format, A1tile_format, A2format, A2tile_format, A1pos_rank, A1pos_ptr, A1crd_rank, A1crd_ptr, A1tile_pos_rank, A1tile_pos_ptr, A1tile_crd_rank, A1tile_crd_ptr, A2pos_rank, A2pos_ptr, A2crd_rank, A2crd_ptr, @@ -1406,7 +1466,100 @@ extern "C" void transpose_2D_f64(int32_t A1format, int32_t A1tile_format, int32_ } /// 3D tensors -extern "C" void transpose_3D_f32(int32_t input_permutation, int32_t output_permutation, +extern "C" void transpose_3D_f32_i32(int32_t input_permutation, int32_t output_permutation, + int32_t A1format, int32_t A1tile_format, + int32_t A2format, int32_t A2tile_format, + int32_t A3format, int32_t A3tile_format, + int A1pos_rank, void *A1pos_ptr, int A1crd_rank, void *A1crd_ptr, + int A1tile_pos_rank, void *A1tile_pos_ptr, int A1tile_crd_rank, void *A1tile_crd_ptr, + int A2pos_rank, void *A2pos_ptr, int A2crd_rank, void *A2crd_ptr, + int A2tile_pos_rank, void *A2tile_pos_ptr, int A2tile_crd_rank, void *A2tile_crd_ptr, + int A3pos_rank, void *A3pos_ptr, int A3crd_rank, void *A3crd_ptr, + int A3tile_pos_rank, void *A3tile_pos_ptr, int A3tile_crd_rank, void *A3tile_crd_ptr, + int Aval_rank, void *Aval_ptr, + int32_t B1format, int32_t B1tile_format, + int32_t B2format, int32_t B2tile_format, + int32_t B3format, int32_t B3tile_format, + int B1pos_rank, void *B1pos_ptr, int B1crd_rank, void *B1crd_ptr, + int B1tile_pos_rank, void *B1tile_pos_ptr, int B1tile_crd_rank, void *B1tile_crd_ptr, + int B2pos_rank, void *B2pos_ptr, int B2crd_rank, void *B2crd_ptr, + int B2tile_pos_rank, void *B2tile_pos_ptr, int B2tile_crd_rank, void *B2tile_crd_ptr, + int B3pos_rank, void *B3pos_ptr, int B3crd_rank, void *B3crd_ptr, + int B3tile_pos_rank, void *B3tile_pos_ptr, int B3tile_crd_rank, void *B3tile_crd_ptr, + int Bval_rank, void *Bval_ptr, int sizes_rank, void *sizes_ptr) +{ + transpose_3D(input_permutation, output_permutation, + A1format, A1tile_format, + A2format, A2tile_format, + A3format, A3tile_format, + A1pos_rank, A1pos_ptr, A1crd_rank, A1crd_ptr, + A1tile_pos_rank, A1tile_pos_ptr, A1tile_crd_rank, A1tile_crd_ptr, + A2pos_rank, A2pos_ptr, A2crd_rank, A2crd_ptr, + A2tile_pos_rank, A2tile_pos_ptr, A2tile_crd_rank, A2tile_crd_ptr, + A3pos_rank, A3pos_ptr, A3crd_rank, A3crd_ptr, + A3tile_pos_rank, A3tile_pos_ptr, A3tile_crd_rank, A3tile_crd_ptr, + Aval_rank, Aval_ptr, + B1format, B1tile_format, + B2format, B2tile_format, + B3format, B3tile_format, + B1pos_rank, B1pos_ptr, B1crd_rank, B1crd_ptr, + B1tile_pos_rank, B1tile_pos_ptr, B1tile_crd_rank, B1tile_crd_ptr, + B2pos_rank, B2pos_ptr, B2crd_rank, B2crd_ptr, + B2tile_pos_rank, B2tile_pos_ptr, B2tile_crd_rank, B2tile_crd_ptr, + B3pos_rank, B3pos_ptr, B3crd_rank, B3crd_ptr, + B3tile_pos_rank, B3tile_pos_ptr, B3tile_crd_rank, B3tile_crd_ptr, + Bval_rank, Bval_ptr, + sizes_rank, sizes_ptr); +} + +extern "C" void transpose_3D_f64_i64(int32_t input_permutation, int32_t output_permutation, + int32_t A1format, int32_t A1tile_format, + int32_t A2format, int32_t A2tile_format, + int32_t A3format, int32_t A3tile_format, + int A1pos_rank, void *A1pos_ptr, int A1crd_rank, void *A1crd_ptr, + int A1tile_pos_rank, void *A1tile_pos_ptr, int A1tile_crd_rank, void *A1tile_crd_ptr, + int A2pos_rank, void *A2pos_ptr, int A2crd_rank, void *A2crd_ptr, + int A2tile_pos_rank, void *A2tile_pos_ptr, int A2tile_crd_rank, void *A2tile_crd_ptr, + int A3pos_rank, void *A3pos_ptr, int A3crd_rank, void *A3crd_ptr, + int A3tile_pos_rank, void *A3tile_pos_ptr, int A3tile_crd_rank, void *A3tile_crd_ptr, + int Aval_rank, void *Aval_ptr, + int32_t B1format, int32_t B1tile_format, + int32_t B2format, int32_t B2tile_format, + int32_t B3format, int32_t B3tile_format, + int B1pos_rank, void *B1pos_ptr, int B1crd_rank, void *B1crd_ptr, + int B1tile_pos_rank, void *B1tile_pos_ptr, int B1tile_crd_rank, void *B1tile_crd_ptr, + int B2pos_rank, void *B2pos_ptr, int B2crd_rank, void *B2crd_ptr, + int B2tile_pos_rank, void *B2tile_pos_ptr, int B2tile_crd_rank, void *B2tile_crd_ptr, + int B3pos_rank, void *B3pos_ptr, int B3crd_rank, void *B3crd_ptr, + int B3tile_pos_rank, void *B3tile_pos_ptr, int B3tile_crd_rank, void *B3tile_crd_ptr, + int Bval_rank, void *Bval_ptr, int sizes_rank, void *sizes_ptr) +{ + transpose_3D(input_permutation, output_permutation, + A1format, A1tile_format, + A2format, A2tile_format, + A3format, A3tile_format, + A1pos_rank, A1pos_ptr, A1crd_rank, A1crd_ptr, + A1tile_pos_rank, A1tile_pos_ptr, A1tile_crd_rank, A1tile_crd_ptr, + A2pos_rank, A2pos_ptr, A2crd_rank, A2crd_ptr, + A2tile_pos_rank, A2tile_pos_ptr, A2tile_crd_rank, A2tile_crd_ptr, + A3pos_rank, A3pos_ptr, A3crd_rank, A3crd_ptr, + A3tile_pos_rank, A3tile_pos_ptr, A3tile_crd_rank, A3tile_crd_ptr, + Aval_rank, Aval_ptr, + B1format, B1tile_format, + B2format, B2tile_format, + B3format, B3tile_format, + B1pos_rank, B1pos_ptr, B1crd_rank, B1crd_ptr, + B1tile_pos_rank, B1tile_pos_ptr, B1tile_crd_rank, B1tile_crd_ptr, + B2pos_rank, B2pos_ptr, B2crd_rank, B2crd_ptr, + B2tile_pos_rank, B2tile_pos_ptr, B2tile_crd_rank, B2tile_crd_ptr, + B3pos_rank, B3pos_ptr, B3crd_rank, B3crd_ptr, + B3tile_pos_rank, B3tile_pos_ptr, B3tile_crd_rank, B3tile_crd_ptr, + Bval_rank, Bval_ptr, + sizes_rank, sizes_ptr); +} + +/// 3D tensors +extern "C" void transpose_3D_f32_i64(int32_t input_permutation, int32_t output_permutation, int32_t A1format, int32_t A1tile_format, int32_t A2format, int32_t A2tile_format, int32_t A3format, int32_t A3tile_format, @@ -1428,7 +1581,7 @@ extern "C" void transpose_3D_f32(int32_t input_permutation, int32_t output_permu int B3tile_pos_rank, void *B3tile_pos_ptr, int B3tile_crd_rank, void *B3tile_crd_ptr, int Bval_rank, void *Bval_ptr, int sizes_rank, void *sizes_ptr) { - transpose_3D(input_permutation, output_permutation, + transpose_3D(input_permutation, output_permutation, A1format, A1tile_format, A2format, A2tile_format, A3format, A3tile_format, @@ -1452,7 +1605,7 @@ extern "C" void transpose_3D_f32(int32_t input_permutation, int32_t output_permu sizes_rank, sizes_ptr); } -extern "C" void transpose_3D_f64(int32_t input_permutation, int32_t output_permutation, +extern "C" void transpose_3D_f64_i32(int32_t input_permutation, int32_t output_permutation, int32_t A1format, int32_t A1tile_format, int32_t A2format, int32_t A2tile_format, int32_t A3format, int32_t A3tile_format, @@ -1474,7 +1627,7 @@ extern "C" void transpose_3D_f64(int32_t input_permutation, int32_t output_permu int B3tile_pos_rank, void *B3tile_pos_ptr, int B3tile_crd_rank, void *B3tile_crd_ptr, int Bval_rank, void *Bval_ptr, int sizes_rank, void *sizes_ptr) { - transpose_3D(input_permutation, output_permutation, + transpose_3D(input_permutation, output_permutation, A1format, A1tile_format, A2format, A2tile_format, A3format, A3tile_format, diff --git a/lib/ExecutionEngine/blis_interface.cpp b/lib/ExecutionEngine/blis_interface.cpp index 38aee759..74e80c12 100644 --- a/lib/ExecutionEngine/blis_interface.cpp +++ b/lib/ExecutionEngine/blis_interface.cpp @@ -32,66 +32,30 @@ #include #if defined(__x86_64__) || defined(_M_X64) || defined(__i386) || defined(_M_IX86) -void bli_dgemm_x86_ukr( - dim_t m, - dim_t n, - dim_t k, - double *restrict alpha, - double *restrict a, - double *restrict b, - double *restrict beta, - double *restrict c, inc_t rs_c0, inc_t cs_c0, - auxinfo_t *restrict data, - cntx_t *restrict cntx) -{ - /// get the micro - arch - const char *arch = bli_arch_string(bli_cpuid_query_id()); - // printf("arch: %s\n", arch); - - if ((strcmp("haswell", arch) == 0) || - (strcmp("zen", arch) == 0) || - (strcmp("zen2", arch) == 0) || - (strcmp("zen3", arch) == 0) || - (strcmp("skx", arch) == 0) || - (strcmp("knl", arch) == 0)) - { - printf("Calling bli_dgemm_haswell_asm_6x8\n"); - bli_dgemm_haswell_asm_6x8(m, n, k, alpha, a, b, beta, c, rs_c0, cs_c0, data, cntx); - } - else - { - llvm::errs() << __FILE__ << " " << __LINE__ << "ERROR: Undefined microkernel" - << "\n"; - } -} +#define bli_dgemm_x86_ukr(\ + m,\ + n,\ + k,\ + alpha,\ + a,\ + b,\ + beta,\ + c, rs_c0, cs_c0,\ + data,\ + cntx) bli_dgemm_haswell_asm_6x8(m, n, k, alpha, a, b, beta, c, rs_c0, cs_c0, data, cntx); #elif defined(__aarch64__) || defined(__arm__) || defined(_M_ARM) || defined(_ARCH_PPC) -void bli_dgemm_arm_ukr( - dim_t m, - dim_t n, - dim_t k, - double *restrict alpha, - double *restrict a, - double *restrict b, - double *restrict beta, - double *restrict c, inc_t rs_c0, inc_t cs_c0, - auxinfo_t *restrict data, - cntx_t *restrict cntx) -{ - // get the micro - arch - const char *arch = bli_arch_string(bli_cpuid_query_id()); - // printf("arch: %s\n", arch); - - if ((strcmp("firestorm", arch) == 0)) - { - bli_dgemm_armv8a_asm_8x6r(m, n, k, alpha, a, b, beta, c, rs_c0, cs_c0, data, cntx); - } - else - { - llvm::errs() << __FILE__ << " " << __LINE__ << "Undefined microkernel" - << "\n"; - } -} +#define bli_dgemm_arm_ukr(\ + m,\ + n,\ + k,\ + alpha,\ + a,\ + b,\ + beta,\ + c, rs_c0, cs_c0,\ + data,\ + cntx) bli_dgemm_armv8a_asm_6x8(m, n, k, alpha, a, b, beta, c, rs_c0, cs_c0, data, cntx); #endif /// generic arch-independent gemm microkernel reference implementation: @@ -116,10 +80,10 @@ void dgemm_generic_noopt_mxn( int64_t m, int64_t n, int64_t k, - double *alpha, - double *a, double *b, - double *beta, - double *c, + double *restrict alpha, + double *restrict a, double *restrict b, + double *restrict beta, + double *restrict c, int64_t rs_c, int64_t cs_c, auxinfo_t *restrict data, cntx_t *restrict cntx) @@ -177,49 +141,9 @@ void dgemm_generic_noopt_mxn( }; extern "C" void _mlir_ciface_linalg_matmul_viewsxs_viewsxs_viewsxs( - StridedMemRefType *A, StridedMemRefType *B, - StridedMemRefType *C, int mr, int nr) + StridedMemRefType *restrict A, StridedMemRefType *restrict B, + StridedMemRefType *restrict C, int mr, int nr, double alpha, double beta) { - if (A->strides[1] != B->strides[1] || A->strides[1] != C->strides[1] || - A->strides[1] != 1 || A->sizes[0] < A->strides[1] || - B->sizes[0] < B->strides[1] || C->sizes[0] < C->strides[1] || - C->sizes[0] != A->sizes[0] || C->sizes[1] != B->sizes[1] || - A->sizes[1] != B->sizes[0]) - { - printMemRefMetaData(std::cerr, *A); - printMemRefMetaData(std::cerr, *B); - printMemRefMetaData(std::cerr, *C); - - return; - } - - // printMemRefMetaData(std::cerr, *A); - // printMemRefMetaData(std::cerr, *B); - // printMemRefMetaData(std::cerr, *C); - - // printf("\n"); - // printf("A->sizes[0]-m: %d\n", A->sizes[0]); - // printf("A->sizes[1]-k: %d\n", A->sizes[1]); - // printf("A->strides[0]: %d\n", A->strides[0]); - // printf("A->strides[1]: %d\n", A->strides[1]); - - // printf("B->sizes[0]: %d\n", B->sizes[0]); - // printf("B->sizes[1]-n: %d\n", B->sizes[1]); - // printf("B->strides[0]: %d\n", B->strides[0]); - // printf("B->strides[1]: %d\n", B->strides[1]); - // printf("\n"); - - // printf("mr: %d", mr); - // printf("nr: %d", nr); - - double alpha = 1.0f; - double beta = 1.0f; - if (beta == -1.0) - { - alpha *= -1.0; - beta = 1.0; - } - auxinfo_t data; bli_auxinfo_set_next_a(A->data + A->offset, &data); bli_auxinfo_set_next_b(B->data + B->offset, &data); @@ -227,10 +151,6 @@ extern "C" void _mlir_ciface_linalg_matmul_viewsxs_viewsxs_viewsxs( /// Partial tile if (A->sizes[0] < mr || B->sizes[1] < nr) { - // printf("A->sizes[0]: %d\n", A->sizes[0]); - // printf("B->sizes[1]: %d\n", B->sizes[1]); - // printf("mr: %d\n", mr); - // printf("nr: %d\n", nr); dgemm_generic_noopt_mxn(A->sizes[0], // m B->sizes[1], // n A->sizes[1], // k diff --git a/llvm b/llvm index f22cde10..61f8a7f6 160000 --- a/llvm +++ b/llvm @@ -1 +1 @@ -Subproject commit f22cde10e7cc711bba9f43d7529ea6c1394c5b48 +Subproject commit 61f8a7f618901797ee8663389a29722f29216a96 diff --git a/runtimes/mcl b/runtimes/mcl new file mode 160000 index 00000000..271562f1 --- /dev/null +++ b/runtimes/mcl @@ -0,0 +1 @@ +Subproject commit 271562f15faff9c2c0b1058d395ae710a6577596 diff --git a/integration_test/CMakeLists.txt b/test/integration/CMakeLists.txt similarity index 83% rename from integration_test/CMakeLists.txt rename to test/integration/CMakeLists.txt index 76c22d91..d7b60b20 100644 --- a/integration_test/CMakeLists.txt +++ b/test/integration/CMakeLists.txt @@ -19,12 +19,20 @@ message(STATUS "Using COMET_UTILITY_LIBRARIES: ${COMET_UTILITY_LIBRARY_DIR}") set(COMET_INTEGRATION_TEST_DATA_DIR ${CMAKE_CURRENT_SOURCE_DIR}/data/) message(STATUS "Using COMET_INTEGRATION_TEST_DATA_DIR: ${COMET_INTEGRATION_TEST_DATA_DIR}") -if(ENABLE_GPU_TARGET) -set(COMET_ENABLE_GPU True) -message(STATUS "Using COMET_ENABLE_GPU: ${COMET_ENABLE_GPU}") -else() set(COMET_ENABLE_GPU False) -message(STATUS "Using COMET_ENABLE_GPU: ${COMET_ENABLE_GPU}") + +if(ENABLE_AMD_GPU_BACKEND) + find_package(hip) + if(${hip_FOUND}) + set(COMET_ENABLE_GPU True) + endif() +endif() + +if(ENABLE_NVIDIA_GPU_BACKEND) + find_package(CUDAToolkit) + if(${CUDAToolkit_FOUND}) + set(COMET_ENABLE_GPU True) + endif() endif() configure_lit_site_cfg( diff --git a/integration_test/compound_exps/CSR_Dense_chain_mult_matrix.ta b/test/integration/compound_exps/CSR_Dense_chain_mult_matrix.ta similarity index 92% rename from integration_test/compound_exps/CSR_Dense_chain_mult_matrix.ta rename to test/integration/compound_exps/CSR_Dense_chain_mult_matrix.ta index 55b9a6b1..3c499b0c 100644 --- a/integration_test/compound_exps/CSR_Dense_chain_mult_matrix.ta +++ b/test/integration/compound_exps/CSR_Dense_chain_mult_matrix.ta @@ -1,6 +1,6 @@ # RUN: export SPARSE_FILE_NAME0=%comet_integration_test_data_dir/test_rank2.mtx # RUN: comet-opt --convert-ta-to-it --convert-to-loops --convert-to-llvm %s &> CSR_Dense_chain_mult_matrix.llvm -# RUN: mlir-cpu-runner CSR_Dense_chain_mult_matrix.llvm -O3 -e main -entry-point-result=void -shared-libs=%mlir_utility_library_dir/libmlir_runner_utils%shlibext,%comet_utility_library_dir/libcomet_runner_utils%shlibext | FileCheck %s +# RUN: mlir-cpu-runner CSR_Dense_chain_mult_matrix.llvm -O3 -e main -entry-point-result=void -shared-libs=%mlir_utility_library_dir/libmlir_runner_utils%shlibext,%comet_utility_library_dir/libcomet_runner_utils%shlibext,%mlir_utility_library_dir/libomp%shlibext | FileCheck %s def main() { diff --git a/integration_test/compound_exps/CSR_mult_dTranspose.ta b/test/integration/compound_exps/CSR_mult_dTranspose.ta similarity index 92% rename from integration_test/compound_exps/CSR_mult_dTranspose.ta rename to test/integration/compound_exps/CSR_mult_dTranspose.ta index d9c614e8..12178064 100644 --- a/integration_test/compound_exps/CSR_mult_dTranspose.ta +++ b/test/integration/compound_exps/CSR_mult_dTranspose.ta @@ -1,7 +1,7 @@ # RUN: export SPARSE_FILE_NAME0=%comet_integration_test_data_dir/test_rank2.mtx # RUN: export SORT_TYPE=SEQ_QSORT # RUN: comet-opt --convert-ta-to-it --convert-to-loops --convert-to-llvm %s &> CSR_mult_dTranspose.llvm -# RUN: mlir-cpu-runner CSR_mult_dTranspose.llvm -O3 -e main -entry-point-result=void -shared-libs=%comet_utility_library_dir/libcomet_runner_utils%shlibext | FileCheck %s +# RUN: mlir-cpu-runner CSR_mult_dTranspose.llvm -O3 -e main -entry-point-result=void -shared-libs=%comet_utility_library_dir/libcomet_runner_utils%shlibext,%mlir_utility_library_dir/libomp%shlibext | FileCheck %s def main() { #IndexLabel Declarations diff --git a/integration_test/compound_exps/Dense_chain_mult_matrix.ta b/test/integration/compound_exps/Dense_chain_mult_matrix.ta similarity index 89% rename from integration_test/compound_exps/Dense_chain_mult_matrix.ta rename to test/integration/compound_exps/Dense_chain_mult_matrix.ta index d0342578..12601e15 100644 --- a/integration_test/compound_exps/Dense_chain_mult_matrix.ta +++ b/test/integration/compound_exps/Dense_chain_mult_matrix.ta @@ -1,5 +1,5 @@ # RUN: comet-opt --convert-ta-to-it --convert-to-loops --convert-to-llvm %s &> Dense_chain_mult_matrix.llvm -# RUN: mlir-cpu-runner Dense_chain_mult_matrix.llvm -O3 -e main -entry-point-result=void -shared-libs=%comet_utility_library_dir/libcomet_runner_utils%shlibext | FileCheck %s +# RUN: mlir-cpu-runner Dense_chain_mult_matrix.llvm -O3 -e main -entry-point-result=void -shared-libs=%comet_utility_library_dir/libcomet_runner_utils%shlibext,%mlir_utility_library_dir/libomp%shlibext | FileCheck %s def main() { diff --git a/integration_test/compound_exps/Dense_eltwise_dTranspose.ta b/test/integration/compound_exps/Dense_eltwise_dTranspose.ta similarity index 90% rename from integration_test/compound_exps/Dense_eltwise_dTranspose.ta rename to test/integration/compound_exps/Dense_eltwise_dTranspose.ta index 60edb07d..53b238f9 100644 --- a/integration_test/compound_exps/Dense_eltwise_dTranspose.ta +++ b/test/integration/compound_exps/Dense_eltwise_dTranspose.ta @@ -1,5 +1,5 @@ # RUN: comet-opt --convert-ta-to-it --convert-to-loops --convert-to-llvm %s &> Dense_eltwise_dTranspose.llvm -# RUN: mlir-cpu-runner Dense_eltwise_dTranspose.llvm -O3 -e main -entry-point-result=void -shared-libs=%comet_utility_library_dir/libcomet_runner_utils%shlibext | FileCheck %s +# RUN: mlir-cpu-runner Dense_eltwise_dTranspose.llvm -O3 -e main -entry-point-result=void -shared-libs=%comet_utility_library_dir/libcomet_runner_utils%shlibext,%mlir_utility_library_dir/libomp%shlibext | FileCheck %s #TODO(gkestor): read dense input from file diff --git a/integration_test/compound_exps/Dense_mult_dTranspose.ta b/test/integration/compound_exps/Dense_mult_dTranspose.ta similarity index 90% rename from integration_test/compound_exps/Dense_mult_dTranspose.ta rename to test/integration/compound_exps/Dense_mult_dTranspose.ta index 81f117cf..03a947e8 100644 --- a/integration_test/compound_exps/Dense_mult_dTranspose.ta +++ b/test/integration/compound_exps/Dense_mult_dTranspose.ta @@ -1,5 +1,5 @@ # RUN: comet-opt --convert-ta-to-it --convert-to-loops --convert-to-llvm %s &> Dense_mult_dTranspose.llvm -# RUN: mlir-cpu-runner Dense_mult_dTranspose.llvm -O3 -e main -entry-point-result=void -shared-libs=%comet_utility_library_dir/libcomet_runner_utils%shlibext | FileCheck %s +# RUN: mlir-cpu-runner Dense_mult_dTranspose.llvm -O3 -e main -entry-point-result=void -shared-libs=%comet_utility_library_dir/libcomet_runner_utils%shlibext,%mlir_utility_library_dir/libomp%shlibext | FileCheck %s def main() { #IndexLabel Declarations diff --git a/integration_test/compound_exps/dTranspose_eltwise_CSR.ta b/test/integration/compound_exps/dTranspose_eltwise_CSR.ta similarity index 90% rename from integration_test/compound_exps/dTranspose_eltwise_CSR.ta rename to test/integration/compound_exps/dTranspose_eltwise_CSR.ta index 5046cf7e..64e4e7d1 100644 --- a/integration_test/compound_exps/dTranspose_eltwise_CSR.ta +++ b/test/integration/compound_exps/dTranspose_eltwise_CSR.ta @@ -1,7 +1,7 @@ # RUN: export SPARSE_FILE_NAME0=%comet_integration_test_data_dir/test_rank2.mtx # RUN: export SORT_TYPE=SEQ_QSORT # RUN: comet-opt --convert-ta-to-it --convert-to-loops --convert-to-llvm %s &> dTranspose_eltwise_CSR.llvm -# RUN: mlir-cpu-runner dTranspose_eltwise_CSR.llvm -O3 -e main -entry-point-result=void -shared-libs=%comet_utility_library_dir/libcomet_runner_utils%shlibext | FileCheck %s +# RUN: mlir-cpu-runner dTranspose_eltwise_CSR.llvm -O3 -e main -entry-point-result=void -shared-libs=%comet_utility_library_dir/libcomet_runner_utils%shlibext,%mlir_utility_library_dir/libomp%shlibext | FileCheck %s #TODO(gkestor): read dense input from file @@ -28,7 +28,7 @@ def main() { # CHECK: data = # CHECK-NEXT: 5, # CHECK-NEXT: data = -# CHECK-NEXT: 0, +# CHECK-NEXT: -1, # CHECK-NEXT: data = # CHECK-NEXT: 0,2,4,5,7,9, # CHECK-NEXT: data = diff --git a/integration_test/compound_exps/dTranspose_eltwise_Dense.ta b/test/integration/compound_exps/dTranspose_eltwise_Dense.ta similarity index 90% rename from integration_test/compound_exps/dTranspose_eltwise_Dense.ta rename to test/integration/compound_exps/dTranspose_eltwise_Dense.ta index 9f0e9c72..1f5fe373 100644 --- a/integration_test/compound_exps/dTranspose_eltwise_Dense.ta +++ b/test/integration/compound_exps/dTranspose_eltwise_Dense.ta @@ -1,5 +1,5 @@ # RUN: comet-opt --convert-ta-to-it --convert-to-loops --convert-to-llvm %s &> dTranspose_eltwise_Dense.llvm -# RUN: mlir-cpu-runner dTranspose_eltwise_Dense.llvm -O3 -e main -entry-point-result=void -shared-libs=%comet_utility_library_dir/libcomet_runner_utils%shlibext | FileCheck %s +# RUN: mlir-cpu-runner dTranspose_eltwise_Dense.llvm -O3 -e main -entry-point-result=void -shared-libs=%comet_utility_library_dir/libcomet_runner_utils%shlibext,%mlir_utility_library_dir/libomp%shlibext | FileCheck %s #TODO(gkestor): read dense input from file diff --git a/integration_test/compound_exps/dTranspose_mult_CSR.ta b/test/integration/compound_exps/dTranspose_mult_CSR.ta similarity index 92% rename from integration_test/compound_exps/dTranspose_mult_CSR.ta rename to test/integration/compound_exps/dTranspose_mult_CSR.ta index 7fdd639d..4d1b844a 100644 --- a/integration_test/compound_exps/dTranspose_mult_CSR.ta +++ b/test/integration/compound_exps/dTranspose_mult_CSR.ta @@ -1,7 +1,7 @@ # RUN: export SPARSE_FILE_NAME0=%comet_integration_test_data_dir/test_rank2.mtx # RUN: export SORT_TYPE=SEQ_QSORT # RUN: comet-opt --convert-ta-to-it --convert-to-loops --convert-to-llvm %s &> dTranspose_mult_CSR.llvm -# RUN: mlir-cpu-runner dTranspose_mult_CSR.llvm -O3 -e main -entry-point-result=void -shared-libs=%comet_utility_library_dir/libcomet_runner_utils%shlibext | FileCheck %s +# RUN: mlir-cpu-runner dTranspose_mult_CSR.llvm -O3 -e main -entry-point-result=void -shared-libs=%comet_utility_library_dir/libcomet_runner_utils%shlibext,%mlir_utility_library_dir/libomp%shlibext | FileCheck %s #TODO(gkestor): read dense input from file diff --git a/integration_test/compound_exps/dTranspose_mult_Dense.ta b/test/integration/compound_exps/dTranspose_mult_Dense.ta similarity index 90% rename from integration_test/compound_exps/dTranspose_mult_Dense.ta rename to test/integration/compound_exps/dTranspose_mult_Dense.ta index 3e3339a3..9c636c8f 100644 --- a/integration_test/compound_exps/dTranspose_mult_Dense.ta +++ b/test/integration/compound_exps/dTranspose_mult_Dense.ta @@ -1,5 +1,5 @@ # RUN: comet-opt --convert-ta-to-it --convert-to-loops --convert-to-llvm %s &> dTranspose_mult_Dense.llvm -# RUN: mlir-cpu-runner dTranspose_mult_Dense.llvm -O3 -e main -entry-point-result=void -shared-libs=%comet_utility_library_dir/libcomet_runner_utils%shlibext | FileCheck %s +# RUN: mlir-cpu-runner dTranspose_mult_Dense.llvm -O3 -e main -entry-point-result=void -shared-libs=%comet_utility_library_dir/libcomet_runner_utils%shlibext,%mlir_utility_library_dir/libomp%shlibext | FileCheck %s #TODO(gkestor): read dense input from file diff --git a/test/integration/data/tc.mtx b/test/integration/data/tc.mtx new file mode 100644 index 00000000..2ebf757c --- /dev/null +++ b/test/integration/data/tc.mtx @@ -0,0 +1,30 @@ +%%MatrixMarket matrix coordinate real general +% +% This is a test sparse matrix in Matrix Market Exchange Format. +% see https://math.nist.gov/MatrixMarket +% +7 7 24 +1 2 1.0 +1 4 1.0 +2 1 1.0 +2 4 1.0 +2 5 1.0 +2 7 1.0 +3 4 1.0 +3 6 1.0 +3 7 1.0 +4 1 1.0 +4 2 1.0 +4 3 1.0 +4 6 1.0 +4 7 1.0 +5 2 1.0 +5 6 1.0 +5 7 1.0 +6 3 1.0 +6 4 1.0 +6 5 1.0 +7 2 1.0 +7 3 1.0 +7 4 1.0 +7 5 1.0 diff --git a/test/integration/data/test_8x6.mtx b/test/integration/data/test_8x6.mtx new file mode 100644 index 00000000..e13b677c --- /dev/null +++ b/test/integration/data/test_8x6.mtx @@ -0,0 +1,14 @@ +%%MatrixMarket matrix coordinate real general +% +% This is a test sparse matrix in Matrix Market Exchange Format. +% see https://math.nist.gov/MatrixMarket +% +% This example comes from https://matteding.github.io/2019/04/25/sparse-matrices/ +8 6 7 +1 1 8.0 +1 3 2.0 +2 3 5.0 +5 3 7.0 +5 4 1.0 +5 5 2.0 +7 4 9.0 diff --git a/test/integration/data/test_rank2.mtx b/test/integration/data/test_rank2.mtx new file mode 100644 index 00000000..13a34fba --- /dev/null +++ b/test/integration/data/test_rank2.mtx @@ -0,0 +1,15 @@ +%%MatrixMarket matrix coordinate real general +% +% This is a test sparse matrix in Matrix Market Exchange Format. +% see https://math.nist.gov/MatrixMarket +% +5 5 9 +1 1 1.0 +1 4 1.4 +2 2 2.0 +2 5 2.5 +3 3 3.0 +4 4 4.0 +4 1 4.1 +5 5 5.0 +5 2 5.2 diff --git a/test/integration/data/test_rank2_denser.mtx b/test/integration/data/test_rank2_denser.mtx new file mode 100644 index 00000000..e22ebc75 --- /dev/null +++ b/test/integration/data/test_rank2_denser.mtx @@ -0,0 +1,18 @@ +%%MatrixMarket matrix coordinate real general +% +% This is a test sparse matrix in Matrix Market Exchange Format. +% see https://math.nist.gov/MatrixMarket +% +5 5 12 +1 1 1.0 +1 4 1.4 +2 2 2.0 +2 5 2.5 +3 3 3.0 +3 1 3.1 +4 4 4.0 +4 1 4.1 +4 2 4.2 +5 5 5.0 +5 2 5.2 +5 3 5.3 \ No newline at end of file diff --git a/test/integration/data/test_rank2_small.mtx b/test/integration/data/test_rank2_small.mtx new file mode 100644 index 00000000..55bc6ba1 --- /dev/null +++ b/test/integration/data/test_rank2_small.mtx @@ -0,0 +1,15 @@ +%%MatrixMarket matrix coordinate real general +% +% This is a test sparse matrix in Matrix Market Exchange Format. +% see https://math.nist.gov/MatrixMarket +% +% This is the same example used in the paper +% https://arxiv.org/pdf/2102.05187.pdf +5 4 7 +1 1 1.0 +1 4 2.0 +2 1 3.0 +2 2 4.0 +4 2 5.0 +5 3 6.0 +5 4 7.0 diff --git a/test/integration/data/test_rank2_transpose.mtx b/test/integration/data/test_rank2_transpose.mtx new file mode 100644 index 00000000..92331cd3 --- /dev/null +++ b/test/integration/data/test_rank2_transpose.mtx @@ -0,0 +1,15 @@ +%%MatrixMarket matrix coordinate real general +% +% This is a test sparse matrix in Matrix Market Exchange Format. +% see https://math.nist.gov/MatrixMarket +% +5 5 9 +1 1 1.0 +1 4 4.1 +2 2 2.0 +2 5 5.2 +3 3 3.0 +4 4 4.0 +4 1 1.4 +5 5 5.0 +5 2 2.5 diff --git a/test/integration/data/test_rank3.tns b/test/integration/data/test_rank3.tns new file mode 100644 index 00000000..195a3cae --- /dev/null +++ b/test/integration/data/test_rank3.tns @@ -0,0 +1,4 @@ +7 7 7 3 +1 3 2 1.3 +2 1 3 2.11 +3 6 5 3.0 diff --git a/test/integration/data/test_rank8.tns b/test/integration/data/test_rank8.tns new file mode 100644 index 00000000..1867edc0 --- /dev/null +++ b/test/integration/data/test_rank8.tns @@ -0,0 +1,17 @@ +7 3 3 3 3 3 7 16 +2 2 2 1 2 2 2 1 2.11 +3 3 3 3 3 3 3 3 3.0 +1 1 1 1 1 1 1 3 1.3 +1 1 1 2 1 1 1 3 1.23 +1 1 1 2 1 1 1 2 1.22 +2 2 2 2 2 2 2 2 2.0 +1 1 1 1 1 1 1 1 1.0 +7 1 1 1 1 1 1 1 7.0 +3 3 3 3 1 3 3 3 3.1 +3 3 3 3 1 2 2 3 3.122 +3 3 3 3 1 1 2 3 3.112 +3 3 3 3 1 2 1 3 3.121 +2 2 2 1 2 2 2 2 2.1 +2 1 2 1 2 2 2 1 2.111 +2 1 2 1 2 2 2 3 2.113 +1 1 1 1 1 1 5 1 1.5 \ No newline at end of file diff --git a/test/integration/data/wide.mtx b/test/integration/data/wide.mtx new file mode 100644 index 00000000..9e0d5f2a --- /dev/null +++ b/test/integration/data/wide.mtx @@ -0,0 +1,23 @@ +%%MatrixMarket matrix coordinate real general +% +% This is a test sparse matrix in Matrix Market Exchange Format. +% see https://math.nist.gov/MatrixMarket +% +4 256 17 +1 1 -1 +1 127 2 +1 128 -3 +1 255 4 +2 2 -5 +2 254 6 +3 3 -7 +4 1 8 +4 2 -9 +4 4 10 +4 99 -11 +4 127 12 +4 128 -13 +4 129 14 +4 250 -15 +4 254 16 +4 256 -17 diff --git a/integration_test/compound_exps/gpu/Dense_chain_mult_matrix.ta b/test/integration/gpu/Dense_chain_mult_matrix.ta similarity index 75% rename from integration_test/compound_exps/gpu/Dense_chain_mult_matrix.ta rename to test/integration/gpu/Dense_chain_mult_matrix.ta index 27f0b673..d73607b7 100644 --- a/integration_test/compound_exps/gpu/Dense_chain_mult_matrix.ta +++ b/test/integration/gpu/Dense_chain_mult_matrix.ta @@ -1,5 +1,5 @@ -# RUN: comet-opt --convert-ta-to-it --convert-to-loops --convert-to-triton --convert-to-llvm %s &> Dense_chain_mult_matrix.llvm -# RUN: mlir-cpu-runner Dense_chain_mult_matrix.llvm -O3 -e main -entry-point-result=void -shared-libs=%comet_utility_library_dir/libcomet_runner_utils%shlibext | FileCheck %s +# RUN: comet-opt --target=GPU --convert-ta-to-it --convert-to-loops --convert-to-triton --convert-to-llvm %s &> Dense_chain_mult_matrix.llvm +# RUN: mlir-cpu-runner Dense_chain_mult_matrix.llvm -O3 -e main -entry-point-result=void -shared-libs=%comet_utility_library_dir/libcomet_runner_utils%shlibext,%mlir_utility_library_dir/libomp%shlibext | FileCheck %s def main() { diff --git a/integration_test/ops/gpu/eltwise_add_dense_matrix.ta b/test/integration/gpu/eltwise_add_dense_matrix.ta similarity index 100% rename from integration_test/ops/gpu/eltwise_add_dense_matrix.ta rename to test/integration/gpu/eltwise_add_dense_matrix.ta diff --git a/integration_test/semiring/gpu/eltwise_monoidMin_DensexDense_oDense.ta b/test/integration/gpu/eltwise_monoidMin_DensexDense_oDense.ta similarity index 100% rename from integration_test/semiring/gpu/eltwise_monoidMin_DensexDense_oDense.ta rename to test/integration/gpu/eltwise_monoidMin_DensexDense_oDense.ta diff --git a/integration_test/semiring/gpu/eltwise_monoidMinus_DensexDense_oDense.ta b/test/integration/gpu/eltwise_monoidMinus_DensexDense_oDense.ta similarity index 100% rename from integration_test/semiring/gpu/eltwise_monoidMinus_DensexDense_oDense.ta rename to test/integration/gpu/eltwise_monoidMinus_DensexDense_oDense.ta diff --git a/integration_test/semiring/gpu/eltwise_monoidPlus_DensexDense_oDense.ta b/test/integration/gpu/eltwise_monoidPlus_DensexDense_oDense.ta similarity index 100% rename from integration_test/semiring/gpu/eltwise_monoidPlus_DensexDense_oDense.ta rename to test/integration/gpu/eltwise_monoidPlus_DensexDense_oDense.ta diff --git a/integration_test/semiring/gpu/eltwise_monoidTimes_DensexDense_oDense.ta b/test/integration/gpu/eltwise_monoidTimes_DensexDense_oDense.ta similarity index 100% rename from integration_test/semiring/gpu/eltwise_monoidTimes_DensexDense_oDense.ta rename to test/integration/gpu/eltwise_monoidTimes_DensexDense_oDense.ta diff --git a/integration_test/ops/gpu/eltwise_mult_DensexDense_oDense.ta b/test/integration/gpu/eltwise_mult_DensexDense_oDense.ta similarity index 100% rename from integration_test/ops/gpu/eltwise_mult_DensexDense_oDense.ta rename to test/integration/gpu/eltwise_mult_DensexDense_oDense.ta diff --git a/integration_test/ops/gpu/eltwise_subtract_dense_matrix.ta b/test/integration/gpu/eltwise_subtract_dense_matrix.ta similarity index 100% rename from integration_test/ops/gpu/eltwise_subtract_dense_matrix.ta rename to test/integration/gpu/eltwise_subtract_dense_matrix.ta diff --git a/integration_test/semiring/gpu/mm_SemiringPlusTimes_DensexDense_oDense.ta b/test/integration/gpu/mm_SemiringPlusTimes_DensexDense_oDense.ta similarity index 100% rename from integration_test/semiring/gpu/mm_SemiringPlusTimes_DensexDense_oDense.ta rename to test/integration/gpu/mm_SemiringPlusTimes_DensexDense_oDense.ta diff --git a/integration_test/ops/gpu/mult_dense_ij-ikj-kj.ta b/test/integration/gpu/mult_dense_ij-ikj-kj.ta similarity index 100% rename from integration_test/ops/gpu/mult_dense_ij-ikj-kj.ta rename to test/integration/gpu/mult_dense_ij-ikj-kj.ta diff --git a/integration_test/ops/gpu/mult_dense_matrix.ta b/test/integration/gpu/mult_dense_matrix.ta similarity index 100% rename from integration_test/ops/gpu/mult_dense_matrix.ta rename to test/integration/gpu/mult_dense_matrix.ta diff --git a/integration_test/ops/gpu/mult_dense_matrix_vector.ta b/test/integration/gpu/mult_dense_matrix_vector.ta similarity index 100% rename from integration_test/ops/gpu/mult_dense_matrix_vector.ta rename to test/integration/gpu/mult_dense_matrix_vector.ta diff --git a/integration_test/semiring/gpu/mv_SemiringPlusTimes_DensexDense_oDense.ta b/test/integration/gpu/mv_SemiringPlusTimes_DensexDense_oDense.ta similarity index 100% rename from integration_test/semiring/gpu/mv_SemiringPlusTimes_DensexDense_oDense.ta rename to test/integration/gpu/mv_SemiringPlusTimes_DensexDense_oDense.ta diff --git a/test/integration/gpu/not-supported/lit.local.cfg b/test/integration/gpu/not-supported/lit.local.cfg new file mode 100644 index 00000000..059ca10b --- /dev/null +++ b/test/integration/gpu/not-supported/lit.local.cfg @@ -0,0 +1 @@ +config.unsupported = True \ No newline at end of file diff --git a/integration_test/ops/gpu/sum_dense_matrix.ta b/test/integration/gpu/not-supported/sum_dense_matrix.ta similarity index 100% rename from integration_test/ops/gpu/sum_dense_matrix.ta rename to test/integration/gpu/not-supported/sum_dense_matrix.ta diff --git a/integration_test/ops/gpu/transpose_dense_matrix.ta b/test/integration/gpu/not-supported/transpose_dense_matrix.ta similarity index 100% rename from integration_test/ops/gpu/transpose_dense_matrix.ta rename to test/integration/gpu/not-supported/transpose_dense_matrix.ta diff --git a/integration_test/kernels/ccsd_t1_21_loops.ta b/test/integration/kernels/ccsd_t1_21_loops.ta similarity index 90% rename from integration_test/kernels/ccsd_t1_21_loops.ta rename to test/integration/kernels/ccsd_t1_21_loops.ta index 12013572..f587c4c0 100644 --- a/integration_test/kernels/ccsd_t1_21_loops.ta +++ b/test/integration/kernels/ccsd_t1_21_loops.ta @@ -1,5 +1,5 @@ # RUN: comet-opt --convert-ta-to-it --convert-to-loops --convert-to-llvm %s &> ccsd_t1_21_loops.llvm -# RUN: mlir-cpu-runner ccsd_t1_21_loops.llvm -O3 -e main -entry-point-result=void -shared-libs=%comet_utility_library_dir/libcomet_runner_utils%shlibext | FileCheck %s +# RUN: mlir-cpu-runner ccsd_t1_21_loops.llvm -O3 -e main -entry-point-result=void -shared-libs=%comet_utility_library_dir/libcomet_runner_utils%shlibext,%mlir_utility_library_dir/libomp%shlibext | FileCheck %s def main() { #IndexLabel Declarations diff --git a/integration_test/kernels/ccsd_t1_3_loops.ta b/test/integration/kernels/ccsd_t1_3_loops.ta similarity index 90% rename from integration_test/kernels/ccsd_t1_3_loops.ta rename to test/integration/kernels/ccsd_t1_3_loops.ta index f940fd70..1e44b867 100644 --- a/integration_test/kernels/ccsd_t1_3_loops.ta +++ b/test/integration/kernels/ccsd_t1_3_loops.ta @@ -1,5 +1,5 @@ # RUN: comet-opt --convert-ta-to-it --convert-to-loops --convert-to-llvm %s &> ccsd_t1_3_loops.llvm -# RUN: mlir-cpu-runner ccsd_t1_3_loops.llvm -O3 -e main -entry-point-result=void -shared-libs=%comet_utility_library_dir/libcomet_runner_utils%shlibext | FileCheck %s +# RUN: mlir-cpu-runner ccsd_t1_3_loops.llvm -O3 -e main -entry-point-result=void -shared-libs=%comet_utility_library_dir/libcomet_runner_utils%shlibext,%mlir_utility_library_dir/libomp%shlibext | FileCheck %s def main() { #IndexLabel Declarations diff --git a/integration_test/kernels/ccsd_t1_4_loops.ta b/test/integration/kernels/ccsd_t1_4_loops.ta similarity index 90% rename from integration_test/kernels/ccsd_t1_4_loops.ta rename to test/integration/kernels/ccsd_t1_4_loops.ta index fb7f2e62..9a2c3b96 100644 --- a/integration_test/kernels/ccsd_t1_4_loops.ta +++ b/test/integration/kernels/ccsd_t1_4_loops.ta @@ -1,5 +1,5 @@ # RUN: comet-opt --convert-ta-to-it --convert-to-loops --convert-to-llvm %s &> ccsd_t1_4_loops.llvm -# RUN: mlir-cpu-runner ccsd_t1_4_loops.llvm -O3 -e main -entry-point-result=void -shared-libs=%comet_utility_library_dir/libcomet_runner_utils%shlibext | FileCheck %s +# RUN: mlir-cpu-runner ccsd_t1_4_loops.llvm -O3 -e main -entry-point-result=void -shared-libs=%comet_utility_library_dir/libcomet_runner_utils%shlibext,%mlir_utility_library_dir/libomp%shlibext | FileCheck %s def main() { #IndexLabel Declarations diff --git a/integration_test/kernels/gnn.ta b/test/integration/kernels/gnn.ta similarity index 94% rename from integration_test/kernels/gnn.ta rename to test/integration/kernels/gnn.ta index a5b003b0..18a6e454 100644 --- a/integration_test/kernels/gnn.ta +++ b/test/integration/kernels/gnn.ta @@ -4,7 +4,7 @@ # RUN: comet-opt --convert-ta-to-it --opt-fusion --convert-to-loops --convert-to-llvm %s &> gnn_loops.llvm # RUN: export SPARSE_FILE_NAME0=%comet_integration_test_data_dir/test_rank2_small.mtx -# RUN: mlir-cpu-runner gnn_loops.llvm -O3 -e main -entry-point-result=void -shared-libs=%comet_utility_library_dir/libcomet_runner_utils%shlibext | FileCheck %s +# RUN: mlir-cpu-runner gnn_loops.llvm -O3 -e main -entry-point-result=void -shared-libs=%comet_utility_library_dir/libcomet_runner_utils%shlibext,%mlir_utility_library_dir/libomp%shlibext | FileCheck %s def main() { diff --git a/integration_test/kernels/gram_matrix.ta b/test/integration/kernels/gram_matrix.ta similarity index 89% rename from integration_test/kernels/gram_matrix.ta rename to test/integration/kernels/gram_matrix.ta index 3a6b5e7a..c429640c 100644 --- a/integration_test/kernels/gram_matrix.ta +++ b/test/integration/kernels/gram_matrix.ta @@ -1,5 +1,5 @@ # RUN: comet-opt --convert-ta-to-it --convert-to-loops --convert-to-llvm %s &> gram_matrix.llvm -# RUN: mlir-cpu-runner gram_matrix.llvm -O3 -e main -entry-point-result=void -shared-libs=%comet_utility_library_dir/libcomet_runner_utils%shlibext | FileCheck %s +# RUN: mlir-cpu-runner gram_matrix.llvm -O3 -e main -entry-point-result=void -shared-libs=%comet_utility_library_dir/libcomet_runner_utils%shlibext,%mlir_utility_library_dir/libomp%shlibext | FileCheck %s def main() { diff --git a/test/integration/kernels/sddmm.ta b/test/integration/kernels/sddmm.ta new file mode 100644 index 00000000..316a1a2f --- /dev/null +++ b/test/integration/kernels/sddmm.ta @@ -0,0 +1,44 @@ +# Sampled dense-dense matrix product (SDDMM) +# A[i,j] = B[i,j] .* (C[i,k] * D[k,j]); +# A and B are sparse; C and D are dense. + +# --opt-fusion pass performs the redundancy-aware fusion on SDDMM kernels +# RUN: comet-opt --convert-ta-to-it --opt-fusion --convert-to-loops --convert-to-llvm %s &> sddmm.llvm +# RUN: export SPARSE_FILE_NAME0=%comet_integration_test_data_dir/test_rank2_small.mtx +# RUN: mlir-cpu-runner sddmm.llvm -O3 -e main -entry-point-result=void -shared-libs=%comet_utility_library_dir/libcomet_runner_utils%shlibext | FileCheck %s + +def main() { + #IndexLabel Declarations + IndexLabel [i] = [?]; + IndexLabel [j] = [?]; + IndexLabel [k] = [4]; + + #Tensor Declarations + Tensor B([i, j], {CSR}); + Tensor C([i, k], {Dense}); + Tensor D([k, j], {Dense}); + Tensor A([i, j], {CSR}); + Tensor T([i, j], {Dense}); + + #Tensor Data Initialization + B[i, j] = comet_read(0); + C[i, k] = 1.2; + D[k, j] = 3.4; + T[i, j] = 0.0; + + T[i, j] = C[i, k] * D[k, j]; + A[i, j] = B[i, j] .* T[i, j]; + print(A); +} + +# Print the result for verification. +# CHECK: data = +# CHECK-NEXT: 5, +# CHECK-NEXT: data = +# CHECK-NEXT: -1, +# CHECK-NEXT: data = +# CHECK-NEXT: 0,2,4,4,5,7, +# CHECK-NEXT: data = +# CHECK-NEXT: 0,3,0,1,1,2,3, +# CHECK-NEXT: data = +# CHECK-NEXT: 16.32,32.64,48.96,65.28,81.6,97.92,114.24, diff --git a/integration_test/kernels/triangleCount_SandiaLL.ta b/test/integration/kernels/triangleCount_SandiaLL.ta similarity index 92% rename from integration_test/kernels/triangleCount_SandiaLL.ta rename to test/integration/kernels/triangleCount_SandiaLL.ta index eabcf445..9d984556 100644 --- a/integration_test/kernels/triangleCount_SandiaLL.ta +++ b/test/integration/kernels/triangleCount_SandiaLL.ta @@ -13,7 +13,7 @@ # RUN: comet-opt --opt-comp-workspace --convert-ta-to-it --convert-to-loops --convert-to-llvm %s &> triangleCount_SandiaLL.llvm # RUN: export SPARSE_FILE_NAME0=%comet_integration_test_data_dir/tc.mtx -# RUN: mlir-cpu-runner triangleCount_SandiaLL.llvm -O3 -e main -entry-point-result=void -shared-libs=%comet_utility_library_dir/libcomet_runner_utils%shlibext | FileCheck %s +# RUN: mlir-cpu-runner triangleCount_SandiaLL.llvm -O3 -e main -entry-point-result=void -shared-libs=%comet_utility_library_dir/libcomet_runner_utils%shlibext,%mlir_utility_library_dir/libmlir_c_runner_utils%shlibext | FileCheck %s def main() { #IndexLabel Declarations @@ -34,5 +34,4 @@ def main() { } # Print the result for verification. -# CHECK: data = -# CHECK-NEXT: 5, \ No newline at end of file +# CHECK: 5 \ No newline at end of file diff --git a/integration_test/kernels/triangleCount_SandiaLL_wMasking.ta b/test/integration/kernels/triangleCount_SandiaLL_wMasking.ta similarity index 89% rename from integration_test/kernels/triangleCount_SandiaLL_wMasking.ta rename to test/integration/kernels/triangleCount_SandiaLL_wMasking.ta index 40ae4ac7..40867678 100644 --- a/integration_test/kernels/triangleCount_SandiaLL_wMasking.ta +++ b/test/integration/kernels/triangleCount_SandiaLL_wMasking.ta @@ -13,7 +13,7 @@ # RUN: comet-opt --opt-comp-workspace --convert-ta-to-it --convert-to-loops --convert-to-llvm %s &> triangleCount_SandiaLL.mask.llvm # RUN: export SPARSE_FILE_NAME0=%comet_integration_test_data_dir/tc.mtx -# RUN: mlir-cpu-runner triangleCount_SandiaLL.mask.llvm -O3 -e main -entry-point-result=void -shared-libs=%comet_utility_library_dir/libcomet_runner_utils%shlibext | FileCheck %s +# RUN: mlir-cpu-runner triangleCount_SandiaLL.mask.llvm -O3 -e main -entry-point-result=void -shared-libs=%comet_utility_library_dir/libcomet_runner_utils%shlibext,%mlir_utility_library_dir/libmlir_c_runner_utils%shlibext | FileCheck %s def main() { #IndexLabel Declarations @@ -31,7 +31,7 @@ def main() { # Sandia_LL method: ntri = sum (sum ((L * L) .* L)) # var ntri = SUM((L[i,k] * L[k,j]) .* L[i,j]); ## - C[i, j] = L1[i, k] * L1[k, j]; # L0 is the mask, using push-based method. + C[i, k] = L1[i, j] * L1[j, k]; # L0 is the mask, using push-based method. # valid options: {push, pull, auto} var ntri = SUM(C[i, j]); @@ -39,5 +39,4 @@ def main() { } # Print the result for verification. -# CHECK: data = -# CHECK-NEXT: 5, \ No newline at end of file +# CHECK: 5 \ No newline at end of file diff --git a/integration_test/lit.cfg.py b/test/integration/lit.cfg.py similarity index 99% rename from integration_test/lit.cfg.py rename to test/integration/lit.cfg.py index 8d835470..664cd352 100644 --- a/integration_test/lit.cfg.py +++ b/test/integration/lit.cfg.py @@ -58,7 +58,7 @@ config.test_source_root = os.path.dirname(__file__) # test_exec_root: The root path where tests should be run. -config.test_exec_root = os.path.join(config.comet_obj_root, 'integration_test') +config.test_exec_root = os.path.join(config.comet_obj_root, 'integration') # Tweak the PATH to include the tools dir. llvm_config.with_environment('PATH', config.llvm_tools_dir, append_path=True) diff --git a/integration_test/lit.site.cfg.py.in b/test/integration/lit.site.cfg.py.in similarity index 96% rename from integration_test/lit.site.cfg.py.in rename to test/integration/lit.site.cfg.py.in index f0fc42d9..ce5222c4 100644 --- a/integration_test/lit.site.cfg.py.in +++ b/test/integration/lit.site.cfg.py.in @@ -58,4 +58,4 @@ import lit.llvm lit.llvm.initialize(lit_config, config) # Let the main config do the real work. -lit_config.load_config(config, "@COMET_SOURCE_DIR@/integration_test/lit.cfg.py") +lit_config.load_config(config, "@COMET_SOURCE_DIR@/test/integration/lit.cfg.py") diff --git a/integration_test/ops/eltwise_add_CSRxCSR_oCSR.ta b/test/integration/not_supported/eltwise_add_CSRxCSR_oCSR.ta similarity index 91% rename from integration_test/ops/eltwise_add_CSRxCSR_oCSR.ta rename to test/integration/not_supported/eltwise_add_CSRxCSR_oCSR.ta index aabf23d7..e486887d 100644 --- a/integration_test/ops/eltwise_add_CSRxCSR_oCSR.ta +++ b/test/integration/not_supported/eltwise_add_CSRxCSR_oCSR.ta @@ -1,9 +1,10 @@ # Sparse matrix sparse matrix elementwise addition # Sparse matrix is in CSR format. Currently workspace transformation on the IndexTree dialect works for only CSR format +# UNSUPPORTED: * # RUN: comet-opt --opt-comp-workspace --convert-ta-to-it --convert-to-loops --convert-to-llvm %s &> eltwise_add_CSRxCSR_oCSR.llvm # RUN: export SPARSE_FILE_NAME0=%comet_integration_test_data_dir/test_rank2.mtx # RUN: export SPARSE_FILE_NAME1=%comet_integration_test_data_dir/test_rank2_transpose.mtx -# RUN: mlir-cpu-runner eltwise_add_CSRxCSR_oCSR.llvm -O3 -e main -entry-point-result=void -shared-libs=%comet_utility_library_dir/libcomet_runner_utils%shlibext | FileCheck %s +# RUN: mlir-cpu-runner eltwise_add_CSRxCSR_oCSR.llvm -O3 -e main -entry-point-result=void -shared-libs=%comet_utility_library_dir/libcomet_runner_utils%shlibext,%mlir_utility_library_dir/libomp%shlibext | FileCheck %s def main() { @@ -29,7 +30,7 @@ def main() { # CHECK: data = # CHECK-NEXT: 5, # CHECK-NEXT: data = -# CHECK-NEXT: 0, +# CHECK-NEXT: -1, # CHECK-NEXT: data = # CHECK-NEXT: 0,2,4,5,7,9, # CHECK-NEXT: data = diff --git a/integration_test/ops/eltwise_mult_DCSRxDense_oDCSR.ta b/test/integration/not_supported/eltwise_mult_DCSRxDense_oDCSR.ta similarity index 90% rename from integration_test/ops/eltwise_mult_DCSRxDense_oDCSR.ta rename to test/integration/not_supported/eltwise_mult_DCSRxDense_oDCSR.ta index b1be02ee..476438a3 100644 --- a/integration_test/ops/eltwise_mult_DCSRxDense_oDCSR.ta +++ b/test/integration/not_supported/eltwise_mult_DCSRxDense_oDCSR.ta @@ -1,6 +1,7 @@ +# UNSUPPORTED: * # RUN: comet-opt --convert-ta-to-it --convert-to-loops --convert-to-llvm %s &> eltwise_DCSRxDense_oDCSR.llvm # RUN: export SPARSE_FILE_NAME0=%comet_integration_test_data_dir/test_rank2.mtx -# RUN: mlir-cpu-runner eltwise_DCSRxDense_oDCSR.llvm -O3 -e main -entry-point-result=void -shared-libs=%comet_utility_library_dir/libcomet_runner_utils%shlibext | FileCheck %s +# RUN: mlir-cpu-runner eltwise_DCSRxDense_oDCSR.llvm -O3 -e main -entry-point-result=void -shared-libs=%comet_utility_library_dir/libcomet_runner_utils%shlibext,%mlir_utility_library_dir/libomp%shlibext | FileCheck %s def main() { #IndexLabel Declarations diff --git a/integration_test/ops/eltwise_subtract_CSRxCSR_oCSR.ta b/test/integration/not_supported/eltwise_subtract_CSRxCSR_oCSR.ta similarity index 90% rename from integration_test/ops/eltwise_subtract_CSRxCSR_oCSR.ta rename to test/integration/not_supported/eltwise_subtract_CSRxCSR_oCSR.ta index e8aa66ae..e9cdaa1a 100644 --- a/integration_test/ops/eltwise_subtract_CSRxCSR_oCSR.ta +++ b/test/integration/not_supported/eltwise_subtract_CSRxCSR_oCSR.ta @@ -1,9 +1,10 @@ # Sparse matrix dense matrix elementwise subtraction # Sparse matrix is in CSR format. Currently workspace transformation on the IndexTree dialect works for only CSR format +# UNSUPPORTED: * # RUN: comet-opt --opt-comp-workspace --convert-ta-to-it --convert-to-loops --convert-to-llvm %s &> eltwise_sub_CSRxCSR_oCSR_sameSpPattern.llvm # RUN: export SPARSE_FILE_NAME0=%comet_integration_test_data_dir/test_rank2.mtx # RUN: export SPARSE_FILE_NAME1=%comet_integration_test_data_dir/test_rank2_transpose.mtx -# RUN: mlir-cpu-runner eltwise_sub_CSRxCSR_oCSR_sameSpPattern.llvm -O3 -e main -entry-point-result=void -shared-libs=%comet_utility_library_dir/libcomet_runner_utils%shlibext | FileCheck %s +# RUN: mlir-cpu-runner eltwise_sub_CSRxCSR_oCSR_sameSpPattern.llvm -O3 -e main -entry-point-result=void -shared-libs=%comet_utility_library_dir/libcomet_runner_utils%shlibext,%mlir_utility_library_dir/libomp%shlibext | FileCheck %s def main() { @@ -29,7 +30,7 @@ def main() { # CHECK: data = # CHECK-NEXT: 5, # CHECK-NEXT: data = -# CHECK-NEXT: 0, +# CHECK-NEXT: -1, # CHECK-NEXT: data = # CHECK-NEXT: 0,2,4,5,7,9, # CHECK-NEXT: data = diff --git a/test/integration/not_supported/lit.local.cfg b/test/integration/not_supported/lit.local.cfg new file mode 100644 index 00000000..059ca10b --- /dev/null +++ b/test/integration/not_supported/lit.local.cfg @@ -0,0 +1 @@ +config.unsupported = True \ No newline at end of file diff --git a/integration_test/ops/eltwise_add_dense_matrix.ta b/test/integration/ops/eltwise_add_dense_matrix.ta similarity index 88% rename from integration_test/ops/eltwise_add_dense_matrix.ta rename to test/integration/ops/eltwise_add_dense_matrix.ta index b0723b94..6c9d4a6c 100644 --- a/integration_test/ops/eltwise_add_dense_matrix.ta +++ b/test/integration/ops/eltwise_add_dense_matrix.ta @@ -1,5 +1,5 @@ # RUN: comet-opt --convert-ta-to-it --convert-to-loops --convert-to-llvm %s &> eltwise_add_dense_matrix.llvm -# RUN: mlir-cpu-runner eltwise_add_dense_matrix.llvm -O3 -e main -entry-point-result=void -shared-libs=%comet_utility_library_dir/libcomet_runner_utils%shlibext | FileCheck %s +# RUN: mlir-cpu-runner eltwise_add_dense_matrix.llvm -O3 -e main -entry-point-result=void -shared-libs=%comet_utility_library_dir/libcomet_runner_utils%shlibext,%mlir_utility_library_dir/libomp%shlibext | FileCheck %s def main() { diff --git a/integration_test/ops/eltwise_mult_COOxDense_oCOO.ta b/test/integration/ops/eltwise_mult_COOxDense_oCOO.ta similarity index 90% rename from integration_test/ops/eltwise_mult_COOxDense_oCOO.ta rename to test/integration/ops/eltwise_mult_COOxDense_oCOO.ta index 999d9bc0..265b4659 100644 --- a/integration_test/ops/eltwise_mult_COOxDense_oCOO.ta +++ b/test/integration/ops/eltwise_mult_COOxDense_oCOO.ta @@ -1,6 +1,6 @@ # RUN: comet-opt --convert-ta-to-it --convert-to-loops --convert-to-llvm %s &> eltwise_COOxDense_oCOO.llvm # RUN: export SPARSE_FILE_NAME0=%comet_integration_test_data_dir/test_rank2.mtx -# RUN: mlir-cpu-runner eltwise_COOxDense_oCOO.llvm -O3 -e main -entry-point-result=void -shared-libs=%comet_utility_library_dir/libcomet_runner_utils%shlibext | FileCheck %s +# RUN: mlir-cpu-runner eltwise_COOxDense_oCOO.llvm -O3 -e main -entry-point-result=void -shared-libs=%comet_utility_library_dir/libcomet_runner_utils%shlibext,%mlir_utility_library_dir/libomp%shlibext | FileCheck %s def main() { #IndexLabel Declarations @@ -27,7 +27,7 @@ def main() { # CHECK-NEXT: data = # CHECK-NEXT: 0,0,1,1,2,3,3,4,4, # CHECK-NEXT: data = -# CHECK-NEXT: 0, +# CHECK-NEXT: -1, # CHECK-NEXT: data = # CHECK-NEXT: 0,3,1,4,2,0,3,1,4, # CHECK-NEXT: data = diff --git a/integration_test/ops/eltwise_mult_CSRxCSR_oCSR.ta b/test/integration/ops/eltwise_mult_CSRxCSR_oCSR.ta similarity index 92% rename from integration_test/ops/eltwise_mult_CSRxCSR_oCSR.ta rename to test/integration/ops/eltwise_mult_CSRxCSR_oCSR.ta index 8d92708a..6498c23c 100644 --- a/integration_test/ops/eltwise_mult_CSRxCSR_oCSR.ta +++ b/test/integration/ops/eltwise_mult_CSRxCSR_oCSR.ta @@ -3,7 +3,7 @@ # In this example, workspace transformations are apply on the input sparse matrices. # RUN: comet-opt --opt-comp-workspace --convert-ta-to-it --convert-to-loops --convert-to-llvm %s &> eltwise_CSRxCSR_oCSR.llvm # RUN: export SPARSE_FILE_NAME0=%comet_integration_test_data_dir/test_rank2.mtx -# RUN: mlir-cpu-runner eltwise_CSRxCSR_oCSR.llvm -O3 -e main -entry-point-result=void -shared-libs=%comet_utility_library_dir/libcomet_runner_utils%shlibext | FileCheck %s +# RUN: mlir-cpu-runner eltwise_CSRxCSR_oCSR.llvm -O3 -e main -entry-point-result=void -shared-libs=%comet_utility_library_dir/libcomet_runner_utils%shlibext,%mlir_utility_library_dir/libomp%shlibext | FileCheck %s def main() { #IndexLabel Declarations @@ -28,7 +28,7 @@ def main() { # CHECK: data = # CHECK-NEXT: 5, # CHECK-NEXT: data = -# CHECK-NEXT: 0, +# CHECK-NEXT: -1, # CHECK-NEXT: data = # CHECK-NEXT: 0,2,4,5,7,9, # CHECK-NEXT: data = diff --git a/integration_test/ops/eltwise_mult_CSRxDense_oCSR.ta b/test/integration/ops/eltwise_mult_CSRxDense_oCSR.ta similarity index 90% rename from integration_test/ops/eltwise_mult_CSRxDense_oCSR.ta rename to test/integration/ops/eltwise_mult_CSRxDense_oCSR.ta index ab3cbc2f..34cd1314 100644 --- a/integration_test/ops/eltwise_mult_CSRxDense_oCSR.ta +++ b/test/integration/ops/eltwise_mult_CSRxDense_oCSR.ta @@ -1,6 +1,6 @@ # RUN: comet-opt --convert-ta-to-it --convert-to-loops --convert-to-llvm %s &> eltwise_CSRxDense_oCSR.llvm # RUN: export SPARSE_FILE_NAME0=%comet_integration_test_data_dir/test_rank2.mtx -# RUN: mlir-cpu-runner eltwise_CSRxDense_oCSR.llvm -O3 -e main -entry-point-result=void -shared-libs=%comet_utility_library_dir/libcomet_runner_utils%shlibext | FileCheck %s +# RUN: mlir-cpu-runner eltwise_CSRxDense_oCSR.llvm -O3 -e main -entry-point-result=void -shared-libs=%comet_utility_library_dir/libcomet_runner_utils%shlibext,%mlir_utility_library_dir/libomp%shlibext | FileCheck %s def main() { #IndexLabel Declarations @@ -26,7 +26,7 @@ def main() { # CHECK: data = # CHECK-NEXT: 5, # CHECK-NEXT: data = -# CHECK-NEXT: 0, +# CHECK-NEXT: -1, # CHECK-NEXT: data = # CHECK-NEXT: 0,2,4,5,7,9, # CHECK-NEXT: data = diff --git a/integration_test/ops/eltwise_mult_CSRxDense_oDense.ta b/test/integration/ops/eltwise_mult_CSRxDense_oDense.ta similarity index 90% rename from integration_test/ops/eltwise_mult_CSRxDense_oDense.ta rename to test/integration/ops/eltwise_mult_CSRxDense_oDense.ta index b6ac97fb..f7c8e99d 100644 --- a/integration_test/ops/eltwise_mult_CSRxDense_oDense.ta +++ b/test/integration/ops/eltwise_mult_CSRxDense_oDense.ta @@ -1,6 +1,6 @@ # RUN: comet-opt --convert-ta-to-it --convert-to-loops --convert-to-llvm %s &> eltwise_CSRxDense_oDense.llvm # RUN: export SPARSE_FILE_NAME0=%comet_integration_test_data_dir/test_rank2.mtx -# RUN: mlir-cpu-runner eltwise_CSRxDense_oDense.llvm -O3 -e main -entry-point-result=void -shared-libs=%comet_utility_library_dir/libcomet_runner_utils%shlibext | FileCheck %s +# RUN: mlir-cpu-runner eltwise_CSRxDense_oDense.llvm -O3 -e main -entry-point-result=void -shared-libs=%comet_utility_library_dir/libcomet_runner_utils%shlibext,%mlir_utility_library_dir/libomp%shlibext | FileCheck %s def main() { #IndexLabel Declarations diff --git a/integration_test/ops/eltwise_mult_DCSRxDense_oDense.ta b/test/integration/ops/eltwise_mult_DCSRxDense_oDense.ta similarity index 90% rename from integration_test/ops/eltwise_mult_DCSRxDense_oDense.ta rename to test/integration/ops/eltwise_mult_DCSRxDense_oDense.ta index 7edfe716..bfb570ae 100644 --- a/integration_test/ops/eltwise_mult_DCSRxDense_oDense.ta +++ b/test/integration/ops/eltwise_mult_DCSRxDense_oDense.ta @@ -1,6 +1,6 @@ # RUN: comet-opt --convert-ta-to-it --convert-to-loops --convert-to-llvm %s &> eltwise_DCSRxDense_oDense.llvm # RUN: export SPARSE_FILE_NAME0=%comet_integration_test_data_dir/test_rank2.mtx -# RUN: mlir-cpu-runner eltwise_DCSRxDense_oDense.llvm -O3 -e main -entry-point-result=void -shared-libs=%comet_utility_library_dir/libcomet_runner_utils%shlibext | FileCheck %s +# RUN: mlir-cpu-runner eltwise_DCSRxDense_oDense.llvm -O3 -e main -entry-point-result=void -shared-libs=%comet_utility_library_dir/libcomet_runner_utils%shlibext,%mlir_utility_library_dir/libomp%shlibext | FileCheck %s def main() { #IndexLabel Declarations diff --git a/integration_test/ops/eltwise_mult_DensexCSR_oDense.ta b/test/integration/ops/eltwise_mult_DensexCSR_oDense.ta similarity index 91% rename from integration_test/ops/eltwise_mult_DensexCSR_oDense.ta rename to test/integration/ops/eltwise_mult_DensexCSR_oDense.ta index 942927cc..59efb5ed 100644 --- a/integration_test/ops/eltwise_mult_DensexCSR_oDense.ta +++ b/test/integration/ops/eltwise_mult_DensexCSR_oDense.ta @@ -1,6 +1,6 @@ # RUN: comet-opt --convert-ta-to-it --convert-to-loops --convert-to-llvm %s &> eltwise_DensexCSR_oDense.llvm # RUN: export SPARSE_FILE_NAME0=%comet_integration_test_data_dir/test_rank2.mtx -# RUN: mlir-cpu-runner eltwise_DensexCSR_oDense.llvm -O3 -e main -entry-point-result=void -shared-libs=%comet_utility_library_dir/libcomet_runner_utils%shlibext | FileCheck %s +# RUN: mlir-cpu-runner eltwise_DensexCSR_oDense.llvm -O3 -e main -entry-point-result=void -shared-libs=%comet_utility_library_dir/libcomet_runner_utils%shlibext,%mlir_utility_library_dir/libomp%shlibext | FileCheck %s def main() { #IndexLabel Declarations diff --git a/integration_test/ops/eltwise_mult_DensexDense_oDense.ta b/test/integration/ops/eltwise_mult_DensexDense_oDense.ta similarity index 89% rename from integration_test/ops/eltwise_mult_DensexDense_oDense.ta rename to test/integration/ops/eltwise_mult_DensexDense_oDense.ta index c9321f01..cc484b8a 100644 --- a/integration_test/ops/eltwise_mult_DensexDense_oDense.ta +++ b/test/integration/ops/eltwise_mult_DensexDense_oDense.ta @@ -1,5 +1,5 @@ # RUN: comet-opt --convert-ta-to-it --convert-to-loops --convert-to-llvm %s &> eltwise_DensexDense_oDense.llvm -# RUN: mlir-cpu-runner eltwise_DensexDense_oDense.llvm -O3 -e main -entry-point-result=void -shared-libs=%comet_utility_library_dir/libcomet_runner_utils%shlibext | FileCheck %s +# RUN: mlir-cpu-runner eltwise_DensexDense_oDense.llvm -O3 -e main -entry-point-result=void -shared-libs=%comet_utility_library_dir/libcomet_runner_utils%shlibext,%mlir_utility_library_dir/libomp%shlibext | FileCheck %s def main() { diff --git a/integration_test/ops/eltwise_mult_dense_4Dtensors.ta b/test/integration/ops/eltwise_mult_dense_4Dtensors.ta similarity index 90% rename from integration_test/ops/eltwise_mult_dense_4Dtensors.ta rename to test/integration/ops/eltwise_mult_dense_4Dtensors.ta index 94f0190a..feac1423 100644 --- a/integration_test/ops/eltwise_mult_dense_4Dtensors.ta +++ b/test/integration/ops/eltwise_mult_dense_4Dtensors.ta @@ -1,5 +1,5 @@ # RUN: comet-opt --convert-ta-to-it --convert-to-loops --convert-to-llvm %s &> eltwise_dense_4Dtensors.llvm -# RUN: mlir-cpu-runner eltwise_dense_4Dtensors.llvm -O3 -e main -entry-point-result=void -shared-libs=%comet_utility_library_dir/libcomet_runner_utils%shlibext | FileCheck %s +# RUN: mlir-cpu-runner eltwise_dense_4Dtensors.llvm -O3 -e main -entry-point-result=void -shared-libs=%comet_utility_library_dir/libcomet_runner_utils%shlibext,%mlir_utility_library_dir/libomp%shlibext | FileCheck %s def main() { # IndexLabel Declarations diff --git a/integration_test/ops/eltwise_subtract_dense_matrix.ta b/test/integration/ops/eltwise_subtract_dense_matrix.ta similarity index 88% rename from integration_test/ops/eltwise_subtract_dense_matrix.ta rename to test/integration/ops/eltwise_subtract_dense_matrix.ta index cb81057c..26a0483b 100644 --- a/integration_test/ops/eltwise_subtract_dense_matrix.ta +++ b/test/integration/ops/eltwise_subtract_dense_matrix.ta @@ -1,5 +1,5 @@ # RUN: comet-opt --convert-ta-to-it --convert-to-loops --convert-to-llvm %s &> eltwise_sub_dense_matrix.llvm -# RUN: mlir-cpu-runner eltwise_sub_dense_matrix.llvm -O3 -e main -entry-point-result=void -shared-libs=%comet_utility_library_dir/libcomet_runner_utils%shlibext | FileCheck %s +# RUN: mlir-cpu-runner eltwise_sub_dense_matrix.llvm -O3 -e main -entry-point-result=void -shared-libs=%comet_utility_library_dir/libcomet_runner_utils%shlibext,%mlir_utility_library_dir/libomp%shlibext | FileCheck %s def main() { diff --git a/integration_test/ops/mult_DenseMatxCOO.ta b/test/integration/ops/mult_DenseMatxCOO.ta similarity index 92% rename from integration_test/ops/mult_DenseMatxCOO.ta rename to test/integration/ops/mult_DenseMatxCOO.ta index 666709a5..28770cef 100644 --- a/integration_test/ops/mult_DenseMatxCOO.ta +++ b/test/integration/ops/mult_DenseMatxCOO.ta @@ -2,7 +2,7 @@ # Sparse matrix is in COO format # RUN: comet-opt --convert-ta-to-it --convert-to-loops --convert-to-llvm %s &> mult_DenseMatxCOO.llvm # RUN: export SPARSE_FILE_NAME0=%comet_integration_test_data_dir/test_rank2.mtx -# RUN: mlir-cpu-runner mult_DenseMatxCOO.llvm -O3 -e main -entry-point-result=void -shared-libs=%comet_utility_library_dir/libcomet_runner_utils%shlibext | FileCheck %s +# RUN: mlir-cpu-runner mult_DenseMatxCOO.llvm -O3 -e main -entry-point-result=void -shared-libs=%comet_utility_library_dir/libcomet_runner_utils%shlibext,%mlir_utility_library_dir/libomp%shlibext | FileCheck %s def main() { #IndexLabel Declarations diff --git a/integration_test/ops/mult_DenseMatxCSR.ta b/test/integration/ops/mult_DenseMatxCSR.ta similarity index 92% rename from integration_test/ops/mult_DenseMatxCSR.ta rename to test/integration/ops/mult_DenseMatxCSR.ta index c6ba0f37..af11df30 100644 --- a/integration_test/ops/mult_DenseMatxCSR.ta +++ b/test/integration/ops/mult_DenseMatxCSR.ta @@ -2,7 +2,7 @@ # Sparse matrix is in CSR format # RUN: comet-opt --convert-ta-to-it --convert-to-loops --convert-to-llvm %s &> mult_DenseMatxCSR.llvm # RUN: export SPARSE_FILE_NAME0=%comet_integration_test_data_dir/test_rank2.mtx -# RUN: mlir-cpu-runner mult_DenseMatxCSR.llvm -O3 -e main -entry-point-result=void -shared-libs=%comet_utility_library_dir/libcomet_runner_utils%shlibext | FileCheck %s +# RUN: mlir-cpu-runner mult_DenseMatxCSR.llvm -O3 -e main -entry-point-result=void -shared-libs=%comet_utility_library_dir/libcomet_runner_utils%shlibext,%mlir_utility_library_dir/libomp%shlibext | FileCheck %s def main() { #IndexLabel Declarations diff --git a/integration_test/ops/mult_DenseMatxDCSR.ta b/test/integration/ops/mult_DenseMatxDCSR.ta similarity index 92% rename from integration_test/ops/mult_DenseMatxDCSR.ta rename to test/integration/ops/mult_DenseMatxDCSR.ta index a08aa718..8f9af07f 100644 --- a/integration_test/ops/mult_DenseMatxDCSR.ta +++ b/test/integration/ops/mult_DenseMatxDCSR.ta @@ -2,7 +2,7 @@ # Sparse matrix is in DCSR format # RUN: comet-opt --convert-ta-to-it --convert-to-loops --convert-to-llvm %s &> mult_DenseMatxDCSR.llvm # RUN: export SPARSE_FILE_NAME0=%comet_integration_test_data_dir/test_rank2.mtx -# RUN: mlir-cpu-runner mult_DenseMatxDCSR.llvm -O3 -e main -entry-point-result=void -shared-libs=%comet_utility_library_dir/libcomet_runner_utils%shlibext | FileCheck %s +# RUN: mlir-cpu-runner mult_DenseMatxDCSR.llvm -O3 -e main -entry-point-result=void -shared-libs=%comet_utility_library_dir/libcomet_runner_utils%shlibext,%mlir_utility_library_dir/libomp%shlibext | FileCheck %s def main() { #IndexLabel Declarations diff --git a/integration_test/ops/mult_DenseVecxCOO.ta b/test/integration/ops/mult_DenseVecxCOO.ta similarity index 90% rename from integration_test/ops/mult_DenseVecxCOO.ta rename to test/integration/ops/mult_DenseVecxCOO.ta index 9a3deb4d..66afdf4f 100644 --- a/integration_test/ops/mult_DenseVecxCOO.ta +++ b/test/integration/ops/mult_DenseVecxCOO.ta @@ -2,7 +2,7 @@ # Sparse matrix is in COO format # RUN: comet-opt --convert-ta-to-it --convert-to-loops --convert-to-llvm %s &> mult_DenseVecxCOO.llvm # RUN: export SPARSE_FILE_NAME0=%comet_integration_test_data_dir/test_rank2.mtx -# RUN: mlir-cpu-runner mult_DenseVecxCOO.llvm -O3 -e main -entry-point-result=void -shared-libs=%comet_utility_library_dir/libcomet_runner_utils%shlibext | FileCheck %s +# RUN: mlir-cpu-runner mult_DenseVecxCOO.llvm -O3 -e main -entry-point-result=void -shared-libs=%comet_utility_library_dir/libcomet_runner_utils%shlibext,%mlir_utility_library_dir/libomp%shlibext | FileCheck %s def main() { diff --git a/integration_test/ops/mult_DenseVecxCSR.ta b/test/integration/ops/mult_DenseVecxCSR.ta similarity index 90% rename from integration_test/ops/mult_DenseVecxCSR.ta rename to test/integration/ops/mult_DenseVecxCSR.ta index 8e875943..d52cb8b9 100644 --- a/integration_test/ops/mult_DenseVecxCSR.ta +++ b/test/integration/ops/mult_DenseVecxCSR.ta @@ -2,7 +2,7 @@ # Sparse matrix is in CSR format # RUN: comet-opt --convert-ta-to-it --convert-to-loops --convert-to-llvm %s &> mult_DenseVecxCSR.llvm # RUN: export SPARSE_FILE_NAME0=%comet_integration_test_data_dir/test_rank2.mtx -# RUN: mlir-cpu-runner mult_DenseVecxCSR.llvm -O3 -e main -entry-point-result=void -shared-libs=%comet_utility_library_dir/libcomet_runner_utils%shlibext | FileCheck %s +# RUN: mlir-cpu-runner mult_DenseVecxCSR.llvm -O3 -e main -entry-point-result=void -shared-libs=%comet_utility_library_dir/libcomet_runner_utils%shlibext,%mlir_utility_library_dir/libomp%shlibext | FileCheck %s def main() { diff --git a/integration_test/ops/mult_DenseVecxDCSR.ta b/test/integration/ops/mult_DenseVecxDCSR.ta similarity index 90% rename from integration_test/ops/mult_DenseVecxDCSR.ta rename to test/integration/ops/mult_DenseVecxDCSR.ta index e7dc759a..037a211b 100644 --- a/integration_test/ops/mult_DenseVecxDCSR.ta +++ b/test/integration/ops/mult_DenseVecxDCSR.ta @@ -2,7 +2,7 @@ # Sparse matrix is in DCSR format # RUN: comet-opt --convert-ta-to-it --convert-to-loops --convert-to-llvm %s &> mult_DenseVecxDCSR.llvm # RUN: export SPARSE_FILE_NAME0=%comet_integration_test_data_dir/test_rank2.mtx -# RUN: mlir-cpu-runner mult_DenseVecxDCSR.llvm -O3 -e main -entry-point-result=void -shared-libs=%comet_utility_library_dir/libcomet_runner_utils%shlibext | FileCheck %s +# RUN: mlir-cpu-runner mult_DenseVecxDCSR.llvm -O3 -e main -entry-point-result=void -shared-libs=%comet_utility_library_dir/libcomet_runner_utils%shlibext,%mlir_utility_library_dir/libomp%shlibext | FileCheck %s def main() { diff --git a/integration_test/ops/mult_dense_4Dtensors.ta b/test/integration/ops/mult_dense_4Dtensors.ta similarity index 91% rename from integration_test/ops/mult_dense_4Dtensors.ta rename to test/integration/ops/mult_dense_4Dtensors.ta index 27f302a2..f755e693 100644 --- a/integration_test/ops/mult_dense_4Dtensors.ta +++ b/test/integration/ops/mult_dense_4Dtensors.ta @@ -1,5 +1,5 @@ # RUN: comet-opt --convert-ta-to-it --convert-to-loops --convert-to-llvm %s &> mult_dense_4Dtensors.llvm -# RUN: mlir-cpu-runner mult_dense_4Dtensors.llvm -O3 -e main -entry-point-result=void -shared-libs=%comet_utility_library_dir/libcomet_runner_utils%shlibext | FileCheck %s +# RUN: mlir-cpu-runner mult_dense_4Dtensors.llvm -O3 -e main -entry-point-result=void -shared-libs=%comet_utility_library_dir/libcomet_runner_utils%shlibext,%mlir_utility_library_dir/libomp%shlibext | FileCheck %s # This example comes from TCCG - ccsd12 def main() { diff --git a/integration_test/ops/mult_dense_ij-ikj-kj.ta b/test/integration/ops/mult_dense_ij-ikj-kj.ta similarity index 92% rename from integration_test/ops/mult_dense_ij-ikj-kj.ta rename to test/integration/ops/mult_dense_ij-ikj-kj.ta index a3a443e0..cde1e858 100644 --- a/integration_test/ops/mult_dense_ij-ikj-kj.ta +++ b/test/integration/ops/mult_dense_ij-ikj-kj.ta @@ -2,7 +2,7 @@ # No assumption that contraction indices should disapper in the output tensor. # RUN: comet-opt --convert-ta-to-it --convert-to-loops --convert-to-llvm %s &> mult_dense_ij-ikj-kj.llvm -# RUN: mlir-cpu-runner mult_dense_ij-ikj-kj.llvm -O3 -e main -entry-point-result=void -shared-libs=%comet_utility_library_dir/libcomet_runner_utils%shlibext | FileCheck %s +# RUN: mlir-cpu-runner mult_dense_ij-ikj-kj.llvm -O3 -e main -entry-point-result=void -shared-libs=%comet_utility_library_dir/libcomet_runner_utils%shlibext,%mlir_utility_library_dir/libomp%shlibext | FileCheck %s def main() { #IndexLabel Declarations diff --git a/integration_test/ops/mult_dense_matrix.ta b/test/integration/ops/mult_dense_matrix.ta similarity index 90% rename from integration_test/ops/mult_dense_matrix.ta rename to test/integration/ops/mult_dense_matrix.ta index 8222a267..4c3f9f19 100644 --- a/integration_test/ops/mult_dense_matrix.ta +++ b/test/integration/ops/mult_dense_matrix.ta @@ -1,5 +1,5 @@ # RUN: comet-opt --convert-ta-to-it --convert-to-loops --convert-to-llvm %s &> mult_dense_matrix.llvm -# RUN: mlir-cpu-runner mult_dense_matrix.llvm -O3 -e main -entry-point-result=void -shared-libs=%comet_utility_library_dir/libcomet_runner_utils%shlibext | FileCheck %s +# RUN: mlir-cpu-runner mult_dense_matrix.llvm -O3 -e main -entry-point-result=void -shared-libs=%comet_utility_library_dir/libcomet_runner_utils%shlibext,%mlir_utility_library_dir/libomp%shlibext | FileCheck %s def main() { diff --git a/integration_test/ops/mult_dense_matrix_vector.ta b/test/integration/ops/mult_dense_matrix_vector.ta similarity index 88% rename from integration_test/ops/mult_dense_matrix_vector.ta rename to test/integration/ops/mult_dense_matrix_vector.ta index 63aa9a28..b27ee7c6 100644 --- a/integration_test/ops/mult_dense_matrix_vector.ta +++ b/test/integration/ops/mult_dense_matrix_vector.ta @@ -1,5 +1,5 @@ # RUN: comet-opt --convert-ta-to-it --convert-to-loops --convert-to-llvm %s &> mult_dense_matrix_vector.llvm -# RUN: mlir-cpu-runner mult_dense_matrix_vector.llvm -O3 -e main -entry-point-result=void -shared-libs=%comet_utility_library_dir/libcomet_runner_utils%shlibext | FileCheck %s +# RUN: mlir-cpu-runner mult_dense_matrix_vector.llvm -O3 -e main -entry-point-result=void -shared-libs=%comet_utility_library_dir/libcomet_runner_utils%shlibext,%mlir_utility_library_dir/libomp%shlibext | FileCheck %s def main() { #IndexLabel Declarations diff --git a/integration_test/ops/mult_spgemm_CSRxCSR_oCSR.ta b/test/integration/ops/mult_spgemm_CSRxCSR_oCSR.ta similarity index 92% rename from integration_test/ops/mult_spgemm_CSRxCSR_oCSR.ta rename to test/integration/ops/mult_spgemm_CSRxCSR_oCSR.ta index 785fbff1..776d25d2 100644 --- a/integration_test/ops/mult_spgemm_CSRxCSR_oCSR.ta +++ b/test/integration/ops/mult_spgemm_CSRxCSR_oCSR.ta @@ -3,7 +3,7 @@ # RUN: comet-opt --opt-comp-workspace --convert-ta-to-it --convert-to-loops --convert-to-llvm %s &> mult_spgemm_CSRxCSR_oCSR.llvm # RUN: export SPARSE_FILE_NAME0=%comet_integration_test_data_dir/test_rank2.mtx # RUN: export SPARSE_FILE_NAME1=%comet_integration_test_data_dir/test_rank2.mtx -# RUN: mlir-cpu-runner mult_spgemm_CSRxCSR_oCSR.llvm -O3 -e main -entry-point-result=void -shared-libs=%comet_utility_library_dir/libcomet_runner_utils%shlibext | FileCheck %s +# RUN: mlir-cpu-runner mult_spgemm_CSRxCSR_oCSR.llvm -O3 -e main -entry-point-result=void -shared-libs=%comet_utility_library_dir/libcomet_runner_utils%shlibext,%mlir_utility_library_dir/libomp%shlibext | FileCheck %s def main() { @@ -30,7 +30,7 @@ def main() { # CHECK: data = # CHECK-NEXT: 5, # CHECK-NEXT: data = -# CHECK-NEXT: 0, +# CHECK-NEXT: -1, # CHECK-NEXT: data = # CHECK-NEXT: 0,2,4,5,7,9, # CHECK-NEXT: data = diff --git a/integration_test/ops/mult_spgemm_CSRxCSR_oCSR_wMasking.ta b/test/integration/ops/mult_spgemm_CSRxCSR_oCSR_wMasking.ta similarity index 92% rename from integration_test/ops/mult_spgemm_CSRxCSR_oCSR_wMasking.ta rename to test/integration/ops/mult_spgemm_CSRxCSR_oCSR_wMasking.ta index 6e937df7..3362bf76 100644 --- a/integration_test/ops/mult_spgemm_CSRxCSR_oCSR_wMasking.ta +++ b/test/integration/ops/mult_spgemm_CSRxCSR_oCSR_wMasking.ta @@ -3,7 +3,7 @@ # RUN: comet-opt --opt-comp-workspace --convert-ta-to-it --convert-to-loops --convert-to-llvm %s &> mult_spgemm_CSRxCSR_oCSR.mask.llvm # RUN: export SPARSE_FILE_NAME0=%comet_integration_test_data_dir/test_rank2.mtx # RUN: export SPARSE_FILE_NAME1=%comet_integration_test_data_dir/test_rank2.mtx -# RUN: mlir-cpu-runner mult_spgemm_CSRxCSR_oCSR.mask.llvm -O3 -e main -entry-point-result=void -shared-libs=%comet_utility_library_dir/libcomet_runner_utils%shlibext | FileCheck %s +# RUN: mlir-cpu-runner mult_spgemm_CSRxCSR_oCSR.mask.llvm -O3 -e main -entry-point-result=void -shared-libs=%comet_utility_library_dir/libcomet_runner_utils%shlibext,%mlir_utility_library_dir/libomp%shlibext | FileCheck %s def main() { @@ -31,7 +31,7 @@ def main() { # CHECK: data = # CHECK-NEXT: 5, # CHECK-NEXT: data = -# CHECK-NEXT: 0, +# CHECK-NEXT: -1, # CHECK-NEXT: data = # CHECK-NEXT: 0,2,4,5,7,9, # CHECK-NEXT: data = diff --git a/integration_test/ops/mult_spmm_COOxDense.ta b/test/integration/ops/mult_spmm_COOxDense.ta similarity index 91% rename from integration_test/ops/mult_spmm_COOxDense.ta rename to test/integration/ops/mult_spmm_COOxDense.ta index 94f4eb0b..9398f12f 100644 --- a/integration_test/ops/mult_spmm_COOxDense.ta +++ b/test/integration/ops/mult_spmm_COOxDense.ta @@ -2,7 +2,7 @@ # Sparse matrix is in COO format # RUN: comet-opt --convert-ta-to-it --convert-to-loops --convert-to-llvm %s &> mult_spmm_COOxDense.llvm # RUN: export SPARSE_FILE_NAME0=%comet_integration_test_data_dir/test_rank2.mtx -# RUN: mlir-cpu-runner mult_spmm_COOxDense.llvm -O3 -e main -entry-point-result=void -shared-libs=%comet_utility_library_dir/libcomet_runner_utils%shlibext | FileCheck %s +# RUN: mlir-cpu-runner mult_spmm_COOxDense.llvm -O3 -e main -entry-point-result=void -shared-libs=%comet_utility_library_dir/libcomet_runner_utils%shlibext,%mlir_utility_library_dir/libomp%shlibext | FileCheck %s def main() { #IndexLabel Declarations diff --git a/integration_test/ops/mult_spmm_CSRxDense.ta b/test/integration/ops/mult_spmm_CSRxDense.ta similarity index 91% rename from integration_test/ops/mult_spmm_CSRxDense.ta rename to test/integration/ops/mult_spmm_CSRxDense.ta index 94010458..9fbf73be 100644 --- a/integration_test/ops/mult_spmm_CSRxDense.ta +++ b/test/integration/ops/mult_spmm_CSRxDense.ta @@ -2,7 +2,7 @@ # Sparse matrix is in CSR format # RUN: comet-opt --convert-ta-to-it --convert-to-loops --convert-to-llvm %s &> mult_spmm_CSRxDense.llvm # RUN: export SPARSE_FILE_NAME0=%comet_integration_test_data_dir/test_rank2.mtx -# RUN: mlir-cpu-runner mult_spmm_CSRxDense.llvm -O3 -e main -entry-point-result=void -shared-libs=%comet_utility_library_dir/libcomet_runner_utils%shlibext | FileCheck %s +# RUN: mlir-cpu-runner mult_spmm_CSRxDense.llvm -O3 -e main -entry-point-result=void -shared-libs=%comet_utility_library_dir/libcomet_runner_utils%shlibext,%mlir_utility_library_dir/libomp%shlibext | FileCheck %s def main() { #IndexLabel Declarations diff --git a/integration_test/ops/mult_spmm_DCSRxDense.ta b/test/integration/ops/mult_spmm_DCSRxDense.ta similarity index 91% rename from integration_test/ops/mult_spmm_DCSRxDense.ta rename to test/integration/ops/mult_spmm_DCSRxDense.ta index ed3142db..166fc957 100644 --- a/integration_test/ops/mult_spmm_DCSRxDense.ta +++ b/test/integration/ops/mult_spmm_DCSRxDense.ta @@ -2,7 +2,7 @@ # Sparse matrix is in DCSR format # RUN: comet-opt --convert-ta-to-it --convert-to-loops --convert-to-llvm %s &> mult_spmm_DCSRxDense.llvm # RUN: export SPARSE_FILE_NAME0=%comet_integration_test_data_dir/test_rank2.mtx -# RUN: mlir-cpu-runner mult_spmm_DCSRxDense.llvm -O3 -e main -entry-point-result=void -shared-libs=%comet_utility_library_dir/libcomet_runner_utils%shlibext | FileCheck %s +# RUN: mlir-cpu-runner mult_spmm_DCSRxDense.llvm -O3 -e main -entry-point-result=void -shared-libs=%comet_utility_library_dir/libcomet_runner_utils%shlibext,%mlir_utility_library_dir/libomp%shlibext | FileCheck %s def main() { #IndexLabel Declarations diff --git a/integration_test/ops/mult_spmv_COOxDense.ta b/test/integration/ops/mult_spmv_COOxDense.ta similarity index 90% rename from integration_test/ops/mult_spmv_COOxDense.ta rename to test/integration/ops/mult_spmv_COOxDense.ta index 72443821..8bb19c9a 100644 --- a/integration_test/ops/mult_spmv_COOxDense.ta +++ b/test/integration/ops/mult_spmv_COOxDense.ta @@ -2,7 +2,7 @@ # Sparse matrix is in COO format # RUN: comet-opt --convert-ta-to-it --convert-to-loops --convert-to-llvm %s &> mult_spmv_COOxDense.llvm # RUN: export SPARSE_FILE_NAME0=%comet_integration_test_data_dir/test_rank2.mtx -# RUN: mlir-cpu-runner mult_spmv_COOxDense.llvm -O3 -e main -entry-point-result=void -shared-libs=%comet_utility_library_dir/libcomet_runner_utils%shlibext | FileCheck %s +# RUN: mlir-cpu-runner mult_spmv_COOxDense.llvm -O3 -e main -entry-point-result=void -shared-libs=%comet_utility_library_dir/libcomet_runner_utils%shlibext,%mlir_utility_library_dir/libomp%shlibext | FileCheck %s def main() { diff --git a/integration_test/ops/mult_spmv_CSRxDense.ta b/test/integration/ops/mult_spmv_CSRxDense.ta similarity index 90% rename from integration_test/ops/mult_spmv_CSRxDense.ta rename to test/integration/ops/mult_spmv_CSRxDense.ta index e81c0ce1..490a9375 100644 --- a/integration_test/ops/mult_spmv_CSRxDense.ta +++ b/test/integration/ops/mult_spmv_CSRxDense.ta @@ -2,7 +2,7 @@ # Sparse matrix is in CSR format # RUN: comet-opt --convert-ta-to-it --convert-to-loops --convert-to-llvm %s &> mult_spmv_CSRxDense.llvm # RUN: export SPARSE_FILE_NAME0=%comet_integration_test_data_dir/test_rank2.mtx -# RUN: mlir-cpu-runner mult_spmv_CSRxDense.llvm -O3 -e main -entry-point-result=void -shared-libs=%comet_utility_library_dir/libcomet_runner_utils%shlibext | FileCheck %s +# RUN: mlir-cpu-runner mult_spmv_CSRxDense.llvm -O3 -e main -entry-point-result=void -shared-libs=%comet_utility_library_dir/libcomet_runner_utils%shlibext,%mlir_utility_library_dir/libomp%shlibext | FileCheck %s def main() { diff --git a/integration_test/ops/mult_spmv_DCSRxDense.ta b/test/integration/ops/mult_spmv_DCSRxDense.ta similarity index 90% rename from integration_test/ops/mult_spmv_DCSRxDense.ta rename to test/integration/ops/mult_spmv_DCSRxDense.ta index f62b1e9c..c2e604eb 100644 --- a/integration_test/ops/mult_spmv_DCSRxDense.ta +++ b/test/integration/ops/mult_spmv_DCSRxDense.ta @@ -2,7 +2,7 @@ # Sparse matrix is in DCSR format # RUN: comet-opt --convert-ta-to-it --convert-to-loops --convert-to-llvm %s &> mult_spmv_DCSRxDense.llvm # RUN: export SPARSE_FILE_NAME0=%comet_integration_test_data_dir/test_rank2.mtx -# RUN: mlir-cpu-runner mult_spmv_DCSRxDense.llvm -O3 -e main -entry-point-result=void -shared-libs=%comet_utility_library_dir/libcomet_runner_utils%shlibext | FileCheck %s +# RUN: mlir-cpu-runner mult_spmv_DCSRxDense.llvm -O3 -e main -entry-point-result=void -shared-libs=%comet_utility_library_dir/libcomet_runner_utils%shlibext,%mlir_utility_library_dir/libomp%shlibext | FileCheck %s def main() { diff --git a/integration_test/ops/sum_COO.ta b/test/integration/ops/sum_COO.ta similarity index 80% rename from integration_test/ops/sum_COO.ta rename to test/integration/ops/sum_COO.ta index e501363c..ded23bba 100644 --- a/integration_test/ops/sum_COO.ta +++ b/test/integration/ops/sum_COO.ta @@ -1,6 +1,6 @@ # RUN: comet-opt --convert-to-loops --convert-to-llvm %s &> sum_COO.llvm # RUN: export SPARSE_FILE_NAME0=%comet_integration_test_data_dir/test_rank2.mtx -# RUN: mlir-cpu-runner sum_COO.llvm -O3 -e main -entry-point-result=void -shared-libs=%comet_utility_library_dir/libcomet_runner_utils%shlibext | FileCheck %s +# RUN: mlir-cpu-runner sum_COO.llvm -O3 -e main -entry-point-result=void -shared-libs=%comet_utility_library_dir/libcomet_runner_utils%shlibext,%mlir_utility_library_dir/libmlir_c_runner_utils%shlibext,%mlir_utility_library_dir/libomp%shlibext | FileCheck %s def main() { #IndexLabel Declarations @@ -19,5 +19,4 @@ def main() { } # Print the result for verification. -# CHECK: data = -# CHECK-NEXT: 28.2, +# CHECK: 28.2 diff --git a/integration_test/ops/sum_CSF.ta b/test/integration/ops/sum_CSF.ta similarity index 81% rename from integration_test/ops/sum_CSF.ta rename to test/integration/ops/sum_CSF.ta index 157caa13..c3ae33c4 100644 --- a/integration_test/ops/sum_CSF.ta +++ b/test/integration/ops/sum_CSF.ta @@ -1,6 +1,6 @@ # RUN: comet-opt --convert-to-loops --convert-to-llvm %s &> sum_CSF.llvm # RUN: export SPARSE_FILE_NAME0=%comet_integration_test_data_dir/test_rank3.tns -# RUN: mlir-cpu-runner sum_CSF.llvm -O3 -e main -entry-point-result=void -shared-libs=%comet_utility_library_dir/libcomet_runner_utils%shlibext | FileCheck %s +# RUN: mlir-cpu-runner sum_CSF.llvm -O3 -e main -entry-point-result=void -shared-libs=%comet_utility_library_dir/libcomet_runner_utils%shlibext,%mlir_utility_library_dir/libmlir_c_runner_utils%shlibext,%mlir_utility_library_dir/libomp%shlibext | FileCheck %s def main() { @@ -21,5 +21,4 @@ def main() { } # Print the result for verification. -# CHECK: data = -# CHECK-NEXT: 6.41, \ No newline at end of file +# CHECK: 6.41 \ No newline at end of file diff --git a/integration_test/ops/sum_CSR.ta b/test/integration/ops/sum_CSR.ta similarity index 80% rename from integration_test/ops/sum_CSR.ta rename to test/integration/ops/sum_CSR.ta index fc0e1dda..3f816088 100644 --- a/integration_test/ops/sum_CSR.ta +++ b/test/integration/ops/sum_CSR.ta @@ -1,6 +1,6 @@ # RUN: comet-opt --convert-to-loops --convert-to-llvm %s &> sum_CSR.llvm # RUN: export SPARSE_FILE_NAME0=%comet_integration_test_data_dir/test_rank2.mtx -# RUN: mlir-cpu-runner sum_CSR.llvm -O3 -e main -entry-point-result=void -shared-libs=%comet_utility_library_dir/libcomet_runner_utils%shlibext | FileCheck %s +# RUN: mlir-cpu-runner sum_CSR.llvm -O3 -e main -entry-point-result=void -shared-libs=%comet_utility_library_dir/libcomet_runner_utils%shlibext,%mlir_utility_library_dir/libmlir_c_runner_utils%shlibext,%mlir_utility_library_dir/libomp%shlibext | FileCheck %s def main() { #IndexLabel Declarations @@ -19,5 +19,4 @@ def main() { } # Print the result for verification. -# CHECK: data = -# CHECK-NEXT: 28.2, \ No newline at end of file +# CHECK: 28.2 \ No newline at end of file diff --git a/integration_test/ops/sum_dense_matrix.ta b/test/integration/ops/sum_dense_matrix.ta similarity index 82% rename from integration_test/ops/sum_dense_matrix.ta rename to test/integration/ops/sum_dense_matrix.ta index 591c9f04..75786054 100644 --- a/integration_test/ops/sum_dense_matrix.ta +++ b/test/integration/ops/sum_dense_matrix.ta @@ -1,5 +1,5 @@ # RUN: comet-opt --convert-to-loops --convert-to-llvm %s &> sum_dense_matrix.llvm -# RUN: mlir-cpu-runner sum_dense_matrix.llvm -O3 -e main -entry-point-result=void -shared-libs=%comet_utility_library_dir/libcomet_runner_utils%shlibext | FileCheck %s +# RUN: mlir-cpu-runner sum_dense_matrix.llvm -O3 -e main -entry-point-result=void -shared-libs=%comet_utility_library_dir/libcomet_runner_utils%shlibext,%mlir_utility_library_dir/libmlir_c_runner_utils%shlibext | FileCheck %s def main() { #IndexLabel Declarations @@ -18,5 +18,4 @@ def main() { } # Print the result for verification. -# CHECK: data = -# CHECK-NEXT: 59.2, \ No newline at end of file +# CHECK: 59.2 \ No newline at end of file diff --git a/integration_test/ops/sum_dense_tensor.ta b/test/integration/ops/sum_dense_tensor.ta similarity index 83% rename from integration_test/ops/sum_dense_tensor.ta rename to test/integration/ops/sum_dense_tensor.ta index 0ae67109..cb8b4e99 100644 --- a/integration_test/ops/sum_dense_tensor.ta +++ b/test/integration/ops/sum_dense_tensor.ta @@ -1,5 +1,5 @@ # RUN: comet-opt --convert-to-loops --convert-to-llvm %s &> sum_dense_tensor.llvm -# RUN: mlir-cpu-runner sum_dense_tensor.llvm -O3 -e main -entry-point-result=void -shared-libs=%comet_utility_library_dir/libcomet_runner_utils%shlibext | FileCheck %s +# RUN: mlir-cpu-runner sum_dense_tensor.llvm -O3 -e main -entry-point-result=void -shared-libs=%comet_utility_library_dir/libcomet_runner_utils%shlibext,%mlir_utility_library_dir/libmlir_c_runner_utils%shlibext | FileCheck %s def main() { #IndexLabel Declarations @@ -19,5 +19,4 @@ def main() { } # Print the result for verification. -# CHECK: data = -# CHECK-NEXT: 236.8, +# CHECK: 236.8 diff --git a/integration_test/ops/transpose_dense_matrix.ta b/test/integration/ops/transpose_dense_matrix.ta similarity index 100% rename from integration_test/ops/transpose_dense_matrix.ta rename to test/integration/ops/transpose_dense_matrix.ta diff --git a/integration_test/ops/transpose_dense_tensor.ta b/test/integration/ops/transpose_dense_tensor.ta similarity index 100% rename from integration_test/ops/transpose_dense_tensor.ta rename to test/integration/ops/transpose_dense_tensor.ta diff --git a/integration_test/ops/utility_getTime.ta b/test/integration/ops/utility_getTime.ta similarity index 100% rename from integration_test/ops/utility_getTime.ta rename to test/integration/ops/utility_getTime.ta diff --git a/integration_test/ops/utility_printCOO.ta b/test/integration/ops/utility_printCOO.ta similarity index 100% rename from integration_test/ops/utility_printCOO.ta rename to test/integration/ops/utility_printCOO.ta diff --git a/integration_test/ops/utility_printCOO_lowerTri.ta b/test/integration/ops/utility_printCOO_lowerTri.ta similarity index 100% rename from integration_test/ops/utility_printCOO_lowerTri.ta rename to test/integration/ops/utility_printCOO_lowerTri.ta diff --git a/integration_test/ops/utility_printCOO_lowerTriStrict.ta b/test/integration/ops/utility_printCOO_lowerTriStrict.ta similarity index 100% rename from integration_test/ops/utility_printCOO_lowerTriStrict.ta rename to test/integration/ops/utility_printCOO_lowerTriStrict.ta diff --git a/integration_test/ops/utility_printCOO_multi.ta b/test/integration/ops/utility_printCOO_multi.ta similarity index 100% rename from integration_test/ops/utility_printCOO_multi.ta rename to test/integration/ops/utility_printCOO_multi.ta diff --git a/integration_test/ops/utility_printCOO_upperTri.ta b/test/integration/ops/utility_printCOO_upperTri.ta similarity index 100% rename from integration_test/ops/utility_printCOO_upperTri.ta rename to test/integration/ops/utility_printCOO_upperTri.ta diff --git a/integration_test/ops/utility_printCOO_upperTriStrict.ta b/test/integration/ops/utility_printCOO_upperTriStrict.ta similarity index 100% rename from integration_test/ops/utility_printCOO_upperTriStrict.ta rename to test/integration/ops/utility_printCOO_upperTriStrict.ta diff --git a/integration_test/ops/utility_printCSF_multi.ta b/test/integration/ops/utility_printCSF_multi.ta similarity index 100% rename from integration_test/ops/utility_printCSF_multi.ta rename to test/integration/ops/utility_printCSF_multi.ta diff --git a/integration_test/ops/utility_printCSR.ta b/test/integration/ops/utility_printCSR.ta similarity index 100% rename from integration_test/ops/utility_printCSR.ta rename to test/integration/ops/utility_printCSR.ta diff --git a/integration_test/ops/utility_printCSR_CSF.ta b/test/integration/ops/utility_printCSR_CSF.ta similarity index 100% rename from integration_test/ops/utility_printCSR_CSF.ta rename to test/integration/ops/utility_printCSR_CSF.ta diff --git a/integration_test/ops/utility_printCSR_lowerTri.ta b/test/integration/ops/utility_printCSR_lowerTri.ta similarity index 100% rename from integration_test/ops/utility_printCSR_lowerTri.ta rename to test/integration/ops/utility_printCSR_lowerTri.ta diff --git a/integration_test/ops/utility_printCSR_lowerTriStrict.ta b/test/integration/ops/utility_printCSR_lowerTriStrict.ta similarity index 100% rename from integration_test/ops/utility_printCSR_lowerTriStrict.ta rename to test/integration/ops/utility_printCSR_lowerTriStrict.ta diff --git a/integration_test/ops/utility_printCSR_multi.ta b/test/integration/ops/utility_printCSR_multi.ta similarity index 100% rename from integration_test/ops/utility_printCSR_multi.ta rename to test/integration/ops/utility_printCSR_multi.ta diff --git a/integration_test/ops/utility_printCSR_upperTri.ta b/test/integration/ops/utility_printCSR_upperTri.ta similarity index 100% rename from integration_test/ops/utility_printCSR_upperTri.ta rename to test/integration/ops/utility_printCSR_upperTri.ta diff --git a/integration_test/ops/utility_printCSR_upperTriStrict.ta b/test/integration/ops/utility_printCSR_upperTriStrict.ta similarity index 100% rename from integration_test/ops/utility_printCSR_upperTriStrict.ta rename to test/integration/ops/utility_printCSR_upperTriStrict.ta diff --git a/integration_test/ops/utility_printDense.ta b/test/integration/ops/utility_printDense.ta similarity index 100% rename from integration_test/ops/utility_printDense.ta rename to test/integration/ops/utility_printDense.ta diff --git a/test/integration/opts/ccsd_t1_11_all_opts.ta b/test/integration/opts/ccsd_t1_11_all_opts.ta new file mode 100644 index 00000000..af2cc9ee --- /dev/null +++ b/test/integration/opts/ccsd_t1_11_all_opts.ta @@ -0,0 +1,26 @@ + +# RUN: comet-opt --opt-matmul-tiling --opt-matmul-mkernel --opt-dense-transpose --convert-tc-to-ttgt --opt-multiop-factorize --convert-to-loops --convert-to-llvm %s &> ccsd_t1_11.llvm +# RUN: mlir-cpu-runner ccsd_t1_11.llvm -O3 -e main -entry-point-result=void -shared-libs=%comet_utility_library_dir/libcomet_runner_utils%shlibext,%mlir_utility_library_dir/libomp%shlibext | FileCheck %s + +def main() { + # IndexLabel Declarations + IndexLabel [i, j, m, n] = [8]; + IndexLabel [a, b, c, d] = [8]; + + # Tensor Declarations + Tensor i0([i, a], {Dense});; + Tensor t1([i, a], {Dense});; + Tensor t2([i, j, a, b], {Dense});; + Tensor v([a, b, i, j], {Dense});; + + i0[i, a] = 0.0; + t1[i, a] = 2.6; + t2[i, j, a, b] = 3.5; + v[i, a, j, b] = 7.3; + + i0[i, a] = v[c, d, m, n] * t2[i, n, a, d] * t1[m, c]; # 11 + print(i0[i, a]); +} + +# CHECK: data = +# CHECK: 272097,272097,272097,272097,272097,272097,272097,272097,272097,272097,272097,272097,272097,272097,272097,272097,272097,272097,272097,272097,272097,272097,272097,272097,272097,272097,272097,272097,272097,272097,272097,272097,272097,272097,272097,272097,272097,272097,272097,272097,272097,272097,272097,272097,272097,272097,272097,272097,272097,272097,272097,272097,272097,272097,272097,272097,272097,272097,272097,272097,272097,272097,272097,272097, \ No newline at end of file diff --git a/integration_test/opts/ccsd_t1_21_ttgt.ta b/test/integration/opts/ccsd_t1_21_ttgt.ta similarity index 100% rename from integration_test/opts/ccsd_t1_21_ttgt.ta rename to test/integration/opts/ccsd_t1_21_ttgt.ta diff --git a/integration_test/opts/ccsd_t1_21_ttgt_all_opts.ta b/test/integration/opts/ccsd_t1_21_ttgt_all_opts.ta similarity index 91% rename from integration_test/opts/ccsd_t1_21_ttgt_all_opts.ta rename to test/integration/opts/ccsd_t1_21_ttgt_all_opts.ta index 553d8333..ce0eebfb 100644 --- a/integration_test/opts/ccsd_t1_21_ttgt_all_opts.ta +++ b/test/integration/opts/ccsd_t1_21_ttgt_all_opts.ta @@ -1,5 +1,5 @@ -# RUN: comet-opt -opt-bestperm-ttgt -opt-matmul-tiling -opt-matmul-mkernel -opt-dense-transpose --convert-tc-to-ttgt --convert-to-llvm %s &> ccsd_t1_21_ttgt_all.llvm -# RUN: mlir-cpu-runner ccsd_t1_21_ttgt_all.llvm -O3 -e main -entry-point-result=void -shared-libs=%comet_utility_library_dir/libcomet_runner_utils%shlibext | FileCheck %s +# RUN: comet-opt -opt-matmul-tiling -opt-matmul-mkernel -opt-dense-transpose --convert-tc-to-ttgt --convert-to-loops --convert-to-llvm %s &> ccsd_t1_21_ttgt_all.llvm +# RUN: mlir-cpu-runner ccsd_t1_21_ttgt_all.llvm -O3 -e main -entry-point-result=void -shared-libs=%comet_utility_library_dir/libcomet_runner_utils%shlibext,%mlir_utility_library_dir/libomp%shlibext | FileCheck %s def main() { #IndexLabel Declarations diff --git a/test/integration/opts/ccsd_t1_21_ttgt_tiling.ta b/test/integration/opts/ccsd_t1_21_ttgt_tiling.ta new file mode 100644 index 00000000..d4f9714f --- /dev/null +++ b/test/integration/opts/ccsd_t1_21_ttgt_tiling.ta @@ -0,0 +1,24 @@ +# RUN: comet-opt --opt-matmul-tiling --convert-tc-to-ttgt --convert-to-loops --convert-to-llvm %s &> ccsd_t1_21_ttgt_tiling.llvm +# RUN: mlir-cpu-runner ccsd_t1_21_ttgt_tiling.llvm -O3 -e main -entry-point-result=void -shared-libs=%comet_utility_library_dir/libcomet_runner_utils%shlibext,%mlir_utility_library_dir/libomp%shlibext | FileCheck %s + +def main() { + #IndexLabel Declarations + IndexLabel [i, c] = [16]; + IndexLabel [m, n, a] = [32]; + + Tensor v([i, c, m, n], {Dense}); + Tensor t2([m, n, c, a], {Dense}); + Tensor i0([i, a], {Dense}); + + v[i, c, m, n] = 2.3; + t2[m, n, c, a] = 3.4; + i0[i, a] = 0.0; + + #Tensor contraction + i0[i, a] = v[i, c, m, n] * t2[m, n, c, a]; #ccsd_t1 21st expression + print(i0); +} + +# Print the result for verification. +# CHECK: data = +# CHECK-NEXT: 128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123,128123, \ No newline at end of file diff --git a/integration_test/opts/ccsd_t1_3_ttgt.ta b/test/integration/opts/ccsd_t1_3_ttgt.ta similarity index 100% rename from integration_test/opts/ccsd_t1_3_ttgt.ta rename to test/integration/opts/ccsd_t1_3_ttgt.ta diff --git a/integration_test/opts/ccsd_t1_4_ttgt.ta b/test/integration/opts/ccsd_t1_4_ttgt.ta similarity index 86% rename from integration_test/opts/ccsd_t1_4_ttgt.ta rename to test/integration/opts/ccsd_t1_4_ttgt.ta index e05fe51e..bc06daa6 100644 --- a/integration_test/opts/ccsd_t1_4_ttgt.ta +++ b/test/integration/opts/ccsd_t1_4_ttgt.ta @@ -1,4 +1,4 @@ -# RUN: comet-opt --convert-tc-to-ttgt --convert-to-loops --convert-to-llvm %s &> ccsd_t1_4_ttgt.llvm +# RUN: comet-opt --convert-tc-to-ttgt=1 --convert-to-loops --convert-to-llvm %s &> ccsd_t1_4_ttgt.llvm # RUN: mlir-cpu-runner ccsd_t1_4_ttgt.llvm -O3 -e main -entry-point-result=void -shared-libs=%comet_utility_library_dir/libcomet_runner_utils%shlibext | FileCheck %s def main() { diff --git a/integration_test/opts/ccsd_t1_4_ttgt_bestperm.ta b/test/integration/opts/ccsd_t1_4_ttgt_bestperm.ta similarity index 76% rename from integration_test/opts/ccsd_t1_4_ttgt_bestperm.ta rename to test/integration/opts/ccsd_t1_4_ttgt_bestperm.ta index 439e9052..0739d7f5 100644 --- a/integration_test/opts/ccsd_t1_4_ttgt_bestperm.ta +++ b/test/integration/opts/ccsd_t1_4_ttgt_bestperm.ta @@ -1,5 +1,5 @@ -# RUN: comet-opt --opt-bestperm-ttgt --convert-tc-to-ttgt --convert-to-llvm %s &> ccsd_t1_4_ttgt_bestperm.llvm -# RUN: mlir-cpu-runner ccsd_t1_4_ttgt_bestperm.llvm -O3 -e main -entry-point-result=void -shared-libs=%comet_utility_library_dir/libcomet_runner_utils%shlibext | FileCheck %s +# RUN: comet-opt --convert-tc-to-ttgt --convert-to-loops --convert-to-llvm %s &> ccsd_t1_4_ttgt_bestperm.llvm +# RUN: mlir-cpu-runner ccsd_t1_4_ttgt_bestperm.llvm -O3 -e main -entry-point-result=void -shared-libs=%comet_utility_library_dir/libcomet_runner_utils%shlibext,%mlir_utility_library_dir/libomp%shlibext | FileCheck %s def main() { #IndexLabel Declarations diff --git a/integration_test/opts/chain_mult_factorize.ta b/test/integration/opts/chain_mult_factorize.ta similarity index 90% rename from integration_test/opts/chain_mult_factorize.ta rename to test/integration/opts/chain_mult_factorize.ta index d2488e7d..338c0cc9 100644 --- a/integration_test/opts/chain_mult_factorize.ta +++ b/test/integration/opts/chain_mult_factorize.ta @@ -1,5 +1,5 @@ # RUN: comet-opt --opt-multiop-factorize --convert-ta-to-it --convert-to-loops --convert-to-llvm %s &> chain_mult_factorize.llvm -# RUN: mlir-cpu-runner chain_mult_factorize.llvm -O3 -e main -entry-point-result=void -shared-libs=%comet_utility_library_dir/libcomet_runner_utils%shlibext | FileCheck %s +# RUN: mlir-cpu-runner chain_mult_factorize.llvm -O3 -e main -entry-point-result=void -shared-libs=%comet_utility_library_dir/libcomet_runner_utils%shlibext,%mlir_utility_library_dir/libomp%shlibext | FileCheck %s def main() { diff --git a/integration_test/opts/fusion.ta b/test/integration/opts/fusion.ta similarity index 94% rename from integration_test/opts/fusion.ta rename to test/integration/opts/fusion.ta index eff4a869..d869f2aa 100644 --- a/integration_test/opts/fusion.ta +++ b/test/integration/opts/fusion.ta @@ -5,7 +5,7 @@ # --opt-fusion pass performs the redundancy-aware fusion on GNN kernels # RUN: comet-opt --convert-ta-to-it --opt-fusion --convert-to-loops --convert-to-llvm %s &> fusion_loops.llvm # RUN: export SPARSE_FILE_NAME0=%comet_integration_test_data_dir/test_rank2_small.mtx -# RUN: mlir-cpu-runner fusion_loops.llvm -O3 -e main -entry-point-result=void -shared-libs=%comet_utility_library_dir/libcomet_runner_utils%shlibext | FileCheck %s +# RUN: mlir-cpu-runner fusion_loops.llvm -O3 -e main -entry-point-result=void -shared-libs=%comet_utility_library_dir/libcomet_runner_utils%shlibext,%mlir_utility_library_dir/libomp%shlibext | FileCheck %s def main() { #IndexLabel Declarations diff --git a/integration_test/opts/opt_dense_transpose.ta b/test/integration/opts/opt_dense_transpose.ta similarity index 92% rename from integration_test/opts/opt_dense_transpose.ta rename to test/integration/opts/opt_dense_transpose.ta index 6533fafe..01850fc3 100644 --- a/integration_test/opts/opt_dense_transpose.ta +++ b/test/integration/opts/opt_dense_transpose.ta @@ -1,5 +1,5 @@ # RUN: comet-opt -opt-dense-transpose --convert-ta-to-it --convert-to-loops --convert-to-llvm %s &> opt_dense_transpose.llvm -# RUN: mlir-cpu-runner opt_dense_transpose.llvm -O3 -e main -entry-point-result=void -shared-libs=%comet_utility_library_dir/libcomet_runner_utils%shlibext | FileCheck %s +# RUN: mlir-cpu-runner opt_dense_transpose.llvm -O3 -e main -entry-point-result=void -shared-libs=%comet_utility_library_dir/libcomet_runner_utils%shlibext,%mlir_utility_library_dir/libomp%shlibext | FileCheck %s #TODO(gkestor): read dense input from file diff --git a/integration_test/semiring/eltwise_monoidMin_DensexDense_oDense.ta b/test/integration/semiring/eltwise_monoidMin_DensexDense_oDense.ta similarity index 88% rename from integration_test/semiring/eltwise_monoidMin_DensexDense_oDense.ta rename to test/integration/semiring/eltwise_monoidMin_DensexDense_oDense.ta index c74636e6..59727d05 100644 --- a/integration_test/semiring/eltwise_monoidMin_DensexDense_oDense.ta +++ b/test/integration/semiring/eltwise_monoidMin_DensexDense_oDense.ta @@ -1,5 +1,5 @@ # RUN: comet-opt --convert-ta-to-it --convert-to-loops --convert-to-llvm %s &> eltwise_monoidMin_DensexDense_oDense.llvm -# RUN: mlir-cpu-runner eltwise_monoidMin_DensexDense_oDense.llvm -O3 -e main -entry-point-result=void -shared-libs=%comet_utility_library_dir/libcomet_runner_utils%shlibext | FileCheck %s +# RUN: mlir-cpu-runner eltwise_monoidMin_DensexDense_oDense.llvm -O3 -e main -entry-point-result=void -shared-libs=%comet_utility_library_dir/libcomet_runner_utils%shlibext,%mlir_utility_library_dir/libomp%shlibext | FileCheck %s def main() { diff --git a/integration_test/semiring/eltwise_monoidMinus_DensexDense_oDense.ta b/test/integration/semiring/eltwise_monoidMinus_DensexDense_oDense.ta similarity index 88% rename from integration_test/semiring/eltwise_monoidMinus_DensexDense_oDense.ta rename to test/integration/semiring/eltwise_monoidMinus_DensexDense_oDense.ta index a46b4747..4a5b4367 100644 --- a/integration_test/semiring/eltwise_monoidMinus_DensexDense_oDense.ta +++ b/test/integration/semiring/eltwise_monoidMinus_DensexDense_oDense.ta @@ -1,5 +1,5 @@ # RUN: comet-opt --convert-ta-to-it --convert-to-loops --convert-to-llvm %s &> eltwise_monoidMinus_DensexDense_oDense.llvm -# RUN: mlir-cpu-runner eltwise_monoidMinus_DensexDense_oDense.llvm -O3 -e main -entry-point-result=void -shared-libs=%comet_utility_library_dir/libcomet_runner_utils%shlibext | FileCheck %s +# RUN: mlir-cpu-runner eltwise_monoidMinus_DensexDense_oDense.llvm -O3 -e main -entry-point-result=void -shared-libs=%comet_utility_library_dir/libcomet_runner_utils%shlibext,%mlir_utility_library_dir/libomp%shlibext | FileCheck %s def main() { diff --git a/integration_test/semiring/eltwise_monoidPlus_COOxDense_oCOO.ta b/test/integration/semiring/eltwise_monoidPlus_COOxDense_oCOO.ta similarity index 98% rename from integration_test/semiring/eltwise_monoidPlus_COOxDense_oCOO.ta rename to test/integration/semiring/eltwise_monoidPlus_COOxDense_oCOO.ta index 58b5aa60..956c3743 100644 --- a/integration_test/semiring/eltwise_monoidPlus_COOxDense_oCOO.ta +++ b/test/integration/semiring/eltwise_monoidPlus_COOxDense_oCOO.ta @@ -27,7 +27,7 @@ def main() { # CHECK-NEXT: data = # CHECK-NEXT: 0,0,1,1,2,3,3,4,4, # CHECK-NEXT: data = -# CHECK-NEXT: 0, +# CHECK-NEXT: -1, # CHECK-NEXT: data = # CHECK-NEXT: 0,3,1,4,2,0,3,1,4, # CHECK-NEXT: data = diff --git a/integration_test/semiring/eltwise_monoidPlus_DensexDense_oDense.ta b/test/integration/semiring/eltwise_monoidPlus_DensexDense_oDense.ta similarity index 88% rename from integration_test/semiring/eltwise_monoidPlus_DensexDense_oDense.ta rename to test/integration/semiring/eltwise_monoidPlus_DensexDense_oDense.ta index ea346f70..bcab5d14 100644 --- a/integration_test/semiring/eltwise_monoidPlus_DensexDense_oDense.ta +++ b/test/integration/semiring/eltwise_monoidPlus_DensexDense_oDense.ta @@ -1,5 +1,5 @@ # RUN: comet-opt --convert-ta-to-it --convert-to-loops --convert-to-llvm %s &> eltwise_monoidPlus_DensexDense_oDense.llvm -# RUN: mlir-cpu-runner eltwise_monoidPlus_DensexDense_oDense.llvm -O3 -e main -entry-point-result=void -shared-libs=%comet_utility_library_dir/libcomet_runner_utils%shlibext | FileCheck %s +# RUN: mlir-cpu-runner eltwise_monoidPlus_DensexDense_oDense.llvm -O3 -e main -entry-point-result=void -shared-libs=%comet_utility_library_dir/libcomet_runner_utils%shlibext,%mlir_utility_library_dir/libomp%shlibext | FileCheck %s def main() { diff --git a/integration_test/semiring/eltwise_monoidTimes_COOxDense_oCOO.ta b/test/integration/semiring/eltwise_monoidTimes_COOxDense_oCOO.ta similarity index 98% rename from integration_test/semiring/eltwise_monoidTimes_COOxDense_oCOO.ta rename to test/integration/semiring/eltwise_monoidTimes_COOxDense_oCOO.ta index 94d6de49..056d1f6c 100644 --- a/integration_test/semiring/eltwise_monoidTimes_COOxDense_oCOO.ta +++ b/test/integration/semiring/eltwise_monoidTimes_COOxDense_oCOO.ta @@ -27,7 +27,7 @@ def main() { # CHECK-NEXT: data = # CHECK-NEXT: 0,0,1,1,2,3,3,4,4, # CHECK-NEXT: data = -# CHECK-NEXT: 0, +# CHECK-NEXT: -1, # CHECK-NEXT: data = # CHECK-NEXT: 0,3,1,4,2,0,3,1,4, # CHECK-NEXT: data = diff --git a/integration_test/semiring/eltwise_monoidTimes_DensexDense_oDense.ta b/test/integration/semiring/eltwise_monoidTimes_DensexDense_oDense.ta similarity index 88% rename from integration_test/semiring/eltwise_monoidTimes_DensexDense_oDense.ta rename to test/integration/semiring/eltwise_monoidTimes_DensexDense_oDense.ta index 883f4ba0..c5025209 100644 --- a/integration_test/semiring/eltwise_monoidTimes_DensexDense_oDense.ta +++ b/test/integration/semiring/eltwise_monoidTimes_DensexDense_oDense.ta @@ -1,5 +1,5 @@ # RUN: comet-opt --convert-ta-to-it --convert-to-loops --convert-to-llvm %s &> eltwise_monoidTimes_DensexDense_oDense.llvm -# RUN: mlir-cpu-runner eltwise_monoidTimes_DensexDense_oDense.llvm -O3 -e main -entry-point-result=void -shared-libs=%comet_utility_library_dir/libcomet_runner_utils%shlibext | FileCheck %s +# RUN: mlir-cpu-runner eltwise_monoidTimes_DensexDense_oDense.llvm -O3 -e main -entry-point-result=void -shared-libs=%comet_utility_library_dir/libcomet_runner_utils%shlibext,%mlir_utility_library_dir/libomp%shlibext | FileCheck %s def main() { diff --git a/integration_test/semiring/eltwise_monoidTimes_dense_4Dtensors.ta b/test/integration/semiring/eltwise_monoidTimes_dense_4Dtensors.ta similarity index 89% rename from integration_test/semiring/eltwise_monoidTimes_dense_4Dtensors.ta rename to test/integration/semiring/eltwise_monoidTimes_dense_4Dtensors.ta index f34e803a..0af574b3 100644 --- a/integration_test/semiring/eltwise_monoidTimes_dense_4Dtensors.ta +++ b/test/integration/semiring/eltwise_monoidTimes_dense_4Dtensors.ta @@ -1,5 +1,5 @@ # RUN: comet-opt --convert-ta-to-it --convert-to-loops --convert-to-llvm %s &> eltwise_monoidTimes_dense_4Dtensors.llvm -# RUN: mlir-cpu-runner eltwise_monoidTimes_dense_4Dtensors.llvm -O3 -e main -entry-point-result=void -shared-libs=%comet_utility_library_dir/libcomet_runner_utils%shlibext | FileCheck %s +# RUN: mlir-cpu-runner eltwise_monoidTimes_dense_4Dtensors.llvm -O3 -e main -entry-point-result=void -shared-libs=%comet_utility_library_dir/libcomet_runner_utils%shlibext,%mlir_utility_library_dir/libomp%shlibext | FileCheck %s def main() { # IndexLabel Declarations diff --git a/integration_test/semiring/mm_SemiringAnyPair_CSRxCSR_oCSR.ta b/test/integration/semiring/mm_SemiringAnyPair_CSRxCSR_oCSR.ta similarity index 98% rename from integration_test/semiring/mm_SemiringAnyPair_CSRxCSR_oCSR.ta rename to test/integration/semiring/mm_SemiringAnyPair_CSRxCSR_oCSR.ta index b4fe8c63..6bc08203 100644 --- a/integration_test/semiring/mm_SemiringAnyPair_CSRxCSR_oCSR.ta +++ b/test/integration/semiring/mm_SemiringAnyPair_CSRxCSR_oCSR.ta @@ -29,7 +29,7 @@ def main() { # CHECK: data = # CHECK-NEXT: 5, # CHECK-NEXT: data = -# CHECK-NEXT: 0, +# CHECK-NEXT: -1, # CHECK-NEXT: data = # CHECK-NEXT: 0,2,4,5,7,9, # CHECK-NEXT: data = diff --git a/integration_test/semiring/mm_SemiringMinFirst_CSRxCSR_oCSR.ta b/test/integration/semiring/mm_SemiringMinFirst_CSRxCSR_oCSR.ta similarity index 98% rename from integration_test/semiring/mm_SemiringMinFirst_CSRxCSR_oCSR.ta rename to test/integration/semiring/mm_SemiringMinFirst_CSRxCSR_oCSR.ta index 64acee8c..ed581294 100644 --- a/integration_test/semiring/mm_SemiringMinFirst_CSRxCSR_oCSR.ta +++ b/test/integration/semiring/mm_SemiringMinFirst_CSRxCSR_oCSR.ta @@ -29,7 +29,7 @@ def main() { # CHECK: data = # CHECK-NEXT: 5, # CHECK-NEXT: data = -# CHECK-NEXT: 0, +# CHECK-NEXT: -1, # CHECK-NEXT: data = # CHECK-NEXT: 0,2,4,5,7,9, # CHECK-NEXT: data = diff --git a/integration_test/semiring/mm_SemiringMinPlus_CSRxCSR_oCSR.ta b/test/integration/semiring/mm_SemiringMinPlus_CSRxCSR_oCSR.ta similarity index 98% rename from integration_test/semiring/mm_SemiringMinPlus_CSRxCSR_oCSR.ta rename to test/integration/semiring/mm_SemiringMinPlus_CSRxCSR_oCSR.ta index c4929325..fab90f28 100644 --- a/integration_test/semiring/mm_SemiringMinPlus_CSRxCSR_oCSR.ta +++ b/test/integration/semiring/mm_SemiringMinPlus_CSRxCSR_oCSR.ta @@ -29,7 +29,7 @@ def main() { # CHECK: data = # CHECK-NEXT: 5, # CHECK-NEXT: data = -# CHECK-NEXT: 0, +# CHECK-NEXT: -1, # CHECK-NEXT: data = # CHECK-NEXT: 0,2,4,5,7,9, # CHECK-NEXT: data = diff --git a/integration_test/semiring/mm_SemiringMinSecond_CSRxCSR_oCSR.ta b/test/integration/semiring/mm_SemiringMinSecond_CSRxCSR_oCSR.ta similarity index 98% rename from integration_test/semiring/mm_SemiringMinSecond_CSRxCSR_oCSR.ta rename to test/integration/semiring/mm_SemiringMinSecond_CSRxCSR_oCSR.ta index 6b6008db..314b0d32 100644 --- a/integration_test/semiring/mm_SemiringMinSecond_CSRxCSR_oCSR.ta +++ b/test/integration/semiring/mm_SemiringMinSecond_CSRxCSR_oCSR.ta @@ -29,7 +29,7 @@ def main() { # CHECK: data = # CHECK-NEXT: 5, # CHECK-NEXT: data = -# CHECK-NEXT: 0, +# CHECK-NEXT: -1, # CHECK-NEXT: data = # CHECK-NEXT: 0,2,4,5,7,9, # CHECK-NEXT: data = diff --git a/integration_test/semiring/mm_SemiringPlusFirst_CSRxCSR_oCSR.ta b/test/integration/semiring/mm_SemiringPlusFirst_CSRxCSR_oCSR.ta similarity index 98% rename from integration_test/semiring/mm_SemiringPlusFirst_CSRxCSR_oCSR.ta rename to test/integration/semiring/mm_SemiringPlusFirst_CSRxCSR_oCSR.ta index c49478e9..3308a5c1 100644 --- a/integration_test/semiring/mm_SemiringPlusFirst_CSRxCSR_oCSR.ta +++ b/test/integration/semiring/mm_SemiringPlusFirst_CSRxCSR_oCSR.ta @@ -29,7 +29,7 @@ def main() { # CHECK: data = # CHECK-NEXT: 5, # CHECK-NEXT: data = -# CHECK-NEXT: 0, +# CHECK-NEXT: -1, # CHECK-NEXT: data = # CHECK-NEXT: 0,2,4,5,7,9, # CHECK-NEXT: data = diff --git a/integration_test/semiring/mm_SemiringPlusPair_CSRxCSR_oCSR.ta b/test/integration/semiring/mm_SemiringPlusPair_CSRxCSR_oCSR.ta similarity index 91% rename from integration_test/semiring/mm_SemiringPlusPair_CSRxCSR_oCSR.ta rename to test/integration/semiring/mm_SemiringPlusPair_CSRxCSR_oCSR.ta index 1b590eec..77ca6477 100644 --- a/integration_test/semiring/mm_SemiringPlusPair_CSRxCSR_oCSR.ta +++ b/test/integration/semiring/mm_SemiringPlusPair_CSRxCSR_oCSR.ta @@ -3,7 +3,7 @@ # RUN: comet-opt --opt-comp-workspace --convert-ta-to-it --convert-to-loops --convert-to-llvm %s &> mm_SemiringPlusPair_CSRxCSR_oCSR.llvm # RUN: export SPARSE_FILE_NAME0=%comet_integration_test_data_dir/test_rank2.mtx # RUN: export SPARSE_FILE_NAME1=%comet_integration_test_data_dir/test_rank2.mtx -# RUN: mlir-cpu-runner mm_SemiringPlusPair_CSRxCSR_oCSR.llvm -O3 -e main -entry-point-result=void -shared-libs=%comet_utility_library_dir/libcomet_runner_utils%shlibext | FileCheck %s +# RUN: mlir-cpu-runner mm_SemiringPlusPair_CSRxCSR_oCSR.llvm -O3 -e main -entry-point-result=void -shared-libs=%comet_utility_library_dir/libcomet_runner_utils%shlibext,%mlir_utility_library_dir/libomp%shlibext | FileCheck %s def main() { #IndexLabel Declarations @@ -29,7 +29,7 @@ def main() { # CHECK: data = # CHECK-NEXT: 5, # CHECK-NEXT: data = -# CHECK-NEXT: 0, +# CHECK-NEXT: -1, # CHECK-NEXT: data = # CHECK-NEXT: 0,2,4,5,7,9, # CHECK-NEXT: data = diff --git a/integration_test/semiring/mm_SemiringPlusSecond_CSRxCSR_oCSR.ta b/test/integration/semiring/mm_SemiringPlusSecond_CSRxCSR_oCSR.ta similarity index 98% rename from integration_test/semiring/mm_SemiringPlusSecond_CSRxCSR_oCSR.ta rename to test/integration/semiring/mm_SemiringPlusSecond_CSRxCSR_oCSR.ta index fab31d69..9cf0fe0c 100644 --- a/integration_test/semiring/mm_SemiringPlusSecond_CSRxCSR_oCSR.ta +++ b/test/integration/semiring/mm_SemiringPlusSecond_CSRxCSR_oCSR.ta @@ -29,7 +29,7 @@ def main() { # CHECK: data = # CHECK-NEXT: 5, # CHECK-NEXT: data = -# CHECK-NEXT: 0, +# CHECK-NEXT: -1, # CHECK-NEXT: data = # CHECK-NEXT: 0,2,4,5,7,9, # CHECK-NEXT: data = diff --git a/integration_test/semiring/mm_SemiringPlusTimes_COOxDense_oDense.ta b/test/integration/semiring/mm_SemiringPlusTimes_COOxDense_oDense.ta similarity index 100% rename from integration_test/semiring/mm_SemiringPlusTimes_COOxDense_oDense.ta rename to test/integration/semiring/mm_SemiringPlusTimes_COOxDense_oDense.ta diff --git a/integration_test/semiring/mm_SemiringPlusTimes_CSRxCSR_oCSR.ta b/test/integration/semiring/mm_SemiringPlusTimes_CSRxCSR_oCSR.ta similarity index 98% rename from integration_test/semiring/mm_SemiringPlusTimes_CSRxCSR_oCSR.ta rename to test/integration/semiring/mm_SemiringPlusTimes_CSRxCSR_oCSR.ta index 8413d573..a7bdc3fe 100644 --- a/integration_test/semiring/mm_SemiringPlusTimes_CSRxCSR_oCSR.ta +++ b/test/integration/semiring/mm_SemiringPlusTimes_CSRxCSR_oCSR.ta @@ -30,7 +30,7 @@ def main() { # CHECK: data = # CHECK-NEXT: 5, # CHECK-NEXT: data = -# CHECK-NEXT: 0, +# CHECK-NEXT: -1, # CHECK-NEXT: data = # CHECK-NEXT: 0,2,4,5,7,9, # CHECK-NEXT: data = diff --git a/integration_test/semiring/mm_SemiringPlusTimes_CSRxDense_oDense.ta b/test/integration/semiring/mm_SemiringPlusTimes_CSRxDense_oDense.ta similarity index 90% rename from integration_test/semiring/mm_SemiringPlusTimes_CSRxDense_oDense.ta rename to test/integration/semiring/mm_SemiringPlusTimes_CSRxDense_oDense.ta index c792f618..7135cdbe 100644 --- a/integration_test/semiring/mm_SemiringPlusTimes_CSRxDense_oDense.ta +++ b/test/integration/semiring/mm_SemiringPlusTimes_CSRxDense_oDense.ta @@ -2,7 +2,7 @@ # Sparse matrix is in CSR format # RUN: comet-opt --convert-ta-to-it --convert-to-loops --convert-to-llvm %s &> mm_SemiringPlusTimes_CSRxDense_oDense.llvm # RUN: export SPARSE_FILE_NAME0=%comet_integration_test_data_dir/test_rank2.mtx -# RUN: mlir-cpu-runner mm_SemiringPlusTimes_CSRxDense_oDense.llvm -O3 -e main -entry-point-result=void -shared-libs=%comet_utility_library_dir/libcomet_runner_utils%shlibext | FileCheck %s +# RUN: mlir-cpu-runner mm_SemiringPlusTimes_CSRxDense_oDense.llvm -O3 -e main -entry-point-result=void -shared-libs=%comet_utility_library_dir/libcomet_runner_utils%shlibext,%mlir_utility_library_dir/libomp%shlibext | FileCheck %s def main() { #IndexLabel Declarations diff --git a/integration_test/semiring/mm_SemiringPlusTimes_DCSRxDense_oDense.ta b/test/integration/semiring/mm_SemiringPlusTimes_DCSRxDense_oDense.ta similarity index 90% rename from integration_test/semiring/mm_SemiringPlusTimes_DCSRxDense_oDense.ta rename to test/integration/semiring/mm_SemiringPlusTimes_DCSRxDense_oDense.ta index c314e4e1..40b39422 100644 --- a/integration_test/semiring/mm_SemiringPlusTimes_DCSRxDense_oDense.ta +++ b/test/integration/semiring/mm_SemiringPlusTimes_DCSRxDense_oDense.ta @@ -2,7 +2,7 @@ # Sparse matrix is in DCSR format # RUN: comet-opt --convert-ta-to-it --convert-to-loops --convert-to-llvm %s &> mm_SemiringPlusTimes_DCSRxDense_oDense.llvm # RUN: export SPARSE_FILE_NAME0=%comet_integration_test_data_dir/test_rank2.mtx -# RUN: mlir-cpu-runner mm_SemiringPlusTimes_DCSRxDense_oDense.llvm -O3 -e main -entry-point-result=void -shared-libs=%comet_utility_library_dir/libcomet_runner_utils%shlibext | FileCheck %s +# RUN: mlir-cpu-runner mm_SemiringPlusTimes_DCSRxDense_oDense.llvm -O3 -e main -entry-point-result=void -shared-libs=%comet_utility_library_dir/libcomet_runner_utils%shlibext,%mlir_utility_library_dir/libomp%shlibext | FileCheck %s def main() { #IndexLabel Declarations diff --git a/integration_test/semiring/mm_SemiringPlusTimes_DensexCOO_oDense.ta b/test/integration/semiring/mm_SemiringPlusTimes_DensexCOO_oDense.ta similarity index 91% rename from integration_test/semiring/mm_SemiringPlusTimes_DensexCOO_oDense.ta rename to test/integration/semiring/mm_SemiringPlusTimes_DensexCOO_oDense.ta index 902331aa..340cd219 100644 --- a/integration_test/semiring/mm_SemiringPlusTimes_DensexCOO_oDense.ta +++ b/test/integration/semiring/mm_SemiringPlusTimes_DensexCOO_oDense.ta @@ -2,7 +2,7 @@ # Sparse matrix is in COO format # RUN: comet-opt --convert-ta-to-it --convert-to-loops --convert-to-llvm %s &> mm_SemiringPlusTimes_DensexCOO_oDense.llvm # RUN: export SPARSE_FILE_NAME0=%comet_integration_test_data_dir/test_rank2.mtx -# RUN: mlir-cpu-runner mm_SemiringPlusTimes_DensexCOO_oDense.llvm -O3 -e main -entry-point-result=void -shared-libs=%comet_utility_library_dir/libcomet_runner_utils%shlibext | FileCheck %s +# RUN: mlir-cpu-runner mm_SemiringPlusTimes_DensexCOO_oDense.llvm -O3 -e main -entry-point-result=void -shared-libs=%comet_utility_library_dir/libcomet_runner_utils%shlibext,%mlir_utility_library_dir/libomp%shlibext | FileCheck %s def main() { #IndexLabel Declarations diff --git a/integration_test/semiring/mm_SemiringPlusTimes_DensexCSR_oDense.ta b/test/integration/semiring/mm_SemiringPlusTimes_DensexCSR_oDense.ta similarity index 91% rename from integration_test/semiring/mm_SemiringPlusTimes_DensexCSR_oDense.ta rename to test/integration/semiring/mm_SemiringPlusTimes_DensexCSR_oDense.ta index 000358d4..5ffd4e3b 100644 --- a/integration_test/semiring/mm_SemiringPlusTimes_DensexCSR_oDense.ta +++ b/test/integration/semiring/mm_SemiringPlusTimes_DensexCSR_oDense.ta @@ -2,7 +2,7 @@ # Sparse matrix is in CSR format # RUN: comet-opt --convert-ta-to-it --convert-to-loops --convert-to-llvm %s &> mm_SemiringPlusTimes_DensexCSR_oDense.llvm # RUN: export SPARSE_FILE_NAME0=%comet_integration_test_data_dir/test_rank2.mtx -# RUN: mlir-cpu-runner mm_SemiringPlusTimes_DensexCSR_oDense.llvm -O3 -e main -entry-point-result=void -shared-libs=%comet_utility_library_dir/libcomet_runner_utils%shlibext | FileCheck %s +# RUN: mlir-cpu-runner mm_SemiringPlusTimes_DensexCSR_oDense.llvm -O3 -e main -entry-point-result=void -shared-libs=%comet_utility_library_dir/libcomet_runner_utils%shlibext,%mlir_utility_library_dir/libomp%shlibext | FileCheck %s def main() { #IndexLabel Declarations diff --git a/integration_test/semiring/mm_SemiringPlusTimes_DensexDCSR_oDense.ta b/test/integration/semiring/mm_SemiringPlusTimes_DensexDCSR_oDense.ta similarity index 90% rename from integration_test/semiring/mm_SemiringPlusTimes_DensexDCSR_oDense.ta rename to test/integration/semiring/mm_SemiringPlusTimes_DensexDCSR_oDense.ta index 5a0c8ae9..a7a09d85 100644 --- a/integration_test/semiring/mm_SemiringPlusTimes_DensexDCSR_oDense.ta +++ b/test/integration/semiring/mm_SemiringPlusTimes_DensexDCSR_oDense.ta @@ -2,7 +2,7 @@ # Sparse matrix is in DCSR format # RUN: comet-opt --convert-ta-to-it --convert-to-loops --convert-to-llvm %s &> mm_SemiringPlusTimes_DensexDCSR_oDense.llvm # RUN: export SPARSE_FILE_NAME0=%comet_integration_test_data_dir/test_rank2.mtx -# RUN: mlir-cpu-runner mm_SemiringPlusTimes_DensexDCSR_oDense.llvm -O3 -e main -entry-point-result=void -shared-libs=%comet_utility_library_dir/libcomet_runner_utils%shlibext | FileCheck %s +# RUN: mlir-cpu-runner mm_SemiringPlusTimes_DensexDCSR_oDense.llvm -O3 -e main -entry-point-result=void -shared-libs=%comet_utility_library_dir/libcomet_runner_utils%shlibext,%mlir_utility_library_dir/libomp%shlibext | FileCheck %s def main() { #IndexLabel Declarations diff --git a/integration_test/semiring/mm_SemiringPlusTimes_DensexDense_oDense.ta b/test/integration/semiring/mm_SemiringPlusTimes_DensexDense_oDense.ta similarity index 88% rename from integration_test/semiring/mm_SemiringPlusTimes_DensexDense_oDense.ta rename to test/integration/semiring/mm_SemiringPlusTimes_DensexDense_oDense.ta index d7689058..cabf4863 100644 --- a/integration_test/semiring/mm_SemiringPlusTimes_DensexDense_oDense.ta +++ b/test/integration/semiring/mm_SemiringPlusTimes_DensexDense_oDense.ta @@ -1,5 +1,5 @@ # RUN: comet-opt --convert-ta-to-it --convert-to-loops --convert-to-llvm %s &> mm_SemiringPlusTimes_DensexDense_oDense.llvm -# RUN: mlir-cpu-runner mm_SemiringPlusTimes_DensexDense_oDense.llvm -O3 -e main -entry-point-result=void -shared-libs=%comet_utility_library_dir/libcomet_runner_utils%shlibext | FileCheck %s +# RUN: mlir-cpu-runner mm_SemiringPlusTimes_DensexDense_oDense.llvm -O3 -e main -entry-point-result=void -shared-libs=%comet_utility_library_dir/libcomet_runner_utils%shlibext,%mlir_utility_library_dir/libomp%shlibext | FileCheck %s def main() { diff --git a/integration_test/semiring/mm_SemiringPlusTimes_dense_4Dtensors.ta b/test/integration/semiring/mm_SemiringPlusTimes_dense_4Dtensors.ta similarity index 90% rename from integration_test/semiring/mm_SemiringPlusTimes_dense_4Dtensors.ta rename to test/integration/semiring/mm_SemiringPlusTimes_dense_4Dtensors.ta index 0f3c2ccc..edbb69fa 100644 --- a/integration_test/semiring/mm_SemiringPlusTimes_dense_4Dtensors.ta +++ b/test/integration/semiring/mm_SemiringPlusTimes_dense_4Dtensors.ta @@ -1,5 +1,5 @@ # RUN: comet-opt --convert-ta-to-it --convert-to-loops --convert-to-llvm %s &> mm_SemiringPlusTimes_dense_4Dtensors.llvm -# RUN: mlir-cpu-runner mm_SemiringPlusTimes_dense_4Dtensors.llvm -O3 -e main -entry-point-result=void -shared-libs=%comet_utility_library_dir/libcomet_runner_utils%shlibext | FileCheck %s +# RUN: mlir-cpu-runner mm_SemiringPlusTimes_dense_4Dtensors.llvm -O3 -e main -entry-point-result=void -shared-libs=%comet_utility_library_dir/libcomet_runner_utils%shlibext,%mlir_utility_library_dir/libomp%shlibext | FileCheck %s # This example comes from TCCG - ccsd12 def main() { diff --git a/integration_test/semiring/mv_SemiringPlusTimes_COOxDense_oDense.ta b/test/integration/semiring/mv_SemiringPlusTimes_COOxDense_oDense.ta similarity index 89% rename from integration_test/semiring/mv_SemiringPlusTimes_COOxDense_oDense.ta rename to test/integration/semiring/mv_SemiringPlusTimes_COOxDense_oDense.ta index 66e103e9..8f0d9bac 100644 --- a/integration_test/semiring/mv_SemiringPlusTimes_COOxDense_oDense.ta +++ b/test/integration/semiring/mv_SemiringPlusTimes_COOxDense_oDense.ta @@ -2,7 +2,7 @@ # Sparse matrix is in COO format # RUN: comet-opt --convert-ta-to-it --convert-to-loops --convert-to-llvm %s &> mv_SemiringPlusTimes_COOxDense_oDense.llvm # RUN: export SPARSE_FILE_NAME0=%comet_integration_test_data_dir/test_rank2.mtx -# RUN: mlir-cpu-runner mv_SemiringPlusTimes_COOxDense_oDense.llvm -O3 -e main -entry-point-result=void -shared-libs=%comet_utility_library_dir/libcomet_runner_utils%shlibext | FileCheck %s +# RUN: mlir-cpu-runner mv_SemiringPlusTimes_COOxDense_oDense.llvm -O3 -e main -entry-point-result=void -shared-libs=%comet_utility_library_dir/libcomet_runner_utils%shlibext,%mlir_utility_library_dir/libomp%shlibext | FileCheck %s def main() { diff --git a/integration_test/semiring/mv_SemiringPlusTimes_CSRxDense_oDense.ta b/test/integration/semiring/mv_SemiringPlusTimes_CSRxDense_oDense.ta similarity index 89% rename from integration_test/semiring/mv_SemiringPlusTimes_CSRxDense_oDense.ta rename to test/integration/semiring/mv_SemiringPlusTimes_CSRxDense_oDense.ta index 221405c9..1a173cf0 100644 --- a/integration_test/semiring/mv_SemiringPlusTimes_CSRxDense_oDense.ta +++ b/test/integration/semiring/mv_SemiringPlusTimes_CSRxDense_oDense.ta @@ -2,7 +2,7 @@ # Sparse matrix is in CSR format # RUN: comet-opt --convert-ta-to-it --convert-to-loops --convert-to-llvm %s &> mv_SemiringPlusTimes_CSRxDense_oDense.llvm # RUN: export SPARSE_FILE_NAME0=%comet_integration_test_data_dir/test_rank2.mtx -# RUN: mlir-cpu-runner mv_SemiringPlusTimes_CSRxDense_oDense.llvm -O3 -e main -entry-point-result=void -shared-libs=%comet_utility_library_dir/libcomet_runner_utils%shlibext | FileCheck %s +# RUN: mlir-cpu-runner mv_SemiringPlusTimes_CSRxDense_oDense.llvm -O3 -e main -entry-point-result=void -shared-libs=%comet_utility_library_dir/libcomet_runner_utils%shlibext,%mlir_utility_library_dir/libomp%shlibext | FileCheck %s def main() { diff --git a/integration_test/semiring/mv_SemiringPlusTimes_DCSRxDense_oDense.ta b/test/integration/semiring/mv_SemiringPlusTimes_DCSRxDense_oDense.ta similarity index 89% rename from integration_test/semiring/mv_SemiringPlusTimes_DCSRxDense_oDense.ta rename to test/integration/semiring/mv_SemiringPlusTimes_DCSRxDense_oDense.ta index 4b58c7c1..0bc1cc5d 100644 --- a/integration_test/semiring/mv_SemiringPlusTimes_DCSRxDense_oDense.ta +++ b/test/integration/semiring/mv_SemiringPlusTimes_DCSRxDense_oDense.ta @@ -2,7 +2,7 @@ # Sparse matrix is in DCSR format # RUN: comet-opt --convert-ta-to-it --convert-to-loops --convert-to-llvm %s &> mv_SemiringPlusTimes_DCSRxDense_oDense.llvm # RUN: export SPARSE_FILE_NAME0=%comet_integration_test_data_dir/test_rank2.mtx -# RUN: mlir-cpu-runner mv_SemiringPlusTimes_DCSRxDense_oDense.llvm -O3 -e main -entry-point-result=void -shared-libs=%comet_utility_library_dir/libcomet_runner_utils%shlibext | FileCheck %s +# RUN: mlir-cpu-runner mv_SemiringPlusTimes_DCSRxDense_oDense.llvm -O3 -e main -entry-point-result=void -shared-libs=%comet_utility_library_dir/libcomet_runner_utils%shlibext,%mlir_utility_library_dir/libomp%shlibext | FileCheck %s def main() { diff --git a/integration_test/semiring/mv_SemiringPlusTimes_DensexCOO_oDense.ta b/test/integration/semiring/mv_SemiringPlusTimes_DensexCOO_oDense.ta similarity index 89% rename from integration_test/semiring/mv_SemiringPlusTimes_DensexCOO_oDense.ta rename to test/integration/semiring/mv_SemiringPlusTimes_DensexCOO_oDense.ta index 6421959d..81ffd664 100644 --- a/integration_test/semiring/mv_SemiringPlusTimes_DensexCOO_oDense.ta +++ b/test/integration/semiring/mv_SemiringPlusTimes_DensexCOO_oDense.ta @@ -2,7 +2,7 @@ # Sparse matrix is in COO format # RUN: comet-opt --convert-ta-to-it --convert-to-loops --convert-to-llvm %s &> mv_SemiringPlusTimes_DensexCOO_oDense.llvm # RUN: export SPARSE_FILE_NAME0=%comet_integration_test_data_dir/test_rank2.mtx -# RUN: mlir-cpu-runner mv_SemiringPlusTimes_DensexCOO_oDense.llvm -O3 -e main -entry-point-result=void -shared-libs=%comet_utility_library_dir/libcomet_runner_utils%shlibext | FileCheck %s +# RUN: mlir-cpu-runner mv_SemiringPlusTimes_DensexCOO_oDense.llvm -O3 -e main -entry-point-result=void -shared-libs=%comet_utility_library_dir/libcomet_runner_utils%shlibext,%mlir_utility_library_dir/libomp%shlibext | FileCheck %s def main() { diff --git a/integration_test/semiring/mv_SemiringPlusTimes_DensexCSR_oDense.ta b/test/integration/semiring/mv_SemiringPlusTimes_DensexCSR_oDense.ta similarity index 89% rename from integration_test/semiring/mv_SemiringPlusTimes_DensexCSR_oDense.ta rename to test/integration/semiring/mv_SemiringPlusTimes_DensexCSR_oDense.ta index 27d3ffbb..6768b841 100644 --- a/integration_test/semiring/mv_SemiringPlusTimes_DensexCSR_oDense.ta +++ b/test/integration/semiring/mv_SemiringPlusTimes_DensexCSR_oDense.ta @@ -2,7 +2,7 @@ # Sparse matrix is in CSR format # RUN: comet-opt --convert-ta-to-it --convert-to-loops --convert-to-llvm %s &> mv_SemiringPlusTimes_DensexCSR_oDense.llvm # RUN: export SPARSE_FILE_NAME0=%comet_integration_test_data_dir/test_rank2.mtx -# RUN: mlir-cpu-runner mv_SemiringPlusTimes_DensexCSR_oDense.llvm -O3 -e main -entry-point-result=void -shared-libs=%comet_utility_library_dir/libcomet_runner_utils%shlibext | FileCheck %s +# RUN: mlir-cpu-runner mv_SemiringPlusTimes_DensexCSR_oDense.llvm -O3 -e main -entry-point-result=void -shared-libs=%comet_utility_library_dir/libcomet_runner_utils%shlibext,%mlir_utility_library_dir/libomp%shlibext | FileCheck %s def main() { diff --git a/integration_test/semiring/mv_SemiringPlusTimes_DensexDCSR_oDense.ta b/test/integration/semiring/mv_SemiringPlusTimes_DensexDCSR_oDense.ta similarity index 89% rename from integration_test/semiring/mv_SemiringPlusTimes_DensexDCSR_oDense.ta rename to test/integration/semiring/mv_SemiringPlusTimes_DensexDCSR_oDense.ta index a13ccab3..070fd40d 100644 --- a/integration_test/semiring/mv_SemiringPlusTimes_DensexDCSR_oDense.ta +++ b/test/integration/semiring/mv_SemiringPlusTimes_DensexDCSR_oDense.ta @@ -2,7 +2,7 @@ # Sparse matrix is in DCSR format # RUN: comet-opt --convert-ta-to-it --convert-to-loops --convert-to-llvm %s &> mv_SemiringPlusTimes_DensexDCSR_oDense.llvm # RUN: export SPARSE_FILE_NAME0=%comet_integration_test_data_dir/test_rank2.mtx -# RUN: mlir-cpu-runner mv_SemiringPlusTimes_DensexDCSR_oDense.llvm -O3 -e main -entry-point-result=void -shared-libs=%comet_utility_library_dir/libcomet_runner_utils%shlibext | FileCheck %s +# RUN: mlir-cpu-runner mv_SemiringPlusTimes_DensexDCSR_oDense.llvm -O3 -e main -entry-point-result=void -shared-libs=%comet_utility_library_dir/libcomet_runner_utils%shlibext,%mlir_utility_library_dir/libomp%shlibext | FileCheck %s def main() { diff --git a/integration_test/semiring/mv_SemiringPlusTimes_DensexDense_oDense.ta b/test/integration/semiring/mv_SemiringPlusTimes_DensexDense_oDense.ta similarity index 87% rename from integration_test/semiring/mv_SemiringPlusTimes_DensexDense_oDense.ta rename to test/integration/semiring/mv_SemiringPlusTimes_DensexDense_oDense.ta index e135207b..c7396104 100644 --- a/integration_test/semiring/mv_SemiringPlusTimes_DensexDense_oDense.ta +++ b/test/integration/semiring/mv_SemiringPlusTimes_DensexDense_oDense.ta @@ -1,5 +1,5 @@ # RUN: comet-opt --convert-ta-to-it --convert-to-loops --convert-to-llvm %s &> mv_SemiringPlusTimes_DensexDense_oDense.llvm -# RUN: mlir-cpu-runner mv_SemiringPlusTimes_DensexDense_oDense.llvm -O3 -e main -entry-point-result=void -shared-libs=%comet_utility_library_dir/libcomet_runner_utils%shlibext | FileCheck %s +# RUN: mlir-cpu-runner mv_SemiringPlusTimes_DensexDense_oDense.llvm -O3 -e main -entry-point-result=void -shared-libs=%comet_utility_library_dir/libcomet_runner_utils%shlibext,%mlir_utility_library_dir/libomp%shlibext | FileCheck %s def main() { #IndexLabel Declarations diff --git a/tools/README.md b/tools/README.md new file mode 100644 index 00000000..09595c90 --- /dev/null +++ b/tools/README.md @@ -0,0 +1,56 @@ +# COMET Compilation Tools + + +## SPIRV-LLVM Translator +1. If not already done, initialize the relevant submodules +```bash +git submodule update --init . +``` +2. Create a link from spirv-llvm-translate to llvm-spirv projects directory +```bash +ln -s ${PWD}/spirv-llvm-translate llvm-spirv/llvm/projects +``` +3. Build LLVM-SPIRV translator +```bash +mkdir llvm-spirv/build +cd llvm-spirv/build +cmake -GNinja ../llvm \ + -DLLVM_ENABLE_PROJECTS=clang \ + -DLLVM_TARGETS_TO_BUILD=host +ninja llvm-spirv +ninja llvm-dis +``` + +## Build MCL runtime +Follow directions in `/runtimes/mcl` directory + +## Running FPGA kernel generated from COMET +1. Generate the runnable from COMET +```bash +comet-op --target=FPGA --convert-ta-to-it --convert-to-loops --convert-to-llvm --xclbin_path= &> run.mlir +``` +This command will generate a `.bin` file at `.bin>` +and print the host code to `run.mlir`. + +Note that the path passed to ``--xclbin-path`` needs to be the path to the final xclbin file generated after all the following steps have been run. + + +2. Converting the `.bin` file to `.xclbin` using the `spirv-to-xclbin.py` script +``` +python3 spirv-to-xclbin.py >.bin -k -l -o +``` +This should generate (among others) the `.xclbin` file at the location .xlcbin + +3. Running on the FPGA + + a. First, start MCL + + b. Run the host code + +```bash +/mcl_scheduler & +``` + +``` +/mlir-cpu-runner run.mlir -O3 -e main -entry-point-result=void -shared-libs=/libcomet_runner_utils.so,/libomp.so +``` diff --git a/tools/llvm-spirv b/tools/llvm-spirv new file mode 160000 index 00000000..4856a933 --- /dev/null +++ b/tools/llvm-spirv @@ -0,0 +1 @@ +Subproject commit 4856a9330ee01d30e9e11b6c2f991662b4c04b07 diff --git a/tools/spirv-llvm-translate b/tools/spirv-llvm-translate new file mode 160000 index 00000000..791d407f --- /dev/null +++ b/tools/spirv-llvm-translate @@ -0,0 +1 @@ +Subproject commit 791d407f7f8376c72584a63f6ffe29cc5759da14 diff --git a/tools/spirv-to-xclbin.py b/tools/spirv-to-xclbin.py new file mode 100644 index 00000000..1a027033 --- /dev/null +++ b/tools/spirv-to-xclbin.py @@ -0,0 +1,130 @@ +import shutil +import subprocess +import sys +import argparse +import os +from pathlib import Path + +libspir_hls = 'libspir64-39-hls.bc' +libsqlite = 'libsqlite3.28.0.so' +clang_version = '3.9' + +def retrieve_xilinx_paths(): + if os.getenv('VPP_PATH'): + vpp_path = os.getenv('VPP_PATH') + else : + vpp_path = shutil.which('v++') + + if vpp_path is None: + raise RuntimeError('Path to v++ not found. Please set the environmental variable "VPP_PATH" to the v++ executable location') + + path = Path(vpp_path) + vpp_path = path.parents[0] + version = path.parents[1].name + vitis_hls_path = path.parents[3].joinpath('Vitis_HLS', version) + + # Usually in ../Xilinx/Vits_HLS//lnx64/lib/ + libspir_hls_path = vitis_hls_path.joinpath('lnx64', 'lib', f'{libspir_hls}') + + if not libspir_hls_path.exists() : + libspir_hls_path = os.getenv('LIBSPIR_HLS_PATH') + if libspir_hls_path is None: + raise RuntimeError(f'Path to libspir-*-hls.bc not found') + + # Usually in ../Xilinx/Vits_HLS//lib/lnx64.o/lib/ + libsqllite_path = vitis_hls_path.joinpath('lib', 'lnx64.o') + + if not libsqllite_path.joinpath(libsqlite).exists() : + libsqllite_path = os.getenv('LIBSQLLITE_PATH') + if libsqllite_path is None: + raise RuntimeError(f'Path to libsqlite not. Please set the environmental variable "LIBSQLLITE_PATH" to the directory containing it') + + llvm_hls_path = vitis_hls_path.joinpath('lnx64', 'tools', f'clang-{clang_version}-csynth', 'bin') + + # Usually in ../Xilinx/Vits_HLS//lnx64/tools/clang--csynth/bin/ + if not llvm_hls_path.joinpath('llvm-as').exists() : + llvm_hls_path = os.getenv('LLVM_HLS_BIN_PATH') + if llvm_hls_path is None: + raise RuntimeError(f'Path to Xilinx llvm binaries not found. Please set the environmental variable "LLVM_HLS_BIN_PATH to the directory containing them') + + # Usually in ../Xilinx/Vits_HLS//lnx64/tools/clang--csynth/bin/ + if not llvm_hls_path.joinpath('llvm-link').exists() : + raise RuntimeError(f'Path to Xilinx llvm-link not found') + + return (libspir_hls_path, libsqllite_path, llvm_hls_path, vpp_path) + +llvm_spirv_path = '/qfs/people/thom895/TwoTruths/COMET/tools/llvm-spirv/build/bin/' +libspir_hls_path, lib_sqllite_path, llvm_hls_path, vpp_path = retrieve_xilinx_paths() +platform = 'xilinx_vck5000_gen3x16_xdma_1_202120_1' + +def convert_ocl_ids_storage_class(input, output): + ret = subprocess.call([f'{llvm_spirv_path}/llvm-spirv', '-to-text', '-o', f'{output}.spt', input]) + if ret != 0 : + print('Error in converting SPIRV binary to text') + + lines = None + with open(f'{output}.spt', 'r') as f: + vars = [] + lines = f.readlines() + for i, line in enumerate(lines): + if 'BuiltIn 26' in line or 'BuiltIn 27' in line or 'BuiltIn 28' in line: + if 'Decorate ' in line: + vars.append(line.split()[2]) + elif '4 Variable 1' in line and line.strip().endswith('1'): + lines[i] = line.strip()[:-1]+'0\n' + + with open(f'{output}.spt', 'w') as f: + f.writelines(lines) + + ret = subprocess.call([f'{llvm_spirv_path}/llvm-spirv', '-to-binary', '-o', f'{output}.spv', f'{output}.spt']) + if ret != 0 : + print('Error in converting modified SPIRV text to binary') + +def generate_xclbin(input, output, kernel, platform): + + convert_ocl_ids_storage_class(input, output) + + ret = subprocess.call([f'{llvm_spirv_path}/llvm-spirv', '-r', '-o', f'{output}.bc', f'{output}.spv']) + if ret != 0 : + print('Error in converting SPIRV binary to bitcode') + + ret = subprocess.call([f'{llvm_spirv_path}/llvm-dis', '-o', f'{output}.ll', f'{output}.bc']) + if ret != 0 : + print('Error in converting SPIRV bitcode to LLVMIR') + + my_env = os.environ.copy() + if 'LD_LIBRARY_PATH' in my_env: + my_env['LD_LIBRARY_PATH'] = f"{lib_sqllite_path}:{my_env['LD_LIBRARY_PATH']}" + else : + my_env['LD_LIBRARY_PATH'] = f'{lib_sqllite_path}' + print(my_env['LD_LIBRARY_PATH']) + + ret = subprocess.call([f'{llvm_hls_path}/llvm-as', '-o', f'{output}.xpirbc', f'{output}.ll'], env=my_env) + if ret != 0 : + print('Error in converting SPIRV LLVIR to .xpirbc') + return + + ret = subprocess.call([f'{llvm_hls_path}/llvm-link', '-o', f'{output}.linked.xpirbc', f'{output}.xpirbc', f'{libspir_hls_path}'], env=my_env) + if ret != 0 : + print(f'Error in linking {output}.xpirbc ') + return + + ret = subprocess.call([f'{vpp_path}/v++', '--platform', platform, '-c', '-k', kernel, '--temp_dir', f'./{output}.temp', '-o', f'{output}.linked.xo', f'{output}.linked.xpirbc'], env=my_env) + if ret != 0 : + print('Error in v++ compilation ') + return + + ret = subprocess.call([f'{vpp_path}/v++', '--platform', platform, '-l', '--temp_dir', f'./{output}.temp_link', '-o', f'{output}.linked.xclbin', f'{output}.linked.xo'], env=my_env) + if ret != 0 : + print('Error in v++ linking') + return + +arg_parser = argparse.ArgumentParser(description='SPIRV-to-XCLBIN Converter') +arg_parser.add_argument( dest='ifile', metavar='FILE', type=str, help='Input file') +arg_parser.add_argument('-o', '--output', dest='ofile', required=True, help='Output file') +arg_parser.add_argument('-k', '--kernel', dest='kernel', required=True, help='Kernel name') +arg_parser.add_argument('-l', '--platform', dest='platform', required=True, help='Xilinx platform to target. You can check the available platforms with platoformid -l') +args = arg_parser.parse_args() + + +generate_xclbin(args.ifile, args.ofile, args.kernel, args.platform) \ No newline at end of file diff --git a/triton b/triton index 6fc969d8..fb903856 160000 --- a/triton +++ b/triton @@ -1 +1 @@ -Subproject commit 6fc969d896bd65e48b2dba1394d92b542652e5a4 +Subproject commit fb9038560bc032b9099aeff52877ab068794a9e2 diff --git a/triton.patch b/triton.patch index f367dee2..17c91eb5 100644 --- a/triton.patch +++ b/triton.patch @@ -1,39 +1,81 @@ diff --git a/CMakeLists.txt b/CMakeLists.txt -index 309855e..4bc4db7 100644 +index a892a666d..5d2099b58 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt -@@ -20,7 +20,7 @@ if(NOT WIN32) - endif() +@@ -55,6 +55,7 @@ endif() + include_directories(${CMAKE_CURRENT_SOURCE_DIR}/include) + set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -D__STDC_FORMAT_MACROS -fPIC -std=gnu++17") ++add_compile_options(-w) -- -+set_directory_properties(PROPERTIES COMPILE_OPTIONS "-w") - # Options - option(TRITON_BUILD_TUTORIALS "Build C++ Triton tutorials" ON) - option(TRITON_BUILD_PYTHON_MODULE "Build Python Triton bindings" OFF) -@@ -103,7 +103,7 @@ endfunction() + # ######### + # LLVM +@@ -107,7 +108,7 @@ endfunction() # Disable warnings that show up in external code (gtest;pybind11) -set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -Werror -Wno-covered-switch-default -fvisibility=hidden") -+set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -Wno-covered-switch-default -fvisibility=hidden") ++set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -Wno-covered-switch-default -fvisibility=hidden") include_directories(".") include_directories(${MLIR_INCLUDE_DIRS}) -@@ -220,6 +220,6 @@ if(TRITON_BUILD_PYTHON_MODULE AND NOT WIN32) - target_link_libraries(triton PRIVATE ${PYTHON_LDFLAGS}) +@@ -290,8 +291,8 @@ if(NOT TRITON_BUILD_PYTHON_MODULE) endif() + add_subdirectory(third_party/f2reduce) -add_subdirectory(bin) -add_subdirectory(test) --add_subdirectory(unittest) -+# add_subdirectory(bin) -+# add_subdirectory(test) -+# add_subdirectory(unittest) -diff --git a/include/triton/Dialect/NVGPU/IR/CMakeLists.txt b/include/triton/Dialect/NVGPU/IR/CMakeLists.txt -index f8932cd..37fb664 100644 ---- a/include/triton/Dialect/NVGPU/IR/CMakeLists.txt -+++ b/include/triton/Dialect/NVGPU/IR/CMakeLists.txt ++#add_subdirectory(bin) ++#add_subdirectory(test) + + if(TRITON_BUILD_UT) + add_subdirectory(unittest) +diff --git a/third_party/amd/CMakeLists.txt b/third_party/amd/CMakeLists.txt +index 8228c3d39..ccc124c80 100644 +--- a/third_party/amd/CMakeLists.txt ++++ b/third_party/amd/CMakeLists.txt +@@ -5,6 +5,6 @@ add_subdirectory(lib) + if(TRITON_BUILD_PYTHON_MODULE) + add_triton_plugin(TritonAMD ${CMAKE_CURRENT_SOURCE_DIR}/python/triton_amd.cc LINK_LIBS TritonAMDGPUToLLVM TritonAMDGPUTransforms TritonAMDGPUDialectToLLVM) + endif() +-if(TRITON_BUILD_UT) +- add_subdirectory(unittest) +-endif() ++#if(TRITON_BUILD_UT) ++# add_subdirectory(unittest) ++#endif() +diff --git a/third_party/amd/include/Dialect/TritonAMDGPU/IR/CMakeLists.txt b/third_party/amd/include/Dialect/TritonAMDGPU/IR/CMakeLists.txt +index 25a57075b..4d23ecf7d 100644 +--- a/third_party/amd/include/Dialect/TritonAMDGPU/IR/CMakeLists.txt ++++ b/third_party/amd/include/Dialect/TritonAMDGPU/IR/CMakeLists.txt +@@ -1,8 +1,8 @@ + set(MLIR_BINARY_DIR ${CMAKE_BINARY_DIR}) + + set(LLVM_TARGET_DEFINITIONS TritonAMDGPUOps.td) +-mlir_tablegen(Dialect.h.inc -gen-dialect-decls -dialect=amdgpu) +-mlir_tablegen(Dialect.cpp.inc -gen-dialect-defs -dialect=amdgpu) ++mlir_tablegen(Dialect.h.inc -gen-dialect-decls -dialect=ttamdgpu) ++mlir_tablegen(Dialect.cpp.inc -gen-dialect-defs -dialect=ttamdgpu) + mlir_tablegen(OpsConversions.inc -gen-llvmir-conversions) + mlir_tablegen(Ops.h.inc -gen-op-decls) + mlir_tablegen(Ops.cpp.inc -gen-op-defs) +diff --git a/third_party/amd/include/Dialect/TritonAMDGPU/IR/TritonAMDGPUDialect.td b/third_party/amd/include/Dialect/TritonAMDGPU/IR/TritonAMDGPUDialect.td +index c0c18b07e..25f56039a 100644 +--- a/third_party/amd/include/Dialect/TritonAMDGPU/IR/TritonAMDGPUDialect.td ++++ b/third_party/amd/include/Dialect/TritonAMDGPU/IR/TritonAMDGPUDialect.td +@@ -27,7 +27,7 @@ + include "mlir/IR/OpBase.td" + + def TritonAMDGPU_Dialect : Dialect { +- let name = "amdgpu"; ++ let name = "ttamdgpu"; + let cppNamespace = "::mlir::triton::amdgpu"; + + let description = [{ +diff --git a/third_party/nvidia/include/Dialect/NVGPU/IR/CMakeLists.txt b/third_party/nvidia/include/Dialect/NVGPU/IR/CMakeLists.txt +index f8932cdc4..37fb66417 100644 +--- a/third_party/nvidia/include/Dialect/NVGPU/IR/CMakeLists.txt ++++ b/third_party/nvidia/include/Dialect/NVGPU/IR/CMakeLists.txt @@ -1,8 +1,8 @@ set(MLIR_BINARY_DIR ${CMAKE_BINARY_DIR}) @@ -45,10 +87,10 @@ index f8932cd..37fb664 100644 mlir_tablegen(OpsConversions.inc -gen-llvmir-conversions) mlir_tablegen(Ops.h.inc -gen-op-decls) mlir_tablegen(Ops.cpp.inc -gen-op-defs) -diff --git a/include/triton/Dialect/NVGPU/IR/NVGPUDialect.td b/include/triton/Dialect/NVGPU/IR/NVGPUDialect.td -index 6978173..3b76e10 100644 ---- a/include/triton/Dialect/NVGPU/IR/NVGPUDialect.td -+++ b/include/triton/Dialect/NVGPU/IR/NVGPUDialect.td +diff --git a/third_party/nvidia/include/Dialect/NVGPU/IR/NVGPUDialect.td b/third_party/nvidia/include/Dialect/NVGPU/IR/NVGPUDialect.td +index 6978173d4..3b76e10c7 100644 +--- a/third_party/nvidia/include/Dialect/NVGPU/IR/NVGPUDialect.td ++++ b/third_party/nvidia/include/Dialect/NVGPU/IR/NVGPUDialect.td @@ -25,7 +25,7 @@ include "mlir/IR/OpBase.td"