Move memory space check to type verifier

Author: adam-smnk

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

@@ -232,9 +232,6 @@ LogicalResult TensorDescType::verify(
   if (rank != 1 && rank != 2)
     return emitError() << "expected 1D or 2D tensor";
 
-  // Scattered attribute imposes extra restriction on tensor descriptor.
-  // Block attribute can only be validated further against data transfer
-  // operations.
   auto scatterAttr = mlir::dyn_cast_if_present<ScatterTensorDescAttr>(encoding);
   if (scatterAttr) {
     // Expected tensor ranks for scattered data:
@@ -248,6 +245,14 @@ LogicalResult TensorDescType::verify(
       return emitError() << "expected chunk blocks for 2D tensor";
   }
 
+  if (auto blockAttr =
+          mlir::dyn_cast_if_present<BlockTensorDescAttr>(encoding)) {
+    MemorySpaceAttr memorySpaceAttr = blockAttr.getMemorySpace();
+    if (rank == 2 && memorySpaceAttr &&
+        memorySpaceAttr.getValue() == MemorySpace::SLM)
+      return emitError() << "SLM is not supported for 2D block tensor";
+  }
+
   if (auto sgMapAttr = llvm::dyn_cast_if_present<SGMapAttr>(sg_map)) {
     ArrayRef<uint32_t> wiLayout = sgMapAttr.getWiLayout();
     ArrayRef<uint32_t> wiData = sgMapAttr.getWiData();

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

@@ -231,10 +231,6 @@ LogicalResult CreateNdDescOp::verify() {
   if (getType().isScattered())
     return emitOpError("Expects a non-scattered TensorDesc.\n");
 
-  if (getType().getRank() == 2 &&
-      tdescMemorySpace == static_cast<unsigned>(MemorySpace::SLM))
-    return emitOpError("SLM is not supported for 2D Block TensorDesc.\n");
-
   return success();
 }
 

mlir/test/Dialect/XeGPU/invalid.mlir

@@ -17,7 +17,7 @@ func.func @test_create_nd_tdesc_vc_2(%src: memref<24x32xf32>) {
 
 // -----
 func.func @test_create_nd_tdesc_vc_3(%src: memref<2x24x32xf32, 3>) {
-  // expected-error@+1 {{SLM is not supported for 2D Block TensorDesc}}
+  // expected-error@+1 {{SLM is not supported for 2D block tensor}}
   %1 = xegpu.create_nd_tdesc %src[0, 0, 0] : memref<2x24x32xf32, 3> -> !xegpu.tensor_desc<8x16xf32, #xegpu.block_tdesc_attr<memory_space = slm>>
   return
 }