save work

Author: charithaintc

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

@@ -608,8 +608,8 @@ def XeGPU_LoadGatherOp : XeGPU_Op<"load", [
       return getElementTypeOrSelf(type);
     }
 
-    Type getValueType() {
-      return getValue().getType();
+    VectorType getValueType() {
+      return llvm::dyn_cast<VectorType>(getValue().getType());
     }
 
     Type getMaskType() {
@@ -668,8 +668,8 @@ def XeGPU_StoreScatterOp : XeGPU_Op<"store", [
       return getTensorDesc().getType();
     }
 
-    Type getValueType() {
-      return getValue().getType();
+    VectorType getValueType() {
+      return llvm::dyn_cast<VectorType>(getValue().getType());
     }
 
     Type getMaskType() {

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

@@ -500,28 +500,8 @@ LogicalResult LoadGatherOp::verify() {
     transpose({1, 0}, tdescShape);
   }
 
-  auto sgMap = tdescTy.getSGMapAttr();
-  // In SIMD mode, sg_map is not present. In this case, value shape must match
-  // the tensor descriptor shape.
-  if (!sgMap)
-    return valueShape == tdescShape
-               ? success()
-               : emitOpError("Unexpected result shape")
-                     << "(Expected shape: " << makeString(tdescShape)
-                     << ", Given shape: " << makeString(valueShape) << ").\n";
-  // In SIMT mode, sg_map, wi_data, and chunk_size determine the value shape.
-  auto distributedVectorShapeOrFailure = tdescTy.getDistributedVectorType();
-  if (failed(distributedVectorShapeOrFailure))
-    return emitOpError("Failed to compute distributed vector shape for "
-                       "tensor descriptor ")
-           << tdescTy;
-  if (cast<VectorType>(valueTy) != distributedVectorShapeOrFailure.value())
-    return emitOpError("Unexpected result shape")
-           << "(Expected shape: "
-           << makeString(distributedVectorShapeOrFailure.value().getShape())
-           << ", Given shape: " << makeString(valueShape) << ").\n";
-
-  return success();
+  return isArgShapesValid(tdescTy, valueTy, tdescShape,
+                          [&]() { return emitOpError(); });
 }
 
 //===----------------------------------------------------------------------===//
@@ -555,28 +535,8 @@ LogicalResult StoreScatterOp::verify() {
     transpose({1, 0}, tdescShape);
   }
 
-  auto sgMap = tdescTy.getSGMapAttr();
-  // In SIMD mode, sg_map is not present. In this case, value shape must match
-  // the tensor descriptor shape.
-  if (!sgMap)
-    return valueShape == tdescShape
-               ? success()
-               : emitOpError("Unexpected value shape")
-                     << "(Expected shape: " << makeString(tdescShape)
-                     << ", Given shape: " << makeString(valueShape) << ").\n";
-  // In SIMT mode, sg_map, wi_data, and chunk_size determine the value shape.
-  auto distributedVectorShapeOrFailure = tdescTy.getDistributedVectorType();
-  if (failed(distributedVectorShapeOrFailure))
-    return emitOpError("Failed to compute distributed vector shape for "
-                       "tensor descriptor ")
-           << tdescTy;
-  if (cast<VectorType>(valueTy) != distributedVectorShapeOrFailure.value())
-    return emitOpError("Unexpected value shape")
-           << "(Expected shape: "
-           << makeString(distributedVectorShapeOrFailure.value().getShape())
-           << ", Given shape: " << makeString(valueShape) << ").\n";
-
-  return success();
+  return isArgShapesValid(tdescTy, valueTy, tdescShape,
+                          [&]() { return emitOpError(); });
 }
 
 //===----------------------------------------------------------------------===//

mlir/test/Dialect/XeGPU/invalid.mlir

@@ -273,7 +273,7 @@ func.func @test_load_gather_sg_map_1(%src: ui64) {
   %0 = arith.constant dense<1>: vector<4xi1>
   %cst = arith.constant dense<[0, 8, 16, 24]> : vector<4xindex>
   %1 = xegpu.create_tdesc %src, %cst : ui64, vector<4xindex> -> !xegpu.tensor_desc<4x2xf32, #xegpu.scatter_tdesc_attr<chunk_size = 2>, #xegpu.sg_map<wi_layout = [4, 1], wi_data = [1, 1]>>
-  // expected-error@+1 {{Unexpected result shape(Expected shape: [2, 1], Given shape: [1, 2])}}
+  // expected-error@+1 {{Result shape [1, 2] is not consistent with distributed vector shape [2, 1] for tensor descriptor}}
   %2 = xegpu.load %1, %0 <{l1_hint = #xegpu.cache_hint<cached>, transpose}> : !xegpu.tensor_desc<4x2xf32, #xegpu.scatter_tdesc_attr<chunk_size = 2>, #xegpu.sg_map<wi_layout = [4, 1], wi_data = [1, 1]>>, vector<4xi1> -> vector<1x2xf32>
   return
 }
@@ -283,7 +283,7 @@ func.func @test_load_gather_sg_map_2(%src: ui64) {
   %0 = arith.constant dense<1>: vector<4xi1>
   %cst = arith.constant dense<[0, 8, 16, 24]> : vector<4xindex>
   %1 = xegpu.create_tdesc %src, %cst : ui64, vector<4xindex> -> !xegpu.tensor_desc<4x2xf32, #xegpu.scatter_tdesc_attr<chunk_size = 2>, #xegpu.sg_map<wi_layout = [4, 1], wi_data = [1, 1]>>
-  // expected-error@+1 {{Unexpected result shape(Expected shape: [2, 1], Given shape: [2])}}
+  // expected-error@+1 {{esult shape [2] is not consistent with distributed vector shape [2, 1] for tensor descriptor}}
   %2 = xegpu.load %1, %0 <{l1_hint = #xegpu.cache_hint<cached>, transpose}> : !xegpu.tensor_desc<4x2xf32, #xegpu.scatter_tdesc_attr<chunk_size = 2>, #xegpu.sg_map<wi_layout = [4, 1], wi_data = [1, 1]>>, vector<4xi1> -> vector<2xf32>
   return
 }
@@ -295,7 +295,7 @@ func.func @test_store_scatter_sg_map_1(%src: ui64) {
   %cst = arith.constant dense<[0, 8, 16, 24]> : vector<4xindex>
   %val = arith.constant dense<2.9>: vector<1x2xf32>
   %1 = xegpu.create_tdesc %src, %cst : ui64, vector<4xindex> -> !xegpu.tensor_desc<4x2xf32, #xegpu.scatter_tdesc_attr<chunk_size = 2>, #xegpu.sg_map<wi_layout = [4, 1], wi_data = [1, 1]>>
-  // expected-error@+1 {{Unexpected value shape(Expected shape: [2, 1], Given shape: [1, 2])}}
+  // expected-error@+1 {{Result shape [1, 2] is not consistent with distributed vector shape [2, 1] for tensor descriptor}}
   xegpu.store %val, %1, %0 <{l1_hint = #xegpu.cache_hint<cached>, transpose}> : vector<1x2xf32>, !xegpu.tensor_desc<4x2xf32, #xegpu.scatter_tdesc_attr<chunk_size = 2>, #xegpu.sg_map<wi_layout = [4, 1], wi_data = [1, 1]>>, vector<4xi1>
   return
 }
@@ -306,7 +306,7 @@ func.func @test_store_scatter_sg_map_2(%src: ui64) {
   %cst = arith.constant dense<[0, 8, 16, 24]> : vector<4xindex>
   %val = arith.constant dense<2.9>: vector<2xf32>
   %1 = xegpu.create_tdesc %src, %cst : ui64, vector<4xindex> -> !xegpu.tensor_desc<4x2xf32, #xegpu.scatter_tdesc_attr<chunk_size = 2>, #xegpu.sg_map<wi_layout = [4, 1], wi_data = [1, 1]>>
-  // expected-error@+1 {{Unexpected value shape(Expected shape: [2, 1], Given shape: [2])}}
+  // expected-error@+1 {{esult shape [2] is not consistent with distributed vector shape [2, 1] for tensor descriptor}}
   xegpu.store %val, %1, %0 <{l1_hint = #xegpu.cache_hint<cached>, transpose}> : vector<2xf32>, !xegpu.tensor_desc<4x2xf32, #xegpu.scatter_tdesc_attr<chunk_size = 2>, #xegpu.sg_map<wi_layout = [4, 1], wi_data = [1, 1]>>, vector<4xi1>
   return
 }