[GPU] Register initial GPU pipeline that uses IMEX

include/gc/Transforms/Passes.h

@@ -115,6 +115,10 @@ std::unique_ptr<Pass> createMergeAllocPass();
 void populateFrontendPasses(mlir::OpPassManager &);
 void populateCPUPipeline(mlir::OpPassManager &);
 
+#ifdef GC_USE_IMEX
+void populateGPUPipeline(mlir::OpPassManager &);
+#endif
+
 #define GEN_PASS_DECL
 #include "gc/Transforms/Passes.h.inc"
 

lib/gc/Transforms/GPU/CMakeLists.txt

@@ -1,5 +1,6 @@
 gc_add_mlir_library(GcGpuPasses
   LinalgToXeGPU.cpp
+  Pipeline.cpp
 
   DEPENDS
     GraphCompilerPassIncGen
@@ -18,3 +19,7 @@ gc_add_mlir_library(GcGpuPasses
     GcUtilsIR
 )
 
+include(imex)
+get_property(IMEX_INCLUDES GLOBAL PROPERTY IMEX_INCLUDES)
+target_include_directories(GcGpuPasses PRIVATE ${IMEX_INCLUDES})
+

lib/gc/Transforms/GPU/LinalgToXeGPU.cpp

@@ -1405,6 +1405,17 @@ LogicalResult createMemoryFillKernel(linalg::LinalgOp linalgOp,
   auto outputType = cast<ShapedType>(output.getType());
   auto outputShape = outputType.getShape();
 
+  if (outputShape.size() != 2) {
+    return rewriter.notifyMatchFailure(
+        linalgOp, "Memory fill operation expects 2D output");
+  }
+
+  // Otherwise 'xegpu-to-vc' pass will fail to convert it to VC
+  if (outputShape[0] * outputShape[1] < 16) {
+    return rewriter.notifyMatchFailure(
+        linalgOp, "Memory fill operation is to small to be converted to xegpu");
+  }
+
   // Extract SIMD sized sub-tiles
   int maxSizeSIMD = 256;
   int64_t subTileCols = outputShape[1];

lib/gc/Transforms/GPU/Pipeline.cpp

@@ -0,0 +1,120 @@
+//===- Pipeline.cpp - Graph Compiler GPU pipeline ---------------*- C++ -*-===//
+//
+// This file is licensed 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
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Conversion/Passes.h"
+#include "mlir/Dialect/Arith/Transforms/Passes.h"
+#include "mlir/Dialect/Bufferization/Transforms/OneShotAnalysis.h"
+#include "mlir/Dialect/Bufferization/Transforms/Passes.h"
+#include "mlir/Dialect/LLVMIR/LLVMDialect.h"
+#include "mlir/Dialect/LLVMIR/Transforms/Passes.h"
+#include "mlir/Dialect/Linalg/Passes.h"
+#include "mlir/Dialect/Math/Transforms/Passes.h"
+#include "mlir/Dialect/MemRef/IR/MemRef.h"
+#include "mlir/Dialect/MemRef/Transforms/Passes.h"
+#include "mlir/Dialect/SCF/IR/SCF.h"
+#include "mlir/Dialect/Tensor/IR/Tensor.h"
+#include "mlir/IR/DialectRegistry.h"
+#include "mlir/InitAllPasses.h"
+#include "mlir/Pass/PassManager.h"
+#include "mlir/Support/LogicalResult.h"
+#include "mlir/Transforms/Passes.h"
+#include <iostream>
+
+#include "mlir/Dialect/GPU/IR/GPUDialect.h"
+#include "mlir/Dialect/GPU/Transforms/Passes.h"
+#include "mlir/Dialect/SPIRV/Transforms/Passes.h"
+
+#include <imex/Conversion/Passes.h>
+#include <imex/Transforms/Passes.h>
+
+#include <string>
+
+#include "gc/Transforms/Passes.h"
+
+namespace mlir::gc {
+
+void populateGPUPipeline(mlir::OpPassManager &pm) {
+  pm.addNestedPass<func::FuncOp>(createIterativeTilingAndFusion());
+
+  pm.addPass(bufferization::createEmptyTensorEliminationPass());
+  pm.addPass(bufferization::createEmptyTensorToAllocTensorPass());
+
+  bufferization::OneShotBufferizationOptions options;
+  options.bufferizeFunctionBoundaries = true;
+  options.setFunctionBoundaryTypeConversion(
+      bufferization::LayoutMapOption::IdentityLayoutMap);
+  pm.addPass(bufferization::createOneShotBufferizePass(options));
+
+  pm.addPass(bufferization::createDropEquivalentBufferResultsPass());
+  pm.addNestedPass<func::FuncOp>(
+      bufferization::createFinalizingBufferizePass());
+  pm.addPass(createCanonicalizerPass());
+  pm.addPass(createCSEPass());
+  pm.addPass(bufferization::createDropEquivalentBufferResultsPass());
+  pm.addPass(memref::createExpandReallocPass());
+  pm.addPass(createCanonicalizerPass());
+  pm.addPass(bufferization::createOwnershipBasedBufferDeallocationPass());
+  pm.addPass(createCanonicalizerPass());
+  pm.addPass(bufferization::createBufferDeallocationSimplificationPass());
+  pm.addPass(bufferization::createLowerDeallocationsPass());
+  pm.addPass(createCSEPass());
+  pm.addPass(createCanonicalizerPass());
+  pm.addPass(createBufferizationToMemRefPass());
+
+  pm.addNestedPass<func::FuncOp>(createForallToParallelLoopPass());
+  pm.addNestedPass<func::FuncOp>(createLinalgToXeGPU(
+      {/*kTile=*/16, /*stages=*/1, /*dpasTiles=*/{8, 16, 16}}));
+
+  pm.addNestedPass<func::FuncOp>(createConvertLinalgToLoopsPass());
+  pm.addPass(xegpu::createXeGPUFoldAliasOps());
+  pm.addPass(memref::createFoldMemRefAliasOpsPass());
+  pm.addNestedPass<func::FuncOp>(createGpuMapParallelLoopsPass());
+  pm.addNestedPass<func::FuncOp>(createParallelLoopToGpuPass());
+
+  pm.addNestedPass<func::FuncOp>(imex::createInsertGPUAllocsPass("opencl"));
+  pm.addPass(createGpuKernelOutliningPass());
+  pm.addPass(createCanonicalizerPass());
+  pm.addPass(imex::createSetSPIRVCapabilitiesPass());
+  pm.addNestedPass<gpu::GPUModuleOp>(
+      imex::createSetSPIRVAbiAttributePass("opencl"));
+  pm.addPass(createLowerAffinePass());
+  pm.addPass(imex::createVectorLinearizePass());
+  pm.addNestedPass<gpu::GPUModuleOp>(imex::createConvertXeGPUToVCPass());
+  pm.addPass(createReconcileUnrealizedCastsPass());
+  pm.addPass(imex::createBF16ToGPUPass());
+  pm.addNestedPass<gpu::GPUModuleOp>(createConvertFuncToSPIRVPass());
+  pm.addNestedPass<gpu::GPUModuleOp>(createConvertVectorToSPIRVPass());
+  pm.addPass(imex::createConvertGPUXToSPIRVPass());
+  pm.addNestedPass<spirv::ModuleOp>(spirv::createSPIRVLowerABIAttributesPass());
+  pm.addNestedPass<spirv::ModuleOp>(spirv::createSPIRVUpdateVCEPass());
+  pm.addNestedPass<func::FuncOp>(LLVM::createRequestCWrappersPass());
+  pm.addPass(imex::createSerializeSPIRVPass());
+  pm.addPass(createConvertVectorToSCFPass());
+  pm.addPass(imex::createConvertGPUToGPUXPass());
+  pm.addPass(createConvertSCFToCFPass());
+  pm.addPass(createConvertControlFlowToLLVMPass());
+  pm.addPass(createConvertVectorToLLVMPass());
+  pm.addPass(createConvertIndexToLLVMPass());
+  pm.addPass(createArithToLLVMConversionPass());
+  pm.addPass(createConvertFuncToLLVMPass());
+  pm.addPass(createConvertMathToLLVMPass());
+  pm.addPass(imex::createConvertGPUXToLLVMPass());
+  pm.addPass(createConvertIndexToLLVMPass());
+  pm.addPass(memref::createExpandStridedMetadataPass());
+  pm.addPass(createLowerAffinePass());
+  pm.addPass(createFinalizeMemRefToLLVMConversionPass());
+  pm.addPass(createReconcileUnrealizedCastsPass());
+}
+
+void registerGPUPipeline() {
+  PassPipelineRegistration<>("gc-gpu-pipeline",
+                             "The GPU pipeline for Graph Compiler with IMEX",
+                             populateGPUPipeline);
+}
+
+} // namespace mlir::gc

src/gc-opt/gc-opt.cpp

@@ -36,6 +36,9 @@
 
 namespace mlir::gc {
 void registerCPUPipeline();
+#ifdef GC_USE_IMEX
+void registerGPUPipeline();
+#endif
 } // namespace mlir::gc
 
 int main(int argc, char *argv[]) {
@@ -47,6 +50,7 @@ int main(int argc, char *argv[]) {
   imex::registerConvertGPUXToSPIRV();
   imex::registerConvertXeGPUToVC();
   imex::registerConvertXeTileToXeGPU();
+  mlir::gc::registerGPUPipeline();
 #endif
   mlir::registerAllPasses();
   mlir::gc::registerCPUPipeline();