Add transpose, reduce and broadcast definitions to xetile (#747)

Author: chencha3

include/imex/Dialect/XeTile/IR/CMakeLists.txt

@@ -2,8 +2,10 @@ add_mlir_dialect(XeTileOps xetile)
 add_mlir_doc(XeTileOps XeTileDialect Dialects/ -gen-dialect-doc -dialect=xetile)
 
 set(LLVM_TARGET_DEFINITIONS XeTileOps.td)
-mlir_tablegen(XeTileOpsAttrs.h.inc -gen-attrdef-decls)
-mlir_tablegen(XeTileOpsAttrs.cpp.inc -gen-attrdef-defs)
+mlir_tablegen(XeTileOpsAttrs.h.inc -gen-attrdef-decls --attrdefs-dialect=xetile)
+mlir_tablegen(XeTileOpsAttrs.cpp.inc -gen-attrdef-defs --attrdefs-dialect=xetile)
+
+set(LLVM_TARGET_DEFINITIONS XeTileAttrs.td)
 mlir_tablegen(XeTileOpsEnums.h.inc -gen-enum-decls)
 mlir_tablegen(XeTileOpsEnums.cpp.inc -gen-enum-defs)
 add_public_tablegen_target(XeTileOpsAttrsIncGen)

include/imex/Dialect/XeTile/IR/XeTileDialect.td

@@ -49,7 +49,8 @@ def XeTile_Dialect : Dialect {
     let cppNamespace = "::imex::xetile";
 
     let dependentDialects = [
-          "::mlir::memref::MemRefDialect"];
+          "::mlir::memref::MemRefDialect",
+          "::mlir::vector::VectorDialect"];
 
     // TODO: temporary disable it.
     let useDefaultTypePrinterParser = true;

include/imex/Dialect/XeTile/IR/XeTileOps.h

@@ -15,6 +15,7 @@
 #ifndef _XETILE_OPS_H_INCLUDED_
 #define _XETILE_OPS_H_INCLUDED_
 
+#include <mlir/Dialect/Vector/IR/VectorOps.h>
 #include <mlir/IR/BuiltinTypeInterfaces.h>
 #include <mlir/IR/BuiltinTypes.h>
 #include <mlir/IR/Dialect.h>

include/imex/Dialect/XeTile/IR/XeTileOps.td

@@ -18,6 +18,8 @@ include "imex/Dialect/XeTile/IR/XeTileDialect.td"
 include "imex/Dialect/XeTile/IR/XeTileTypes.td"
 include "imex/Dialect/XeTile/IR/XeTileAttrs.td"
 
+include "mlir/Dialect/Vector/IR/VectorAttributes.td"
+
 // Base class for dialect operations. This operation inherits from the base
 // `Op` class in OpBase.td, and provides:
 //   * The parent dialect of the operation.
@@ -519,5 +521,52 @@ def XeTile_AtomicRMWOp : XeTile_Op<"atomic_rmw", []> {
     }];
 }
 
+def XeTile_TransposeOp: XeTile_Op<"transpose", []> {
+    let summary = "transpose a 2D vector.";
+    let description = [{
+        It has the same semantic with `vector.transpose`, but limits the vector to be 2D.
+    }];
+
+    let arguments = (ins XeTile_2DVector: $source,
+                         DenseI64ArrayAttr:$permutation);
+    let results = (outs XeTile_2DVector: $result);
+    let assemblyFormat = [{
+        $source $permutation attr-dict `:` type($source) `->` type($result)
+    }];
+    let hasVerifier = 1;
+}
+
+def XeTile_ReduceOp: XeTile_Op<"reduce", []> {
+    let summary = "performs a reduction operation over a 2D vector.";
+    let description = [{
+        It has the same semantics as the `vector.multi_reduction`,
+        but restricts the vector dimension to 2D, and also the result
+        is 2D too, with the reduced axis being 1.
+    }];
+
+    let arguments = (ins Vector_CombiningKindAttr: $kind,
+                         XeTile_2DVector: $source,
+                         DenseI64ArrayAttr: $reduction_dim);
+    let results = (outs XeTile_2DVector: $result);
+    let assemblyFormat = [{
+        $kind `,` $source $reduction_dim attr-dict `:` type($source) `->` type($result)
+    }];
+
+    let hasVerifier = 1;
+}
+
+def XeTile_BroadCastOp: XeTile_Op<"broadcast", []> {
+    let summary = "broadcast a vector from 1D to 2D.";
+
+    let arguments = (ins XeTile_2DVector: $source,
+                         DenseI64ArrayAttr: $broadcast_dim);
+    let results = (outs XeTile_2DVector: $result);
+    let assemblyFormat = [{
+        $source $broadcast_dim attr-dict `:` type($source) `->` type($result)
+    }];
+    let hasVerifier = 1;
+}
+
+
 
 #endif // _XETILE_OPS_TD_INCLUDED_

lib/Dialect/XeTile/IR/XeTileOps.cpp

@@ -170,7 +170,7 @@ mlir::LogicalResult InitTileOp::verify() {
   }
 
   if (isSourceMemRef() && sourceMemRefHasStaticShape()) {
-    auto memrefType = getSourceType().dyn_cast<mlir::MemRefType>();
+    auto memrefType = mlir::dyn_cast<mlir::MemRefType>(getSourceType());
 
     // Checks for memrefs with format: memref<[shape], strided<[strides],
     // offsets:[offset]>>
@@ -905,6 +905,42 @@ mlir::OpFoldResult TileUnpackOp::fold(FoldAdaptor /*adaptor*/) {
   return nullptr;
 }
 
+mlir::LogicalResult TransposeOp::verify() {
+  auto srcShape = getSource().getType().getShape();
+  auto resShape = getResult().getType().getShape();
+  auto permutation = getPermutation();
+
+  auto transposeShape = srcShape.vec();
+  for (auto [i, j] : llvm::enumerate(permutation)) {
+    if (j >= (int64_t)srcShape.size())
+      return emitOpError("permutation index out of bounds");
+    transposeShape[i] = srcShape[j];
+  }
+
+  if (transposeShape != resShape.vec())
+    return emitOpError("Incorrect transpose permutation");
+
+  return mlir::success();
+}
+
+mlir::LogicalResult ReduceOp::verify() {
+  auto dims = getReductionDim();
+  auto resShape = getResult().getType().getShape();
+  for (auto i : dims)
+    if (resShape[i] != 1)
+      return emitOpError("reduction dimension of result must have size 1");
+  return mlir::success();
+}
+
+mlir::LogicalResult BroadCastOp::verify() {
+  auto dims = getBroadcastDim();
+  auto srcShape = getSource().getType().getShape();
+  for (auto i : dims)
+    if (srcShape[i] != 1)
+      return emitOpError("broadcast dimension of source must have size 1");
+  return mlir::success();
+}
+
 } // namespace xetile
 } // namespace imex
 

test/Dialect/XeTile/IR/invalid.mlir

@@ -228,3 +228,25 @@ func.func @tile_unpack_invalid_output_shape(%in : vector<4x4x16x16xf16>) {
 #wg_map_4 = #xetile.tile_attr<order = [0, 1, 2]>
 // expected-error@+1 {{expect integer array of size 2 for wg_data}}
 #wg_map_5 = #xetile.tile_attr<wg_data = [32, 64, 128]>
+
+
+// -----
+func.func @test_transpose(%source: vector<8x16xf16>) {
+  // expected-error@+1 {{Incorrect transpose permutation}}
+  %1 = xetile.transpose %source [0, 1] : vector<8x16xf16> -> vector<16x8xf16>
+  return
+}
+
+// -----
+func.func @test_reduce(%source: vector<8x16xf16>) {
+  // expected-error@+1 {{reduction dimension of result must have size 1}}
+  %1 = xetile.reduce <add>, %source [0] : vector<8x16xf16> -> vector<2x16xf16>
+  return
+}
+
+// -----
+func.func @test_broadcast(%source: vector<2x16xf16>) {
+  // expected-error@+1 {{broadcast dimension of source must have size 1}}
+  %1 = xetile.broadcast %source [0] : vector<2x16xf16> -> vector<8x16xf16>
+  return
+}

test/Dialect/XeTile/IR/ops.mlir

@@ -309,3 +309,21 @@ func.func @test_atomic_rmw(%tile : !xetile.tile<8x16xf16>, %value : vector<8x16x
   %1 = xetile.atomic_rmw "addf" %value, %tile : vector<8x16xf16>, !xetile.tile<8x16xf16> -> vector<8x16xf16>
   return
 }
+
+func.func @test_transpose(%source: vector<8x16xf16>) {
+  // CHECK: xetile.transpose {{.*}} [1, 0] : vector<8x16xf16> -> vector<16x8xf16>
+  %1 = xetile.transpose %source [1, 0] : vector<8x16xf16> -> vector<16x8xf16>
+  return
+}
+
+func.func @test_reduce(%source: vector<8x16xf16>) {
+  // CHECK: xetile.reduce {{.*}} [0] : vector<8x16xf16> -> vector<1x16xf16>
+  %1 = xetile.reduce <add>, %source [0] : vector<8x16xf16> -> vector<1x16xf16>
+  return
+}
+
+func.func @test_broadcast(%source: vector<1x16xf16>) {
+  // CHECK: xetile.broadcast {{.*}} [0] : vector<1x16xf16> -> vector<8x16xf16>
+  %1 = xetile.broadcast %source [0] : vector<1x16xf16> -> vector<8x16xf16>
+  return
+}