Create an MLIR environment with matrix multiplication

.github/actions/install-mlir-env-runtime-dependencies/action.yaml

@@ -0,0 +1,18 @@
+---
+name: Install runtime dependencies
+description: Install the runtime dependencies for CompilerGym.
+runs:
+    using: composite
+    steps:
+        - uses: ./.github/actions/install-runtime-dependencies
+
+        - name: Install MLIR env runtime dependencies
+          run: |
+              if [ "$(uname)" = "Darwin" ]; then
+                  # TODO(boian): figure out how to install the C++ std library.
+                  echo "Error: Darwin is not supported"
+                  exit 1
+              else
+                 sudo apt-get install libstdc++-9-dev
+              fi
+          shell: bash

.github/workflows/ci.yaml

@@ -92,6 +92,61 @@ jobs:
                   if-no-files-found: error
                   retention-days: 7
 
+    build-mlir-env-linux-cmake:
+        runs-on: ubuntu-latest
+        steps:
+            - uses: actions/checkout@v2
+
+            - uses: actions/setup-python@v2
+              with:
+                  python-version: 3.9
+
+            - name: Install build dependencies
+              uses: ./.github/actions/install-cmake-build-dependencies
+
+            - name: Download LLVM 14.0.0 release
+              run: |
+                  if [ "$(uname)" = "Darwin" ]; then
+                    wget https://github.com/llvm/llvm-project/releases/download/llvmorg-14.0.0/clang+llvm-14.0.0-x86_64-apple-darwin.tar.xz -O ~/llvm.tar.xz
+                  else
+                    wget https://github.com/llvm/llvm-project/releases/download/llvmorg-14.0.0/clang+llvm-14.0.0-x86_64-linux-gnu-ubuntu-18.04.tar.xz -O ~/llvm.tar.xz
+                  fi
+                  mkdir ~/llvm && tar xvf ~/llvm.tar.xz --strip-components 1 -C ~/llvm
+                  rm ~/llvm.tar.xz
+                  echo "Unpacked, testing for expected file:"
+                  test -d ~/llvm/lib/cmake/llvm
+
+            - name: CMake Build
+              run: |
+                  cmake \
+                      -GNinja \
+                      -DCMAKE_C_COMPILER=clang \
+                      -DCMAKE_CXX_COMPILER=clang++ \
+                      -DCMAKE_EXE_LINKER_FLAGS_INIT="-fuse-ld=lld" \
+                      -DCMAKE_MODULE_LINKER_FLAGS_INIT="-fuse-ld=lld" \
+                      -DCMAKE_SHARED_LINKER_FLAGS_INIT="-fuse-ld=lld" \
+                      -DCOMPILER_GYM_LLVM_PROVIDER=external \
+                      -DLLVM_DIR=$HOME/llvm/lib/cmake/llvm \
+                      -DMLIR_DIR=$HOME/llvm/lib/cmake/mlir \
+                      -DClang_DIR=$HOME/llvm/lib/cmake/clang \
+                      -DPython3_FIND_VIRTUALENV=FIRST \
+                      -DCOMPILER_GYM_ENABLE_MLIR_ENV=ON \
+                      -DCOMPILER_GYM_ENABLE_LLVM_ENV=OFF \
+                      -DCOMPILER_GYM_BUILD_TESTS=ON \
+                      -DCOMPILER_GYM_BUILD_EXAMPLES=ON \
+                      -S . \
+                      -B ~/cmake_build
+                  cmake --build ~/cmake_build
+              shell: bash
+
+            - name: Upload Python wheel
+              uses: actions/upload-artifact@v2
+              with:
+                  name: mlir-env-linux-wheel-cmake
+                  path: ~/cmake_build/py_pkg/dist/*.whl
+                  if-no-files-found: error
+                  retention-days: 7
+
     build-macos:
         runs-on: macos-latest
         steps:
@@ -149,7 +204,9 @@ jobs:
               run: python -m pip install -r tests/requirements.txt
 
             - name: Run the test suite
-              run: make install-test-cov PYTEST_ARGS="--ignore tests/llvm --ignore tests/gcc --ignore tests/loop_tool"
+              run: |
+                  make install-test-cov \
+                  PYTEST_ARGS="--ignore tests/llvm --ignore tests/gcc --ignore tests/loop_tool --ignore tests/mlir"
 
             - name: Upload coverage report to Codecov
               uses: codecov/codecov-action@v2
@@ -184,7 +241,39 @@ jobs:
               run: python -m pip install -r tests/requirements.txt
 
             - name: Run the test suite
-              run: make install-test
+              run: make install-test PYTEST_ARGS="--ignore tests/mlir"
+
+    test-mlir-env-linux-cmake:
+        needs: build-mlir-env-linux-cmake
+        runs-on: ubuntu-latest
+        strategy:
+            fail-fast: false
+            matrix:
+                python: [3.6, 3.7, 3.8, 3.9]
+        steps:
+            - uses: actions/checkout@v2
+
+            - name: Set up Python ${{ matrix.python }}
+              uses: actions/setup-python@v2
+              with:
+                  python-version: ${{ matrix.python }}
+
+            - name: Download Python wheel
+              uses: actions/download-artifact@v2
+              with:
+                  name: mlir-env-linux-wheel-cmake
+
+            - name: Install wheel
+              run: python -m pip install *.whl
+
+            - name: Install runtime dependencies
+              uses: ./.github/actions/install-mlir-env-runtime-dependencies
+
+            - name: Install test dependencies
+              run: python -m pip install -r tests/requirements.txt
+
+            - name: Run the test suite
+              run: make install-test-cov TEST_TARGET="tests/mlir"
 
     test-macos:
         needs: build-macos
@@ -216,7 +305,9 @@ jobs:
               run: python -m pip install -r tests/requirements.txt
 
             - name: Run the test suite
-              run: make install-test-cov PYTEST_ARGS="--ignore tests/llvm --ignore tests/gcc --ignore tests/loop_tool"
+              run: |
+                  make install-test-cov \
+                  PYTEST_ARGS="--ignore tests/llvm --ignore tests/gcc --ignore tests/loop_tool --ignore tests/mlir"
 
             - name: Upload coverage report to Codecov
               uses: codecov/codecov-action@v2

CMakeLists.txt

@@ -14,6 +14,8 @@ endif()
 
 project(compiler_gym ASM C CXX)
 
+set(CMAKE_EXECUTE_PROCESS_COMMAND_ECHO STDOUT)
+
 set(CMAKE_C_STANDARD 11 CACHE STRING "C standard to be used.")
 set(CMAKE_CXX_STANDARD 17 CACHE STRING "C++ standard to be used.")
 
@@ -66,6 +68,17 @@ if(NOT CMAKE_BUILD_TYPE AND NOT CMAKE_CONFIGURATION_TYPES)
     )
 endif()
 
+option(
+    COMPILER_GYM_ENABLE_MLIR_ENV
+    "Whether or not to include the MLIR environment"
+    OFF
+)
+option(
+    COMPILER_GYM_ENABLE_LLVM_ENV
+    "Whether or not to include the LLVM environment"
+    ON
+)
+
 set(CMAKE_POSITION_INDEPENDENT_CODE TRUE)
 
 find_package(Python3 REQUIRED COMPONENTS Interpreter)

WORKSPACE

@@ -17,9 +17,9 @@ http_archive(
 
 http_archive(
     name = "com_github_google_benchmark",
-    sha256 = "616f252f37d61b15037e3c2ef956905baf9c9eecfeab400cb3ad25bae714e214",
-    strip_prefix = "benchmark-1.4.0",
-    url = "https://github.com/google/benchmark/archive/v1.4.0.tar.gz",
+    sha256 = "6132883bc8c9b0df5375b16ab520fac1a85dc9e4cf5be59480448ece74b278d4",
+    strip_prefix = "benchmark-1.6.1",
+    url = "https://github.com/google/benchmark/archive/v1.6.1.tar.gz",
 )
 
 # === Google flags ===

build_tools/cmake/build_external_cmake_project.cmake

@@ -10,7 +10,7 @@ include(write_cache_script)
 function(build_external_cmake_project)
     cmake_parse_arguments(
         _RULE
-        ""
+        "NO_INSTALL"
         "NAME;SRC_DIR;INSTALL_PREFIX"
         "CONFIG_ARGS"
         ${ARGN}
@@ -38,10 +38,25 @@ function(build_external_cmake_project)
         COMMAND "${CMAKE_COMMAND}" --build "${_BIN_DIR}"
         COMMAND_ERROR_IS_FATAL ANY
     )
-    execute_process(
-        COMMAND "${CMAKE_COMMAND}" --install "${_BIN_DIR}"
-        COMMAND_ERROR_IS_FATAL ANY
-    )
+    if(NOT _RULE_NO_INSTALL)
+        execute_process(
+            COMMAND "${CMAKE_COMMAND}" --install "${_BIN_DIR}"
+            COMMAND_ERROR_IS_FATAL ANY
+        )
+    endif()
     list(PREPEND CMAKE_PREFIX_PATH "${_INSTALL_PREFIX}")
     set(CMAKE_PREFIX_PATH ${CMAKE_PREFIX_PATH} PARENT_SCOPE)
+
+    file(
+        GLOB_RECURSE _CONFIG_DEPS
+        FOLLOW_SYMLINKS
+        LIST_DIRECTORIES
+        true
+        "${_RULE_SRC_DIR}/*"
+    )
+    list(APPEND _CONFOG_DEPS "${_RULE_SRC_DIR}")
+    set_property(
+        DIRECTORY APPEND
+        PROPERTY CMAKE_CONFIGURE_DEPENDS ${_CONFIG_DEPS}
+    )
 endfunction()

build_tools/cmake/cg_cc_binary.cmake

@@ -57,7 +57,7 @@ function(cg_cc_binary)
         _RULE
         "HOSTONLY;TESTONLY"
         "NAME"
-        "SRCS;COPTS;DEFINES;LINKOPTS;DATA;DEPS;ABS_DEPS;INCLUDES"
+        "SRCS;COPTS;DEFINES;LINKOPTS;DATA;DEPS;ABS_DEPS;NON_LIB_DEPS;INCLUDES"
         ${ARGN}
     )
 
@@ -122,6 +122,11 @@ function(cg_cc_binary)
 
     cg_add_data_dependencies(NAME ${_RULE_NAME} DATA ${_RULE_DATA})
 
+    if(_RULE_NON_LIB_DEPS)
+        rename_bazel_targets(_RULE_NON_LIB_DEPS "${_RULE_NON_LIB_DEPS}")
+        add_dependencies(${_NAME} ${_RULE_NON_LIB_DEPS})
+    endif()
+
     # Add all targets to a folder in the IDE for organization.
     set_target_properties(
         ${_NAME}

build_tools/cmake/cg_cc_library.cmake

@@ -61,7 +61,7 @@ function(cg_cc_library)
         _RULE
         "PUBLIC;TESTONLY;SHARED"
         "NAME"
-        "HDRS;TEXTUAL_HDRS;SRCS;COPTS;DEFINES;LINKOPTS;DATA;DEPS;ABS_DEPS;NON_LIB_DEPS;INCLUDES"
+        "HDRS;TEXTUAL_HDRS;SRCS;COPTS;DEFINES;LINKOPTS;DATA;DEPS;ABS_DEPS;NON_LIB_DEPS;ABS_NON_LIB_DEPS;INCLUDES"
         ${ARGN}
     )
 
@@ -183,6 +183,9 @@ function(cg_cc_library)
         rename_bazel_targets(_NON_LIB_DEPS "${_RULE_NON_LIB_DEPS}")
         add_dependencies(${_NAME} ${_NON_LIB_DEPS})
     endif()
+    if(_RULE_ABS_NON_LIB_DEPS)
+        add_dependencies(${_NAME} ${_RULE_ABS_NON_LIB_DEPS})
+    endif()
 
     add_library(${_PACKAGE_NS}::${_RULE_NAME} ALIAS ${_NAME})
 

build_tools/cmake/cg_genrule.cmake

@@ -27,6 +27,7 @@ function(cg_genrule)
 
     # TODO(boian): remove this renaming when call sites do not include ":" in target dependency names
     rename_bazel_targets(_DEPS "${_RULE_DEPENDS}")
+    list(APPEND _DEPS ${_RULE_ABS_DEPENDS})
 
     rename_bazel_targets(_NAME "${_RULE_NAME}")
 

build_tools/cmake/cg_macros.cmake

@@ -273,53 +273,6 @@ function(cg_add_data_dependencies)
     add_dependencies(${_NAME} ${_DEPS})
 endfunction()
 
-#-------------------------------------------------------------------------------
-# Tool symlinks
-#-------------------------------------------------------------------------------
-
-# cg_symlink_tool
-#
-# Adds a command to TARGET which symlinks a tool from elsewhere
-# (FROM_TOOL_TARGET_NAME) to a local file name (TO_EXE_NAME) in the current
-# binary directory.
-#
-# Parameters:
-#   TARGET: Local target to which to add the symlink command (i.e. an
-#     cg_py_library, etc).
-#   FROM_TOOL_TARGET: Target of the tool executable that is the source of the
-#     link.
-#   TO_EXE_NAME: The executable name to output in the current binary dir.
-function(cg_symlink_tool)
-    cmake_parse_arguments(
-        ARG
-        ""
-        "TARGET;FROM_TOOL_TARGET;TO_EXE_NAME"
-        ""
-        ${ARGN}
-    )
-
-    # Transform TARGET
-    cg_package_ns(_PACKAGE_NS)
-    cg_package_name(_PACKAGE_NAME)
-    set(_TARGET "${_PACKAGE_NAME}_${ARG_TARGET}")
-    set(_FROM_TOOL_TARGET ${ARG_FROM_TOOL_TARGET})
-    set(_TO_TOOL_PATH
-        "${CMAKE_CURRENT_BINARY_DIR}/${ARG_TO_EXE_NAME}${CMAKE_EXECUTABLE_SUFFIX}"
-    )
-    get_filename_component(_TO_TOOL_DIR "${_TO_TOOL_PATH}" DIRECTORY)
-
-    add_custom_command(
-        TARGET "${_TARGET}"
-        BYPRODUCTS
-            "${CMAKE_CURRENT_BINARY_DIR}/${ARG_TO_EXE_NAME}${CMAKE_EXECUTABLE_SUFFIX}"
-        COMMAND ${CMAKE_COMMAND} -E make_directory "${_TO_TOOL_DIR}"
-        COMMAND
-            ${CMAKE_COMMAND} -E create_symlink
-            "$<TARGET_FILE:${_FROM_TOOL_TARGET}>" "${_TO_TOOL_PATH}"
-        VERBATIM
-    )
-endfunction()
-
 #-------------------------------------------------------------------------------
 # Tests
 #-------------------------------------------------------------------------------

build_tools/cmake/write_cache_script.cmake

@@ -25,6 +25,7 @@ function(write_cache_script _DST_FILE)
         CMAKE_STATIC_LINKER_FLAGS
         CMAKE_SHARED_LINKER_FLAGS_INIT
         CMAKE_SHARED_LINKER_FLAGS
+        CMAKE_EXECUTE_PROCESS_COMMAND_ECHO
         CMAKE_EXE_LINKER_FLAGS_INIT
         CMAKE_EXE_LINKER_FLAGS
     )
@@ -33,7 +34,7 @@ function(write_cache_script _DST_FILE)
             file(
                 APPEND
                 "${_DST_FILE}"
-                "set(${_VAR} \"${${_VAR}}\" CACHE STRING \"\")\n"
+                "set(${_VAR} \"${${_VAR}}\" CACHE STRING \"\" FORCE)\n"
             )
         endif()
     endforeach()

compiler_gym/BUILD

@@ -18,6 +18,7 @@ py_library(
         "//compiler_gym/bin",
         "//compiler_gym/datasets",
         "//compiler_gym/envs",
+        "//compiler_gym/errors",
         "//compiler_gym/leaderboard",
         "//compiler_gym/service",
         "//compiler_gym/spaces",
@@ -51,11 +52,11 @@ py_library(
 py_library(
     name = "random_search",
     srcs = ["random_search.py"],
-    data = ["//compiler_gym/envs/llvm/service"],
     visibility = ["//visibility:public"],
     deps = [
+        ":config",
         "//compiler_gym/envs",
-        "//compiler_gym/service:connection",
+        "//compiler_gym/errors",
         "//compiler_gym/util",
     ],
 )
@@ -77,6 +78,7 @@ py_library(
     name = "validation_error",
     srcs = ["validation_error.py"],
     visibility = ["//compiler_gym:__subpackages__"],
+    deps = ["//compiler_gym/errors"],
 )
 
 py_library(
@@ -86,6 +88,20 @@ py_library(
     deps = [
         ":compiler_env_state",
         ":validation_error",
+        "//compiler_gym/errors",
         "//compiler_gym/util",
     ],
 )
+
+genrule(
+    name = "make_config",
+    srcs = ["make_config.py"],
+    outs = ["config.py"],
+    cmd = "\"$<\" --out-file-path \"$@\" --enable-llvm-env",
+)
+
+py_library(
+    name = "config",
+    srcs = [":config.py"],
+    visibility = ["//compiler_gym:__subpackages__"],
+)