XeVM tests enabling (#1069)

Author: Garra1980

lib/Conversion/XeGPUToXeVM/XeGPUToXeVM.cpp

@@ -121,13 +121,6 @@ class CreateNdDescToXeVMPattern
                   xegpu::CreateNdDescOp::Adaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override {
     auto loc = op.getLoc();
-    auto resultDesc = cast<TensorDescType>(op.getResult().getType());
-    auto sgMap = resultDesc.getLayoutAttr();
-    if (!sgMap) {
-      op.emitError() << "XeVM expects SGMap attribute to be present for tensor "
-                        "descriptors";
-      return mlir::failure();
-    }
     auto source = op.getSource();
     Type payloadElemTy = rewriter.getI32Type();
     Type i64Ty = rewriter.getI64Type();
@@ -292,8 +285,7 @@ class LoadStorePrefetchNdToXeVMPattern : public OpConversionPattern<OpType> {
       auto l3 = translateStoreXeGPUCacheHint(op.getL3Hint());
       VectorType srcFlatVecTy =
           VectorType::get(srcVecTy.getNumElements(), srcVecTy.getElementType());
-      Value srcFlatVec = rewriter.create<vector::ShapeCastOp>(loc, srcFlatVecTy,
-                                                              op.getValue());
+      Value srcFlatVec = op.getValue();
       srcFlatVecTy = encodeVectorTypeTo(srcFlatVecTy,
                                         rewriter.getIntegerType(elemBitSize));
       srcFlatVec =
@@ -327,9 +319,7 @@ class LoadStorePrefetchNdToXeVMPattern : public OpConversionPattern<OpType> {
         resultFlatVec = rewriter.create<vector::BitCastOp>(
             loc, encodeVectorTypeTo(loadedTy, dstVecTy.getElementType()),
             resultFlatVec);
-        auto newOp =
-            rewriter.create<vector::ShapeCastOp>(loc, dstVecTy, resultFlatVec);
-        rewriter.replaceOp(op, newOp);
+        rewriter.replaceOp(op, resultFlatVec);
       }
     }
     return success();
@@ -548,14 +538,8 @@ class DpasToXeVMPattern : public OpConversionPattern<xegpu::DpasOp> {
     }
     auto rc = IntegerAttr::get(rewriter.getI32Type(), 8);
 
-    VectorType aNty =
-        VectorType::get(aTy.getNumElements(), aTy.getElementType());
-    Value aVec = rewriter.create<vector::ShapeCastOp>(loc, aNty, op.getLhs());
-
-    VectorType bNty =
-        VectorType::get(bTy.getNumElements(), bTy.getElementType());
-    Value bVec = rewriter.create<vector::ShapeCastOp>(loc, bNty, op.getRhs());
-
+    Value aVec = op.getLhs();
+    Value bVec = op.getRhs();
     auto cvecty = cast<VectorType>(c.getType());
     VectorType cNty =
         VectorType::get(cvecty.getNumElements(), cvecty.getElementType());

test/Conversion/XeGPUToXeVM/dpas.mlir

@@ -5,16 +5,11 @@
 #sg_map_c_f32 = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>
 
 gpu.module @load_store_check {
-    func.func @dpas(%a_loaded: vector<8x1xf16>, %b_loaded: vector<8x2xf16>, %c_loaded: vector<8x1xf32>) -> vector<8x1xf32> {
+    //CHECK: func.func @dpas(%[[arg0:.*]]: vector<8xf16>, %[[arg1:.*]]: vector<16xf16>, %[[arg2:.*]]: vector<8xf32>) -> vector<8xf32>
+    func.func @dpas(%a_loaded: vector<8xf16>, %b_loaded: vector<16xf16>, %c_loaded: vector<8xf32>) -> vector<8xf32> {
         // Loads are checked in a separate test.
-        // Cast arguments to SIMT-style vectors.
-        //CHECK: %[[CAST_A:.*]] = vector.shape_cast %arg0 : vector<8x1xf16> to vector<8xf16>
-        //CHECK-NEXT: %[[CAST_B:.*]] = vector.shape_cast %arg1 : vector<8x2xf16> to vector<16xf16>
-        //CHECK-NEXT: %[[CAST_C:.*]] = vector.shape_cast %arg2 : vector<8x1xf32> to vector<8xf32>
-        //CHECK-NEXT: %[[D:.*]] = xevm.dpas %[[CAST_C]], %[[CAST_A]], %[[CAST_B]] {pa = f16, pb = f16, rc = 8} : (vector<8xf32>, vector<8xf16>, vector<16xf16>) -> vector<8xf32>
-        // Cast result back to expected shape
-        //CHECK-NEXT: %[[CAST_D:.*]] = vector.shape_cast %[[D]] : vector<8xf32> to vector<8x1xf32>
-        %d = xegpu.dpas %a_loaded, %b_loaded, %c_loaded {a_layout = #sg_map_a_f16, b_layout = #sg_map_b_f16, c_layout = #sg_map_c_f32} : vector<8x1xf16>, vector<8x2xf16>, vector<8x1xf32> -> vector<8x1xf32>
-        return %d : vector<8x1xf32>
+        //CHECK: %[[D:.*]] = xevm.dpas %[[arg2]], %[[arg0]], %[[arg1]] {pa = f16, pb = f16, rc = 8} : (vector<8xf32>, vector<8xf16>, vector<16xf16>) -> vector<8xf32>
+        %d = xegpu.dpas %a_loaded, %b_loaded, %c_loaded {a_layout = #sg_map_a_f16, b_layout = #sg_map_b_f16, c_layout = #sg_map_c_f32} : vector<8xf16>, vector<16xf16>, vector<8xf32> -> vector<8xf32>
+        return %d : vector<8xf32>
     }
 }

test/Conversion/XeGPUToXeVM/lit.local.cfg

@@ -13,7 +13,8 @@ gpu.module @load_store_check {
         // CHECK: %[[LD_DESC_3:.*]] = vector.insert {{.*}}, %[[LD_DESC_2]] [3] : i32 into vector<8xi32>
         // CHECK: %[[LD_DESC_4:.*]] = vector.insert {{.*}}, %[[LD_DESC_3]] [4] : i32 into vector<8xi32>
         // CHECK: %[[LD_DESC:.*]] = vector.insert {{.*}}, %[[LD_DESC_4]] [5] : i32 into vector<8xi32>
-        %src_tdesc = xegpu.create_nd_tdesc %srcce[0, 0] : memref<8x16xf32> -> !xegpu.tensor_desc<8x16xf32, #xegpu.block_tdesc_attr<memory_space = global>, #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>>
+        %src_tdesc = xegpu.create_nd_tdesc %srcce[0, 0] : memref<8x16xf32> -> !xegpu.tensor_desc<8x16xf32>
+
 
         //CHECK: %[[LD_DESC_I64:.*]] = vector.bitcast %[[LD_DESC]] : vector<8xi32> to vector<4xi64>
         //CHECK: %[[LD_INTPTR:.*]] = vector.extract %[[LD_DESC_I64]][0] : i64 from vector<4xi64>
@@ -25,15 +26,14 @@ gpu.module @load_store_check {
         //CHECK: %[[LD_SIZEOF_F32:.*]] = arith.constant 4 : i32
         //CHECK: %[[LD_BASE_ROW_IN_BYTES:.*]] = arith.muli %[[LD_BASE_W]], %[[LD_SIZEOF_F32]] : i32
         //CHECK: %[[LD_LOADED_I32:.*]] = xevm.blockload2d %[[LD_LLVMPTR]], %[[LD_BASE_ROW_IN_BYTES]], %[[LD_BASE_H]], %[[LD_BASE_ROW_IN_BYTES]], %[[LD_TILE_W]], %[[LD_TILE_H]] {elem_size_in_bits = 32, tile_width = 16, tile_height = 8, v_blocks = 1, transpose = false, vnni_transform = false, l1_cache_control = C, l3_cache_control = UC} : (!llvm.ptr<1>, i32, i32, i32, i32, i32) -> vector<8xi32>
-        %loaded = xegpu.load_nd %src_tdesc <{l1_hint = #xegpu.cache_hint<cached>, l2_hint = #xegpu.cache_hint<uncached>}> : !xegpu.tensor_desc<8x16xf32, #xegpu.block_tdesc_attr<memory_space = global>, #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>> -> vector<8x1xf32>
+        %loaded = xegpu.load_nd %src_tdesc <{l1_hint = #xegpu.cache_hint<cached>, l2_hint = #xegpu.cache_hint<uncached>}> : !xegpu.tensor_desc<8x16xf32> -> vector<8xf32>
         //CHECK: %[[LD_LOADED_F32:.*]] = vector.bitcast %[[LD_LOADED_I32]] : vector<8xi32> to vector<8xf32>
-        //CHECK: %[[LD_LOADED_F32_DISTRIBUTED:.*]] = vector.shape_cast %[[LD_LOADED_F32]] : vector<8xf32> to vector<8x1xf32>
 
         %tid_x = gpu.thread_id x
         %tid_x_i32 = arith.index_cast %tid_x : index to i32
         %tid_x_f32 = arith.sitofp %tid_x_i32 : i32 to f32
-        //CHECK: %[[LOADED_F32_DISTRIBUTED_MODIFIED:.*]] = vector.insert %{{.*}}, %[[LD_LOADED_F32_DISTRIBUTED]] [0, 0] : f32 into vector<8x1xf32>
-        %loaded_modified = vector.insert %tid_x_f32, %loaded[0, 0] : f32 into vector<8x1xf32>
+        //CHECK: %[[LOADED_F32_MODIFIED:.*]] = vector.insert %{{.*}}, %[[LD_LOADED_F32]] [0] : f32 into vector<8xf32>
+        %loaded_modified = vector.insert %tid_x_f32, %loaded[0] : f32 into vector<8xf32>
 
         // CHECK: %[[PTR_AS_I64:.*]] = arith.index_castui {{.*}} : index to i64
         // CHECK: %[[CREATE_DESC_I64:.*]] = vector.bitcast {{.*}} : vector<8xi32> to vector<4xi64>
@@ -43,7 +43,7 @@ gpu.module @load_store_check {
         // CHECK: %[[DESC_3:.*]] = vector.insert {{.*}}, %[[DESC_2]] [3] : i32 into vector<8xi32>
         // CHECK: %[[DESC_4:.*]] = vector.insert {{.*}}, %[[DESC_3]] [4] : i32 into vector<8xi32>
         // CHECK: %[[DESC:.*]] = vector.insert {{.*}}, %[[DESC_4]] [5] : i32 into vector<8xi32>
-        %dst_tdesc = xegpu.create_nd_tdesc %dstte[0, 0] : memref<8x16xf32> -> !xegpu.tensor_desc<8x16xf32, #xegpu.block_tdesc_attr<memory_space = global>, #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>>
+        %dst_tdesc = xegpu.create_nd_tdesc %dstte[0, 0] : memref<8x16xf32> -> !xegpu.tensor_desc<8x16xf32, #xegpu.block_tdesc_attr<memory_space = global>>
 
         //CHECK: %[[DESC_I64:.*]] = vector.bitcast %[[DESC]] : vector<8xi32> to vector<4xi64>
         //CHECK: %[[INTPTR:.*]] = vector.extract %[[DESC_I64]][0] : i64 from vector<4xi64>
@@ -54,10 +54,9 @@ gpu.module @load_store_check {
         //CHECK: %[[LLVMPTR:.*]] = llvm.inttoptr %[[INTPTR]] : i64 to !llvm.ptr<1>
         //CHECK: %[[SIZEOF_F32:.*]] = arith.constant 4 : i32
         //CHECK: %[[BASE_ROW_IN_BYTES:.*]] = arith.muli %[[BASE_W]], %[[SIZEOF_F32]] : i32
-        //CHECK: %[[FLAT_VALUE:.*]] = vector.shape_cast %[[LOADED_F32_DISTRIBUTED_MODIFIED]] : vector<8x1xf32> to vector<8xf32>
-        //CHECK: %[[FLAT_VALUE_I32:.*]] = vector.bitcast %[[FLAT_VALUE]] : vector<8xf32> to vector<8xi32>
+        //CHECK: %[[FLAT_VALUE_I32:.*]] = vector.bitcast %[[LOADED_F32_MODIFIED]] : vector<8xf32> to vector<8xi32>
         //CHECK: xevm.blockstore2d %[[LLVMPTR]], %[[BASE_ROW_IN_BYTES]], %[[BASE_H]], %[[BASE_ROW_IN_BYTES]], %[[TILE_W]], %[[TILE_H]], %[[FLAT_VALUE_I32]] {elem_size_in_bits = 32, tile_width = 16, tile_height = 8, v_blocks = 1, l1_cache_control = WB, l3_cache_control = UC} : (!llvm.ptr<1>, i32, i32, i32, i32, i32, vector<8xi32>)
-        xegpu.store_nd %loaded_modified, %dst_tdesc <{l1_hint = #xegpu.cache_hint<write_back>, l2_hint = #xegpu.cache_hint<uncached>}>: vector<8x1xf32>, !xegpu.tensor_desc<8x16xf32, #xegpu.block_tdesc_attr<memory_space = global>, #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>>
+        xegpu.store_nd %loaded_modified, %dst_tdesc <{l1_hint = #xegpu.cache_hint<write_back>, l2_hint = #xegpu.cache_hint<uncached>}>: vector<8xf32>, !xegpu.tensor_desc<8x16xf32, #xegpu.block_tdesc_attr<memory_space = global>>
         gpu.return
     }
 }

test/Conversion/XeGPUToXeVM/loadstore_nd.mlir

@@ -2,27 +2,25 @@
 // RUN:                                       --runner imex-cpu-runner -e main \
 // RUN:                                       --entry-point-result=void \
 // RUN:                                       --shared-libs=%irunner_utils,%mlir_runner_utils,%mlir_c_runner_utils,%levelzero_runtime --filecheck
-
-#sg_map_a_bf16 = #xegpu.layout<lane_layout = [16, 1], lane_data = [1, 1]>
 module @gemm attributes {gpu.container_module} {
   gpu.module @kernel {
     gpu.func @load_store_2d(%src: memref<128xf32, 1>, %dst: memref<128xf32, 1>) kernel {
         %srcce = memref.memory_space_cast %src : memref<128xf32, 1> to memref<128xf32>
         %dstte = memref.memory_space_cast %dst : memref<128xf32, 1> to memref<128xf32>
 
         %offsets = arith.constant dense<[0,2,4,6,8,10,12,14,16,18,20,22,24,26,28,30]> : vector<16xindex>
-        %src_tdesc = xegpu.create_tdesc %srcce, %offsets : memref<128xf32>, vector<16xindex> -> !xegpu.tensor_desc<16x2xf32, #xegpu.scatter_tdesc_attr<chunk_size = 2>, #sg_map_a_bf16>
-        %dst_tdesc = xegpu.create_tdesc %dstte, %offsets : memref<128xf32>, vector<16xindex> -> !xegpu.tensor_desc<16x2xf32, #xegpu.scatter_tdesc_attr<chunk_size = 2>, #sg_map_a_bf16>
+        %src_tdesc = xegpu.create_tdesc %srcce, %offsets : memref<128xf32>, vector<16xindex> -> !xegpu.tensor_desc<16x2xf32, #xegpu.scatter_tdesc_attr<chunk_size = 2>>
+        %dst_tdesc = xegpu.create_tdesc %dstte, %offsets : memref<128xf32>, vector<16xindex> -> !xegpu.tensor_desc<16x2xf32, #xegpu.scatter_tdesc_attr<chunk_size = 2>>
 
         %mask = arith.constant dense<[1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1]> : vector<16xi1>
-        %loaded = xegpu.load %src_tdesc, %mask <{l1_hint = #xegpu.cache_hint<cached>, l2_hint = #xegpu.cache_hint<uncached>, transpose}> : !xegpu.tensor_desc<16x2xf32, #xegpu.scatter_tdesc_attr<chunk_size = 2>, #sg_map_a_bf16>, vector<16xi1> -> vector<2x1xf32>
+        %loaded = xegpu.load %src_tdesc, %mask <{l1_hint = #xegpu.cache_hint<cached>, l2_hint = #xegpu.cache_hint<uncached>}> : !xegpu.tensor_desc<16x2xf32, #xegpu.scatter_tdesc_attr<chunk_size = 2>>, vector<16xi1> -> vector<2xf32>
 
         %tid_x = gpu.thread_id x
         %tid_x_i32 = arith.index_cast %tid_x : index to i32
         %tid_x_f32 = arith.sitofp %tid_x_i32 : i32 to f32
-        %loaded_modified = vector.insert %tid_x_f32, %loaded[0,0] : f32 into vector<2x1xf32>
+        %loaded_modified = vector.insert %tid_x_f32, %loaded[0] : f32 into vector<2xf32>
 
-        xegpu.store %loaded_modified, %dst_tdesc, %mask <{l1_hint = #xegpu.cache_hint<write_back>, l2_hint = #xegpu.cache_hint<uncached>, transpose}> : vector<2x1xf32>, !xegpu.tensor_desc<16x2xf32, #xegpu.scatter_tdesc_attr<chunk_size = 2>, #sg_map_a_bf16>, vector<16xi1>
+        xegpu.store %loaded_modified, %dst_tdesc, %mask <{l1_hint = #xegpu.cache_hint<write_back>, l2_hint = #xegpu.cache_hint<uncached>}> : vector<2xf32>, !xegpu.tensor_desc<16x2xf32, #xegpu.scatter_tdesc_attr<chunk_size = 2>>, vector<16xi1>
         gpu.return
     }
   }