Skip to content

Commit 9797190

Browse files
drprajapdrprajap
committed
[mlir][spirv] Add gpu printf op lowering to spirv.CL.printf op
This change contains following: - adds lowering of printf op to spirv.CL.printf op in GPUToSPIRV pass. - Fixes Constant decoration parsing for spirv GlobalVariable. - minor modification to spirv.CL.printf op assembly format.
1 parent ee40ffd commit 9797190

File tree

6 files changed

+190
-5
lines changed

6 files changed

+190
-5
lines changed

‎mlir/include/mlir/Dialect/SPIRV/IR/SPIRVCLOps.td

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -875,7 +875,7 @@ def SPIRV_CLPrintfOp : SPIRV_CLOp<"printf", 184, []> {
875875
#### Example:
876876

877877
```mlir
878-
%0 = spirv.CL.printf %0 %1 %2 : (!spirv.ptr<i8, UniformConstant>, (i32, i32)) -> i32
878+
%0 = spirv.CL.printf %0 : !spirv.ptr<i8, UniformConstant>(%1, %2 : i32, i32) -> i32
879879
```
880880
}];
881881

@@ -889,7 +889,7 @@ def SPIRV_CLPrintfOp : SPIRV_CLOp<"printf", 184, []> {
889889
);
890890

891891
let assemblyFormat = [{
892-
$format `,` $arguments attr-dict `:` `(` type($format) `,` `(` type($arguments) `)` `)` `->` type($result)
892+
$format `:` type($format) ( `(` $arguments^ `:` type($arguments) `)`)? attr-dict `->` type($result)
893893
}];
894894

895895
let hasVerifier = 0;

‎mlir/lib/Conversion/GPUToSPIRV/GPUToSPIRV.cpp

Lines changed: 113 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -121,6 +121,15 @@ class GPUShuffleConversion final : public OpConversionPattern<gpu::ShuffleOp> {
121121
ConversionPatternRewriter &rewriter) const override;
122122
};
123123

124+
class GPUPrintfConversion final : public OpConversionPattern<gpu::PrintfOp> {
125+
public:
126+
using OpConversionPattern<gpu::PrintfOp>::OpConversionPattern;
127+
128+
LogicalResult
129+
matchAndRewrite(gpu::PrintfOp gpuPrintfOp, OpAdaptor adaptor,
130+
ConversionPatternRewriter &rewriter) const override;
131+
};
132+
124133
} // namespace
125134

126135
//===----------------------------------------------------------------------===//
@@ -597,6 +606,108 @@ class GPUSubgroupReduceConversion final
597606
}
598607
};
599608

609+
LogicalResult GPUPrintfConversion::matchAndRewrite(
610+
gpu::PrintfOp gpuPrintfOp, OpAdaptor adaptor,
611+
ConversionPatternRewriter &rewriter) const {
612+
613+
auto loc = gpuPrintfOp.getLoc();
614+
615+
auto funcOp =
616+
rewriter.getBlock()->getParent()->getParentOfType<mlir::spirv::FuncOp>();
617+
618+
auto moduleOp = funcOp->getParentOfType<mlir::spirv::ModuleOp>();
619+
620+
const char formatStringPrefix[] = "printfMsg";
621+
unsigned stringNumber = 0;
622+
mlir::SmallString<16> globalVarName;
623+
mlir::spirv::GlobalVariableOp globalVar;
624+
625+
// formulate spirv global variable name
626+
do {
627+
globalVarName.clear();
628+
(formatStringPrefix + llvm::Twine(stringNumber++))
629+
.toStringRef(globalVarName);
630+
} while (moduleOp.lookupSymbol(globalVarName));
631+
632+
auto i8Type = rewriter.getI8Type();
633+
auto i32Type = rewriter.getI32Type();
634+
635+
unsigned scNum = 0;
636+
auto createSpecConstant = [&](unsigned value) {
637+
auto attr = rewriter.getI8IntegerAttr(value);
638+
mlir::SmallString<16> specCstName;
639+
(llvm::Twine(globalVarName) + "_sc" + llvm::Twine(scNum++))
640+
.toStringRef(specCstName);
641+
642+
return rewriter.create<mlir::spirv::SpecConstantOp>(
643+
loc, rewriter.getStringAttr(specCstName), attr);
644+
};
645+
646+
// define GlobalVarOp with printf format string using SpecConstants
647+
// and make composite of SpecConstants
648+
{
649+
mlir::Operation *parent =
650+
mlir::SymbolTable::getNearestSymbolTable(gpuPrintfOp->getParentOp());
651+
652+
mlir::ConversionPatternRewriter::InsertionGuard guard(rewriter);
653+
654+
mlir::Block &entryBlock = *parent->getRegion(0).begin();
655+
rewriter.setInsertionPointToStart(
656+
&entryBlock); // insertion point at module level
657+
658+
// Create Constituents with SpecConstant to construct
659+
// SpecConstantCompositeOp
660+
llvm::SmallString<20> formatString(gpuPrintfOp.getFormat());
661+
formatString.push_back('\0'); // Null terminate for C
662+
mlir::SmallVector<mlir::Attribute, 4> constituents;
663+
for (auto c : formatString) {
664+
auto cSpecConstantOp = createSpecConstant(c);
665+
constituents.push_back(mlir::SymbolRefAttr::get(cSpecConstantOp));
666+
}
667+
668+
// Create specialization constant composite defined via spirv.SpecConstant
669+
size_t contentSize = constituents.size();
670+
auto globalType = mlir::spirv::ArrayType::get(i8Type, contentSize);
671+
mlir::spirv::SpecConstantCompositeOp specCstComposite;
672+
mlir::SmallString<16> specCstCompositeName;
673+
(llvm::Twine(globalVarName) + "_scc").toStringRef(specCstCompositeName);
674+
specCstComposite = rewriter.create<mlir::spirv::SpecConstantCompositeOp>(
675+
loc, mlir::TypeAttr::get(globalType),
676+
rewriter.getStringAttr(specCstCompositeName),
677+
rewriter.getArrayAttr(constituents));
678+
679+
// Define GlobalVariable initialized from Constant Composite
680+
globalVar = rewriter.create<mlir::spirv::GlobalVariableOp>(
681+
loc,
682+
mlir::spirv::PointerType::get(
683+
globalType, mlir::spirv::StorageClass::UniformConstant),
684+
globalVarName, mlir::FlatSymbolRefAttr::get(specCstComposite));
685+
686+
globalVar->setAttr("Constant", rewriter.getUnitAttr());
687+
}
688+
689+
// Get SSA value of Global variable
690+
mlir::Value globalPtr =
691+
rewriter.create<mlir::spirv::AddressOfOp>(loc, globalVar);
692+
mlir::Value fmtStr = rewriter.create<mlir::spirv::BitcastOp>(
693+
loc,
694+
mlir::spirv::PointerType::get(i8Type,
695+
mlir::spirv::StorageClass::UniformConstant),
696+
globalPtr);
697+
698+
// Get printf arguments
699+
auto argsRange = adaptor.getArgs();
700+
mlir::SmallVector<mlir::Value, 4> printfArgs;
701+
printfArgs.reserve(argsRange.size() + 1);
702+
printfArgs.append(argsRange.begin(), argsRange.end());
703+
704+
rewriter.create<mlir::spirv::CLPrintfOp>(loc, i32Type, fmtStr, printfArgs);
705+
706+
rewriter.eraseOp(gpuPrintfOp);
707+
708+
return mlir::success();
709+
}
710+
600711
//===----------------------------------------------------------------------===//
601712
// GPU To SPIRV Patterns.
602713
//===----------------------------------------------------------------------===//
@@ -620,5 +731,6 @@ void mlir::populateGPUToSPIRVPatterns(SPIRVTypeConverter &typeConverter,
620731
SingleDimLaunchConfigConversion<gpu::SubgroupSizeOp,
621732
spirv::BuiltIn::SubgroupSize>,
622733
WorkGroupSizeConversion, GPUAllReduceConversion,
623-
GPUSubgroupReduceConversion>(typeConverter, patterns.getContext());
734+
GPUSubgroupReduceConversion, GPUPrintfConversion>(typeConverter,
735+
patterns.getContext());
624736
}

‎mlir/lib/Target/SPIRV/Deserialization/Deserializer.cpp

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -319,6 +319,7 @@ LogicalResult spirv::Deserializer::processDecoration(ArrayRef<uint32_t> words) {
319319
case spirv::Decoration::Restrict:
320320
case spirv::Decoration::RestrictPointer:
321321
case spirv::Decoration::NoContraction:
322+
case spirv::Decoration::Constant:
322323
if (words.size() != 2) {
323324
return emitError(unknownLoc, "OpDecoration with ")
324325
<< decorationName << "needs a single target <id>";

‎mlir/lib/Target/SPIRV/Serialization/Serializer.cpp

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -286,6 +286,7 @@ LogicalResult Serializer::processDecorationAttr(Location loc, uint32_t resultID,
286286
case spirv::Decoration::Restrict:
287287
case spirv::Decoration::RestrictPointer:
288288
case spirv::Decoration::NoContraction:
289+
case spirv::Decoration::Constant:
289290
// For unit attributes and decoration attributes, the args list
290291
// has no values so we do nothing.
291292
if (isa<UnitAttr, DecorationAttr>(attr))

‎mlir/test/Conversion/GPUToSPIRV/printf.mlir

Lines changed: 71 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,71 @@
1+
// RUN: mlir-opt -allow-unregistered-dialect -split-input-file -convert-gpu-to-spirv -verify-diagnostics %s -o - | FileCheck %s
2+
3+
module attributes {
4+
gpu.container_module,
5+
spirv.target_env = #spirv.target_env<#spirv.vce<v1.0, [Addresses, Int8, Kernel], []>, #spirv.resource_limits<>>
6+
} {
7+
func.func @main() {
8+
%c1 = arith.constant 1 : index
9+
10+
gpu.launch_func @kernels::@printf
11+
blocks in (%c1, %c1, %c1) threads in (%c1, %c1, %c1)
12+
args()
13+
return
14+
}
15+
16+
gpu.module @kernels {
17+
// CHECK: spirv.module @{{.*}} Physical32 OpenCL {
18+
// CHECK-DAG: spirv.SpecConstant [[SPECCST:@.*]] = {{.*}} : i8
19+
// CHECK-DAG: spirv.SpecConstantComposite [[SPECCSTCOMPOSITE:@.*]] ([[SPECCST]], {{.*}}) : !spirv.array<[[ARRAYSIZE:.*]] x i8>
20+
// CHECK-DAG: spirv.GlobalVariable [[PRINTMSG:@.*]] initializer([[SPECCSTCOMPOSITE]]) : !spirv.ptr<!spirv.array<[[ARRAYSIZE]] x i8>, UniformConstant>
21+
gpu.func @printf() kernel
22+
attributes
23+
{spirv.entry_point_abi = #spirv.entry_point_abi<>} {
24+
// CHECK: [[FMTSTR_ADDR:%.*]] = spirv.mlir.addressof [[PRINTMSG]] : !spirv.ptr<!spirv.array<[[ARRAYSIZE]] x i8>, UniformConstant>
25+
// CHECK-NEXT: [[FMTSTR_PTR:%.*]] = spirv.Bitcast [[FMTSTR_ADDR]] : !spirv.ptr<!spirv.array<[[ARRAYSIZE]] x i8>, UniformConstant> to !spirv.ptr<i8, UniformConstant>
26+
// CHECK-NEXT {{%.*}} = spirv.CL.printf [[FMTSTR_PTR]] : (!spirv.ptr<i8, UniformConstant>) -> i32
27+
gpu.printf "\nHello\n"
28+
// CHECK: spirv.Return
29+
gpu.return
30+
}
31+
}
32+
}
33+
34+
// -----
35+
36+
module attributes {
37+
gpu.container_module,
38+
spirv.target_env = #spirv.target_env<#spirv.vce<v1.0, [Addresses, Int8, Kernel], []>, #spirv.resource_limits<>>
39+
} {
40+
func.func @main() {
41+
%c1 = arith.constant 1 : index
42+
%c100 = arith.constant 100: i32
43+
%cst_f32 = arith.constant 314.4: f32
44+
45+
gpu.launch_func @kernels1::@printf_args
46+
blocks in (%c1, %c1, %c1) threads in (%c1, %c1, %c1)
47+
args(%c100: i32, %cst_f32: f32)
48+
return
49+
}
50+
51+
gpu.module @kernels1 {
52+
// CHECK: spirv.module @{{.*}} Physical32 OpenCL {
53+
// CHECK-DAG: spirv.SpecConstant [[SPECCST:@.*]] = {{.*}} : i8
54+
// CHECK-DAG: spirv.SpecConstantComposite [[SPECCSTCOMPOSITE:@.*]] ([[SPECCST]], {{.*}}) : !spirv.array<[[ARRAYSIZE:.*]] x i8>
55+
// CHECK-DAG: spirv.GlobalVariable [[PRINTMSG:@.*]] initializer([[SPECCSTCOMPOSITE]]) : !spirv.ptr<!spirv.array<[[ARRAYSIZE]] x i8>, UniformConstant>
56+
gpu.func @printf_args(%arg0: i32, %arg1: f32) kernel
57+
attributes {spirv.entry_point_abi = #spirv.entry_point_abi<>} {
58+
%0 = gpu.block_id x
59+
%1 = gpu.block_id y
60+
%2 = gpu.thread_id x
61+
62+
// CHECK: [[FMTSTR_ADDR:%.*]] = spirv.mlir.addressof [[PRINTMSG]] : !spirv.ptr<!spirv.array<[[ARRAYSIZE]] x i8>, UniformConstant>
63+
// CHECK-NEXT: [[FMTSTR_PTR1:%.*]] = spirv.Bitcast [[FMTSTR_ADDR]] : !spirv.ptr<!spirv.array<[[ARRAYSIZE]] x i8>, UniformConstant> to !spirv.ptr<i8, UniformConstant>
64+
// CHECK-NEXT: {{%.*}} = spirv.CL.printf [[FMTSTR_PTR1]], {{%.*}}, {{%.*}}, {{%.*}} : (!spirv.ptr<i8, UniformConstant>, (i32, f32, i32)) -> i32
65+
gpu.printf "\nHello, world : %d %f \n Thread id: %d\n" %arg0, %arg1, %2: i32, f32, index
66+
67+
// CHECK: spirv.Return
68+
gpu.return
69+
}
70+
}
71+
}

‎mlir/test/Dialect/SPIRV/IR/ocl-ops.mlir

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -275,8 +275,8 @@ func.func @rintvec(%arg0 : vector<3xf16>) -> () {
275275
//===----------------------------------------------------------------------===//
276276
// CHECK-LABEL: func.func @printf(
277277
func.func @printf(%arg0 : !spirv.ptr<i8, UniformConstant>, %arg1 : i32, %arg2 : i32) -> i32 {
278-
// CHECK: spirv.CL.printf {{%.*}}, {{%.*}}, {{%.*}} : (!spirv.ptr<i8, UniformConstant>, (i32, i32)) -> i32
279-
%0 = spirv.CL.printf %arg0, %arg1, %arg2 : (!spirv.ptr<i8, UniformConstant>, (i32, i32)) -> i32
278+
// CHECK: spirv.CL.printf {{%.*}} : !spirv.ptr<i8, UniformConstant>({{%.*}}, {{%.*}} : i32, i32) -> i32
279+
%0 = spirv.CL.printf %arg0 : !spirv.ptr<i8, UniformConstant>(%arg1, %arg2 : i32, i32) -> i32
280280
return %0 : i32
281281
}
282282

0 commit comments

Comments
 (0)