xegpu: setDescLayout retains TensorDesc BlockTensorDescAttrs

Author: tkarna

mlir/lib/Dialect/XeGPU/TransformOps/XeGPUTransformOps.cpp

@@ -125,13 +125,12 @@ xegpu::LayoutAttr createLayoutAttr(MLIRContext *ctx, ArrayRef<int32_t> sgLayout,
 xegpu::CreateNdDescOp setDescLayout(transform::TransformRewriter &rewriter,
                                     xegpu::CreateNdDescOp descOp,
                                     xegpu::LayoutAttr layout) {
-  auto oldTensorDesc = descOp.getResult();
-  auto descShapedType = cast<ShapedType>(oldTensorDesc.getType());
+  auto oldTensorDesc = descOp.getType();
   auto descType = xegpu::TensorDescType::get(
-      descShapedType.getShape(), descShapedType.getElementType(),
-      /*array_length=*/1,
-      /*boundary_check=*/true,
-      /*memory_space=*/xegpu::MemorySpace::Global,
+      oldTensorDesc.getShape(), oldTensorDesc.getElementType(),
+      /*array_length=*/oldTensorDesc.getArrayLength(),
+      /*boundary_check=*/oldTensorDesc.getBoundaryCheck(),
+      /*memory_space=*/oldTensorDesc.getMemorySpace(),
       /*layout=*/layout);
 
   rewriter.setInsertionPointAfter(descOp);

mlir/test/Dialect/XeGPU/transform-ops.mlir

@@ -3,8 +3,9 @@
 // CHECK-LABEL: @set_desc_layout
 func.func @set_desc_layout(%arg0: memref<4096x4096xf16>) {
   // CHECK: %[[V0:.+]] = xegpu.create_nd_tdesc %arg0
+  // CHECK-SAME: #xegpu.block_tdesc_attr<boundary_check = false>
   // CHECK-SAME: #xegpu.layout<sg_layout = [8, 4], sg_data = [32, 32], inst_data = [8, 16]>>
-  %0 = xegpu.create_nd_tdesc %arg0 : memref<4096x4096xf16> -> !xegpu.tensor_desc<256x32xf16>
+  %0 = xegpu.create_nd_tdesc %arg0 : memref<4096x4096xf16> -> !xegpu.tensor_desc<256x32xf16, #xegpu.block_tdesc_attr<boundary_check = false>>
   return
 }