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
|
//===- TensorToSPIRV.cpp - Tensor to SPIR-V Patterns ----------------------===//
//
// 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 patterns to convert Tensor dialect to SPIR-V dialect.
//
//===----------------------------------------------------------------------===//
#include "mlir/Conversion/TensorToSPIRV/TensorToSPIRV.h"
#include "../SPIRVCommon/Pattern.h"
#include "mlir/Dialect/SPIRV/IR/SPIRVDialect.h"
#include "mlir/Dialect/SPIRV/IR/SPIRVOps.h"
#include "mlir/Dialect/SPIRV/Transforms/SPIRVConversion.h"
#include "mlir/Dialect/SPIRV/Utils/LayoutUtils.h"
#include "mlir/Dialect/Tensor/IR/Tensor.h"
#include "mlir/IR/AffineMap.h"
#include "mlir/Support/LogicalResult.h"
#include "llvm/Support/Debug.h"
#define DEBUG_TYPE "tensor-to-spirv-pattern"
using namespace mlir;
//===----------------------------------------------------------------------===//
// Operation conversion
//===----------------------------------------------------------------------===//
namespace {
/// Converts tensor.extract into loading using access chains from SPIR-V local
/// variables.
class TensorExtractPattern final
: public OpConversionPattern<tensor::ExtractOp> {
public:
TensorExtractPattern(TypeConverter &typeConverter, MLIRContext *context,
int64_t threshold, PatternBenefit benefit = 1)
: OpConversionPattern(typeConverter, context, benefit),
byteCountThreshold(threshold) {}
LogicalResult
matchAndRewrite(tensor::ExtractOp extractOp, OpAdaptor adaptor,
ConversionPatternRewriter &rewriter) const override {
auto tensorType = cast<RankedTensorType>(extractOp.getTensor().getType());
if (!tensorType.hasStaticShape())
return rewriter.notifyMatchFailure(extractOp, "non-static tensor");
if (tensorType.getNumElements() * tensorType.getElementTypeBitWidth() >
byteCountThreshold * 8)
return rewriter.notifyMatchFailure(extractOp,
"exceeding byte count threshold");
Location loc = extractOp.getLoc();
int64_t rank = tensorType.getRank();
SmallVector<int64_t, 4> strides(rank, 1);
for (int i = rank - 2; i >= 0; --i) {
strides[i] = strides[i + 1] * tensorType.getDimSize(i + 1);
}
Type varType = spirv::PointerType::get(adaptor.getTensor().getType(),
spirv::StorageClass::Function);
spirv::VariableOp varOp;
if (adaptor.getTensor().getDefiningOp<spirv::ConstantOp>()) {
// We could use the initializer directly; but certain driver compilers
// have bugs dealing with that. So for now, use spirv.Store for
// initialization.
varOp = rewriter.create<spirv::VariableOp>(loc, varType,
spirv::StorageClass::Function,
/*initializer=*/nullptr);
rewriter.create<spirv::StoreOp>(loc, varOp, adaptor.getTensor());
} else {
// Need to store the value to the local variable. It's questionable
// whether we want to support such case though.
return failure();
}
auto &typeConverter = *getTypeConverter<SPIRVTypeConverter>();
auto indexType = typeConverter.getIndexType();
Value index = spirv::linearizeIndex(adaptor.getIndices(), strides,
/*offset=*/0, indexType, loc, rewriter);
auto acOp = rewriter.create<spirv::AccessChainOp>(loc, varOp, index);
rewriter.replaceOpWithNewOp<spirv::LoadOp>(extractOp, acOp);
return success();
}
private:
int64_t byteCountThreshold;
};
} // namespace
//===----------------------------------------------------------------------===//
// Pattern population
//===----------------------------------------------------------------------===//
void mlir::populateTensorToSPIRVPatterns(SPIRVTypeConverter &typeConverter,
int64_t byteCountThreshold,
RewritePatternSet &patterns) {
patterns.add<TensorExtractPattern>(typeConverter, patterns.getContext(),
byteCountThreshold);
}
|
