forked from iree-org/iree
-
Notifications
You must be signed in to change notification settings - Fork 0
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
[metal] Enable compiling to Metal library when possible
This commit extends Metal compilation to additionally invoke Metal shader compilers on macOS to further compile MSL into Metal library, so we don't need to JIT compile during runtime. This reduces runtime overhead. For other platforms where we don't have access to offline Metal compilers, we still embed MSL source code in IREE flatbuffer.
- Loading branch information
1 parent
6e93f00
commit 59e67a7
Showing
7 changed files
with
275 additions
and
38 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
76 changes: 76 additions & 0 deletions
76
compiler/src/iree/compiler/Dialect/HAL/Target/MetalSPIRV/MSLToMetalLib.cpp
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,76 @@ | ||
| // Copyright 2023 The IREE Authors | ||
| // | ||
| // Licensed under the Apache License v2.0 with LLVM Exceptions. | ||
| // See https://llvm.org/LICENSE.txt for license information. | ||
| // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception | ||
|
|
||
| #include "iree/compiler/Dialect/HAL/Target/MetalSPIRV/MSLToMetalLib.h" | ||
|
|
||
| #include <stdlib.h> | ||
|
|
||
| #include "llvm/ADT/Twine.h" | ||
| #include "llvm/Support/Debug.h" | ||
| #include "llvm/Support/FileUtilities.h" | ||
| #include "llvm/Support/MemoryBuffer.h" | ||
| #include "llvm/Support/raw_ostream.h" | ||
| #include "mlir/Support/LogicalResult.h" | ||
|
|
||
| #define DEBUG_TYPE "iree-msl-to-metal-lib" | ||
|
|
||
| namespace mlir { | ||
| namespace iree_compiler { | ||
| namespace IREE { | ||
| namespace HAL { | ||
|
|
||
| /// Returns the command to compile the given MSL source file into Metal library. | ||
| static std::string getMetalCompileCommand(StringRef mslFile, | ||
| StringRef libFile) { | ||
| return llvm::Twine("xcrun -sdk macosx metal -c ") | ||
| .concat(mslFile) | ||
| .concat(" -o - | xcrun -sdk macosx metallib - -o ") | ||
| .concat(libFile) | ||
| .str(); | ||
| } | ||
|
|
||
| /// Returns the given command via system shell. | ||
| static LogicalResult runSystemCommand(StringRef command) { | ||
| LLVM_DEBUG(llvm::dbgs() << "Running system command: '" << command << "'\n"); | ||
| int exitCode = system(command.data()); | ||
| if (exitCode == 0) return success(); | ||
| llvm::errs() << "Failed to run system command '" << command | ||
| << "' with error code: " << exitCode << "\n"; | ||
| return failure(); | ||
| } | ||
|
|
||
| std::unique_ptr<llvm::MemoryBuffer> compileMSLToMetalLib(StringRef mslCode, | ||
| StringRef entryPoint) { | ||
| SmallString<32> mslFile, airFile, libFile; | ||
| int mslFd = 0; | ||
| llvm::sys::fs::createTemporaryFile(entryPoint, "metal", mslFd, mslFile); | ||
| llvm::sys::fs::createTemporaryFile(entryPoint, "metallib", libFile); | ||
| llvm::FileRemover mslRemover(mslFile.c_str()); | ||
| llvm::FileRemover libRemover(libFile.c_str()); | ||
|
|
||
| { // Write input MSL code to the temporary file. | ||
| llvm::raw_fd_ostream inputStream(mslFd, /*shouldClose=*/true); | ||
| inputStream << mslCode << "\n"; | ||
| } | ||
|
|
||
| std::string command = getMetalCompileCommand(mslFile, libFile); | ||
| if (failed(runSystemCommand(command))) return nullptr; | ||
|
|
||
| auto fileOrErr = | ||
| llvm::MemoryBuffer::getFileOrSTDIN(libFile, /*isText=*/false); | ||
| if (std::error_code error = fileOrErr.getError()) { | ||
| llvm::errs() << "Failed to open generated metallib file '" << libFile | ||
| << "' with error: " << error.message(); | ||
| return nullptr; | ||
| } | ||
|
|
||
| return std::move(*fileOrErr); | ||
| } | ||
|
|
||
| } // namespace HAL | ||
| } // namespace IREE | ||
| } // namespace iree_compiler | ||
| } // namespace mlir |
29 changes: 29 additions & 0 deletions
29
compiler/src/iree/compiler/Dialect/HAL/Target/MetalSPIRV/MSLToMetalLib.h
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,29 @@ | ||
| // Copyright 2023 The IREE Authors | ||
| // | ||
| // Licensed under the Apache License v2.0 with LLVM Exceptions. | ||
| // See https://llvm.org/LICENSE.txt for license information. | ||
| // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception | ||
|
|
||
| #ifndef IREE_COMPILER_DIALECT_HAL_TARGET_METALSPIRV_MSLTOMETALLIB_H_ | ||
| #define IREE_COMPILER_DIALECT_HAL_TARGET_METALSPIRV_MSLTOMETALLIB_H_ | ||
|
|
||
| #include "llvm/ADT/StringRef.h" | ||
| #include "llvm/Support/MemoryBuffer.h" | ||
|
|
||
| namespace mlir { | ||
| namespace iree_compiler { | ||
| namespace IREE { | ||
| namespace HAL { | ||
|
|
||
| // Invokes system commands to compile the given |mslCode| into a Metal library | ||
| // and returns the library binary code. |fileName| will be used as a hint for | ||
| // creating intermediate files. | ||
| std::unique_ptr<llvm::MemoryBuffer> compileMSLToMetalLib( | ||
| llvm::StringRef mslCode, llvm::StringRef fileName); | ||
|
|
||
| } // namespace HAL | ||
| } // namespace IREE | ||
| } // namespace iree_compiler | ||
| } // namespace mlir | ||
|
|
||
| #endif // IREE_COMPILER_DIALECT_HAL_TARGET_METALSPIRV_MSLTOMETALLIB_H_ |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Oops, something went wrong.