-
Notifications
You must be signed in to change notification settings - Fork 10.8k
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
[mlir][VectorOps] Support string literals in
vector.print
(#68695)
Printing strings within integration tests is currently quite annoyingly verbose, and can't be tucked into shared helpers as the types depend on the length of the string: ``` llvm.mlir.global internal constant @hello_world("Hello, World!\0") func.func @entry() { %0 = llvm.mlir.addressof @hello_world : !llvm.ptr<array<14 x i8>> %1 = llvm.mlir.constant(0 : index) : i64 %2 = llvm.getelementptr %0[%1, %1] : (!llvm.ptr<array<14 x i8>>, i64, i64) -> !llvm.ptr<i8> llvm.call @printCString(%2) : (!llvm.ptr<i8>) -> () return } ``` So this patch adds a simple extension to `vector.print` to simplify this: ``` func.func @entry() { // Print a vector of characters ;) vector.print str "Hello, World!" return } ``` Most of the logic for this is now shared with `cf.assert` which already does something similar. Depends on #68694
- Loading branch information
Showing
14 changed files
with
204 additions
and
58 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
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,35 @@ | ||
//===- PrintCallHelper.h - Helper to emit runtime print calls ---*- C++ -*-===// | ||
// | ||
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. | ||
// See https://llvm.org/LICENSE.txt for license information. | ||
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception | ||
// | ||
//===----------------------------------------------------------------------===// | ||
|
||
#ifndef MLIR_DIALECT_LLVMIR_PRINTCALLHELPER_H_ | ||
#define MLIR_DIALECT_LLVMIR_PRINTCALLHELPER_H_ | ||
|
||
#include "mlir/Dialect/LLVMIR/LLVMDialect.h" | ||
#include "llvm/ADT/StringRef.h" | ||
#include <optional> | ||
|
||
namespace mlir { | ||
|
||
class OpBuilder; | ||
class LLVMTypeConverter; | ||
|
||
namespace LLVM { | ||
|
||
/// Generate IR that prints the given string to stdout. | ||
/// If a custom runtime function is defined via `runtimeFunctionName`, it must | ||
/// have the signature void(char const*). The default function is `printString`. | ||
void createPrintStrCall(OpBuilder &builder, Location loc, ModuleOp moduleOp, | ||
StringRef symbolName, StringRef string, | ||
const LLVMTypeConverter &typeConverter, | ||
bool addNewline = true, | ||
std::optional<StringRef> runtimeFunctionName = {}); | ||
} // namespace LLVM | ||
|
||
} // namespace mlir | ||
|
||
#endif |
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
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
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
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,66 @@ | ||
//===- PrintCallHelper.cpp - Helper to emit runtime print calls -----------===// | ||
// | ||
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. | ||
// See https://llvm.org/LICENSE.txt for license information. | ||
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception | ||
// | ||
//===----------------------------------------------------------------------===// | ||
|
||
#include "mlir/Conversion/LLVMCommon/PrintCallHelper.h" | ||
#include "mlir/Conversion/LLVMCommon/TypeConverter.h" | ||
#include "mlir/Dialect/LLVMIR/FunctionCallUtils.h" | ||
#include "mlir/Dialect/LLVMIR/LLVMDialect.h" | ||
#include "mlir/IR/Builders.h" | ||
#include "mlir/IR/BuiltinOps.h" | ||
#include "llvm/ADT/ArrayRef.h" | ||
|
||
using namespace mlir; | ||
using namespace llvm; | ||
|
||
static std::string ensureSymbolNameIsUnique(ModuleOp moduleOp, | ||
StringRef symbolName) { | ||
static int counter = 0; | ||
std::string uniqueName = std::string(symbolName); | ||
while (moduleOp.lookupSymbol(uniqueName)) { | ||
uniqueName = std::string(symbolName) + "_" + std::to_string(counter++); | ||
} | ||
return uniqueName; | ||
} | ||
|
||
void mlir::LLVM::createPrintStrCall( | ||
OpBuilder &builder, Location loc, ModuleOp moduleOp, StringRef symbolName, | ||
StringRef string, const LLVMTypeConverter &typeConverter, bool addNewline, | ||
std::optional<StringRef> runtimeFunctionName) { | ||
auto ip = builder.saveInsertionPoint(); | ||
builder.setInsertionPointToStart(moduleOp.getBody()); | ||
MLIRContext *ctx = builder.getContext(); | ||
|
||
// Create a zero-terminated byte representation and allocate global symbol. | ||
SmallVector<uint8_t> elementVals; | ||
elementVals.append(string.begin(), string.end()); | ||
if (addNewline) | ||
elementVals.push_back('\n'); | ||
elementVals.push_back('\0'); | ||
auto dataAttrType = RankedTensorType::get( | ||
{static_cast<int64_t>(elementVals.size())}, builder.getI8Type()); | ||
auto dataAttr = | ||
DenseElementsAttr::get(dataAttrType, llvm::ArrayRef(elementVals)); | ||
auto arrayTy = | ||
LLVM::LLVMArrayType::get(IntegerType::get(ctx, 8), elementVals.size()); | ||
auto globalOp = builder.create<LLVM::GlobalOp>( | ||
loc, arrayTy, /*constant=*/true, LLVM::Linkage::Private, | ||
ensureSymbolNameIsUnique(moduleOp, symbolName), dataAttr); | ||
|
||
// Emit call to `printStr` in runtime library. | ||
builder.restoreInsertionPoint(ip); | ||
auto msgAddr = builder.create<LLVM::AddressOfOp>( | ||
loc, typeConverter.getPointerType(arrayTy), globalOp.getName()); | ||
SmallVector<LLVM::GEPArg> indices(1, 0); | ||
Value gep = builder.create<LLVM::GEPOp>( | ||
loc, typeConverter.getPointerType(builder.getI8Type()), arrayTy, msgAddr, | ||
indices); | ||
Operation *printer = LLVM::lookupOrCreatePrintStringFn( | ||
moduleOp, typeConverter.useOpaquePointers(), runtimeFunctionName); | ||
builder.create<LLVM::CallOp>(loc, TypeRange(), SymbolRefAttr::get(printer), | ||
gep); | ||
} |
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
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
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
Oops, something went wrong.