Add attributes required for work group level XeTile programming.

Following changes are introduced.
- WG and SG levels layout attributes for XeTile
- Custom printers/parsers for XeTile ops (needed by XeTileToXeGPU lowering)
- Additional verifications for tile_mma that were missing in initial version
- Updated test cases

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

@@ -1,2 +1,7 @@
 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)
+add_public_tablegen_target(XeTileOpsAttrsIncGen)

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

@@ -0,0 +1,96 @@
+//===------------ XeTileAttr.td - XeTile dialect  -------*- tablegen -*-===//
+//
+// Copyright 2022 Intel Corporation
+// Part of the IMEX Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+///
+/// \file
+/// This file defines custom attributes used by XeTile dialect.
+///
+//===----------------------------------------------------------------------===//
+
+#ifndef _XETILE_ATTR_DEF_TD_INCLUDED_
+#define _XETILE_ATTR_DEF_TD_INCLUDED_
+
+include "mlir/IR/AttrTypeBase.td"
+include "imex/Dialect/XeTile/IR/XeTileDialect.td"
+
+class XeTile_Attr<string name, string attrMnemonic, list<Trait> traits = []>
+  : AttrDef<XeTile_Dialect, name, traits> {
+    let mnemonic = attrMnemonic;
+}
+
+def XeTile_SubGroupMapAttr : XeTile_Attr<"SubGroupMap", "sg_map"> {
+  let parameters = (ins
+    OptionalParameter<"mlir::DenseI32ArrayAttr">:$mma_block_size,
+    "mlir::DenseI32ArrayAttr":$wi_layout,
+    "mlir::DenseI32ArrayAttr":$wi_data
+  );
+  let assemblyFormat = "`<` struct(params) `>`";
+  let genVerifyDecl = true;
+  let builders = [
+    AttrBuilder<(ins "llvm::ArrayRef<int32_t>":$mma_block_size,
+                      "llvm::ArrayRef<int32_t>":$wi_layout,
+                      "llvm::ArrayRef<int32_t>":$wi_data),
+    [{
+        return $_get($_ctxt, mlir::DenseI32ArrayAttr::get($_ctxt, mma_block_size),
+                            mlir::DenseI32ArrayAttr::get($_ctxt, wi_layout),
+                            mlir::DenseI32ArrayAttr::get($_ctxt, wi_data));
+    }]>,
+    AttrBuilder<(ins "llvm::ArrayRef<int32_t>":$wi_layout,
+                      "llvm::ArrayRef<int32_t>":$wi_data),
+    [{
+        return $_get($_ctxt, mlir::DenseI32ArrayAttr(), mlir::DenseI32ArrayAttr::get($_ctxt, wi_layout),
+                            mlir::DenseI32ArrayAttr::get($_ctxt, wi_data));
+    }]>
+  ];
+}
+
+def XeTile_WorkGroupMapAttr : XeTile_Attr<"WorkGroupMap", "wg_map"> {
+  let parameters = (ins
+    "mlir::DenseI32ArrayAttr":$sg_layout,
+    "mlir::DenseI32ArrayAttr":$sg_data
+  );
+  let assemblyFormat = "`<` struct(params) `>`";
+  let genVerifyDecl = true;
+  let builders = [
+    AttrBuilder<(ins "llvm::ArrayRef<int32_t>":$sg_layout,
+                      "llvm::ArrayRef<int32_t>":$sg_data),
+    [{
+        return $_get($_ctxt, mlir::DenseI32ArrayAttr::get($_ctxt, sg_layout),
+                            mlir::DenseI32ArrayAttr::get($_ctxt, sg_data));
+    }]>
+  ];
+}
+
+def XeTile_XeMapAttr : XeTile_Attr<"XeMap", "xe_map"> {
+  let parameters = (ins
+    XeTile_WorkGroupMapAttr:$wg,
+    XeTile_SubGroupMapAttr:$sg
+  );
+  let assemblyFormat = "`<` struct(params) `>`";
+  let builders = [
+    AttrBuilder<(ins "llvm::ArrayRef<int32_t>":$mma_block_size,
+                      "llvm::ArrayRef<int32_t>":$wi_layout,
+                      "llvm::ArrayRef<int32_t>":$wi_data,
+                      "llvm::ArrayRef<int32_t>":$sg_layout,
+                      "llvm::ArrayRef<int32_t>":$sg_data),
+    [{
+        return $_get($_ctxt, WorkGroupMapAttr::get($_ctxt, sg_layout, sg_data),
+          SubGroupMapAttr::get($_ctxt, mma_block_size, wi_layout, wi_data)) ;
+    }]>,
+    AttrBuilder<(ins  "llvm::ArrayRef<int32_t>":$wi_layout,
+                      "llvm::ArrayRef<int32_t>":$wi_data,
+                      "llvm::ArrayRef<int32_t>":$sg_layout,
+                      "llvm::ArrayRef<int32_t>":$sg_data),
+    [{
+        return $_get($_ctxt, WorkGroupMapAttr::get($_ctxt, sg_layout, sg_data),
+          SubGroupMapAttr::get($_ctxt, wi_layout, wi_data)) ;
+    }]>
+  ];
+}
+
+#endif // _XETILE_ATTR_DEF_TD_INCLUDED_

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

@@ -0,0 +1,60 @@
+//===--------------- XeTileOps.td - XeTile dialect  -------*- tablegen -*-===//
+//
+// Copyright 2022 Intel Corporation
+// Part of the IMEX Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+///
+/// \file
+/// This file defines the XeTile dialect.
+///
+//===----------------------------------------------------------------------===//
+
+#ifndef _XETILE_BASE_TD_INCLUDED_
+#define _XETILE_BASE_TD_INCLUDED_
+
+include "mlir/IR/OpBase.td"
+include "mlir/IR/OpAsmInterface.td"
+include "mlir/IR/AttrTypeBase.td"
+include "mlir/IR/BuiltinTypes.td"
+include "mlir/IR/BuiltinTypeInterfaces.td"
+include "mlir/Interfaces/SideEffectInterfaces.td"
+include "mlir/Interfaces/CastInterfaces.td"
+include "mlir/Interfaces/ControlFlowInterfaces.td"
+include "mlir/Interfaces/CopyOpInterface.td"
+include "mlir/Interfaces/InferTypeOpInterface.td"
+include "mlir/Interfaces/ShapedOpInterfaces.td"
+
+// Provide a definition of the 'XeTile' dialect in the ODS framework so that we
+// can define our operations.
+def XeTile_Dialect : Dialect {
+    // The namespace of our dialect
+    let name = "xetile";
+
+    // A short one-line summary
+    let summary = "A dialect for enabling tile-base programming at subgroup level";
+
+    // A longer description
+    let description = [{
+        XeTile provides an abstraction supporting tile-based computation to simplify the
+        lowering of DNN operation like matrix multiplication. XeTile dialect works at tile sizes
+        that are larger than the tile sizes supported by the hardware. XeTile dilaect also hides
+        the auto-padding requirements for out-of-bound memory accesses and, supports arbitrary
+        input matrix sizes.
+    }];
+
+    // The C++ namespace that the dialect class definition resides in.
+    let cppNamespace = "::imex::xetile";
+
+    let dependentDialects = [
+          "::mlir::memref::MemRefDialect"];
+
+    // TODO: temporary disable it.
+    let useDefaultTypePrinterParser = true;
+    let useDefaultAttributePrinterParser = true;
+}
+
+
+#endif // _XETILE_BASE_TD_INCLUDED_

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

@@ -1,4 +1,4 @@
-//===- XeTileOps.h - XeTile dialect  -------*- C++ -*-===//
+//===---------------------- XeTileOps.h - XeTile dialect  -------*- C++ -*-===//
 //
 // Copyright 2022 Intel Corporation
 // Part of the IMEX Project, under the Apache License v2.0 with LLVM Exceptions.
@@ -12,8 +12,8 @@
 ///
 //===----------------------------------------------------------------------===//
 
-#ifndef _XeTile_OPS_H_INCLUDED_
-#define _XeTile_OPS_H_INCLUDED_
+#ifndef _XETILE_OPS_H_INCLUDED_
+#define _XETILE_OPS_H_INCLUDED_
 
 #include <mlir/IR/BuiltinTypeInterfaces.h>
 #include <mlir/IR/BuiltinTypes.h>
@@ -54,16 +54,18 @@ class TileBase : public mlir::Type, public mlir::ShapedType::Trait<TileBase> {
   /// Clone this type with the given shape and element type. If the
   /// provided shape is `None`, the current shape of the type is used.
   TileBase cloneWith(std::optional<llvm::ArrayRef<int64_t>> shape,
-                     mlir::Type elementType) const;
+                     mlir::Type elementType, mlir::Attribute encoding) const;
 };
 
 } // namespace xetile
 } // namespace imex
 
 #include <imex/Dialect/XeTile/IR/XeTileOpsDialect.h.inc>
+#define GET_ATTRDEF_CLASSES
+#include <imex/Dialect/XeTile/IR/XeTileOpsAttrs.h.inc>
 #define GET_TYPEDEF_CLASSES
 #include <imex/Dialect/XeTile/IR/XeTileOpsTypes.h.inc>
 #define GET_OP_CLASSES
 #include <imex/Dialect/XeTile/IR/XeTileOps.h.inc>
 
-#endif // _XeTile_OPS_H_INCLUDED_
+#endif // _XETILE_OPS_H_INCLUDED_