[Mlir-commits] [mlir] [MLIR][XeGPU] Remove offsets from create_nd_tdesc & remove update_nd_offset, move offsets to load/store/prefetch ops (PR #193330)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Tue Apr 21 19:54:47 PDT 2026
llvmbot wrote:
<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-mlir-gpu
Author: Nishant Patel (nbpatel)
<details>
<summary>Changes</summary>
This PR removes the optional offsets/const_offsets operands on xegpu.create_nd_tdesc and instead mandates offsets directly on the consuming load, store, and prefetch ops. It also deprecates the update_nd_offset op.
---
Patch is 517.08 KiB, truncated to 20.00 KiB below, full version: https://github.com/llvm/llvm-project/pull/193330.diff
26 Files Affected:
- (modified) mlir/include/mlir/Dialect/XeGPU/IR/XeGPUOps.td (+19-139)
- (modified) mlir/lib/Conversion/XeGPUToXeVM/XeGPUToXeVM.cpp (-3)
- (modified) mlir/lib/Dialect/XeGPU/IR/XeGPUOps.cpp (+22-216)
- (modified) mlir/lib/Dialect/XeGPU/TransformOps/XeGPUTransformOps.cpp (+1-6)
- (modified) mlir/lib/Dialect/XeGPU/Transforms/XeGPUBlocking.cpp (+1-2)
- (modified) mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp (-21)
- (modified) mlir/lib/Dialect/XeGPU/Transforms/XeGPUSubgroupDistribute.cpp (-5)
- (modified) mlir/lib/Dialect/XeGPU/Transforms/XeGPUUnroll.cpp (+45-283)
- (modified) mlir/lib/Dialect/XeGPU/Transforms/XeGPUWgToSgDistribute.cpp (+31-168)
- (modified) mlir/test/Dialect/XeGPU/invalid.mlir (+45-70)
- (modified) mlir/test/Dialect/XeGPU/layout.mlir (+10-10)
- (modified) mlir/test/Dialect/XeGPU/ops.mlir (+96-134)
- (modified) mlir/test/Dialect/XeGPU/propagate-layout-inst-data.mlir (+38-56)
- (modified) mlir/test/Dialect/XeGPU/propagate-layout-subgroup.mlir (+14-14)
- (modified) mlir/test/Dialect/XeGPU/propagate-layout.mlir (+164-174)
- (modified) mlir/test/Dialect/XeGPU/xegpu-blocking.mlir (+108-149)
- (modified) mlir/test/Dialect/XeGPU/xegpu-recover-layout.mlir (+16-18)
- (removed) mlir/test/Dialect/XeGPU/xegpu-unroll-patterns-no-desc-offsets.mlir (-61)
- (modified) mlir/test/Dialect/XeGPU/xegpu-unroll-patterns.mlir (+56-100)
- (modified) mlir/test/Dialect/XeGPU/xegpu-vector-linearize.mlir (+12-13)
- (modified) mlir/test/Dialect/XeGPU/xegpu-wg-to-sg-elemwise.mlir (+30-30)
- (modified) mlir/test/Dialect/XeGPU/xegpu-wg-to-sg-rr.mlir (+245-107)
- (removed) mlir/test/Dialect/XeGPU/xegpu-wg-to-sg-unify-ops-rr.mlir (-246)
- (removed) mlir/test/Dialect/XeGPU/xegpu-wg-to-sg-unify-ops.mlir (-963)
- (modified) mlir/test/Dialect/XeGPU/xegpu-wg-to-sg.mlir (+928-202)
- (modified) mlir/test/lib/Dialect/XeGPU/TestXeGPUTransforms.cpp (+3-6)
``````````diff
diff --git a/mlir/include/mlir/Dialect/XeGPU/IR/XeGPUOps.td b/mlir/include/mlir/Dialect/XeGPU/IR/XeGPUOps.td
index 31fe93d209a6d..6995125465d90 100644
--- a/mlir/include/mlir/Dialect/XeGPU/IR/XeGPUOps.td
+++ b/mlir/include/mlir/Dialect/XeGPU/IR/XeGPUOps.td
@@ -77,11 +77,6 @@ def XeGPU_CreateNdDescOp: XeGPU_Op<"create_nd_tdesc", [Pure, ViewLikeOpInterface
For the case of dynamic memrefs or pointer, the shape and layout information of the
memory region should be explicitly passed via `shape` and `strides` parameters.
- - `offsets`: [optional] index values represents offsets from the "source" at the each dimension
- at which the subview of the target memory will be created. It is encoded via
- "offsets" and "const_offsets", such that it can accept various forms, such as,
- operands (e.g., [%c0, %c]) and attributes (e.g., [2, 4]). Offsets is optional and may be set at load_nd, store_nd, and prefetch_nd.
-
- `shape`: the shape information of the memory region pointed by the "source". It is
typically encoded via the MemRefType of the source, e.g., memref<4096x4096xf16>.
But if "source" is simply a pointer represented as uint64_t type, or a memref
@@ -100,41 +95,34 @@ def XeGPU_CreateNdDescOp: XeGPU_Op<"create_nd_tdesc", [Pure, ViewLikeOpInterface
Example 1 (suppose the tensor shape inferred by the compiler is 8x16):
```mlir
%0 = memref.alloc() : memref<1024x1024xf32>
- %c0 = arith.constant 0 : index
- %c1 = arith.constant 1 : index
- %1 = xegpu.create_nd_tdesc %0[%c0, %c0]: memref<1024x1024xf32> -> TensorDesc<8x16xf32>
+ %1 = xegpu.create_nd_tdesc %0 : memref<1024x1024xf32> -> TensorDesc<8x16xf32>
```
Example 2 (suppose the tensor shape inferred by the compiler is 8x16):
```mlir
%0 = memref.alloc(%h, %w) : memref<?x?xf32>
- %c0 = arith.constant 0 : index
%c1 = arith.constant 1 : index
- %1 = xegpu.create_nd_tdesc %0[%c0, %c0], [%h, %w], [%w, %c1]: memref<?x?xf32> -> TensorDesc<8x16xf32>
+ %1 = xegpu.create_nd_tdesc %0, shape:[%h, %w], strides:[%w, %c1]: memref<?x?xf32> -> TensorDesc<8x16xf32>
```
Example 3 (suppose the tensor shape inferred by the compiler is 8x16):
```mlir
%0 = ... : ui64
- %c0 = arith.constant 0 : index
%c1 = arith.constant 1 : index
- %1 = xegpu.create_nd_tdesc %0[%c0, %c0], [%h, %w], [%w, %c1]: ui64 -> TensorDesc<8x16xf32>
+ %1 = xegpu.create_nd_tdesc %0, shape:[%h, %w], strides:[%w, %c1]: ui64 -> TensorDesc<8x16xf32>
```
}];
let arguments = (ins
XeGPU_BaseAddrType: $source,
- Variadic<Index>: $offsets,
Variadic<Index>: $shape,
Variadic<Index>: $strides,
- OptionalAttr<DenseI64ArrayAttr>: $const_offsets,
OptionalAttr<DenseI64ArrayAttr>: $const_shape,
OptionalAttr<DenseI64ArrayAttr>: $const_strides
);
let assemblyFormat = [{
$source ``
- custom<OptionalDynamicIndexList>($offsets, $const_offsets)
(`,` `shape` `:` custom<DynamicIndexList>($shape, $const_shape)^
`,` `strides``:` custom<DynamicIndexList>($strides, $const_strides))?
attr-dict `:` type($source) `->` qualified(type($TensorDesc))
@@ -148,14 +136,6 @@ def XeGPU_CreateNdDescOp: XeGPU_Op<"create_nd_tdesc", [Pure, ViewLikeOpInterface
OpBuilder<(ins "Type": $tdesc, "TypedValue<MemRefType>": $source)>,
OpBuilder<(ins "Type": $tdesc, "Value ": $source,
- "llvm::ArrayRef<OpFoldResult>": $shape,
- "llvm::ArrayRef<OpFoldResult>": $strides)>,
-
- OpBuilder<(ins "Type": $tdesc, "TypedValue<MemRefType>": $source,
- "llvm::ArrayRef<OpFoldResult>": $offsets)>,
-
- OpBuilder<(ins "Type": $tdesc, "Value": $source,
- "llvm::ArrayRef<OpFoldResult>": $offsets,
"llvm::ArrayRef<OpFoldResult>": $shape,
"llvm::ArrayRef<OpFoldResult>": $strides)>
];
@@ -181,14 +161,6 @@ def XeGPU_CreateNdDescOp: XeGPU_Op<"create_nd_tdesc", [Pure, ViewLikeOpInterface
return getType().getShape();
}
- SmallVector<OpFoldResult> getMixedOffsets() {
- auto statics = getConstOffsets().value_or(SmallVector<int64_t>());
- auto dynamics = getOffsets();
- if (statics.size() == 0 && dynamics.size() == 0)
- return {};
- return getMixedValues(statics, dynamics, getContext());
- }
-
SmallVector<OpFoldResult> getMixedSizes() {
SmallVector<int64_t> statics;
@@ -261,7 +233,7 @@ def XeGPU_PrefetchNdOp : XeGPU_Op<"prefetch_nd", [AnchorLayoutInterface]> {
- `TensorDesc`: A tensor descriptor specifying the base nd-region of
memory and tensor tile to be prefetched.
- - `offsets`: [optional] index values representing per-dimension offsets from the
+ - `offsets`: index values representing per-dimension offsets from the
base position encoded in `TensorDesc`. It is encoded via "offsets"
and "const_offsets".
@@ -286,7 +258,7 @@ def XeGPU_PrefetchNdOp : XeGPU_Op<"prefetch_nd", [AnchorLayoutInterface]> {
let arguments = (ins XeGPU_TensorDesc: $TensorDesc,
Variadic<Index>: $offsets,
- OptionalAttr<DenseI64ArrayAttr>: $const_offsets,
+ DenseI64ArrayAttr: $const_offsets,
OptionalAttr<XeGPU_CacheHintAttr>: $l1_hint,
OptionalAttr<XeGPU_CacheHintAttr>: $l2_hint,
OptionalAttr<XeGPU_CacheHintAttr>: $l3_hint,
@@ -306,11 +278,7 @@ def XeGPU_PrefetchNdOp : XeGPU_Op<"prefetch_nd", [AnchorLayoutInterface]> {
}
SmallVector<OpFoldResult> getMixedOffsets() {
- auto statics = getConstOffsets().value_or(SmallVector<int64_t>());
- auto dynamics = getOffsets();
- if (statics.size() == 0 && dynamics.size() == 0)
- return {};
- return getMixedValues(statics, dynamics, getContext());
+ return getMixedValues(getConstOffsets(), getOffsets(), getContext());
}
xegpu::DistributeLayoutAttr getDescLayoutAttr() {
@@ -325,15 +293,11 @@ def XeGPU_PrefetchNdOp : XeGPU_Op<"prefetch_nd", [AnchorLayoutInterface]> {
let assemblyFormat = [{
$TensorDesc ``
- custom<OptionalDynamicIndexList>($offsets, $const_offsets)
+ custom<DynamicIndexList>($offsets, $const_offsets)
prop-dict attr-dict `:` qualified(type($TensorDesc))
}];
let builders = [
- OpBuilder<(ins "Value": $TensorDesc,
- "xegpu::CachePolicyAttr": $l1_hint,
- "xegpu::CachePolicyAttr": $l2_hint,
- "xegpu::CachePolicyAttr": $l3_hint)>,
OpBuilder<(ins "Value": $TensorDesc,
"ArrayRef<OpFoldResult>": $offsets,
"xegpu::CachePolicyAttr": $l1_hint,
@@ -409,7 +373,7 @@ def XeGPU_LoadNdOp : XeGPU_Op<"load_nd", [
let arguments = (ins XeGPU_TensorDesc: $TensorDesc,
Variadic<Index>: $offsets,
- OptionalAttr<DenseI64ArrayAttr>: $const_offsets,
+ DenseI64ArrayAttr: $const_offsets,
OptionalAttr<UnitAttr>: $packed,
OptionalAttr<DenseI64ArrayAttr>: $transpose,
OptionalAttr<XeGPU_CacheHintAttr>: $l1_hint,
@@ -437,11 +401,7 @@ def XeGPU_LoadNdOp : XeGPU_Op<"load_nd", [
}
SmallVector<OpFoldResult> getMixedOffsets() {
- auto statics = getConstOffsets().value_or(SmallVector<int64_t>());
- auto dynamics = getOffsets();
- if (statics.size() == 0 && dynamics.size() == 0)
- return {};
- return getMixedValues(statics, dynamics, getContext());
+ return getMixedValues(getConstOffsets(), getOffsets(), getContext());
}
xegpu::DistributeLayoutAttr getDescLayoutAttr() {
@@ -456,16 +416,11 @@ def XeGPU_LoadNdOp : XeGPU_Op<"load_nd", [
let assemblyFormat = [{
$TensorDesc ``
- custom<OptionalDynamicIndexList>($offsets, $const_offsets)
+ custom<DynamicIndexList>($offsets, $const_offsets)
prop-dict attr-dict `:` qualified(type($TensorDesc)) `->` type($value)
}];
let builders = [
- OpBuilder<(ins "Type": $value, "Value": $TensorDesc,
- "UnitAttr": $packed, "DenseI64ArrayAttr": $transpose,
- "xegpu::CachePolicyAttr": $l1_hint,
- "xegpu::CachePolicyAttr": $l2_hint,
- "xegpu::CachePolicyAttr": $l3_hint)>,
OpBuilder<(ins "Type": $value, "Value": $TensorDesc,
"ArrayRef<OpFoldResult>": $offsets,
"UnitAttr": $packed, "DenseI64ArrayAttr": $transpose,
@@ -534,7 +489,7 @@ def XeGPU_StoreNdOp : XeGPU_Op<"store_nd", [
let arguments = (ins XeGPU_ValueType: $value,
XeGPU_TensorDesc: $TensorDesc,
Variadic<Index>: $offsets,
- OptionalAttr<DenseI64ArrayAttr>: $const_offsets,
+ DenseI64ArrayAttr: $const_offsets,
OptionalAttr<XeGPU_CacheHintAttr>: $l1_hint,
OptionalAttr<XeGPU_CacheHintAttr>: $l2_hint,
OptionalAttr<XeGPU_CacheHintAttr>: $l3_hint,
@@ -558,11 +513,7 @@ def XeGPU_StoreNdOp : XeGPU_Op<"store_nd", [
}
SmallVector<OpFoldResult> getMixedOffsets() {
- auto statics = getConstOffsets().value_or(SmallVector<int64_t>());
- auto dynamics = getOffsets();
- if (statics.size() == 0 && dynamics.size() == 0)
- return {};
- return getMixedValues(statics, dynamics, getContext());
+ return getMixedValues(getConstOffsets(), getOffsets(), getContext());
}
xegpu::DistributeLayoutAttr getDescLayoutAttr() {
@@ -578,15 +529,11 @@ def XeGPU_StoreNdOp : XeGPU_Op<"store_nd", [
let assemblyFormat = [{
$value `,`
$TensorDesc ``
- custom<OptionalDynamicIndexList>($offsets, $const_offsets)
+ custom<DynamicIndexList>($offsets, $const_offsets)
prop-dict attr-dict `:` type($value) `,` qualified(type($TensorDesc))
}];
let builders = [
- OpBuilder<(ins "Value": $value, "Value": $TensorDesc,
- "xegpu::CachePolicyAttr": $l1_hint,
- "xegpu::CachePolicyAttr": $l2_hint,
- "xegpu::CachePolicyAttr": $l3_hint)>,
OpBuilder<(ins "Value": $value, "Value": $TensorDesc,
"ArrayRef<OpFoldResult>": $offsets,
"xegpu::CachePolicyAttr": $l1_hint,
@@ -599,55 +546,6 @@ def XeGPU_StoreNdOp : XeGPU_Op<"store_nd", [
let hasVerifier = 1;
}
-def XeGPU_UpdateNdOffsetOp : XeGPU_Op<"update_nd_offset",
- [Pure, AllTypesMatch<["TensorDesc", "result"]>]> {
- let summary = "It updates the offsets for the TensorDesc.";
- let description = [{The op updates the offset of the given TensorDesc.
- The offsets are relative offset to the current position in the number
- of elements. It will result in a same type TensorDesc as the input.
-
- Example:
- ```
- %2 = xegpu.update_nd_offset %1, [0, 16]: !xegpu.tensor_desc<8x16xf32>
- ```
- }];
-
- let arguments = (ins
- XeGPU_TensorDesc: $TensorDesc,
- Variadic<Index>: $offsets,
- DenseI64ArrayAttr: $const_offsets);
-
- let results = (outs XeGPU_TensorDesc: $result);
-
- let extraClassDeclaration = extraBaseClassDeclaration # [{
- xegpu::TensorDescType getTensorDescType() {
- return getTensorDesc().getType();
- }
-
- SmallVector<OpFoldResult> getMixedOffsets() {
- Builder b(getContext());
- return getMixedValues(getConstOffsets(), getOffsets(), b);
- }
-
- size_t getNumOffsets() {
- return getMixedOffsets().size();
- }
-
- OpFoldResult getOffset(unsigned idx) {
- assert(idx < getNumOffsets() && "Invalid out of bound access.");
- return getMixedOffsets()[idx];
- }
- }];
-
- let assemblyFormat = [{
- $TensorDesc `,`
- custom<DynamicIndexList>($offsets, $const_offsets)
- attr-dict `:` qualified(type($result))
- }];
-
- let hasVerifier = 1;
-}
-
def XeGPU_PrefetchOp : XeGPU_Op<"prefetch", [AnchorLayoutInterface]> {
let summary = "prefetches a set of scattered data points to cache";
@@ -721,7 +619,7 @@ def XeGPU_PrefetchOp : XeGPU_Op<"prefetch", [AnchorLayoutInterface]> {
}];
let arguments = (ins XeGPU_GatherScatterSourceType:$source,
- Optional<AnyTypeOf<[XeGPU_OffsetType, Index]>>:$offsets,
+ AnyTypeOf<[XeGPU_OffsetType, Index]>:$offsets,
OptionalAttr<XeGPU_CacheHintAttr>:$l1_hint,
OptionalAttr<XeGPU_CacheHintAttr>:$l2_hint,
OptionalAttr<XeGPU_CacheHintAttr>:$l3_hint,
@@ -755,19 +653,11 @@ def XeGPU_PrefetchOp : XeGPU_Op<"prefetch", [AnchorLayoutInterface]> {
}];
let assemblyFormat = [{
- $source
- (`[` $offsets^ `]`)?
+ $source `[` $offsets `]`
prop-dict
attr-dict `:` type(operands)
}];
- let builders = [
- OpBuilder<(ins "Value": $source,
- "xegpu::CachePolicyAttr": $l1_hint,
- "xegpu::CachePolicyAttr": $l2_hint,
- "xegpu::CachePolicyAttr": $l3_hint)>
- ];
-
let hasVerifier = 1;
}
@@ -838,7 +728,7 @@ def XeGPU_LoadGatherOp : XeGPU_Op<"load", [MemoryEffects<[MemRead]>, AnchorLayou
}];
let arguments = (ins XeGPU_GatherScatterSourceType:$source,
- Optional<AnyTypeOf<[XeGPU_OffsetType, Index]>>:$offsets,
+ AnyTypeOf<[XeGPU_OffsetType, Index]>:$offsets,
AnyTypeOf<[XeGPU_MaskType, I1]>:$mask, OptionalAttr<I64Attr>:$chunk_size,
OptionalAttr<XeGPU_CacheHintAttr>:$l1_hint,
OptionalAttr<XeGPU_CacheHintAttr>:$l2_hint,
@@ -887,17 +777,12 @@ def XeGPU_LoadGatherOp : XeGPU_Op<"load", [MemoryEffects<[MemRead]>, AnchorLayou
}];
let assemblyFormat = [{
- $source
- (`[` $offsets^ `]`)? `,`
+ $source `[` $offsets `]` `,`
$mask prop-dict
attr-dict `:` type(operands) `->` type($value)
}];
let builders = [
- OpBuilder<(ins "Type": $value, "Value": $source, "Value": $mask,
- "xegpu::CachePolicyAttr": $l1_hint,
- "xegpu::CachePolicyAttr": $l2_hint,
- "xegpu::CachePolicyAttr": $l3_hint)>,
OpBuilder<(ins "Type": $value, "Value": $source,
"ArrayRef<OpFoldResult>": $offsets, "Value": $mask,
"IntegerAttr": $chunk_size,
@@ -984,7 +869,7 @@ def XeGPU_StoreScatterOp : XeGPU_Op<"store", [MemoryEffects<[MemWrite]>, AnchorL
let arguments = (ins XeGPU_ValueOrScalarType:$value,
XeGPU_GatherScatterSourceType:$dest,
- Optional<AnyTypeOf<[XeGPU_OffsetType, Index]>>:$offsets,
+ AnyTypeOf<[XeGPU_OffsetType, Index]>:$offsets,
AnyTypeOf<[XeGPU_MaskType, I1]>:$mask, OptionalAttr<I64Attr>:$chunk_size,
OptionalAttr<XeGPU_CacheHintAttr>:$l1_hint,
OptionalAttr<XeGPU_CacheHintAttr>:$l2_hint,
@@ -1031,18 +916,13 @@ def XeGPU_StoreScatterOp : XeGPU_Op<"store", [MemoryEffects<[MemWrite]>, AnchorL
let assemblyFormat = [{
$value `,`
- $dest
- (`[` $offsets^ `]`)? `,`
+ $dest `[` $offsets `]` `,`
$mask
prop-dict
attr-dict `:` type(operands)
}];
let builders = [
- OpBuilder<(ins "Value": $value, "Value": $dest, "Value": $mask,
- "xegpu::CachePolicyAttr": $l1_hint,
- "xegpu::CachePolicyAttr": $l2_hint,
- "xegpu::CachePolicyAttr": $l3_hint)>,
OpBuilder<(ins "Value": $value, "Value": $dest,
"ArrayRef<OpFoldResult>": $offsets, "Value": $mask,
"IntegerAttr": $chunk_size,
diff --git a/mlir/lib/Conversion/XeGPUToXeVM/XeGPUToXeVM.cpp b/mlir/lib/Conversion/XeGPUToXeVM/XeGPUToXeVM.cpp
index 50eba56a16080..2d472ba163ff8 100644
--- a/mlir/lib/Conversion/XeGPUToXeVM/XeGPUToXeVM.cpp
+++ b/mlir/lib/Conversion/XeGPUToXeVM/XeGPUToXeVM.cpp
@@ -186,9 +186,6 @@ class CreateNdDescToXeVMPattern
matchAndRewrite(xegpu::CreateNdDescOp op,
xegpu::CreateNdDescOp::Adaptor adaptor,
ConversionPatternRewriter &rewriter) const override {
- SmallVector<OpFoldResult> mixedOffsets = op.getMixedOffsets();
- if (mixedOffsets.size() != 0)
- return rewriter.notifyMatchFailure(op, "Offsets not supported.");
auto loc = op.getLoc();
auto source = op.getSource();
// Op is lowered to a code sequence that populates payload.
diff --git a/mlir/lib/Dialect/XeGPU/IR/XeGPUOps.cpp b/mlir/lib/Dialect/XeGPU/IR/XeGPUOps.cpp
index 51ce6ce53a2fe..78e60ad815985 100644
--- a/mlir/lib/Dialect/XeGPU/IR/XeGPUOps.cpp
+++ b/mlir/lib/Dialect/XeGPU/IR/XeGPUOps.cpp
@@ -205,10 +205,8 @@ void CreateNdDescOp::build(OpBuilder &builder, OperationState &state,
[[maybe_unused]] auto ty = source.getType();
assert(ty.hasStaticShape() && "expecting a memref with static shape");
- build(builder, state, tdesc, source, ValueRange({}) /* dynamic offsets */,
- ValueRange({}) /* empty dynamic shape */,
+ build(builder, state, tdesc, source, ValueRange({}) /* empty dynamic shape */,
ValueRange({}) /* empty dynamic strides */,
- DenseI64ArrayAttr({}) /* const offsets */,
DenseI64ArrayAttr({}) /* empty const shape*/,
DenseI64ArrayAttr({}) /* empty const strides*/);
}
@@ -247,72 +245,8 @@ void CreateNdDescOp::build(OpBuilder &builder, OperationState &state,
}
}
- build(builder, state, tdesc, source, ValueRange({}), dynamicShape,
- dynamicStrides, builder.getDenseI64ArrayAttr({}), staticShapeAttr,
- staticStridesAttr);
-}
-
-void CreateNdDescOp::build(OpBuilder &builder, OperationState &state,
- Type tdesc, TypedValue<MemRefType> source,
- llvm::ArrayRef<OpFoldResult> offsets) {
- [[maybe_unused]] auto ty = source.getType();
- assert(ty.hasStaticShape() && offsets.size() == (size_t)ty.getRank());
-
- llvm::SmallVector<int64_t> staticOffsets;
- llvm::SmallVector<Value> dynamicOffsets;
- dispatchIndexOpFoldResults(offsets, dynamicOffsets, staticOffsets);
-
- build(builder, state, tdesc, source, dynamicOffsets /* dynamic offsets */,
- ValueRange({}) /* empty dynamic shape */,
- ValueRange({}) /* empty dynamic strides */,
- builder.getDenseI64ArrayAttr(staticOffsets) /* const offsets */,
- {} /* empty const shape*/, {} /* empty const strides*/);
-}
-
-void CreateNdDescOp::build(OpBuilder &builder, OperationState &state,
- Type tdesc, Value source,
- llvm::ArrayRef<OpFoldResult> offsets,
- llvm::ArrayRef<OpFoldResult> shape,
- llvm::ArrayRef<OpFoldResult> strides) {
- assert(!shape.empty() && !offsets.empty() && !strides.empty() &&
- shape.size() == strides.size() && shape.size() == offsets.size());
-
- Type srcTy = source.getType();
- assert((isa<IntegerType, MemRefType>(srcTy)) &&
- "Source has to be either int or memref.");
-
- llvm::SmallVector<Value> dynamicOffsets;
- llvm::SmallVector<Value> dynamicShape;
- llvm::SmallVector<Value> dynamicStrides;
-
- llvm::SmallVector<int64_t> staticOffsets;
- llvm::SmallVector<int64_t> staticShape;
- llvm::SmallVector<int64_t> staticStrides;
-
- dispatchIndexOpFoldResults(offsets, dynamicOffsets, staticOffsets);
- dispatchIndexOpFoldResults(shape, dynamicShape, staticShape);
- dispatchIndexOpFoldResults(strides, dynamicStrides, staticStrides);
-
- auto staticOffsetsAttr = builder.getDenseI64ArrayAttr(staticOffsets);
- auto staticShapeAttr = builder.getDenseI64ArrayAttr(staticShape);
- auto staticStridesAttr = builder.getDenseI64ArrayAttr(staticStrides);
-
- if (auto memrefTy = dyn_cast<MemRefType>(srcTy)) {
- auto memrefShape = memrefTy.getShape();
- auto [memrefStrides, _] = memrefTy.getStridesAndOffset();
-
- // if shape and strides are from Memref, we don't need attributes for them
- // to keep the IR print clean (only do so for full-static case, otherwise
- // printer would fail trying to print empty array-attr).
- if (staticShape == memrefShape && staticStrides == memrefStrides &&
- dynamicShape.empty() && dynamicStrides.empty()) {
- staticShapeAttr = DenseI64ArrayAttr();
- staticStridesAttr = DenseI64ArrayAttr();
- }
- }
-
- build(builder, state, tdesc, source, dynamicOffsets, dynamicShape,
- dynamicStrides, staticOffsetsAttr, staticShapeAttr, staticStridesAttr);
+ build(builder, state, tdesc, source, dynamicShape, dynamicStr...
[truncated]
``````````
</details>
https://github.com/llvm/llvm-project/pull/193330
More information about the Mlir-commits
mailing list