MLIR  21.0.0git
XeGPUUtils.h
Go to the documentation of this file.
1 //===- XeGPUUtils.h - Vector Utilities --------------------------*- C++ -*-===//
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 #ifndef MLIR_DIALECT_XEGPU_UTILS_XEGPUUTILS_H_
10 #define MLIR_DIALECT_XEGPU_UTILS_XEGPUUTILS_H_
11 
12 #include "mlir/IR/BuiltinTypes.h"
13 namespace mlir {
14 
15 class VectorType;
16 class OpOperand;
17 class OpResult;
18 class OpBuilder;
19 class ValueRange;
20 class TypeConverter;
21 
22 namespace xegpu {
23 class LayoutAttr;
24 class TensorDescType;
25 } // namespace xegpu
26 
27 namespace xegpu {
28 
29 /// If tensor descriptor has a layout attribute it is used in SIMT mode.
30 /// In this mode, the distributed vector shape is determined as follows:
31 /// Definitions:
32 /// lane_data_size = lane_data[0] × lane_data[1]
33 /// subgroup_size = lane_layout[0] × lane_layout[1]
34 /// distribution_unit_size = subgroup_size × lane_data_size
35 ///
36 /// Case 1: Regular loads/stores.
37 /// The following conditions must be met:
38 /// * tensor_desc[0] == lane_layout[0]
39 /// Distributed vector is a 1D vector with shape:
40 /// [chunk_size]
41 ///
42 /// Case 2: Block loads/stores
43 /// Additional definitions:
44 /// tensor_size = tensor_desc[0] * .. * tensor_desc[r-1] * array_length
45 /// n_distribution_units = tensor_size / distribution_unit_size
46 /// fragment_size = n_distribution_units * lane_data_size
47 /// Given above definitions, the following conditions must be met:
48 /// * tensor_desc[0] % (lane_layout[0] × lane_data[0]) == 0
49 /// * tensor_desc[1] % (lane_layout[1] × lane_data[1]) == 0
50 /// Distributed vector is a 1D vector with shape:
51 /// [fragment_size]
52 FailureOr<VectorType> getDistributedVectorType(xegpu::TensorDescType tdescTy);
53 
54 /// Helper to get the distributed vector type for a given vector type according
55 /// to a given LayoutAttr.
56 FailureOr<VectorType> getDistributedVectorType(VectorType originalType,
57  LayoutAttr layout);
58 
59 /// Return the attribute name for the OpOperand to attach LayoutAttr
60 std::string getLayoutName(const OpOperand &operand);
61 
62 /// Return the attribute name for the OpResult to attach LayoutAttr
63 std::string getLayoutName(const OpResult result);
64 
65 /// Retrieves the LayoutAttr associated with a given Value. For TensorDescType
66 /// values, the LayoutAttr is extracted from the TensorDescType itself. For
67 /// other values, it is obtained from the attributes of the defining operation.
68 /// Returns nullptr if no LayoutAttr is found.
69 LayoutAttr getLayoutAttr(const Value value);
70 
71 /// Retrieves the LayoutAttr associated with a given OpOperand. It will
72 /// first check the operand_layout_{id} of the owner operation. If not found,
73 /// it will check the operand itself and its defining op.
74 LayoutAttr getLayoutAttr(const OpOperand &opr);
75 
76 /// Sets the LayoutAttr for a given OpOperand or OpResult by attaching
77 /// it to the owner's dictionary attributes
78 template <typename T,
79  typename = std::enable_if_t<std::is_same_v<T, OpOperand> ||
80  std::is_same_v<T, OpResult>>>
81 void setLayoutAttr(const T &operandOrResult, const LayoutAttr layout);
82 
83 /// Set the LayoutAttr for each OpOperand and OpResult of the given operation.
84 /// If the operation contains regions, it is also applied recursively to the
85 /// contained operations
86 void setLayoutAttrs(Operation *op,
87  function_ref<LayoutAttr(Value)> getLayoutImpl);
88 
89 /// Extract a set of small vectors from a value with a given shape using
90 /// vector.extract_stride_slice
92  Location loc, Value value,
93  ArrayRef<int64_t> shape);
94 
95 /// Create a vector of shape from a set of values using
96 /// vector.insert_stride_slice.
98  ValueRange values,
99  ArrayRef<int64_t> shape);
100 
101 /// Do type conversion for SCF structural ops, e.g., scf.for using SCF structure
102 /// type convertion patterns. Since VectorType cannot carry the layout
103 /// attribute, which is needed to guide the type conversion for XeGPU, they are
104 /// first converted into RankedTensorType, where the layout attribute can be
105 /// attached. And then upstream SCF structural type conversion patterns are
106 /// applied with the provided converter.
107 /// TODO: This is a temporary solution. We should refactor it when context-aware
108 /// type conversion is available.
110  TypeConverter converter);
111 
112 } // namespace xegpu
113 
114 } // namespace mlir
115 
116 #endif // MLIR_DIALECT_XEGPU_UTILS_XEGPUUTILS_H_
This class defines the main interface for locations in MLIR and acts as a non-nullable wrapper around...
Definition: Location.h:76
This class helps build Operations.
Definition: Builders.h:205
This class represents an operand of an operation.
Definition: Value.h:257
This is a value defined by a result of an operation.
Definition: Value.h:447
Operation is the basic unit of execution within MLIR.
Definition: Operation.h:88
Type conversion class.
This class provides an abstraction over the different types of ranges over Values.
Definition: ValueRange.h:387
This class represents an instance of an SSA value in the MLIR system, representing a computable value...
Definition: Value.h:96
Value createVectorWithShapeFromValues(OpBuilder &builder, Location loc, ValueRange values, ArrayRef< int64_t > shape)
Create a vector of shape from a set of values using vector.insert_stride_slice.
Definition: XeGPUUtils.cpp:208
LayoutAttr getLayoutAttr(const Value value)
Retrieves the LayoutAttr associated with a given Value.
Definition: XeGPUUtils.cpp:115
void setLayoutAttr(const T &operandOrResult, const LayoutAttr layout)
Sets the LayoutAttr for a given OpOperand or OpResult by attaching it to the owner's dictionary attri...
Definition: XeGPUUtils.cpp:156
std::string getLayoutName(const OpOperand &operand)
Return the attribute name for the OpOperand to attach LayoutAttr.
Definition: XeGPUUtils.cpp:104
void doSCFStructuralTypeConversionWithTensorType(Operation *op, TypeConverter converter)
Do type conversion for SCF structural ops, e.g., scf.for using SCF structure type convertion patterns...
Definition: XeGPUUtils.cpp:233
void setLayoutAttrs(Operation *op, function_ref< LayoutAttr(Value)> getLayoutImpl)
Set the LayoutAttr for each OpOperand and OpResult of the given operation.
Definition: XeGPUUtils.cpp:173
SmallVector< Value > extractVectorsWithShapeFromValue(OpBuilder &builder, Location loc, Value value, ArrayRef< int64_t > shape)
Extract a set of small vectors from a value with a given shape using vector.extract_stride_slice.
Definition: XeGPUUtils.cpp:188
FailureOr< VectorType > getDistributedVectorType(xegpu::TensorDescType tdescTy)
If tensor descriptor has a layout attribute it is used in SIMT mode.
Definition: XeGPUUtils.cpp:38
Include the generated interface declarations.