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/SPIRV/SPIRVLowering.h"
10 #include "mlir/Dialect/SPIRV/SPIRVOps.h"
11 #include "mlir/Dialect/SPIRV/SPIRVTypes.h"
12 #include "mlir/IR/Function.h"
13 #include "mlir/Pass/Pass.h"
14 
15 using namespace mlir;
16 
17 //===----------------------------------------------------------------------===//
18 // Printing op availability pass
19 //===----------------------------------------------------------------------===//
20 
21 namespace {
22 /// A pass for testing SPIR-V op availability.
23 struct PrintOpAvailability : public FunctionPass<PrintOpAvailability> {
24   void runOnFunction() override;
25 };
26 } // end anonymous namespace
27 
28 void PrintOpAvailability::runOnFunction() {
29   auto f = getFunction();
30   llvm::outs() << f.getName() << "\n";
31 
32   Dialect *spvDialect = getContext().getRegisteredDialect("spv");
33 
34   f.getOperation()->walk([&](Operation *op) {
35     if (op->getDialect() != spvDialect)
36       return WalkResult::advance();
37 
38     auto opName = op->getName();
39     auto &os = llvm::outs();
40 
41     if (auto minVersion = dyn_cast<spirv::QueryMinVersionInterface>(op))
42       os << opName << " min version: "
43          << spirv::stringifyVersion(minVersion.getMinVersion()) << "\n";
44 
45     if (auto maxVersion = dyn_cast<spirv::QueryMaxVersionInterface>(op))
46       os << opName << " max version: "
47          << spirv::stringifyVersion(maxVersion.getMaxVersion()) << "\n";
48 
49     if (auto extension = dyn_cast<spirv::QueryExtensionInterface>(op)) {
50       os << opName << " extensions: [";
51       for (const auto &exts : extension.getExtensions()) {
52         os << " [";
53         interleaveComma(exts, os, [&](spirv::Extension ext) {
54           os << spirv::stringifyExtension(ext);
55         });
56         os << "]";
57       }
58       os << " ]\n";
59     }
60 
61     if (auto capability = dyn_cast<spirv::QueryCapabilityInterface>(op)) {
62       os << opName << " capabilities: [";
63       for (const auto &caps : capability.getCapabilities()) {
64         os << " [";
65         interleaveComma(caps, os, [&](spirv::Capability cap) {
66           os << spirv::stringifyCapability(cap);
67         });
68         os << "]";
69       }
70       os << " ]\n";
71     }
72     os.flush();
73 
74     return WalkResult::advance();
75   });
76 }
77 
78 namespace mlir {
79 void registerPrintOpAvailabilityPass() {
80   PassRegistration<PrintOpAvailability> printOpAvailabilityPass(
81       "test-spirv-op-availability", "Test SPIR-V op availability");
82 }
83 } // namespace mlir
84 
85 //===----------------------------------------------------------------------===//
86 // Converting target environment pass
87 //===----------------------------------------------------------------------===//
88 
89 namespace {
90 /// A pass for testing SPIR-V op availability.
91 struct ConvertToTargetEnv : public FunctionPass<ConvertToTargetEnv> {
92   void runOnFunction() override;
93 };
94 
95 struct ConvertToAtomCmpExchangeWeak : public RewritePattern {
96   ConvertToAtomCmpExchangeWeak(MLIRContext *context);
97   PatternMatchResult matchAndRewrite(Operation *op,
98                                      PatternRewriter &rewriter) const override;
99 };
100 
101 struct ConvertToBitReverse : public RewritePattern {
102   ConvertToBitReverse(MLIRContext *context);
103   PatternMatchResult matchAndRewrite(Operation *op,
104                                      PatternRewriter &rewriter) const override;
105 };
106 
107 struct ConvertToGroupNonUniformBallot : public RewritePattern {
108   ConvertToGroupNonUniformBallot(MLIRContext *context);
109   PatternMatchResult matchAndRewrite(Operation *op,
110                                      PatternRewriter &rewriter) const override;
111 };
112 
113 struct ConvertToModule : public RewritePattern {
114   ConvertToModule(MLIRContext *context);
115   PatternMatchResult matchAndRewrite(Operation *op,
116                                      PatternRewriter &rewriter) const override;
117 };
118 
119 struct ConvertToSubgroupBallot : public RewritePattern {
120   ConvertToSubgroupBallot(MLIRContext *context);
121   PatternMatchResult matchAndRewrite(Operation *op,
122                                      PatternRewriter &rewriter) const override;
123 };
124 } // end anonymous namespace
125 
126 void ConvertToTargetEnv::runOnFunction() {
127   MLIRContext *context = &getContext();
128   FuncOp fn = getFunction();
129 
130   auto targetEnv = fn.getOperation()
131                        ->getAttr(spirv::getTargetEnvAttrName())
132                        .cast<spirv::TargetEnvAttr>();
133   auto target = spirv::SPIRVConversionTarget::get(targetEnv, context);
134 
135   OwningRewritePatternList patterns;
136   patterns.insert<ConvertToAtomCmpExchangeWeak, ConvertToBitReverse,
137                   ConvertToGroupNonUniformBallot, ConvertToModule,
138                   ConvertToSubgroupBallot>(context);
139 
140   if (failed(applyPartialConversion(fn, *target, patterns)))
141     return signalPassFailure();
142 }
143 
144 ConvertToAtomCmpExchangeWeak::ConvertToAtomCmpExchangeWeak(MLIRContext *context)
145     : RewritePattern("test.convert_to_atomic_compare_exchange_weak_op",
146                      {"spv.AtomicCompareExchangeWeak"}, 1, context) {}
147 
148 PatternMatchResult
149 ConvertToAtomCmpExchangeWeak::matchAndRewrite(Operation *op,
150                                               PatternRewriter &rewriter) const {
151   Value ptr = op->getOperand(0);
152   Value value = op->getOperand(1);
153   Value comparator = op->getOperand(2);
154 
155   // Create a spv.AtomicCompareExchangeWeak op with AtomicCounterMemory bits in
156   // memory semantics to additionally require AtomicStorage capability.
157   rewriter.replaceOpWithNewOp<spirv::AtomicCompareExchangeWeakOp>(
158       op, value.getType(), ptr, spirv::Scope::Workgroup,
159       spirv::MemorySemantics::AcquireRelease |
160           spirv::MemorySemantics::AtomicCounterMemory,
161       spirv::MemorySemantics::Acquire, value, comparator);
162   return matchSuccess();
163 }
164 
165 ConvertToBitReverse::ConvertToBitReverse(MLIRContext *context)
166     : RewritePattern("test.convert_to_bit_reverse_op", {"spv.BitReverse"}, 1,
167                      context) {}
168 
169 PatternMatchResult
170 ConvertToBitReverse::matchAndRewrite(Operation *op,
171                                      PatternRewriter &rewriter) const {
172   Value predicate = op->getOperand(0);
173 
174   rewriter.replaceOpWithNewOp<spirv::BitReverseOp>(
175       op, op->getResult(0).getType(), predicate);
176   return matchSuccess();
177 }
178 
179 ConvertToGroupNonUniformBallot::ConvertToGroupNonUniformBallot(
180     MLIRContext *context)
181     : RewritePattern("test.convert_to_group_non_uniform_ballot_op",
182                      {"spv.GroupNonUniformBallot"}, 1, context) {}
183 
184 PatternMatchResult ConvertToGroupNonUniformBallot::matchAndRewrite(
185     Operation *op, PatternRewriter &rewriter) const {
186   Value predicate = op->getOperand(0);
187 
188   rewriter.replaceOpWithNewOp<spirv::GroupNonUniformBallotOp>(
189       op, op->getResult(0).getType(), spirv::Scope::Workgroup, predicate);
190   return matchSuccess();
191 }
192 
193 ConvertToModule::ConvertToModule(MLIRContext *context)
194     : RewritePattern("test.convert_to_module_op", {"spv.module"}, 1, context) {}
195 
196 PatternMatchResult
197 ConvertToModule::matchAndRewrite(Operation *op,
198                                  PatternRewriter &rewriter) const {
199   rewriter.replaceOpWithNewOp<spirv::ModuleOp>(
200       op, spirv::AddressingModel::PhysicalStorageBuffer64,
201       spirv::MemoryModel::Vulkan);
202   return matchSuccess();
203 }
204 
205 ConvertToSubgroupBallot::ConvertToSubgroupBallot(MLIRContext *context)
206     : RewritePattern("test.convert_to_subgroup_ballot_op",
207                      {"spv.SubgroupBallotKHR"}, 1, context) {}
208 
209 PatternMatchResult
210 ConvertToSubgroupBallot::matchAndRewrite(Operation *op,
211                                          PatternRewriter &rewriter) const {
212   Value predicate = op->getOperand(0);
213 
214   rewriter.replaceOpWithNewOp<spirv::SubgroupBallotKHROp>(
215       op, op->getResult(0).getType(), predicate);
216   return matchSuccess();
217 }
218 
219 namespace mlir {
220 void registerConvertToTargetEnvPass() {
221   PassRegistration<ConvertToTargetEnv> convertToTargetEnvPass(
222       "test-spirv-target-env", "Test SPIR-V target environment");
223 }
224 } // namespace mlir
225