├── .github └── workflows │ ├── buildandtest.yml │ └── cuda-test-manual.yml ├── .gitignore ├── .gitmodules ├── CMakeLists.txt ├── Doxyfile ├── LICENSE ├── README.md ├── apps ├── CMakeLists.txt └── tensor_times_vector │ ├── CMakeLists.txt │ ├── README.md │ └── tensor_times_vector.cpp ├── ci ├── test-pr.md └── test-pr.py ├── include ├── taco.h └── taco │ ├── README.md │ ├── codegen │ └── module.h │ ├── component.h │ ├── coordinate.h │ ├── cuda.h │ ├── error.h │ ├── error │ └── error_messages.h │ ├── format.h │ ├── index_notation │ ├── README.md │ ├── index_notation.h │ ├── index_notation_nodes.h │ ├── index_notation_nodes_abstract.h │ ├── index_notation_printer.h │ ├── index_notation_rewriter.h │ ├── index_notation_visitor.h │ ├── intrinsic.h │ ├── iteration_algebra.h │ ├── iteration_algebra_printer.h │ ├── kernel.h │ ├── properties.h │ ├── property_pointers.h │ ├── provenance_graph.h │ ├── schedule.h │ ├── tensor_operator.h │ └── transformations.h │ ├── ir │ ├── ir.h │ ├── ir_generators.h │ ├── ir_printer.h │ ├── ir_rewriter.h │ ├── ir_verifier.h │ ├── ir_visitor.h │ ├── simplify.h │ └── workspace_rewriter.h │ ├── ir_tags.h │ ├── lower │ ├── iterator.h │ ├── lower.h │ ├── lowerer_impl.h │ ├── lowerer_impl_imperative.h │ ├── merge_lattice.h │ ├── mode.h │ ├── mode_format_compressed.h │ ├── mode_format_dense.h │ ├── mode_format_impl.h │ └── mode_format_singleton.h │ ├── parser │ ├── einsum_parser.h │ ├── lexer.h │ ├── parser.h │ └── schedule_parser.h │ ├── storage │ ├── array.h │ ├── coordinate.h │ ├── file_io_mtx.h │ ├── file_io_rb.h │ ├── file_io_tns.h │ ├── index.h │ ├── pack.h │ ├── storage.h │ ├── typed_index.h │ ├── typed_value.h │ └── typed_vector.h │ ├── taco_tensor_t.h │ ├── target.h │ ├── tensor.h │ ├── type.h │ ├── util │ ├── collections.h │ ├── comparable.h │ ├── env.h │ ├── files.h │ ├── fill.h │ ├── functions.h │ ├── intrusive_ptr.h │ ├── name_generator.h │ ├── printable.h │ ├── scopedmap.h │ ├── scopedset.h │ ├── singleton.h │ ├── strings.h │ ├── timers.h │ ├── uncopyable.h │ └── variadic.h │ └── version.h.in ├── misc └── mainpage.h ├── python_bindings ├── CMakeLists.txt ├── Makefile ├── include │ ├── pyDatatypes.h │ ├── pyFormat.h │ ├── pyIndexNotation.h │ ├── pyParsers.h │ ├── pyTensor.h │ ├── pyTensorIO.h │ └── pyeinsum.h ├── pytaco │ ├── __init__.py │ ├── numpy_license.txt │ └── pytensor │ │ ├── __init__.py │ │ ├── taco_tensor.py │ │ └── tensorIO.py ├── source │ ├── conf.py │ ├── index.rst │ └── rst_files │ │ ├── as_np_dt.rst │ │ ├── datatype.rst │ │ ├── dtype_object.rst │ │ ├── expr_funcs.rst │ │ ├── file_io.rst │ │ ├── format.rst │ │ ├── format_class.rst │ │ ├── format_funcs.rst │ │ ├── idx_exp_obj.rst │ │ ├── index_expressions.rst │ │ ├── index_vars.rst │ │ ├── iv_funcs.rst │ │ ├── mode_format.rst │ │ ├── parsers.rst │ │ ├── sched.rst │ │ ├── tensor_class.rst │ │ ├── tensor_info.rst │ │ ├── tensors.rst │ │ └── udfs.rst ├── src │ ├── pyDatatypes.cpp │ ├── pyFormat.cpp │ ├── pyIndexNotation.cpp │ ├── pyParsers.cpp │ ├── pyTensor.cpp │ ├── pyTensorIO.cpp │ └── pytaco.cpp └── unit_tests.py ├── src ├── CMakeLists.txt ├── codegen │ ├── codegen.cpp │ ├── codegen.h │ ├── codegen_c.cpp │ ├── codegen_c.h │ ├── codegen_cuda.cpp │ ├── codegen_cuda.h │ └── module.cpp ├── cuda.cpp ├── error.cpp ├── error │ ├── error_checks.cpp │ ├── error_checks.h │ └── error_messages.cpp ├── format.cpp ├── index_notation │ ├── README.md │ ├── index_notation.cpp │ ├── index_notation_nodes.cpp │ ├── index_notation_nodes_abstract.cpp │ ├── index_notation_printer.cpp │ ├── index_notation_rewriter.cpp │ ├── index_notation_visitor.cpp │ ├── intrinsic.cpp │ ├── iteration_algebra.cpp │ ├── iteration_algebra_printer.cpp │ ├── kernel.cpp │ ├── properties.cpp │ ├── property_pointers.cpp │ ├── provenance_graph.cpp │ ├── schedule.cpp │ ├── tensor_operator.cpp │ └── transformations.cpp ├── ir │ ├── README.md │ ├── ir.cpp │ ├── ir_generators.cpp │ ├── ir_printer.cpp │ ├── ir_rewriter.cpp │ ├── ir_verifier.cpp │ ├── ir_visitor.cpp │ ├── simplify.cpp │ └── workspace_rewriter.cpp ├── ir_tags.cpp ├── lower │ ├── README.md │ ├── expr_tools.cpp │ ├── expr_tools.h │ ├── iteration_forest.cpp │ ├── iteration_forest.h │ ├── iteration_graph.cpp │ ├── iteration_graph.h │ ├── iterator.cpp │ ├── lower.cpp │ ├── lowerer_impl.cpp │ ├── lowerer_impl_imperative.cpp │ ├── merge_lattice.cpp │ ├── mode.cpp │ ├── mode_access.cpp │ ├── mode_access.h │ ├── mode_format_compressed.cpp │ ├── mode_format_dense.cpp │ ├── mode_format_impl.cpp │ ├── mode_format_singleton.cpp │ ├── tensor_path.cpp │ └── tensor_path.h ├── parser │ ├── README.md │ ├── einsum_parser.cpp │ ├── lexer.cpp │ ├── parser.cpp │ └── schedule_parser.cpp ├── storage │ ├── README.md │ ├── array.cpp │ ├── file_io_mtx.cpp │ ├── file_io_rb.cpp │ ├── file_io_tns.cpp │ ├── index.cpp │ ├── pack.cpp │ ├── storage.cpp │ ├── typed_index.cpp │ └── typed_value.cpp ├── taco_tensor_t.cpp ├── target.cpp ├── tensor.cpp ├── type.cpp └── util │ ├── env.cpp │ ├── files.cpp │ ├── name_generator.cpp │ └── strings.cpp ├── test ├── CMakeLists.txt ├── README_testing.md ├── data │ ├── 2tensor.mtx │ ├── 3tensor.tns │ ├── d33.mtx │ ├── d33.ttx │ ├── d432.ttx │ ├── d567.ttx │ ├── ds33.mtx │ ├── qcd │ │ ├── theta.ttx │ │ ├── theta_simple.ttx │ │ ├── z.ttx │ │ └── z_simple.ttx │ ├── rua_32.mtx │ └── rua_32.rb ├── expr_factory.cpp ├── expr_factory.h ├── gtest │ ├── LICENSE │ ├── README.TAC │ ├── gtest-all.cc │ └── gtest.h ├── op_factory.h ├── taco-cli-test.bats ├── test-iteration_algebra.cpp ├── test-typedcomponent-memory.cpp ├── test.cpp ├── test.h ├── test_properties.cpp ├── test_tensors.cpp ├── test_tensors.h ├── tests-api.cpp ├── tests-copyprop.cpp ├── tests-error.cpp ├── tests-expr-reduction.cpp ├── tests-expr.cpp ├── tests-expr_storage.cpp ├── tests-format-mm-permute.cpp ├── tests-format.cpp ├── tests-index.cpp ├── tests-index_notation.cpp ├── tests-indexexpr.cpp ├── tests-indexstmt.cpp ├── tests-io.cpp ├── tests-lower.cpp ├── tests-merge_lattice.cpp ├── tests-parafac.cpp ├── tests-qcd.cpp ├── tests-regression.cpp ├── tests-schedule-parser.cpp ├── tests-scheduling-eval.cpp ├── tests-scheduling.cpp ├── tests-storage.cpp ├── tests-storage_alloc.cpp ├── tests-tensor.cpp ├── tests-tensor_types.cpp ├── tests-transformation.cpp ├── tests-transpose.cpp ├── tests-type.cpp ├── tests-windowing.cpp └── tests-workspaces.cpp └── tools ├── CMakeLists.txt └── taco.cpp /.github/workflows/buildandtest.yml: -------------------------------------------------------------------------------- 1 | name: Build and Test 2 | 3 | on: 4 | push: 5 | branches: 6 | - master 7 | - array_algebra 8 | pull_request: 9 | branches: 10 | - master 11 | - array_algebra 12 | 13 | jobs: 14 | build-test-cpu: 15 | name: builds taco with no options for cpu and runs all tests 16 | runs-on: ubuntu-18.04 17 | 18 | steps: 19 | - uses: actions/checkout@v2 20 | - name: create_build 21 | run: mkdir build 22 | - name: cmake 23 | run: cmake .. 24 | working-directory: build 25 | - name: make 26 | run: make -j2 27 | working-directory: build 28 | - name: test 29 | run: make test 30 | env: 31 | CTEST_OUTPUT_ON_FAILURE: 1 32 | CTEST_PARALLEL_LEVEL: 2 33 | working-directory: build 34 | 35 | build-test-cpu-release: 36 | name: builds taco release for cpu and runs all tests 37 | runs-on: ubuntu-18.04 38 | 39 | steps: 40 | - uses: actions/checkout@v2 41 | - name: create_build 42 | run: mkdir build 43 | - name: cmake 44 | run: cmake -DCMAKE_BUILD_TYPE=Release .. 45 | working-directory: build 46 | - name: make 47 | run: make -j2 48 | working-directory: build 49 | - name: test 50 | run: make test 51 | env: 52 | CTEST_OUTPUT_ON_FAILURE: 1 53 | CTEST_PARALLEL_LEVEL: 2 54 | working-directory: build 55 | 56 | build-test-cpu-openmp-python-asserts: 57 | name: builds taco with compile-time asserts, openmp, and python and runs all tests 58 | runs-on: ubuntu-18.04 59 | 60 | steps: 61 | - uses: actions/checkout@v2 62 | - name: apt-get update 63 | run: sudo apt-get update 64 | - name: install numpy and scipy 65 | run: sudo DEBIAN_FRONTEND=noninteractive apt-get install -y python3-numpy python3-scipy 66 | - name: create_build 67 | run: mkdir build 68 | - name: cmake 69 | run: cmake -DCMAKE_BUILD_TYPE=RelWithDebInfo -DOPENMP=ON -DPYTHON=ON .. 70 | working-directory: build 71 | - name: make 72 | run: make -j2 73 | working-directory: build 74 | - name: test 75 | run: make test 76 | env: 77 | CTEST_OUTPUT_ON_FAILURE: 1 78 | CTEST_PARALLEL_LEVEL: 2 79 | working-directory: build 80 | 81 | build-test-python-macos-clang: 82 | name: builds taco and pytaco on macos with clang and runs all tests 83 | runs-on: macos-10.15 84 | 85 | steps: 86 | - uses: actions/checkout@v2 87 | - name: install numpy and scipy 88 | run: pip3 install numpy scipy 89 | - name: create_build 90 | run: mkdir build 91 | - name: cmake 92 | run: cmake -DCMAKE_BUILD_TYPE=Debug -DPYTHON=ON .. 93 | working-directory: build 94 | - name: make 95 | run: make -j2 96 | working-directory: build 97 | - name: test 98 | run: make test 99 | env: 100 | CTEST_OUTPUT_ON_FAILURE: 1 101 | CTEST_PARALLEL_LEVEL: 2 102 | working-directory: build 103 | -------------------------------------------------------------------------------- /.github/workflows/cuda-test-manual.yml: -------------------------------------------------------------------------------- 1 | name: "CUDA build and test (manual)" 2 | 3 | # Note: This workflow is triggered by hand by TACO developers. 4 | # It should be run after the code has been reviewed by humans. 5 | # This review step is important to ensure the safety of the 6 | # self-hosted runner. 7 | 8 | on: 9 | workflow_dispatch: 10 | inputs: 11 | CMAKE_BUILD_TYPE: 12 | description: CMAKE_BUILD_TYPE 13 | required: true 14 | default: Debug 15 | OPENMP: 16 | description: OPENMP 17 | required: true 18 | default: 'ON' 19 | PYTHON: 20 | description: PYTHON 21 | required: true 22 | default: 'OFF' 23 | jobs: 24 | ubuntu1604-cuda: 25 | name: tests ubuntu 16.04 with CUDA 9 26 | runs-on: [self-hosted, ubuntu-16.04, cuda] 27 | steps: 28 | - uses: actions/checkout@v2 29 | - name: create_build 30 | run: mkdir build 31 | - name: cmake 32 | run: cmake -DCMAKE_BUILD_TYPE=${{ github.event.inputs.CMAKE_BUILD_TYPE }} -DCUDA=ON -DOPENMP=${{ github.event.inputs.OPENMP }} -DPYTHON=${{ github.event.inputs.PYTHON }} .. 33 | working-directory: build 34 | - name: make 35 | run: make -j8 36 | working-directory: build 37 | - name: test 38 | run: make test 39 | env: 40 | CTEST_OUTPUT_ON_FAILURE: 1 41 | CTEST_PARALLEL_LEVEL: 8 42 | working-directory: build 43 | -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | build/ 2 | bin/ 3 | lib/ 4 | *__pycache__* 5 | *.swp 6 | *.swo 7 | *.DS_Store 8 | *Makefile 9 | *CMakeFiles/ 10 | *cmake_install.cmake 11 | CMakeCache.txt 12 | doc 13 | .idea/ 14 | apps/tensor_times_vector/tensor_times_vector 15 | -------------------------------------------------------------------------------- /.gitmodules: -------------------------------------------------------------------------------- 1 | [submodule "python_bindings/pybind11"] 2 | path = python_bindings/pybind11 3 | url = https://github.com/pybind/pybind11 4 | [submodule "test/bats"] 5 | path = test/bats 6 | url = https://github.com/bats-core/bats-core 7 | -------------------------------------------------------------------------------- /Doxyfile: -------------------------------------------------------------------------------- 1 | L 2 | 3 | # Project related configuration options 4 | PROJECT_NAME = taco 5 | OUTPUT_DIRECTORY = doc 6 | OUTPUT_LANGUAGE = English 7 | BRIEF_MEMBER_DESC = YES 8 | REPEAT_BRIEF = YES 9 | ALWAYS_DETAILED_SEC = NO 10 | INLINE_INHERITED_MEMB = NO 11 | JAVADOC_AUTOBRIEF = YES 12 | TAB_SIZE = 4 13 | MARKDOWN_SUPPORT = YES 14 | AUTOLINK_SUPPORT = YES 15 | BUILTIN_STL_SUPPORT = YES 16 | DISTRIBUTE_GROUP_DOC = YES 17 | 18 | # Build related configuration options 19 | EXTRACT_ALL = YES 20 | CLASS_DIAGRAMS = YES 21 | HAVE_DOT = YES 22 | COLLABORATION_GRAPH = NO 23 | GENERATE_TODOLIST = YES 24 | GENERATE_TESTLIST = YES 25 | GENERATE_BUGLIST = NO 26 | GENERATE_DEPRECATEDLIST= YES 27 | SHOW_FILES = YES 28 | SHOW_NAMESPACES = YES 29 | 30 | # Configuration options related to the input files 31 | INPUT = include src misc/mainpage.h 32 | FILE_PATTERNS = *.h 33 | RECURSIVE = YES 34 | EXCLUDE = build xcode test 35 | EXAMPLE_PATH = apps 36 | EXAMPLE_RECURSIVE = YES 37 | 38 | # Configuration options related to source browsing 39 | SOURCE_BROWSER = NO 40 | INLINE_SOURCES = NO 41 | STRIP_CODE_COMMENTS = YES 42 | REFERENCED_BY_RELATION = NO 43 | REFERENCES_RELATION = NO 44 | REFERENCES_LINK_SOURCE = YES 45 | SOURCE_TOOLTIPS = YES 46 | 47 | # Configuration options related to the alphabetical class index 48 | ALPHABETICAL_INDEX = NO 49 | 50 | # Configuration options related to the HTML output 51 | GENERATE_HTML = YES 52 | HTML_OUTPUT = . 53 | HTML_FILE_EXTENSION = .html 54 | HTML_TIMESTAMP = NO 55 | HTML_DYNAMIC_SECTIONS = NO 56 | GENERATE_TREEVIEW = YES 57 | SEARCHENGINE = YES 58 | 59 | # Configuration options related to the LaTeX output 60 | GENERATE_LATEX = NO 61 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2017-2019 MIT CSAIL, Adobe Systems, and other contributors. 4 | 5 | Developed by: 6 | 7 | The taco team 8 | http://tensor-compiler.org 9 | 10 | Permission is hereby granted, free of charge, to any person obtaining a copy 11 | of this software and associated documentation files (the "Software"), to deal 12 | in the Software without restriction, including without limitation the rights 13 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 14 | copies of the Software, and to permit persons to whom the Software is 15 | furnished to do so, subject to the following conditions: 16 | 17 | The above copyright notice and this permission notice shall be included in all 18 | copies or substantial portions of the Software. 19 | 20 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 21 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 22 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 23 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 24 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 25 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 26 | SOFTWARE. 27 | -------------------------------------------------------------------------------- /apps/CMakeLists.txt: -------------------------------------------------------------------------------- 1 | add_subdirectory(tensor_times_vector) 2 | -------------------------------------------------------------------------------- /apps/tensor_times_vector/CMakeLists.txt: -------------------------------------------------------------------------------- 1 | cmake_minimum_required(VERSION 2.8.12) 2 | if(POLICY CMP0048) 3 | cmake_policy(SET CMP0048 NEW) 4 | endif() 5 | project(tensor_times_vector) 6 | 7 | set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -std=c++11") 8 | file(GLOB SOURCE_CODE ${PROJECT_SOURCE_DIR}/*.cpp) 9 | add_executable(${PROJECT_NAME} ${SOURCE_CODE}) 10 | 11 | # To let the app be a standalone project 12 | if (NOT TACO_INCLUDE_DIR) 13 | if (NOT DEFINED ENV{TACO_INCLUDE_DIR} OR NOT DEFINED ENV{TACO_LIBRARY_DIR}) 14 | message(FATAL_ERROR "Set the environment variables TACO_INCLUDE_DIR and TACO_LIBRARY_DIR") 15 | endif () 16 | set(TACO_INCLUDE_DIR $ENV{TACO_INCLUDE_DIR}) 17 | set(TACO_LIBRARY_DIR $ENV{TACO_LIBRARY_DIR}) 18 | find_library(taco taco ${TACO_LIBRARY_DIR}) 19 | target_link_libraries(${PROJECT_NAME} LINK_PUBLIC ${taco}) 20 | else() 21 | set_target_properties("${PROJECT_NAME}" PROPERTIES OUTPUT_NAME "taco-${PROJECT_NAME}") 22 | target_link_libraries(${PROJECT_NAME} LINK_PUBLIC taco) 23 | endif () 24 | 25 | # Include taco headers 26 | include_directories(${TACO_INCLUDE_DIR}) 27 | -------------------------------------------------------------------------------- /apps/tensor_times_vector/README.md: -------------------------------------------------------------------------------- 1 | If you want to use it as a standalone app, 2 | Point the cmake build system to taco like so: 3 | 4 | export TACO_INCLUDE_DIR= 5 | export TACO_LIBRARY_DIR= 6 | 7 | Build the tensor_times_vector example like so: 8 | 9 | mkdir build 10 | cd build 11 | cmake .. 12 | make 13 | 14 | Run the explicit tensor_times_vector example like so: 15 | 16 | ./tensor_times_vector 17 | -------------------------------------------------------------------------------- /apps/tensor_times_vector/tensor_times_vector.cpp: -------------------------------------------------------------------------------- 1 | #include 2 | #include "taco.h" 3 | 4 | using namespace taco; 5 | 6 | int main(int argc, char* argv[]) { 7 | Format csr({Dense,Sparse}); 8 | Format csf({Sparse,Sparse,Sparse}); 9 | Format sv({Sparse}); 10 | 11 | Tensor A("A", {2,3}, csr); 12 | Tensor B("B", {2,3,4}, csf); 13 | Tensor c("c", {4}, sv); 14 | 15 | // Insert data into B and c 16 | B(0,0,0) = 1.0; 17 | B(1,2,0) = 2.0; 18 | B(1,2,1) = 3.0; 19 | c(0) = 4.0; 20 | c(1) = 5.0; 21 | 22 | IndexVar i, j, k; 23 | A(i,j) = B(i,j,k) * c(k); 24 | 25 | std::cout << A << std::endl; 26 | } 27 | -------------------------------------------------------------------------------- /include/taco.h: -------------------------------------------------------------------------------- 1 | #ifndef TACO_TACO_H 2 | #define TACO_TACO_H 3 | 4 | #include "taco/tensor.h" 5 | #include "taco/format.h" 6 | #include "taco/index_notation/tensor_operator.h" 7 | #include "taco/index_notation/index_notation.h" 8 | 9 | #endif 10 | -------------------------------------------------------------------------------- /include/taco/README.md: -------------------------------------------------------------------------------- 1 | taco consists of two intertwined processes: one that compiles index 2 | expressions to executable code and one that invokes executable code on 3 | inputs to produce outputs. 4 | 5 | The main modules are layered the following way: 6 | 7 | ![Modules](https://github.com/tensor-compiler/taco/wiki/images/modules.png) 8 | 9 | The modules work together the following way to compile and execute code: 10 | 11 | ![Flow](https://github.com/tensor-compiler/taco/wiki/images/flow.png) 12 | -------------------------------------------------------------------------------- /include/taco/codegen/module.h: -------------------------------------------------------------------------------- 1 | #ifndef TACO_MODULE_H 2 | #define TACO_MODULE_H 3 | 4 | #include 5 | #include 6 | #include 7 | #include 8 | #include 9 | 10 | #include "taco/target.h" 11 | #include "taco/ir/ir.h" 12 | 13 | namespace taco { 14 | namespace ir { 15 | 16 | class Module { 17 | public: 18 | /// Create a module for some target 19 | Module(Target target=getTargetFromEnvironment()) 20 | : lib_handle(nullptr), moduleFromUserSource(false), target(target) { 21 | setJITLibname(); 22 | setJITTmpdir(); 23 | } 24 | 25 | /// Compile the source into a library, returning its full path 26 | std::string compile(); 27 | 28 | /// Compile the module into a source file located at the specified location 29 | /// path and prefix. The generated source will be path/prefix.{.c|.bc, .h} 30 | void compileToSource(std::string path, std::string prefix); 31 | 32 | /// Compile the module into a static library located at the specified location 33 | /// path and prefix. The generated library will be path/prefix.a 34 | void compileToStaticLibrary(std::string path, std::string prefix); 35 | 36 | /// Add a lowered function to this module */ 37 | void addFunction(Stmt func); 38 | 39 | /// Get the source of the module as a string */ 40 | std::string getSource(); 41 | 42 | /// Get a function pointer to a compiled function. This returns a void* 43 | /// pointer, which the caller is required to cast to the correct function type 44 | /// before calling. If there's no function of this name then a nullptr is 45 | /// returned. 46 | void* getFuncPtr(std::string name); 47 | 48 | /// Call a raw function in this module and return the result 49 | int callFuncPackedRaw(std::string name, void** args); 50 | 51 | /// Call a raw function in this module and return the result 52 | int callFuncPackedRaw(std::string name, std::vector args) { 53 | return callFuncPackedRaw(name, args.data()); 54 | } 55 | 56 | /// Call a function using the taco_tensor_t interface and return the result 57 | int callFuncPacked(std::string name, void** args) { 58 | return callFuncPackedRaw("_shim_"+name, args); 59 | } 60 | 61 | /// Call a function using the taco_tensor_t interface and return the result 62 | int callFuncPacked(std::string name, std::vector args) { 63 | return callFuncPacked(name, args.data()); 64 | } 65 | 66 | /// Set the source of the module 67 | void setSource(std::string source); 68 | 69 | private: 70 | std::stringstream source; 71 | std::stringstream header; 72 | std::string libname; 73 | std::string tmpdir; 74 | void* lib_handle; 75 | std::vector funcs; 76 | 77 | // true iff the module was created from user-provided source 78 | bool moduleFromUserSource; 79 | 80 | Target target; 81 | 82 | void setJITLibname(); 83 | void setJITTmpdir(); 84 | 85 | static std::string chars; 86 | static std::default_random_engine gen; 87 | static std::uniform_int_distribution randint; 88 | }; 89 | 90 | } // namespace ir 91 | } // namespace taco 92 | #endif 93 | -------------------------------------------------------------------------------- /include/taco/component.h: -------------------------------------------------------------------------------- 1 | #ifndef TACO_COMPONENT_H 2 | #define TACO_COMPONENT_H 3 | 4 | #include 5 | #include 6 | #include 7 | #include 8 | #include "taco/coordinate.h" 9 | #include "taco/util/comparable.h" 10 | #include "taco/util/strings.h" 11 | #include "taco/error.h" 12 | #include "taco/error/error_messages.h" 13 | 14 | namespace taco { 15 | 16 | /// Structure to hold a non zero as a tuple (coordinate, value). 17 | /// 18 | /// CType the type of the value stored. 19 | /// Order the number of dimensions of the component 20 | template 21 | class Component { 22 | public: 23 | Component() : coord(), val(0) {} 24 | 25 | Component(Coordinate coordinate, CType value) : coord(coordinate), val(value) { 26 | taco_uassert(coord.order() == Order) << 27 | "Wrong number of indices"; 28 | } 29 | 30 | size_t coordinate(int mode) const { 31 | taco_uassert(mode < Order) << "requested mode coordinate exceeds order of component."; 32 | return coord[mode]; 33 | } 34 | 35 | const Coordinate coordinate() const { 36 | return coord; 37 | } 38 | 39 | const CType& value() const { return val; } 40 | 41 | private: 42 | Coordinate coord; 43 | CType val; 44 | }; 45 | 46 | } 47 | #endif -------------------------------------------------------------------------------- /include/taco/coordinate.h: -------------------------------------------------------------------------------- 1 | #ifndef TACO_COORDINATE_H 2 | #define TACO_COORDINATE_H 3 | 4 | #include 5 | #include 6 | #include 7 | #include 8 | #include "taco/util/comparable.h" 9 | #include "taco/util/strings.h" 10 | #include "taco/error.h" 11 | 12 | namespace taco { 13 | 14 | /// Structure to represent a multidimensional coordinate tuple (i,j,k,...). 15 | template 16 | class Coordinate : util::Comparable> { 17 | public: 18 | Coordinate() {} 19 | 20 | template 21 | Coordinate(T... coordinates) : coordinates{{coordinates...}} {} 22 | 23 | Type& operator[](size_t idx) { 24 | taco_iassert(idx < Order); 25 | return coordinates[idx]; 26 | } 27 | 28 | const Type& operator[](size_t idx) const { 29 | taco_iassert(idx < Order); 30 | return coordinates[idx]; 31 | } 32 | 33 | operator std::vector() const { 34 | std::vector vec; 35 | for (auto coord : coordinates) { 36 | vec.push_back(coord); 37 | } 38 | return vec; 39 | } 40 | 41 | size_t order() { 42 | return Order; 43 | } 44 | 45 | // friend methods 46 | 47 | template 48 | friend bool operator==(const Coordinate& a, const Coordinate& b) { 49 | for (size_t i = 0; i < Order; i++) { 50 | if (a[i] != b[i]) return false; 51 | } 52 | return true; 53 | } 54 | 55 | template 56 | friend bool operator<(const Coordinate& a, const Coordinate& b) { 57 | for (size_t i = 0; i < Order; i++) { 58 | if (a[i] < b[i]) return true; 59 | if (a[i] > b[i]) return false; 60 | } 61 | return false; 62 | } 63 | 64 | template 65 | friend std::ostream& operator<<(std::ostream& os, const Coordinate& c) { 66 | return os << util::join(c.coordinates); 67 | } 68 | 69 | template 70 | friend std::ostream& operator<<(std::ostream& os, 71 | const std::pair, V>& c) { 72 | return os << "(" << util::join(c.first.coordinates, ",") << "):" << c.second; 73 | } 74 | 75 | private: 76 | std::array coordinates; 77 | }; 78 | 79 | } 80 | #endif -------------------------------------------------------------------------------- /include/taco/cuda.h: -------------------------------------------------------------------------------- 1 | #ifndef CUDA_H 2 | #define CUDA_H 3 | 4 | #include 5 | #include 6 | #include 7 | 8 | #ifndef CUDA_BUILT 9 | #define CUDA_BUILT false 10 | #endif 11 | 12 | namespace taco { 13 | /// Functions used by taco to interface with CUDA (especially unified memory) 14 | /// Check if should use CUDA codegen 15 | bool should_use_CUDA_codegen(); 16 | /// Check if should use CUDA unified memory 17 | bool should_use_CUDA_unified_memory(); 18 | /// Enable/Disable CUDA codegen 19 | void set_CUDA_codegen_enabled(bool enabled); 20 | /// Enable/Disable CUDA unified memory 21 | void set_CUDA_unified_memory_enabled(bool enabled); 22 | /// Gets default compiler flags by checking current gpu model 23 | std::string get_default_CUDA_compiler_flags(); 24 | /// Allocates memory using unified memory (and checks for errors) 25 | void* cuda_unified_alloc(size_t size); 26 | /// Frees memory from unified memory (and checks for errors) 27 | void cuda_unified_free(void *ptr); 28 | 29 | } 30 | 31 | #endif 32 | -------------------------------------------------------------------------------- /include/taco/error/error_messages.h: -------------------------------------------------------------------------------- 1 | #ifndef TACO_ERROR_MESSAGES_H 2 | #define TACO_ERROR_MESSAGES_H 3 | 4 | #include 5 | 6 | namespace taco { 7 | namespace error { 8 | 9 | // unsupported type bit width error 10 | extern const std::string type_mismatch; 11 | extern const std::string type_bitwidt; 12 | 13 | // TensorVar::setIndexExpression error messages 14 | extern const std::string expr_dimension_mismatch; 15 | extern const std::string expr_transposition; 16 | extern const std::string expr_distribution; 17 | extern const std::string expr_einsum_missformed; 18 | 19 | // compile error messages 20 | extern const std::string compile_without_expr; 21 | extern const std::string compile_tensor_name_collision; 22 | 23 | // assemble error messages 24 | extern const std::string assemble_without_compile; 25 | 26 | // compute error messages 27 | extern const std::string compute_without_compile; 28 | 29 | // factory function error messages 30 | extern const std::string requires_matrix; 31 | 32 | #define INIT_REASON(reason) \ 33 | string reason_; \ 34 | do { \ 35 | if (reason == nullptr) { \ 36 | reason = &reason_; \ 37 | } \ 38 | *reason = ""; \ 39 | } while (0) 40 | 41 | }} 42 | #endif 43 | -------------------------------------------------------------------------------- /include/taco/index_notation/index_notation_nodes_abstract.h: -------------------------------------------------------------------------------- 1 | #ifndef TACO_INDEX_NOTATION_NODES_ABSTRACT_H 2 | #define TACO_INDEX_NOTATION_NODES_ABSTRACT_H 3 | 4 | #include 5 | #include 6 | 7 | #include "taco/type.h" 8 | #include "taco/util/uncopyable.h" 9 | #include "taco/util/intrusive_ptr.h" 10 | 11 | namespace taco { 12 | 13 | class TensorVar; 14 | class IndexVar; 15 | class IndexExprVisitorStrict; 16 | class IndexStmtVisitorStrict; 17 | class OperatorSplit; 18 | class Precompute; 19 | 20 | /// A node of a scalar index expression tree. 21 | struct IndexExprNode : public util::Manageable, 22 | private util::Uncopyable { 23 | public: 24 | IndexExprNode(); 25 | IndexExprNode(Datatype type); 26 | virtual ~IndexExprNode() = default; 27 | virtual void accept(IndexExprVisitorStrict*) const = 0; 28 | 29 | /// Return the scalar data type of the index expression. 30 | Datatype getDataType() const; 31 | 32 | /// Store the index expression's result to the given workspace w.r.t. index 33 | /// variable `i` and replace the index expression (in the enclosing 34 | /// expression) with a workspace access expression. The index variable `i` is 35 | /// retained in the enclosing expression and used to access the workspace, 36 | /// while `iw` replaces `i` in the index expression that computes workspace 37 | /// results. 38 | void setWorkspace(IndexVar i, IndexVar iw, TensorVar workspace) const; 39 | 40 | /// Return a workspace scheduling construct that describes the workspace to 41 | /// store expression to. 42 | Precompute getWorkspace() const; 43 | 44 | private: 45 | Datatype dataType; 46 | 47 | mutable std::shared_ptr> workspace; 48 | }; 49 | 50 | 51 | /// A node in a tensor index expression tree 52 | struct IndexStmtNode : public util::Manageable, 53 | private util::Uncopyable { 54 | public: 55 | IndexStmtNode(); 56 | IndexStmtNode(Type type); 57 | virtual ~IndexStmtNode() = default; 58 | virtual void accept(IndexStmtVisitorStrict*) const = 0; 59 | 60 | Type getType() const; 61 | 62 | private: 63 | Type type; 64 | }; 65 | 66 | } 67 | #endif 68 | -------------------------------------------------------------------------------- /include/taco/index_notation/index_notation_printer.h: -------------------------------------------------------------------------------- 1 | #ifndef TACO_INDEX_NOTATION_PRINTER_H 2 | #define TACO_INDEX_NOTATION_PRINTER_H 3 | 4 | #include 5 | #include "taco/index_notation/index_notation_visitor.h" 6 | 7 | namespace taco { 8 | 9 | class IndexNotationPrinter : public IndexNotationVisitorStrict { 10 | public: 11 | IndexNotationPrinter(std::ostream& os); 12 | 13 | void print(const IndexExpr& expr); 14 | void print(const IndexStmt& expr); 15 | 16 | using IndexNotationVisitorStrict::visit; 17 | 18 | // Scalar Expressions 19 | void visit(const AccessNode*); 20 | void visit(const LiteralNode*); 21 | void visit(const NegNode*); 22 | void visit(const SqrtNode*); 23 | void visit(const AddNode*); 24 | void visit(const SubNode*); 25 | void visit(const MulNode*); 26 | void visit(const DivNode*); 27 | void visit(const CastNode*); 28 | void visit(const CallNode*); 29 | void visit(const CallIntrinsicNode*); 30 | void visit(const ReductionNode*); 31 | void visit(const IndexVarNode*); 32 | 33 | // Tensor Expressions 34 | void visit(const AssignmentNode*); 35 | void visit(const YieldNode*); 36 | void visit(const ForallNode*); 37 | void visit(const WhereNode*); 38 | void visit(const MultiNode*); 39 | void visit(const SequenceNode*); 40 | void visit(const AssembleNode*); 41 | void visit(const SuchThatNode*); 42 | 43 | private: 44 | std::ostream& os; 45 | 46 | enum class Precedence { 47 | ACCESS = 2, 48 | FUNC = 2, 49 | CAST = 2, 50 | REDUCTION = 2, 51 | NEG = 3, 52 | MUL = 5, 53 | DIV = 5, 54 | ADD = 6, 55 | SUB = 6, 56 | TOP = 20 57 | }; 58 | Precedence parentPrecedence; 59 | 60 | template void visitBinary(Node op, Precedence p); 61 | template void visitImmediate(Node op); 62 | }; 63 | 64 | } 65 | #endif 66 | -------------------------------------------------------------------------------- /include/taco/index_notation/intrinsic.h: -------------------------------------------------------------------------------- 1 | #ifndef TACO_INTRINSIC_H 2 | #define TACO_INTRINSIC_H 3 | 4 | #include 5 | #include 6 | 7 | #include "taco/type.h" 8 | #include "taco/ir/ir.h" 9 | 10 | namespace taco { 11 | 12 | class IndexExpr; 13 | 14 | class Intrinsic { 15 | public: 16 | virtual ~Intrinsic() {} 17 | 18 | /// Returns the name of the instrinsic function. 19 | virtual std::string getName() const = 0; 20 | 21 | /// Infers the type of the return value based on the types of the arguments. 22 | virtual Datatype inferReturnType(const std::vector&) const = 0; 23 | 24 | /// Emits IR to compute the value of the intrinsic function. 25 | virtual ir::Expr lower(const std::vector&) const = 0; 26 | 27 | /// Returns a set ZP of zero-preserving argument sets ZP_i. Each ZP_i 28 | /// identifies a set of arguments to the intrinsic function that, if they are 29 | /// all zero, forces the result to also be zero. The zero-preserving 30 | /// argument sets must be disjoint (i.e., i != j --> ZP_i \intersect ZP_j == 0). 31 | virtual std::vector> 32 | zeroPreservingArgs(const std::vector&) const = 0; 33 | }; 34 | 35 | #define DECLARE_INTRINSIC(NAME) \ 36 | class NAME##Intrinsic : public Intrinsic { \ 37 | public: \ 38 | std::string getName() const; \ 39 | Datatype inferReturnType(const std::vector&) const; \ 40 | ir::Expr lower(const std::vector&) const; \ 41 | std::vector> zeroPreservingArgs(const std::vector&) const; \ 42 | }; 43 | 44 | DECLARE_INTRINSIC(Mod) 45 | DECLARE_INTRINSIC(Abs) 46 | DECLARE_INTRINSIC(Pow) 47 | DECLARE_INTRINSIC(Square) 48 | DECLARE_INTRINSIC(Cube) 49 | DECLARE_INTRINSIC(Sqrt) 50 | DECLARE_INTRINSIC(Cbrt) 51 | DECLARE_INTRINSIC(Exp) 52 | DECLARE_INTRINSIC(Log) 53 | DECLARE_INTRINSIC(Log10) 54 | DECLARE_INTRINSIC(Sin) 55 | DECLARE_INTRINSIC(Cos) 56 | DECLARE_INTRINSIC(Tan) 57 | DECLARE_INTRINSIC(Asin) 58 | DECLARE_INTRINSIC(Acos) 59 | DECLARE_INTRINSIC(Atan) 60 | DECLARE_INTRINSIC(Atan2) 61 | DECLARE_INTRINSIC(Sinh) 62 | DECLARE_INTRINSIC(Cosh) 63 | DECLARE_INTRINSIC(Tanh) 64 | DECLARE_INTRINSIC(Asinh) 65 | DECLARE_INTRINSIC(Acosh) 66 | DECLARE_INTRINSIC(Atanh) 67 | DECLARE_INTRINSIC(Gt) 68 | DECLARE_INTRINSIC(Lt) 69 | DECLARE_INTRINSIC(Gte) 70 | DECLARE_INTRINSIC(Lte) 71 | DECLARE_INTRINSIC(Eq) 72 | DECLARE_INTRINSIC(Neq) 73 | DECLARE_INTRINSIC(Max) 74 | DECLARE_INTRINSIC(Min) 75 | DECLARE_INTRINSIC(Heaviside) 76 | DECLARE_INTRINSIC(Not) 77 | 78 | } 79 | 80 | #endif 81 | 82 | -------------------------------------------------------------------------------- /include/taco/index_notation/iteration_algebra_printer.h: -------------------------------------------------------------------------------- 1 | #ifndef TACO_ITERATION_ALGEBRA_PRINTER_H 2 | #define TACO_ITERATION_ALGEBRA_PRINTER_H 3 | 4 | #include 5 | #include "taco/index_notation/iteration_algebra.h" 6 | 7 | namespace taco { 8 | 9 | // Iteration Algebra Printer 10 | class IterationAlgebraPrinter : IterationAlgebraVisitorStrict { 11 | public: 12 | IterationAlgebraPrinter(std::ostream& os); 13 | void print(const IterationAlgebra& alg); 14 | void visit(const RegionNode* n); 15 | void visit(const ComplementNode* n); 16 | void visit(const IntersectNode* n); 17 | void visit(const UnionNode* n); 18 | 19 | private: 20 | std::ostream& os; 21 | enum class Precedence { 22 | COMPLEMENT = 3, 23 | INTERSECT = 4, 24 | UNION = 5, 25 | TOP = 20 26 | }; 27 | 28 | Precedence parentPrecedence; 29 | 30 | template 31 | void visitBinary(Node n, Precedence precedence); 32 | }; 33 | } 34 | #endif //TACO_ITERATION_ALGEBRA_PRINTER_H 35 | -------------------------------------------------------------------------------- /include/taco/index_notation/kernel.h: -------------------------------------------------------------------------------- 1 | #ifndef TACO_KERNEL_H 2 | #define TACO_KERNEL_H 3 | 4 | #include 5 | #include 6 | 7 | namespace taco { 8 | 9 | class Function; 10 | class IndexStmt; 11 | class TensorStorage; 12 | namespace ir { 13 | class Module; 14 | } 15 | 16 | /// A tensor compute kernel is a runnable object that executes a concrete index 17 | /// notation statement. Kernels allocate memory, assemble indices, and compute 18 | /// component values of the result tensors in the concrete index statement. 19 | /// They can be called to do all these things at once (`evaluate`), to only 20 | /// allocate memory and assemble indices (`assemble`), or to only compute 21 | /// component values (`compute`). 22 | class Kernel { 23 | public: 24 | /// Construct an undefined kernel. 25 | Kernel(); 26 | 27 | /// Construct a kernel from relevant function pointers and a module. 28 | Kernel(IndexStmt stmt, std::shared_ptr module, 29 | void* evaluate, void* assemble, void* compute); 30 | 31 | /// Evaluate the kernel on the given tensor storage arguments, which includes 32 | /// allocating memory, assembling indices, and computing component values. 33 | /// @{ 34 | bool operator()(const std::vector& args) const; 35 | template bool operator()(const Args&... args) const { 36 | return operator()({args...}); 37 | } 38 | /// @} 39 | 40 | /// Execute the kernel to assemble the indices of the results. 41 | /// @{ 42 | bool assemble(const std::vector& args) const; 43 | template bool assemble(const Args&... args) const { 44 | return assemble({args...}); 45 | } 46 | /// @} 47 | 48 | /// Execute the kernel to compute the component values of the results, but 49 | /// do not allocate result memory or assemble result indices. 50 | /// @{ 51 | bool compute(const std::vector& args) const; 52 | template bool compute(const Args&... args) const { 53 | return compute({args...}); 54 | } 55 | /// @} 56 | 57 | /// Check whether the kernel is defined. 58 | bool defined(); 59 | 60 | /// Print the tensor compute kernel. 61 | friend std::ostream& operator<<(std::ostream&, const Kernel&); 62 | 63 | private: 64 | struct Content; 65 | std::shared_ptr content; 66 | size_t numResults; 67 | void* evaluateFunction; 68 | void* assembleFunction; 69 | void* computeFunction; 70 | }; 71 | 72 | /// Compile a concrete index notation statement to a runnable kernel. 73 | Kernel compile(IndexStmt stmt); 74 | 75 | } 76 | #endif 77 | -------------------------------------------------------------------------------- /include/taco/index_notation/properties.h: -------------------------------------------------------------------------------- 1 | #ifndef TACO_PROPERTIES_H 2 | #define TACO_PROPERTIES_H 3 | 4 | #include "taco/index_notation/property_pointers.h" 5 | #include "taco/util/intrusive_ptr.h" 6 | 7 | namespace taco { 8 | 9 | class IndexExpr; 10 | 11 | /// A class containing properties about an operation 12 | class Property : public util::IntrusivePtr { 13 | public: 14 | Property(); 15 | explicit Property(const PropertyPtr* p); 16 | 17 | bool equals(const Property& p) const; 18 | std::ostream& print(std::ostream&) const; 19 | }; 20 | 21 | std::ostream& operator<<(std::ostream&, const Property&); 22 | 23 | /// A class wrapping the annihilator property pointer 24 | class Annihilator : public Property { 25 | public: 26 | explicit Annihilator(Literal); 27 | Annihilator(Literal, std::vector&); 28 | explicit Annihilator(const PropertyPtr*); 29 | 30 | const Literal& annihilator() const; 31 | const std::vector& positions() const; 32 | IndexExpr annihilates(const std::vector&) const; 33 | 34 | typedef AnnihilatorPtr Ptr; 35 | }; 36 | 37 | /// A class wrapping an identity property pointer 38 | class Identity : public Property { 39 | public: 40 | explicit Identity(Literal); 41 | Identity(Literal, std::vector&); 42 | explicit Identity(const PropertyPtr*); 43 | 44 | const Literal& identity() const; 45 | const std::vector& positions() const; 46 | IndexExpr simplify(const std::vector&) const; 47 | 48 | typedef IdentityPtr Ptr; 49 | }; 50 | 51 | /// A class wrapping an associative property pointer 52 | class Associative : public Property { 53 | public: 54 | Associative(); 55 | explicit Associative(const PropertyPtr*); 56 | 57 | typedef AssociativePtr Ptr; 58 | }; 59 | 60 | /// A class wrapping a commutative property pointer 61 | class Commutative : public Property { 62 | public: 63 | Commutative(); 64 | explicit Commutative(const std::vector&); 65 | explicit Commutative(const PropertyPtr*); 66 | 67 | const std::vector& ordering() const; 68 | 69 | typedef CommutativePtr Ptr; 70 | }; 71 | 72 | /// Returns true if property p is of type P. 73 | template bool isa(const Property& p); 74 | 75 | /// Casts the Property p to type P. 76 | template P to(const Property& p); 77 | 78 | /// Finds and returns the property of type P if it exists in the vector. If 79 | /// the property does not exist, returns an undefined instance of the property 80 | /// requested. 81 | /// The vector of properties should not contain duplicates so this is sufficient. 82 | template 83 | inline const P findProperty(const std::vector &properties) { 84 | for (const auto &p: properties) { 85 | if (isa

(p)) return to

(p); 86 | } 87 | return P(nullptr); 88 | } 89 | 90 | } 91 | 92 | #endif //TACO_PROPERTIES_H 93 | -------------------------------------------------------------------------------- /include/taco/index_notation/property_pointers.h: -------------------------------------------------------------------------------- 1 | #ifndef TACO_PROPERTY_POINTERS_H 2 | #define TACO_PROPERTY_POINTERS_H 3 | 4 | #include 5 | #include 6 | #include 7 | #include 8 | #include 9 | 10 | #include "taco/error.h" 11 | #include "taco/util/comparable.h" 12 | 13 | namespace taco { 14 | 15 | class Literal; 16 | struct PropertyPtr; 17 | 18 | /// A pointer to the property data. This will be wrapped in an auxillary class 19 | /// to allow a user to create a vector of properties. Needed since properties 20 | /// have different methods and data 21 | struct PropertyPtr : public util::Manageable, 22 | private util::Uncopyable { 23 | public: 24 | PropertyPtr(); 25 | virtual ~PropertyPtr(); 26 | virtual std::ostream& print(std::ostream& os) const; 27 | virtual bool equals(const PropertyPtr* p) const; 28 | }; 29 | 30 | /// Pointer class for annihilators 31 | struct AnnihilatorPtr : public PropertyPtr { 32 | AnnihilatorPtr(); 33 | AnnihilatorPtr(Literal); 34 | AnnihilatorPtr(Literal, std::vector&); 35 | 36 | const Literal& annihilator() const; 37 | const std::vector& positions() const; 38 | 39 | virtual std::ostream& print(std::ostream& os) const; 40 | virtual bool equals(const PropertyPtr* p) const; 41 | 42 | struct Content; 43 | std::shared_ptr content; 44 | }; 45 | 46 | /// Pointer class for identities 47 | struct IdentityPtr : public PropertyPtr { 48 | public: 49 | IdentityPtr(); 50 | IdentityPtr(Literal); 51 | IdentityPtr(Literal, std::vector&); 52 | 53 | const Literal& identity() const; 54 | const std::vector& positions() const; 55 | 56 | virtual std::ostream& print(std::ostream& os) const; 57 | virtual bool equals(const PropertyPtr* p) const; 58 | 59 | struct Content; 60 | std::shared_ptr content; 61 | }; 62 | 63 | /// Pointer class for associativity 64 | struct AssociativePtr : public PropertyPtr { 65 | AssociativePtr(); 66 | virtual std::ostream& print(std::ostream& os) const; 67 | virtual bool equals(const PropertyPtr* p) const; 68 | }; 69 | 70 | /// Pointer class for commutativity 71 | struct CommutativePtr : public PropertyPtr { 72 | CommutativePtr(); 73 | CommutativePtr(const std::vector&); 74 | const std::vector ordering_; 75 | virtual std::ostream& print(std::ostream& os) const; 76 | virtual bool equals(const PropertyPtr* p) const; 77 | }; 78 | 79 | template 80 | inline bool isa(const PropertyPtr* p) { 81 | return p != nullptr && dynamic_cast(p) != nullptr; 82 | } 83 | 84 | template 85 | inline const P* to(const PropertyPtr* p) { 86 | taco_iassert(isa

(p)) << 87 | "Cannot convert " << typeid(p).name() << " to " << typeid(P).name();; 88 | return static_cast(p); 89 | } 90 | 91 | template 92 | inline const typename P::Ptr* getPtr(const P& propertyPtr) { 93 | taco_iassert(isa(propertyPtr.ptr)); 94 | return static_cast(propertyPtr.ptr); 95 | } 96 | 97 | 98 | } 99 | 100 | #endif //TACO_PROPERTY_POINTERS_H 101 | -------------------------------------------------------------------------------- /include/taco/index_notation/schedule.h: -------------------------------------------------------------------------------- 1 | #ifndef TACO_SCHEDULE_H 2 | #define TACO_SCHEDULE_H 3 | 4 | #include 5 | #include 6 | #include 7 | #include 8 | 9 | namespace taco { 10 | 11 | class IndexExpr; 12 | 13 | class Reorder; 14 | class Precompute; 15 | 16 | /// A schedule controls code generation and determines how index expression 17 | /// should be computed. 18 | class Schedule { 19 | public: 20 | Schedule(); 21 | 22 | /// Returns the workspace commands in the schedule. 23 | std::vector getPrecomputes() const; 24 | 25 | /// Returns the workspace of `expr`. The result is undefined if `expr` is not 26 | /// stored to a workspace. 27 | Precompute getPrecompute(IndexExpr expr) const; 28 | 29 | /// Add a workspace command to the schedule. 30 | void addPrecompute(Precompute precompute); 31 | 32 | /// Removes workspace commands from the schedule. 33 | void clearPrecomputes(); 34 | 35 | private: 36 | struct Content; 37 | std::shared_ptr content; 38 | }; 39 | 40 | /// Print a schedule. 41 | std::ostream& operator<<(std::ostream&, const Schedule&); 42 | 43 | } 44 | #endif 45 | -------------------------------------------------------------------------------- /include/taco/ir/ir_generators.h: -------------------------------------------------------------------------------- 1 | #ifndef TACO_IR_CODEGEN_H 2 | #define TACO_IR_CODEGEN_H 3 | 4 | #include 5 | #include "taco/ir_tags.h" 6 | 7 | namespace taco { 8 | 9 | namespace ir { 10 | class Expr; 11 | class Stmt; 12 | 13 | /// Generate `a[i] += val;` 14 | Stmt compoundStore(Expr a, Expr i, Expr val, bool use_atomics=false, ParallelUnit atomic_parallel_unit=ParallelUnit::NotParallel); 15 | 16 | /// Generate `a += val;` 17 | Stmt compoundAssign(Expr a, Expr val, bool use_atomics=false, ParallelUnit atomic_parallel_unit=ParallelUnit::NotParallel); 18 | 19 | /// Generate `exprs_0 && ... && exprs_n` 20 | Expr conjunction(std::vector exprs); 21 | 22 | /// Generate a statement that doubles the size of `a` if it is full (loc cannot 23 | /// be written to). 24 | Stmt doubleSizeIfFull(Expr a, Expr size, Expr loc); 25 | 26 | /// Generate a statement that resizes `a` to be double its original size or at 27 | /// least equal to `loc` if it is full (loc cannot be written to). 28 | Stmt atLeastDoubleSizeIfFull(Expr a, Expr size, Expr loc); 29 | 30 | }} 31 | #endif 32 | -------------------------------------------------------------------------------- /include/taco/ir/ir_rewriter.h: -------------------------------------------------------------------------------- 1 | #ifndef TACO_IR_REWRITER_H 2 | #define TACO_IR_REWRITER_H 3 | 4 | #include "taco/ir/ir_visitor.h" 5 | #include "taco/ir/ir.h" 6 | 7 | namespace taco { 8 | namespace ir { 9 | 10 | /// Extend this class to rewrite IR nodes. 11 | class IRRewriter : public IRVisitorStrict { 12 | public: 13 | virtual ~IRRewriter(); 14 | 15 | Expr rewrite(Expr); 16 | Stmt rewrite(Stmt); 17 | 18 | protected: 19 | /// visit methods that take Exprs assign to this to return their value. 20 | Expr expr; 21 | 22 | /// visit methods that take Stmts assign to this to return their value. 23 | Stmt stmt; 24 | 25 | using IRVisitorStrict::visit; 26 | virtual void visit(const Literal* op); 27 | virtual void visit(const Var* op); 28 | virtual void visit(const Neg* op); 29 | virtual void visit(const Sqrt* op); 30 | virtual void visit(const Add* op); 31 | virtual void visit(const Sub* op); 32 | virtual void visit(const Mul* op); 33 | virtual void visit(const Div* op); 34 | virtual void visit(const Rem* op); 35 | virtual void visit(const Min* op); 36 | virtual void visit(const Max* op); 37 | virtual void visit(const BitAnd* op); 38 | virtual void visit(const BitOr* op); 39 | virtual void visit(const Eq* op); 40 | virtual void visit(const Neq* op); 41 | virtual void visit(const Gt* op); 42 | virtual void visit(const Lt* op); 43 | virtual void visit(const Gte* op); 44 | virtual void visit(const Lte* op); 45 | virtual void visit(const And* op); 46 | virtual void visit(const Or* op); 47 | virtual void visit(const BinOp* op); 48 | virtual void visit(const Cast* op); 49 | virtual void visit(const Call* op); 50 | virtual void visit(const IfThenElse* op); 51 | virtual void visit(const Case* op); 52 | virtual void visit(const Switch* op); 53 | virtual void visit(const Load* op); 54 | virtual void visit(const Malloc* op); 55 | virtual void visit(const Sizeof* op); 56 | virtual void visit(const Store* op); 57 | virtual void visit(const For* op); 58 | virtual void visit(const While* op); 59 | virtual void visit(const Block* op); 60 | virtual void visit(const Scope* op); 61 | virtual void visit(const Function* op); 62 | virtual void visit(const VarDecl* op); 63 | virtual void visit(const Assign* op); 64 | virtual void visit(const Yield* op); 65 | virtual void visit(const Allocate* op); 66 | virtual void visit(const Free* op); 67 | virtual void visit(const Comment* op); 68 | virtual void visit(const BlankLine* op); 69 | virtual void visit(const Continue* op); 70 | virtual void visit(const Print* op); 71 | virtual void visit(const GetProperty* op); 72 | virtual void visit(const Sort *op); 73 | virtual void visit(const Break *op); 74 | }; 75 | 76 | }} 77 | #endif 78 | -------------------------------------------------------------------------------- /include/taco/ir/ir_verifier.h: -------------------------------------------------------------------------------- 1 | #ifndef TACO_IR_VERIFIER_H 2 | #define TACO_IR_VERIFIER_H 3 | 4 | #include "taco/ir/ir.h" 5 | #include "taco/ir/ir_visitor.h" 6 | 7 | namespace taco { 8 | namespace ir { 9 | /** Verify an Stmt is well-formed for codegen */ 10 | bool verify(const Stmt s, std::string *message); 11 | 12 | /** Verify an Stmt is well-formed for codegen */ 13 | bool verify(const Expr e, std::string *message); 14 | 15 | } // namespace ir 16 | } // namespace taco 17 | 18 | #endif // TACO_IR_VERIFIER_H 19 | -------------------------------------------------------------------------------- /include/taco/ir/simplify.h: -------------------------------------------------------------------------------- 1 | #ifndef TACO_IR_SIMPLIFY_H 2 | #define TACO_IR_SIMPLIFY_H 3 | 4 | namespace taco { 5 | namespace ir { 6 | class Expr; 7 | class Stmt; 8 | 9 | /// Simplifies an expression (e.g. by applying algebraic identities). 10 | ir::Expr simplify(const ir::Expr& expr); 11 | 12 | /// Simplifies a statement (e.g. by applying constant copy propagation). 13 | ir::Stmt simplify(const ir::Stmt& stmt); 14 | 15 | }} 16 | #endif 17 | -------------------------------------------------------------------------------- /include/taco/ir/workspace_rewriter.h: -------------------------------------------------------------------------------- 1 | #ifndef TACO_WORKSPACE_REWRITER_H 2 | #define TACO_WORKSPACE_REWRITER_H 3 | 4 | #include 5 | #include 6 | 7 | 8 | namespace taco { 9 | class TensorVar; 10 | 11 | namespace ir { 12 | class Stmt; 13 | class Expr; 14 | } 15 | 16 | /// Rewrite a post-lowered IR statement to take into account multidimensional temporaries. 17 | /// Replaces Dimension GetProperty nodes that correspond to temporary workspaces with 18 | /// their corresponding dimension found in the temporarySizeMap. 19 | ir::Stmt rewriteTemporaryGP(const ir::Stmt& stmt, std::vector whereTemps, 20 | std::map> temporarySizeMap); 21 | 22 | } 23 | #endif 24 | -------------------------------------------------------------------------------- /include/taco/ir_tags.h: -------------------------------------------------------------------------------- 1 | #ifndef TACO_IR_TAGS_H 2 | #define TACO_IR_TAGS_H 3 | 4 | namespace taco { 5 | 6 | /// ParallelUnit::CPUThread generates a pragma to parallelize over CPU threads 7 | /// ParallelUnit::CPUVector generates a pragma to utilize a CPU vector unit 8 | /// ParallelUnit::GPUBlock must be used with GPUThread to create blocks of GPU threads 9 | /// ParallelUnit::GPUWarp can be optionally used to allow for GPU warp-level primitives 10 | /// ParallelUnit::GPUThread causes for every iteration to be executed on a separate GPU thread 11 | enum class ParallelUnit { 12 | NotParallel, DefaultUnit, GPUBlock, GPUWarp, GPUThread, CPUThread, CPUVector, CPUThreadGroupReduction, GPUBlockReduction, GPUWarpReduction 13 | }; 14 | extern const char *ParallelUnit_NAMES[]; 15 | 16 | /// OutputRaceStrategy::NoRaces raises a compile-time error if an output race exists 17 | /// OutputRaceStrategy::Atomics replace racing instructions with atomics 18 | /// OutputRaceStrategy::Temporary uses a temporary array for outputs that is serially reduced 19 | /// OutputRaceStrategy::ParallelReduction uses reduction operations across a warp/vector 20 | /// OutputRaceStrategy::IgnoreRaces allows the user to specify that races can be safely ignored 21 | enum class OutputRaceStrategy { 22 | IgnoreRaces, NoRaces, Atomics, Temporary, ParallelReduction 23 | }; 24 | extern const char *OutputRaceStrategy_NAMES[]; 25 | 26 | enum class BoundType { 27 | MinExact, MinConstraint, MaxExact, MaxConstraint 28 | }; 29 | extern const char *BoundType_NAMES[]; 30 | 31 | enum class AssembleStrategy { 32 | Append, Insert 33 | }; 34 | extern const char *AssembleStrategy_NAMES[]; 35 | 36 | /// MergeStrategy::TwoFinger merges iterators by incrementing one at a time 37 | /// MergeStrategy::Galloping merges iterators by exponential search (galloping) 38 | enum class MergeStrategy { 39 | TwoFinger, Gallop 40 | }; 41 | extern const char *MergeStrategy_NAMES[]; 42 | 43 | } 44 | 45 | #endif //TACO_IR_TAGS_H 46 | -------------------------------------------------------------------------------- /include/taco/lower/lower.h: -------------------------------------------------------------------------------- 1 | #ifndef TACO_LOWER_H 2 | #define TACO_LOWER_H 3 | 4 | #include 5 | #include 6 | #include 7 | 8 | namespace taco { 9 | 10 | class IndexStmt; 11 | class LowererImpl; 12 | 13 | namespace ir { 14 | class Stmt; 15 | } 16 | 17 | 18 | /// A `Lowerer` lowers concrete index notation statements as specified by a 19 | /// `LowererImpl`. The default `Lowerer`/`LowererImpl` lowers to sequential, 20 | /// multithreaded, GPU, and vectorized code as specified by the concrete index 21 | /// notation. `LowererImpl`, however, can be extended and it's methods 22 | /// overridden to insert custom lowering code to e.g. target specialized 23 | /// hardware. See `lowerer_impl.h` for information about how to create custom 24 | /// lowerers. 25 | class Lowerer { 26 | public: 27 | /// Construct a default lowerer that lowers to imperative multi-threaded code. 28 | Lowerer(); 29 | 30 | /// Construct a lowerer that lowers as specified by the lowerer impl. The 31 | /// lowerer will delete the impl object. 32 | Lowerer(LowererImpl* impl); 33 | 34 | /// Retrieve the lowerer implementation. 35 | std::shared_ptr getLowererImpl(); 36 | 37 | private: 38 | std::shared_ptr impl; 39 | }; 40 | 41 | 42 | /// Lower a concrete index notation statement to a function in the low-level 43 | /// imperative IR. You may specify whether the lowered function should 44 | /// assemble result indices, compute result values, or both. You may optionally 45 | /// also provide a custom `Lowerer` to specify custom ways to lower some or all 46 | /// parts of a concrete index notation statement. 47 | ir::Stmt lower(IndexStmt stmt, std::string functionName, 48 | bool assemble=true, bool compute=true, bool pack=false, bool unpack=false, 49 | Lowerer lowerer=Lowerer()); 50 | 51 | /// Check whether the an index statement can be lowered to C code. If the 52 | /// statement cannot be lowered and a `reason` string is provided then it is 53 | /// filled with the a reason. 54 | bool isLowerable(IndexStmt stmt, std::string* reason=nullptr); 55 | 56 | } 57 | #endif 58 | -------------------------------------------------------------------------------- /include/taco/lower/mode.h: -------------------------------------------------------------------------------- 1 | #ifndef TACO_MODE_H 2 | #define TACO_MODE_H 3 | 4 | #include 5 | 6 | #include "taco/format.h" 7 | 8 | namespace taco { 9 | 10 | class ModePack; 11 | class Dimension; 12 | 13 | namespace ir { 14 | class Expr; 15 | } 16 | 17 | /// A mode of a tensor. 18 | class Mode { 19 | public: 20 | /// Construct an undefined mode. 21 | Mode(); 22 | 23 | /// Construct a tensor mode. 24 | Mode(ir::Expr tensor, Dimension size, int mode, ModeFormat modeFormat, 25 | ModePack modePack, size_t packLoc, ModeFormat parentModeFormat); 26 | 27 | /// Retrieve the name of the tensor mode. 28 | std::string getName() const; 29 | 30 | /// Retrieve the tensor that contains the mode. 31 | ir::Expr getTensorExpr() const; 32 | 33 | /// Retrieve the size of the tensor mode. 34 | Dimension getSize() const; 35 | 36 | /// Retrieve the mode of this mode in its the mode hierarchy. The first 37 | /// mode in a mode hierarchy is at mode 1, and mode 0 is the root mode. 38 | int getLevel() const; 39 | 40 | /// Retrieve the format of the mode. 41 | ModeFormat getModeFormat() const; 42 | 43 | /// Retrieve the mode pack that stores the mode. 44 | ModePack getModePack() const; 45 | 46 | /// Retrieve the location of the mode in its mode pack. 47 | size_t getPackLocation() const; 48 | 49 | /// Retrieve the mode type of the parent mode in the mode hierarchy. 50 | ModeFormat getParentModeType() const; 51 | 52 | /// Store temporary variables that may be needed to access or modify a mode 53 | /// @{ 54 | ir::Expr getVar(std::string varName) const; 55 | bool hasVar(std::string varName) const; 56 | void addVar(std::string varName, ir::Expr var); 57 | /// @} 58 | 59 | /// Check whether the mode is defined. 60 | bool defined() const; 61 | 62 | private: 63 | struct Content; 64 | std::shared_ptr content; 65 | friend class ModePack; 66 | }; 67 | 68 | std::ostream& operator<<(std::ostream&, const Mode&); 69 | 70 | 71 | /// A ModePack is a set of physical arrays, that can be used by one mode or 72 | /// shared by multiple modes (e.g., modes of an array-of-structs COO tensor). 73 | class ModePack { 74 | public: 75 | ModePack(); 76 | ModePack(size_t numModes, ModeFormat modeType, ir::Expr tensor, int mode, 77 | int level); 78 | 79 | /// Returns number of tensor modes belonging to mode pack. 80 | size_t getNumModes() const; 81 | 82 | /// Returns arrays shared by tensor modes. 83 | ir::Expr getArray(size_t i) const; 84 | 85 | private: 86 | struct Content; 87 | std::shared_ptr content; 88 | }; 89 | 90 | } 91 | #endif 92 | -------------------------------------------------------------------------------- /include/taco/lower/mode_format_dense.h: -------------------------------------------------------------------------------- 1 | #ifndef TACO_MODE_FORMAT_DENSE_H 2 | #define TACO_MODE_FORMAT_DENSE_H 3 | 4 | #include "taco/lower/mode_format_impl.h" 5 | 6 | namespace taco { 7 | 8 | class DenseModeFormat : public ModeFormatImpl { 9 | public: 10 | using ModeFormatImpl::getInsertCoord; 11 | 12 | DenseModeFormat(); 13 | DenseModeFormat(const bool isOrdered, const bool isUnique, const bool isZeroless); 14 | 15 | ~DenseModeFormat() override {} 16 | 17 | ModeFormat copy(std::vector properties) const override; 18 | 19 | ModeFunction locate(ir::Expr parentPos, std::vector coords, 20 | Mode mode) const override; 21 | 22 | ir::Stmt getInsertCoord(ir::Expr p, const std::vector& i, 23 | Mode mode) const override; 24 | ir::Expr getWidth(Mode mode) const override; 25 | ir::Stmt getInsertInitCoords(ir::Expr pBegin, ir::Expr pEnd, 26 | Mode mode) const override; 27 | ir::Stmt getInsertInitLevel(ir::Expr szPrev, ir::Expr sz, 28 | Mode mode) const override; 29 | ir::Stmt getInsertFinalizeLevel(ir::Expr szPrev, ir::Expr sz, 30 | Mode mode) const override; 31 | 32 | ir::Expr getAssembledSize(ir::Expr prevSize, Mode mode) const override; 33 | ModeFunction getYieldPos(ir::Expr parentPos, std::vector coords, 34 | Mode mode) const override; 35 | 36 | std::vector getArrays(ir::Expr tensor, int mode, 37 | int level) const override; 38 | 39 | protected: 40 | ir::Expr getSizeArray(ModePack pack) const; 41 | }; 42 | 43 | } 44 | 45 | #endif 46 | -------------------------------------------------------------------------------- /include/taco/lower/mode_format_singleton.h: -------------------------------------------------------------------------------- 1 | #ifndef TACO_MODE_FORMAT_SINGLETON_H 2 | #define TACO_MODE_FORMAT_SINGLETON_H 3 | 4 | #include "taco/lower/mode_format_impl.h" 5 | 6 | namespace taco { 7 | 8 | class SingletonModeFormat : public ModeFormatImpl { 9 | public: 10 | using ModeFormatImpl::getInsertCoord; 11 | 12 | SingletonModeFormat(); 13 | SingletonModeFormat(bool isFull, bool isOrdered, bool isUnique, 14 | bool isZeroless, bool isPadded, 15 | long long allocSize = DEFAULT_ALLOC_SIZE); 16 | 17 | ~SingletonModeFormat() override {} 18 | 19 | ModeFormat copy(std::vector properties) const override; 20 | 21 | ModeFunction posIterBounds(ir::Expr parentPos, Mode mode) const override; 22 | ModeFunction posIterAccess(ir::Expr pos, std::vector coords, 23 | Mode mode) const override; 24 | 25 | ir::Stmt getAppendCoord(ir::Expr pos, ir::Expr coord, 26 | Mode mode) const override; 27 | ir::Expr getSize(ir::Expr parentSize, Mode mode) const override; 28 | ir::Stmt getAppendInitLevel(ir::Expr parentSize, ir::Expr size, 29 | Mode mode) const override; 30 | ir::Stmt getAppendFinalizeLevel(ir::Expr parentSize, ir::Expr size, 31 | Mode mode) const override; 32 | 33 | ir::Expr getAssembledSize(ir::Expr prevSize, Mode mode) const override; 34 | ir::Stmt getInitCoords(ir::Expr prevSize, 35 | std::vector queries, 36 | Mode mode) const override; 37 | ModeFunction getYieldPos(ir::Expr parentPos, std::vector coords, 38 | Mode mode) const override; 39 | ir::Stmt getInsertCoord(ir::Expr parentPos, ir::Expr pos, 40 | std::vector coords, 41 | Mode mode) const override; 42 | 43 | std::vector getArrays(ir::Expr tensor, int mode, 44 | int level) const override; 45 | 46 | protected: 47 | ir::Expr getCoordArray(ModePack pack) const; 48 | 49 | ir::Expr getCoordCapacity(Mode mode) const; 50 | 51 | bool equals(const ModeFormatImpl& other) const override; 52 | 53 | const long long allocSize; 54 | }; 55 | 56 | } 57 | 58 | #endif 59 | -------------------------------------------------------------------------------- /include/taco/parser/einsum_parser.h: -------------------------------------------------------------------------------- 1 | #ifndef TACO_EINSUM_PARSER_H 2 | #define TACO_EINSUM_PARSER_H 3 | 4 | #include 5 | #include 6 | #include 7 | #include 8 | #include 9 | 10 | #include "taco/tensor.h" 11 | #include "taco/util/uncopyable.h" 12 | #include "taco/type.h" 13 | 14 | namespace taco { 15 | class TensorBase; 16 | namespace parser { 17 | 18 | class EinsumParser : public util::Uncopyable { 19 | 20 | public: 21 | 22 | /// Create a parser object from einsum notation 23 | /// @throws ParserError is there is an error with parsing the einsum string 24 | EinsumParser(const std::string &expression, std::vector &tensors, 25 | Format &format, Datatype outType); 26 | 27 | /// Returns true if the expression passed in has an output specified and false otherwise 28 | /// @throws ParserError if output is not specified correctly 29 | bool exprHasOutput(const std::string& subscripts); 30 | 31 | /// Parses the einsum expression and sets the result tensor to the result of that expression 32 | /// @throws ParserError is there is an error with parsing the einsum stirng 33 | void parse(); 34 | 35 | /// Gets the result tensor after parsing is complete. 36 | TensorBase& getResultTensor(); 37 | 38 | private: 39 | std::string einsumSymbols; 40 | std::set einSumSymbolsSet; 41 | std::string einsumPunctuation; 42 | Datatype outType; 43 | Format format; 44 | 45 | 46 | std::string subscripts; 47 | TensorBase resultTensor; 48 | std::vector &tensors; 49 | 50 | /// Replaces the ellipses in an expression 51 | std::string replaceEllipse(std::string inp, std::string &newString); 52 | 53 | /// Returns a sorted string of unique elements for a given einsum expression in implicit notation 54 | std::string findUniqueIndices(const std::string &subscripts); 55 | 56 | /// Builds the result tensor given the explicit input and output substrings 57 | void buildResult(std::vector inputAndOutput); 58 | 59 | /// splits a string by , but keeps the empty string 60 | std::vector splitSubscriptInput(std::string &inp); 61 | 62 | }; 63 | 64 | 65 | }} 66 | 67 | #endif //TACO_EINSUM_PARSER_H 68 | -------------------------------------------------------------------------------- /include/taco/parser/lexer.h: -------------------------------------------------------------------------------- 1 | #ifndef TACO_LEXER_H 2 | #define TACO_LEXER_H 3 | 4 | #include 5 | 6 | namespace taco { 7 | namespace parser { 8 | 9 | enum class Token { 10 | identifier, 11 | int_scalar, 12 | uint_scalar, 13 | float_scalar, 14 | complex_scalar, 15 | comma, 16 | lparen, 17 | rparen, 18 | underscore, 19 | lcurly, 20 | rcurly, 21 | add, 22 | sub, 23 | mul, 24 | div, 25 | eq, 26 | eot, // End of tokens 27 | error 28 | }; 29 | 30 | 31 | // A simple index expression lexer. 32 | class Lexer { 33 | public: 34 | Lexer() {} 35 | Lexer(std::string expr) : expr(expr) {} 36 | 37 | /// Retrieve the next token. 38 | Token getToken(); 39 | 40 | std::string getIdentifier() const; 41 | std::string getLastChar() const; 42 | 43 | /// Convert a token to a string. 44 | std::string tokenString(const Token& token); 45 | 46 | private: 47 | std::string expr; 48 | int lastChar = ' '; 49 | int lastCharPos = -1; 50 | std::string identifier; 51 | 52 | int getNextChar(); 53 | }; 54 | 55 | }} 56 | #endif 57 | -------------------------------------------------------------------------------- /include/taco/parser/schedule_parser.h: -------------------------------------------------------------------------------- 1 | #ifndef TACO_SCHEDULE_PARSER_H 2 | #define TACO_SCHEDULE_PARSER_H 3 | 4 | #include 5 | #include 6 | 7 | namespace taco { 8 | namespace parser { 9 | 10 | // parse a string of the form: "reorder(i,j),precompute(D(i,j)*E(j,k),j,j_pre)" 11 | // into string vectors of the form: 12 | // [ [ "reorder", "i", "j" ], [ "precompute", "D(i,j)*E(j,k)", "j", "j_pre" ] ] 13 | std::vector> ScheduleParser(const std::string); 14 | 15 | std::vector varListParser(const std::string); 16 | 17 | // serialize the result of a parse (for debugging) 18 | std::string serializeParsedSchedule(std::vector>); 19 | 20 | }} 21 | 22 | #endif //TACO_EINSUM_PARSER_H 23 | -------------------------------------------------------------------------------- /include/taco/storage/array.h: -------------------------------------------------------------------------------- 1 | #ifndef TACO_STORAGE_ARRAY_H 2 | #define TACO_STORAGE_ARRAY_H 3 | 4 | #include 5 | #include 6 | #include 7 | #include 8 | #include "taco/util/collections.h" 9 | 10 | namespace taco { 11 | 12 | /// An array is a smart pointer to raw memory together with an element type, 13 | /// a size (number of elements) and a reclamation policy. 14 | class Array { 15 | public: 16 | /// The memory reclamation policy of Array objects. UserOwns means the Array 17 | /// object will not free its data, free means it will reclaim data with the 18 | /// C free function and delete means it will reclaim data with delete[]. 19 | enum Policy {UserOwns, Free, Delete}; 20 | 21 | /// Construct an empty array of undefined elements. 22 | Array(); 23 | 24 | /// Construct an array of elements of the given type. 25 | Array(Datatype type, void* data, size_t size, Policy policy=Free); 26 | 27 | /// Returns the type of the array elements 28 | const Datatype& getType() const; 29 | 30 | /// Returns the number of array elements 31 | size_t getSize() const; 32 | 33 | /// Returns the array data. 34 | /// @{ 35 | const void* getData() const; 36 | void* getData(); 37 | /// @} 38 | 39 | /// Gets the value at a given index 40 | TypedComponentRef get(size_t index) const; 41 | /// Gets the value at a given index 42 | TypedComponentRef operator[] (const int index) const; 43 | 44 | /// Zero the array content 45 | void zero(); 46 | 47 | private: 48 | struct Content; 49 | std::shared_ptr content; 50 | }; 51 | 52 | /// Print the array. 53 | std::ostream& operator<<(std::ostream&, const Array&); 54 | 55 | /// Print the array policy. 56 | std::ostream& operator<<(std::ostream&, Array::Policy); 57 | 58 | /// Construct an index array. The ownership policy determines whether the 59 | /// mode index will free/delete the memory or leave the responsibility for 60 | /// freeing to the user. 61 | template 62 | Array makeArray(T* data, size_t size, Array::Policy policy=Array::UserOwns) { 63 | return Array(type(), data, size, policy); 64 | } 65 | 66 | /// Construct an array of elements of the given type. 67 | Array makeArray(Datatype type, size_t size); 68 | 69 | /// Construct an Array from the values. 70 | template 71 | Array makeArray(const std::vector& values) { 72 | Array array = makeArray(type(), values.size()); 73 | memcpy(array.getData(), values.data(), values.size() * sizeof(T)); 74 | return array; 75 | } 76 | 77 | /// Construct an Array from the values. 78 | template 79 | Array makeArray(const std::initializer_list& values) { 80 | return makeArray(std::vector(values)); 81 | } 82 | 83 | } 84 | #endif 85 | 86 | -------------------------------------------------------------------------------- /include/taco/storage/coordinate.h: -------------------------------------------------------------------------------- 1 | #ifndef TACO_STORAGE_COORDINATE_H 2 | #define TACO_STORAGE_COORDINATE_H 3 | 4 | #include 5 | #include 6 | #include 7 | #include 8 | #include "taco/util/comparable.h" 9 | #include "taco/util/strings.h" 10 | #include "taco/error.h" 11 | 12 | namespace taco { 13 | 14 | template 15 | class Coordinates : util::Comparable> { 16 | public: 17 | template 18 | Coordinates(T... coordinates) : coordinates{{coordinates...}} {} 19 | 20 | Type& operator[](size_t idx) { 21 | taco_iassert(idx < Order); 22 | return coordinates[idx]; 23 | } 24 | 25 | const Type& operator[](size_t idx) const { 26 | taco_iassert(idx < Order); 27 | return coordinates[idx]; 28 | } 29 | 30 | template 31 | friend bool operator==(const Coordinates& a, const Coordinates& b) { 32 | for (size_t i = 0; i < Order; i++) { 33 | if (a[i] != b[i]) return false; 34 | } 35 | return true; 36 | } 37 | 38 | template 39 | friend bool operator<(const Coordinates& a, const Coordinates& b) { 40 | for (size_t i = 0; i < Order; i++) { 41 | if (a[i] < b[i]) return true; 42 | if (a[i] > b[i]) return false; 43 | } 44 | return false; 45 | } 46 | 47 | template 48 | friend std::ostream& operator<<(std::ostream& os, const Coordinates& c) { 49 | return os << util::join(c.coordinates); 50 | } 51 | 52 | template 53 | friend std::ostream& operator<<(std::ostream& os, 54 | const std::pair, V>& c) { 55 | return os << "(" << util::join(c.first.coordinates, ",") << "):" << c.second; 56 | } 57 | 58 | private: 59 | std::array coordinates; 60 | }; 61 | 62 | } 63 | #endif 64 | -------------------------------------------------------------------------------- /include/taco/storage/file_io_mtx.h: -------------------------------------------------------------------------------- 1 | #ifndef TACO_FILE_IO_MTX_H 2 | #define TACO_FILE_IO_MTX_H 3 | 4 | #include 5 | #include 6 | #include 7 | 8 | #include "taco/format.h" 9 | 10 | namespace taco { 11 | class TensorBase; 12 | class Format; 13 | 14 | /// Read an mtx matrix from a file. 15 | TensorBase readMTX(std::string filename, const ModeFormat& modetype, 16 | bool pack=true); 17 | 18 | /// Read an mtx matrix from a file. 19 | TensorBase readMTX(std::string filename, const Format& format, bool pack=true); 20 | 21 | /// Read an mtx matrix from a stream. 22 | TensorBase readMTX(std::istream& stream, const ModeFormat& modetype, 23 | bool pack=true); 24 | 25 | /// Read an mtx matrix from a stream. 26 | TensorBase readMTX(std::istream& stream, const Format& format, bool pack=true); 27 | 28 | TensorBase readSparse(std::istream& stream, const ModeFormat& modetype, 29 | bool symm = false); 30 | TensorBase readDense(std::istream& stream, const ModeFormat& modetype, 31 | bool symm = false); 32 | 33 | TensorBase readSparse(std::istream& stream, const Format& format, 34 | bool symm = false); 35 | TensorBase readDense(std::istream& stream, const Format& format, 36 | bool symm = false); 37 | 38 | /// Write an mtx matrix to a file. 39 | void writeMTX(std::string filename, const TensorBase& tensor); 40 | 41 | /// Write an mtx matrix to a stream. 42 | void writeMTX(std::ostream& stream, const TensorBase& tensor); 43 | void writeSparse(std::ostream& stream, const TensorBase& tensor); 44 | void writeDense(std::ostream& stream, const TensorBase& tensor); 45 | 46 | } 47 | 48 | #endif 49 | -------------------------------------------------------------------------------- /include/taco/storage/file_io_rb.h: -------------------------------------------------------------------------------- 1 | /// Read and write the HB Harwell-Boeing Sparse File Format 2 | 3 | #ifndef TACO_FILE_IO_RB_H 4 | #define TACO_FILE_IO_RB_H 5 | 6 | #include 7 | #include 8 | #include 9 | 10 | #include "taco/format.h" 11 | 12 | namespace taco { 13 | class TensorBase; 14 | class Format; 15 | 16 | void readFile(std::istream &hbfile, 17 | int* nrow, int* ncol, 18 | int** colptr, int** rowind, double** values); 19 | template 20 | void writeFile(std::ostream &hbfile, std::string key, 21 | int nrow, int ncol, int nnzero, 22 | int ptrsize, int indsize, int valsize, 23 | int* colptr, int* rowind, T* values); 24 | 25 | void readHeader(std::istream &hbfile, 26 | std::string* title, std::string* key, 27 | int* totcrd, int* ptrcrd, int* indcrd, int* valcrd, int* rhscrd, 28 | std::string* mxtype, int* nrow, 29 | int* ncol, int* nnzero, int* neltvl, 30 | std::string* ptrfmt, std::string* indfmt, 31 | std::string* valfmt, std::string* rhsfmt); 32 | void writeHeader(std::ostream &hbfile, 33 | std::string title, std::string key, 34 | int totcrd, int ptrcrd, int indcrd, int valcrd, int rhscrd, 35 | std::string mxtype, int nrow, int ncol, int nnzero, int neltvl, 36 | std::string ptrfmt, std::string indfmt, 37 | std::string valfmt, std::string rhsfmt); 38 | void readIndices(std::istream &hbfile, int linesize, int indices[]); 39 | void writeIndices(std::ostream &hbfile, int indsize, 40 | int linesize, int indices[]); 41 | void readValues(std::istream &hbfile, int linesize, double values[]); 42 | 43 | template 44 | void writeValues(std::ostream &hbfile, int valuesize, 45 | int valperline, T values[]); 46 | 47 | // Useless for Taco 48 | void readRHS(); 49 | void writeRHS(); 50 | 51 | /// Read an rb matrix from a file. 52 | TensorBase readRB(std::string filename, const ModeFormat& modetype, 53 | bool pack=true); 54 | 55 | /// Read an rb matrix from a file. 56 | TensorBase readRB(std::string filename, const Format& format, bool pack=true); 57 | 58 | /// Read an rb matrix from a stream 59 | TensorBase readRB(std::istream& stream, const ModeFormat& modetype, 60 | bool pack=true); 61 | 62 | /// Read an rb matrix from a stream 63 | TensorBase readRB(std::istream& stream, const Format& format, bool pack=true); 64 | 65 | /// Write an rb matrix to a file 66 | void writeRB(std::string filename, const TensorBase& tensor); 67 | 68 | /// Write an rb matrix to a stream 69 | void writeRB(std::ostream& stream, const TensorBase& tensor); 70 | 71 | } 72 | #endif 73 | -------------------------------------------------------------------------------- /include/taco/storage/file_io_tns.h: -------------------------------------------------------------------------------- 1 | #ifndef TACO_FILE_IO_TNS_H 2 | #define TACO_FILE_IO_TNS_H 3 | 4 | #include 5 | #include 6 | #include 7 | 8 | #include "taco/format.h" 9 | 10 | namespace taco { 11 | class TensorBase; 12 | class Format; 13 | 14 | /// Read a tns tensor from a file. 15 | TensorBase readTNS(std::string filename, const ModeFormat& modetype, 16 | bool pack=true); 17 | 18 | /// Read a tns tensor from a file. 19 | TensorBase readTNS(std::string filename, const Format& format, bool pack=true); 20 | 21 | /// Read a tns tensor from a stream. 22 | TensorBase readTNS(std::istream& stream, const ModeFormat& modetype, 23 | bool pack=true); 24 | 25 | /// Read a tns tensor from a stream. 26 | TensorBase readTNS(std::istream& stream, const Format& format, bool pack=true); 27 | 28 | /// Write a tns tensor to a file. 29 | void writeTNS(std::string filename, const TensorBase& tensor); 30 | 31 | /// Write a tns tensor to a stream. 32 | void writeTNS(std::ostream& stream, const TensorBase& tensor); 33 | 34 | } 35 | 36 | #endif 37 | -------------------------------------------------------------------------------- /include/taco/storage/index.h: -------------------------------------------------------------------------------- 1 | #ifndef TACO_STORAGE_INDEX_H 2 | #define TACO_STORAGE_INDEX_H 3 | 4 | #include 5 | #include 6 | #include 7 | 8 | #include "taco/format.h" 9 | #include "taco/taco_tensor_t.h" 10 | 11 | namespace taco { 12 | class ModeIndex; 13 | class Array; 14 | 15 | /// An index contains the index data structures of a tensor, but not its values. 16 | /// Thus, an index has a format and zero or more mode indices that describes the 17 | /// non-empty coordinates in each mode. 18 | class Index { 19 | public: 20 | /// Construct an empty index. 21 | Index(); 22 | 23 | /// Construct an index with the given format. 24 | Index(const Format& format); 25 | 26 | /// Construct an index with the given format and data. 27 | /// TODO DEPRECATE 28 | Index(const Format& format, const std::vector& indices); 29 | 30 | /// Returns the index's format. 31 | const Format& getFormat() const; 32 | 33 | /// Returns the number of indices (same as order in format); 34 | int numModeIndices() const; 35 | 36 | /// Returns the ith mode sub-index. 37 | /// @{ 38 | const ModeIndex& getModeIndex(int i) const; 39 | ModeIndex getModeIndex(int i); 40 | /// @} 41 | 42 | /// Returns the index size, which is the number of values it describes. 43 | size_t getSize() const; 44 | 45 | private: 46 | struct Content; 47 | std::shared_ptr content; 48 | }; 49 | 50 | std::ostream& operator<<(std::ostream&, const Index&); 51 | 52 | 53 | /// A mode sub-index of an Index. The type of the mode index is determined by 54 | /// the Format of the Index it is part of. 55 | class ModeIndex { 56 | public: 57 | /// Construct an empty mode index. 58 | ModeIndex(); 59 | 60 | /// Construct a mode index from a set of index arrays. 61 | ModeIndex(const std::vector& indexArrays); 62 | 63 | /// Returns the number of index arrays in this mode index. 64 | int numIndexArrays() const; 65 | 66 | /// Returns the ith index array. The number of index arrays are dictated by 67 | /// the ModeIndex's format in its parent Index. 68 | /// @{ 69 | const Array& getIndexArray(int i) const; 70 | Array getIndexArray(int i); 71 | /// @} 72 | 73 | private: 74 | struct Content; 75 | std::shared_ptr content; 76 | }; 77 | 78 | 79 | /// Factory functions to construct a compressed sparse rows (CSR) index. 80 | /// @{ 81 | Index makeCSRIndex(size_t numrows, int* rowptr, int* colidx); 82 | Index makeCSRIndex(const std::vector& rowptr, 83 | const std::vector& colidx); 84 | /// @} 85 | 86 | /// Factory functions to construct a compressed sparse columns (CSC) index. 87 | /// @{ 88 | Index makeCSCIndex(size_t numrows, int* colptr, int* rowidx); 89 | Index makeCSCIndex(const std::vector& colptr, 90 | const std::vector& rowidx); 91 | /// @} 92 | 93 | } 94 | #endif 95 | -------------------------------------------------------------------------------- /include/taco/storage/pack.h: -------------------------------------------------------------------------------- 1 | /// The pack machinery packs a tensor's non-zero components according to the 2 | /// tensor's storage format. The machinery is available both as an interpreter 3 | /// that can pack into any format, and as a code generator that generates 4 | /// specialized packing code for one format. 5 | 6 | #ifndef TACO_STORAGE_PACK_H 7 | #define TACO_STORAGE_PACK_H 8 | 9 | #include 10 | #include 11 | 12 | #include "taco/type.h" 13 | #include "taco/format.h" 14 | #include "taco/storage/typed_vector.h" 15 | #include "taco/storage/storage.h" 16 | #include "taco/storage/coordinate.h" 17 | 18 | namespace taco { 19 | 20 | class Literal; 21 | 22 | namespace ir { 23 | class Stmt; 24 | } 25 | 26 | TensorStorage pack(Datatype datatype, 27 | const std::vector& dimensions, 28 | const Format& format, 29 | const std::vector& coordinates, 30 | const void* values, 31 | const Literal& fill); 32 | 33 | template 34 | TensorStorage pack(std::vector dimensions, Format format, 35 | const std::vector,V>>& components, 36 | const Literal& fill){ 37 | size_t order = dimensions.size(); 38 | size_t nnz = components.size(); 39 | 40 | std::vector coordinates(order, 41 | TypedIndexVector(type(), nnz)); 42 | std::vector values(nnz); 43 | for (size_t i = 0; i < nnz; i++) { 44 | values[i] = components[i].second; 45 | auto& coords = components[i].first; 46 | for (size_t j = 0; j < order; j++) { 47 | coordinates[j][i] = coords[j]; 48 | } 49 | } 50 | 51 | return pack(type(), dimensions, format, coordinates, values.data(), fill); 52 | } 53 | 54 | } 55 | #endif 56 | -------------------------------------------------------------------------------- /include/taco/storage/storage.h: -------------------------------------------------------------------------------- 1 | #ifndef TACO_STORAGE_STORAGE_H 2 | #define TACO_STORAGE_STORAGE_H 3 | 4 | #include 5 | #include 6 | 7 | #include "taco/format.h" 8 | #include "taco/storage/index.h" 9 | #include "taco/storage/array.h" 10 | #include "taco/storage/typed_vector.h" 11 | #include "taco/storage/typed_index.h" 12 | 13 | struct taco_tensor_t; 14 | 15 | namespace taco { 16 | class Format; 17 | class Type; 18 | class Datatype; 19 | class Index; 20 | class Array; 21 | class Literal; 22 | 23 | /// Storage for a tensor object. Tensor storage consists of a value array that 24 | /// contains the tensor values and one index per mode. The type of each 25 | /// mode index is determined by the mode type in the format, and the 26 | /// ordering of the mode indices is determined by the format mode ordering. 27 | class TensorStorage { 28 | public: 29 | 30 | /// Construct tensor storage for the given format. 31 | TensorStorage(Datatype componentType, const std::vector& dimensions, 32 | Format format, Literal fillVal); 33 | 34 | /// Returns the tensor storage format. 35 | const Format& getFormat() const; 36 | 37 | /// Returns the component type of the tensor storage. 38 | Datatype getComponentType() const; 39 | 40 | /// Returns the order of the tensor storage. 41 | int getOrder() const; 42 | 43 | /// Returns the dimensions of the tensor storage. 44 | const std::vector& getDimensions() const; 45 | 46 | /// Get the tensor index, which describes the non-zero values. 47 | /// @{ 48 | const Index& getIndex() const; 49 | Index getIndex(); 50 | /// @} 51 | 52 | /// Returns the value array that contains the tensor components. 53 | const Array& getValues() const; 54 | 55 | /// Returns the tensor component value array. 56 | Array getValues(); 57 | 58 | /// Returns the full value attached to the tensor storage 59 | Literal getFillValue(); 60 | 61 | /// Returns the size of the storage in bytes. 62 | size_t getSizeInBytes(); 63 | 64 | /// Convert to a taco_tensor_t, whose lifetime is the same as the storage. 65 | operator struct taco_tensor_t*() const; 66 | 67 | /// Set the tensor index, which describes the non-zero values. 68 | void setIndex(const Index& index); 69 | 70 | /// Set the tensor component value array. 71 | void setValues(const Array& values); 72 | 73 | 74 | private: 75 | struct Content; 76 | std::shared_ptr content; 77 | }; 78 | 79 | 80 | 81 | /// Compare tensor storage objects. 82 | bool equals(TensorStorage a, TensorStorage b); 83 | 84 | /// Print Storage objects to a stream. 85 | std::ostream& operator<<(std::ostream&, const TensorStorage&); 86 | 87 | } 88 | #endif 89 | -------------------------------------------------------------------------------- /include/taco/taco_tensor_t.h: -------------------------------------------------------------------------------- 1 | /// This file defines the runtime struct used to pass raw tensors to generated 2 | /// code. Note: this file must be valid C99, not C++. 3 | /// This *must* be kept in sync with the version used in codegen_c.cpp 4 | /// TODO: Remove `vals_size` after old lowering machinery has been replaced. 5 | 6 | #ifndef TACO_TENSOR_T_DEFINED 7 | #define TACO_TENSOR_T_DEFINED 8 | 9 | #include 10 | 11 | typedef enum { taco_mode_dense, taco_mode_sparse } taco_mode_t; 12 | 13 | typedef struct taco_tensor_t { 14 | int32_t order; // tensor order (number of modes) 15 | int32_t* dimensions; // tensor dimensions 16 | int32_t csize; // component size 17 | int32_t* mode_ordering; // mode storage ordering 18 | taco_mode_t* mode_types; // mode storage types 19 | uint8_t*** indices; // tensor index data (per mode) 20 | uint8_t* vals; // tensor values 21 | uint8_t* fill_value; // tensor fill value 22 | int32_t vals_size; // values array size 23 | } taco_tensor_t; 24 | 25 | taco_tensor_t *init_taco_tensor_t(int32_t order, int32_t csize, 26 | int32_t* dimensions, int32_t* modeOrdering, 27 | taco_mode_t* mode_types, void* fill_ptr); 28 | 29 | void deinit_taco_tensor_t(taco_tensor_t* t); 30 | 31 | #endif 32 | -------------------------------------------------------------------------------- /include/taco/target.h: -------------------------------------------------------------------------------- 1 | #ifndef TACO_TARGET_H 2 | #define TACO_TARGET_H 3 | 4 | #include "taco/error.h" 5 | 6 | namespace taco { 7 | 8 | /// This struct represents the machine & OS to generate code for, both for 9 | /// JIT and AOT code generation. 10 | struct Target { 11 | /// Architectures. If C99, we generate C code, and if it is a specific 12 | /// machine arch (e.g. x86 or arm) we use LLVM. 13 | enum Arch {C99=0, X86} arch; 14 | 15 | /// Operating System. Used when deciding which OS-specific calls to use. 16 | enum OS {OSUnknown=0, Linux, MacOS, Windows} os; 17 | 18 | std::string compiler_env = "TACO_CC"; 19 | 20 | std::string compiler = "cc"; 21 | 22 | // As we support them, we'll stick in optional features into the target as 23 | // well, including things like parallelism model (e.g. openmp, cilk) for 24 | // C code generation, and hardware features (e.g. AVX) for LLVM code gen. 25 | 26 | /// Given a string of the form arch-os-features, construct the corresponding 27 | /// Target object. 28 | Target(const std::string &s); 29 | 30 | Target(Arch a, OS o) : arch(a), os(o) { 31 | taco_tassert(a == C99 && o != Windows && o != OSUnknown) 32 | << "Unsupported target."; 33 | } 34 | 35 | /// Validate a target string 36 | static bool validateTargetString(const std::string &s); 37 | 38 | }; 39 | 40 | /// Gets the target from the environment. If this is not set in the 41 | /// environment, it uses the default C99 backend with the current OS 42 | Target getTargetFromEnvironment(); 43 | 44 | } // namespace taco 45 | 46 | #endif 47 | -------------------------------------------------------------------------------- /include/taco/util/comparable.h: -------------------------------------------------------------------------------- 1 | #ifndef TACO_UTIL_COMPARABLE_H 2 | #define TACO_UTIL_COMPARABLE_H 3 | 4 | namespace taco { 5 | namespace util { 6 | 7 | /// Interface for classes that can be compared to each other. Classes that 8 | /// implement this interfaces define `==` and `<` to get `!=`, `>`, `<=`, and 9 | /// `>=` for free. 10 | template 11 | class Comparable {}; 12 | 13 | template 14 | bool operator!=(const Comparable &lhs, const Comparable &rhs) { 15 | return !(*static_cast(&lhs) == *static_cast(&rhs)); 16 | } 17 | 18 | template 19 | bool operator>(const Comparable &lhs, const Comparable &rhs) { 20 | return !(*static_cast(&lhs) < *static_cast(&rhs) || 21 | *static_cast(&lhs) == *static_cast(&rhs)); 22 | } 23 | 24 | template 25 | bool operator<=(const Comparable &lhs, const Comparable &rhs) { 26 | return !(*static_cast(&lhs) > *static_cast(&rhs)); 27 | } 28 | 29 | template 30 | bool operator>=(const Comparable &lhs, const Comparable &rhs) { 31 | return !(*static_cast(&lhs) < *static_cast(&rhs)); 32 | } 33 | 34 | }} 35 | #endif 36 | -------------------------------------------------------------------------------- /include/taco/util/env.h: -------------------------------------------------------------------------------- 1 | #ifndef SRC_UTIL_ENV_H_ 2 | #define SRC_UTIL_ENV_H_ 3 | 4 | #include 5 | #include 6 | #include 7 | 8 | #include "taco/error.h" 9 | 10 | namespace taco { 11 | namespace util { 12 | std::string getFromEnv(std::string flag, std::string dflt); 13 | std::string getTmpdir(); 14 | extern std::string cachedtmpdir; 15 | extern void cachedtmpdirCleanup(void); 16 | 17 | inline std::string getFromEnv(std::string flag, std::string dflt) { 18 | char const *ret = getenv(flag.c_str()); 19 | if (!ret) { 20 | return dflt; 21 | } else { 22 | return std::string(ret); 23 | } 24 | } 25 | 26 | inline std::string getTmpdir() { 27 | if (cachedtmpdir == ""){ 28 | // use posix logic for finding a temp dir 29 | auto tmpdir = getFromEnv("TMPDIR", "/tmp/"); 30 | 31 | // if the directory does not have a trailing slash, add one 32 | if (tmpdir.back() != '/') { 33 | tmpdir += '/'; 34 | } 35 | 36 | // ensure it is an absolute path 37 | taco_uassert(tmpdir.front() == '/') << 38 | "The TMPDIR environment variable must be an absolute path"; 39 | 40 | taco_uassert(access(tmpdir.c_str(), W_OK) == 0) << 41 | "Unable to write to temporary directory for code generation. " 42 | "Please set the environment variable TMPDIR to somewhere writable"; 43 | 44 | // ensure that we use a taco tmpdir unique to this process. 45 | auto tacotmpdirtemplate = tmpdir + "taco_tmp_XXXXXX"; 46 | char *ctacotmpdirtemplate = new char[tacotmpdirtemplate.length() + 1]; 47 | std::strcpy(ctacotmpdirtemplate, tacotmpdirtemplate.c_str()); 48 | char *ctacotmpdir = mkdtemp(ctacotmpdirtemplate); 49 | taco_uassert(ctacotmpdir != NULL) << 50 | "Unable to create taco temporary directory for code generation. Please set" 51 | "the environment variable TMPDIR to somewhere searchable and writable"; 52 | std::string tacotmpdir(ctacotmpdir); 53 | delete [] ctacotmpdirtemplate; 54 | 55 | // if the directory does not have a trailing slash, add one 56 | if (tacotmpdir.back() != '/') { 57 | tacotmpdir += '/'; 58 | } 59 | 60 | cachedtmpdir = tacotmpdir; 61 | 62 | //cleanup unless we are in debug mode 63 | #ifndef TACO_DEBUG 64 | atexit(cachedtmpdirCleanup); 65 | #endif 66 | } 67 | return cachedtmpdir; 68 | } 69 | 70 | }} 71 | 72 | #endif /* SRC_UTIL_ENV_H_ */ 73 | -------------------------------------------------------------------------------- /include/taco/util/files.h: -------------------------------------------------------------------------------- 1 | #ifndef TACO_UTIL_FILES_H 2 | #define TACO_UTIL_FILES_H 3 | 4 | #include 5 | #include 6 | 7 | namespace taco { 8 | namespace util { 9 | 10 | std::string sanitizePath(std::string path); 11 | 12 | void openStream(std::fstream& stream, std::string path, std::fstream::openmode mode); 13 | 14 | }} 15 | #endif 16 | -------------------------------------------------------------------------------- /include/taco/util/functions.h: -------------------------------------------------------------------------------- 1 | #ifndef TACO_FUNCTIONAL_H 2 | #define TACO_FUNCTIONAL_H 3 | 4 | #include 5 | 6 | namespace taco { 7 | namespace util { 8 | 9 | template 10 | Fnptr functorAddress(std::function f) { 11 | return *f.template target(); 12 | } 13 | 14 | template 15 | bool targetPtrEqual(std::function f, std::function g) { 16 | return functorAddress(f) != nullptr && functorAddress(f) == functorAddress(g); 17 | } 18 | 19 | } 20 | } 21 | #endif //TACO_FUNCTIONAL_H 22 | -------------------------------------------------------------------------------- /include/taco/util/name_generator.h: -------------------------------------------------------------------------------- 1 | #ifndef TACO_UTIL_NAME_GENERATOR_H 2 | #define TACO_UTIL_NAME_GENERATOR_H 3 | 4 | #include 5 | #include 6 | #include 7 | 8 | namespace taco { 9 | namespace util { 10 | 11 | std::string uniqueName(char prefix); 12 | std::string uniqueName(const std::string& prefix); 13 | 14 | int getUniqueId(); 15 | 16 | class NameGenerator { 17 | public: 18 | NameGenerator(); 19 | NameGenerator(std::vector reserved); 20 | 21 | std::string getUniqueName(std::string name); 22 | 23 | private: 24 | std::map nameCounters; 25 | }; 26 | 27 | }} 28 | #endif 29 | -------------------------------------------------------------------------------- /include/taco/util/printable.h: -------------------------------------------------------------------------------- 1 | #ifndef TACO_UTIL_PRINTABLE_H 2 | #define TACO_UTIL_PRINTABLE_H 3 | 4 | #include 5 | 6 | namespace taco { 7 | namespace util { 8 | 9 | class Printable { 10 | public: 11 | virtual void print(std::ostream &os) const = 0; 12 | virtual ~Printable() {}; 13 | }; 14 | 15 | inline std::ostream &operator<<(std::ostream &os, const Printable &printable) { 16 | printable.print(os); 17 | return os; 18 | } 19 | 20 | }} 21 | #endif 22 | -------------------------------------------------------------------------------- /include/taco/util/scopedmap.h: -------------------------------------------------------------------------------- 1 | #ifndef TACO_UTIL_SCOPEDMAP_H 2 | #define TACO_UTIL_SCOPEDMAP_H 3 | 4 | #include 5 | #include 6 | #include 7 | #include 8 | #include "taco/util/collections.h" 9 | 10 | #include "taco/error.h" 11 | 12 | namespace taco { 13 | namespace util { 14 | 15 | template 16 | class ScopedMap { 17 | public: 18 | ScopedMap() { 19 | scope(); 20 | } 21 | 22 | ~ScopedMap() { 23 | unscope(); 24 | } 25 | 26 | /// Add a level of scoping. 27 | void scope() { 28 | scopes.push_front(std::map()); 29 | } 30 | 31 | /// Remove a level of scoping. 32 | void unscope() { 33 | scopes.pop_front(); 34 | } 35 | 36 | void insert(const std::pair& value) { 37 | scopes.front().insert(value); 38 | } 39 | 40 | void remove(const Key& key) { 41 | for (auto& scope : scopes) { 42 | const auto it = scope.find(key); 43 | if (it != scope.end()) { 44 | scope.erase(it); 45 | return; 46 | } 47 | } 48 | taco_ierror << "Not in scope"; 49 | } 50 | 51 | const Value& get(const Key& key) const { 52 | for (auto& scope : scopes) { 53 | if (scope.find(key) != scope.end()) { 54 | return scope.at(key); 55 | } 56 | } 57 | taco_ierror << "Not in scope"; 58 | return scopes.front().begin()->second; // silence warnings 59 | } 60 | 61 | bool contains(const Key& key) { 62 | for (auto& scope : scopes) { 63 | if (scope.find(key) != scope.end()) { 64 | return true; 65 | } 66 | } 67 | return false; 68 | } 69 | 70 | friend std::ostream& operator<<(std::ostream& os, ScopedMap smap) { 71 | os << "ScopedMap:" << std::endl; 72 | for (auto& scope : util::reverse(smap.scopes)) { 73 | os << " - "; 74 | if (scope.size() > 0) { 75 | auto val = *scope.begin(); 76 | os << val.first << " -> " << val.second << std::endl; 77 | } 78 | for (auto& val : excludeFirst(scope)) { 79 | os << " " << val.first << " -> " << val.second << std::endl; 80 | } 81 | std::cout << std::endl; 82 | } 83 | return os; 84 | } 85 | 86 | private: 87 | std::list> scopes; 88 | }; 89 | 90 | }} 91 | #endif 92 | -------------------------------------------------------------------------------- /include/taco/util/scopedset.h: -------------------------------------------------------------------------------- 1 | #ifndef TACO_UTIL_SCOPEDSET_H 2 | #define TACO_UTIL_SCOPEDSET_H 3 | 4 | #include 5 | #include 6 | #include 7 | 8 | #include "taco/error.h" 9 | 10 | namespace taco { 11 | namespace util { 12 | 13 | template 14 | class ScopedSet { 15 | public: 16 | ScopedSet() { 17 | scope(); 18 | } 19 | 20 | ~ScopedSet() { 21 | unscope(); 22 | } 23 | 24 | /// Add a level of scoping. 25 | void scope() { 26 | scopes.push_front(std::set()); 27 | } 28 | 29 | /// Remove a level of scoping. 30 | void unscope() { 31 | scopes.pop_front(); 32 | } 33 | 34 | void insert(const Key& key) { 35 | scopes.front().insert(key); 36 | } 37 | 38 | void remove(const Key& key) { 39 | for (auto& scope : scopes) { 40 | const auto it = scope.find(key); 41 | if (it != scope.end()) { 42 | scope.erase(it); 43 | return; 44 | } 45 | } 46 | taco_ierror << "Not in scope"; 47 | } 48 | 49 | bool contains(const Key& key) { 50 | for (auto& scope : scopes) { 51 | if (scope.find(key) != scope.end()) { 52 | return true; 53 | } 54 | } 55 | return false; 56 | } 57 | 58 | friend std::ostream& operator<<(std::ostream& os, ScopedSet sset) { 59 | os << "ScopedSet:" << std::endl; 60 | return os; 61 | } 62 | 63 | private: 64 | std::list> scopes; 65 | }; 66 | 67 | }} 68 | #endif 69 | -------------------------------------------------------------------------------- /include/taco/util/singleton.h: -------------------------------------------------------------------------------- 1 | /// Define a singleton design pattern 2 | #define SINGLETON(NAME) \ 3 | public: \ 4 | static NAME& get##NAME() { \ 5 | static NAME NAME_singleton; \ 6 | return NAME_singleton; \ 7 | }; \ 8 | NAME(NAME const&) = delete; \ 9 | void operator=(NAME const&) = delete; \ 10 | private: \ 11 | NAME() {}; 12 | -------------------------------------------------------------------------------- /include/taco/util/strings.h: -------------------------------------------------------------------------------- 1 | #ifndef TACO_UTIL_STRINGS_H 2 | #define TACO_UTIL_STRINGS_H 3 | 4 | #include 5 | #include 6 | #include 7 | #include 8 | #include 9 | #include 10 | #include 11 | 12 | // To get the value of a compiler macro variable 13 | #define STRINGIFY(x) #x 14 | #define TO_STRING(x) STRINGIFY(x) 15 | 16 | namespace taco { 17 | namespace util { 18 | 19 | /// Turn anything except floating points that can be written to a stream 20 | /// into a string. 21 | template 22 | typename std::enable_if::value, std::string>::type 23 | toString(const T &val) { 24 | std::stringstream sstream; 25 | sstream << val; 26 | return sstream.str(); 27 | } 28 | 29 | /// Turn any floating point that can be written to a stream into a string, 30 | /// forcing full precision and inclusion of the decimal point. 31 | template 32 | typename std::enable_if::value, std::string>::type 33 | toString(const T &val) { 34 | if (std::isinf(val)) { 35 | return (val < 0) ? "-INFINITY" : "INFINITY"; 36 | } 37 | std::stringstream sstream; 38 | sstream << std::setprecision(std::numeric_limits::max_digits10) << std::showpoint << val; 39 | return sstream.str(); 40 | } 41 | 42 | /// Join the elements between begin and end in a sep-separated string. 43 | template 44 | std::string join(Iterator begin, Iterator end, const std::string &sep=", ") { 45 | std::ostringstream result; 46 | if (begin != end) { 47 | result << *begin++; 48 | } 49 | while (begin != end) { 50 | result << sep << *begin++; 51 | } 52 | return result.str(); 53 | } 54 | 55 | /// Join the elements in the collection in a sep-separated string. 56 | template 57 | std::string join(const Collection &collection, const std::string &sep=", ") { 58 | return join(collection.begin(), collection.end(), sep); 59 | } 60 | 61 | /// Join the elements in the map into a sep-separated string. 62 | template 63 | std::string join(const std::map &collection, const std::string &sep=", ") { 64 | std::ostringstream result; 65 | auto begin = collection.begin(); 66 | auto end = collection.end(); 67 | if (begin != end) { 68 | result << begin->first << " -> " << begin->second; 69 | begin++; 70 | } 71 | while (begin != end) { 72 | result << sep << begin->first << " -> " << begin->second; 73 | begin++; 74 | } 75 | return result.str(); 76 | } 77 | 78 | /// Split the string. 79 | std::vector split(const std::string &str, const std::string &delim, 80 | bool keepDelim = false); 81 | 82 | /// Returns the text repeated n times 83 | std::string repeat(std::string text, size_t n); 84 | 85 | /// Returns a string of `n` characters where `text` is centered and the rest 86 | /// is filled with `fill`. 87 | std::string fill(std::string text, char fill, size_t n); 88 | 89 | }} 90 | #endif 91 | -------------------------------------------------------------------------------- /include/taco/util/uncopyable.h: -------------------------------------------------------------------------------- 1 | #ifndef TACO_UTIL_UNCOPYABLE_H 2 | #define TACO_UTIL_UNCOPYABLE_H 3 | 4 | namespace taco { 5 | namespace util { 6 | 7 | class Uncopyable { 8 | protected: 9 | Uncopyable() = default; 10 | ~Uncopyable() = default; 11 | 12 | private: 13 | Uncopyable(const Uncopyable&) = delete; 14 | Uncopyable& operator=(const Uncopyable&) = delete; 15 | }; 16 | 17 | }} 18 | #endif 19 | -------------------------------------------------------------------------------- /include/taco/util/variadic.h: -------------------------------------------------------------------------------- 1 | #ifndef TACO_UTIL_VARIADIC_H 2 | #define TACO_UTIL_VARIADIC_H 3 | 4 | #include 5 | #include 6 | 7 | namespace taco { 8 | namespace util { 9 | 10 | // Compare types 11 | template 12 | struct areSame : std::true_type {}; 13 | 14 | template 15 | struct areSame 16 | : std::integral_constant{} && areSame{}> 17 | {}; 18 | 19 | 20 | // Product 21 | namespace { 22 | 23 | template 24 | struct productHelper; 25 | 26 | template 27 | struct productHelper { 28 | static const int value = productHelper::value; 29 | }; 30 | 31 | 32 | template 33 | struct productHelper { 34 | static const int value = val; 35 | }; 36 | 37 | } // unnamed namespace 38 | 39 | template struct product; 40 | 41 | template <> 42 | struct product <> { 43 | static constexpr int value = 1; 44 | }; 45 | 46 | template 47 | struct product { 48 | static const int value = productHelper<1, val, rest...>::value; 49 | }; 50 | 51 | template 52 | struct product { 53 | static const int value = product::value; 54 | }; 55 | 56 | 57 | // Machinery to compute array offsets 58 | template struct seq {}; 59 | 60 | /// Remove first value from int variadic template 61 | template 62 | struct removeFirst { 63 | typedef seq type; 64 | }; 65 | 66 | 67 | /// Compute product of static sequence 68 | template inline constexpr 69 | int computeProduct(seq seq) { 70 | return product::value; 71 | } 72 | 73 | 74 | /// Compute the offset into an n-dimensional array 75 | template inline 76 | int computeOffset(seq dims, int index, Indices... rest) { 77 | typename removeFirst::type innerDims; 78 | return index * computeProduct(innerDims) + computeOffset(innerDims, rest...); 79 | } 80 | 81 | template inline constexpr 82 | int computeOffset(const seq &dims, int i) { 83 | return i; 84 | } 85 | 86 | 87 | /// Compute the offset into an n-dimensional array 88 | template inline 89 | int computeOffset(seq dims, const std::vector& indices) { 90 | return computeOffset(dims, indices.begin(), indices.end()); 91 | } 92 | 93 | template inline 94 | int computeOffset(seq dims, 95 | const std::vector::const_iterator& begin, 96 | const std::vector::const_iterator& end) { 97 | typename removeFirst::type innerDims; 98 | const size_t i = *begin; 99 | constexpr size_t stride = computeProduct(innerDims); 100 | constexpr size_t rest = computeOffset(innerDims, begin+1, end); 101 | return i * stride + rest; 102 | } 103 | 104 | template inline 105 | int computeOffset(const seq<> &dims, 106 | const std::vector::const_iterator& begin, 107 | const std::vector::const_iterator& end) { 108 | return 0; 109 | } 110 | 111 | inline constexpr 112 | int computeOffset(seq<> dims, const std::vector& indices) { 113 | return 0; 114 | } 115 | 116 | }} 117 | #endif 118 | -------------------------------------------------------------------------------- /include/taco/version.h.in: -------------------------------------------------------------------------------- 1 | #ifndef TACO_VERSION_H 2 | #define TACO_VERSION_H 3 | 4 | // This file contains version/config info, gathered at cmake config time. 5 | 6 | #define TACO_BUILD_TYPE "@CMAKE_BUILD_TYPE@" 7 | #define TACO_BUILD_DATE "@TACO_BUILD_DATE@" 8 | 9 | #define TACO_BUILD_COMPILER_ID "@CMAKE_CXX_COMPILER_ID@" 10 | #define TACO_BUILD_COMPILER_VERSION "@CMAKE_CXX_COMPILER_VERSION@" 11 | 12 | #define TACO_VERSION_MAJOR "@PROJECT_VERSION_MAJOR@" 13 | #define TACO_VERSION_MINOR "@PROJECT_VERSION_MINOR@" 14 | // if taco starts using a patch number, add @PROJECT_VERSION_PATCH@ here 15 | // if taco starts using a tweak number, add @PROJECT_VERSION_TWEAK@ here 16 | 17 | // For non-git builds, this will be an empty string. 18 | #define TACO_VERSION_GIT_SHORTHASH "@TACO_GIT_SHORTHASH@" 19 | 20 | #define TACO_FEATURE_OPENMP @TACO_FEATURE_OPENMP@ 21 | #define TACO_FEATURE_PYTHON @TACO_FEATURE_PYTHON@ 22 | #define TACO_FEATURE_CUDA @TACO_FEATURE_CUDA@ 23 | 24 | #endif /* TACO_VERSION_H */ 25 | -------------------------------------------------------------------------------- /misc/mainpage.h: -------------------------------------------------------------------------------- 1 | /** \file 2 | * This file only exists to contain the front-page of the documentation 3 | */ 4 | 5 | /** \mainpage The Tensor Algebra Compiler (taco) 6 | * 7 | * The Tensor Algebra Compiler (taco) compiles and computes expressions on 8 | * sparse and dense tensors. 9 | * 10 | * To learn about taco, we recommend you start by reading the documentation at 11 | * [http://tensor-compiler.org/docs/]. 12 | * 13 | * You can also look in the apps directory for examples of how to use the API. 14 | * 15 | */ 16 | 17 | /** 18 | * \example apps/tensor_times_vector/tensor_times_vector.cpp 19 | */ 20 | -------------------------------------------------------------------------------- /python_bindings/CMakeLists.txt: -------------------------------------------------------------------------------- 1 | set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -flto -fvisibility=hidden -Wno-gnu-zero-variadic-macro-arguments") 2 | 3 | add_subdirectory(pybind11) 4 | 5 | include_directories(${PYTHON_INCLUDE_DIRS}) 6 | include_directories(${TACO_INCLUDE_DIR}) 7 | include_directories(${TACO_PROJECT_DIR}/python_bindings/include) 8 | if(CUDA) 9 | include_directories(${CUDA_INCLUDE_DIRS}) 10 | endif(CUDA) 11 | 12 | set(PY_SRC_DIRS src) 13 | 14 | foreach(dir ${PY_SRC_DIRS}) 15 | file(GLOB PY_SOURCES ${PY_SOURCES} ${dir}/*.cpp) 16 | endforeach() 17 | 18 | file(GLOB PYTACO_FILES ${CMAKE_CURRENT_SOURCE_DIR}/pytaco/*.py) 19 | file(GLOB PYTENSOR_FILES ${CMAKE_CURRENT_SOURCE_DIR}/pytaco/pytensor/*.py) 20 | 21 | set(PY_SOURCES ${PY_SOURCES}) 22 | pybind11_add_module(core_modules ${PY_SOURCES} ${TACO_SOURCES}) 23 | 24 | set_target_properties(core_modules PROPERTIES LIBRARY_OUTPUT_DIRECTORY ${CMAKE_LIBRARY_OUTPUT_DIRECTORY}/pytaco/core) 25 | target_link_libraries(core_modules LINK_PRIVATE taco) 26 | file(COPY ${CMAKE_CURRENT_SOURCE_DIR}/pytaco DESTINATION ${CMAKE_LIBRARY_OUTPUT_DIRECTORY}) 27 | 28 | set(PY_INSTALL_PATH ${CMAKE_INSTALL_PREFIX}/lib/pytaco) 29 | install(FILES ${PYTACO_FILES} DESTINATION ${PY_INSTALL_PATH}) 30 | install(FILES ${PYTENSOR_FILES} DESTINATION ${PY_INSTALL_PATH}/pytensor) 31 | install(TARGETS core_modules DESTINATION ${PY_INSTALL_PATH}/core) 32 | 33 | configure_file("unit_tests.py" "unit_tests.py" @ONLY) 34 | add_test(NAME pytaco COMMAND python3 unit_tests.py) 35 | -------------------------------------------------------------------------------- /python_bindings/Makefile: -------------------------------------------------------------------------------- 1 | # Minimal makefile for Sphinx documentation 2 | # 3 | 4 | # You can set these variables from the command line. 5 | SPHINXOPTS = 6 | SPHINXBUILD = sphinx-build 7 | SOURCEDIR = source 8 | BUILDDIR = build 9 | 10 | # Put it first so that "make" without argument is like "make help". 11 | help: 12 | @$(SPHINXBUILD) -M help "$(SOURCEDIR)" "$(BUILDDIR)" $(SPHINXOPTS) $(O) 13 | 14 | .PHONY: help Makefile 15 | 16 | # Catch-all target: route all unknown targets to Sphinx using the new 17 | # "make mode" option. $(O) is meant as a shortcut for $(SPHINXOPTS). 18 | %: Makefile 19 | @$(SPHINXBUILD) -M $@ "$(SOURCEDIR)" "$(BUILDDIR)" $(SPHINXOPTS) $(O) -------------------------------------------------------------------------------- /python_bindings/include/pyDatatypes.h: -------------------------------------------------------------------------------- 1 | #ifndef TACO_DATATYPES_H 2 | #define TACO_DATATYPES_H 3 | #include "taco/type.h" 4 | #include 5 | 6 | namespace py = pybind11; 7 | 8 | namespace taco{ 9 | namespace pythonBindings{ 10 | 11 | std::string getNpType(const taco::Datatype &type); 12 | py::object asNpDtype(const taco::Datatype& obj); 13 | void defineTacoTypes(py::module &m); 14 | 15 | }} 16 | 17 | #endif //TACO_DATATYPES_H 18 | -------------------------------------------------------------------------------- /python_bindings/include/pyFormat.h: -------------------------------------------------------------------------------- 1 | #ifndef TACO_PYFORMAT_H 2 | #define TACO_PYFORMAT_H 3 | 4 | #include "taco/format.h" 5 | #include 6 | 7 | namespace py = pybind11; 8 | 9 | namespace taco{ 10 | namespace pythonBindings{ 11 | 12 | std::size_t hashFormat(const taco::Format& format); 13 | void defineModeFormats(py::module& m); 14 | void defineModeFormatPack(py::module &m); 15 | void defineFormat(py::module &m); 16 | 17 | }} 18 | 19 | #endif //TACO_PYFORMAT_H 20 | -------------------------------------------------------------------------------- /python_bindings/include/pyIndexNotation.h: -------------------------------------------------------------------------------- 1 | #ifndef TACO_PYINDEXNOTATION_H 2 | #define TACO_PYINDEXNOTATION_H 3 | 4 | 5 | #include "taco/tensor.h" 6 | #include 7 | 8 | namespace py = pybind11; 9 | 10 | namespace taco{ 11 | namespace pythonBindings{ 12 | 13 | void defineIndexNotation(py::module &m); 14 | 15 | }} 16 | 17 | 18 | #endif //TACO_PYINDEXNOTATION_H 19 | -------------------------------------------------------------------------------- /python_bindings/include/pyParsers.h: -------------------------------------------------------------------------------- 1 | #ifndef TACO_PYPARSERS_H 2 | #define TACO_PYPARSERS_H 3 | 4 | #include "taco/parser/parser.h" 5 | #include "taco/parser/einsum_parser.h" 6 | #include 7 | 8 | namespace py = pybind11; 9 | 10 | namespace taco{ 11 | namespace pythonBindings{ 12 | 13 | void defineParser(py::module& m); 14 | 15 | }} 16 | 17 | #endif //TACO_PYPARSERS_H 18 | -------------------------------------------------------------------------------- /python_bindings/include/pyTensor.h: -------------------------------------------------------------------------------- 1 | #ifndef TACO_PY_TENSOR_H 2 | #define TACO_PY_TENSOR_H 3 | 4 | #include "taco/tensor.h" 5 | #include 6 | 7 | namespace py = pybind11; 8 | 9 | namespace taco{ 10 | namespace pythonBindings{ 11 | 12 | void defineTensor(py::module& m); 13 | 14 | }} 15 | 16 | #endif //TACO_PY_TENSOR_H 17 | -------------------------------------------------------------------------------- /python_bindings/include/pyTensorIO.h: -------------------------------------------------------------------------------- 1 | #ifndef TACO_PYTENSORIO_H 2 | #define TACO_PYTENSORIO_H 3 | 4 | #include "taco/tensor.h" 5 | #include 6 | 7 | namespace py = pybind11; 8 | 9 | namespace taco{ 10 | namespace pythonBindings{ 11 | 12 | void defineIOFuncs(py::module &m); 13 | 14 | 15 | }} 16 | 17 | #endif //TACO_PYTENSORIO_H 18 | -------------------------------------------------------------------------------- /python_bindings/include/pyeinsum.h: -------------------------------------------------------------------------------- 1 | #ifndef TACO_PY_EINSUM_H 2 | #define TACO_PY_EINSUM_H 3 | 4 | #include 5 | #include "taco/parser/parser.h" 6 | 7 | namespace py = pybind11; 8 | 9 | 10 | #endif //TACO_PY_EINSUM_H 11 | -------------------------------------------------------------------------------- /python_bindings/pytaco/__init__.py: -------------------------------------------------------------------------------- 1 | from .core.core_modules import * 2 | from .pytensor import * 3 | -------------------------------------------------------------------------------- /python_bindings/pytaco/numpy_license.txt: -------------------------------------------------------------------------------- 1 | Copyright (c) 2005-2019, NumPy Developers. 2 | All rights reserved. 3 | 4 | Redistribution and use in source and binary forms, with or without 5 | modification, are permitted provided that the following conditions are 6 | met: 7 | 8 | * Redistributions of source code must retain the above copyright 9 | notice, this list of conditions and the following disclaimer. 10 | 11 | * Redistributions in binary form must reproduce the above 12 | copyright notice, this list of conditions and the following 13 | disclaimer in the documentation and/or other materials provided 14 | with the distribution. 15 | 16 | * Neither the name of the NumPy Developers nor the names of any 17 | contributors may be used to endorse or promote products derived 18 | from this software without specific prior written permission. 19 | 20 | THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS 21 | "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT 22 | LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR 23 | A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT 24 | OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, 25 | SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT 26 | LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, 27 | DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY 28 | THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT 29 | (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE 30 | OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. -------------------------------------------------------------------------------- /python_bindings/pytaco/pytensor/__init__.py: -------------------------------------------------------------------------------- 1 | from .taco_tensor import * 2 | from .tensorIO import * 3 | -------------------------------------------------------------------------------- /python_bindings/pytaco/pytensor/tensorIO.py: -------------------------------------------------------------------------------- 1 | from ..core import core_modules as _cm 2 | from ..pytensor.taco_tensor import tensor 3 | 4 | 5 | def write(filename, t): 6 | """ 7 | Writes a tensor to a file. 8 | 9 | Writes the input tensor t to the file specified by filename using the extension given in the filename. This 10 | function only supports writing the extensions listed in the :ref:`io` section. 11 | 12 | Notes 13 | --------- 14 | This function forces tensor compilation. 15 | 16 | Parameters 17 | ------------- 18 | filename: string 19 | The name of the file the tensor will be written to. The file must be given an extension of the .mtx, .tns or 20 | .rb. 21 | 22 | t: tensor 23 | The tensor to write to filename. 24 | 25 | 26 | Examples 27 | ---------- 28 | To write to a .tns file we can do the following: 29 | 30 | .. doctest:: 31 | 32 | >>> import pytaco as pt 33 | >>> a = pt.tensor([2, 2], [pt.dense, pt.compressed], dtype=pt.float32) 34 | >>> a.insert([0, 0], 10) 35 | >>> pt.write("simple_test.tns", a) 36 | 37 | """ 38 | _cm._write(filename, t._tensor) 39 | 40 | 41 | def read(filename, fmt, pack=True): 42 | """ 43 | Reads a tensor from a file. 44 | 45 | Reads a tensor from the file filename. The extension must be one of the those supported by PyTaco listed in the 46 | :ref:`io` section. 47 | 48 | Parameters 49 | ------------- 50 | filename: string 51 | The name of the file containing the tensor to read. 52 | 53 | fmt: pytaco.format 54 | The :class:`~format` that PyTaco should use to store the tensor. 55 | 56 | pack: boolean, optional 57 | If true, by taco will immediately pack the tensor in the specified format after reading the file. Otherwise, 58 | it will keep the data in a buffer until either the user explicitly calls :func:`~pytaco.tensor.pack` or 59 | the tensor is implicitly packed by taco. 60 | 61 | Examples 62 | ---------- 63 | >>> import pytaco as pt 64 | >>> a = pt.tensor([2, 2], [pt.dense, pt.compressed], dtype=pt.float32) 65 | >>> a.insert([0, 0], 10) 66 | >>> pt.write("simple_test.tns", a) 67 | >>> t = pt.read("simple_test.tns", pt.csr) 68 | >>> t[0, 0] 69 | 10.0 70 | 71 | Returns 72 | --------- 73 | tensor 74 | A :class:`tensor` of type double, with :class:`format` fmt containing the data read from the file filename. 75 | """ 76 | cppTensor = _cm._read(filename, fmt, pack) 77 | return tensor._fromCppTensor(cppTensor) -------------------------------------------------------------------------------- /python_bindings/source/conf.py: -------------------------------------------------------------------------------- 1 | # Configuration file for the Sphinx documentation builder. 2 | # 3 | # This file only contains a selection of the most common options. For a full 4 | # list see the documentation: 5 | # http://www.sphinx-doc.org/en/master/config 6 | 7 | # -- Path setup -------------------------------------------------------------- 8 | 9 | # If extensions (or modules to document with autodoc) are in another directory, 10 | # add these directories to sys.path here. If the directory is relative to the 11 | # documentation root, use os.path.abspath to make it absolute, like shown here. 12 | # 13 | import os 14 | import sys 15 | sys.path.insert(0, os.path.abspath('../../build/lib/')) 16 | 17 | 18 | # -- Project information ----------------------------------------------------- 19 | 20 | project = 'TACO Python API Reference' 21 | copyright = '2022 TACO Development Team' 22 | author = 'TACO Development Team' 23 | 24 | 25 | # -- General configuration --------------------------------------------------- 26 | 27 | # Add any Sphinx extension module names here, as strings. They can be 28 | # extensions coming with Sphinx (named 'sphinx.ext.*') or your custom 29 | # ones. 30 | extensions = [ 31 | 'sphinx.ext.autodoc', 32 | 'sphinx.ext.autosummary', 33 | 'sphinx.ext.doctest', 34 | 'sphinx.ext.intersphinx', 35 | 'sphinx.ext.mathjax', 36 | 'sphinx.ext.ifconfig', 37 | 'numpydoc', 38 | 39 | ] 40 | 41 | 42 | # Add any paths that contain templates here, relative to this directory. 43 | templates_path = ['_templates'] 44 | 45 | # List of patterns, relative to source directory, that match files and 46 | # directories to ignore when looking for source files. 47 | # This pattern also affects html_static_path and html_extra_path. 48 | exclude_patterns = [] 49 | 50 | autosummary_generate = True 51 | # -- Options for HTML output ------------------------------------------------- 52 | 53 | # The theme to use for HTML and HTML Help pages. See the documentation for 54 | # a list of builtin themes. 55 | # 56 | html_theme = 'sphinx_rtd_theme' 57 | 58 | # Add any paths that contain custom static files (such as style sheets) here, 59 | # relative to this directory. They are copied after the builtin static files, 60 | # so a file named "default.css" will overwrite the builtin "default.css". 61 | # html_static_path = ['_static'] 62 | 63 | html_title = 'TACO Python API Reference' 64 | 65 | 66 | # -- Extension configuration ------------------------------------------------- 67 | 68 | 69 | numpydoc_attributes_as_param_list = True 70 | numpydoc_class_members_toctree = False 71 | numpydoc_show_class_members = False 72 | 73 | 74 | # -- Options for intersphinx extension --------------------------------------- 75 | 76 | # Example configuration for intersphinx: refer to the Python standard library. 77 | intersphinx_mapping = {'https://docs.python.org/': None} 78 | -------------------------------------------------------------------------------- /python_bindings/source/index.rst: -------------------------------------------------------------------------------- 1 | .. PyTaco documentation master file, created by 2 | sphinx-quickstart on Sat May 25 19:06:12 2019. 3 | You can adapt this file completely to your liking, but it should at least 4 | contain the root `toctree` directive. 5 | 6 | Table of Contents 7 | ================================== 8 | 9 | .. toctree:: 10 | :maxdepth: 3 11 | 12 | rst_files/datatype 13 | rst_files/format 14 | rst_files/file_io 15 | rst_files/sched 16 | rst_files/index_expressions 17 | rst_files/tensors 18 | rst_files/parsers 19 | rst_files/udfs 20 | 21 | 22 | 23 | 24 | 25 | 26 | -------------------------------------------------------------------------------- /python_bindings/source/rst_files/as_np_dt.rst: -------------------------------------------------------------------------------- 1 | Converting PyTaco Data Type to NumPy 2 | ====================================== 3 | 4 | .. autofunction:: pytaco.as_np_dtype 5 | -------------------------------------------------------------------------------- /python_bindings/source/rst_files/datatype.rst: -------------------------------------------------------------------------------- 1 | Data Types 2 | ======================== 3 | 4 | .. toctree:: 5 | :maxdepth: 2 6 | 7 | dtype_object 8 | as_np_dt 9 | 10 | 11 | 12 | 13 | 14 | 15 | 16 | -------------------------------------------------------------------------------- /python_bindings/source/rst_files/dtype_object.rst: -------------------------------------------------------------------------------- 1 | PyTaco Data Type Object 2 | ================================= 3 | 4 | .. autoclass:: pytaco.dtype 5 | -------------------------------------------------------------------------------- /python_bindings/source/rst_files/expr_funcs.rst: -------------------------------------------------------------------------------- 1 | .. _expr_funcs: 2 | 3 | 4 | Index Expression Functions 5 | ============================= 6 | 7 | .. currentmodule:: pytaco 8 | 9 | .. autosummary:: 10 | :toctree: functions 11 | 12 | cast 13 | remainder 14 | abs 15 | pow 16 | square 17 | cube 18 | sqrt 19 | cube_root 20 | exp 21 | log 22 | log10 23 | sin 24 | cos 25 | tan 26 | asin 27 | acos 28 | atan 29 | atan2 30 | sinh 31 | cosh 32 | tanh 33 | asinh 34 | acosh 35 | atanh 36 | logical_not 37 | gt 38 | lt 39 | ge 40 | le 41 | eq 42 | ne 43 | max 44 | min 45 | heaviside 46 | sum 47 | add 48 | mul 49 | sub 50 | div -------------------------------------------------------------------------------- /python_bindings/source/rst_files/file_io.rst: -------------------------------------------------------------------------------- 1 | .. _io: 2 | 3 | Tensor I/O 4 | ================================= 5 | 6 | PyTaco can read and write tensors natively from three different file formats. 7 | The formats PyTaco supports are: 8 | 9 | - `Matrix Market (Coordinate) Format `_ - .mtx 10 | - `Rutherford-Boeing Format `_ - .rb 11 | - `FROSTT Format `_ - .tns 12 | 13 | For both the read and write functions, the file format is inferred from the 14 | file extension. This means that the extension must be given when specifying a 15 | file name. 16 | 17 | .. currentmodule:: pytaco 18 | 19 | .. autosummary:: 20 | :toctree: functions 21 | 22 | read 23 | write 24 | from_array 25 | from_sp_csc 26 | from_sp_csr 27 | to_array 28 | to_sp_csc 29 | to_sp_csr 30 | as_tensor 31 | -------------------------------------------------------------------------------- /python_bindings/source/rst_files/format.rst: -------------------------------------------------------------------------------- 1 | Formats 2 | ======================= 3 | .. toctree:: 4 | :maxdepth: 3 5 | 6 | mode_format 7 | format_class 8 | format_funcs 9 | 10 | 11 | -------------------------------------------------------------------------------- /python_bindings/source/rst_files/format_class.rst: -------------------------------------------------------------------------------- 1 | Tensor Format Object 2 | ======================= 3 | .. autoclass:: pytaco.format -------------------------------------------------------------------------------- /python_bindings/source/rst_files/format_funcs.rst: -------------------------------------------------------------------------------- 1 | Format Functions 2 | =================== 3 | 4 | .. autofunction:: pytaco.is_dense -------------------------------------------------------------------------------- /python_bindings/source/rst_files/idx_exp_obj.rst: -------------------------------------------------------------------------------- 1 | Index Expression Objects 2 | ========================== 3 | 4 | .. autoclass:: pytaco.index_expression -------------------------------------------------------------------------------- /python_bindings/source/rst_files/index_expressions.rst: -------------------------------------------------------------------------------- 1 | .. _iexpr: 2 | 3 | Index Expressions 4 | ======================= 5 | 6 | Index expressions are used to express computations in taco. They are formed by indexing into :class:`pytaco.tensor` s with 7 | :class:`~pytaco.index_var` s and using the variety of arithmetic operations PyTaco provides for index vars. 8 | 9 | It should be noted that index expressions describe to taco how a computation should be performed. Thus, they must be 10 | assigned to a tensor before taco does any computations. 11 | 12 | The documentation in this section displays all the functions taco supports with index expressions as well as a wide 13 | number of examples for using them. 14 | 15 | Index expressions provide a convenient way for taco to construct outputs without making temporaries and are essential 16 | to getting leveraging taco's ability to fuse operations. 17 | 18 | There is an equivalent tensor function for each index expression function. The user should note that the tensor functions 19 | all use these under the hood. 20 | 21 | .. toctree:: 22 | :maxdepth: 3 23 | 24 | index_vars 25 | idx_exp_obj 26 | expr_funcs 27 | iv_funcs 28 | 29 | -------------------------------------------------------------------------------- /python_bindings/source/rst_files/index_vars.rst: -------------------------------------------------------------------------------- 1 | PyTaco Index Variables 2 | ======================= 3 | .. autoclass:: pytaco.index_var 4 | -------------------------------------------------------------------------------- /python_bindings/source/rst_files/iv_funcs.rst: -------------------------------------------------------------------------------- 1 | Convenience Functions 2 | ====================== 3 | 4 | .. currentmodule:: pytaco 5 | 6 | .. autosummary:: 7 | :toctree: functions 8 | 9 | get_index_vars -------------------------------------------------------------------------------- /python_bindings/source/rst_files/mode_format.rst: -------------------------------------------------------------------------------- 1 | Mode Formats 2 | ======================= 3 | .. autoclass:: pytaco.mode_format 4 | -------------------------------------------------------------------------------- /python_bindings/source/rst_files/parsers.rst: -------------------------------------------------------------------------------- 1 | Evaluating String Expressions 2 | ============================== 3 | 4 | .. currentmodule:: pytaco 5 | 6 | .. autosummary:: 7 | :toctree: functions 8 | 9 | evaluate 10 | einsum 11 | -------------------------------------------------------------------------------- /python_bindings/source/rst_files/sched.rst: -------------------------------------------------------------------------------- 1 | Scheduling Commands 2 | ====================== 3 | 4 | .. currentmodule:: pytaco 5 | 6 | .. autosummary:: 7 | :toctree: functions 8 | 9 | get_num_threads 10 | set_num_threads 11 | set_parallel_schedule 12 | get_parallel_schedule -------------------------------------------------------------------------------- /python_bindings/source/rst_files/tensor_class.rst: -------------------------------------------------------------------------------- 1 | Tensor Class 2 | ================= 3 | 4 | .. autoclass:: pytaco.tensor 5 | 6 | -------------------------------------------------------------------------------- /python_bindings/source/rst_files/tensor_info.rst: -------------------------------------------------------------------------------- 1 | .. toctree:: 2 | :maxdepth: 3 3 | 4 | tensor_class 5 | 6 | .. _tfuncs: 7 | 8 | Tensor Compute Functions 9 | ========================== 10 | 11 | .. currentmodule:: pytaco 12 | 13 | .. autosummary:: 14 | :toctree: functions 15 | 16 | tensor_abs 17 | tensor_square 18 | tensor_cube 19 | tensor_sqrt 20 | tensor_cube_root 21 | tensor_exp 22 | tensor_log 23 | tensor_log10 24 | tensor_sin 25 | tensor_cos 26 | tensor_tan 27 | tensor_asin 28 | tensor_acos 29 | tensor_atan 30 | tensor_atan2 31 | tensor_sinh 32 | tensor_cosh 33 | tensor_tanh 34 | tensor_asinh 35 | tensor_acosh 36 | tensor_atanh 37 | tensor_logical_not 38 | tensor_sum 39 | tensor_add 40 | tensor_mul 41 | tensor_div 42 | tensor_sub 43 | tensor_pow 44 | tensor_gt 45 | tensor_lt 46 | tensor_ge 47 | tensor_le 48 | tensor_eq 49 | tensor_ne 50 | tensor_max 51 | tensor_min 52 | tensor_heaviside 53 | as_type 54 | matmul 55 | inner 56 | dot 57 | outer 58 | tensordot 59 | 60 | Tensor Manipulation Functions 61 | =============================== 62 | .. currentmodule:: pytaco 63 | 64 | .. autosummary:: 65 | :toctree: functions 66 | 67 | tensor.transpose 68 | tensor.pack 69 | tensor.compile 70 | tensor.assemble 71 | tensor.compute 72 | tensor.evaluate 73 | tensor.to_dense 74 | tensor.to_array 75 | tensor.toarray 76 | tensor.to_sp_csr 77 | tensor.to_sp_csc 78 | tensor.copy 79 | tensor.insert 80 | tensor.remove_explicit_zeros 81 | remove_explicit_zeros 82 | -------------------------------------------------------------------------------- /python_bindings/source/rst_files/tensors.rst: -------------------------------------------------------------------------------- 1 | Tensors 2 | ========= 3 | 4 | .. toctree:: 5 | :maxdepth: 3 6 | 7 | tensor_info -------------------------------------------------------------------------------- /python_bindings/source/rst_files/udfs.rst: -------------------------------------------------------------------------------- 1 | User-Defined Functions 2 | ========================= 3 | 4 | Users can defined their own functions to execute with PyTaco. These functions 5 | must be implemented in a C header file and can be called using the following 6 | functions. 7 | 8 | .. currentmodule:: pytaco 9 | 10 | .. autosummary:: 11 | :toctree: functions 12 | 13 | apply 14 | set_udf_dir 15 | 16 | -------------------------------------------------------------------------------- /python_bindings/src/pyTensorIO.cpp: -------------------------------------------------------------------------------- 1 | #include "../include/pyTensorIO.h" 2 | #include "pybind11/stl.h" 3 | 4 | namespace taco{ 5 | namespace pythonBindings { 6 | 7 | template 8 | static Tensor tensorRead(std::string filename, T modeType, bool pack = true){ 9 | return Tensor(taco::read(filename, modeType, pack)); 10 | } 11 | 12 | void defineIOFuncs(py::module &m){ 13 | m.def("_read", tensorRead, py::arg("filename"), py::arg("format").noconvert(), 14 | py::arg("pack")=true); 15 | 16 | m.def("_read", tensorRead, py::arg("filename"), py::arg("modeType").noconvert(), 17 | py::arg("pack")=true); 18 | 19 | m.def("_write",[](std::string s, TensorBase& t) -> void { 20 | // force tensor evaluation 21 | t.pack(); 22 | if(t.needsCompute()) { 23 | t.evaluate(); 24 | } 25 | write(s, t); 26 | }, py::arg("filename"), py::arg("tensor").noconvert()); 27 | } 28 | 29 | }} 30 | -------------------------------------------------------------------------------- /src/CMakeLists.txt: -------------------------------------------------------------------------------- 1 | if (TACO_SHARED_LIBRARY) 2 | set(TACO_LIBRARY_TYPE SHARED) 3 | message("-- Shared library") 4 | else() 5 | set(TACO_LIBRARY_TYPE STATIC) 6 | message("-- Static library") 7 | endif() 8 | 9 | set(TACO_SRC_DIRS . parser index_notation lower ir codegen storage error util) 10 | 11 | foreach(dir ${TACO_SRC_DIRS}) 12 | file(GLOB TACO_HEADERS ${TACO_HEADERS} ${dir}/*.h) 13 | file(GLOB TACO_HEADERS ${TACO_HEADERS} ${TACO_INCLUDE_DIR}/taco/${dir}/*.h) 14 | file(GLOB TACO_SOURCES ${TACO_SOURCES} ${dir}/*.cpp) 15 | endforeach() 16 | 17 | set(TACO_HEADERS ${TACO_HEADERS} ../include/taco/ir_tags.h) 18 | set(TACO_SOURCES ${TACO_SOURCES} ir_tags.cpp) 19 | 20 | add_definitions(${TACO_DEFINITIONS}) 21 | include_directories(${TACO_SRC_DIR}) 22 | add_library(taco ${TACO_LIBRARY_TYPE} ${TACO_HEADERS} ${TACO_SOURCES}) 23 | if (CUDA) 24 | include_directories(${CUDA_INCLUDE_DIRS}) 25 | target_link_libraries(taco PUBLIC ${CUDA_LIBRARIES}) 26 | endif (CUDA) 27 | install(TARGETS taco DESTINATION lib) 28 | 29 | if (LINUX) 30 | target_link_libraries(taco PRIVATE ${TACO_LIBRARIES} dl) 31 | else() 32 | target_link_libraries(taco PRIVATE ${TACO_LIBRARIES}) 33 | endif() 34 | -------------------------------------------------------------------------------- /src/codegen/codegen.h: -------------------------------------------------------------------------------- 1 | #ifndef TACO_CODEGEN_H 2 | #define TACO_CODEGEN_H 3 | 4 | #include 5 | #include "taco/ir/ir.h" 6 | #include "taco/ir/ir_printer.h" 7 | 8 | namespace taco { 9 | namespace ir { 10 | 11 | 12 | class CodeGen : public IRPrinter { 13 | public: 14 | /// Kind of output: header or implementation 15 | enum OutputKind { HeaderGen, ImplementationGen }; 16 | enum CodeGenType { C, CUDA }; 17 | 18 | CodeGen(std::ostream& stream, CodeGenType type) : IRPrinter(stream), codeGenType(type) {}; 19 | CodeGen(std::ostream& stream, bool color, bool simplify, CodeGenType type) : IRPrinter(stream, color, simplify), codeGenType(type) {}; 20 | /// Initialize the default code generator 21 | static std::shared_ptr init_default(std::ostream &dest, OutputKind outputKind); 22 | 23 | /// Compile a lowered function 24 | virtual void compile(Stmt stmt, bool isFirst=false) =0; 25 | 26 | protected: 27 | static bool checkForAlloc(const Function *func); 28 | static int countYields(const Function *func); 29 | 30 | static std::string printCType(Datatype type, bool is_ptr); 31 | static std::string printCUDAType(Datatype type, bool is_ptr); 32 | 33 | static std::string printCAlloc(std::string pointer, std::string size); 34 | static std::string printCUDAAlloc(std::string pointer, std::string size); 35 | std::string printAlloc(std::string pointer, std::string size); 36 | 37 | static std::string printCFree(std::string pointer); 38 | static std::string printCUDAFree(std::string pointer); 39 | std::string printFree(std::string pointer); 40 | 41 | std::string printType(Datatype type, bool is_ptr); 42 | std::string printContextDeclAndInit(std::map varMap, 43 | std::vector localVars, int labels, 44 | std::string funcName); 45 | std::string printDecls(std::map varMap, 46 | std::vector inputs, std::vector outputs); 47 | std::string printPack(std::map, 48 | std::string> outputProperties, std::vector outputs); 49 | std::string printCoroutineFinish(int numYields, std::string funcName); 50 | void printYield(const Yield* op, std::vector localVars, 51 | std::map varMap, int labelCount, std::string funcName); 52 | std::string printFuncName(const Function *func, 53 | std::map inputMap={}, 54 | std::map outputMap={}); 55 | 56 | void resetUniqueNameCounters(); 57 | std::string genUniqueName(std::string name); 58 | void doIndentStream(std::stringstream &stream); 59 | CodeGenType codeGenType; 60 | 61 | private: 62 | virtual std::string restrictKeyword() const { return ""; } 63 | 64 | std::string printTensorProperty(std::string varname, const GetProperty* op, bool is_ptr); 65 | std::string unpackTensorProperty(std::string varname, const GetProperty* op, 66 | bool is_output_prop); 67 | std::string packTensorProperty(std::string varname, Expr tnsr, TensorProperty property, 68 | int mode, int index); 69 | std::string pointTensorProperty(std::string varname); 70 | }; 71 | 72 | 73 | 74 | 75 | } // namespace ir 76 | } // namespace taco 77 | #endif 78 | -------------------------------------------------------------------------------- /src/codegen/codegen_c.h: -------------------------------------------------------------------------------- 1 | #ifndef TACO_BACKEND_C_H 2 | #define TACO_BACKEND_C_H 3 | #include 4 | #include 5 | 6 | #include "taco/ir/ir.h" 7 | #include "taco/ir/ir_printer.h" 8 | #include "codegen.h" 9 | 10 | namespace taco { 11 | namespace ir { 12 | 13 | 14 | class CodeGen_C : public CodeGen { 15 | public: 16 | /// Initialize a code generator that generates code to an 17 | /// output stream. 18 | CodeGen_C(std::ostream &dest, OutputKind outputKind, bool simplify=true); 19 | ~CodeGen_C(); 20 | 21 | /// Compile a lowered function 22 | void compile(Stmt stmt, bool isFirst=false); 23 | 24 | /// Generate shims that unpack an array of pointers representing 25 | /// a mix of taco_tensor_t* and scalars into a function call 26 | static void generateShim(const Stmt& func, std::stringstream &stream); 27 | 28 | protected: 29 | using IRPrinter::visit; 30 | 31 | void visit(const Function*); 32 | void visit(const VarDecl*); 33 | void visit(const Yield*); 34 | void visit(const Var*); 35 | void visit(const For*); 36 | void visit(const While*); 37 | void visit(const GetProperty*); 38 | void visit(const Min*); 39 | void visit(const Max*); 40 | void visit(const Allocate*); 41 | void visit(const Sqrt*); 42 | void visit(const Store*); 43 | void visit(const Assign*); 44 | 45 | std::map varMap; 46 | std::vector localVars; 47 | std::ostream &out; 48 | 49 | OutputKind outputKind; 50 | 51 | std::string funcName; 52 | int labelCount; 53 | bool emittingCoroutine; 54 | 55 | class FindVars; 56 | 57 | private: 58 | virtual std::string restrictKeyword() const { return "restrict"; } 59 | }; 60 | 61 | } // namespace ir 62 | } // namespace taco 63 | #endif 64 | -------------------------------------------------------------------------------- /src/codegen/codegen_cuda.h: -------------------------------------------------------------------------------- 1 | #ifndef TACO_BACKEND_CUDA_H 2 | #define TACO_BACKEND_CUDA_H 3 | 4 | #include 5 | #include 6 | 7 | #include "taco/ir/ir.h" 8 | #include "taco/ir/ir_printer.h" 9 | #include "codegen.h" 10 | 11 | namespace taco { 12 | namespace ir { 13 | 14 | 15 | class CodeGen_CUDA : public CodeGen { 16 | public: 17 | /// Initialize a code generator that generates code to an 18 | /// output stream. 19 | CodeGen_CUDA(std::ostream &dest, OutputKind outputKind); 20 | ~CodeGen_CUDA(); 21 | 22 | /// Compile a lowered function 23 | void compile(Stmt stmt, bool isFirst=false); 24 | 25 | /// Generate shims that unpack an array of pointers representing 26 | /// a mix of taco_tensor_t* and scalars into a function call 27 | static void generateShim(const Stmt& func, std::stringstream &ret); 28 | protected: 29 | using IRPrinter::visit; 30 | void visit(const Function*); 31 | void visit(const Var*); 32 | void visit(const For*); 33 | void visit(const While*); 34 | void visit(const GetProperty*); 35 | void visit(const Min*); 36 | void visit(const Max*); 37 | void visit(const Allocate*); 38 | void visit(const Sqrt*); 39 | void visit(const Add*); 40 | void visit(const Sub*); 41 | void visit(const Mul*); 42 | void visit(const Div*); 43 | void visit(const VarDecl*); 44 | void visit(const Literal*); 45 | void visit(const Yield*); 46 | void visit(const Call*); 47 | void visit(const Store*); 48 | void visit(const Assign*); 49 | void visit(const Continue*); 50 | void visit(const Free* op); 51 | std::string printDeviceFuncName(const std::vector> currentParameters, int index); 52 | void printDeviceFuncCall(const std::vector> currentParameters, Expr blockSize, int index, Expr gridSize); 53 | void printThreadIDVariable(std::pair threadIDVar, Expr start, Expr increment, Expr numThreads); 54 | void printBlockIDVariable(std::pair blockIDVar, Expr start, Expr increment); 55 | void printWarpIDVariable(std::pair warpIDVar, Expr start, Expr increment, Expr warpSize); 56 | void printThreadBoundCheck(Expr end); 57 | void printDeviceFunctions(const Function* func); 58 | void printBinCastedOp(Expr a, Expr b, std::string op, Precedence precedence); 59 | Stmt simplifyFunctionBodies(Stmt stmt); 60 | 61 | bool isHostFunction=true; 62 | 63 | std::map varMap; 64 | std::vector localVars; 65 | 66 | std::vector>> deviceFunctionParameters; 67 | std::vector deviceFunctionBlockSizes; 68 | std::vector deviceFunctionGridSizes; 69 | std::vector deviceFunctions; // expressions to replace to calls of device function 70 | std::set scalarVarsPassedToDeviceFunction; // need to be allocated in uvm 71 | int deviceFunctionLoopDepth; 72 | std::set parentParallelUnits; 73 | std::map parallelUnitSizes; 74 | std::map parallelUnitIDVars; 75 | 76 | bool emittedTimerStartCode = false; 77 | 78 | std::ostream &out; 79 | 80 | OutputKind outputKind; 81 | 82 | std::string funcName; 83 | int labelCount; 84 | bool emittingCoroutine; 85 | 86 | class FindVars; 87 | class DeviceFunctionCollector; 88 | 89 | private: 90 | virtual std::string restrictKeyword() const { return "__restrict__"; } 91 | }; 92 | 93 | } // namespace ir 94 | } // namespace taco 95 | #endif 96 | -------------------------------------------------------------------------------- /src/cuda.cpp: -------------------------------------------------------------------------------- 1 | #include "taco/cuda.h" 2 | #include "taco/error.h" 3 | 4 | #if CUDA_BUILT 5 | #include 6 | #endif 7 | 8 | using namespace std; 9 | namespace taco { 10 | /// Functions used by taco to interface with CUDA (especially unified memory) 11 | static bool CUDA_codegen_enabled = CUDA_BUILT; 12 | static bool CUDA_unified_memory_enabled = CUDA_BUILT; 13 | bool should_use_CUDA_codegen() { 14 | return CUDA_codegen_enabled; 15 | } 16 | 17 | bool should_use_CUDA_unified_memory() { 18 | return CUDA_unified_memory_enabled; 19 | } 20 | 21 | void set_CUDA_codegen_enabled(bool enabled) { 22 | CUDA_codegen_enabled = enabled; 23 | } 24 | 25 | void set_CUDA_unified_memory_enabled(bool enabled) { 26 | taco_iassert(CUDA_BUILT); 27 | CUDA_unified_memory_enabled = enabled; 28 | } 29 | 30 | string get_default_CUDA_compiler_flags() { 31 | #if CUDA_BUILT 32 | cudaDeviceProp prop; 33 | cudaGetDeviceProperties(&prop, 0); 34 | string computeCap = to_string(prop.major) + to_string(prop.minor); 35 | return "-w -O3 -Xcompiler \"-fPIC -shared -ffast-math -O3\" --generate-code arch=compute_" + computeCap + ",code=sm_" + computeCap; 36 | #else 37 | taco_ierror; 38 | return ""; 39 | #endif 40 | } 41 | 42 | #if CUDA_BUILT 43 | // https://stackoverflow.com/questions/14038589/what-is-the-canonical-way-to-check-for-errors-using-the-cuda-runtime-api 44 | #define gpuErrchk(ans) { gpuAssert((ans), __FILE__, __LINE__); } 45 | inline void gpuAssert(cudaError_t code, const char *file, int line, bool abort=true) 46 | { 47 | if (code != cudaSuccess && code != cudaErrorCudartUnloading) // cudart unloading is normal behavior 48 | { 49 | taco_ierror << "GPUassert: " << code << " " << cudaGetErrorString(code) << " " << file << " " << line; 50 | } 51 | } 52 | #endif 53 | 54 | void* cuda_unified_alloc(size_t size) { 55 | #if CUDA_BUILT 56 | if (size == 0) return nullptr; 57 | void *ptr; 58 | gpuErrchk(cudaMallocManaged(&ptr, size, 1)); 59 | return ptr; 60 | #else 61 | taco_ierror; 62 | return nullptr; 63 | #endif 64 | } 65 | 66 | void cuda_unified_free(void *ptr) { 67 | #if CUDA_BUILT 68 | gpuErrchk(cudaFree(ptr)); 69 | #else 70 | taco_ierror; 71 | #endif 72 | } 73 | } 74 | -------------------------------------------------------------------------------- /src/error.cpp: -------------------------------------------------------------------------------- 1 | #include "taco/error.h" 2 | 3 | #include 4 | #include 5 | 6 | using namespace std; 7 | 8 | namespace taco { 9 | 10 | TacoException::TacoException(std::string msg) : message(msg) {} 11 | 12 | const char* TacoException::what() const noexcept { 13 | return message.c_str(); 14 | } 15 | 16 | 17 | ErrorReport::ErrorReport(const char *file, const char *func, int line, 18 | bool condition, const char *conditionString, 19 | Kind kind, bool warning) 20 | : msg(NULL), file(file), func(func), line(line), condition(condition), 21 | conditionString(conditionString), kind(kind), warning(warning) { 22 | if (condition) { 23 | return; 24 | } 25 | msg = new std::ostringstream; 26 | 27 | switch (kind) { 28 | case User: 29 | if (warning) { 30 | (*msg) << "Warning"; 31 | } else { 32 | (*msg) << "Error"; 33 | } 34 | (*msg) << " at " << file << ":" << line << " in " << func << ":" << endl; 35 | break; 36 | case Internal: 37 | (*msg) << "Compiler bug"; 38 | if (warning) { 39 | (*msg) << "(warning)"; 40 | } 41 | (*msg) << " at " << file << ":" << line << " in " << func; 42 | (*msg) << endl << "Please report it to developers"; 43 | 44 | if (conditionString) { 45 | (*msg) << endl << " Condition failed: " << conditionString; 46 | } 47 | (*msg) << endl; 48 | break; 49 | case Temporary: 50 | (*msg) << "Temporary assumption broken"; 51 | (*msg) << " at " << file << ":" << line << endl; 52 | (*msg) << " Not supported yet, but planned for the future"; 53 | if (conditionString) { 54 | (*msg) << endl << " Condition failed: " << conditionString; 55 | } 56 | (*msg) << endl; 57 | break; 58 | } 59 | (*msg) << " "; 60 | } 61 | 62 | void ErrorReport::explodeWithException() { 63 | TacoException e = TacoException(msg->str()); 64 | delete msg; 65 | throw e; 66 | } 67 | 68 | } 69 | -------------------------------------------------------------------------------- /src/error/error_checks.h: -------------------------------------------------------------------------------- 1 | #ifndef TACO_ERROR_CHECKS_H 2 | #define TACO_ERROR_CHECKS_H 3 | 4 | #include 5 | #include 6 | #include 7 | 8 | namespace taco { 9 | class IndexVar; 10 | class IndexExpr; 11 | class Format; 12 | class Shape; 13 | 14 | namespace error { 15 | 16 | /// Check whether all dimensions indexed by the same variable are the same. 17 | /// If they are not, then the first element of the returned tuple will be false, 18 | /// and a human readable error will be returned in the second component. 19 | std::pair dimensionsTypecheck(const std::vector& resultVars, 20 | const IndexExpr& expr, 21 | const Shape& shape); 22 | 23 | /// Returns true iff the index expression contains a transposition. 24 | bool containsTranspose(const Format& resultFormat, 25 | const std::vector& resultVars, 26 | const IndexExpr& expr); 27 | 28 | }} 29 | #endif 30 | -------------------------------------------------------------------------------- /src/error/error_messages.cpp: -------------------------------------------------------------------------------- 1 | #include "taco/error/error_messages.h" 2 | 3 | namespace taco { 4 | namespace error { 5 | 6 | const std::string type_mismatch = 7 | "Type mismatch."; 8 | 9 | const std::string type_bitwidt = 10 | "The given bit width is not supported for this type."; 11 | 12 | const std::string expr_dimension_mismatch = 13 | "Dimension size mismatch."; 14 | 15 | const std::string expr_transposition = 16 | "Computations with transpositions are not supported, but are planned for the " 17 | "future."; 18 | 19 | const std::string expr_distribution = 20 | "Expressions with free variables that do not appear on the right hand side " 21 | "of the expression are not supported, but are planned for the future"; 22 | 23 | const std::string expr_einsum_missformed = 24 | "Summations/reductions are not specified and the Einstein summation " 25 | "convention cannot be applied, because the expression is not of the form " 26 | "a*...*b + ... + c*...*d."; 27 | 28 | const std::string compile_without_expr = 29 | "The tensor must be assigned to before compile is called."; 30 | 31 | const std::string compile_tensor_name_collision = 32 | "Tensor name collision."; 33 | 34 | const std::string assemble_without_compile = 35 | "The compile method must be called before assemble."; 36 | 37 | const std::string compute_without_compile = 38 | "The compile method must be called before compute."; 39 | 40 | const std::string requires_matrix = 41 | "The argument must be a matrix."; 42 | 43 | }} 44 | -------------------------------------------------------------------------------- /src/index_notation/README.md: -------------------------------------------------------------------------------- 1 | This directory contains the API and data structures for the index 2 | notation language that describes computations on tensors. It also 3 | contains utilities such as visitors, rewriters, and printers. For 4 | more information about the index notation API, see the [include 5 | directory](https://github.com/tensor-compiler/taco/tree/master/include/taco/index_notation) 6 | -------------------------------------------------------------------------------- /src/index_notation/index_notation_nodes.cpp: -------------------------------------------------------------------------------- 1 | #include "taco/index_notation/index_notation_nodes.h" 2 | #include "taco/index_notation/intrinsic.h" 3 | 4 | #include 5 | #include 6 | 7 | #include "taco/util/collections.h" 8 | 9 | using namespace std; 10 | 11 | namespace taco { 12 | 13 | template 14 | static std::vector getDataTypes(const std::vector args) { 15 | std::function getType = [](T arg) { return arg.getDataType(); }; 16 | return util::map(args, getType); 17 | } 18 | 19 | 20 | // class CastNode 21 | CastNode::CastNode(IndexExpr a, Datatype newType) 22 | : IndexExprNode(newType), a(a) { 23 | } 24 | 25 | // class CallIntrinsicNode 26 | CallIntrinsicNode::CallIntrinsicNode(const std::shared_ptr& func, 27 | const std::vector& args) 28 | : IndexExprNode(func->inferReturnType(getDataTypes(args))), 29 | func(func), args(args) { 30 | } 31 | 32 | // class CallNode 33 | CallNode::CallNode(std::string name, const std::vector& args, OpImpl defaultLowerFunc, 34 | const IterationAlgebra &iterAlg, const std::vector &properties, 35 | const std::map, OpImpl>& regionDefinitions) 36 | : CallNode(name, args, defaultLowerFunc, iterAlg, properties, regionDefinitions, definedIndices(args)){ 37 | } 38 | 39 | // class CallNode 40 | CallNode::CallNode(std::string name, const std::vector& args, OpImpl defaultLowerFunc, 41 | const IterationAlgebra &iterAlg, const std::vector &properties, 42 | const std::map, OpImpl>& regionDefinitions, 43 | const std::vector& definedRegions) 44 | : IndexExprNode(inferReturnType(defaultLowerFunc, args)), name(name), args(args), defaultLowerFunc(defaultLowerFunc), 45 | iterAlg(applyDemorgan(iterAlg)), properties(properties), regionDefinitions(regionDefinitions), 46 | definedRegions(definedRegions) { 47 | 48 | taco_iassert(defaultLowerFunc != nullptr); 49 | for (const auto& pair: regionDefinitions) { 50 | taco_iassert(args.size() >= pair.first.size()); 51 | } 52 | } 53 | 54 | // class ReductionNode 55 | ReductionNode::ReductionNode(IndexExpr op, IndexVar var, IndexExpr a) 56 | : IndexExprNode(a.getDataType()), op(op), var(var), a(a) { 57 | taco_iassert(isa(op.ptr) || isa(op.ptr)); 58 | } 59 | 60 | IndexVarNode::IndexVarNode(const std::string& name, const Datatype& type) 61 | : IndexExprNode(type), content(new Content) { 62 | 63 | if (!type.isInt() && !type.isUInt()) { 64 | taco_not_supported_yet << ". IndexVars must be integral type."; 65 | } 66 | 67 | content->name = name; 68 | } 69 | 70 | std::string IndexVarNode::getName() const { 71 | return content->name; 72 | } 73 | 74 | bool operator==(const IndexVarNode& a, const IndexVarNode& b) { 75 | return a.content->name == b.content->name; 76 | } 77 | 78 | bool operator<(const IndexVarNode& a, const IndexVarNode& b) { 79 | return a.content->name < b.content->name; 80 | } 81 | 82 | } 83 | -------------------------------------------------------------------------------- /src/index_notation/index_notation_nodes_abstract.cpp: -------------------------------------------------------------------------------- 1 | #include "taco/index_notation/index_notation_nodes_abstract.h" 2 | 3 | #include "taco/index_notation/index_notation.h" 4 | #include "taco/index_notation/schedule.h" 5 | #include "taco/index_notation/transformations.h" 6 | 7 | #include 8 | 9 | using namespace std; 10 | 11 | namespace taco { 12 | 13 | // class ExprNode 14 | IndexExprNode::IndexExprNode() : workspace(nullptr) { 15 | } 16 | 17 | IndexExprNode::IndexExprNode(Datatype type) 18 | : dataType(type), workspace(nullptr) { 19 | } 20 | 21 | Datatype IndexExprNode::getDataType() const { 22 | return dataType; 23 | } 24 | 25 | void IndexExprNode::setWorkspace(IndexVar i, IndexVar iw, 26 | TensorVar workspace) const { 27 | this->workspace = 28 | make_shared>(i,iw,workspace); 29 | } 30 | 31 | Precompute IndexExprNode::getWorkspace() const { 32 | if (workspace == nullptr) { 33 | return Precompute(); 34 | } 35 | return Precompute(this, get<0>(*workspace), get<1>(*workspace), 36 | get<2>(*workspace)); 37 | } 38 | 39 | 40 | // class TensorExprNode 41 | IndexStmtNode::IndexStmtNode() { 42 | } 43 | 44 | IndexStmtNode::IndexStmtNode(Type type) : type (type) { 45 | } 46 | 47 | } 48 | -------------------------------------------------------------------------------- /src/index_notation/iteration_algebra_printer.cpp: -------------------------------------------------------------------------------- 1 | #include "taco/index_notation/iteration_algebra_printer.h" 2 | 3 | namespace taco { 4 | 5 | // Iteration Algebra Printer 6 | IterationAlgebraPrinter::IterationAlgebraPrinter(std::ostream& os) : os(os) {} 7 | 8 | void IterationAlgebraPrinter::print(const IterationAlgebra& alg) { 9 | parentPrecedence = Precedence::TOP; 10 | alg.accept(this); 11 | } 12 | 13 | void IterationAlgebraPrinter::visit(const RegionNode* n) { 14 | os << n->expr(); 15 | } 16 | 17 | void IterationAlgebraPrinter::visit(const ComplementNode* n) { 18 | Precedence precedence = Precedence::COMPLEMENT; 19 | bool parenthesize = precedence > parentPrecedence; 20 | parentPrecedence = precedence; 21 | os << "~"; 22 | if (parenthesize) { 23 | os << "("; 24 | } 25 | n->a.accept(this); 26 | if (parenthesize) { 27 | os << ")"; 28 | } 29 | } 30 | 31 | void IterationAlgebraPrinter::visit(const IntersectNode* n) { 32 | visitBinary(n, Precedence::INTERSECT); 33 | } 34 | 35 | void IterationAlgebraPrinter::visit(const UnionNode* n) { 36 | visitBinary(n, Precedence::UNION); 37 | } 38 | 39 | template 40 | void IterationAlgebraPrinter::visitBinary(Node n, Precedence precedence) { 41 | bool parenthesize = precedence > parentPrecedence; 42 | if (parenthesize) { 43 | os << "("; 44 | } 45 | parentPrecedence = precedence; 46 | n->a.accept(this); 47 | os << " " << n->algebraString() << " "; 48 | parentPrecedence = precedence; 49 | n->b.accept(this); 50 | if (parenthesize) { 51 | os << ")"; 52 | } 53 | } 54 | 55 | } -------------------------------------------------------------------------------- /src/index_notation/schedule.cpp: -------------------------------------------------------------------------------- 1 | #include "taco/index_notation/schedule.h" 2 | 3 | #include 4 | 5 | #include "taco/index_notation/index_notation.h" 6 | #include "taco/index_notation/transformations.h" 7 | #include "taco/util/collections.h" 8 | #include "taco/util/strings.h" 9 | #include "taco/error/error_messages.h" 10 | 11 | using namespace std; 12 | 13 | namespace taco { 14 | 15 | // class Schedule 16 | struct Schedule::Content { 17 | map precomputes; 18 | }; 19 | 20 | Schedule::Schedule() : content(new Content) { 21 | } 22 | 23 | std::vector Schedule::getPrecomputes() const { 24 | vector workspaces; 25 | for (auto& workspace : content->precomputes) { 26 | workspaces.push_back(workspace.second); 27 | } 28 | return workspaces; 29 | } 30 | 31 | Precompute Schedule::getPrecompute(IndexExpr expr) const { 32 | if (!util::contains(content->precomputes, expr)) { 33 | return Precompute(); 34 | } 35 | return content->precomputes.at(expr); 36 | } 37 | 38 | void Schedule::addPrecompute(Precompute workspace) { 39 | if (!util::contains(content->precomputes, workspace.getExpr())) { 40 | content->precomputes.insert({workspace.getExpr(), workspace}); 41 | } 42 | else { 43 | content->precomputes.at(workspace.getExpr()) = workspace; 44 | } 45 | } 46 | 47 | void Schedule::clearPrecomputes() { 48 | content->precomputes.clear(); 49 | } 50 | 51 | std::ostream& operator<<(std::ostream& os, const Schedule& schedule) { 52 | auto workspaces = schedule.getPrecomputes(); 53 | if (workspaces.size() > 0) { 54 | os << "Workspace Commands:" << endl << util::join(workspaces, "\n"); 55 | } 56 | return os; 57 | } 58 | 59 | } 60 | -------------------------------------------------------------------------------- /src/ir/README.md: -------------------------------------------------------------------------------- 1 | The ir module describes a low-level C-like imperative and immutable intermediate representation (IR), including tools such as visitors, printers and rewriters. This low-level IR is the input to the code generators in the `codegen` module. 2 | -------------------------------------------------------------------------------- /src/ir/ir_generators.cpp: -------------------------------------------------------------------------------- 1 | #include "taco/ir/ir_generators.h" 2 | 3 | #include "taco/ir/ir.h" 4 | #include "taco/error.h" 5 | #include "taco/util/strings.h" 6 | 7 | namespace taco { 8 | namespace ir { 9 | 10 | Stmt compoundStore(Expr a, Expr i, Expr val, bool use_atomics, ParallelUnit atomic_parallel_unit) { 11 | Expr add = (val.type().getKind() == Datatype::Bool) 12 | ? Or::make(Load::make(a, i), val) 13 | : Add::make(Load::make(a, i), val); 14 | return Store::make(a, i, add, use_atomics, atomic_parallel_unit); 15 | } 16 | 17 | Stmt compoundAssign(Expr a, Expr val, bool use_atomics, ParallelUnit atomic_parallel_unit) { 18 | Expr add = (val.type().getKind() == Datatype::Bool) 19 | ? Or::make(a, val) : Add::make(a, val); 20 | return Assign::make(a, add, use_atomics, atomic_parallel_unit); 21 | } 22 | 23 | Expr conjunction(std::vector exprs) { 24 | taco_iassert(exprs.size() > 0) << "No expressions to and"; 25 | Expr conjunction = exprs[0]; 26 | for (size_t i = 1; i < exprs.size(); i++) { 27 | conjunction = And::make(conjunction, exprs[i]); 28 | } 29 | return conjunction; 30 | } 31 | 32 | Stmt doubleSizeIfFull(Expr a, Expr size, Expr needed) { 33 | Stmt realloc = Allocate::make(a, Mul::make(size, 2), true, size); 34 | Stmt resize = Assign::make(size, Mul::make(size, 2)); 35 | Stmt ifBody = Block::make({realloc, resize}); 36 | return IfThenElse::make(Lte::make(size, needed), ifBody); 37 | } 38 | 39 | Stmt atLeastDoubleSizeIfFull(Expr a, Expr size, Expr needed) { 40 | Expr newSizeVar = Var::make(util::toString(a) + "_new_size", Int()); 41 | Expr newSize = Max::make(Mul::make(size, 2), Add::make(needed, 1)); 42 | Stmt computeNewSize = VarDecl::make(newSizeVar, newSize); 43 | Stmt realloc = Allocate::make(a, newSizeVar, true, size); 44 | Stmt updateSize = Assign::make(size, newSizeVar); 45 | Stmt ifBody = Block::make({computeNewSize, realloc, updateSize}); 46 | return IfThenElse::make(Lte::make(size, needed), ifBody); 47 | } 48 | 49 | }} 50 | -------------------------------------------------------------------------------- /src/ir/workspace_rewriter.cpp: -------------------------------------------------------------------------------- 1 | #include "taco/ir/workspace_rewriter.h" 2 | 3 | #include 4 | #include 5 | 6 | #include "taco/ir/ir.h" 7 | #include "taco/ir/ir_rewriter.h" 8 | #include "taco/util/collections.h" 9 | #include "taco/index_notation/index_notation.h" 10 | #include "taco/index_notation/index_notation_nodes.h" 11 | #include "taco/index_notation/index_notation_visitor.h" 12 | 13 | using namespace std; 14 | using namespace taco::ir; 15 | namespace taco { 16 | 17 | struct WorkspaceRewriter : ir::IRRewriter { 18 | WorkspaceRewriter(std::vector whereTemps, std::map> temporarySizeMap) : whereTemps(whereTemps), 20 | temporarySizeMap(temporarySizeMap) {} 21 | std::vector whereTemps; 22 | std::map> temporarySizeMap; 23 | 24 | using IRRewriter::visit; 25 | void visit(const ir::GetProperty* op) { 26 | Expr tensor = rewrite(op->tensor); 27 | 28 | if (op->property == TensorProperty::Dimension && !whereTemps.empty()) { 29 | for (auto& temp : whereTemps) { 30 | string gpName = temp.getName() + to_string(op->mode + 1) + "_dimension"; 31 | 32 | if (temp.defined() && gpName == op->name && (temporarySizeMap.find(temp) != temporarySizeMap.end())) { 33 | //taco_iassert(temporarySizeMap.find(temp) != temporarySizeMap.end()) << "Cannot rewrite workspace " 34 | // "Dimension GetProperty due " 35 | // "to tensorVar not in " 36 | // "expression map"; 37 | auto tempExprList = temporarySizeMap.at(temp); 38 | 39 | taco_iassert((int)tempExprList.size() > op->mode) << "Cannot rewrite workspace (" 40 | << op->tensor 41 | << ") Dimension GetProperty due to mode (" 42 | << op->mode 43 | << ") not in expression map (size = " 44 | << tempExprList.size() << ")"; 45 | expr = tempExprList.at(op->mode); 46 | return; 47 | } 48 | } 49 | } 50 | expr = op; 51 | } 52 | 53 | void visit(const VarDecl* decl) { 54 | Expr rhs = rewrite(decl->rhs); 55 | stmt = (rhs == decl->rhs) ? decl : VarDecl::make(decl->var, rhs); 56 | } 57 | }; 58 | 59 | ir::Stmt rewriteTemporaryGP(const ir::Stmt& stmt, std::vector whereTemps, 60 | std::map> temporarySizeMap) { 61 | return WorkspaceRewriter(whereTemps, temporarySizeMap).rewrite(stmt); 62 | } 63 | 64 | } 65 | -------------------------------------------------------------------------------- /src/ir_tags.cpp: -------------------------------------------------------------------------------- 1 | #include "taco/ir_tags.h" 2 | 3 | namespace taco { 4 | 5 | const char *ParallelUnit_NAMES[] = {"NotParallel", "DefaultUnit", "GPUBlock", "GPUWarp", "GPUThread", "CPUThread", "CPUVector", "CPUThreadGroupReduction", "GPUBlockReduction", "GPUWarpReduction"}; 6 | const char *OutputRaceStrategy_NAMES[] = {"IgnoreRaces", "NoRaces", "Atomics", "Temporary", "ParallelReduction"}; 7 | const char *BoundType_NAMES[] = {"MinExact", "MinConstraint", "MaxExact", "MaxConstraint"}; 8 | const char *AssembleStrategy_NAMES[] = {"Append", "Insert"}; 9 | const char *MergeStrategy_NAMES[] = {"TwoFinger", "Gallop"}; 10 | 11 | } 12 | -------------------------------------------------------------------------------- /src/lower/README.md: -------------------------------------------------------------------------------- 1 | The `lower` machinery turns index expressions (`expr`) into low-level C-like imperative IR (`ir`). The lowering concepts and algorithms are described in the paper [The Tensor Algebra Compiler](http://tensor-compiler.org/kjolstad-oopsla17-tensor-compiler.pdf). 2 | -------------------------------------------------------------------------------- /src/lower/expr_tools.h: -------------------------------------------------------------------------------- 1 | #ifndef TACO_AVAILABLE_EXPRS_H 2 | #define TACO_AVAILABLE_EXPRS_H 3 | 4 | #include 5 | 6 | namespace taco { 7 | class IndexVar; 8 | class IndexExpr; 9 | 10 | /// Retrieves available sub-expression, which are the maximal sub-expressions 11 | /// whose operands are only indexed by the given index variables. 12 | std::vector 13 | getAvailableExpressions(const IndexExpr& expr, 14 | const std::vector& vars); 15 | 16 | /// Retrieves the minimal sub-expression that covers all the index variables 17 | IndexExpr getSubExpr(IndexExpr expr, const std::vector& vars); 18 | 19 | /// Retrieves the minimal sub-expression that covers all the index variables 20 | /// DEPRECATED: This is a deprecated function to keep functionality while 21 | /// redesigning the index expression IR with reduction nodes. 22 | IndexExpr getSubExprOld(IndexExpr expr, const std::vector& vars); 23 | 24 | } 25 | #endif 26 | -------------------------------------------------------------------------------- /src/lower/iteration_forest.h: -------------------------------------------------------------------------------- 1 | #ifndef TACO_ITERATION_FOREST_H 2 | #define TACO_ITERATION_FOREST_H 3 | 4 | #include 5 | #include 6 | 7 | #include "taco/index_notation/index_notation.h" 8 | #include "tensor_path.h" 9 | 10 | namespace taco { 11 | 12 | /// An iteration forest is a tree decomposition of a tensor path graph, where 13 | /// all tensor path edges move from an index variable higher in the tree to one 14 | /// strictly lower in the tree. 15 | class IterationForest { 16 | public: 17 | IterationForest() {} 18 | 19 | IterationForest(const std::vector& paths); 20 | 21 | const std::vector& getRoots() const {return roots;} 22 | 23 | bool hasParent(const IndexVar&) const; 24 | 25 | const IndexVar& getParent(const IndexVar&) const; 26 | 27 | const std::vector& getChildren(const IndexVar&) const; 28 | 29 | std::vector getNodes() const; 30 | 31 | friend std::ostream& operator<<(std::ostream&,const IterationForest&); 32 | 33 | private: 34 | std::vector roots; 35 | std::map> children; 36 | std::map parents; 37 | }; 38 | 39 | } 40 | #endif 41 | -------------------------------------------------------------------------------- /src/lower/iteration_graph.h: -------------------------------------------------------------------------------- 1 | #ifndef TACO_ITERATION_GRAPH_H 2 | #define TACO_ITERATION_GRAPH_H 3 | 4 | #include 5 | #include 6 | 7 | namespace taco { 8 | 9 | class TensorVar; 10 | class IndexVar; 11 | class IndexExpr; 12 | class Assignment; 13 | 14 | class TensorPath; 15 | 16 | enum class IndexVarType { 17 | Free, 18 | Sum 19 | }; 20 | 21 | /// An iteration graph constist of index variables arranged in a forest 22 | /// together with tensor paths super-imposed on the forest. 23 | /// - The iteration graph is arranged in a forest decomposition where all 24 | /// tensor paths move from index variables higher in the tree to index 25 | /// variables strictly lower in the tree. 26 | /// - The tensor paths describe how to iterate over the index variables through 27 | /// the indices of the corresponding (sparse or dense) tensors. 28 | class IterationGraph { 29 | public: 30 | IterationGraph(); 31 | 32 | /// Creates an iteration graph for a tensor with a defined expression. 33 | static IterationGraph make(Assignment); 34 | 35 | /// Returns the iteration graph roots; the index variables with no parents. 36 | const std::vector& getRoots() const; 37 | 38 | /// Returns the children of the index variable 39 | const std::vector& getChildren(const IndexVar&) const; 40 | 41 | /// Returns the parent of the index variable 42 | const IndexVar& getParent(const IndexVar&) const; 43 | 44 | /// Returns the ancestors of the index variable including itself. 45 | std::vector getAncestors(const IndexVar&) const; 46 | 47 | /// Returns all descendant of the index variable, including itself. 48 | std::vector getDescendants(const IndexVar&) const; 49 | 50 | 51 | /// Returns the tensor paths of the operand tensors in the iteration graph. 52 | const std::vector& getTensorPaths() const; 53 | 54 | /// Returns the tensor path corresponding to a tensor read expression. 55 | const TensorPath& getTensorPath(const IndexExpr&) const; 56 | 57 | /// Returns the tensor path of the result tensor. 58 | const TensorPath& getResultTensorPath() const; 59 | 60 | 61 | /// Returns the index variable type. 62 | IndexVarType getIndexVarType(const IndexVar&) const; 63 | 64 | /// Returns true iff the index variable is free. 65 | bool isFree(const IndexVar&) const; 66 | 67 | /// Returns true iff the index variable is a reduction. 68 | bool isReduction(const IndexVar&) const; 69 | 70 | /// Returns true if the index variable is the only free var in its subtree. 71 | bool isLastFreeVariable(const IndexVar&) const; 72 | 73 | /// Returns true if the index variable is the ancestor of any free variable. 74 | bool hasFreeVariableDescendant(const IndexVar&) const; 75 | 76 | /// Returns true if the index variable has a reduction variable ancestor. 77 | bool hasReductionVariableAncestor(const IndexVar&) const; 78 | 79 | 80 | /// Returns the index expression at the given index variable. 81 | const IndexExpr& getIndexExpr(const IndexVar&) const; 82 | 83 | 84 | /// Print an iteration graph as a dot file. 85 | void printAsDot(std::ostream&); 86 | 87 | /// Print an iteration graph. 88 | friend std::ostream& operator<<(std::ostream&, const IterationGraph&); 89 | 90 | private: 91 | struct Content; 92 | std::shared_ptr content; 93 | }; 94 | 95 | } 96 | #endif 97 | -------------------------------------------------------------------------------- /src/lower/lower.cpp: -------------------------------------------------------------------------------- 1 | #include "taco/lower/lower.h" 2 | 3 | #include 4 | #include 5 | #include 6 | #include 7 | #include 8 | 9 | #include "taco/index_notation/index_notation.h" 10 | #include "taco/index_notation/index_notation_nodes.h" 11 | 12 | #include "taco/ir/ir.h" 13 | #include "taco/ir/simplify.h" 14 | #include "taco/ir/ir_generators.h" 15 | #include "taco/ir/ir_printer.h" 16 | 17 | #include "taco/lower/lowerer_impl.h" 18 | #include "taco/lower/lowerer_impl_imperative.h" 19 | #include "taco/lower/iterator.h" 20 | #include "mode_access.h" 21 | 22 | #include "error/error_checks.h" 23 | #include "taco/error/error_messages.h" 24 | #include "taco/util/name_generator.h" 25 | #include "taco/util/collections.h" 26 | #include "taco/util/strings.h" 27 | 28 | #include "taco/ir/ir_verifier.h" 29 | 30 | using namespace std; 31 | using namespace taco::ir; 32 | 33 | namespace taco { 34 | 35 | 36 | // class Lowerer 37 | Lowerer::Lowerer() : impl(new LowererImplImperative()) { 38 | } 39 | 40 | Lowerer::Lowerer(LowererImpl* impl) : impl(impl) { 41 | } 42 | 43 | std::shared_ptr Lowerer::getLowererImpl() { 44 | return impl; 45 | } 46 | 47 | ir::Stmt lower(IndexStmt stmt, std::string name, 48 | bool assemble, bool compute, bool pack, bool unpack, 49 | Lowerer lowerer) { 50 | string reason; 51 | taco_iassert(isLowerable(stmt, &reason)) 52 | << "Not lowerable, because " << reason << ": " << stmt; 53 | 54 | ir::Stmt lowered = lowerer.getLowererImpl()->lower(stmt, name, assemble, compute, pack, unpack); 55 | 56 | // TODO: re-enable this 57 | // std::string messages; 58 | // verify(lowered, &messages); 59 | // if (!messages.empty()) { 60 | // std::cerr << "Verifier messages:\n" << messages << "\n"; 61 | // } 62 | 63 | return lowered; 64 | } 65 | 66 | 67 | bool isLowerable(IndexStmt stmt, std::string* reason) { 68 | INIT_REASON(reason); 69 | 70 | // Must be concrete index notation 71 | string r; 72 | if (!isConcreteNotation(stmt, &r)) { 73 | *reason = "the index statement is not in concrete index notation, because " 74 | + r; 75 | return false; 76 | } 77 | 78 | // Check for transpositions 79 | // if (!error::containsTranspose(this->getFormat(), freeVars, indexExpr)) { 80 | // *reason = error::expr_transposition; 81 | // } 82 | 83 | return true; 84 | } 85 | 86 | } 87 | -------------------------------------------------------------------------------- /src/lower/lowerer_impl.cpp: -------------------------------------------------------------------------------- 1 | #include 2 | #include "taco/lower/lowerer_impl.h" 3 | 4 | #include "taco/index_notation/index_notation.h" 5 | #include "taco/index_notation/index_notation_nodes.h" 6 | #include "taco/index_notation/index_notation_visitor.h" 7 | #include "taco/index_notation/index_notation_rewriter.h" 8 | #include "taco/index_notation/provenance_graph.h" 9 | #include "taco/ir/ir.h" 10 | #include "taco/ir/ir_generators.h" 11 | #include "taco/ir/ir_visitor.h" 12 | #include "taco/ir/simplify.h" 13 | #include "taco/lower/iterator.h" 14 | #include "taco/lower/merge_lattice.h" 15 | #include "mode_access.h" 16 | #include "taco/util/collections.h" 17 | #include "taco/ir/workspace_rewriter.h" 18 | 19 | using namespace std; 20 | using namespace taco::ir; 21 | using taco::util::combine; 22 | 23 | namespace taco { 24 | 25 | class LowererImpl::Visitor : public IndexNotationVisitorStrict { 26 | public: 27 | Visitor(LowererImpl* impl) : impl(impl) {} 28 | Stmt lower(IndexStmt stmt) { 29 | this->stmt = Stmt(); 30 | IndexStmtVisitorStrict::visit(stmt); 31 | return this->stmt; 32 | } 33 | Expr lower(IndexExpr expr) { 34 | this->expr = Expr(); 35 | IndexExprVisitorStrict::visit(expr); 36 | return this->expr; 37 | } 38 | private: 39 | LowererImpl* impl; 40 | Expr expr; 41 | Stmt stmt; 42 | using IndexNotationVisitorStrict::visit; 43 | void visit(const AssignmentNode* node) { stmt = impl->lowerAssignment(node); } 44 | void visit(const YieldNode* node) { stmt = impl->lowerYield(node); } 45 | void visit(const ForallNode* node) { stmt = impl->lowerForall(node); } 46 | void visit(const WhereNode* node) { stmt = impl->lowerWhere(node); } 47 | void visit(const MultiNode* node) { stmt = impl->lowerMulti(node); } 48 | void visit(const SuchThatNode* node) { stmt = impl->lowerSuchThat(node); } 49 | void visit(const SequenceNode* node) { stmt = impl->lowerSequence(node); } 50 | void visit(const AssembleNode* node) { stmt = impl->lowerAssemble(node); } 51 | void visit(const AccessNode* node) { expr = impl->lowerAccess(node); } 52 | void visit(const LiteralNode* node) { expr = impl->lowerLiteral(node); } 53 | void visit(const NegNode* node) { expr = impl->lowerNeg(node); } 54 | void visit(const AddNode* node) { expr = impl->lowerAdd(node); } 55 | void visit(const SubNode* node) { expr = impl->lowerSub(node); } 56 | void visit(const MulNode* node) { expr = impl->lowerMul(node); } 57 | void visit(const DivNode* node) { expr = impl->lowerDiv(node); } 58 | void visit(const SqrtNode* node) { expr = impl->lowerSqrt(node); } 59 | void visit(const CastNode* node) { expr = impl->lowerCast(node); } 60 | void visit(const CallIntrinsicNode* node) { expr = impl->lowerCallIntrinsic(node); } 61 | void visit(const CallNode* node) { expr = impl->lowerTensorOp(node); } 62 | void visit(const ReductionNode* node) { 63 | taco_ierror << "Reduction nodes not supported in concrete index notation"; 64 | } 65 | void visit(const IndexVarNode* node) { expr = impl->lowerIndexVar(node); } 66 | }; 67 | 68 | LowererImpl::LowererImpl() : visitor(new Visitor(this)) { 69 | } 70 | 71 | 72 | Stmt LowererImpl::lower(IndexStmt stmt) { 73 | return visitor->lower(stmt); 74 | } 75 | 76 | 77 | Expr LowererImpl::lower(IndexExpr expr) { 78 | return visitor->lower(expr); 79 | } 80 | 81 | } 82 | -------------------------------------------------------------------------------- /src/lower/mode.cpp: -------------------------------------------------------------------------------- 1 | #include "taco/lower/mode.h" 2 | 3 | #include "taco/lower/mode_format_impl.h" 4 | #include "taco/ir/ir.h" 5 | #include "taco/type.h" 6 | #include "taco/util/strings.h" 7 | #include "taco/util/collections.h" 8 | 9 | using namespace std; 10 | 11 | namespace taco { 12 | 13 | // class Mode 14 | struct Mode::Content { 15 | ir::Expr tensor; /// the tensor containing mode 16 | Dimension size; /// the size of the mode 17 | int mode; /// the location of mode in a mode hierarchy 18 | ModeFormat modeFormat; /// the type of the mode 19 | 20 | ModePack modePack; /// the pack that contains the mode 21 | size_t packLoc; /// position within pack containing mode 22 | 23 | ModeFormat parentModeFormat; /// type of previous mode in the tensor 24 | 25 | std::map vars; 26 | }; 27 | 28 | Mode::Mode() : content(nullptr) { 29 | } 30 | 31 | Mode::Mode(ir::Expr tensor, Dimension size, int mode, ModeFormat modeFormat, 32 | ModePack modePack, size_t packLoc, ModeFormat parentModeFormat) 33 | : content(new Content) { 34 | taco_iassert(modeFormat.defined()); 35 | content->tensor = tensor; 36 | content->size = size; 37 | content->mode = mode; 38 | content->modeFormat = modeFormat; 39 | content->modePack = modePack; 40 | content->packLoc = packLoc; 41 | content->parentModeFormat = parentModeFormat; 42 | } 43 | 44 | std::string Mode::getName() const { 45 | return util::toString(getTensorExpr()) + std::to_string(getLevel()); 46 | } 47 | 48 | ir::Expr Mode::getTensorExpr() const { 49 | return content->tensor; 50 | } 51 | 52 | Dimension Mode::getSize() const { 53 | return content->size; 54 | } 55 | 56 | int Mode::getLevel() const { 57 | return content->mode; 58 | } 59 | 60 | ModeFormat Mode::getModeFormat() const { 61 | return content->modeFormat; 62 | } 63 | 64 | ModePack Mode::getModePack() const { 65 | return content->modePack; 66 | } 67 | 68 | size_t Mode::getPackLocation() const { 69 | return content->packLoc; 70 | } 71 | 72 | ModeFormat Mode::getParentModeType() const { 73 | return content->parentModeFormat; 74 | } 75 | 76 | ir::Expr Mode::getVar(std::string varName) const { 77 | taco_iassert(hasVar(varName)); 78 | return content->vars.at(varName); 79 | } 80 | 81 | bool Mode::hasVar(std::string varName) const { 82 | return util::contains(content->vars, varName); 83 | } 84 | 85 | void Mode::addVar(std::string varName, ir::Expr var) { 86 | taco_iassert(ir::isa(var)); 87 | content->vars[varName] = var; 88 | } 89 | 90 | bool Mode::defined() const { 91 | return content != nullptr; 92 | } 93 | 94 | std::ostream& operator<<(std::ostream& os, const Mode& mode) { 95 | return os << mode.getName(); 96 | } 97 | 98 | 99 | // class ModePack 100 | struct ModePack::Content { 101 | size_t numModes = 0; 102 | vector arrays; 103 | }; 104 | 105 | ModePack::ModePack() : content(new Content) { 106 | } 107 | 108 | ModePack::ModePack(size_t numModes, ModeFormat modeType, ir::Expr tensor, 109 | int mode, int level) : ModePack() { 110 | content->numModes = numModes; 111 | content->arrays = modeType.impl->getArrays(tensor, mode, level); 112 | } 113 | 114 | size_t ModePack::getNumModes() const { 115 | return content->numModes; 116 | } 117 | 118 | ir::Expr ModePack::getArray(size_t i) const { 119 | return content->arrays[i]; 120 | } 121 | 122 | } 123 | -------------------------------------------------------------------------------- /src/lower/mode_access.cpp: -------------------------------------------------------------------------------- 1 | #include "lower/mode_access.h" 2 | 3 | namespace taco { 4 | 5 | ModeAccess::ModeAccess(Access access, int mode) : access(access), mode(mode){ 6 | } 7 | 8 | Access ModeAccess::getAccess() const { 9 | return access; 10 | } 11 | 12 | size_t ModeAccess::getModePos() const { 13 | return mode; 14 | } 15 | 16 | bool operator==(const ModeAccess& a, const ModeAccess& b) { 17 | return a.getAccess() == b.getAccess() && a.getModePos() == b.getModePos(); 18 | } 19 | 20 | bool operator<(const ModeAccess& a, const ModeAccess& b) { 21 | // First break on the mode position. 22 | if (a.getModePos() != b.getModePos()) { 23 | return a.getModePos() < b.getModePos(); 24 | } 25 | 26 | // Then, return a deep comparison of the underlying access. 27 | return a.getAccess() { 12 | public: 13 | 14 | /// Create a mode access. 15 | ModeAccess(Access access, int mode); 16 | 17 | /// Get the `Access` expression of the mode access. 18 | Access getAccess() const; 19 | 20 | /// Get the position of the accessed mode. 21 | size_t getModePos() const; 22 | 23 | private: 24 | Access access; 25 | int mode; 26 | }; 27 | 28 | bool operator==(const ModeAccess&, const ModeAccess&); 29 | bool operator<(const ModeAccess&, const ModeAccess&); 30 | std::ostream &operator<<(std::ostream &os, const ModeAccess &); 31 | 32 | } 33 | #endif 34 | -------------------------------------------------------------------------------- /src/lower/mode_format_dense.cpp: -------------------------------------------------------------------------------- 1 | #include "taco/lower/mode_format_dense.h" 2 | 3 | using namespace std; 4 | using namespace taco::ir; 5 | 6 | namespace taco { 7 | 8 | DenseModeFormat::DenseModeFormat() : DenseModeFormat(true, true, false) { 9 | } 10 | 11 | DenseModeFormat::DenseModeFormat(const bool isOrdered, const bool isUnique, 12 | const bool isZeroless) : 13 | ModeFormatImpl("dense", true, isOrdered, isUnique, false, true, isZeroless, 14 | true, false, false, true, true, false, false, false, true) { 15 | } 16 | 17 | ModeFormat DenseModeFormat::copy( 18 | std::vector properties) const { 19 | bool isOrdered = this->isOrdered; 20 | bool isUnique = this->isUnique; 21 | bool isZeroless = this->isZeroless; 22 | for (const auto property : properties) { 23 | switch (property) { 24 | case ModeFormat::ORDERED: 25 | isOrdered = true; 26 | break; 27 | case ModeFormat::NOT_ORDERED: 28 | isOrdered = false; 29 | break; 30 | case ModeFormat::UNIQUE: 31 | isUnique = true; 32 | break; 33 | case ModeFormat::NOT_UNIQUE: 34 | isUnique = false; 35 | break; 36 | case ModeFormat::ZEROLESS: 37 | isZeroless = true; 38 | break; 39 | case ModeFormat::NOT_ZEROLESS: 40 | isZeroless = false; 41 | break; 42 | default: 43 | break; 44 | } 45 | } 46 | return ModeFormat( 47 | std::make_shared(isOrdered, isUnique, isZeroless)); 48 | } 49 | 50 | ModeFunction DenseModeFormat::locate(ir::Expr parentPos, 51 | std::vector coords, 52 | Mode mode) const { 53 | Expr pos = ir::Add::make(ir::Mul::make(parentPos, getWidth(mode)), coords.back()); 54 | return ModeFunction(Stmt(), {pos, true}); 55 | } 56 | 57 | Stmt DenseModeFormat::getInsertCoord(Expr p, 58 | const std::vector& i, Mode mode) const { 59 | return Stmt(); 60 | } 61 | 62 | Expr DenseModeFormat::getWidth(Mode mode) const { 63 | return (mode.getSize().isFixed() && mode.getSize().getSize() < 16) ? 64 | (int)mode.getSize().getSize() : 65 | getSizeArray(mode.getModePack()); 66 | } 67 | 68 | Stmt DenseModeFormat::getInsertInitCoords(Expr pBegin, 69 | Expr pEnd, Mode mode) const { 70 | return Stmt(); 71 | } 72 | 73 | Stmt DenseModeFormat::getInsertInitLevel(Expr szPrev, Expr sz, 74 | Mode mode) const { 75 | return Stmt(); 76 | } 77 | 78 | Stmt DenseModeFormat::getInsertFinalizeLevel(Expr szPrev, 79 | Expr sz, Mode mode) const { 80 | return Stmt(); 81 | } 82 | 83 | Expr DenseModeFormat::getAssembledSize(Expr prevSize, Mode mode) const { 84 | return ir::Mul::make(prevSize, getWidth(mode)); 85 | } 86 | 87 | ModeFunction DenseModeFormat::getYieldPos(Expr parentPos, 88 | std::vector coords, Mode mode) const { 89 | return locate(parentPos, coords, mode); 90 | } 91 | 92 | vector DenseModeFormat::getArrays(Expr tensor, int mode, 93 | int level) const { 94 | return {GetProperty::make(tensor, TensorProperty::Dimension, mode)}; 95 | } 96 | 97 | Expr DenseModeFormat::getSizeArray(ModePack pack) const { 98 | return pack.getArray(0); 99 | } 100 | 101 | } 102 | -------------------------------------------------------------------------------- /src/lower/tensor_path.h: -------------------------------------------------------------------------------- 1 | #ifndef TACO_TENSOR_PATH_H 2 | #define TACO_TENSOR_PATH_H 3 | 4 | #include 5 | #include 6 | 7 | #include "taco/util/comparable.h" 8 | 9 | namespace taco { 10 | 11 | class TensorVar; 12 | class IndexVar; 13 | class Access; 14 | 15 | class TensorPathStep; 16 | 17 | /// A tensor Access expression such as A(i,j,k) results in a path in an 18 | /// iteration graph through i,j,k. The exact path (i->j->k, j->k->i, etc.) is 19 | /// dictated by the ordering of the levels in the tensor storage tree. The index 20 | /// variable that indexes into the mode at the first level is the first index 21 | /// variable in the path, and so forth. 22 | class TensorPath : public util::Comparable { 23 | public: 24 | TensorPath(); 25 | TensorPath(const std::vector& path, const Access& access); 26 | 27 | /// Returns the Access expression that the path represents. 28 | const Access& getAccess() const; 29 | 30 | /// Returns the variables along the path. 31 | const std::vector& getVariables() const; 32 | 33 | /// Returns the size (number of steps) of the path. 34 | size_t getSize() const; 35 | 36 | /// Returns the ith tensor step along the path. 37 | TensorPathStep getStep(size_t i) const; 38 | 39 | /// Returns the last step along this path. 40 | TensorPathStep getLastStep() const; 41 | 42 | /// Returns the step incident on var. 43 | TensorPathStep getStep(const IndexVar& var) const; 44 | 45 | /// True if the path is define, false otherwise 46 | bool defined() const; 47 | 48 | friend bool operator==(const TensorPath&, const TensorPath&); 49 | friend bool operator<(const TensorPath&, const TensorPath&); 50 | 51 | private: 52 | struct Content; 53 | std::shared_ptr content; 54 | }; 55 | 56 | std::ostream& operator<<(std::ostream&, const TensorPath&); 57 | 58 | 59 | /// A step along a tensor path. 60 | class TensorPathStep : public util::Comparable { 61 | public: 62 | /// Return the path the tensor belongs to. 63 | const TensorPath& getPath() const; 64 | 65 | /// Returns the location of this step in the path. 66 | int getStep() const; 67 | 68 | private: 69 | TensorPath path; 70 | int step; 71 | 72 | TensorPathStep(); 73 | TensorPathStep(const TensorPath& path, int step); 74 | friend TensorPath; 75 | }; 76 | 77 | bool operator==(const TensorPathStep&, const TensorPathStep&); 78 | bool operator<(const TensorPathStep&, const TensorPathStep&); 79 | std::ostream& operator<<(std::ostream&, const TensorPathStep&); 80 | 81 | } 82 | #endif 83 | -------------------------------------------------------------------------------- /src/parser/README.md: -------------------------------------------------------------------------------- 1 | The parser turns a textual index expression into an index expression data 2 | structure from the `expr` module. It is usually more convenient to create the 3 | index expressions directly through the `expr` API, but the parser is useful 4 | when building tools that receive the index expressions as text. 5 | -------------------------------------------------------------------------------- /src/storage/README.md: -------------------------------------------------------------------------------- 1 | The `storage` module manages tensor index data structure, packing and iterating 2 | over these, inserting values, and reading/writing values from/to files. 3 | -------------------------------------------------------------------------------- /src/taco_tensor_t.cpp: -------------------------------------------------------------------------------- 1 | #include "taco/taco_tensor_t.h" 2 | #include "taco/cuda.h" 3 | #include 4 | #include 5 | 6 | void * alloc_mem(size_t size); 7 | void free_mem(void *ptr); 8 | 9 | // Allocates from unified memory or using malloc depending on what memory is being used 10 | void * alloc_mem(size_t size) { 11 | if (taco::should_use_CUDA_unified_memory()) { 12 | return taco::cuda_unified_alloc(size); 13 | } 14 | else { 15 | return malloc(size); 16 | } 17 | } 18 | 19 | // Free from unified memory or using free depending on what memory is being used 20 | void free_mem(void *ptr) { 21 | if (taco::should_use_CUDA_unified_memory()) { 22 | taco::cuda_unified_free(ptr); 23 | } 24 | else { 25 | free(ptr); 26 | } 27 | } 28 | 29 | // Note about fill: 30 | // It is planned to allow the fill of a result to be null and for TACO to set this when it does compute. This is the 31 | // case we currently expect the fill pointer to be null. 32 | taco_tensor_t* init_taco_tensor_t(int32_t order, int32_t csize, 33 | int32_t* dimensions, int32_t* modeOrdering, 34 | taco_mode_t* mode_types, void* fill_ptr) { 35 | taco_tensor_t* t = (taco_tensor_t *) alloc_mem(sizeof(taco_tensor_t)); 36 | t->order = order; 37 | t->dimensions = (int32_t *) alloc_mem(order * sizeof(int32_t)); 38 | t->mode_ordering = (int32_t *) alloc_mem(order * sizeof(int32_t)); 39 | t->mode_types = (taco_mode_t *) alloc_mem(order * sizeof(taco_mode_t)); 40 | t->indices = (uint8_t ***) alloc_mem(order * sizeof(uint8_t***)); 41 | t->csize = csize; 42 | 43 | int fill_bytes = csize / 8; 44 | t->fill_value = (uint8_t*) alloc_mem(fill_bytes); 45 | 46 | if (fill_ptr) { 47 | uint8_t* fill_inp = (uint8_t*) fill_ptr; 48 | for (int i = 0; i < fill_bytes; ++i) { 49 | t->fill_value[i] = fill_inp[i]; 50 | } 51 | } 52 | 53 | for (int32_t i = 0; i < order; i++) { 54 | t->dimensions[i] = dimensions[i]; 55 | t->mode_ordering[i] = modeOrdering[i]; 56 | t->mode_types[i] = mode_types[i]; 57 | switch (t->mode_types[i]) { 58 | case taco_mode_dense: 59 | t->indices[i] = (uint8_t **) alloc_mem(1 * sizeof(uint8_t **)); 60 | break; 61 | case taco_mode_sparse: 62 | t->indices[i] = (uint8_t **) alloc_mem(2 * sizeof(uint8_t **)); 63 | break; 64 | } 65 | } 66 | return t; 67 | } 68 | 69 | void deinit_taco_tensor_t(taco_tensor_t* t) { 70 | for (int i = 0; i < t->order; i++) { 71 | free_mem(t->indices[i]); 72 | } 73 | free_mem(t->indices); 74 | 75 | free_mem(t->dimensions); 76 | free_mem(t->mode_ordering); 77 | free_mem(t->mode_types); 78 | free_mem(t); 79 | } -------------------------------------------------------------------------------- /src/target.cpp: -------------------------------------------------------------------------------- 1 | #include 2 | #include 3 | #include 4 | 5 | #include "taco/target.h" 6 | 7 | using namespace std; 8 | 9 | namespace taco { 10 | 11 | namespace { 12 | map archMap = {{"c99", Target::C99}, 13 | {"x86", Target::X86}}; 14 | 15 | map osMap = {{"unknown", Target::OSUnknown}, 16 | {"linux", Target::Linux}, 17 | {"macos", Target::MacOS}, 18 | {"windows", Target::Windows}}; 19 | 20 | bool parseTargetString(Target& target, string target_string) { 21 | string rest = target_string; 22 | vector tokens; 23 | auto current_pos = rest.find('-'); 24 | 25 | while (current_pos != string::npos) { 26 | tokens.push_back(rest.substr(0, current_pos)); 27 | rest = rest.substr(current_pos+1); 28 | } 29 | 30 | // now parse the tokens 31 | taco_uassert(tokens.size() >= 2) << 32 | "Invalid target string: " << target_string; 33 | 34 | // first must be architecture 35 | if (archMap.count(tokens[0]) == 0) { 36 | return false; 37 | } 38 | target.arch = archMap[tokens[0]]; 39 | 40 | // next must be os 41 | if (osMap.count(tokens[1]) == 0) { 42 | return false; 43 | } 44 | target.os = osMap[tokens[1]]; 45 | 46 | return true; 47 | } 48 | 49 | } // anonymous namespace 50 | 51 | Target::Target(const std::string &s) { 52 | parseTargetString(*this, s); 53 | } 54 | 55 | 56 | 57 | bool Target::validateTargetString(const string &s) { 58 | string::size_type arch_end = string::npos; 59 | string::size_type os_end = string::npos; 60 | 61 | // locate arch 62 | for (auto res : archMap) { 63 | if (s.find(res.first) != string::npos) { 64 | arch_end = s.find(res.first); 65 | } 66 | } 67 | 68 | // locate os 69 | for (auto res : osMap) { 70 | if (s.find(res.first, arch_end+1) != string::npos) { 71 | os_end = s.find(res.first, arch_end+1); 72 | } 73 | } 74 | 75 | return (arch_end != string::npos) && (os_end != string::npos); 76 | } 77 | 78 | Target getTargetFromEnvironment() { 79 | return Target(Target::Arch::C99, Target::OS::MacOS); 80 | } 81 | } // namespace taco 82 | -------------------------------------------------------------------------------- /src/util/env.cpp: -------------------------------------------------------------------------------- 1 | #include "taco/util/env.h" 2 | #include 3 | #include 4 | #include 5 | 6 | namespace taco { 7 | namespace util { 8 | 9 | std::string cachedtmpdir = ""; 10 | 11 | static int unlink_cb(const char *fpath, const struct stat *sb, int typeflag, struct FTW *ftwbuf) 12 | { 13 | int rv = remove(fpath); 14 | taco_uassert(rv == 0) << 15 | "Unable to create cleanup taco temporary directory. Sorry."; 16 | return rv; 17 | } 18 | 19 | void cachedtmpdirCleanup(void) { 20 | if (cachedtmpdir != ""){ 21 | int rv = nftw(cachedtmpdir.c_str(), unlink_cb, 64, FTW_DEPTH | FTW_PHYS); 22 | taco_uassert(rv == 0) << 23 | "Unable to create cleanup taco temporary directory. Sorry."; 24 | } 25 | } 26 | }} 27 | -------------------------------------------------------------------------------- /src/util/files.cpp: -------------------------------------------------------------------------------- 1 | #include "taco/util/files.h" 2 | 3 | #include "taco/error.h" 4 | 5 | #include 6 | #include 7 | #include 8 | 9 | using namespace std; 10 | 11 | namespace taco { 12 | namespace util { 13 | 14 | std::string sanitizePath(std::string path) { 15 | if (path[0] == '~') { 16 | path = path.replace(0, 1, std::getenv("HOME")); 17 | } 18 | 19 | return path; 20 | } 21 | 22 | void openStream(std::fstream& stream, std::string path, fstream::openmode mode) { 23 | stream.open(sanitizePath(path), mode); 24 | taco_uassert(stream.is_open()) << "Error opening file: " << path; 25 | } 26 | 27 | }} 28 | -------------------------------------------------------------------------------- /src/util/name_generator.cpp: -------------------------------------------------------------------------------- 1 | #include "taco/util/name_generator.h" 2 | 3 | #include 4 | #include 5 | 6 | using namespace std; 7 | 8 | namespace taco { 9 | namespace util { 10 | 11 | atomic uniqueNameCounter; 12 | 13 | static inline int uniqueCount() { 14 | return uniqueNameCounter++; 15 | } 16 | 17 | string uniqueName(char prefix) { 18 | return prefix + to_string(uniqueCount()); 19 | } 20 | 21 | string uniqueName(const string& prefix) { 22 | return prefix + to_string(uniqueCount()); 23 | } 24 | 25 | int getUniqueId() { 26 | return uniqueCount(); 27 | } 28 | 29 | 30 | // class NameGenerator 31 | NameGenerator::NameGenerator() { 32 | } 33 | 34 | NameGenerator::NameGenerator(std::vector reserved) { 35 | for (auto& str : reserved) { 36 | nameCounters.insert({str,0}); 37 | } 38 | } 39 | 40 | std::string NameGenerator::getUniqueName(std::string name) { 41 | stringstream os; 42 | os << name; 43 | if (nameCounters.count(name) > 0) { 44 | os << nameCounters[name]++; 45 | } else { 46 | nameCounters[name] = 0; 47 | } 48 | return os.str(); 49 | } 50 | 51 | }} 52 | -------------------------------------------------------------------------------- /src/util/strings.cpp: -------------------------------------------------------------------------------- 1 | #include "taco/util/strings.h" 2 | 3 | #include 4 | 5 | using namespace std; 6 | 7 | namespace taco { 8 | namespace util { 9 | 10 | vector split(const string &str, const string &delim, bool keepDelim) { 11 | vector results; 12 | size_t prev = 0; 13 | size_t next = 0; 14 | 15 | while ((next = str.find(delim, prev)) != std::string::npos) { 16 | if (next - prev != 0) { 17 | std::string substr = ((keepDelim) ? delim : "") 18 | + str.substr(prev, next-prev); 19 | results.push_back(substr); 20 | } 21 | prev = next + delim.size(); 22 | } 23 | 24 | if (prev < str.size()) { 25 | string substr = ((keepDelim) ? delim : "") + str.substr(prev); 26 | results.push_back(substr); 27 | } 28 | 29 | return results; 30 | } 31 | 32 | std::string repeat(std::string text, size_t n) { 33 | string str; 34 | for (size_t i = 0; i < n; i++) { 35 | str += text; 36 | } 37 | return str; 38 | } 39 | 40 | string fill(string text, char fill, size_t n) { 41 | size_t numfills = n - (text.size()+2); 42 | size_t prefix = numfills/2; 43 | size_t suffix = numfills/2 + (numfills % 2); 44 | return string(prefix,fill) + " " + text + " " + string(suffix,fill); 45 | } 46 | 47 | }} 48 | -------------------------------------------------------------------------------- /test/CMakeLists.txt: -------------------------------------------------------------------------------- 1 | file(GLOB GTEST_SOURCE gtest/gtest-all.cc) 2 | file(GLOB GTEST_HEADER gtest/gtest.h) 3 | include_directories(${TACO_TEST_DIR} ${TACO_SRC_DIR}) 4 | add_library(taco-gtest ${GTEST_HEADER} ${GTEST_SOURCE}) 5 | set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -Wno-missing-declarations -DTACO_TEST_DIR=\"${TACO_TEST_DIR}\"") 6 | 7 | file(GLOB TEST_HEADERS *.h) 8 | file(GLOB TEST_SOURCES *.cpp) 9 | 10 | add_executable(taco-test ${TEST_SOURCES} ${TEST_HEADERS}) 11 | target_link_libraries(taco-test taco-gtest) 12 | target_link_libraries(taco-test pthread) 13 | target_link_libraries(taco-test taco) 14 | 15 | if(${CMAKE_VERSION} VERSION_LESS "3.9.0") 16 | add_test(NAME taco-test COMMAND taco-test) 17 | else() 18 | cmake_policy(SET CMP0054 NEW) 19 | cmake_policy(SET CMP0057 NEW) 20 | include(GoogleTest) 21 | gtest_add_tests(TARGET taco-test) 22 | endif() 23 | 24 | # This will run all *.bats test files in the source test/ folder. 25 | add_test(NAME taco-cli-test COMMAND ${CMAKE_SOURCE_DIR}/test/bats/bin/bats ${CMAKE_SOURCE_DIR}/test/) 26 | -------------------------------------------------------------------------------- /test/data/2tensor.mtx: -------------------------------------------------------------------------------- 1 | %%MatrixMarket matrix coordinate real general 2 | % comment 3 | % comment 4 | 5 | % comment 6 | 32 32 3 7 | 1 1 101.0 8 | 2 1 102.0 9 | 6 3 307.1 10 | -------------------------------------------------------------------------------- /test/data/3tensor.tns: -------------------------------------------------------------------------------- 1 | 736 1 1 1 2 | 1073 1 6 1.1 3 | 881 1 7 1 4 | -------------------------------------------------------------------------------- /test/data/d33.mtx: -------------------------------------------------------------------------------- 1 | %%MatrixMarket matrix coordinate real general 2 | %------------------------------------------------------------------------------- 3 | % name: D33 4 | % [UNSYMMETRIC MATRIX 5 | %------------------------------------------------------------------------------- 6 | 3 3 3 7 | 3 1 3 8 | 1 2 2 9 | 3 3 4 10 | -------------------------------------------------------------------------------- /test/data/d33.ttx: -------------------------------------------------------------------------------- 1 | %%MatrixMarket matrix array real general 2 | % 3 | 3 3 4 | 0 5 | 2 6 | 0 7 | 0 8 | 0 9 | 0 10 | 3 11 | 0 12 | 4 13 | -------------------------------------------------------------------------------- /test/data/d432.ttx: -------------------------------------------------------------------------------- 1 | %%MatrixMarket tensor array real general 2 | % 3 | 4 3 2 4 | 1 5 | 2 6 | 3 7 | 4 8 | 5 9 | 6 10 | 7 11 | 8 12 | 9 13 | 10 14 | 11 15 | 12 16 | 13 17 | 14 18 | 15 19 | 16 20 | 17 21 | 18 22 | 19 23 | 20 24 | 21 25 | 22 26 | 23 27 | 24 28 | 25 29 | 26 30 | 27 31 | 28 32 | 29 33 | 30 34 | 31 35 | 32 36 | -------------------------------------------------------------------------------- /test/data/d567.ttx: -------------------------------------------------------------------------------- 1 | %%MatrixMarket tensor coordinate real general 2 | % 3 | 5 6 7 4 4 | 1 1 1 1 5 | 2 3 1 2 6 | 5 1 7 3 7 | 3 6 3 4 8 | -------------------------------------------------------------------------------- /test/data/ds33.mtx: -------------------------------------------------------------------------------- 1 | %%MatrixMarket matrix coordinate real symmetric 2 | %------------------------------------------------------------------------------- 3 | % name: DS33 4 | % [SYMMETRIC MATRIX 5 | %------------------------------------------------------------------------------- 6 | 3 3 3 7 | 2 1 1 8 | 2 2 2 9 | 3 1 3 10 | -------------------------------------------------------------------------------- /test/data/qcd/theta_simple.ttx: -------------------------------------------------------------------------------- 1 | %%MatrixMarket tensor array real general 2 | 3 3 3 | 1 4 | 2 5 | 3 6 | 4 7 | 5 8 | 6 9 | 7 10 | 8 11 | 9 12 | -------------------------------------------------------------------------------- /test/data/qcd/z.ttx: -------------------------------------------------------------------------------- 1 | %%MatrixMarket tensor array real general 2 | 25 3 | 0.018094644466314386 4 | 0.03216122251007098 5 | 0.05175549310070531 6 | 0.03732396989966932 7 | 0.03983200779523393 8 | 0.0283994291602892 9 | 0.015053426036896383 10 | 0.025641196890290593 11 | 0.006318332965082153 12 | 0.04482299242303558 13 | 0.047809640091653435 14 | 0.054136275416494926 15 | 0.0448179146885659 16 | 0.04431866589371729 17 | 0.04660130003657021 18 | 0.052290489578555095 19 | 0.06411858343777134 20 | 0.042423247653834525 21 | 0.05636732371230492 22 | 0.045870417295565596 23 | 0.044666765679272954 24 | 0.008901025385598236 25 | 0.041728542511581417 26 | 0.052938987360283435 27 | 0.05360810601064291 28 | -------------------------------------------------------------------------------- /test/data/qcd/z_simple.ttx: -------------------------------------------------------------------------------- 1 | %%MatrixMarket tensor array real general 2 | 3 3 | 10 4 | 20 5 | 30 6 | -------------------------------------------------------------------------------- /test/data/rua_32.mtx: -------------------------------------------------------------------------------- 1 | %%MatrixMarket matrix coordinate real general 2 | %------------------------------------------------------------------------------- 3 | % name: RUA_32 4 | %------------------------------------------------------------------------------- 5 | 32 32 126 6 | 1 1 101 7 | 2 1 102 8 | 3 1 103 9 | 4 1 104 10 | 7 1 107 11 | 26 1 126 12 | 1 2 201 13 | 2 2 202 14 | 9 2 209 15 | 21 2 221 16 | 28 2 228 17 | 2 3 302 18 | 3 3 303 19 | 6 3 306 20 | 8 3 308 21 | 9 3 309 22 | 29 3 329 23 | 3 4 403 24 | 4 4 404 25 | 5 4 405 26 | 12 4 412 27 | 3 5 503 28 | 5 5 505 29 | 23 5 523 30 | 27 5 527 31 | 1 6 601 32 | 6 6 606 33 | 16 6 616 34 | 3 7 703 35 | 7 7 707 36 | 14 7 714 37 | 21 7 721 38 | 31 7 731 39 | 1 8 801 40 | 8 8 808 41 | 12 8 812 42 | 17 8 817 43 | 27 8 827 44 | 7 9 907 45 | 9 9 909 46 | 10 9 910 47 | 13 9 913 48 | 19 9 919 49 | 23 9 923 50 | 27 9 927 51 | 1 10 1001 52 | 10 10 1010 53 | 11 10 1011 54 | 21 10 1021 55 | 23 10 1023 56 | 25 10 1025 57 | 27 10 1027 58 | 2 11 1102 59 | 11 11 1111 60 | 15 11 1115 61 | 18 11 1118 62 | 29 11 1129 63 | 6 12 1206 64 | 12 12 1212 65 | 24 12 1224 66 | 11 13 1311 67 | 13 13 1313 68 | 3 14 1403 69 | 14 14 1414 70 | 2 15 1502 71 | 15 15 1515 72 | 20 15 1520 73 | 4 16 1604 74 | 16 16 1616 75 | 22 16 1622 76 | 4 17 1704 77 | 16 17 1716 78 | 17 17 1717 79 | 6 18 1806 80 | 10 18 1810 81 | 18 18 1818 82 | 20 18 1820 83 | 30 18 1830 84 | 1 19 1901 85 | 19 19 1919 86 | 26 19 1926 87 | 8 20 2008 88 | 16 20 2016 89 | 20 20 2020 90 | 3 21 2103 91 | 21 21 2121 92 | 32 21 2132 93 | 11 22 2211 94 | 22 22 2222 95 | 2 23 2302 96 | 17 23 2317 97 | 21 23 2321 98 | 23 23 2323 99 | 12 24 2412 100 | 24 24 2424 101 | 26 24 2426 102 | 6 25 2506 103 | 15 25 2515 104 | 18 25 2518 105 | 24 25 2524 106 | 25 25 2525 107 | 13 26 2613 108 | 18 26 2618 109 | 22 26 2622 110 | 26 26 2626 111 | 5 27 2705 112 | 24 27 2724 113 | 26 27 2726 114 | 27 27 2727 115 | 9 28 2809 116 | 28 28 2828 117 | 3 29 2903 118 | 5 29 2905 119 | 27 29 2927 120 | 29 29 2929 121 | 32 29 2932 122 | 12 30 3012 123 | 17 30 3017 124 | 23 30 3023 125 | 30 30 3030 126 | 13 31 3113 127 | 14 31 3114 128 | 31 31 3131 129 | 24 32 3224 130 | 28 32 3228 131 | 32 32 3232 132 | -------------------------------------------------------------------------------- /test/data/rua_32.rb: -------------------------------------------------------------------------------- 1 | 1Real unsymmetric assembled matrix based on IBM32 RUA_32 2 | 36 3 8 13 0 3 | RUA 32 32 126 0 4 | (16I5) (16I5) (10F7.1) 5 | 1 7 12 18 22 26 29 34 39 46 53 58 61 63 65 68 6 | 71 74 79 82 85 88 90 94 97 102 106 110 112 117 121 124 7 | 127 8 | 1 2 3 4 7 26 1 2 9 21 28 2 3 6 8 9 9 | 29 3 4 5 12 3 5 23 27 1 6 16 3 7 14 21 10 | 31 1 8 12 17 27 7 9 10 13 19 23 27 1 10 11 11 | 21 23 25 27 2 11 15 18 29 6 12 24 11 13 3 14 12 | 2 15 20 4 16 22 4 16 17 6 10 18 20 30 1 19 13 | 26 8 16 20 3 21 32 11 22 2 17 21 23 12 24 26 14 | 6 15 18 24 25 13 18 22 26 5 24 26 27 9 28 3 15 | 5 27 29 32 12 17 23 30 13 14 31 24 28 32 16 | 101.0 102.0 103.0 104.0 107.0 126.0 201.0 202.0 209.0 221.0 17 | 228.0 302.0 303.0 306.0 308.0 309.0 329.0 403.0 404.0 405.0 18 | 412.0 503.0 505.0 523.0 527.0 601.0 606.0 616.0 703.0 707.0 19 | 714.0 721.0 731.0 801.0 808.0 812.0 817.0 827.0 907.0 909.0 20 | 910.0 913.0 919.0 923.0 927.0 1001.0 1010.0 1011.0 1021.0 1023.0 21 | 1025.0 1027.0 1102.0 1111.0 1115.0 1118.0 1129.0 1206.0 1212.0 1224.0 22 | 1311.0 1313.0 1403.0 1414.0 1502.0 1515.0 1520.0 1604.0 1616.0 1622.0 23 | 1704.0 1716.0 1717.0 1806.0 1810.0 1818.0 1820.0 1830.0 1901.0 1919.0 24 | 1926.0 2008.0 2016.0 2020.0 2103.0 2121.0 2132.0 2211.0 2222.0 2302.0 25 | 2317.0 2321.0 2323.0 2412.0 2424.0 2426.0 2506.0 2515.0 2518.0 2524.0 26 | 2525.0 2613.0 2618.0 2622.0 2626.0 2705.0 2724.0 2726.0 2727.0 2809.0 27 | 2828.0 2903.0 2905.0 2927.0 2929.0 2932.0 3012.0 3017.0 3023.0 3030.0 28 | 3113.0 3114.0 3131.0 3224.0 3228.0 3232.0 29 | -------------------------------------------------------------------------------- /test/expr_factory.h: -------------------------------------------------------------------------------- 1 | #ifndef TACO_EXPR_FACTORY_H 2 | #define TACO_EXPR_FACTORY_H 3 | 4 | #include 5 | 6 | #include "taco/tensor.h" 7 | #include "taco/format.h" 8 | 9 | namespace taco { 10 | namespace test { 11 | 12 | struct ExprFactory { 13 | virtual Tensor operator()(std::vector>&, Format) = 0; 14 | virtual ~ExprFactory() {}; 15 | }; 16 | 17 | struct VectorElwiseSqrtFactory : public ExprFactory { 18 | virtual Tensor operator()(std::vector>&, Format); 19 | virtual ~VectorElwiseSqrtFactory() {}; 20 | }; 21 | 22 | struct MatrixElwiseMultiplyFactory : public ExprFactory { 23 | virtual Tensor operator()(std::vector>&, Format); 24 | virtual ~MatrixElwiseMultiplyFactory() {}; 25 | }; 26 | 27 | struct MatrixMultiplyFactory : public ExprFactory { 28 | virtual Tensor operator()(std::vector>&, Format); 29 | virtual ~MatrixMultiplyFactory() {}; 30 | }; 31 | 32 | struct MatrixTransposeMultiplyFactory : public ExprFactory { 33 | virtual Tensor operator()(std::vector>&, Format); 34 | virtual ~MatrixTransposeMultiplyFactory() {}; 35 | }; 36 | 37 | struct MatrixColumnSquaredNormFactory : public ExprFactory { 38 | virtual Tensor operator()(std::vector>&, Format); 39 | virtual ~MatrixColumnSquaredNormFactory() {}; 40 | }; 41 | 42 | struct MatrixColumnNormalizeFactory : public ExprFactory { 43 | virtual Tensor operator()(std::vector>&, Format); 44 | virtual ~MatrixColumnNormalizeFactory() {}; 45 | }; 46 | 47 | struct MTTKRP1Factory : public ExprFactory { 48 | virtual Tensor operator()(std::vector>&, Format); 49 | virtual ~MTTKRP1Factory() {}; 50 | }; 51 | 52 | struct MTTKRP2Factory : public ExprFactory { 53 | virtual Tensor operator()(std::vector>&, Format); 54 | virtual ~MTTKRP2Factory() {}; 55 | }; 56 | 57 | struct MTTKRP3Factory : public ExprFactory { 58 | virtual Tensor operator()(std::vector>&, Format); 59 | virtual ~MTTKRP3Factory() {}; 60 | }; 61 | 62 | struct TensorSquaredNormFactory : public ExprFactory { 63 | virtual Tensor operator()(std::vector>&, Format); 64 | virtual ~TensorSquaredNormFactory() {}; 65 | }; 66 | 67 | struct FactorizedTensorSquaredNormFactory : public ExprFactory { 68 | virtual Tensor operator()(std::vector>&, Format); 69 | virtual ~FactorizedTensorSquaredNormFactory() {}; 70 | }; 71 | 72 | struct FactorizedTensorInnerProductFactory : public ExprFactory { 73 | virtual Tensor operator()(std::vector>&, Format); 74 | virtual ~FactorizedTensorInnerProductFactory() {}; 75 | }; 76 | 77 | struct KroneckerFactory : public ExprFactory { 78 | virtual Tensor operator()(std::vector>&, Format); 79 | virtual ~KroneckerFactory() {}; 80 | }; 81 | 82 | }} 83 | #endif 84 | -------------------------------------------------------------------------------- /test/gtest/LICENSE: -------------------------------------------------------------------------------- 1 | Copyright 2008, Google Inc. 2 | All rights reserved. 3 | 4 | Redistribution and use in source and binary forms, with or without 5 | modification, are permitted provided that the following conditions are 6 | met: 7 | 8 | * Redistributions of source code must retain the above copyright 9 | notice, this list of conditions and the following disclaimer. 10 | * Redistributions in binary form must reproduce the above 11 | copyright notice, this list of conditions and the following disclaimer 12 | in the documentation and/or other materials provided with the 13 | distribution. 14 | * Neither the name of Google Inc. nor the names of its 15 | contributors may be used to endorse or promote products derived from 16 | this software without specific prior written permission. 17 | 18 | THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS 19 | "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT 20 | LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR 21 | A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT 22 | OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, 23 | SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT 24 | LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, 25 | DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY 26 | THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT 27 | (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE 28 | OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. 29 | -------------------------------------------------------------------------------- /test/gtest/README.TAC: -------------------------------------------------------------------------------- 1 | This directory contains Google Test 1.7.0 (SVN r690) after the 2 | source code has been fused using the fuse_gtest_files.py tool 3 | that ships with Google Test. 4 | -------------------------------------------------------------------------------- /test/test-typedcomponent-memory.cpp: -------------------------------------------------------------------------------- 1 | #include 2 | #include 3 | #include 4 | 5 | #include "test.h" 6 | 7 | #include "taco/tensor.h" 8 | 9 | using namespace taco; 10 | 11 | TEST(makecsr, access_past_pos) { 12 | int pagesize = getpagesize(); 13 | int matrixsize = pagesize / sizeof(int) - 1; 14 | // carefully arrange pos[] array memory so that pos[1024] will segfault 15 | int *fence = (int*)mmap(NULL , pagesize*2, PROT_NONE , MAP_PRIVATE | MAP_ANON , -1, 0); 16 | int *pos = (int*)mmap(fence, pagesize , PROT_READ | PROT_WRITE, MAP_PRIVATE | MAP_ANON | MAP_FIXED, -1, 0); 17 | int *crd = new int[matrixsize]; 18 | double *value = new double[matrixsize]; 19 | ASSERT_EQ(pos, fence) << "Unable to create guard zone for test"; 20 | // 1023x1023 CSR identity matrix 21 | pos[matrixsize] = matrixsize; 22 | for(int i = 0; i < matrixsize; i++) { 23 | pos[i] = i; 24 | crd[i] = i; 25 | value[i] = 1.0; 26 | } 27 | { 28 | taco::Tensor A = makeCSR("A", {matrixsize, matrixsize}, pos, crd, value); 29 | } 30 | delete []crd; 31 | delete []value; 32 | munmap(pos, pagesize); 33 | munmap(fence, pagesize*2); 34 | printf("Success!\n"); 35 | } 36 | -------------------------------------------------------------------------------- /test/test.cpp: -------------------------------------------------------------------------------- 1 | #include 2 | 3 | #include "test.h" 4 | #include "taco/tensor.h" 5 | 6 | #include "taco/util/strings.h" 7 | #include "taco/storage/storage.h" 8 | 9 | int main(int argc, char **argv) { 10 | // If there is just one argument and it is not a gtest option, then filter 11 | // the tests using that argument surrounded by wildcards. 12 | std::string filter; 13 | if (argc == 2 && std::string(argv[argc-1]).substr(0,2) != "--") { 14 | filter = std::string(argv[1]); 15 | filter = "*" + filter + "*"; 16 | filter = std::string("--gtest_filter=") + filter; 17 | argv[1] = (char*)filter.c_str(); 18 | } 19 | 20 | ::testing::InitGoogleTest(&argc, argv); 21 | return RUN_ALL_TESTS(); 22 | } 23 | 24 | #define STRINGIFY(x) #x 25 | 26 | namespace taco { 27 | namespace test { 28 | 29 | void ASSERT_STORAGE_EQ(TensorStorage expected, TensorStorage actual) { 30 | SCOPED_TRACE("\nexpected:\n" + util::toString(expected) + 31 | "\nactual:\n" + util::toString(actual)); 32 | ASSERT_TRUE(equals(expected, actual)); 33 | } 34 | 35 | void ASSERT_TENSOR_EQ(TensorBase expected, TensorBase actual) { 36 | SCOPED_TRACE(string("expected: ") + util::toString(expected)); 37 | SCOPED_TRACE(string(" actual: ") + util::toString(actual)); 38 | ASSERT_TRUE(equals(expected, actual)); 39 | } 40 | 41 | std::string testDirectory() { 42 | return TO_STRING(TACO_TEST_DIR); 43 | } 44 | 45 | std::string testDataDirectory() { 46 | return testDirectory() + "/data/"; 47 | } 48 | 49 | ostream& operator<<(ostream& os, const NotationTest& test) { 50 | os << endl; 51 | os << "Expected: " << test.expected << endl; 52 | os << "Actual: " << test.actual << endl; 53 | return os; 54 | } 55 | 56 | void ASSERT_THROWS_EXCEPTION_WITH_ERROR(std::function f, std::string err) { 57 | EXPECT_THROW({ 58 | try { 59 | f(); 60 | } catch (TacoException& e) { 61 | // Catch and inspect the exception to make sure that err is within it. 62 | auto s = std::string(e.what()); 63 | ASSERT_TRUE(s.find(err) != std::string::npos); 64 | // Throw the exception back up to gtest. 65 | throw; 66 | } 67 | }, TacoException); 68 | } 69 | 70 | }} 71 | -------------------------------------------------------------------------------- /test/tests-copyprop.cpp: -------------------------------------------------------------------------------- 1 | #include "test.h" 2 | 3 | #include "taco/ir/ir.h" 4 | #include "taco/ir/simplify.h" 5 | 6 | using taco::ir::Var; 7 | using taco::ir::VarDecl; 8 | using taco::ir::Assign; 9 | using taco::ir::Block; 10 | using taco::Int32; 11 | using taco::ir::simplify; 12 | using taco::ir::Neg; 13 | using taco::ir::Add; 14 | using taco::ir::While; 15 | using taco::ir::Scope; 16 | 17 | TEST(expr, simplify_copy) { 18 | auto a = Var::make("a", Int32), 19 | b = Var::make("b", Int32), 20 | c = Var::make("c", Int32), 21 | d = Var::make("d", Int32); 22 | 23 | auto aDecl = VarDecl::make(a, 42), 24 | bDecl = VarDecl::make(b, a), 25 | cDecl = VarDecl::make(c, b), 26 | // assign `d' a non-trivial expression so it won't get optimized away 27 | dDecl = VarDecl::make(d, Neg::make(c)), 28 | dInc = Assign::make(d, Add::make(d, 1)); 29 | auto block = Block::make(aDecl, bDecl, cDecl, dDecl, dInc); 30 | 31 | auto simplified = simplify(block); 32 | auto *simplifiedBlock = simplified.as(); 33 | ASSERT_EQ(simplifiedBlock->contents.size(), size_t(3)); 34 | 35 | const VarDecl *simplifiedDDecl = simplifiedBlock->contents[1].as(); 36 | ASSERT_EQ(simplifiedDDecl->rhs.as()->a, a); 37 | } 38 | 39 | TEST(expr, simplify_copy_before_update) { 40 | auto a = Var::make("a", Int32), 41 | b = Var::make("b", Int32), 42 | c = Var::make("c", Int32); 43 | 44 | auto aDecl = VarDecl::make(a, 42), 45 | bDecl = VarDecl::make(b, a), 46 | cDecl = VarDecl::make(c, Neg::make(b)), 47 | cInc = Assign::make(c, Add::make(c, 1)), 48 | bInc = Assign::make(b, Add::make(b, 1)); 49 | auto block = Block::make(aDecl, bDecl, cDecl, cInc, bInc); 50 | 51 | auto simplified = simplify(block); 52 | auto *simplifiedBlock = simplified.as(); 53 | ASSERT_EQ(simplifiedBlock->contents.size(), size_t(5)); 54 | 55 | const VarDecl *simplifiedDDecl = simplifiedBlock->contents[2].as(); 56 | ASSERT_EQ(simplifiedDDecl->rhs.as()->a, a); 57 | } 58 | 59 | TEST(expr, simplify_dont_copy_loop_dependent_var) { 60 | auto a = Var::make("a", Int32), 61 | b = Var::make("b", Int32); 62 | 63 | auto aDecl = VarDecl::make(a, 42), 64 | bDecl = VarDecl::make(b, a), 65 | bInc = Assign::make(b, Add::make(b, 1)), 66 | loop = While::make(true, bInc); 67 | auto block = Block::make(aDecl, bDecl, loop); 68 | 69 | auto simplified = simplify(block); 70 | auto *simplifiedBlock = simplified.as(); 71 | ASSERT_EQ(simplifiedBlock->contents.size(), size_t(3)); 72 | 73 | auto *simplifiedLoop = block.as()->contents[2].as(); 74 | auto *simplifiedBody = simplifiedLoop->contents.as(); 75 | const Assign *simplifiedInc = simplifiedBody->scopedStmt.as(); 76 | ASSERT_EQ(simplifiedInc->lhs, b); 77 | ASSERT_EQ(simplifiedInc->rhs.as()->a, b); 78 | } 79 | -------------------------------------------------------------------------------- /test/tests-error.cpp: -------------------------------------------------------------------------------- 1 | #include "test.h" 2 | #include "test_tensors.h" 3 | 4 | #include "taco/tensor.h" 5 | #include "taco/error/error_messages.h" 6 | 7 | using namespace taco; 8 | 9 | const IndexVar i("i"), j("j"), k("k"); 10 | 11 | TEST(error, expr_dimension_mismatch_freevar) { 12 | Tensor a({3}, Format({Sparse}, {0})); 13 | Tensor b({5}, Format({Sparse}, {0})); 14 | ASSERT_THROW(a(i) = b(i), taco::TacoException); 15 | } 16 | 17 | TEST(error, expr_dimension_mismatch_sumvar) { 18 | Tensor a({5}, Format({Sparse}, {0})); 19 | Tensor B({5,4}, Format({Sparse,Sparse}, {0,1})); 20 | Tensor c({3}, Format({Sparse}, {0})); 21 | ASSERT_THROW(a(i) = B(i,j)*c(j), taco::TacoException); 22 | } 23 | 24 | TEST(error, compile_without_expr) { 25 | Tensor a({5}, Sparse); 26 | ASSERT_THROW(a.compile(), taco::TacoException); 27 | } 28 | 29 | TEST(error, compile_tensor_name_collision) { 30 | Tensor a("a", {5}, Sparse); 31 | Tensor b("a", {5}, Sparse); // name should be "b" 32 | a(i) = b(i); 33 | ASSERT_THROW(a.compile(), taco::TacoException); 34 | } 35 | 36 | TEST(error, assemble_without_compile) { 37 | Tensor a({5}, Sparse); 38 | Tensor b({5}, Sparse); 39 | a(i) = b(i); 40 | ASSERT_THROW(a.assemble(), taco::TacoException); 41 | } 42 | 43 | TEST(error, compute_without_compile) { 44 | Tensor a({5}, Sparse); 45 | Tensor b({5}, Sparse); 46 | a(i) = b(i); 47 | ASSERT_THROW(a.compute(), taco::TacoException); 48 | } 49 | -------------------------------------------------------------------------------- /test/tests-expr-reduction.cpp: -------------------------------------------------------------------------------- 1 | #include "test.h" 2 | #include "test_tensors.h" 3 | #include "taco/tensor.h" 4 | 5 | using namespace taco; 6 | 7 | const IndexVar i("i"), j("j"), k("k"); 8 | 9 | TEST(expr, reduction0) { 10 | Tensor a("a"); 11 | 12 | Tensor b = d5a("b", Dense); 13 | Tensor c = d5b("c", Dense); 14 | b.pack(); 15 | c.pack(); 16 | 17 | a = sum(i, b(i)*c(i)); 18 | a.evaluate(); 19 | 20 | Tensor expected("expected"); 21 | expected.insert({}, 40.0); 22 | expected.pack(); 23 | ASSERT_TRUE(equals(expected,a)) << endl << expected << endl << endl << a;; 24 | } 25 | -------------------------------------------------------------------------------- /test/tests-format.cpp: -------------------------------------------------------------------------------- 1 | #include "test.h" 2 | #include "test_tensors.h" 3 | 4 | #include 5 | 6 | #include "taco/tensor.h" 7 | #include "taco/format.h" 8 | #include "taco/index_notation/index_notation.h" 9 | #include "taco/storage/storage.h" 10 | #include "taco/util/strings.h" 11 | 12 | using namespace taco; 13 | 14 | typedef std::tuple>, 15 | std::vector, 16 | std::vector> TestData; 17 | 18 | struct format : public TestWithParam {}; 19 | 20 | TEST_P(format, pack) { 21 | const TensorData& data = std::get<0>(GetParam())[0]; 22 | 23 | Format format(std::get<1>(GetParam()), std::get<2>(GetParam())); 24 | Tensor tensor = data.makeTensor("tensor", format); 25 | tensor.pack(); 26 | 27 | EXPECT_TRUE(data.compare(tensor)); 28 | } 29 | 30 | template 31 | std::vector> packageInputs(Ts... inputs) { 32 | return {inputs...}; 33 | } 34 | 35 | const auto modeTypes1 = generateModeTypes(1); 36 | const auto modeTypes2 = generateModeTypes(2); 37 | const auto modeTypes3 = generateModeTypes(3); 38 | 39 | const auto modeOrderings1 = generateModeOrderings(1); 40 | const auto modeOrderings2 = generateModeOrderings(2); 41 | const auto modeOrderings3 = generateModeOrderings(3); 42 | 43 | INSTANTIATE_TEST_CASE_P(vector, format, Combine( 44 | Values( 45 | packageInputs(d1a_data()), 46 | packageInputs(d1b_data()), 47 | packageInputs(d5a_data()), 48 | packageInputs(d5b_data()), 49 | packageInputs(d5c_data()) 50 | ), ValuesIn(modeTypes1), ValuesIn(modeOrderings1))); 51 | 52 | INSTANTIATE_TEST_CASE_P(matrix, format, Combine( 53 | Values( 54 | packageInputs(d33a_data()), 55 | packageInputs(d33b_data()) 56 | ), ValuesIn(modeTypes2), ValuesIn(modeOrderings2))); 57 | 58 | INSTANTIATE_TEST_CASE_P(tensor3, format, Combine( 59 | Values( 60 | packageInputs(d233a_data()), 61 | packageInputs(d233b_data()) 62 | ), ValuesIn(modeTypes3), ValuesIn(modeOrderings3))); 63 | 64 | TEST(format, sparse) { 65 | Tensor A = d33a("A", Sparse); 66 | A.pack(); 67 | ASSERT_COMPONENTS_EQUALS({{{0,2}, {0,2}}, {{0,1,3}, {1,0,2}}}, {2,3,4}, A); 68 | } 69 | 70 | TEST(format, dense) { 71 | Tensor A = d33a("A", Dense); 72 | A.pack(); 73 | ASSERT_COMPONENTS_EQUALS({{{3}}, {{3}}}, {0,2,0, 0,0,0, 3,0,4}, A); 74 | } 75 | -------------------------------------------------------------------------------- /test/tests-index.cpp: -------------------------------------------------------------------------------- 1 | #include "test.h" 2 | #include "test_tensors.h" 3 | 4 | #include "taco/format.h" 5 | #include "taco/storage/index.h" 6 | #include "taco/storage/array.h" 7 | 8 | using namespace taco; 9 | 10 | TEST(index, size) { 11 | Index index = makeCSRIndex({0, 1, 3, 4, 6}, {0, 0, 3, 1, 1, 2}); 12 | ASSERT_EQ(6u, index.getSize()); 13 | 14 | index = Index(Format({Dense,Dense}), 15 | {ModeIndex({makeArray({3})}), 16 | ModeIndex({makeArray({4})})}); 17 | ASSERT_EQ(12u, index.getSize()); 18 | } 19 | 20 | TEST(index, makeCSR) { 21 | vector rowptr = {0, 1, 3, 4, 6}; 22 | vector colidx = {0, 0, 3, 1, 1, 2}; 23 | 24 | Index index = makeCSRIndex(rowptr, colidx); 25 | ASSERT_EQ(2, index.numModeIndices()); 26 | ASSERT_EQ(1, index.getModeIndex(0).numIndexArrays()); 27 | ASSERT_EQ(2, index.getModeIndex(1).numIndexArrays()); 28 | 29 | auto rowptrarray = index.getModeIndex(1).getIndexArray(0); 30 | ASSERT_ARRAY_EQ(rowptr, {(int*)rowptrarray.getData(), rowptrarray.getSize()}); 31 | 32 | auto colidxarray = index.getModeIndex(1).getIndexArray(1); 33 | ASSERT_ARRAY_EQ(colidx, {(int*)colidxarray.getData(), colidxarray.getSize()}); 34 | } 35 | -------------------------------------------------------------------------------- /test/tests-indexexpr.cpp: -------------------------------------------------------------------------------- 1 | #include "test.h" 2 | #include "test_tensors.h" 3 | #include "taco/tensor.h" 4 | #include "taco/index_notation/index_notation.h" 5 | #include "taco/index_notation/index_notation_nodes.h" 6 | 7 | using namespace taco; 8 | const IndexVar i("i"), j("j"), k("k"); 9 | 10 | Type vectype(type(), {3}); 11 | TensorVar a("a", vectype), b("b", vectype), c("c", vectype); 12 | 13 | TEST(indexexpr, access) { 14 | IndexExpr expr = b(i); 15 | ASSERT_TRUE(isa(expr)); 16 | ASSERT_TRUE(isa(expr.ptr)); 17 | Access access = to(expr); 18 | ASSERT_EQ(access.getTensorVar(), b); 19 | ASSERT_NE(access.getTensorVar(), c); 20 | ASSERT_EQ(access.getIndexVars()[0], i); 21 | ASSERT_NE(access.getIndexVars()[0], j); 22 | } 23 | 24 | TEST(indexexpr, literal) { 25 | IndexExpr expr = 20; 26 | ASSERT_TRUE(isa(expr)); 27 | ASSERT_TRUE(isa(expr.ptr)); 28 | Literal literal = to(expr); 29 | ASSERT_EQ(type(), literal.getDataType()); 30 | ASSERT_EQ(20, literal.getVal()); 31 | } 32 | 33 | TEST(indexexpr, neg) { 34 | IndexExpr expr = -b(i); 35 | ASSERT_TRUE(isa(expr)); 36 | ASSERT_TRUE(isa(expr.ptr)); 37 | Neg neg = to(expr); 38 | ASSERT_TRUE(equals(neg.getA(), b(i))); 39 | } 40 | 41 | TEST(indexexpr, add) { 42 | IndexExpr expr = b(i) + c(i); 43 | ASSERT_TRUE(isa(expr)); 44 | ASSERT_TRUE(isa(expr.ptr)); 45 | Add add = to(expr); 46 | ASSERT_TRUE(equals(add.getA(), b(i))); 47 | ASSERT_TRUE(equals(add.getB(), c(i))); 48 | } 49 | 50 | TEST(indexexpr, sub) { 51 | IndexExpr expr = b(i) - c(i); 52 | ASSERT_TRUE(isa(expr)); 53 | ASSERT_TRUE(isa(expr.ptr)); 54 | Sub sub = to(expr); 55 | ASSERT_TRUE(equals(sub.getA(), b(i))); 56 | ASSERT_TRUE(equals(sub.getB(), c(i))); 57 | } 58 | 59 | TEST(indexexpr, mul) { 60 | IndexExpr expr = b(i) * c(i); 61 | ASSERT_TRUE(isa(expr)); 62 | ASSERT_TRUE(isa(expr.ptr)); 63 | Mul mul = to(expr); 64 | ASSERT_TRUE(equals(mul.getA(), b(i))); 65 | ASSERT_TRUE(equals(mul.getB(), c(i))); 66 | } 67 | 68 | TEST(indexexpr, div) { 69 | IndexExpr expr = b(i) / 2; 70 | ASSERT_TRUE(isa

(expr)); 71 | ASSERT_TRUE(isa(expr.ptr)); 72 | Div div = to
(expr); 73 | ASSERT_TRUE(equals(div.getA(), b(i))); 74 | ASSERT_TRUE(equals(div.getB(), Literal(2))); 75 | } 76 | 77 | TEST(indexexpr, indexvar) { 78 | IndexExpr expr = i; 79 | ASSERT_TRUE(isa(expr)); 80 | ASSERT_TRUE(isa(expr.ptr)); 81 | IndexVar var = to(expr); 82 | ASSERT_EQ(type(), var.getDataType()); 83 | ASSERT_EQ("i", var.getName()); 84 | } 85 | -------------------------------------------------------------------------------- /test/tests-indexstmt.cpp: -------------------------------------------------------------------------------- 1 | #include "test.h" 2 | #include "test_tensors.h" 3 | #include "taco/tensor.h" 4 | #include "taco/index_notation/index_notation.h" 5 | 6 | using namespace taco; 7 | const IndexVar i("i"), j("j"), k("k"); 8 | 9 | TEST(indexstmt, assignment) { 10 | Type t(type(), {3}); 11 | TensorVar a("a", t), b("b", t), c("c", t); 12 | 13 | IndexStmt stmt = a(i) = b(i) + c(i); 14 | ASSERT_TRUE(isa(stmt)); 15 | Assignment assignment = to(stmt); 16 | ASSERT_TRUE(equals(a(i), assignment.getLhs())); 17 | ASSERT_TRUE(equals(b(i) + c(i), assignment.getRhs())); 18 | ASSERT_EQ(IndexExpr(), assignment.getOperator()); 19 | } 20 | 21 | TEST(indexstmt, forall) { 22 | Type t(type(), {3}); 23 | TensorVar a("a", t), b("b", t), c("c", t); 24 | 25 | IndexStmt stmt = forall(i, a(i) = b(i) + c(i)); 26 | ASSERT_TRUE(isa(stmt)); 27 | Forall forallstmt = to(stmt); 28 | ASSERT_EQ(i, forallstmt.getIndexVar()); 29 | ASSERT_TRUE(equals(a(i) = b(i) + c(i), forallstmt.getStmt())); 30 | ASSERT_TRUE(equals(forall(i, a(i) = b(i) + c(i)), forallstmt)); 31 | } 32 | 33 | TEST(indexstmt, where) { 34 | Type t(type(), {3}); 35 | TensorVar a("a", t, Sparse), b("b", t, Sparse), c("c", t, Sparse); 36 | TensorVar w("w", t, Dense); 37 | 38 | IndexStmt stmt = where(forall(i, a(i)=w(i)*c(i)), forall(i, w(i)=b(i))); 39 | ASSERT_TRUE(isa(stmt)); 40 | Where wherestmt = to(stmt); 41 | ASSERT_TRUE(equals(forall(i, a(i)=w(i)*c(i)), wherestmt.getConsumer())); 42 | ASSERT_TRUE(equals(forall(i, w(i)=b(i)), wherestmt.getProducer())); 43 | ASSERT_TRUE(equals(where(forall(i, a(i)=w(i)*c(i)), forall(i, w(i)=b(i))), 44 | wherestmt)); 45 | } 46 | 47 | TEST(indexstmt, multi) { 48 | Type t(type(), {3}); 49 | TensorVar a("a", t, Sparse), b("b", t, Sparse), c("c", t, Sparse); 50 | 51 | IndexStmt stmt = multi(a(i)=c(i), b(i)=c(i)); 52 | ASSERT_TRUE(isa(stmt)); 53 | Multi multistmt = to(stmt); 54 | ASSERT_TRUE(equals(multistmt.getStmt1(), a(i) = c(i))); 55 | ASSERT_TRUE(equals(multistmt.getStmt2(), b(i) = c(i))); 56 | ASSERT_TRUE(equals(multistmt, multi(a(i)=c(i), b(i)=c(i)))); 57 | } 58 | 59 | TEST(indexstmt, sequence) { 60 | Type t(type(), {3}); 61 | TensorVar a("a", t, Sparse), b("b", t, Sparse), c("c", t, Sparse); 62 | 63 | IndexStmt stmt = sequence(a(i) = b(i), a(i) += c(i)); 64 | ASSERT_TRUE(isa(stmt)); 65 | Sequence sequencestmt = to(stmt); 66 | ASSERT_TRUE(equals(a(i) = b(i), sequencestmt.getDefinition())); 67 | ASSERT_TRUE(equals(a(i) += c(i), sequencestmt.getMutation())); 68 | ASSERT_TRUE(equals(sequence(a(i) = b(i), a(i) += c(i)), 69 | sequencestmt)); 70 | } 71 | 72 | TEST(indexstmt, spmm) { 73 | Type t(type(), {3,3}); 74 | TensorVar A("A", t, Sparse), B("B", t, Sparse), C("C", t, Sparse); 75 | TensorVar w("w", Type(type(),{3}), Dense); 76 | 77 | auto spmm = forall(i, 78 | forall(k, 79 | where(forall(j, A(i,j) = w(j)), 80 | forall(j, w(j) += B(i,k)*C(k,j)) 81 | ) 82 | ) 83 | ); 84 | } 85 | 86 | 87 | 88 | -------------------------------------------------------------------------------- /test/tests-qcd.cpp: -------------------------------------------------------------------------------- 1 | /// This file tests expressions from quantum chromodynamics (QCD). 2 | #include "test.h" 3 | #include "taco/tensor.h" 4 | 5 | using namespace std; 6 | using namespace taco; 7 | 8 | static string qcdTestData(string name) { 9 | return testDataDirectory() + "qcd/" + name; 10 | } 11 | 12 | double getScalarValue(Tensor tensor) { 13 | return ((double*)tensor.getStorage().getValues().getData())[0]; 14 | } 15 | 16 | IndexVar i("i"), j("j"), k("k"); 17 | 18 | TEST(qcd, mul0) { 19 | Tensor tau("tau"); 20 | Tensor z = read(qcdTestData("z.ttx"), Dense); 21 | Tensor theta = read(qcdTestData("theta.ttx"), Dense); 22 | 23 | tau = z(i) * z(j) * theta(i,j); 24 | 25 | tau.evaluate(); 26 | ASSERT_DOUBLE_EQ(0.5274167972947047, getScalarValue(tau)); 27 | } 28 | 29 | TEST(qcd, mul1) { 30 | Tensor tau("tau"); 31 | Tensor z = read(qcdTestData("z.ttx"), Dense); 32 | Tensor theta = read(qcdTestData("theta.ttx"), Dense); 33 | 34 | tau = z(i) * z(j) * theta(i,j) * theta(i,j); 35 | 36 | tau.evaluate(); 37 | // Using -O0 to compile the generated kernel yields a slightly different 38 | // answer than using -O3. 39 | #ifdef TACO_DEBUG 40 | ASSERT_DOUBLE_EQ(0.41212798763234648, getScalarValue(tau)); 41 | #else 42 | ASSERT_DOUBLE_EQ(0.41212798763234737, getScalarValue(tau)); 43 | #endif 44 | } 45 | 46 | TEST(qcd, mul2) { 47 | Tensor tau("tau"); 48 | Tensor z = read(qcdTestData("z.ttx"), Dense); 49 | Tensor theta = read(qcdTestData("theta.ttx"), Dense); 50 | 51 | tau = z(i) * z(j) * theta(i,j) * theta(i,j) * theta(i,j); 52 | 53 | tau.evaluate(); 54 | ASSERT_DOUBLE_EQ(0.4120590379120669, getScalarValue(tau)); 55 | } 56 | 57 | TEST(DISABLED_qcd, mul3) { 58 | Tensor tau("tau"); 59 | Tensor z = read(qcdTestData("z.ttx"), Dense); 60 | Tensor theta = read(qcdTestData("theta.ttx"), Dense); 61 | 62 | tau = z(i) * z(j) * z(k) * theta(i,j) * theta(i,k); 63 | 64 | tau.evaluate(); 65 | ASSERT_TRUE(abs(0.3223971010027149 - getScalarValue(tau)) / 0.3223971010027149 < .00000000001); 66 | } 67 | 68 | -------------------------------------------------------------------------------- /test/tests-regression.cpp: -------------------------------------------------------------------------------- 1 | #include "test.h" 2 | 3 | #include "taco/tensor.h" 4 | 5 | using namespace taco; 6 | 7 | -------------------------------------------------------------------------------- /tools/CMakeLists.txt: -------------------------------------------------------------------------------- 1 | file(GLOB TOOL_SOURCES "${TACO_TOOLS_DIR}/*.cpp") 2 | 3 | foreach(TOOL_SOURCE ${TOOL_SOURCES}) 4 | get_filename_component(TOOL ${TOOL_SOURCE} NAME_WE) 5 | add_executable("${TOOL}-tool" ${TOOL_SOURCE}) 6 | target_link_libraries("${TOOL}-tool" taco) 7 | target_include_directories("${TOOL}-tool" PRIVATE "${CMAKE_BINARY_DIR}/include") 8 | SET_TARGET_PROPERTIES("${TOOL}-tool" PROPERTIES OUTPUT_NAME ${TOOL}) 9 | install(TARGETS "${TOOL}-tool" DESTINATION bin) 10 | endforeach() 11 | 12 | include_directories(${TACO_SRC_DIR}) 13 | --------------------------------------------------------------------------------