MLIR 24.0.0git
ComplexToSPIRV.cpp
Go to the documentation of this file.
1//===- ComplexToSPIRV.cpp - Complex to SPIR-V Patterns --------------------===//
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// This file implements patterns to convert Complex dialect to SPIR-V dialect.
10//
11//===----------------------------------------------------------------------===//
12
18
19#define DEBUG_TYPE "complex-to-spirv-pattern"
20
21using namespace mlir;
22
23//===----------------------------------------------------------------------===//
24// Operation conversion
25//===----------------------------------------------------------------------===//
26
27namespace {
28/// Creates a scalar floating-point constant of the given value.
29Value createFPConstant(OpBuilder &builder, Location loc, Type type,
30 double value) {
31 return spirv::ConstantOp::create(
32 builder, loc, type, builder.getFloatAttr(cast<FloatType>(type), value));
33}
34
35/// Computes abs(re + i*im) as max * sqrt(1 + (min / max)^2) to avoid overflow
36/// in squaring. The result is NaN when both components are zero.
37template <typename SqrtOp, typename FAbsOp, typename FMaxOp, typename FMinOp>
38Value createComplexAbs(OpBuilder &builder, Location loc, Value re, Value im) {
39 Value absRe = FAbsOp::create(builder, loc, re);
40 Value absIm = FAbsOp::create(builder, loc, im);
41 Value max = FMaxOp::create(builder, loc, absRe, absIm);
42 Value min = FMinOp::create(builder, loc, absRe, absIm);
43 Value ratio = spirv::FDivOp::create(builder, loc, min, max);
44 Value ratioSq = spirv::FMulOp::create(builder, loc, ratio, ratio);
45 Value one = createFPConstant(builder, loc, re.getType(), 1.0);
46 Value ratioSqPlusOne = spirv::FAddOp::create(builder, loc, ratioSq, one);
47 Value sqrt = SqrtOp::create(builder, loc, ratioSqPlusOne);
48 return spirv::FMulOp::create(builder, loc, max, sqrt);
49}
50
51/// Returns true when both components are zero.
52Value createIsZero(OpBuilder &builder, Location loc, Value re, Value im) {
53 Value zero = createFPConstant(builder, loc, re.getType(), 0.0);
54 Value reIsZero = spirv::FOrdEqualOp::create(builder, loc, re, zero);
55 Value imIsZero = spirv::FOrdEqualOp::create(builder, loc, im, zero);
56 return spirv::LogicalAndOp::create(builder, loc, reIsZero, imIsZero);
57}
58
59struct ConstantOpPattern final : OpConversionPattern<complex::ConstantOp> {
60 using Base::Base;
61
62 LogicalResult
63 matchAndRewrite(complex::ConstantOp constOp, OpAdaptor adaptor,
64 ConversionPatternRewriter &rewriter) const override {
65 auto spirvType =
66 getTypeConverter()->convertType<ShapedType>(constOp.getType());
67 if (!spirvType)
68 return rewriter.notifyMatchFailure(constOp,
69 "unable to convert result type");
70
71 rewriter.replaceOpWithNewOp<spirv::ConstantOp>(
72 constOp, spirvType,
73 DenseElementsAttr::get(spirvType, constOp.getValue().getValue()));
74 return success();
75 }
76};
77
78struct CreateOpPattern final : OpConversionPattern<complex::CreateOp> {
79 using Base::Base;
80
81 LogicalResult
82 matchAndRewrite(complex::CreateOp createOp, OpAdaptor adaptor,
83 ConversionPatternRewriter &rewriter) const override {
84 Type spirvType = getTypeConverter()->convertType(createOp.getType());
85 if (!spirvType)
86 return rewriter.notifyMatchFailure(createOp,
87 "unable to convert result type");
88
89 rewriter.replaceOpWithNewOp<spirv::CompositeConstructOp>(
90 createOp, spirvType, adaptor.getOperands());
91 return success();
92 }
93};
94
95struct ReOpPattern final : OpConversionPattern<complex::ReOp> {
96 using Base::Base;
97
98 LogicalResult
99 matchAndRewrite(complex::ReOp reOp, OpAdaptor adaptor,
100 ConversionPatternRewriter &rewriter) const override {
101 Type spirvType = getTypeConverter()->convertType(reOp.getType());
102 if (!spirvType)
103 return rewriter.notifyMatchFailure(reOp, "unable to convert result type");
104
105 rewriter.replaceOpWithNewOp<spirv::CompositeExtractOp>(
106 reOp, adaptor.getComplex(), llvm::ArrayRef(0));
107 return success();
108 }
109};
110
111struct ImOpPattern final : OpConversionPattern<complex::ImOp> {
112 using Base::Base;
113
114 LogicalResult
115 matchAndRewrite(complex::ImOp imOp, OpAdaptor adaptor,
116 ConversionPatternRewriter &rewriter) const override {
117 Type spirvType = getTypeConverter()->convertType(imOp.getType());
118 if (!spirvType)
119 return rewriter.notifyMatchFailure(imOp, "unable to convert result type");
120
121 rewriter.replaceOpWithNewOp<spirv::CompositeExtractOp>(
122 imOp, adaptor.getComplex(), llvm::ArrayRef(1));
123 return success();
124 }
125};
126
127template <typename ComplexOp, typename SPIRVOp>
128struct ElementwiseBinaryOpPattern final : OpConversionPattern<ComplexOp> {
129 using OpConversionPattern<ComplexOp>::OpConversionPattern;
130 using OpAdaptor = typename ComplexOp::Adaptor;
131
132 LogicalResult
133 matchAndRewrite(ComplexOp op, OpAdaptor adaptor,
134 ConversionPatternRewriter &rewriter) const override {
135 Type spirvType =
136 this->getTypeConverter()->convertType(op.getResult().getType());
137 if (!spirvType)
138 return rewriter.notifyMatchFailure(op, "unable to convert result type");
139
140 Location loc = op.getLoc();
141 Value lhs = adaptor.getLhs();
142 Value rhs = adaptor.getRhs();
143
144 Value lhsRe = spirv::CompositeExtractOp::create(rewriter, loc, lhs, {0});
145 Value lhsIm = spirv::CompositeExtractOp::create(rewriter, loc, lhs, {1});
146 Value rhsRe = spirv::CompositeExtractOp::create(rewriter, loc, rhs, {0});
147 Value rhsIm = spirv::CompositeExtractOp::create(rewriter, loc, rhs, {1});
148
149 Value resultRe = SPIRVOp::create(rewriter, loc, lhsRe, rhsRe);
150 Value resultIm = SPIRVOp::create(rewriter, loc, lhsIm, rhsIm);
151
152 rewriter.replaceOpWithNewOp<spirv::CompositeConstructOp>(
153 op, spirvType, llvm::ArrayRef<Value>{resultRe, resultIm});
154 return success();
155 }
156};
157
158template <typename ComplexOp, typename SPIRVCompareOp, typename SPIRVCombinerOp>
159struct ComparisonOpPattern final : OpConversionPattern<ComplexOp> {
160 using OpConversionPattern<ComplexOp>::OpConversionPattern;
161 using OpAdaptor = typename ComplexOp::Adaptor;
162
163 LogicalResult
164 matchAndRewrite(ComplexOp op, OpAdaptor adaptor,
165 ConversionPatternRewriter &rewriter) const override {
166 Type spirvType =
167 this->getTypeConverter()->convertType(op.getResult().getType());
168 if (!spirvType)
169 return rewriter.notifyMatchFailure(op, "unable to convert result type");
170
171 Location loc = op.getLoc();
172 Value lhs = adaptor.getLhs();
173 Value rhs = adaptor.getRhs();
174
175 Value lhsRe = spirv::CompositeExtractOp::create(rewriter, loc, lhs, {0});
176 Value lhsIm = spirv::CompositeExtractOp::create(rewriter, loc, lhs, {1});
177 Value rhsRe = spirv::CompositeExtractOp::create(rewriter, loc, rhs, {0});
178 Value rhsIm = spirv::CompositeExtractOp::create(rewriter, loc, rhs, {1});
179
180 Value cmpRe = SPIRVCompareOp::create(rewriter, loc, lhsRe, rhsRe);
181 Value cmpIm = SPIRVCompareOp::create(rewriter, loc, lhsIm, rhsIm);
182
183 rewriter.replaceOpWithNewOp<SPIRVCombinerOp>(op, spirvType, cmpRe, cmpIm);
184 return success();
185 }
186};
187
188struct MulOpPattern final : OpConversionPattern<complex::MulOp> {
189 using Base::Base;
190
191 LogicalResult
192 matchAndRewrite(complex::MulOp op, OpAdaptor adaptor,
193 ConversionPatternRewriter &rewriter) const override {
194 Type spirvType = getTypeConverter()->convertType(op.getResult().getType());
195 if (!spirvType)
196 return rewriter.notifyMatchFailure(op, "unable to convert result type");
197
198 Location loc = op.getLoc();
199 Value lhs = adaptor.getLhs();
200 Value rhs = adaptor.getRhs();
201
202 Value a = spirv::CompositeExtractOp::create(rewriter, loc, lhs, {0});
203 Value b = spirv::CompositeExtractOp::create(rewriter, loc, lhs, {1});
204 Value c = spirv::CompositeExtractOp::create(rewriter, loc, rhs, {0});
205 Value d = spirv::CompositeExtractOp::create(rewriter, loc, rhs, {1});
206
207 Value ac = spirv::FMulOp::create(rewriter, loc, a, c);
208 Value bd = spirv::FMulOp::create(rewriter, loc, b, d);
209 Value ad = spirv::FMulOp::create(rewriter, loc, a, d);
210 Value bc = spirv::FMulOp::create(rewriter, loc, b, c);
211 Value resultRe = spirv::FSubOp::create(rewriter, loc, ac, bd);
212 Value resultIm = spirv::FAddOp::create(rewriter, loc, ad, bc);
213
214 rewriter.replaceOpWithNewOp<spirv::CompositeConstructOp>(
215 op, spirvType, llvm::ArrayRef<Value>{resultRe, resultIm});
216 return success();
217 }
218};
219
220template <typename SqrtOp>
221struct AbsOpPattern final : OpConversionPattern<complex::AbsOp> {
222 using OpConversionPattern<complex::AbsOp>::OpConversionPattern;
223
224 LogicalResult
225 matchAndRewrite(complex::AbsOp op, OpAdaptor adaptor,
226 ConversionPatternRewriter &rewriter) const override {
227 Type spirvType =
228 this->getTypeConverter()->convertType(op.getResult().getType());
229 if (!spirvType)
230 return rewriter.notifyMatchFailure(op, "unable to convert result type");
231
232 Location loc = op.getLoc();
233 Value complexVal = adaptor.getComplex();
234
235 Value re =
236 spirv::CompositeExtractOp::create(rewriter, loc, complexVal, {0});
237 Value im =
238 spirv::CompositeExtractOp::create(rewriter, loc, complexVal, {1});
239
240 Value reSq = spirv::FMulOp::create(rewriter, loc, re, re);
241 Value imSq = spirv::FMulOp::create(rewriter, loc, im, im);
242 Value sum = spirv::FAddOp::create(rewriter, loc, reSq, imSq);
243
244 rewriter.replaceOpWithNewOp<SqrtOp>(op, sum);
245 return success();
246 }
247};
248
249template <typename ComplexOp, bool NegateReal>
250struct NegationOpPattern final : OpConversionPattern<ComplexOp> {
251 using OpConversionPattern<ComplexOp>::OpConversionPattern;
252 using OpAdaptor = typename ComplexOp::Adaptor;
253
254 LogicalResult
255 matchAndRewrite(ComplexOp op, OpAdaptor adaptor,
256 ConversionPatternRewriter &rewriter) const override {
257 Type spirvType =
258 this->getTypeConverter()->convertType(op.getResult().getType());
259 if (!spirvType)
260 return rewriter.notifyMatchFailure(op, "unable to convert result type");
261
262 Location loc = op.getLoc();
263 Value complexVal = adaptor.getComplex();
264
265 Value re =
266 spirv::CompositeExtractOp::create(rewriter, loc, complexVal, {0});
267 Value im =
268 spirv::CompositeExtractOp::create(rewriter, loc, complexVal, {1});
269
270 Value resultRe =
271 NegateReal ? spirv::FNegateOp::create(rewriter, loc, re) : re;
272 Value resultIm = spirv::FNegateOp::create(rewriter, loc, im);
273
274 rewriter.replaceOpWithNewOp<spirv::CompositeConstructOp>(
275 op, spirvType, llvm::ArrayRef<Value>{resultRe, resultIm});
276 return success();
277 }
278};
279
280struct DivOpPattern final : OpConversionPattern<complex::DivOp> {
281 using Base::Base;
282
283 LogicalResult
284 matchAndRewrite(complex::DivOp op, OpAdaptor adaptor,
285 ConversionPatternRewriter &rewriter) const override {
286 Type spirvType = getTypeConverter()->convertType(op.getResult().getType());
287 if (!spirvType)
288 return rewriter.notifyMatchFailure(op, "unable to convert result type");
289
290 Location loc = op.getLoc();
291 Value lhs = adaptor.getLhs();
292 Value rhs = adaptor.getRhs();
293
294 Value a = spirv::CompositeExtractOp::create(rewriter, loc, lhs, {0});
295 Value b = spirv::CompositeExtractOp::create(rewriter, loc, lhs, {1});
296 Value c = spirv::CompositeExtractOp::create(rewriter, loc, rhs, {0});
297 Value d = spirv::CompositeExtractOp::create(rewriter, loc, rhs, {1});
298
299 Value ac = spirv::FMulOp::create(rewriter, loc, a, c);
300 Value bd = spirv::FMulOp::create(rewriter, loc, b, d);
301 Value bc = spirv::FMulOp::create(rewriter, loc, b, c);
302 Value ad = spirv::FMulOp::create(rewriter, loc, a, d);
303 Value cc = spirv::FMulOp::create(rewriter, loc, c, c);
304 Value dd = spirv::FMulOp::create(rewriter, loc, d, d);
305 Value denom = spirv::FAddOp::create(rewriter, loc, cc, dd);
306 Value numRe = spirv::FAddOp::create(rewriter, loc, ac, bd);
307 Value numIm = spirv::FSubOp::create(rewriter, loc, bc, ad);
308 Value resultRe = spirv::FDivOp::create(rewriter, loc, numRe, denom);
309 Value resultIm = spirv::FDivOp::create(rewriter, loc, numIm, denom);
310
311 rewriter.replaceOpWithNewOp<spirv::CompositeConstructOp>(
312 op, spirvType, llvm::ArrayRef<Value>{resultRe, resultIm});
313 return success();
314 }
315};
316
317template <typename Atan2Op>
318struct AngleOpPattern final : OpConversionPattern<complex::AngleOp> {
319 using OpConversionPattern<complex::AngleOp>::OpConversionPattern;
320
321 LogicalResult
322 matchAndRewrite(complex::AngleOp op, OpAdaptor adaptor,
323 ConversionPatternRewriter &rewriter) const override {
324 Type spirvType =
325 this->getTypeConverter()->convertType(op.getResult().getType());
326 if (!spirvType)
327 return rewriter.notifyMatchFailure(op, "unable to convert result type");
328
329 Location loc = op.getLoc();
330 Value complexVal = adaptor.getComplex();
331
332 Value re =
333 spirv::CompositeExtractOp::create(rewriter, loc, complexVal, {0});
334 Value im =
335 spirv::CompositeExtractOp::create(rewriter, loc, complexVal, {1});
336
337 rewriter.replaceOpWithNewOp<Atan2Op>(op, im, re);
338 return success();
339 }
340};
341
342/// Computes complex.sign as z / abs(z), selecting a zero result when z is
343/// zero to avoid a division by zero. abs(z) is computed as
344/// max * sqrt(1 + (min / max)^2) to avoid overflow in squaring.
345template <typename SqrtOp, typename FAbsOp, typename FMaxOp, typename FMinOp>
346struct SignOpPattern final : OpConversionPattern<complex::SignOp> {
347 using OpConversionPattern<complex::SignOp>::OpConversionPattern;
348
349 LogicalResult
350 matchAndRewrite(complex::SignOp op, OpAdaptor adaptor,
351 ConversionPatternRewriter &rewriter) const override {
352 Type spirvType =
353 this->getTypeConverter()->convertType(op.getResult().getType());
354 if (!spirvType)
355 return rewriter.notifyMatchFailure(op, "unable to convert result type");
356
357 Location loc = op.getLoc();
358 Value complexVal = adaptor.getComplex();
359
360 Value re =
361 spirv::CompositeExtractOp::create(rewriter, loc, complexVal, {0});
362 Value im =
363 spirv::CompositeExtractOp::create(rewriter, loc, complexVal, {1});
364
365 Value abs =
366 createComplexAbs<SqrtOp, FAbsOp, FMaxOp, FMinOp>(rewriter, loc, re, im);
367
368 Value signRe = spirv::FDivOp::create(rewriter, loc, re, abs);
369 Value signIm = spirv::FDivOp::create(rewriter, loc, im, abs);
370
371 Value isZero = createIsZero(rewriter, loc, re, im);
372
373 // Select per component: spirv.Select with a scalar condition and a
374 // composite result requires SPIR-V 1.4, so operate on scalars here and
375 // assemble the result afterwards.
376 Value resultRe = spirv::SelectOp::create(rewriter, loc, isZero, re, signRe);
377 Value resultIm = spirv::SelectOp::create(rewriter, loc, isZero, im, signIm);
378
379 rewriter.replaceOpWithNewOp<spirv::CompositeConstructOp>(
380 op, spirvType, llvm::ArrayRef<Value>{resultRe, resultIm});
381 return success();
382 }
383};
384
385enum class SinCos { Sin, Cos };
386
387/// Lowers complex.sin/cos following the identities
388/// sin(x + iy) = sin(x) * cosh(y) + i * cos(x) * sinh(y)
389/// cos(x + iy) = cos(x) * cosh(y) - i * sin(x) * sinh(y)
390template <SinCos Kind, typename SinOp, typename CosOp, typename SinhOp,
391 typename CoshOp,
392 typename ComplexOp = std::conditional_t<
393 Kind == SinCos::Sin, complex::SinOp, complex::CosOp>>
394struct TrigonometricOpPattern final : OpConversionPattern<ComplexOp> {
395 using OpConversionPattern<ComplexOp>::OpConversionPattern;
396 using OpAdaptor = typename ComplexOp::Adaptor;
397
398 LogicalResult
399 matchAndRewrite(ComplexOp op, OpAdaptor adaptor,
400 ConversionPatternRewriter &rewriter) const override {
401 Type spirvType =
402 this->getTypeConverter()->convertType(op.getResult().getType());
403 if (!spirvType)
404 return rewriter.notifyMatchFailure(op, "unable to convert result type");
405
406 Location loc = op.getLoc();
407 Value complexVal = adaptor.getComplex();
408
409 Value re =
410 spirv::CompositeExtractOp::create(rewriter, loc, complexVal, {0});
411 Value im =
412 spirv::CompositeExtractOp::create(rewriter, loc, complexVal, {1});
413
414 Value sin = SinOp::create(rewriter, loc, re);
415 Value cos = CosOp::create(rewriter, loc, re);
416 Value sinh = SinhOp::create(rewriter, loc, im);
417 Value cosh = CoshOp::create(rewriter, loc, im);
418
419 bool isSin = Kind == SinCos::Sin;
420 Value resultRe =
421 spirv::FMulOp::create(rewriter, loc, isSin ? sin : cos, cosh);
422 Value resultIm =
423 spirv::FMulOp::create(rewriter, loc, isSin ? cos : sin, sinh);
424 if (!isSin)
425 resultIm = spirv::FNegateOp::create(rewriter, loc, resultIm);
426
427 rewriter.replaceOpWithNewOp<spirv::CompositeConstructOp>(
428 op, spirvType, llvm::ArrayRef<Value>{resultRe, resultIm});
429 return success();
430 }
431};
432} // namespace
433
434//===----------------------------------------------------------------------===//
435// Pattern population
436//===----------------------------------------------------------------------===//
437
439 const SPIRVTypeConverter &typeConverter, RewritePatternSet &patterns) {
440 MLIRContext *context = patterns.getContext();
441
442 patterns
443 .add<ConstantOpPattern, CreateOpPattern, ReOpPattern, ImOpPattern,
444 ElementwiseBinaryOpPattern<complex::AddOp, spirv::FAddOp>,
445 ElementwiseBinaryOpPattern<complex::SubOp, spirv::FSubOp>,
446 ComparisonOpPattern<complex::EqualOp, spirv::FOrdEqualOp,
447 spirv::LogicalAndOp>,
448 ComparisonOpPattern<complex::NotEqualOp, spirv::FUnordNotEqualOp,
449 spirv::LogicalOrOp>,
450 MulOpPattern, DivOpPattern,
451 NegationOpPattern<complex::NegOp, /*NegateReal=*/true>,
452 NegationOpPattern<complex::ConjOp, /*NegateReal=*/false>,
453 AbsOpPattern<spirv::GLSqrtOp>, AbsOpPattern<spirv::CLSqrtOp>,
454 AngleOpPattern<spirv::GLAtan2Op>, AngleOpPattern<spirv::CLAtan2Op>,
455 SignOpPattern<spirv::GLSqrtOp, spirv::GLFAbsOp, spirv::GLFMaxOp,
456 spirv::GLFMinOp>,
457 SignOpPattern<spirv::CLSqrtOp, spirv::CLFAbsOp, spirv::CLFMaxOp,
458 spirv::CLFMinOp>,
459 TrigonometricOpPattern<SinCos::Sin, spirv::GLSinOp, spirv::GLCosOp,
460 spirv::GLSinhOp, spirv::GLCoshOp>,
461 TrigonometricOpPattern<SinCos::Sin, spirv::CLSinOp, spirv::CLCosOp,
462 spirv::CLSinhOp, spirv::CLCoshOp>,
463 TrigonometricOpPattern<SinCos::Cos, spirv::GLSinOp, spirv::GLCosOp,
464 spirv::GLSinhOp, spirv::GLCoshOp>,
465 TrigonometricOpPattern<SinCos::Cos, spirv::CLSinOp, spirv::CLCosOp,
466 spirv::CLSinhOp, spirv::CLCoshOp>>(
467 typeConverter, context);
468}
return success()
b
Return true if permutation is a valid permutation of the outer_dims_perm (case OuterOrInnerPerm::Oute...
static Value max(ImplicitLocOpBuilder &builder, Value value, Value bound)
static Value min(ImplicitLocOpBuilder &builder, Value value, Value bound)
static Value createFPConstant(Location loc, Type srcType, Type dstType, PatternRewriter &rewriter, double value)
Creates llvm.mlir.constant with a floating-point scalar or vector value.
FloatAttr getFloatAttr(Type type, double value)
Definition Builders.cpp:263
static DenseElementsAttr get(ShapedType type, ArrayRef< Attribute > values)
Constructs a dense elements attribute from an array of element values.
This class defines the main interface for locations in MLIR and acts as a non-nullable wrapper around...
Definition Location.h:76
MLIRContext is the top-level object for a collection of MLIR operations.
Definition MLIRContext.h:71
This class helps build Operations.
Definition Builders.h:210
MLIRContext * getContext() const
RewritePatternSet & add(ConstructorArg &&arg, ConstructorArgs &&...args)
Add an instance of each of the pattern types 'Ts' to the pattern list with the given arguments.
Type conversion from builtin types to SPIR-V types for shader interface.
Instances of the Type class are uniqued, have an immutable identifier and an optional mutable compone...
Definition Types.h:74
This class represents an instance of an SSA value in the MLIR system, representing a computable value...
Definition Value.h:96
Type getType() const
Return the type of this value.
Definition Value.h:105
Include the generated interface declarations.
void populateComplexToSPIRVPatterns(const SPIRVTypeConverter &typeConverter, RewritePatternSet &patterns)
Appends to a pattern list additional patterns for translating Complex ops to SPIR-V ops.