1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
|
//===- SPIRVWebGPUTransforms.cpp - WebGPU-specific transforms -------------===//
//
// 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
//
//===----------------------------------------------------------------------===//
//
// This file implements SPIR-V transforms used when targetting WebGPU.
//
//===----------------------------------------------------------------------===//
#include "mlir/Dialect/SPIRV/Transforms/SPIRVWebGPUTransforms.h"
#include "mlir/Dialect/SPIRV/IR/SPIRVOps.h"
#include "mlir/Dialect/SPIRV/Transforms/Passes.h"
#include "mlir/IR/BuiltinAttributes.h"
#include "mlir/IR/Location.h"
#include "mlir/IR/PatternMatch.h"
#include "mlir/IR/TypeUtilities.h"
#include "mlir/Support/LogicalResult.h"
#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
#include "llvm/ADT/ArrayRef.h"
#include "llvm/ADT/STLExtras.h"
#include "llvm/Support/FormatVariadic.h"
#include <array>
#include <cstdint>
namespace mlir {
namespace spirv {
#define GEN_PASS_DEF_SPIRVWEBGPUPREPAREPASS
#include "mlir/Dialect/SPIRV/Transforms/Passes.h.inc"
} // namespace spirv
} // namespace mlir
namespace mlir {
namespace spirv {
namespace {
//===----------------------------------------------------------------------===//
// Helpers
//===----------------------------------------------------------------------===//
Attribute getScalarOrSplatAttr(Type type, int64_t value) {
APInt sizedValue(getElementTypeOrSelf(type).getIntOrFloatBitWidth(), value);
if (auto intTy = dyn_cast<IntegerType>(type))
return IntegerAttr::get(intTy, sizedValue);
return SplatElementsAttr::get(cast<ShapedType>(type), sizedValue);
}
Value lowerExtendedMultiplication(Operation *mulOp, PatternRewriter &rewriter,
Value lhs, Value rhs,
bool signExtendArguments) {
Location loc = mulOp->getLoc();
Type argTy = lhs.getType();
// Emulate 64-bit multiplication by splitting each input element of type i32
// into 2 16-bit digits of type i32. This is so that the intermediate
// multiplications and additions do not overflow. We extract these 16-bit
// digits from i32 vector elements by masking (low digit) and shifting right
// (high digit).
//
// The multiplication algorithm used is the standard (long) multiplication.
// Multiplying two i32 integers produces 64 bits of result, i.e., 4 16-bit
// digits.
// - With zero-extended arguments, we end up emitting only 4 multiplications
// and 4 additions after constant folding.
// - With sign-extended arguments, we end up emitting 8 multiplications and
// and 12 additions after CSE.
Value cstLowMask = rewriter.create<ConstantOp>(
loc, lhs.getType(), getScalarOrSplatAttr(argTy, (1 << 16) - 1));
auto getLowDigit = [&rewriter, loc, cstLowMask](Value val) {
return rewriter.create<BitwiseAndOp>(loc, val, cstLowMask);
};
Value cst16 = rewriter.create<ConstantOp>(loc, lhs.getType(),
getScalarOrSplatAttr(argTy, 16));
auto getHighDigit = [&rewriter, loc, cst16](Value val) {
return rewriter.create<ShiftRightLogicalOp>(loc, val, cst16);
};
auto getSignDigit = [&rewriter, loc, cst16, &getHighDigit](Value val) {
// We only need to shift arithmetically by 15, but the extra
// sign-extension bit will be truncated by the logical shift, so this is
// fine. We do not have to introduce an extra constant since any
// value in [15, 32) would do.
return getHighDigit(
rewriter.create<ShiftRightArithmeticOp>(loc, val, cst16));
};
Value cst0 = rewriter.create<ConstantOp>(loc, lhs.getType(),
getScalarOrSplatAttr(argTy, 0));
Value lhsLow = getLowDigit(lhs);
Value lhsHigh = getHighDigit(lhs);
Value lhsExt = signExtendArguments ? getSignDigit(lhs) : cst0;
Value rhsLow = getLowDigit(rhs);
Value rhsHigh = getHighDigit(rhs);
Value rhsExt = signExtendArguments ? getSignDigit(rhs) : cst0;
std::array<Value, 4> lhsDigits = {lhsLow, lhsHigh, lhsExt, lhsExt};
std::array<Value, 4> rhsDigits = {rhsLow, rhsHigh, rhsExt, rhsExt};
std::array<Value, 4> resultDigits = {cst0, cst0, cst0, cst0};
for (auto [i, lhsDigit] : llvm::enumerate(lhsDigits)) {
for (auto [j, rhsDigit] : llvm::enumerate(rhsDigits)) {
if (i + j >= resultDigits.size())
continue;
if (lhsDigit == cst0 || rhsDigit == cst0)
continue;
Value &thisResDigit = resultDigits[i + j];
Value mul = rewriter.create<IMulOp>(loc, lhsDigit, rhsDigit);
Value current = rewriter.createOrFold<IAddOp>(loc, thisResDigit, mul);
thisResDigit = getLowDigit(current);
if (i + j + 1 != resultDigits.size()) {
Value &nextResDigit = resultDigits[i + j + 1];
Value carry = rewriter.createOrFold<IAddOp>(loc, nextResDigit,
getHighDigit(current));
nextResDigit = carry;
}
}
}
auto combineDigits = [loc, cst16, &rewriter](Value low, Value high) {
Value highBits = rewriter.create<ShiftLeftLogicalOp>(loc, high, cst16);
return rewriter.create<BitwiseOrOp>(loc, low, highBits);
};
Value low = combineDigits(resultDigits[0], resultDigits[1]);
Value high = combineDigits(resultDigits[2], resultDigits[3]);
return rewriter.create<CompositeConstructOp>(
loc, mulOp->getResultTypes().front(), llvm::ArrayRef({low, high}));
}
//===----------------------------------------------------------------------===//
// Rewrite Patterns
//===----------------------------------------------------------------------===//
template <typename MulExtendedOp, bool SignExtendArguments>
struct ExpandMulExtendedPattern final : OpRewritePattern<MulExtendedOp> {
using OpRewritePattern<MulExtendedOp>::OpRewritePattern;
LogicalResult matchAndRewrite(MulExtendedOp op,
PatternRewriter &rewriter) const override {
Location loc = op->getLoc();
Value lhs = op.getOperand1();
Value rhs = op.getOperand2();
// Currently, WGSL only supports 32-bit integer types. Any other integer
// types should already have been promoted/demoted to i32.
auto elemTy = cast<IntegerType>(getElementTypeOrSelf(lhs.getType()));
if (elemTy.getIntOrFloatBitWidth() != 32)
return rewriter.notifyMatchFailure(
loc,
llvm::formatv("Unexpected integer type for WebGPU: '{0}'", elemTy));
Value mul = lowerExtendedMultiplication(op, rewriter, lhs, rhs,
SignExtendArguments);
rewriter.replaceOp(op, mul);
return success();
}
};
using ExpandSMulExtendedPattern =
ExpandMulExtendedPattern<SMulExtendedOp, true>;
using ExpandUMulExtendedPattern =
ExpandMulExtendedPattern<UMulExtendedOp, false>;
//===----------------------------------------------------------------------===//
// Passes
//===----------------------------------------------------------------------===//
class WebGPUPreparePass
: public impl::SPIRVWebGPUPreparePassBase<WebGPUPreparePass> {
public:
void runOnOperation() override {
RewritePatternSet patterns(&getContext());
populateSPIRVExpandExtendedMultiplicationPatterns(patterns);
if (failed(
applyPatternsAndFoldGreedily(getOperation(), std::move(patterns))))
signalPassFailure();
}
};
} // namespace
//===----------------------------------------------------------------------===//
// Public Interface
//===----------------------------------------------------------------------===//
void populateSPIRVExpandExtendedMultiplicationPatterns(
RewritePatternSet &patterns) {
// WGSL currently does not support extended multiplication ops, see:
// https://github.com/gpuweb/gpuweb/issues/1565.
patterns.add<ExpandSMulExtendedPattern, ExpandUMulExtendedPattern>(
patterns.getContext());
}
} // namespace spirv
} // namespace mlir
|
