commit 5f9a96364afdf4e7fa8891a250204c25f768a225

Author: Chao Chen <[email protected]>
Date:   Thu Sep 21 10:22:40 2023 -0700

    Update XeGPU Dialect.

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

@@ -1,2 +1,9 @@
 add_mlir_dialect(XeGPUOps xegpu)
 add_mlir_doc(XeGPUOps XeGPUDialect Dialects/ -gen-dialect-doc)
+
+set(LLVM_TARGET_DEFINITIONS XeGPUOps.td)
+mlir_tablegen(XeGPUOpsAttrs.h.inc -gen-attrdef-decls)
+mlir_tablegen(XeGPUOpsAttrs.cpp.inc -gen-attrdef-defs)
+mlir_tablegen(XeGPUOpsEnums.h.inc -gen-enum-decls)
+mlir_tablegen(XeGPUOpsEnums.cpp.inc -gen-enum-defs)
+add_public_tablegen_target(XeGPUOpsAttrsIncGen)

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

@@ -0,0 +1,73 @@
+
+#ifndef _XEGPU_ATTRS_TD_INCLUDED_
+#define _XEGPU_ATTRS_TD_INCLUDED_
+
+include "imex/Dialect/XeGPU/IR/XeGPUDialect.td"
+
+include "mlir/IR/EnumAttr.td"
+
+class XeGPUAttrDef<string name, string attrMnemonic, list<Trait> traits = [], string baseCppClass = "::mlir::Attribute">
+    : AttrDef<XeGPUDialect, name, traits, baseCppClass> {
+  let mnemonic = attrMnemonic;
+}
+
+def XeGPU_ScatteredAttr : XeGPUAttrDef<"Scattered", "scattered", []> {
+  let summary = "Scattered attribute for scattered read and write operation.";
+  let description = [{An attribute represent scattered read and write operation.
+    It does not (need to) have meaningful input values. The existence of itself
+    implies scattered read/write.}];
+
+  let assemblyFormat = "";
+}
+
+
+def XeGPU_MemoryScopeAttr : I32EnumAttr<
+    "MemoryScope", "", [ I32EnumAttrCase<"GLOBAL", 1, "global">,
+                         I32EnumAttrCase<"SLM", 2, "slm"> ]> {
+    let cppNamespace = "::imex::xegpu";
+}
+
+def XeGPU_CacheReadAttr : I32EnumAttr<
+    "CacheReadHint", "", [ I32EnumAttrCase<"UNCACHED", 0, "uncached">,
+                           I32EnumAttrCase<"CACHED", 1, "cached">,
+                           I32EnumAttrCase<"STREAMING", 2, "streaming">,
+                           I32EnumAttrCase<"READ_INVALDIATE", 3, "read_invalidiate"> ]> {
+
+    let cppNamespace = "::imex::xegpu";
+}
+
+def XeGPU_CacheWriteAttr : I32EnumAttr<
+    "CacheWriteHint", "", [ I32EnumAttrCase<"UNCACHED", 0, "uncached">,
+                            I32EnumAttrCase<"WRITE_THROUGH", 1, "write_through">,
+                            I32EnumAttrCase<"WRITE_BACK", 2, "write_back">,
+                            I32EnumAttrCase<"STREAMING", 3, "streaming"> ]> {
+
+    let cppNamespace = "::imex::xegpu";
+}
+
+// RMW kind attribute
+def ATOMIC_RMW_KIND_ADDF    : I64EnumAttrCase<"addf", 0>;
+def ATOMIC_RMW_KIND_ADDI    : I64EnumAttrCase<"addi", 1>;
+def ATOMIC_RMW_KIND_ASSIGN  : I64EnumAttrCase<"assign", 2>;
+def ATOMIC_RMW_KIND_MAXF    : I64EnumAttrCase<"maxf", 3>;
+def ATOMIC_RMW_KIND_MAXS    : I64EnumAttrCase<"maxs", 4>;
+def ATOMIC_RMW_KIND_MAXU    : I64EnumAttrCase<"maxu", 5>;
+def ATOMIC_RMW_KIND_MINF    : I64EnumAttrCase<"minf", 6>;
+def ATOMIC_RMW_KIND_MINS    : I64EnumAttrCase<"mins", 7>;
+def ATOMIC_RMW_KIND_MINU    : I64EnumAttrCase<"minu", 8>;
+def ATOMIC_RMW_KIND_MULF    : I64EnumAttrCase<"mulf", 9>;
+def ATOMIC_RMW_KIND_MULI    : I64EnumAttrCase<"muli", 10>;
+def ATOMIC_RMW_KIND_ORI     : I64EnumAttrCase<"ori", 11>;
+def ATOMIC_RMW_KIND_ANDI    : I64EnumAttrCase<"andi", 12>;
+
+def XeGPU_AtomicRMWKindAttr : I64EnumAttr<
+    "AtomicRMWKind", "",
+    [ATOMIC_RMW_KIND_ADDF, ATOMIC_RMW_KIND_ADDI, ATOMIC_RMW_KIND_ASSIGN,
+     ATOMIC_RMW_KIND_MAXF, ATOMIC_RMW_KIND_MAXS, ATOMIC_RMW_KIND_MAXU,
+     ATOMIC_RMW_KIND_MINF, ATOMIC_RMW_KIND_MINS, ATOMIC_RMW_KIND_MINU,
+     ATOMIC_RMW_KIND_MULF, ATOMIC_RMW_KIND_MULI, ATOMIC_RMW_KIND_ORI,
+     ATOMIC_RMW_KIND_ANDI]> {
+  let cppNamespace = "::imex::xegpu";
+}
+
+#endif // _XEGPU_ATTRS_TD_INCLUDED_

include/imex/Dialect/XeGPU/IR/XeGPUDialect.td

@@ -0,0 +1,46 @@
+#ifndef _XEGPU_DIALECT_TD_INCLUDED_
+#define _XEGPU_DIALECT_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/ViewLikeInterface.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 'XeGPU' dialect in the ODS framework so that we
+// can define our operations.
+def XeGPUDialect : Dialect {
+    // The namespace of our dialect
+    let name = "xegpu";
+
+    // A short one-line summary of our dialect.
+    let summary = "The XeGPU dialect that models Intel GPU's ISA";
+
+    // A longer description of our dialect.
+    let description = [{
+            The XeGPU dialect models Intel Xe ISA semantics but works at vector and
+            TensorDesc data type. It provides 1:1 mappings to match Xe instructions like
+            DPAS and 2D block load. The matrix size being processed at this level
+            exactly matches the hardware instructions or the intrinsic supported by
+            the lower-level GPU compiler.
+        }];
+
+    // The C++ namespace that the dialect class definition resides in.
+    let cppNamespace = "::imex::xegpu";
+
+    let dependentDialects = ["::mlir::memref::MemRefDialect"];
+
+    // TODO: temporary disable it.
+    let useDefaultTypePrinterParser = true;
+    let useDefaultAttributePrinterParser = true;
+}
+
+#endif //XEGPU_DIALECT

include/imex/Dialect/XeGPU/IR/XeGPUOps.h

@@ -41,15 +41,16 @@ SmallVector<Range, 8> getOrCreateRanges(OffsetSizeAndStrideOpInterface op,
 namespace imex {
 namespace xegpu {
 
-class TileType;
+class TensorDescType;
 
 } // namespace xegpu
 } // namespace imex
 
 namespace imex {
 namespace xegpu {
 
-class TileBase : public mlir::Type, public mlir::ShapedType::Trait<TileBase> {
+class BaseTensorDescType : public mlir::Type,
+                           public mlir::ShapedType::Trait<BaseTensorDescType> {
 public:
   using Type::Type;
 
@@ -64,8 +65,8 @@ 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;
+  BaseTensorDescType cloneWith(std::optional<llvm::ArrayRef<int64_t>> shape,
+                               mlir::Type elementType) const;
 
   /// Return true if the specified element type is ok in a tensor.
   static bool isValidElementType(Type type);
@@ -81,6 +82,9 @@ class TileBase : public mlir::Type, public mlir::ShapedType::Trait<TileBase> {
 } // namespace imex
 
 #include <imex/Dialect/XeGPU/IR/XeGPUOpsDialect.h.inc>
+#include <imex/Dialect/XeGPU/IR/XeGPUOpsEnums.h.inc>
+#define GET_ATTRDEF_CLASSES
+#include <imex/Dialect/XeGPU/IR/XeGPUOpsAttrs.h.inc>
 #define GET_TYPEDEF_CLASSES
 #include <imex/Dialect/XeGPU/IR/XeGPUOpsTypes.h.inc>
 #define GET_OP_CLASSES