MLIR  22.0.0git
XeGPUDialect.cpp
Go to the documentation of this file.
1 //===- XeGPUDialect.cpp - MLIR XeGPU dialect implementation -----*- 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 
15 #include "mlir/IR/Builders.h"
17 #include "llvm/ADT/TypeSwitch.h"
18 #include "llvm/Support/Debug.h"
19 
20 using std::optional;
21 
22 namespace mlir {
23 namespace xegpu {
24 
25 void XeGPUDialect::initialize() {
26  addTypes<
27 #define GET_TYPEDEF_LIST
28 #include <mlir/Dialect/XeGPU/IR/XeGPUTypes.cpp.inc>
29  >();
30  addOperations<
31 #define GET_OP_LIST
32 #include <mlir/Dialect/XeGPU/IR/XeGPU.cpp.inc>
33  >();
34  addAttributes<
35 #define GET_ATTRDEF_LIST
36 #include <mlir/Dialect/XeGPU/IR/XeGPUAttrs.cpp.inc>
37  >();
38 }
39 
40 /// Generates instructions to compute offsets for a subgroup identified by
41 /// its multidimensional indices (sgId), using the specified subgroup layout
42 /// (sgLayout), subgroup data dimensions (sizePerSg), and the overall data
43 /// dimensions (sizePerWg).
44 static SmallVector<SmallVector<Value>>
47  ArrayRef<int64_t> sizePerSg,
48  ArrayRef<int64_t> sizePerWg) {
49 
51 
52  // nd local offset, localOffset[i] = sgId[i] * sizePerSg[i]
53  SmallVector<Value> localOffsets = llvm::map_to_vector(
54  llvm::zip(sgId, sizePerSg), [&](const auto &t) -> Value {
55  return builder.createOrFold<index::MulOp>(
56  loc, std::get<0>(t),
57  builder.createOrFold<arith::ConstantIndexOp>(loc, std::get<1>(t)));
58  });
59 
60  // distUnit[i] is the minimum value between sizePerWg[i] and
61  // sgLayout[i] * sizePerSg[i]
62  SmallVector<int64_t> distUnit = llvm::map_to_vector(
63  llvm::zip_equal(sizePerWg, computeElementwiseMul(sgLayout, sizePerSg)),
64  [](const auto &t) { return std::min(std::get<0>(t), std::get<1>(t)); });
65 
66  for (SmallVector<int64_t> unitOffs :
67  StaticTileOffsetRange(sizePerWg, distUnit)) {
68  SmallVector<Value> base =
69  llvm::map_to_vector(unitOffs, [&](int64_t d) -> Value {
70  return arith::ConstantIndexOp::create(builder, loc, d);
71  });
72 
73  SmallVector<Value> adds = llvm::map_to_vector(
74  llvm::zip_equal(base, localOffsets), [&](const auto &t) -> Value {
75  return builder.createOrFold<arith::AddIOp>(loc, std::get<0>(t),
76  std::get<1>(t));
77  });
78 
79  SmallVector<Value> mods = llvm::map_to_vector(
80  llvm::zip_equal(adds, sizePerWg), [&](const auto &t) -> Value {
81  return builder.createOrFold<index::RemUOp>(
82  loc, std::get<0>(t),
83  arith::ConstantIndexOp::create(builder, loc, std::get<1>(t)));
84  });
85 
86  offsets.push_back(mods);
87  }
88  return offsets;
89 }
90 
91 // Checks if the given shape can be evenly distributed based on the layout
92 // and data factors provided by the LayoutAttr.
93 bool XeGPUDialect::isEvenlyDistributable(llvm::ArrayRef<int64_t> shape,
94  xegpu::DistributeLayoutAttr attr) {
95  assert(attr && "Layout attribute is missing.");
96 
97  // Checks whether the given shape can be evenly distributed using the
98  // specified layout and data attributes. If successful, it returns the work
99  // size for each compute unit; otherwise, it returns `std::nullopt`. The work
100  // size per compute unit is calculated as follows:
101  // - If `data` is null: newShape[i] = shape[i] / layout[i]
102  // - If `data` is not null: newShape[i] = data[i]
103  // When round-robin distribution (`rr`) is enabled, `shape[i]` can be
104  // smaller than `layout[i] * data[i]`, allowing multiple compute units to
105  // share the data.
106  auto tryDistribute = [&](llvm::ArrayRef<int64_t> shape,
107  SmallVector<int64_t> layout,
109  bool rr = true) -> optional<SmallVector<int64_t>> {
110  llvm::SmallVector<int64_t> newShape(shape);
111  if (layout.size()) {
112  if (layout.size() != shape.size())
113  return std::nullopt;
114  auto ratio = computeShapeRatio(shape, layout);
115  if (!ratio.has_value())
116  return std::nullopt;
117  newShape = ratio.value();
118  }
119 
120  if (data.size()) {
121  if (data.size() != shape.size())
122  return std::nullopt;
123  auto ratio = computeShapeRatio(newShape, data);
124  if (!ratio.has_value() && rr)
125  ratio = computeShapeRatio(data, newShape);
126  if (!ratio.has_value())
127  return std::nullopt;
128 
129  // if data is not null, we always return it for next phase.
130  newShape = data;
131  }
132  return newShape;
133  };
134 
135  // check the sgLayout and sgData
136  auto maybeSgShape =
137  tryDistribute(shape, attr.getSgLayoutAsInt(), attr.getSgDataAsInt());
138  if (!maybeSgShape)
139  return false;
140  auto sgShape = maybeSgShape.value();
141 
142  // check InstData, it neither have layout nor need round-robin
143  auto maybeInstShape =
144  tryDistribute(sgShape, {}, attr.getInstDataAsInt(), false);
145  if (!maybeInstShape)
146  return false;
147  auto instShape = maybeInstShape.value();
148 
149  // check LaneLayout and LaneData
150  auto maybeLaneShape = tryDistribute(instShape, attr.getLaneLayoutAsInt(),
151  attr.getLaneDataAsInt(), false);
152  return maybeLaneShape.has_value();
153 }
154 
155 //===----------------------------------------------------------------------===//
156 // XeGPU_BlockTensorDescAttr
157 //===----------------------------------------------------------------------===//
158 BlockTensorDescAttr BlockTensorDescAttr::get(mlir::MLIRContext *context,
159  xegpu::MemorySpace memory_space,
160  int array_length,
161  bool boundary_check) {
162  auto scopeAttr = MemorySpaceAttr::get(context, memory_space);
163  auto lengthAttr =
164  IntegerAttr::get(IntegerType::get(context, 64), array_length);
165  auto boundaryAttr = BoolAttr::get(context, boundary_check);
166  return Base::get(context, scopeAttr, lengthAttr, boundaryAttr);
167 }
168 
169 bool BlockTensorDescAttr::hasDefaultsOnly() {
170  return getMemorySpace().getValue() == xegpu::MemorySpace::Global &&
171  getArrayLength().getInt() == 1 && getBoundaryCheck().getValue();
172 }
173 
174 //===----------------------------------------------------------------------===//
175 // XeGPU_ScatterTensorDescAttr
176 //===----------------------------------------------------------------------===//
177 ScatterTensorDescAttr
179  xegpu::MemorySpace memory_space, int chunk_size) {
180  auto scopeAttr = MemorySpaceAttr::get(context, memory_space);
181  auto chunkSizeAttr =
182  IntegerAttr::get(IntegerType::get(context, 64), chunk_size);
183  return Base::get(context, scopeAttr, chunkSizeAttr);
184 }
185 
186 LogicalResult ScatterTensorDescAttr::verify(
188  MemorySpaceAttr memory_space, IntegerAttr chunk_size) {
189  int64_t chunkSize = chunk_size.getInt();
190  if (chunkSize <= 0)
191  return emitError() << "invalid chunk size";
192 
193  return success();
194 }
195 
196 //===----------------------------------------------------------------------===//
197 // XeGPU_LayoutAttr
198 //===----------------------------------------------------------------------===//
199 LogicalResult
201  DenseI32ArrayAttr sg_layout, DenseI32ArrayAttr sg_data,
202  DenseI32ArrayAttr inst_data, DenseI32ArrayAttr lane_layout,
203  DenseI32ArrayAttr lane_data, DenseI32ArrayAttr order) {
204 
205  // A valid layout must include at least one of sg_layout and lane_layout.
206  // sg_layout is essential for Workgroup layout, while lane_layout is
207  // required for Subgroup layout.
208  if (!sg_layout && !inst_data && !lane_layout) {
209  return emitError()
210  << "expected at least one of sg_layout, inst_data or lane_layout";
211  }
212 
213  // generate code to check sg_laout, inst_data and lane_layout having the same
214  // rank if they are not null.
215 
216  if (sg_layout && inst_data && sg_layout.size() != inst_data.size()) {
217  return emitError()
218  << "expected sg_layout and inst_data to have the same rank";
219  }
220 
221  if (sg_layout && lane_layout && sg_layout.size() != lane_layout.size()) {
222  return emitError()
223  << "expected sg_layout and lane_layout to have the same rank";
224  }
225 
226  if (inst_data && lane_layout && inst_data.size() != lane_layout.size()) {
227  return emitError()
228  << "expected inst_data and lane_layout to have the same rank";
229  }
230 
231  // sg_data is optional for Workgroup layout, but its presence requires
232  // sg_layout.
233  if (sg_data) {
234  if (!sg_layout)
235  return emitError() << "expected sg_layout being used with sg_data";
236  if (sg_data.size() != sg_layout.size())
237  return emitError()
238  << "expected sg_data and sg_layout to have the same rank";
239  }
240 
241  // lane_data is optional for Subgroup layout, but its presence requires
242  // lane_layout.
243  if (lane_data) {
244  if (!lane_layout)
245  return emitError() << "expected lane_layout being used with lane_data";
246  if (lane_data.size() != lane_layout.size())
247  return emitError()
248  << "expected lane_data and lane_layout to have the same rank";
249  }
250 
251  if (order) {
252  if (!sg_layout && !lane_layout)
253  return emitError()
254  << "expected sg_layout/lane_layout being used with order";
255 
256  if (sg_layout && order.size() != sg_layout.size())
257  return emitError()
258  << "expected order and sg_layout to have the same rank";
259 
260  if (lane_layout && order.size() != lane_layout.size())
261  return emitError()
262  << "expected order and lane_layout to have the same rank";
263  }
264 
265  return success();
266 }
267 
268 FailureOr<SmallVector<Value>>
269 LayoutAttr::delinearizeSubgroupId(OpBuilder &builder, Location loc,
270  Value linearId) {
271  // delinearizeSubgroupId is only available for
272  // workgroup-level layout attribute
273  if (!isForWorkgroup())
274  return failure();
275 
276  // TODO: handle order attribute
277  auto hasDefaultOrder = [&]() {
278  DenseI32ArrayAttr order = getOrder();
279  return !order || isIdentityPermutation(llvm::to_vector_of<int64_t>(
280  llvm::reverse(order.asArrayRef())));
281  };
282  if (!hasDefaultOrder())
283  return mlir::emitError(loc, "order attribute is currently not supported.");
284 
285  auto dims = llvm::map_to_vector(getSgLayoutAsInt(), [&](int64_t d) -> Value {
286  return builder.createOrFold<arith::ConstantIndexOp>(loc, d);
287  });
288 
289  return affine::delinearizeIndex(builder, loc, linearId, dims);
290 }
291 
292 /// Implements DistributeLayoutAttr::getOffsets to generate
293 /// instructions for computing multi-dimensional offsets when distributed by
294 /// LayoutAttr.
295 FailureOr<SmallVector<SmallVector<Value>>>
296 LayoutAttr::getOffsets(OpBuilder &builder, Location loc, Value linearId,
297  ArrayRef<int64_t> shape) {
298  if (!isForWorkgroup())
299  return failure();
300 
301  SmallVector<int64_t> sgLayout = getSgLayoutAsInt();
302  SmallVector<int64_t> sgShape = getSgDataAsInt();
303  if (sgShape.empty()) {
304  if (auto derivedShape = computeShapeRatio(shape, sgLayout))
305  sgShape = derivedShape.value();
306  else
307  return failure();
308  }
309 
310  // delinearize Ids
311  auto maybeIds = delinearizeSubgroupId(builder, loc, linearId);
312  if (failed(maybeIds))
313  return failure();
314  SmallVector<Value> sgIds = *maybeIds;
315 
316  return genOffsetsComputingInsts(builder, loc, sgIds, sgLayout, sgShape,
317  shape);
318 }
319 
320 //===----------------------------------------------------------------------===//
321 // XeGPU_SliceAttr
322 //===----------------------------------------------------------------------===//
323 LogicalResult
324 SliceAttr::verify(llvm::function_ref<InFlightDiagnostic()> emitError,
325  xegpu::DistributeLayoutAttr parent, DenseI64ArrayAttr dims) {
326  if (!parent || !dims)
327  return emitError() << "expected parent layout and dims attribute";
328 
329  int64_t rank = parent.getRank();
330 
331  // check every element in dims is unique and smaller than rank
332  llvm::SmallDenseSet<int64_t> seen;
333  for (int64_t dim : dims.asArrayRef()) {
334  if (dim < 0 || dim >= rank)
335  return emitError() << "invalid dim (" << dim << ") in slice attribute.";
336  if (!seen.insert(dim).second)
337  return emitError() << "repeated dim (" << dim << ") in slice attribute.";
338  }
339  return success();
340 }
341 
342 SliceAttr SliceAttr::flatten() const {
343  xegpu::DistributeLayoutAttr parent = getParent();
344  SmallVector<DenseI64ArrayAttr> slicedDims({getDims()});
345 
346  while (auto sliceAttr = dyn_cast<xegpu::SliceAttr>(parent)) {
347  parent = sliceAttr.getParent();
348  slicedDims.push_back(sliceAttr.getDims());
349  }
350 
351  auto layoutAttr = dyn_cast<xegpu::LayoutAttr>(parent);
352  SmallVector<int64_t> indices =
353  llvm::to_vector(llvm::seq<int64_t>(0, layoutAttr.getRank()));
354 
355  // get remaining dims (flattend) by applying slice ops with all slicedDims
356  SmallVector<int64_t> remainingDims(indices);
357  for (auto dim : llvm::reverse(slicedDims))
358  remainingDims = XeGPUDialect::slice(llvm::ArrayRef<int64_t>(remainingDims),
359  dim.asArrayRef());
360 
361  // get flattend sliced dims by applying slice ops with the remaining dims
362  SmallVector<int64_t> flattendDims = XeGPUDialect::slice(
363  llvm::ArrayRef<int64_t>(indices), llvm::ArrayRef<int64_t>(remainingDims));
364 
365  return xegpu::SliceAttr::get(
366  getContext(), layoutAttr,
367  DenseI64ArrayAttr::get(getContext(), flattendDims));
368 }
369 
370 FailureOr<SmallVector<Value>>
371 SliceAttr::delinearizeSubgroupId(OpBuilder &builder, Location loc,
372  Value linearId) {
373  SliceAttr attr = flatten();
374  auto parent = dyn_cast<LayoutAttr>(attr.getParent());
375  return parent.delinearizeSubgroupId(builder, loc, linearId);
376 }
377 
378 /// Implements DistributeLayoutAttr::getOffsets to generate
379 /// instructions for computing multi-dimensional offsets when distributed by
380 /// SliceAttr.
381 FailureOr<SmallVector<SmallVector<Value>>>
382 SliceAttr::getOffsets(OpBuilder &builder, Location loc, Value linearId,
383  ArrayRef<int64_t> shape) {
384  assert(getRank() == static_cast<int64_t>(shape.size()) && "invalid shape.");
385  if (!isForWorkgroup())
386  return failure();
387 
388  SmallVector<int64_t> sgLayout = getSgLayoutAsInt();
389  SmallVector<int64_t> sgShape = getSgDataAsInt();
390  if (sgShape.empty()) {
391  if (auto derivedShape = computeShapeRatio(shape, sgLayout))
392  sgShape = derivedShape.value();
393  else
394  return failure();
395  }
396 
397  // delinearize Ids
398  auto maybeIds = delinearizeSubgroupId(builder, loc, linearId);
399  if (failed(maybeIds))
400  return failure();
401 
402  // The effective sgIds for offsets computing correspond
403  // to the dims that are not sliced.
404  ArrayRef<int64_t> dims = flatten().getDims().asArrayRef();
405  SmallVector<Value> sgIds =
406  XeGPUDialect::slice(ArrayRef<Value>(*maybeIds), dims);
407 
408  return genOffsetsComputingInsts(builder, loc, sgIds, sgLayout, sgShape,
409  shape);
410 }
411 
412 //===----------------------------------------------------------------------===//
413 // XeGPU_RangeAttr
414 //===----------------------------------------------------------------------===//
415 
416 LogicalResult
418  IntegerAttr startOfRange, IntegerAttr endOfRange) {
419  if (startOfRange.getInt() >= endOfRange.getInt())
420  return emitError() << "'end' : " << endOfRange.getInt()
421  << " must be greater than 'start' : "
422  << startOfRange.getInt();
423 
424  return success();
425 }
426 
427 //===----------------------------------------------------------------------===//
428 // XeGPU_TensorDescType
429 //===----------------------------------------------------------------------===//
430 
431 mlir::Type TensorDescType::parse(AsmParser &parser) {
433  mlir::Type elementType;
434  mlir::FailureOr<mlir::Attribute> encoding;
435  mlir::FailureOr<mlir::Attribute> layout;
436 
437  // Parse literal '<'
438  if (parser.parseLess())
439  return {};
440 
441  auto shapeLoc = parser.getCurrentLocation();
442  if (mlir::failed(parser.parseDimensionList(shape))) {
443  parser.emitError(shapeLoc, "failed to parse parameter 'shape'");
444  return {};
445  }
446 
447  auto elemTypeLoc = parser.getCurrentLocation();
448  if (mlir::failed(parser.parseType(elementType))) {
449  parser.emitError(elemTypeLoc, "failed to parse parameter 'elementType'");
450  return {};
451  }
452 
453  // parse optional attributes
454  while (mlir::succeeded(parser.parseOptionalComma())) {
455  mlir::Attribute attr;
456  ParseResult res = parser.parseAttribute(attr);
457  if (mlir::succeeded(res)) {
458  if (mlir::isa<LayoutAttr>(attr)) {
459  layout = attr;
460  continue;
461  }
462  if (mlir::isa<BlockTensorDescAttr, ScatterTensorDescAttr>(attr)) {
463  encoding = attr;
464  continue;
465  }
466  }
467  return {};
468  }
469 
470  // Parse literal '>'
471  if (parser.parseGreater())
472  return {};
473 
474  MLIRContext *ctxt = parser.getContext();
475  return TensorDescType::getChecked(
476  [&]() { return parser.emitError(parser.getNameLoc()); }, ctxt, shape,
477  elementType, encoding.value_or(BlockTensorDescAttr::get(ctxt)),
478  layout.value_or(mlir::Attribute()));
479 }
480 
481 void TensorDescType::print(AsmPrinter &printer) const {
482  printer << "<";
483 
484  auto shape = getShape();
485  for (int64_t dim : shape) {
486  if (mlir::ShapedType::isDynamic(dim))
487  printer << '?';
488  else
489  printer << dim;
490  printer << 'x';
491  }
492 
493  printer << getElementType();
494 
495  auto encoding = getEncoding();
496  auto blockAttr = llvm::dyn_cast_if_present<BlockTensorDescAttr>(encoding);
497  if (encoding && (!blockAttr || !blockAttr.hasDefaultsOnly()))
498  printer << ", " << encoding;
499 
500  if (auto layout = getLayout())
501  printer << ", " << layout;
502 
503  printer << ">";
504 }
505 
506 TensorDescType TensorDescType::get(llvm::ArrayRef<int64_t> shape,
507  mlir::Type elementType, int array_length,
508  bool boundary_check,
509  MemorySpace memory_space,
510  mlir::Attribute layout) {
511  auto context = elementType.getContext();
512  auto attr = BlockTensorDescAttr::get(context, memory_space, array_length,
513  boundary_check);
514  return Base::get(context, shape, elementType, attr, layout);
515 }
516 
517 TensorDescType TensorDescType::get(llvm::ArrayRef<int64_t> shape,
518  mlir::Type elementType, int chunk_size,
519  MemorySpace memory_space,
520  mlir::Attribute layout) {
521  auto context = elementType.getContext();
522  auto attr = ScatterTensorDescAttr::get(context, memory_space, chunk_size);
523  return Base::get(context, shape, elementType, attr, layout);
524 }
525 
526 LogicalResult
528  llvm::ArrayRef<int64_t> shape, mlir::Type elementType,
529  mlir::Attribute encoding, mlir::Attribute layout) {
530  size_t rank = shape.size();
531 
532  if (rank == 0)
533  return emitError() << "expected non-zero rank tensor";
534 
535  auto blockAttr = mlir::dyn_cast_if_present<BlockTensorDescAttr>(encoding);
536  if (blockAttr) {
537  MemorySpaceAttr memorySpaceAttr = blockAttr.getMemorySpace();
538  if (rank > 1 && memorySpaceAttr &&
539  memorySpaceAttr.getValue() == MemorySpace::SLM)
540  return emitError() << "SLM is only supported for 1D block tensor";
541  }
542 
543  // for gather and scatter ops, Low-precision types are packed in 32-bit units.
544  unsigned bitWidth = elementType.getIntOrFloatBitWidth();
545  int chunkAlignmentFactor =
548  : 1;
549  auto scatterAttr = mlir::dyn_cast_if_present<ScatterTensorDescAttr>(encoding);
550  if (scatterAttr) {
551  int64_t chunkSize = scatterAttr.getChunkSizeAsInt();
552  if (rank == 1 && chunkSize != 1)
553  return emitError() << "expected non-contiguous elements for 1D tensor";
554 
555  // If chunk size > 1, the second dimension of the tensor shape must be
556  // equal to chunk size and it must be a multiple of the
557  // chunkAlignmentFactor.
558  if (chunkSize > 1) {
559  if (shape.back() != chunkSize)
560  return emitError() << "expected last dim of tensor to match chunk size";
561  if (shape.back() % chunkAlignmentFactor != 0)
562  return emitError() << "expected last dim of tensor to be a multiple of "
563  << chunkAlignmentFactor;
564  }
565  }
566 
567  auto layoutAttr = llvm::dyn_cast_if_present<LayoutAttr>(layout);
568  if (layoutAttr) {
569  if (rank != (size_t)layoutAttr.getRank())
570  return emitError() << "expected layout rank to match tensor rank";
571 
572  auto laneData = layoutAttr.getLaneData();
573  if (scatterAttr && laneData) {
574  // Validate subgroup mapping rules for scattered tensors.
575  // if chunkSize > 1, the last dimension of the tensor should
576  // be distributed in the units divisible by chunkAlignmentFactor.
577  int64_t chunkSize = scatterAttr.getChunkSizeAsInt();
578  if (chunkSize > 1 && laneData[rank - 1] % chunkAlignmentFactor)
579  return emitError()
580  << "expected last dim of lane_data to be a multiple of: "
581  << chunkAlignmentFactor;
582  }
583 
584  if (!XeGPUDialect::isEvenlyDistributable(shape, layoutAttr)) {
585  std::string shapeStr;
586  llvm::raw_string_ostream stream(shapeStr);
587  llvm::interleaveComma(shape, stream);
588  return emitError() << "cannot distribute [" << shapeStr << "] using "
589  << layoutAttr;
590  }
591  }
592  return success();
593 }
594 
595 //===----------------------------------------------------------------------===//
596 // XeGPU_MemDescType
597 //===----------------------------------------------------------------------===//
598 mlir::Type MemDescType::parse(AsmParser &parser) {
600  mlir::Type elementType;
601  mlir::FailureOr<MemLayoutAttr> layout;
602 
603  // Parse literal '<'
604  if (parser.parseLess())
605  return {};
606 
607  auto shapeLoc = parser.getCurrentLocation();
608  if (mlir::failed(parser.parseDimensionList(shape, false, true))) {
609  parser.emitError(shapeLoc, "failed to parse parameter 'shape'");
610  return {};
611  }
612 
613  auto elemTypeLoc = parser.getCurrentLocation();
614  if (mlir::failed(parser.parseType(elementType))) {
615  parser.emitError(elemTypeLoc, "failed to parse parameter 'elementType'");
616  return {};
617  }
618 
619  // parse optional attributes
620  if (mlir::succeeded(parser.parseOptionalComma())) {
621  MemLayoutAttr attr;
622  ParseResult res = parser.parseAttribute(attr);
623  if (mlir::failed(res))
624  return {};
625  layout = attr;
626  }
627 
628  // Parse literal '>'
629  if (parser.parseGreater())
630  return {};
631 
632  MLIRContext *ctxt = parser.getContext();
633  return MemDescType::getChecked(
634  [&]() { return parser.emitError(parser.getNameLoc()); }, ctxt, shape,
635  elementType, layout.value_or(MemLayoutAttr()));
636 }
637 
638 void MemDescType::print(AsmPrinter &printer) const {
639  printer << "<";
640 
641  printer.printDimensionList(getShape());
642  printer << 'x';
643  printer << getElementType();
644 
645  if (auto layout = getMemLayout())
646  printer << ", " << layout;
647 
648  printer << ">";
649 }
650 
651 //===----------------------------------------------------------------------===//
652 // XeGPU_MemDescType
653 //===----------------------------------------------------------------------===//
654 
655 Attribute MemLayoutAttr::parse(AsmParser &parser, Type type) {
656 
657  auto context = parser.getContext();
658  llvm::SMLoc loc = parser.getCurrentLocation();
659 
660  llvm::SmallDenseSet<StringRef> seenKeys;
661  SmallVector<NamedAttribute> attributes;
662 
663  auto parseElt = [&]() -> ParseResult {
664  StringRef nameId;
665  if (failed(parser.parseKeyword(&nameId)))
666  return parser.emitError(loc, "expected valid attribute name");
667 
668  if (!seenKeys.insert(nameId).second)
669  return parser.emitError(loc, "duplicate key '")
670  << nameId << " in mem layout attribute";
671 
672  if (failed(parser.parseEqual()))
673  return failure();
674 
675  Attribute attr;
676  if (failed(parser.parseAttribute(attr)))
677  return failure();
678  attributes.emplace_back(nameId, attr);
679  return success();
680  };
681 
682  // Parse literal '<'
683  if (parser.parseLess())
684  return {};
685 
686  if (failed(parser.parseCommaSeparatedList(parseElt)))
687  return {};
688 
689  // Parse literal '>'
690  if (parser.parseGreater())
691  return {};
692 
693  return parser.getChecked<MemLayoutAttr>(
694  loc, context, DictionaryAttr::get(context, attributes));
695 }
696 
697 void MemLayoutAttr::print(AsmPrinter &printer) const {
698  printer << "<";
699  ArrayRef<NamedAttribute> attrs = getAttrs().getValue();
700  for (size_t i = 0; i < attrs.size(); i++) {
701  printer << attrs[i].getName().str() << " = " << attrs[i].getValue();
702  if (i < attrs.size() - 1)
703  printer << ", ";
704  }
705  printer << ">";
706 }
707 
708 } // namespace xegpu
709 } // namespace mlir
710 
711 #include <mlir/Dialect/XeGPU/IR/XeGPUDialect.cpp.inc>
712 #define GET_ATTRDEF_CLASSES
713 #include <mlir/Dialect/XeGPU/IR/XeGPUAttrs.cpp.inc>
714 #define GET_TYPEDEF_CLASSES
715 #include <mlir/Dialect/XeGPU/IR/XeGPUTypes.cpp.inc>
static MLIRContext * getContext(OpFoldResult val)
static Type getElementType(Type type)
Determine the element type of type.
static Value min(ImplicitLocOpBuilder &builder, Value value, Value bound)
static void print(spirv::VerCapExtAttr triple, DialectAsmPrinter &printer)
static ArrayRef< int64_t > getShape(Type type)
Returns the shape of the given type.
Definition: Traits.cpp:117
Attributes are known-constant values of operations.
Definition: Attributes.h:25
MLIRContext * getContext() const
Return the context this attribute belongs to.
Definition: Attributes.cpp:37
static BoolAttr get(MLIRContext *context, bool value)
This class represents a diagnostic that is inflight and set to be reported.
Definition: Diagnostics.h:314
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
A range-style iterator that allows for iterating over the offsets of all potential tiles of size tile...
Instances of the Type class are uniqued, have an immutable identifier and an optional mutable compone...
Definition: Types.h:74
MLIRContext * getContext() const
Return the MLIRContext in which this type was uniqued.
Definition: Types.cpp:35
unsigned getIntOrFloatBitWidth() const
Return the bit width of an integer or a float type, assert failure on other types.
Definition: Types.cpp:122
This class represents an instance of an SSA value in the MLIR system, representing a computable value...
Definition: Value.h:96
Specialization of arith.constant op that returns an integer of index type.
Definition: Arith.h:113
static ConstantIndexOp create(OpBuilder &builder, Location location, int64_t value)
Definition: ArithOps.cpp:359
static DenseArrayAttrImpl get(MLIRContext *context, ArrayRef< T > content)
Builder from ArrayRef<T>.
FailureOr< SmallVector< Value > > delinearizeIndex(OpBuilder &b, Location loc, Value linearIndex, ArrayRef< Value > basis, bool hasOuterBound=true)
Generate the IR to delinearize linearIndex given the basis and return the multi-index.
Definition: Utils.cpp:1967
QueryRef parse(llvm::StringRef line, const QuerySession &qs)
Definition: Query.cpp:21
detail::InFlightRemark failed(Location loc, RemarkOpts opts)
Report an optimization remark that failed.
Definition: Remarks.h:491
auto getDims(VectorType vType)
Returns a range over the dims (size and scalability) of a VectorType.
Definition: VectorUtils.h:130
constexpr unsigned packedSizeInBitsForGatherScatter
static SmallVector< SmallVector< Value > > genOffsetsComputingInsts(OpBuilder &builder, Location loc, SmallVector< Value > sgId, ArrayRef< int64_t > sgLayout, ArrayRef< int64_t > sizePerSg, ArrayRef< int64_t > sizePerWg)
Generates instructions to compute offsets for a subgroup identified by its multidimensional indices (...
Include the generated interface declarations.
SmallVector< int64_t > computeElementwiseMul(ArrayRef< int64_t > v1, ArrayRef< int64_t > v2)
Return a vector containing llvm::zip_equal(v1, v2) multiplied elementwise.
InFlightDiagnostic emitError(Location loc)
Utility method to emit an error message using this location.
detail::DenseArrayAttrImpl< int64_t > DenseI64ArrayAttr
bool isIdentityPermutation(ArrayRef< int64_t > permutation)
Returns true if permutation is an identity permutation.
auto get(MLIRContext *context, Ts &&...params)
Helper method that injects context only if needed, this helps unify some of the attribute constructio...
detail::DenseArrayAttrImpl< int32_t > DenseI32ArrayAttr
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 verify(Operation *op, bool verifyRecursively=true)
Perform (potentially expensive) checks of invariants, used to detect compiler bugs,...
Definition: Verifier.cpp:423