[Conversion] lowering XeGPU.ops to VC-Intrinsics (#669)

load1d/store1d/barrier/gather/scatter/atomic/udpateoffset are all lowered; related xegpu tests run correctly
intel · Nov 10, 2023 · c6de5d1 · c6de5d1
1 parent 1577e15
commit c6de5d1
Show file tree

Hide file tree

Showing 18 changed files with 1,183 additions and 237 deletions.
diff --git a/include/imex/Conversion/Passes.td b/include/imex/Conversion/Passes.td
@@ -250,6 +250,10 @@ memref, arith and math.
   }];
   let constructor = "imex::createConvertGPUXToSPIRVPass()";
   let dependentDialects = ["::mlir::spirv::SPIRVDialect"];
+  let options = [
+    Option<"enableSimtIntrinsic", "enable-simt-intrinsic","bool", "false",
+           "Enable XeGPU.simt Ops lowered to intel genISA simt Intrinsics">
+  ];
 }
 
 //===----------------------------------------------------------------------===//

diff --git a/include/imex/Dialect/XeGPU/IR/XeGPUOps.td b/include/imex/Dialect/XeGPU/IR/XeGPUOps.td
@@ -575,38 +575,28 @@ def XeGPU_InvokeSIMDOp : XeGPU_Op<"invoke_SIMD", []> {
 
 }
 
-def XeGPU_AtomicRMWOp
-  : XeGPU_Op<"atomic_rmw", []> {
-    let summary = "performa ready-modify-write operation that is free from data races.";
-
-    let arguments = (ins
-      XeGPU_AtomicRMWKindAttr:$kind,
-      XeGPU_Vector2DType:$value,
-      XeGPU_TensorDesc:$tensorDesc,
-      XeGPU_MaskType:$mask,
-      DefaultValuedAttr<XeGPU_ModeAttr, "imex::xegpu::Mode::SIMT">: $mode
-    );
-
-    let assemblyFormat = [{
-      $kind $value `,` $tensorDesc `,` $mask (`{` `mode` `=` $mode^ `}`)?
-      attr-dict
-        `:` `(` qualified(type($value)) `,` qualified(type($tensorDesc)) `,` qualified(type($mask)) `)`
-    }];
-  }
+def XeGPU_AtomicRMWOp: XeGPU_Op<"atomic_rmw", []> {
+  let summary = "perform ready-modify-write operation that is free from data races.";
+  let arguments = (ins
+    XeGPU_AtomicRMWKindAttr:$kind,
+    XeGPU_TensorDesc:$tensorDesc,
+    XeGPU_MaskType:$mask,
+    Optional<XeGPU_ValueType>:$value,
+    DefaultValuedAttr<XeGPU_ModeAttr, "imex::xegpu::Mode::SIMT">: $mode
+  );
+  let results = (outs XeGPU_ValueType:$result);
+  let assemblyFormat = [{
+    $kind $tensorDesc `,` $mask (`,` $value^)? (`{` `mode` `=` $mode^ `}`)? attr-dict `:` qualified(type(operands)) `->` type($result)
+  }];
+}
 
 
-def XeGPU_AllocNbarrierOp
-  : XeGPU_Op<"alloc_nbarrier", []> {
+def XeGPU_AllocNbarrierOp: XeGPU_Op<"alloc_nbarrier", []> {
       let summary = "allocate a specific number of named barriers.";
+      let arguments = (ins I32Attr: $nbarrierCount);
+      let assemblyFormat = "$nbarrierCount attr-dict";
+}
 
-      let arguments = (ins
-        I8: $nbarrier_count
-      );
-
-      let assemblyFormat = [{
-        $nbarrier_count attr-dict `:` qualified(type($nbarrier_count))
-      }];
-  }
 
 def XeGPU_CreateNbarrierOp
   : XeGPU_Op<"create_nbarrier", []> {

diff --git a/lib/Conversion/GPUToSPIRV/GPUToSPIRVPass.cpp b/lib/Conversion/GPUToSPIRV/GPUToSPIRVPass.cpp
@@ -158,7 +158,8 @@ void GPUXToSPIRVPass::runOnOperation() {
     target->addIllegalDialect<imex::xegpu::XeGPUDialect>();
     typeConverter.addConversion(
         [&](xegpu::TensorDescType type) -> ::mlir::Type {
-          return ::mlir::IntegerType::get(context, 64);
+          auto i64Type = ::mlir::IntegerType::get(context, 64);
+          return ::mlir::VectorType::get(2, i64Type);
         });
     typeConverter.addConversion([&](::mlir::VectorType type) -> ::mlir::Type {
       unsigned rank = type.getRank();
@@ -175,6 +176,8 @@ void GPUXToSPIRVPass::runOnOperation() {
         for (unsigned i = 0; i < rank; i++) {
           sum *= type.getShape()[i];
         }
+        if (llvm::isa<mlir::IndexType>(elemType))
+          elemType = ::mlir::IntegerType::get(context, 64);
         return ::mlir::VectorType::get(sum, elemType);
       }
     });