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  // for LoadMatrixOp, the layout is attached to the property of the op
138  if (auto loadOp = dyn_cast<xegpu::LoadMatrixOp>(defOp))
139  return loadOp.getLayoutAttr();
140 
141  // for StoreMatrixOp, the layout is attached to the property of the op
142  if (auto storeOp = dyn_cast<xegpu::StoreMatrixOp>(defOp))
143  return storeOp.getLayoutAttr();
144 
145  std::string layoutName = getLayoutName(result);
146  if (defOp->hasAttr(layoutName))
147  return defOp->getAttrOfType<xegpu::DistributeLayoutAttr>(layoutName);
148  }
149 
150  if (auto arg = dyn_cast<BlockArgument>(value)) {
151  auto parentOp = arg.getOwner()->getParentOp();
152  if (auto loop = dyn_cast<LoopLikeOpInterface>(parentOp)) {
153  OpOperand *tiedInit = loop.getTiedLoopInit(arg);
154  if (tiedInit)
155  return getDistributeLayoutAttr(tiedInit->get());
156  }
157  }
158 
159  return nullptr;
160 }
161 
162 xegpu::DistributeLayoutAttr
164  Operation *op = opr.getOwner();
165 
166  if (auto loadOp = dyn_cast<xegpu::LoadMatrixOp>(op))
167  return loadOp.getLayoutAttr();
168 
169  if (auto storeOp = dyn_cast<xegpu::StoreMatrixOp>(op))
170  return storeOp.getLayoutAttr();
171 
172  std::string layoutName = xegpu::getLayoutName(opr);
173  if (op->hasAttr(layoutName))
174  return op->getAttrOfType<xegpu::DistributeLayoutAttr>(layoutName);
175  return getDistributeLayoutAttr(opr.get());
176 }
177 
178 template <typename T, typename>
179 void xegpu::setDistributeLayoutAttr(const T &operandOrResult,
180  const DistributeLayoutAttr layout) {
181  Operation *owner = operandOrResult.getOwner();
182  std::string name = xegpu::getLayoutName(operandOrResult);
183  if (layout && !owner->hasAttrOfType<DistributeLayoutAttr>(name))
184  owner->setAttr(name, layout);
185 }
186 
187 // Explicit instantiation for OpResult
188 template void xegpu::setDistributeLayoutAttr<mlir::OpResult>(
189  const mlir::OpResult &result,
190  const mlir::xegpu::DistributeLayoutAttr layout);
191 
192 // Explicit instantiation for OpOperand
193 template void xegpu::setDistributeLayoutAttr<mlir::OpOperand>(
194  const mlir::OpOperand &operand,
195  const mlir::xegpu::DistributeLayoutAttr layout);
196 
198  Operation *op, function_ref<DistributeLayoutAttr(Value)> getLayoutImpl) {
199  op->walk([&](Operation *nestOp) {
200  if (isa<xegpu::LoadMatrixOp, xegpu::StoreMatrixOp>(nestOp))
201  return;
202 
203  for (OpOperand &opr : nestOp->getOpOperands()) {
204  auto layout = getLayoutImpl(opr.get());
205  setDistributeLayoutAttr(opr, layout);
206  }
207  for (OpResult result : nestOp->getOpResults()) {
208  auto layout = getLayoutImpl(result);
209  setDistributeLayoutAttr(result, layout);
210  }
211  });
212 }
213 
214 template <typename T, typename>
215 void xegpu::removeLayoutAttr(const T &operandOrResult) {
216  Operation *owner = operandOrResult.getOwner();
217  std::string name = xegpu::getLayoutName(operandOrResult);
218  if (owner->hasAttrOfType<DistributeLayoutAttr>(name))
219  owner->removeAttr(name);
220 }
221 
222 // Explicit instantiation for OpResult
223 template void
224 xegpu::removeLayoutAttr<mlir::OpResult>(const mlir::OpResult &result);
225 
226 // Explicit instantiation for OpOperand
227 template void
228 xegpu::removeLayoutAttr<mlir::OpOperand>(const mlir::OpOperand &operand);
229 
231  op->walk([&](Operation *nestOp) {
232  for (OpOperand &opr : nestOp->getOpOperands())
233  removeLayoutAttr(opr);
234  for (OpResult result : nestOp->getOpResults())
235  removeLayoutAttr(result);
236  });
237 }
238 
241  Value value, ArrayRef<int64_t> shape) {
242  auto vecTy = dyn_cast<VectorType>(value.getType());
243  if (!vecTy)
244  return {value};
245 
246  ArrayRef<int64_t> srcShape = vecTy.getShape();
247  if (!computeShapeRatio(srcShape, shape))
248  return {value};
249 
250  SmallVector<Value> result;
251  for (SmallVector<int64_t> offsets : StaticTileOffsetRange(srcShape, shape)) {
252  SmallVector<int64_t> staticStrides(offsets.size(), 1);
253  result.push_back(vector::ExtractStridedSliceOp::create(
254  builder, loc, value, offsets, shape, staticStrides));
255  }
256 
257  return result;
258 }
259 
261  ValueRange values,
262  ArrayRef<int64_t> shape) {
263  VectorType inputTy = dyn_cast<VectorType>(values[0].getType());
264  assert(llvm::all_of(values.getTypes(),
265  [&](Type type) { return type == inputTy; }) &&
266  "values must be of the same VectorType");
267 
268  Type elemTy = inputTy.getElementType();
269  ArrayRef<int64_t> tileShape = inputTy.getShape();
270 
271  VectorType resultTy = VectorType::get(shape, elemTy);
272  auto zeroAttr = builder.getZeroAttr(elemTy);
273  Value result = arith::ConstantOp::create(
274  builder, loc, resultTy, DenseElementsAttr::get(resultTy, zeroAttr));
275 
276  for (auto [src, offsets] :
277  llvm::zip_equal(values, StaticTileOffsetRange(shape, tileShape))) {
278  SmallVector<int64_t> staticStrides(offsets.size(), 1);
279  result = vector::InsertStridedSliceOp::create(builder, loc, src, result,
280  offsets, staticStrides);
281  }
282  return result;
283 }
284 
286  Operation *op, TypeConverter converter) {
287  MLIRContext *context = op->getContext();
288 
289  auto materializeCast = [](OpBuilder &builder, Type type, ValueRange inputs,
290  Location loc) -> Value {
291  return UnrealizedConversionCastOp::create(builder, loc, type, inputs)
292  .getResult(0);
293  };
294 
295  { // convert VectorType to RankedTensorType for SCF Structural ops
296  TypeConverter converter;
297  converter.addConversion([](Type type) -> Type { return type; });
298  converter.addConversion([](VectorType type) -> Type {
299  return RankedTensorType::get(type.getShape(), type.getElementType());
300  });
301  converter.addSourceMaterialization(materializeCast);
302  converter.addTargetMaterialization(materializeCast);
303 
304  mlir::ConversionTarget target(*context);
305  target.addLegalOp<UnrealizedConversionCastOp>();
306 
309  target);
310  (void)mlir::applyPartialConversion(op, target, std::move(patterns));
311  }
312 
313  { // propagate the layout attribute to RankedTensorType by checking
314  // BuiltInUnrealizedCastOps
315  // for VectorType to RankedTensorType cast.
316  op->walk([](UnrealizedConversionCastOp castOp) {
317  if (castOp.getNumOperands() != 1 || castOp.getNumResults() != 1)
318  return WalkResult::skip();
319 
320  Value input = castOp.getInputs()[0];
321  Value result = castOp.getResults()[0];
322  auto inputTy = dyn_cast<VectorType>(input.getType());
323  auto resultTy = dyn_cast<RankedTensorType>(result.getType());
324 
325  // Only look at ops casting from VectorType to RankedTensorType
326  if (!inputTy || !resultTy)
327  return WalkResult::skip();
328 
329  xegpu::DistributeLayoutAttr layout =
331  if (!layout)
332  return WalkResult::skip();
333 
334  RankedTensorType newTy = resultTy.cloneWithEncoding(layout);
335  result.setType(newTy);
336 
337  // update the arguments if user is a LoopLike op.
338  for (OpOperand &use : result.getUses()) {
339  if (auto loop = dyn_cast<LoopLikeOpInterface>(use.getOwner())) {
340  BlockArgument arg = loop.getTiedLoopRegionIterArg(&use);
341  arg.setType(newTy);
342  }
343  // whileOp has two regions, the BlockArgument of the after region
344  // is not exposed by LoopLikeOpInterface
345  if (auto whileOp = dyn_cast<scf::WhileOp>(use.getOwner())) {
346  unsigned idx = use.getOperandNumber();
347  BlockArgument arg = whileOp.getAfterArguments()[idx];
348  arg.setType(newTy);
349  }
350  }
351  return WalkResult::advance();
352  });
353 
354  // using yieldOp as anchor to update the result type of its ParentOp
355  op->walk([](scf::YieldOp yieldOp) {
356  Operation *parentOp = yieldOp->getParentOp();
357  for (OpResult r : parentOp->getOpResults()) {
358  unsigned idx = r.getResultNumber();
359  Type resultTy = r.getType();
360  Type yieldTy = yieldOp.getResults()[idx].getType();
361  if (isa<RankedTensorType>(resultTy) && yieldTy != resultTy)
362  r.setType(yieldTy);
363  }
364  });
365  }
366 
367  { // perform the conversion from RankedTensorType to VectorType based on the
368  // DistributeLayoutAttr
369 
370  // Handle the UnrealizedConversionCastOp introduced by the first step.
371  // For vector->RankedTensorType, it will simply forward the inputs.
372  // For RankedTensorType->vector, it will update the inputs with the
373  // one from the adaptor.
374  class UnrealizedConversionCastOpPattern
375  : public OpConversionPattern<mlir::UnrealizedConversionCastOp> {
376  using OpConversionPattern<
377  mlir::UnrealizedConversionCastOp>::OpConversionPattern;
378 
379  mlir::LogicalResult
380  matchAndRewrite(mlir::UnrealizedConversionCastOp op,
381  OneToNOpAdaptor adaptor,
382  ConversionPatternRewriter &rewriter) const override {
383  auto inputs = op.getOperands();
384  auto outputs = op.getOutputs();
385 
386  if (inputs.size() != 1 || outputs.size() != 1)
387  return failure();
388 
389  auto inputTy = inputs[0].getType();
390  auto outputTy = outputs[0].getType();
391 
392  if (isa<VectorType>(inputTy) && isa<RankedTensorType>(outputTy)) {
393  rewriter.replaceOpWithMultiple(op, adaptor.getInputs());
394  return success();
395  }
396 
397  if (isa<RankedTensorType>(inputTy) && isa<VectorType>(outputTy)) {
398  SmallVector<Value> values = xegpu::flattenValues(adaptor.getInputs());
399  auto newOp = UnrealizedConversionCastOp::create(rewriter, op.getLoc(),
400  outputTy, values);
401  rewriter.replaceOp(op, newOp);
402  return success();
403  }
404  return failure();
405  }
406  };
407 
408  converter.addSourceMaterialization(materializeCast);
409  converter.addTargetMaterialization([&](OpBuilder &builder, TypeRange type,
410  ValueRange inputs, Location loc) {
411  return UnrealizedConversionCastOp::create(builder, loc, type, inputs)
412  .getResults();
413  });
414 
415  mlir::ConversionTarget target(*context);
416  target.addDynamicallyLegalOp<UnrealizedConversionCastOp>(
417  [](UnrealizedConversionCastOp op) {
418  auto isTensorTy = [](Type type) {
419  return isa<RankedTensorType>(type);
420  };
421  return llvm::none_of(op->getOperandTypes(), isTensorTy) &&
422  llvm::none_of(op->getResultTypes(), isTensorTy);
423  });
425  patterns.insert<UnrealizedConversionCastOpPattern>(context);
427  target);
428  (void)mlir::applyPartialConversion(op, target, std::move(patterns));
429  }
430 }
431 
432 std::optional<std::string> xegpu::getChipStr(Operation *op) {
433  auto gpuModuleOp = op->getParentOfType<gpu::GPUModuleOp>();
434 
435  if (!gpuModuleOp)
436  return std::nullopt;
437 
438  auto targetAttrs = gpuModuleOp.getTargets();
439  if (targetAttrs) {
440  for (auto &attr : *targetAttrs) {
441  auto xevmAttr = llvm::dyn_cast<xevm::XeVMTargetAttr>(attr);
442  if (xevmAttr)
443  return xevmAttr.getChip().str();
444  }
445  }
446 
447  return std::nullopt;
448 }
449 
450 /// Generates element-wise addition ops of two arrays with same length.
452  Location loc,
455  assert(lhs.size() == rhs.size() && "lhs and rhs must have the same size");
457  for (auto [l, r] : llvm::zip_equal(lhs, rhs)) {
458  auto lval = getValueOrCreateConstantIndexOp(builder, loc, l);
459  auto rval = getValueOrCreateConstantIndexOp(builder, loc, r);
460  results.push_back(builder.createOrFold<index::AddOp>(loc, lval, rval));
461  }
462  return results;
463 }
464 
465 /// Generates element-wise addition ops of two arrays with automatic alignment.
466 /// When the input arrays have different sizes, the shorter array is
467 /// right-aligned with the longer array, and the unmatched leading elements from
468 /// the longer array are preserved unchanged. This is commonly used for offset
469 /// computation where higher-dimensional offsets need to be added to
470 /// lower-dimensional adjustments.
471 ///
472 /// Example:
473 /// lhs = [l1, l2, l3], rhs = [r1, r2]
474 /// Result: [11, l2+r1, l3+r2]
479  // ensure a is longer than b
480  ArrayRef<OpFoldResult> a = lhs.size() >= rhs.size() ? lhs : rhs;
481  ArrayRef<OpFoldResult> b = lhs.size() >= rhs.size() ? rhs : lhs;
482  SmallVector<OpFoldResult> results(a.take_front(a.size() - b.size()));
483  a = a.slice(a.size() - b.size());
484  results.append(addElementwise(builder, loc, a, b));
485  return results;
486 }
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, PatternBenefit benefit=1)
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:260
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:179
void setDistributeLayoutAttrs(Operation *op, function_ref< DistributeLayoutAttr(Value)> getLayoutImpl)
Set the DistributeLayoutAttr for each OpOperand and OpResult of the given operation.
Definition: XeGPUUtils.cpp:197
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:215
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:285
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:432
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:240
void removeLayoutAttrs(Operation *op)
Removes the DistributeLayoutAttr for each OpOperand and OpResult of the given operation if they exist...
Definition: XeGPUUtils.cpp:230
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:476
SmallVector< OpFoldResult > addElementwise(OpBuilder &builder, Location loc, ArrayRef< OpFoldResult > lhs, ArrayRef< OpFoldResult > rhs)
Generates element-wise addition ops of two arrays with same length.
Definition: XeGPUUtils.cpp:451
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.