Order support in xetile.init_tile/load_tile (#707)

order support

Author: Hardcode84

lib/Conversion/XeTileToXeGPU/XeTileOpConversion.cpp

@@ -294,16 +294,23 @@ class SgInitTileOpPattern
     auto loc = op.getLoc();
     auto source = op.getSource();
     auto tileTy = op.getType();
-    auto innerBlk = tileTy.getInnerBlocks();
-    auto shape = tileTy.getShape();
+    auto innerBlocks = tileTy.getInnerBlocks();
+    auto shape = llvm::to_vector(tileTy.getShape());
     auto indexType = rewriter.getIndexType();
 
     if (tileTy.getRank() != 2)
       return op.emitOpError("The tile shape should be 2D.");
 
-    if (!innerBlk || innerBlk.size() != 2)
+    if (!innerBlocks || innerBlocks.size() != 2)
       return op.emitOpError("Missing valid innerBlock for the tile in op.");
 
+    // Need to make a copy, so we can swap values.
+    auto innerBlk = llvm::to_vector(innerBlocks.asArrayRef());
+    if (tileTy.getOrder().asArrayRef() == mlir::ArrayRef({0, 1})) {
+      std::swap(innerBlk[0], innerBlk[1]);
+      std::swap(shape[0], shape[1]);
+    }
+
     // using array_length for load if dim1 of innerBlocks
     // is smaller than dim 1 of shape.
     auto array_length =
@@ -330,6 +337,7 @@ class SgInitTileOpPattern
 
     auto offsetsX = offsets[0];
     auto offsetsY = offsets[1];
+
     auto tDescTy = xegpu::TensorDescType::get(
         innerBlk, tileTy.getElementType(), xegpu::MemoryScope::GLOBAL,
         array_length, true /*boundary_check*/, {} /*scattered*/,
@@ -450,8 +458,26 @@ struct SgLoadTileOpPattern
       vnniAttr = rewriter.getI32IntegerAttr(axis);
     }
 
-    // TODO: add transpose info
     mlir::DenseI64ArrayAttr transposeAttr;
+    auto srcOrder = tileTy.getOrder();
+    if (srcOrder.asArrayRef() == mlir::ArrayRef({1, 0})) {
+      // Nothing to do
+    } else if (srcOrder.asArrayRef() == mlir::ArrayRef({0, 1})) {
+      auto elemWidth = elemTy.getIntOrFloatBitWidth();
+      if (elemWidth == 32) {
+        transposeAttr = rewriter.getDenseI64ArrayAttr({1, 0});
+      } else if (elemWidth == 16 && vnniAttr && vnniAttr.getInt() == 0) {
+        transposeAttr = rewriter.getDenseI64ArrayAttr({1, 0});
+        transposeBitWidthAttr = rewriter.getI32IntegerAttr(32);
+        vnniAttr = nullptr;
+      } else {
+        return ((mlir::PatternRewriter &)rewriter)
+            .notifyMatchFailure(op, "Unsupported element type for transpose");
+      }
+    } else {
+      return ((mlir::PatternRewriter &)rewriter)
+          .notifyMatchFailure(op, "Unsupported order");
+    }
 
     rewriter.setInsertionPoint(op);
     llvm::SmallVector<::mlir::Value> xegpuOps;
@@ -461,10 +487,17 @@ struct SgLoadTileOpPattern
       auto shape = tdescTy.getShape().vec();
       auto array_length = tdescTy.getArrayLength();
 
+      if (transposeAttr)
+        std::swap(shape[0], shape[1]);
+
       if (vnniAttr) {
         auto axis = vnniAttr.getInt();
         shape[axis] /= factor;
         shape.push_back(factor);
+      } else if (transposeBitWidthAttr) {
+        auto axis = 0;
+        shape[axis] /= factor;
+        shape.push_back(factor);
       }
 
       if (array_length != 1)

lib/Dialect/XeTile/Transforms/Blocking.cpp

@@ -483,12 +483,14 @@ struct InitTileOpPattern : public XeTileConversion<xetile::InitTileOp> {
 
         int factor = 32 / elementSize;
         vnni = false;
-        innerBlocks = mlir::DenseI64ArrayAttr::get(
-            getContext(),
-            getInnerBlockSizes<Load>(
-                op.getOperation(), mlir::FloatType::getF32(getContext()),
-                tileTy.getShape()[0], (tileTy.getShape()[1]) * factor,
-                this->uArchInterface, vnni, transpose));
+        llvm::SmallVector<int64_t, 2> innerBlock = getInnerBlockSizes<Load>(
+            op.getOperation(), mlir::FloatType::getF32(getContext()),
+            tileTy.getShape()[1], (tileTy.getShape()[0]) / factor,
+            this->uArchInterface, vnni, transpose);
+        std::swap(innerBlock[0], innerBlock[1]);
+        innerBlock[0] *= factor;
+        innerBlocks = mlir::DenseI64ArrayAttr::get(getContext(), innerBlock);
+
       } else if (transpose && elementSize < 32) {
         return rewriter.notifyMatchFailure(op, "Invalid transpose.");
       } else {

lib/Utils/XeArch.cpp

@@ -95,12 +95,14 @@ XePVCuArch::get2DLoadConfig(mlir::Operation *op, int element_data_size,
            << "transpose and transform are not supported together";
   }
 
+  // FIXME: We do support transpose on f16 wtih transpose_bit_width==32,
+  // disable check for now.
   // only d32 and d64 is supported for transpose operations
-  if ((transpose) && (element_data_size != 32 && element_data_size != 64)) {
-    return op->emitOpError()
-           << "transposed load only supports d32 and d64 data sizes. "
-           << "Given element data size: d" << element_data_size;
-  }
+  // if ((transpose) && (element_data_size != 32 && element_data_size != 64)) {
+  //   return op->emitOpError()
+  //          << "transposed load only supports d32 and d64 data sizes. "
+  //          << "Given element data size: d" << element_data_size;
+  // }
 
   // only d8 and d16 are suported for VNNI transform operations
   if ((vnni) && (element_data_size != 8 && element_data_size != 16)) {

test/Conversion/XeTileToXeGPU/test_order.mlir

@@ -0,0 +1,38 @@
+// RUN: imex-opt --split-input-file --xetile-blocking --convert-xetile-to-xegpu --cse %s -verify-diagnostics -o -| FileCheck %s
+
+// CHECK-LABEL: @test_func
+// CHECK-SAME: (%[[A:.*]]: memref<128x64xf16>, %[[B:.*]]: memref<64x128xf16, strided<[1, 64]>>)
+// CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index
+// CHECK-DAG: %[[C16:.*]] = arith.constant 16 : index
+// CHECK: %[[D1:.*]] = xegpu.create_nd_tdesc %[[A]][%[[C0]], %[[C0]]] {mode = vc} : memref<128x64xf16> -> !xegpu.tensor_desc<32x16xf16>
+// CHECK: %[[D2:.*]] = xegpu.create_nd_tdesc %[[B]][%[[C0]], %[[C0]]] {mode = vc} : memref<64x128xf16, strided<[1, 64]>> -> !xegpu.tensor_desc<32x16xf16>
+// CHECK: %{{.*}} = xegpu.load_nd %[[D1]] {mode = vc, vnni_axis = 1, l1_hint = cached, l2_hint = cached, l3_hint = cached} : !xegpu.tensor_desc<32x16xf16> -> vector<32x8x2xf16>
+// CHECK: %{{.*}} = xegpu.load_nd %[[D2]] {mode = vc, transpose = [1, 0], transpose_bit_width = 32, l1_hint = cached, l2_hint = cached, l3_hint = cached} : !xegpu.tensor_desc<32x16xf16> -> vector<8x32x2xf16>
+// CHECK: %[[D3:.*]] = xegpu.update_nd_offset %[[D1]], [%[[C0]], %[[C16]]] {mode = vc} : !xegpu.tensor_desc<32x16xf16> -> !xegpu.tensor_desc<32x16xf16>
+// CHECK: %[[D4:.*]] = xegpu.update_nd_offset %[[D2]], [%[[C16]], %[[C0]]] {mode = vc} : !xegpu.tensor_desc<32x16xf16> -> !xegpu.tensor_desc<32x16xf16>
+// CHECK: %{{.*}} = xegpu.load_nd %[[D3]] {mode = vc, vnni_axis = 1, l1_hint = cached, l2_hint = cached, l3_hint = cached} : !xegpu.tensor_desc<32x16xf16> -> vector<32x8x2xf16>
+// CHECK: %{{.*}} = xegpu.load_nd %[[D4]] {mode = vc, transpose = [1, 0], transpose_bit_width = 32, l1_hint = cached, l2_hint = cached, l3_hint = cached} : !xegpu.tensor_desc<32x16xf16> -> vector<8x32x2xf16>
+gpu.module @test_kernel {
+func.func @test_func(%A : memref<128x64xf16>, %B : memref<64x128xf16, strided<[1, 64], offset: 0>>) {
+  %c0 = arith.constant 0 : index
+  %c32 = arith.constant 32 : index
+  %c16  = arith.constant 16 : index
+  %A_block_iter0 = xetile.init_tile %A[%c0, %c0] : memref<128x64xf16> -> !xetile.tile<32x16xf16>
+  %B_block_iter0 = xetile.init_tile %B[%c0, %c0] : memref<64x128xf16, strided<[1, 64], offset: 0>> -> !xetile.tile<16x32xf16, #xetile.tile_attr<order = [0, 1]>>
+
+  %A_block_value0 = xetile.load_tile %A_block_iter0 : !xetile.tile<32x16xf16> -> vector<32x16xf16>
+  %B_block_value0 = xetile.load_tile %B_block_iter0 : !xetile.tile<16x32xf16, #xetile.tile_attr<order = [0,1]>> -> vector<16x32xf16>
+
+  %mma_out0 = xetile.tile_mma %A_block_value0, %B_block_value0 : vector<32x16xf16>, vector<16x32xf16> -> vector<32x32xf32>
+
+  %A_block_iter1 = xetile.update_tile_offset %A_block_iter0, [%c0, %c16] : !xetile.tile<32x16xf16>, index, index -> !xetile.tile<32x16xf16>
+  %B_block_iter1 = xetile.update_tile_offset %B_block_iter0, [%c16, %c0] : !xetile.tile<16x32xf16, #xetile.tile_attr<order = [0,1]>>, index, index -> !xetile.tile<16x32xf16, #xetile.tile_attr<order = [0,1]>>
+
+  %A_block_value1 = xetile.load_tile %A_block_iter1 : !xetile.tile<32x16xf16> -> vector<32x16xf16>
+  %B_block_value1 = xetile.load_tile %B_block_iter1  : !xetile.tile<16x32xf16, #xetile.tile_attr<order = [0,1]>> -> vector<16x32xf16>
+
+  %mma_out1 = xetile.tile_mma %A_block_value1, %B_block_value1, %mma_out0 : vector<32x16xf16>, vector<16x32xf16>, vector<32x32xf32> -> vector<32x32xf32>
+
+  return
+}
+}

test/Dialect/XeTile/Transforms/blocking.mlir

@@ -152,7 +152,7 @@ gpu.module @test_kernel {
     //CHECK: gpu.func @tile_mma(%[[arg0:.*]]: memref<128x128xf16>, %[[arg1:.*]]: memref<128x128xf16>)
     //CHECK: %[[c0:.*]] = arith.constant 0 : index
     //CHECK: %[[R0:.*]] = xetile.init_tile %[[arg0]][%[[c0]], %[[c0]]] : memref<128x128xf16> -> !xetile.tile<90x76xf16, #xetile.tile_attr<inner_blocks = [30, 19]>>
-    //CHECK: %[[R1:.*]] = xetile.init_tile %[[arg1]][%[[c0]], %[[c0]]] : memref<128x128xf16> -> !xetile.tile<76x90xf16, #xetile.tile_attr<order = [0, 1], inner_blocks = [19, 6]>>
+    //CHECK: %[[R1:.*]] = xetile.init_tile %[[arg1]][%[[c0]], %[[c0]]] : memref<128x128xf16> -> !xetile.tile<76x90xf16, #xetile.tile_attr<order = [0, 1], inner_blocks = [4, 30]>>
     gpu.func @tile_mma(%a: memref<128x128xf16>, %b: memref<128x128xf16>) {
       %c0 = arith.constant 0 : index
   	  %1 = xetile.init_tile %a[%c0, %c0] : memref<128x128xf16> -> !xetile.tile<90x76xf16>