MLIR  22.0.0git
XeGPUUtils.cpp
Go to the documentation of this file.
1 //===---- XeGPUUtils.cpp - MLIR Utilities for XeGPUOps ------------------===//
2 //
3 // Part of the MLIR 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 utility methods for working with the XeGPU dialect.
10 //
11 //===----------------------------------------------------------------------===//
12 
20 #include "mlir/IR/Builders.h"
21 #include "mlir/IR/Operation.h"
22 #include "mlir/IR/ValueRange.h"
25 #include "llvm/Support/FormatVariadic.h"
26 #include <cstdint>
27 #include <numeric>
28 
29 using namespace mlir;
30 
31 /// convert ArrayRef<ValueRange> into SmallVector<Value>
33  SmallVector<Value> result;
34  for (const auto &vals : values)
35  llvm::append_range(result, vals);
36  return result;
37 }
38 
39 FailureOr<VectorType>
40 mlir::xegpu::getDistributedVectorType(xegpu::TensorDescType tdescTy) {
41  auto layout = llvm::dyn_cast_if_present<LayoutAttr>(tdescTy.getLayout());
42  // It only works for subgroup level layout, which only has lane_layout
43  // and lane_data, and is to distribute a SIMD code into SIMT code.
44  if (!layout || !layout.isForSubgroup())
45  return failure();
46 
47  SmallVector<int64_t> laneData(layout.getLaneData().asArrayRef());
48  SmallVector<int64_t> laneLayout(layout.getLaneLayout().asArrayRef());
49  auto tdescShape = tdescTy.getShape();
50  auto elementType = tdescTy.getElementType();
51 
52  // compute sgSize by multiply elements of laneLayout
53  // e.g. for 2D layout, sgSize = laneLayout[0] * laneLayout[1]
54  // e.g. for 1D layout, sgSize = laneLayout[0]
55  auto sgSize = std::accumulate(laneLayout.begin(), laneLayout.end(), 1,
56  std::multiplies<int64_t>());
57 
58  // Case 1: regular loads/stores
59  auto scatterAttr = tdescTy.getEncodingOfType<ScatterTensorDescAttr>();
60  if (scatterAttr) {
61  auto chunkSize = scatterAttr.getChunkSize().getInt();
62  // Verify if the first dimension of the tensor descriptor shape is
63  // distributable.
64  assert(tdescShape[0] == laneLayout[0] &&
65  "tensor descriptor shape is not distributable");
66  return VectorType::get({chunkSize}, elementType);
67  }
68 
69  // Case 2: block loads/stores
70  // Check if the tensor descriptor shape is distributable.
71  int64_t tensorSize = 1;
72  for (auto [tdescDim, laneDim, laneDataDim] :
73  llvm::zip_equal(tdescShape, laneLayout, laneData)) {
74  assert((tdescDim % (laneDim * laneDataDim) == 0) &&
75  "tensor descriptor shape is not distributable");
76  tensorSize *= tdescDim;
77  }
78  // tensorSize must be adjusted for array_length.
79  tensorSize *= tdescTy.getArrayLength();
80 
81  return VectorType::get({tensorSize / sgSize}, elementType);
82 }
83 
84 FailureOr<VectorType>
85 mlir::xegpu::getDistributedVectorType(VectorType originalType,
86  xegpu::LayoutAttr layout) {
87  int64_t rank = originalType.getRank();
88  // Distributed vector type is only supported for 1D, 2D and 3D vectors.
89  if (rank < 1 || rank > 3)
90  return failure();
91  ArrayRef<int64_t> shape = originalType.getShape();
92  // arrayLength is 1 for 1D and 2D vectors, and equal to the first dimension
93  // of the 3D vector.
94  int arrayLength = 1;
95  if (rank == 3) {
96  arrayLength = shape[0];
97  shape = shape.drop_front();
98  }
99  auto helperTdescTy = xegpu::TensorDescType::get(
100  shape, originalType.getElementType(), arrayLength,
101  /*boundary_check=*/true,
102  /*memory_space=*/xegpu::MemorySpace::Global, layout);
103  return xegpu::getDistributedVectorType(helperTdescTy);
104 }
105 
106 std::string xegpu::getLayoutName(const OpOperand &operand) {
107  const StringRef prefix("layout_operand_");
108  unsigned idx = const_cast<OpOperand &>(operand).getOperandNumber();
109  return llvm::formatv("{0}{1}", prefix, idx).str();
110 }
111 
112 std::string xegpu::getLayoutName(const OpResult result) {
113  const StringRef prefix = "layout_result_";
114  return llvm::formatv("{0}{1}", prefix, result.getResultNumber()).str();
115 }
116 
117 xegpu::DistributeLayoutAttr xegpu::getDistributeLayoutAttr(const Value value) {
118  if (!value)
119  return nullptr;
120 
121  if (auto tdescTy =
122  dyn_cast_if_present<xegpu::TensorDescType>(value.getType()))
123  return tdescTy.getLayoutAttr();
124 
125  if (auto result = dyn_cast<OpResult>(value)) {
126  Operation *defOp = result.getDefiningOp();
127  assert(defOp && "result must have a defining op");
128 
129  // For ConvertLayoutOp, the layout is stored in the targetLayoutAttr
130  if (auto convertOp = dyn_cast<xegpu::ConvertLayoutOp>(defOp))
131  return convertOp.getTargetLayoutAttr();
132 
133  // for LoadNdOp, the layout is stored in the tensor descriptor
134  if (auto loadNd = dyn_cast<xegpu::LoadNdOp>(defOp))
135  return getDistributeLayoutAttr(loadNd.getTensorDesc());
136 
137  std::string layoutName = getLayoutName(result);
138  if (defOp->hasAttr(layoutName))
139  return defOp->getAttrOfType<xegpu::DistributeLayoutAttr>(layoutName);
140  }
141 
142  if (auto arg = dyn_cast<BlockArgument>(value)) {
143  auto parentOp = arg.getOwner()->getParentOp();
144  if (auto loop = dyn_cast<LoopLikeOpInterface>(parentOp)) {
145  OpOperand *tiedInit = loop.getTiedLoopInit(arg);
146  if (tiedInit)
147  return getDistributeLayoutAttr(tiedInit->get());
148  }
149  }
150 
151  return nullptr;
152 }
153 
154 xegpu::DistributeLayoutAttr
156  Operation *op = opr.getOwner();
157  std::string layoutName = xegpu::getLayoutName(opr);
158  if (op->hasAttr(layoutName))
159  return op->getAttrOfType<xegpu::DistributeLayoutAttr>(layoutName);
160  return getDistributeLayoutAttr(opr.get());
161 }
162 
163 template <typename T, typename>
164 void xegpu::setDistributeLayoutAttr(const T &operandOrResult,
165  const DistributeLayoutAttr layout) {
166  Operation *owner = operandOrResult.getOwner();
167  std::string name = xegpu::getLayoutName(operandOrResult);
168  if (layout && !owner->hasAttrOfType<DistributeLayoutAttr>(name))
169  owner->setAttr(name, layout);
170 }
171 
172 // Explicit instantiation for OpResult
173 template void xegpu::setDistributeLayoutAttr<mlir::OpResult>(
174  const mlir::OpResult &result,
175  const mlir::xegpu::DistributeLayoutAttr layout);
176 
177 // Explicit instantiation for OpOperand
178 template void xegpu::setDistributeLayoutAttr<mlir::OpOperand>(
179  const mlir::OpOperand &operand,
180  const mlir::xegpu::DistributeLayoutAttr layout);
181 
183  Operation *op, function_ref<DistributeLayoutAttr(Value)> getLayoutImpl) {
184  op->walk([&](Operation *nestOp) {
185  for (OpOperand &opr : nestOp->getOpOperands()) {
186  auto layout = getLayoutImpl(opr.get());
187  setDistributeLayoutAttr(opr, layout);
188  }
189  for (OpResult result : nestOp->getOpResults()) {
190  auto layout = getLayoutImpl(result);
191  setDistributeLayoutAttr(result, layout);
192  }
193  });
194 }
195 
196 template <typename T, typename>
197 void xegpu::removeLayoutAttr(const T &operandOrResult) {
198  Operation *owner = operandOrResult.getOwner();
199  std::string name = xegpu::getLayoutName(operandOrResult);
200  if (owner->hasAttrOfType<DistributeLayoutAttr>(name))
201  owner->removeAttr(name);
202 }
203 
204 // Explicit instantiation for OpResult
205 template void
206 xegpu::removeLayoutAttr<mlir::OpResult>(const mlir::OpResult &result);
207 
208 // Explicit instantiation for OpOperand
209 template void
210 xegpu::removeLayoutAttr<mlir::OpOperand>(const mlir::OpOperand &operand);
211 
213  op->walk([&](Operation *nestOp) {
214  for (OpOperand &opr : nestOp->getOpOperands())
215  removeLayoutAttr(opr);
216  for (OpResult result : nestOp->getOpResults())
217  removeLayoutAttr(result);
218  });
219 }
220 
223  Value value, ArrayRef<int64_t> shape) {
224  auto vecTy = dyn_cast<VectorType>(value.getType());
225  if (!vecTy)
226  return {value};
227 
228  ArrayRef<int64_t> srcShape = vecTy.getShape();
229  if (!computeShapeRatio(srcShape, shape))
230  return {value};
231 
232  SmallVector<Value> result;
233  for (SmallVector<int64_t> offsets : StaticTileOffsetRange(srcShape, shape)) {
234  SmallVector<int64_t> staticStrides(offsets.size(), 1);
235  result.push_back(vector::ExtractStridedSliceOp::create(
236  builder, loc, value, offsets, shape, staticStrides));
237  }
238 
239  return result;
240 }
241 
243  ValueRange values,
244  ArrayRef<int64_t> shape) {
245  VectorType inputTy = dyn_cast<VectorType>(values[0].getType());
246  assert(llvm::all_of(values.getTypes(),
247  [&](Type type) { return type == inputTy; }) &&
248  "values must be of the same VectorType");
249 
250  Type elemTy = inputTy.getElementType();
251  ArrayRef<int64_t> tileShape = inputTy.getShape();
252 
253  VectorType resultTy = VectorType::get(shape, elemTy);
254  auto zeroAttr = builder.getZeroAttr(elemTy);
255  Value result = arith::ConstantOp::create(
256  builder, loc, resultTy, DenseElementsAttr::get(resultTy, zeroAttr));
257 
258  for (auto [src, offsets] :
259  llvm::zip_equal(values, StaticTileOffsetRange(shape, tileShape))) {
260  SmallVector<int64_t> staticStrides(offsets.size(), 1);
261  result = vector::InsertStridedSliceOp::create(builder, loc, src, result,
262  offsets, staticStrides);
263  }
264  return result;
265 }
266 
268  Operation *op, TypeConverter converter) {
269  MLIRContext *context = op->getContext();
270 
271  auto materializeCast = [](OpBuilder &builder, Type type, ValueRange inputs,
272  Location loc) -> Value {
273  return UnrealizedConversionCastOp::create(builder, loc, type, inputs)
274  .getResult(0);
275  };
276 
277  { // convert VectorType to RankedTensorType for SCF Structural ops
278  TypeConverter converter;
279  converter.addConversion([](Type type) -> Type { return type; });
280  converter.addConversion([](VectorType type) -> Type {
281  return RankedTensorType::get(type.getShape(), type.getElementType());
282  });
283  converter.addSourceMaterialization(materializeCast);
284  converter.addTargetMaterialization(materializeCast);
285 
286  mlir::ConversionTarget target(*context);
287  target.addLegalOp<UnrealizedConversionCastOp>();
288 
291  target);
292  (void)mlir::applyPartialConversion(op, target, std::move(patterns));
293  }
294 
295  { // propagate the layout attribute to RankedTensorType by checking
296  // BuiltInUnrealizedCastOps
297  // for VectorType to RankedTensorType cast.
298  op->walk([](UnrealizedConversionCastOp castOp) {
299  if (castOp.getNumOperands() != 1 || castOp.getNumResults() != 1)
300  return WalkResult::skip();
301 
302  Value input = castOp.getInputs()[0];
303  Value result = castOp.getResults()[0];
304  auto inputTy = dyn_cast<VectorType>(input.getType());
305  auto resultTy = dyn_cast<RankedTensorType>(result.getType());
306 
307  // Only look at ops casting from VectorType to RankedTensorType
308  if (!inputTy || !resultTy)
309  return WalkResult::skip();
310 
311  xegpu::DistributeLayoutAttr layout =
313  if (!layout)
314  return WalkResult::skip();
315 
316  RankedTensorType newTy = resultTy.cloneWithEncoding(layout);
317  result.setType(newTy);
318 
319  // update the arguments if user is a LoopLike op.
320  for (OpOperand &use : result.getUses()) {
321  if (auto loop = dyn_cast<LoopLikeOpInterface>(use.getOwner())) {
322  BlockArgument arg = loop.getTiedLoopRegionIterArg(&use);
323  arg.setType(newTy);
324  }
325  // whileOp has two regions, the BlockArgument of the after region
326  // is not exposed by LoopLikeOpInterface
327  if (auto whileOp = dyn_cast<scf::WhileOp>(use.getOwner())) {
328  unsigned idx = use.getOperandNumber();
329  BlockArgument arg = whileOp.getAfterArguments()[idx];
330  arg.setType(newTy);
331  }
332  }
333  return WalkResult::advance();
334  });
335 
336  // using yieldOp as anchor to update the result type of its ParentOp
337  op->walk([](scf::YieldOp yieldOp) {
338  Operation *parentOp = yieldOp->getParentOp();
339  for (OpResult r : parentOp->getOpResults()) {
340  unsigned idx = r.getResultNumber();
341  Type resultTy = r.getType();
342  Type yieldTy = yieldOp.getResults()[idx].getType();
343  if (isa<RankedTensorType>(resultTy) && yieldTy != resultTy)
344  r.setType(yieldTy);
345  }
346  });
347  }
348 
349  { // perform the conversion from RankedTensorType to VectorType based on the
350  // DistributeLayoutAttr
351 
352  // Handle the UnrealizedConversionCastOp introduced by the first step.
353  // For vector->RankedTensorType, it will simply forward the inputs.
354  // For RankedTensorType->vector, it will update the inputs with the
355  // one from the adaptor.
356  class UnrealizedConversionCastOpPattern
357  : public OpConversionPattern<mlir::UnrealizedConversionCastOp> {
358  using OpConversionPattern<
359  mlir::UnrealizedConversionCastOp>::OpConversionPattern;
360 
361  mlir::LogicalResult
362  matchAndRewrite(mlir::UnrealizedConversionCastOp op,
363  OneToNOpAdaptor adaptor,
364  ConversionPatternRewriter &rewriter) const override {
365  auto inputs = op.getOperands();
366  auto outputs = op.getOutputs();
367 
368  if (inputs.size() != 1 || outputs.size() != 1)
369  return failure();
370 
371  auto inputTy = inputs[0].getType();
372  auto outputTy = outputs[0].getType();
373 
374  if (isa<VectorType>(inputTy) && isa<RankedTensorType>(outputTy)) {
375  rewriter.replaceOpWithMultiple(op, adaptor.getInputs());
376  return success();
377  }
378 
379  if (isa<RankedTensorType>(inputTy) && isa<VectorType>(outputTy)) {
380  SmallVector<Value> values = xegpu::flattenValues(adaptor.getInputs());
381  auto newOp = UnrealizedConversionCastOp::create(rewriter, op.getLoc(),
382  outputTy, values);
383  rewriter.replaceOp(op, newOp);
384  return success();
385  }
386  return failure();
387  }
388  };
389 
390  converter.addSourceMaterialization(materializeCast);
391  converter.addTargetMaterialization([&](OpBuilder &builder, TypeRange type,
392  ValueRange inputs, Location loc) {
393  return UnrealizedConversionCastOp::create(builder, loc, type, inputs)
394  .getResults();
395  });
396 
397  mlir::ConversionTarget target(*context);
398  target.addDynamicallyLegalOp<UnrealizedConversionCastOp>(
399  [](UnrealizedConversionCastOp op) {
400  auto isTensorTy = [](Type type) {
401  return isa<RankedTensorType>(type);
402  };
403  return llvm::none_of(op->getOperandTypes(), isTensorTy) &&
404  llvm::none_of(op->getResultTypes(), isTensorTy);
405  });
407  patterns.insert<UnrealizedConversionCastOpPattern>(context);
409  target);
410  (void)mlir::applyPartialConversion(op, target, std::move(patterns));
411  }
412 }
413 
414 std::optional<std::string> xegpu::getChipStr(Operation *op) {
415  auto gpuModuleOp = op->getParentOfType<gpu::GPUModuleOp>();
416 
417  if (!gpuModuleOp)
418  return std::nullopt;
419 
420  auto targetAttrs = gpuModuleOp.getTargets();
421  if (targetAttrs) {
422  for (auto &attr : *targetAttrs) {
423  auto xevmAttr = llvm::dyn_cast<xevm::XeVMTargetAttr>(attr);
424  if (xevmAttr)
425  return xevmAttr.getChip().str();
426  }
427  }
428 
429  return std::nullopt;
430 }
431 
432 /// Generates element-wise addition ops of two arrays with automatic alignment.
433 /// When the input arrays have different sizes, the shorter array is
434 /// right-aligned with the longer array, and the unmatched leading elements from
435 /// the longer array are preserved unchanged. This is commonly used for offset
436 /// computation where higher-dimensional offsets need to be added to
437 /// lower-dimensional adjustments.
438 ///
439 /// Example:
440 /// lhs = [l1, l2, l3], rhs = [r1, r2]
441 /// Result: [11, l2+r1, l3+r2]
446  // ensure a is longer than b
447  ArrayRef<OpFoldResult> a = lhs.size() >= rhs.size() ? lhs : rhs;
448  ArrayRef<OpFoldResult> b = lhs.size() >= rhs.size() ? rhs : lhs;
449  SmallVector<OpFoldResult> results(a.take_front(a.size() - b.size()));
450  a = a.slice(a.size() - b.size());
451  for (auto [l, r] : llvm::zip(a, b)) {
452  auto lval = getValueOrCreateConstantIndexOp(builder, loc, l);
453  auto rval = getValueOrCreateConstantIndexOp(builder, loc, r);
454  results.push_back(builder.createOrFold<index::AddOp>(loc, lval, rval));
455  }
456  return results;
457  return {};
458 }
This class represents an argument of a Block.
Definition: Value.h:309
TypedAttr getZeroAttr(Type type)
Definition: Builders.cpp:323
This class implements a pattern rewriter for use with ConversionPatterns.
void replaceOp(Operation *op, ValueRange newValues) override
Replace the given operation with the new values.
void replaceOpWithMultiple(Operation *op, SmallVector< SmallVector< Value >> &&newValues)
Replace the given operation with the new value ranges.
This class describes a specific conversion target.
void addLegalOp(OperationName op)
Register the given operations as legal.
void addDynamicallyLegalOp(OperationName op, const DynamicLegalityCallbackFn &callback)
Register the given operation as dynamically legal and set the dynamic legalization callback to the on...
static DenseElementsAttr get(ShapedType type, ArrayRef< Attribute > values)
Constructs a dense elements attribute from an array of element values.
IRValueT get() const
Return the current value being used by this operand.
Definition: UseDefLists.h:160
This class defines the main interface for locations in MLIR and acts as a non-nullable wrapper around...
Definition: Location.h:76
MLIRContext is the top-level object for a collection of MLIR operations.
Definition: MLIRContext.h:63
This class helps build Operations.
Definition: Builders.h:207
void createOrFold(SmallVectorImpl< Value > &results, Location location, Args &&...args)
Create an operation of specific op type at the current insertion point, and immediately try to fold i...
Definition: Builders.h:519
OpConversionPattern is a wrapper around ConversionPattern that allows for matching and rewriting agai...
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
unsigned getResultNumber() const
Returns the number of this result.
Definition: Value.h:459
Operation is the basic unit of execution within MLIR.
Definition: Operation.h:88
AttrClass getAttrOfType(StringAttr name)
Definition: Operation.h:550
bool hasAttrOfType(NameT &&name)
Definition: Operation.h:575
bool hasAttr(StringAttr name)
Return true if the operation has an attribute with the provided name, false otherwise.
Definition: Operation.h:560
std::enable_if_t< llvm::function_traits< std::decay_t< FnT > >::num_args==1, RetT > walk(FnT &&callback)
Walk the operation by calling the callback for each nested operation (including this one),...
Definition: Operation.h:797
MLIRContext * getContext()
Return the context this operation is associated with.
Definition: Operation.h:216
Operation * getParentOp()
Returns the closest surrounding operation that contains this operation or nullptr if this is a top-le...
Definition: Operation.h:234
OpTy getParentOfType()
Return the closest surrounding parent operation that is of type 'OpTy'.
Definition: Operation.h:238
void setAttr(StringAttr name, Attribute value)
If the an attribute exists with the specified name, change it to the new value.
Definition: Operation.h:582
operand_type_range getOperandTypes()
Definition: Operation.h:397
MutableArrayRef< OpOperand > getOpOperands()
Definition: Operation.h:383
result_type_range getResultTypes()
Definition: Operation.h:428
result_range getOpResults()
Definition: Operation.h:420
Attribute removeAttr(StringAttr name)
Remove the attribute with the specified name if it exists.
Definition: Operation.h:600
A range-style iterator that allows for iterating over the offsets of all potential tiles of size tile...
Type conversion class.
void addConversion(FnT &&callback)
Register a conversion function.
void addSourceMaterialization(FnT &&callback)
All of the following materializations require function objects that are convertible to the following ...
void addTargetMaterialization(FnT &&callback)
This method registers a materialization that will be called when converting a value to a target type ...
This class provides an abstraction over the various different ranges of value types.
Definition: TypeRange.h:37
Instances of the Type class are uniqued, have an immutable identifier and an optional mutable compone...
Definition: Types.h:74
This class provides an abstraction over the different types of ranges over Values.
Definition: ValueRange.h:387
type_range getTypes() const
This class represents an instance of an SSA value in the MLIR system, representing a computable value...
Definition: Value.h:96
void setType(Type newType)
Mutate the type of this Value to be of the specified type.
Definition: Value.h:116
Type getType() const
Return the type of this value.
Definition: Value.h:105
use_range getUses() const
Returns a range of all uses, which is useful for iterating over all uses.
Definition: Value.h:188
static WalkResult skip()
Definition: WalkResult.h:48
static WalkResult advance()
Definition: WalkResult.h:47
Operation * getOwner() const
Return the owner of this operand.
Definition: UseDefLists.h:38
void populateSCFStructuralTypeConversionsAndLegality(const TypeConverter &typeConverter, RewritePatternSet &patterns, ConversionTarget &target)
Populates patterns for SCF structural type conversions and sets up the provided ConversionTarget with...
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:242
void setDistributeLayoutAttr(const T &operandOrResult, const DistributeLayoutAttr layout)
Sets the DistributeLayoutAttr for a given OpOperand or OpResult by attaching it to the owner's dictio...
Definition: XeGPUUtils.cpp:164
void setDistributeLayoutAttrs(Operation *op, function_ref< DistributeLayoutAttr(Value)> getLayoutImpl)
Set the DistributeLayoutAttr for each OpOperand and OpResult of the given operation.
Definition: XeGPUUtils.cpp:182
std::string getLayoutName(const OpOperand &operand)
Return the attribute name for the OpOperand to attach DistributeLayoutAttr.
Definition: XeGPUUtils.cpp:106
void removeLayoutAttr(const T &operandOrResult)
Removes the LayoutAttr for a given OpOperand or OpResult if it exists.
Definition: XeGPUUtils.cpp:197
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:267
DistributeLayoutAttr getDistributeLayoutAttr(const Value value)
Retrieves the DistributeLayoutAttr associated with a given Value.
Definition: XeGPUUtils.cpp:117
std::optional< std::string > getChipStr(Operation *op)
Retrieves the chip string from the XeVM target attribute of the parent GPU module operation.
Definition: XeGPUUtils.cpp:414
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:222
void removeLayoutAttrs(Operation *op)
Removes the DistributeLayoutAttr for each OpOperand and OpResult of the given operation if they exist...
Definition: XeGPUUtils.cpp:212
SmallVector< Value > flattenValues(ArrayRef< ValueRange > values)
Flatten a set of ValueRange into a single SmallVector<Value>
Definition: XeGPUUtils.cpp:32
SmallVector< OpFoldResult > addWithRightAligned(OpBuilder &builder, Location loc, ArrayRef< OpFoldResult > lhs, ArrayRef< OpFoldResult > rhs)
Generates element-wise addition ops of two arrays with automatic alignment.
Definition: XeGPUUtils.cpp:443
FailureOr< VectorType > getDistributedVectorType(xegpu::TensorDescType tdescTy)
If tensor descriptor has a layout attribute it is used in SIMT mode.
Definition: XeGPUUtils.cpp:40
Include the generated interface declarations.
Type getType(OpFoldResult ofr)
Returns the int type of the integer in ofr.
Definition: Utils.cpp:304
const FrozenRewritePatternSet & patterns
Value getValueOrCreateConstantIndexOp(OpBuilder &b, Location loc, OpFoldResult ofr)
Converts an OpFoldResult to a Value.
Definition: Utils.cpp:111
auto get(MLIRContext *context, Ts &&...params)
Helper method that injects context only if needed, this helps unify some of the attribute constructio...
std::optional< SmallVector< int64_t > > computeShapeRatio(ArrayRef< int64_t > shape, ArrayRef< int64_t > subShape)
Return the multi-dimensional integral ratio of subShape to the trailing dimensions of shape.
LogicalResult applyPartialConversion(ArrayRef< Operation * > ops, const ConversionTarget &target, const FrozenRewritePatternSet &patterns, ConversionConfig config=ConversionConfig())
Below we define several entry points for operation conversion.