remove vector direction and lenght attirbutes

Author: Jianhui-Li

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

@@ -724,22 +724,4 @@ def XeGPU_MemLayoutAttr : XeGPUAttr<"MemLayout", "mem_layout"> {
 
 }
 
-def RowOriented : I32EnumAttrCase<"ROW", 0, "row">;
-def ColOriented : I32EnumAttrCase<"COL", 1, "col">;
-def MatrixAccessDirection : 
-    I32EnumAttr<"MatrixAccessDirection", 
-    "Matrix elements/vectors can have row or column direction", [
-    RowOriented, ColOriented
-]> {
-  let genSpecializedAttr = 0;
-  let cppNamespace = "::mlir::xegpu";
-}
-def MatrixAccessDirectionAttr : 
-    EnumAttr<XeGPU_Dialect, 
-    MatrixAccessDirection, 
-    "matrix_access_direction">{
-  let summary = [{Describe the direction of memory access for load_matrix and store_matrix.}];
-  let assemblyFormat = "`<` $value `>`";
-}
-
 #endif // MLIR_DIALECT_XEGPU_IR_XEGPUATTRS_TD

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

@@ -1302,8 +1302,6 @@ def XeGPU_LoadMatrixOp: XeGPU_Op<"load_matrix", [MemoryEffects<[MemRead]>,
   let arguments = (ins XeGPU_MemDesc:$mem_desc,
     Variadic<Index>: $offsets,
     DenseI64ArrayAttr: $const_offsets,
-    OptionalAttr<I32Attr>:$vec_length,
-    OptionalAttr<MatrixAccessDirectionAttr>:$vec_direction,
     OptionalAttr<UnitAttr>:$subgroup_block_io,
     OptionalAttr<DistributeLayoutAttr>:$layout
   );
@@ -1355,8 +1353,6 @@ def XeGPU_StoreMatrixOp: XeGPU_Op<"store_matrix", [MemoryEffects<[MemWrite]>,
     XeGPU_MemDesc:$mem_desc,
     Variadic<Index>: $offsets,
     DenseI64ArrayAttr: $const_offsets,
-    OptionalAttr<I32Attr>:$vec_length,
-    OptionalAttr<MatrixAccessDirectionAttr>:$vec_direction,
     OptionalAttr<UnitAttr>:$subgroup_block_io,
     OptionalAttr<DistributeLayoutAttr>:$layout
   );

mlir/lib/Conversion/XeGPUToXeVM/XeGPUToXeVM.cpp

@@ -184,9 +184,9 @@ class CreateNdDescToXeVMPattern
     SmallVector<OpFoldResult> mixedSizes = op.getMixedSizes();
     // Descriptor shape is expected to be 2D.
     int64_t rank = mixedSizes.size();
-    if (rank != 2) {
+    if (rank != 2)
       return rewriter.notifyMatchFailure(op, "Expected 2D shape.");
-    }
+
     auto sourceTy = source.getType();
     auto sourceMemrefTy = dyn_cast<MemRefType>(sourceTy);
     // If source is a memref, we need to extract the aligned pointer as index.
@@ -658,20 +658,6 @@ class LoadStoreMatrixToXeVMPattern : public OpConversionPattern<OpType> {
           return rewriter.notifyMatchFailure(
               op, "The lowering is specific to pvc or bmg.");
         }
-        xegpu::MatrixAccessDirectionAttr vecDirection =
-            op.getVecDirectionAttr();
-        if (vecDirection &&
-            vecDirection.getValue() == xegpu::MatrixAccessDirection::COL &&
-            !mdescTy.isColMajor())
-          return rewriter.notifyMatchFailure(
-              op, "mem_desc should be column major when "
-                  "vec_direction is COLUMN for 1D result.");
-        if (vecDirection &&
-            vecDirection.getValue() == xegpu::MatrixAccessDirection::ROW &&
-            mdescTy.isColMajor())
-          return rewriter.notifyMatchFailure(
-              op, "mem_desc should be row major when "
-                  "vec_direction is ROW for 1D result.");
 
         if constexpr (std::is_same_v<OpType, xegpu::LoadMatrixOp>) {
           Value loadOp =

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

@@ -173,17 +173,18 @@ isValidGatherScatterBufferParams(Type offsetsTy, Type maskTy,
   return success();
 }
 
-LogicalResult IsValidStoreMatrixParams(
-    VectorType dataTy, MemDescType mdescTy, UnitAttr subgroup_block_io,
-    MatrixAccessDirectionAttr vecDirection, IntegerAttr vecLength,
-    function_ref<InFlightDiagnostic()> emitError) {
-
-  if (!dataTy)
-    if (subgroup_block_io || vecDirection || vecLength)
-      return emitError() << "vec_length, vec_direction and subgroup_block_io "
+LogicalResult
+IsValidStoreMatrixParams(VectorType dataTy, MemDescType mdescTy,
+                         UnitAttr subgroup_block_io,
+                         function_ref<InFlightDiagnostic()> emitError) {
+
+  if (!dataTy) {
+    if (subgroup_block_io)
+      return emitError() << "subgroup_block_io "
                             "are only allowed when result is a 1D VectorType.";
     else
       return success();
+  }
 
   if (mdescTy.getRank() != 2)
     return emitError() << "mem_desc must be 2D.";
@@ -192,8 +193,8 @@ LogicalResult IsValidStoreMatrixParams(
   ArrayRef<int64_t> mdescShape = mdescTy.getShape();
 
   if (dataShape.size() == 2) {
-    if (subgroup_block_io || vecDirection || vecLength)
-      return emitError() << "vec_length, vec_direction and subgroup_block_io "
+    if (subgroup_block_io)
+      return emitError() << "subgroup_block_io "
                             "are only allowed when result is a 1D VectorType.";
     if (llvm::any_of(llvm::zip_equal(dataShape, mdescShape),
                      [](auto p) { return std::get<0>(p) > std::get<1>(p); }))
@@ -1097,20 +1098,16 @@ void LoadMatrixOp::build(OpBuilder &builder, OperationState &state, Type res,
   // Call the generated builder with all parameters (including optional ones as
   // nullptr/empty)
   build(builder, state, res, memDesc, dynamicOffsets, staticOffsetsAttr,
-        /*vec_length=*/nullptr, /*vec_direction=*/nullptr,
         /*subgroup_block_io=*/nullptr, layout);
 }
 
 LogicalResult LoadMatrixOp::verify() {
 
   auto resTy = dyn_cast<VectorType>(getRes().getType());
   UnitAttr subgroup_block_io = getSubgroupBlockIoAttr();
-  MatrixAccessDirectionAttr vecDirection = getVecDirectionAttr();
-  IntegerAttr vecLength = getVecLengthAttr();
   MemDescType mdescTy = getMemDesc().getType();
 
   return IsValidStoreMatrixParams(resTy, mdescTy, subgroup_block_io,
-                                  vecDirection, vecLength,
                                   [&]() { return emitError(); });
 }
 
@@ -1126,19 +1123,15 @@ void StoreMatrixOp::build(OpBuilder &builder, OperationState &state, Value data,
   dispatchIndexOpFoldResults(offsets, dynamicOffsets, staticOffsets);
   auto staticOffsetsAttr = builder.getDenseI64ArrayAttr(staticOffsets);
   build(builder, state, data, memDesc, dynamicOffsets, staticOffsetsAttr,
-        /*vec_length=*/nullptr, /*vec_direction=*/nullptr,
         /*subgroup_block_io=*/nullptr, layout);
 }
 
 LogicalResult StoreMatrixOp::verify() {
 
   auto dataTy = dyn_cast<VectorType>(getData().getType());
   UnitAttr subgroup_block_io = getSubgroupBlockIoAttr();
-  MatrixAccessDirectionAttr vecDirection = getVecDirectionAttr();
-  IntegerAttr vecLength = getVecLengthAttr();
   MemDescType mdescTy = getMemDesc().getType();
   return IsValidStoreMatrixParams(dataTy, mdescTy, subgroup_block_io,
-                                  vecDirection, vecLength,
                                   [&]() { return emitError(); });
 }
 

mlir/test/Conversion/XeGPUToXeVM/loadstore_matrix.mlir

@@ -139,10 +139,10 @@ gpu.module @test_kernel [#xevm.target<chip = "pvc">] {
 
     %tid_x = gpu.thread_id x
     %c16 = arith.constant 16 : index
-    %1 = xegpu.load_matrix %0[%c16, %tid_x] {vec_length = 8 : i32, vec_direction = #xegpu.matrix_access_direction<col>}: !xegpu.mem_desc<32x64xf16, #xegpu.mem_layout<stride = [1, 32], block = [16, 16]>>, index, index -> vector<8xf16>
+    %1 = xegpu.load_matrix %0[%c16, %tid_x] : !xegpu.mem_desc<32x64xf16, #xegpu.mem_layout<stride = [1, 32], block = [16, 16]>>, index, index -> vector<8xf16>
 
     //CHECK: llvm.store %[[loaded]], {{.*}} : vector<8xf16>, !llvm.ptr<3>
-    xegpu.store_matrix %1, %0[%c16, %tid_x] {vec_length = 8 : i32, vec_direction = #xegpu.matrix_access_direction<col>}: vector<8xf16>, !xegpu.mem_desc<32x64xf16, #xegpu.mem_layout<stride = [1, 32], block = [16, 16]>>, index, index
+    xegpu.store_matrix %1, %0[%c16, %tid_x] : vector<8xf16>, !xegpu.mem_desc<32x64xf16, #xegpu.mem_layout<stride = [1, 32], block = [16, 16]>>, index, index
 
     gpu.return %1: vector<8xf16>
   }

mlir/test/Dialect/XeGPU/invalid.mlir

@@ -870,16 +870,9 @@ func.func @load_mem_desc_invalid_rank(%arg0: !xegpu.mem_desc<64xf16>) {
   return
 }
 
-// -----
-func.func @load_mem_desc_invalid_attr1(%arg0: !xegpu.mem_desc<16x64xf16>) {
-  // expected-error@+1 {{vec_length, vec_direction and subgroup_block_io are only allowed when result is a 1D VectorType.}}
-  %data1 = xegpu.load_matrix %arg0[8, 8]<{vec_length = 8 : i32, vec_direction = #xegpu.matrix_access_direction<col>}>: !xegpu.mem_desc<16x64xf16> -> vector<16x16xf16>
-  return
-}
-
 // -----
 func.func @load_mem_desc_invalid_attr2(%arg0: !xegpu.mem_desc<16x64xf16>) {
-  // expected-error@+1 {{vec_length, vec_direction and subgroup_block_io are only allowed when result is a 1D VectorType.}}
+  // expected-error@+1 {{subgroup_block_io are only allowed when result is a 1D VectorType.}}
   %data2 = xegpu.load_matrix %arg0[8, 8] <{subgroup_block_io}>: !xegpu.mem_desc<16x64xf16> -> vector<16x16xf16>
   return
 }
@@ -908,14 +901,14 @@ func.func @store_mem_desc_invalid_rank(%arg0: !xegpu.mem_desc<64xf16>, %arg1: ve
 
 // -----
 func.func @store_mem_desc_invalid_attr2(%arg0: !xegpu.mem_desc<16x64xf16>, %data: vector<16x16xf16>) {
-  // expected-error@+1 {{vec_length, vec_direction and subgroup_block_io are only allowed when result is a 1D VectorType.}}
+  // expected-error@+1 {{subgroup_block_io are only allowed when result is a 1D VectorType.}}
   xegpu.store_matrix %data,  %arg0[8, 8] <{subgroup_block_io}>: vector<16x16xf16>, !xegpu.mem_desc<16x64xf16>
   return
 }
 
 // -----
 func.func @store_mem_desc_invalid_attr2(%arg0: !xegpu.mem_desc<16x64xf16>, %data: vector<16x16xf16>) {
-  // expected-error@+1 {{vec_length, vec_direction and subgroup_block_io are only allowed when result is a 1D VectorType.}}
+  // expected-error@+1 {{subgroup_block_io are only allowed when result is a 1D VectorType.}}
   xegpu.store_matrix %data,  %arg0[8, 8] <{subgroup_block_io}>: vector<16x16xf16>, !xegpu.mem_desc<16x64xf16>
   return
 }

mlir/test/Dialect/XeGPU/ops.mlir

@@ -855,8 +855,8 @@ gpu.func @simt_load_matrix_subgroup_block_io(%arg0: !xegpu.mem_desc<16x64xf16, #
 
 // CHECK: gpu.func @simt_load_matrix_vector(%arg0: !xegpu.mem_desc<16x64xf16, #xegpu.mem_layout<stride = [1, 16]>>)
 gpu.func @simt_load_matrix_vector(%arg0: !xegpu.mem_desc<16x64xf16, #xegpu.mem_layout<stride = [1, 16]>>) {
-  // CHECK: xegpu.load_matrix [[ARG0]][8, 8] <{vec_direction = #xegpu.matrix_access_direction<col>, vec_length = 8 : i32}>: !xegpu.mem_desc<16x64xf16, #xegpu.mem_layout<stride = [1, 16]>> -> vector<8xf16> 
-  %data = xegpu.load_matrix %arg0[8, 8] <{vec_direction = #xegpu.matrix_access_direction<col>, vec_length = 8 : i32}>: !xegpu.mem_desc<16x64xf16, #xegpu.mem_layout<stride = [1, 16]>> -> vector<8xf16>
+  // CHECK: xegpu.load_matrix [[ARG0]][8, 8] : !xegpu.mem_desc<16x64xf16, #xegpu.mem_layout<stride = [1, 16]>> -> vector<8xf16> 
+  %data = xegpu.load_matrix %arg0[8, 8] : !xegpu.mem_desc<16x64xf16, #xegpu.mem_layout<stride = [1, 16]>> -> vector<8xf16>
   gpu.return
 }
 
@@ -890,8 +890,8 @@ gpu.func @simt_store_matrix_subgroup_block_io(%arg0: !xegpu.mem_desc<16x64xf16,
 
 // CHECK: gpu.func @simt_store_matrix_vector(%arg0: !xegpu.mem_desc<16x64xf16, #xegpu.mem_layout<stride = [1, 16]>>, %arg1: vector<8xf16>) {
 gpu.func @simt_store_matrix_vector(%arg0: !xegpu.mem_desc<16x64xf16, #xegpu.mem_layout<stride = [1, 16]>>, %arg1: vector<8xf16>) {
-  // CHECK: xegpu.store_matrix [[ARG1]], [[ARG0]][8, 8] <{vec_direction = #xegpu.matrix_access_direction<col>, vec_length = 8 : i32}>: vector<8xf16>, !xegpu.mem_desc<16x64xf16, #xegpu.mem_layout<stride = [1, 16]>> 
-  xegpu.store_matrix %arg1, %arg0[8, 8] <{vec_length = 8 : i32, vec_direction = #xegpu.matrix_access_direction<col>}>: vector<8xf16>, !xegpu.mem_desc<16x64xf16, #xegpu.mem_layout<stride = [1, 16]>>
+  // CHECK: xegpu.store_matrix [[ARG1]], [[ARG0]][8, 8] : vector<8xf16>, !xegpu.mem_desc<16x64xf16, #xegpu.mem_layout<stride = [1, 16]>> 
+  xegpu.store_matrix %arg1, %arg0[8, 8] : vector<8xf16>, !xegpu.mem_desc<16x64xf16, #xegpu.mem_layout<stride = [1, 16]>>
   gpu.return
 }