1 //===- TensorToSPIRV.cpp - Tensor 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 Tensor dialect to SPIR-V dialect.
10 //
11 //===----------------------------------------------------------------------===//
12 
13 #include "mlir/Conversion/TensorToSPIRV/TensorToSPIRV.h"
14 #include "../SPIRVCommon/Pattern.h"
15 #include "mlir/Dialect/SPIRV/IR/SPIRVDialect.h"
16 #include "mlir/Dialect/SPIRV/IR/SPIRVOps.h"
17 #include "mlir/Dialect/SPIRV/Transforms/SPIRVConversion.h"
18 #include "mlir/Dialect/SPIRV/Utils/LayoutUtils.h"
19 #include "mlir/Dialect/Tensor/IR/Tensor.h"
20 #include "mlir/IR/AffineMap.h"
21 #include "mlir/Support/LogicalResult.h"
22 #include "llvm/ADT/SetVector.h"
23 #include "llvm/Support/Debug.h"
24 
25 #define DEBUG_TYPE "tensor-to-spirv-pattern"
26 
27 using namespace mlir;
28 
29 //===----------------------------------------------------------------------===//
30 // Operation conversion
31 //===----------------------------------------------------------------------===//
32 
33 namespace {
34 
35 /// Converts tensor.extract into loading using access chains from SPIR-V local
36 /// variables.
37 class TensorExtractPattern final
38     : public OpConversionPattern<tensor::ExtractOp> {
39 public:
40   TensorExtractPattern(TypeConverter &typeConverter, MLIRContext *context,
41                        int64_t threshold, PatternBenefit benefit = 1)
42       : OpConversionPattern(typeConverter, context, benefit),
43         byteCountThreshold(threshold) {}
44 
45   LogicalResult
46   matchAndRewrite(tensor::ExtractOp extractOp, OpAdaptor adaptor,
47                   ConversionPatternRewriter &rewriter) const override {
48     TensorType tensorType = extractOp.tensor().getType().cast<TensorType>();
49 
50     if (!tensorType.hasStaticShape())
51       return rewriter.notifyMatchFailure(extractOp, "non-static tensor");
52 
53     if (tensorType.getNumElements() * tensorType.getElementTypeBitWidth() >
54         byteCountThreshold * 8)
55       return rewriter.notifyMatchFailure(extractOp,
56                                          "exceeding byte count threshold");
57 
58     Location loc = extractOp.getLoc();
59 
60     int64_t rank = tensorType.getRank();
61     SmallVector<int64_t, 4> strides(rank, 1);
62     for (int i = rank - 2; i >= 0; --i) {
63       strides[i] = strides[i + 1] * tensorType.getDimSize(i + 1);
64     }
65 
66     Type varType = spirv::PointerType::get(adaptor.tensor().getType(),
67                                            spirv::StorageClass::Function);
68 
69     spirv::VariableOp varOp;
70     if (adaptor.tensor().getDefiningOp<spirv::ConstantOp>()) {
71       varOp = rewriter.create<spirv::VariableOp>(
72           loc, varType, spirv::StorageClass::Function,
73           /*initializer=*/adaptor.tensor());
74     } else {
75       // Need to store the value to the local variable. It's questionable
76       // whether we want to support such case though.
77       return failure();
78     }
79 
80     auto &typeConverter = *getTypeConverter<SPIRVTypeConverter>();
81     auto indexType = typeConverter.getIndexType();
82 
83     Value index = spirv::linearizeIndex(adaptor.indices(), strides,
84                                         /*offset=*/0, indexType, loc, rewriter);
85     auto acOp = rewriter.create<spirv::AccessChainOp>(loc, varOp, index);
86 
87     rewriter.replaceOpWithNewOp<spirv::LoadOp>(extractOp, acOp);
88 
89     return success();
90   }
91 
92 private:
93   int64_t byteCountThreshold;
94 };
95 
96 } // namespace
97 
98 //===----------------------------------------------------------------------===//
99 // Pattern population
100 //===----------------------------------------------------------------------===//
101 
102 void mlir::populateTensorToSPIRVPatterns(SPIRVTypeConverter &typeConverter,
103                                          int64_t byteCountThreshold,
104                                          RewritePatternSet &patterns) {
105   patterns.add<TensorExtractPattern>(typeConverter, patterns.getContext(),
106                                      byteCountThreshold);
107 }
108