Add a pass option to provide the XeGPU code level.

XeGPU allows worgroup, subgroup, and workitem level programming.
This options lets the pass manager know at which level the
XeGPU ops belong to.

Author: mshahneo

mlir/include/mlir/Dialect/GPU/Pipelines/Passes.h

@@ -63,11 +63,17 @@ struct GPUToNVVMPipelineOptions
 // Options for the gpu to xevm pipeline.
 struct GPUToXeVMPipelineOptions
     : public PassPipelineOptions<GPUToXeVMPipelineOptions> {
+  // XeGPU op granularity selection: workgroup | subgroup | workitem
+  PassOptions::Option<std::string> xegpuOpLevel{
+      *this, "xegpu-op-level",
+      llvm::cl::desc("Granularity of XeGPU operations to target: workgroup | "
+                     "subgroup | workitem"),
+      llvm::cl::init("workgroup")};
   // General lowering controls.
-  PassOptions::Option<int64_t> indexBitWidth{
-      *this, "index-bitwidth",
+  PassOptions::Option<bool> use64bitIndex{
+      *this, "use-64bit-index",
       llvm::cl::desc("Bitwidth of the index type (host & device)"),
-      llvm::cl::init(64)};
+      llvm::cl::init(true)};
   PassOptions::Option<bool> kernelBarePtrCallConv{
       *this, "kernel-bare-ptr-calling-convention",
       llvm::cl::desc("Use bare pointer calling convention for device kernels"),

mlir/lib/Dialect/GPU/Pipelines/GPUToXeVMPipeline.cpp

@@ -62,24 +62,30 @@ void buildCommonPassPipeline(
 //===----------------------------------------------------------------------===//
 void buildGpuPassPipeline(OpPassManager &pm,
                           const mlir::gpu::GPUToXeVMPipelineOptions &options) {
-  pm.addNestedPass<gpu::GPUModuleOp>(xegpu::createXeGPUWgToSgDistribute());
-  pm.addNestedPass<gpu::GPUModuleOp>(createCSEPass());
-  pm.addNestedPass<gpu::GPUModuleOp>(createLowerAffinePass());
-  pm.addNestedPass<gpu::GPUModuleOp>(xegpu::createXeGPUBlocking());
-  pm.addNestedPass<gpu::GPUModuleOp>(createCanonicalizerPass());
-  pm.addNestedPass<gpu::GPUModuleOp>(createCSEPass());
-  pm.addNestedPass<gpu::GPUModuleOp>(xegpu::createXeGPUPropagateLayout());
-  pm.addNestedPass<gpu::GPUModuleOp>(xegpu::createXeGPUSubgroupDistribute());
-  pm.addNestedPass<gpu::GPUModuleOp>(createCanonicalizerPass());
-  pm.addNestedPass<gpu::GPUModuleOp>(createCSEPass());
-  pm.addNestedPass<gpu::GPUModuleOp>(createLoopInvariantCodeMotionPass());
-  pm.addNestedPass<gpu::GPUModuleOp>(createCSEPass());
-  pm.addNestedPass<gpu::GPUModuleOp>(xegpu::createXeGPUVectorLinearize());
+  if (options.xegpuOpLevel == "workgroup") {
+    pm.addNestedPass<gpu::GPUModuleOp>(xegpu::createXeGPUWgToSgDistribute());
+    pm.addNestedPass<gpu::GPUModuleOp>(createCSEPass());
+    pm.addNestedPass<gpu::GPUModuleOp>(createLowerAffinePass());
+    pm.addNestedPass<gpu::GPUModuleOp>(xegpu::createXeGPUBlocking());
+    pm.addNestedPass<gpu::GPUModuleOp>(createCanonicalizerPass());
+    pm.addNestedPass<gpu::GPUModuleOp>(createCSEPass());
+  }
+  if (options.xegpuOpLevel == "subgroup") {
+    pm.addNestedPass<gpu::GPUModuleOp>(xegpu::createXeGPUPropagateLayout());
+    pm.addNestedPass<gpu::GPUModuleOp>(xegpu::createXeGPUSubgroupDistribute());
+    pm.addNestedPass<gpu::GPUModuleOp>(createCanonicalizerPass());
+    pm.addNestedPass<gpu::GPUModuleOp>(createCSEPass());
+    pm.addNestedPass<gpu::GPUModuleOp>(createLoopInvariantCodeMotionPass());
+    pm.addNestedPass<gpu::GPUModuleOp>(createCSEPass());
+    pm.addNestedPass<gpu::GPUModuleOp>(xegpu::createXeGPUVectorLinearize());
+  }
   pm.addNestedPass<gpu::GPUModuleOp>(createConvertXeGPUToXeVMPass());
-  ConvertGpuOpsToLLVMSPVOpsOptions gpuToLLVMSPVOptions;
-  gpuToLLVMSPVOptions.use64bitIndex = options.indexBitWidth;
-  pm.addNestedPass<gpu::GPUModuleOp>(
-      createConvertGpuOpsToLLVMSPVOps(gpuToLLVMSPVOptions));
+  {
+    ConvertGpuOpsToLLVMSPVOpsOptions gpuToLLVMSPVOptions;
+    gpuToLLVMSPVOptions.use64bitIndex = options.use64bitIndex;
+    pm.addNestedPass<gpu::GPUModuleOp>(
+        createConvertGpuOpsToLLVMSPVOps(gpuToLLVMSPVOptions));
+  }
   pm.addNestedPass<gpu::GPUModuleOp>(createConvertXeVMToLLVMPass());
   pm.addNestedPass<gpu::GPUModuleOp>(createCSEPass());
 }
@@ -104,14 +110,14 @@ void buildHostPostPipeline(OpPassManager &pm,
   }
   pm.addPass(createConvertToLLVMPass());
   pm.addPass(createLowerAffinePass());
+  pm.addPass(createReconcileUnrealizedCastsPass());
   // gpu-module-to-binary
   {
     GpuModuleToBinaryPassOptions gpuToModuleBinOptions;
     gpuToModuleBinOptions.compilationTarget = options.binaryFormat;
     gpuToModuleBinOptions.cmdOptions = options.cmdOptions;
     pm.addPass(createGpuModuleToBinaryPass(gpuToModuleBinOptions));
   }
-  pm.addPass(createReconcileUnrealizedCastsPass());
 }
 } // namespace