1 //===- TestAvailability.cpp - Pass to test SPIR-V op availability ---------===//
2 //
3 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
4 // See https://llvm.org/LICENSE.txt for license information.
5 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
6 //
7 //===----------------------------------------------------------------------===//
8 
9 #include "mlir/Dialect/Func/IR/FuncOps.h"
10 #include "mlir/Dialect/SPIRV/IR/SPIRVOps.h"
11 #include "mlir/Dialect/SPIRV/IR/SPIRVTypes.h"
12 #include "mlir/Dialect/SPIRV/Transforms/SPIRVConversion.h"
13 #include "mlir/IR/BuiltinOps.h"
14 #include "mlir/Pass/Pass.h"
15 
16 using namespace mlir;
17 
18 //===----------------------------------------------------------------------===//
19 // Printing op availability pass
20 //===----------------------------------------------------------------------===//
21 
22 namespace {
23 /// A pass for testing SPIR-V op availability.
24 struct PrintOpAvailability
25     : public PassWrapper<PrintOpAvailability, OperationPass<FuncOp>> {
26   void runOnOperation() override;
27   StringRef getArgument() const final { return "test-spirv-op-availability"; }
28   StringRef getDescription() const final {
29     return "Test SPIR-V op availability";
30   }
31 };
32 } // namespace
33 
34 void PrintOpAvailability::runOnOperation() {
35   auto f = getOperation();
36   llvm::outs() << f.getName() << "\n";
37 
38   Dialect *spvDialect = getContext().getLoadedDialect("spv");
39 
40   f->walk([&](Operation *op) {
41     if (op->getDialect() != spvDialect)
42       return WalkResult::advance();
43 
44     auto opName = op->getName();
45     auto &os = llvm::outs();
46 
47     if (auto minVersionIfx = dyn_cast<spirv::QueryMinVersionInterface>(op)) {
48       Optional<spirv::Version> minVersion = minVersionIfx.getMinVersion();
49       os << opName << " min version: ";
50       if (minVersion)
51         os << spirv::stringifyVersion(*minVersion) << "\n";
52       else
53         os << "None\n";
54     }
55 
56     if (auto maxVersionIfx = dyn_cast<spirv::QueryMaxVersionInterface>(op)) {
57       Optional<spirv::Version> maxVersion = maxVersionIfx.getMaxVersion();
58       os << opName << " max version: ";
59       if (maxVersion)
60         os << spirv::stringifyVersion(*maxVersion) << "\n";
61       else
62         os << "None\n";
63     }
64 
65     if (auto extension = dyn_cast<spirv::QueryExtensionInterface>(op)) {
66       os << opName << " extensions: [";
67       for (const auto &exts : extension.getExtensions()) {
68         os << " [";
69         llvm::interleaveComma(exts, os, [&](spirv::Extension ext) {
70           os << spirv::stringifyExtension(ext);
71         });
72         os << "]";
73       }
74       os << " ]\n";
75     }
76 
77     if (auto capability = dyn_cast<spirv::QueryCapabilityInterface>(op)) {
78       os << opName << " capabilities: [";
79       for (const auto &caps : capability.getCapabilities()) {
80         os << " [";
81         llvm::interleaveComma(caps, os, [&](spirv::Capability cap) {
82           os << spirv::stringifyCapability(cap);
83         });
84         os << "]";
85       }
86       os << " ]\n";
87     }
88     os.flush();
89 
90     return WalkResult::advance();
91   });
92 }
93 
94 namespace mlir {
95 void registerPrintSpirvAvailabilityPass() {
96   PassRegistration<PrintOpAvailability>();
97 }
98 } // namespace mlir
99 
100 //===----------------------------------------------------------------------===//
101 // Converting target environment pass
102 //===----------------------------------------------------------------------===//
103 
104 namespace {
105 /// A pass for testing SPIR-V op availability.
106 struct ConvertToTargetEnv
107     : public PassWrapper<ConvertToTargetEnv, OperationPass<FuncOp>> {
108   StringRef getArgument() const override { return "test-spirv-target-env"; }
109   StringRef getDescription() const override {
110     return "Test SPIR-V target environment";
111   }
112   void runOnOperation() override;
113 };
114 
115 struct ConvertToAtomCmpExchangeWeak : public RewritePattern {
116   ConvertToAtomCmpExchangeWeak(MLIRContext *context);
117   LogicalResult matchAndRewrite(Operation *op,
118                                 PatternRewriter &rewriter) const override;
119 };
120 
121 struct ConvertToBitReverse : public RewritePattern {
122   ConvertToBitReverse(MLIRContext *context);
123   LogicalResult matchAndRewrite(Operation *op,
124                                 PatternRewriter &rewriter) const override;
125 };
126 
127 struct ConvertToGroupNonUniformBallot : public RewritePattern {
128   ConvertToGroupNonUniformBallot(MLIRContext *context);
129   LogicalResult matchAndRewrite(Operation *op,
130                                 PatternRewriter &rewriter) const override;
131 };
132 
133 struct ConvertToModule : public RewritePattern {
134   ConvertToModule(MLIRContext *context);
135   LogicalResult matchAndRewrite(Operation *op,
136                                 PatternRewriter &rewriter) const override;
137 };
138 
139 struct ConvertToSubgroupBallot : public RewritePattern {
140   ConvertToSubgroupBallot(MLIRContext *context);
141   LogicalResult matchAndRewrite(Operation *op,
142                                 PatternRewriter &rewriter) const override;
143 };
144 } // namespace
145 
146 void ConvertToTargetEnv::runOnOperation() {
147   MLIRContext *context = &getContext();
148   FuncOp fn = getOperation();
149 
150   auto targetEnv = fn.getOperation()
151                        ->getAttr(spirv::getTargetEnvAttrName())
152                        .cast<spirv::TargetEnvAttr>();
153   if (!targetEnv) {
154     fn.emitError("missing 'spv.target_env' attribute");
155     return signalPassFailure();
156   }
157 
158   auto target = SPIRVConversionTarget::get(targetEnv);
159 
160   RewritePatternSet patterns(context);
161   patterns.add<ConvertToAtomCmpExchangeWeak, ConvertToBitReverse,
162                ConvertToGroupNonUniformBallot, ConvertToModule,
163                ConvertToSubgroupBallot>(context);
164 
165   if (failed(applyPartialConversion(fn, *target, std::move(patterns))))
166     return signalPassFailure();
167 }
168 
169 ConvertToAtomCmpExchangeWeak::ConvertToAtomCmpExchangeWeak(MLIRContext *context)
170     : RewritePattern("test.convert_to_atomic_compare_exchange_weak_op", 1,
171                      context, {"spv.AtomicCompareExchangeWeak"}) {}
172 
173 LogicalResult
174 ConvertToAtomCmpExchangeWeak::matchAndRewrite(Operation *op,
175                                               PatternRewriter &rewriter) const {
176   Value ptr = op->getOperand(0);
177   Value value = op->getOperand(1);
178   Value comparator = op->getOperand(2);
179 
180   // Create a spv.AtomicCompareExchangeWeak op with AtomicCounterMemory bits in
181   // memory semantics to additionally require AtomicStorage capability.
182   rewriter.replaceOpWithNewOp<spirv::AtomicCompareExchangeWeakOp>(
183       op, value.getType(), ptr, spirv::Scope::Workgroup,
184       spirv::MemorySemantics::AcquireRelease |
185           spirv::MemorySemantics::AtomicCounterMemory,
186       spirv::MemorySemantics::Acquire, value, comparator);
187   return success();
188 }
189 
190 ConvertToBitReverse::ConvertToBitReverse(MLIRContext *context)
191     : RewritePattern("test.convert_to_bit_reverse_op", 1, context,
192                      {"spv.BitReverse"}) {}
193 
194 LogicalResult
195 ConvertToBitReverse::matchAndRewrite(Operation *op,
196                                      PatternRewriter &rewriter) const {
197   Value predicate = op->getOperand(0);
198 
199   rewriter.replaceOpWithNewOp<spirv::BitReverseOp>(
200       op, op->getResult(0).getType(), predicate);
201   return success();
202 }
203 
204 ConvertToGroupNonUniformBallot::ConvertToGroupNonUniformBallot(
205     MLIRContext *context)
206     : RewritePattern("test.convert_to_group_non_uniform_ballot_op", 1, context,
207                      {"spv.GroupNonUniformBallot"}) {}
208 
209 LogicalResult ConvertToGroupNonUniformBallot::matchAndRewrite(
210     Operation *op, PatternRewriter &rewriter) const {
211   Value predicate = op->getOperand(0);
212 
213   rewriter.replaceOpWithNewOp<spirv::GroupNonUniformBallotOp>(
214       op, op->getResult(0).getType(), spirv::Scope::Workgroup, predicate);
215   return success();
216 }
217 
218 ConvertToModule::ConvertToModule(MLIRContext *context)
219     : RewritePattern("test.convert_to_module_op", 1, context, {"spv.module"}) {}
220 
221 LogicalResult
222 ConvertToModule::matchAndRewrite(Operation *op,
223                                  PatternRewriter &rewriter) const {
224   rewriter.replaceOpWithNewOp<spirv::ModuleOp>(
225       op, spirv::AddressingModel::PhysicalStorageBuffer64,
226       spirv::MemoryModel::Vulkan);
227   return success();
228 }
229 
230 ConvertToSubgroupBallot::ConvertToSubgroupBallot(MLIRContext *context)
231     : RewritePattern("test.convert_to_subgroup_ballot_op", 1, context,
232                      {"spv.SubgroupBallotKHR"}) {}
233 
234 LogicalResult
235 ConvertToSubgroupBallot::matchAndRewrite(Operation *op,
236                                          PatternRewriter &rewriter) const {
237   Value predicate = op->getOperand(0);
238 
239   rewriter.replaceOpWithNewOp<spirv::SubgroupBallotKHROp>(
240       op, op->getResult(0).getType(), predicate);
241   return success();
242 }
243 
244 namespace mlir {
245 void registerConvertToTargetEnvPass() {
246   PassRegistration<ConvertToTargetEnv>();
247 }
248 } // namespace mlir
249