MLIR  19.0.0git
TensorToSPIRV.cpp
Go to the documentation of this file.
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 
14 #include "../SPIRVCommon/Pattern.h"
20 #include "mlir/IR/AffineMap.h"
22 #include "llvm/Support/Debug.h"
23 
24 #define DEBUG_TYPE "tensor-to-spirv-pattern"
25 
26 using namespace mlir;
27 
28 //===----------------------------------------------------------------------===//
29 // Operation conversion
30 //===----------------------------------------------------------------------===//
31 
32 namespace {
33 
34 /// Converts tensor.extract into loading using access chains from SPIR-V local
35 /// variables.
36 class TensorExtractPattern final
37  : public OpConversionPattern<tensor::ExtractOp> {
38 public:
39  TensorExtractPattern(TypeConverter &typeConverter, MLIRContext *context,
40  int64_t threshold, PatternBenefit benefit = 1)
41  : OpConversionPattern(typeConverter, context, benefit),
42  byteCountThreshold(threshold) {}
43 
45  matchAndRewrite(tensor::ExtractOp extractOp, OpAdaptor adaptor,
46  ConversionPatternRewriter &rewriter) const override {
47  auto tensorType = cast<RankedTensorType>(extractOp.getTensor().getType());
48 
49  if (!tensorType.hasStaticShape())
50  return rewriter.notifyMatchFailure(extractOp, "non-static tensor");
51 
52  if (tensorType.getNumElements() * tensorType.getElementTypeBitWidth() >
53  byteCountThreshold * 8)
54  return rewriter.notifyMatchFailure(extractOp,
55  "exceeding byte count threshold");
56 
57  Location loc = extractOp.getLoc();
58 
59  int64_t rank = tensorType.getRank();
60  SmallVector<int64_t, 4> strides(rank, 1);
61  for (int i = rank - 2; i >= 0; --i) {
62  strides[i] = strides[i + 1] * tensorType.getDimSize(i + 1);
63  }
64 
65  Type varType = spirv::PointerType::get(adaptor.getTensor().getType(),
66  spirv::StorageClass::Function);
67 
68  spirv::VariableOp varOp;
69  if (adaptor.getTensor().getDefiningOp<spirv::ConstantOp>()) {
70  // We could use the initializer directly; but certain driver compilers
71  // have bugs dealing with that. So for now, use spirv.Store for
72  // initialization.
73  varOp = rewriter.create<spirv::VariableOp>(loc, varType,
74  spirv::StorageClass::Function,
75  /*initializer=*/nullptr);
76  rewriter.create<spirv::StoreOp>(loc, varOp, adaptor.getTensor());
77  } else {
78  // Need to store the value to the local variable. It's questionable
79  // whether we want to support such case though.
80  return failure();
81  }
82 
83  auto &typeConverter = *getTypeConverter<SPIRVTypeConverter>();
84  auto indexType = typeConverter.getIndexType();
85 
86  Value index = spirv::linearizeIndex(adaptor.getIndices(), strides,
87  /*offset=*/0, indexType, loc, rewriter);
88  auto acOp = rewriter.create<spirv::AccessChainOp>(loc, varOp, index);
89 
90  rewriter.replaceOpWithNewOp<spirv::LoadOp>(extractOp, acOp);
91 
92  return success();
93  }
94 
95 private:
96  int64_t byteCountThreshold;
97 };
98 
99 } // namespace
100 
101 //===----------------------------------------------------------------------===//
102 // Pattern population
103 //===----------------------------------------------------------------------===//
104 
106  int64_t byteCountThreshold,
107  RewritePatternSet &patterns) {
108  patterns.add<TensorExtractPattern>(typeConverter, patterns.getContext(),
109  byteCountThreshold);
110 }
This class implements a pattern rewriter for use with ConversionPatterns.
This class defines the main interface for locations in MLIR and acts as a non-nullable wrapper around...
Definition: Location.h:63
MLIRContext is the top-level object for a collection of MLIR operations.
Definition: MLIRContext.h:60
Operation * create(const OperationState &state)
Creates an operation given the fields represented as an OperationState.
Definition: Builders.cpp:464
OpConversionPattern is a wrapper around ConversionPattern that allows for matching and rewriting agai...
This class represents the benefit of a pattern match in a unitless scheme that ranges from 0 (very li...
Definition: PatternMatch.h:34
MLIRContext * getContext() const
Definition: PatternMatch.h:823
RewritePatternSet & add(ConstructorArg &&arg, ConstructorArgs &&...args)
Add an instance of each of the pattern types 'Ts' to the pattern list with the given arguments.
Definition: PatternMatch.h:847
std::enable_if_t<!std::is_convertible< CallbackT, Twine >::value, LogicalResult > notifyMatchFailure(Location loc, CallbackT &&reasonCallback)
Used to notify the listener that the IR failed to be rewritten because of a match failure,...
Definition: PatternMatch.h:718
OpTy replaceOpWithNewOp(Operation *op, Args &&...args)
Replace the results of the given (original) op with a new op that is created without verification (re...
Definition: PatternMatch.h:536
Type conversion from builtin types to SPIR-V types for shader interface.
Type conversion class.
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
static PointerType get(Type pointeeType, StorageClass storageClass)
Definition: SPIRVTypes.cpp:481
Value linearizeIndex(ValueRange indices, ArrayRef< int64_t > strides, int64_t offset, Type integerType, Location loc, OpBuilder &builder)
Generates IR to perform index linearization with the given indices and their corresponding strides,...
Include the generated interface declarations.
LogicalResult failure(bool isFailure=true)
Utility function to generate a LogicalResult.
Definition: LogicalResult.h:62
LogicalResult success(bool isSuccess=true)
Utility function to generate a LogicalResult.
Definition: LogicalResult.h:56
void populateTensorToSPIRVPatterns(SPIRVTypeConverter &typeConverter, int64_t byteCountThreshold, RewritePatternSet &patterns)
Appends to a pattern list additional patterns for translating tensor ops to SPIR-V ops.
This class represents an efficient way to signal success or failure.
Definition: LogicalResult.h:26