revert xegpu def changes

Signed-off-by: dchigarev <[email protected]>

Author: dchigarev

mlir/include/mlir/Dialect/XeGPU/IR/XeGPUOps.td

@@ -261,21 +261,6 @@ def XeGPU_PrefetchNdOp : XeGPU_Op<"prefetch_nd", []> {
         : !xegpu.tensor_desc<8x16xf16>
     ```
 
-    The operation may take optional offsets for the tensor descriptor.
-    The number of offsets must be greater than or equal to the rank of the tensor
-    descriptor and less than or equal to the rank of the source memref.
-    The offsets are applied to the innermost dimensions of the source memref.
-
-    Examples:
-    ```mlir
-      %tdesc = xegpu.create_nd_tdesc %0: memref<2x8x32x32xf32> -> TensorDesc<8x16xf32>
-      // memref[0, 0, %off0, %off1]
-      xegpu.prefetch_nd %tdesc[%off0, %off1] : !xegpu.tensor_desc<8x16xf16>
-      // memref[0, %off0, %off1, %off2]
-      xegpu.prefetch_nd %tdesc[%off0, %off1, %off2] : !xegpu.tensor_desc<8x16xf16>
-      // memref[%off0, %off1, %off2, %off3]
-      xegpu.prefetch_nd %tdesc[%off0, %off1, %off2, %off3] : !xegpu.tensor_desc<8x16xf16>
-    ```
   }];
 
   let arguments = (ins XeGPU_TensorDesc: $TensorDesc,
@@ -365,21 +350,6 @@ def XeGPU_LoadNdOp : XeGPU_Op<"load_nd", [
         : !xegpu.tensor_desc<8x16xf32> -> vector<8xf32>
     ```
 
-    The operation may take optional offsets for the tensor descriptor.
-    The number of offsets must be greater than or equal to the rank of the tensor
-    descriptor and less than or equal to the rank of the source memref.
-    The offsets are applied to the innermost dimensions of the source memref.
-
-    Examples:
-    ```mlir
-      %1 = xegpu.create_nd_tdesc %0: memref<2x8x32x32xf32> -> TensorDesc<8x16xf32>
-      // memref[0, 0, %off0, %off1]
-      xegpu.load_nd %1[%off0, %off1] : !xegpu.tensor_desc<8x16xf32> -> vector<8xf32>
-      // memref[0, %off0, %off1, %off2]
-      xegpu.load_nd %1[%off0, %off1, %off2] : !xegpu.tensor_desc<8x16xf32> -> vector<8xf32>
-      // memref[%off0, %off1, %off2, %off3]
-      xegpu.load_nd %1[%off0, %off1, %off2, %off3] : !xegpu.tensor_desc<8x16xf32> -> vector<8xf32>
-    ```
 
   }];
 
@@ -475,21 +445,6 @@ def XeGPU_StoreNdOp : XeGPU_Op<"store_nd", [
                              : vector<8xf16>, !xegpu.tensor_desc<8x16xf16>
     ```
 
-    The operation may take optional offsets for the tensor descriptor.
-    The number of offsets must be greater than or equal to the rank of the tensor
-    descriptor and less than or equal to the rank of the source memref.
-    The offsets are applied to the innermost dimensions of the source memref.
-
-    Examples:
-    ```mlir
-      %2 = xegpu.create_nd_tdesc %0: memref<2x8x32x32xf32> -> TensorDesc<8x16xf32>
-      // memref[0, 0, %off0, %off1]
-      xegpu.store_nd %3, %2[%off0, %off1] : vector<8x16xf16>, !xegpu.tensor_desc<8x16xf16>
-      // memref[0, %off0, %off1, %off2]
-      xegpu.store_nd %3, %2[%off0, %off1, %off2] : vector<8x16xf16>, !xegpu.tensor_desc<8x16xf16>
-      // memref[%off0, %off1, %off2, %off3]
-      xegpu.store_nd %3, %2[%off0, %off1, %off2, %off3] : vector<8x16xf16>, !xegpu.tensor_desc<8x16xf16>
-    ```
 
   }];
 

mlir/lib/Dialect/XeGPU/IR/XeGPUOps.cpp

@@ -121,22 +121,6 @@ isValidGatherScatterParams(Type maskTy, VectorType valueTy,
   return success();
 }
 
-// Verify that number of offsets matches either the source rank or the tdesc
-// rank.
-static LogicalResult
-isValidNdOffset(TypedValue<TensorDescType> tDesc,
-                std::optional<llvm::ArrayRef<int64_t>> constOffsets,
-                int64_t offsetSize,
-                function_ref<InFlightDiagnostic()> emitError) {
-  int64_t constOffsetSize = constOffsets ? constOffsets->size() : 0;
-  auto tDescRank = tDesc.getType().getRank();
-  if (((offsetSize != 0) && (offsetSize < tDescRank)) ||
-      ((constOffsetSize != 0) && (constOffsetSize < tDescRank)))
-    return emitError() << "Offsets rank cannot be smaller than tensor "
-                          "descriptor rank.";
-  return success();
-}
-
 static LogicalResult
 isValidGatherScatterBufferParams(Type offsetsTy, Type maskTy,
                                  VectorType valueTy, int64_t chunkSize,
@@ -274,10 +258,8 @@ void CreateNdDescOp::build(OpBuilder &builder, OperationState &state,
     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()) {
+    // to keep the IR print clean.
+    if (staticShape == memrefShape && staticStrides == memrefStrides) {
       staticShapeAttr = DenseI64ArrayAttr();
       staticStridesAttr = DenseI64ArrayAttr();
     }
@@ -338,10 +320,8 @@ void CreateNdDescOp::build(OpBuilder &builder, OperationState &state,
     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()) {
+    // to keep the IR print clean.
+    if (staticShape == memrefShape && staticStrides == memrefStrides) {
       staticShapeAttr = DenseI64ArrayAttr();
       staticStridesAttr = DenseI64ArrayAttr();
     }
@@ -491,9 +471,16 @@ LogicalResult PrefetchNdOp::verify() {
   if (!isReadHintOrNone(getL3HintAttr()))
     return emitOpError("invalid l3_hint: ") << getL3HintAttr();
 
-  auto tDesc = getTensorDesc();
-  return isValidNdOffset(tDesc, getConstOffsets(), getMixedOffsets().size(),
-                         [&]() { return emitOpError(); });
+  int64_t tDescRank = tdescTy.getRank();
+  int64_t offsetSize = static_cast<int64_t>(getOffsets().size());
+  int64_t constOffsetSize =
+      getConstOffsetsAttr() ? getConstOffsetsAttr().size() : 0;
+  if (((offsetSize != 0) && (offsetSize != tDescRank)) ||
+      ((constOffsetSize != 0) && (constOffsetSize != tDescRank)))
+    return emitOpError(
+        "Mismatched ranks between offsets and tensor descriptor");
+
+  return success();
 }
 
 //===----------------------------------------------------------------------===//
@@ -609,9 +596,16 @@ LogicalResult LoadNdOp::verify() {
                          << " is not consistent with tensor descriptor "
                          << tdescTy;
 
-  auto tDesc = getTensorDesc();
-  return isValidNdOffset(tDesc, getConstOffsets(), getMixedOffsets().size(),
-                         [&]() { return emitOpError(); });
+  int64_t tDescRank = tdescTy.getRank();
+  int64_t offsetSize = static_cast<int64_t>(getOffsets().size());
+  int64_t constOffsetSize =
+      getConstOffsetsAttr() ? getConstOffsetsAttr().size() : 0;
+  if (((offsetSize != 0) && (offsetSize != tDescRank)) ||
+      ((constOffsetSize != 0) && (constOffsetSize != tDescRank)))
+    return emitOpError(
+        "Mismatched ranks between offsets and tensor descriptor");
+
+  return success();
 }
 
 //===----------------------------------------------------------------------===//
@@ -696,9 +690,16 @@ LogicalResult StoreNdOp::verify() {
                          << " is not consistent with tensor descriptor "
                          << dstTy;
 
-  auto tDesc = getTensorDesc();
-  return isValidNdOffset(tDesc, getConstOffsets(), getMixedOffsets().size(),
-                         [&]() { return emitOpError(); });
+  int64_t tDescRank = dstTy.getRank();
+  int64_t offsetSize = static_cast<int64_t>(getOffsets().size());
+  int64_t constOffsetSize =
+      getConstOffsetsAttr() ? getConstOffsetsAttr().size() : 0;
+  if (((offsetSize != 0) && (offsetSize != tDescRank)) ||
+      ((constOffsetSize != 0) && (constOffsetSize != tDescRank)))
+    return emitOpError(
+        "Mismatched ranks between offsets and tensor descriptor");
+
+  return success();
 }
 
 //===----------------------------------------------------------------------===//

mlir/test/Dialect/XeGPU/invalid.mlir

@@ -132,10 +132,18 @@ func.func @subgroup_load_nd_9(%src: memref<4x8x16xf16>) {
   return
 }
 
+// -----
+func.func @subgroup_load_nd_offset_1(%src: memref<4x8x16xf16>, %x : index) {
+  %1 = xegpu.create_nd_tdesc %src: memref<4x8x16xf16> -> !xegpu.tensor_desc<16xf16>
+// expected-error@+1 {{Mismatched ranks between offsets and tensor descriptor}}
+  %2 = xegpu.load_nd %1[0, 0] : !xegpu.tensor_desc<16xf16> -> vector<16xf16>
+  return
+}
+
 // -----
 func.func @subgroup_load_nd_offset_2(%src: memref<4x8x16xf16>, %x : index) {
   %3 = xegpu.create_nd_tdesc %src: memref<4x8x16xf16> -> !xegpu.tensor_desc<8x16xf16>
-    // expected-error@+1 {{Offsets rank cannot be smaller than tensor descriptor rank.}}
+    // expected-error@+1 {{Mismatched ranks between offsets and tensor descriptor}}
   xegpu.prefetch_nd %3[0] : !xegpu.tensor_desc<8x16xf16>
   return
 }
@@ -144,7 +152,7 @@ func.func @subgroup_load_nd_offset_2(%src: memref<4x8x16xf16>, %x : index) {
 func.func @subgroup_load_nd_offset_3(%src: memref<4x8x16xf16>, %x : index) {
   %3 = xegpu.create_nd_tdesc %src: memref<4x8x16xf16> -> !xegpu.tensor_desc<8x16xf16>
   %5 = xegpu.load_nd %3[0, 0] : !xegpu.tensor_desc<8x16xf16> -> vector<8x16xf16>
-    // expected-error@+1 {{Offsets rank cannot be smaller than tensor descriptor rank.}}
+    // expected-error@+1 {{Mismatched ranks between offsets and tensor descriptor}}
   xegpu.store_nd %5, %3[%x] : vector<8x16xf16>, !xegpu.tensor_desc<8x16xf16>
   return
 }