[xla:gpu] Add support for capturing lmhlo_gpu.conv operations into cuda graphs

PiperOrigin-RevId: 501950550

Author: ezhulenev

xla/mlir/backends/gpu/transforms/outline_cuda_graphs.cc

@@ -37,6 +37,7 @@ limitations under the License.
 #include "xla/mlir/runtime/ir/rt_dialect.h"
 #include "xla/mlir/runtime/ir/rt_ops.h"
 #include "xla/mlir/runtime/utils/custom_calls.h"
+#include "xla/mlir_hlo/lhlo_gpu/IR/lhlo_gpu_ops.h"
 
 namespace xla {
 namespace gpu {
@@ -80,23 +81,29 @@ using CaptureSequence =
 
 //===----------------------------------------------------------------------===//
 
-template <typename T, OpCapturePattern::Capture capture>
+template <OpCapturePattern::Capture capture, typename T, typename... Ts>
 struct OpCapture : public OpCapturePattern {
   FailureOr<OpCapturePattern::Capture> match(Operation* op) final {
-    if (isa<T>(op)) return capture;
+    if (isa<T, Ts...>(op)) return capture;
     return failure();
   }
 };
 
 static constexpr auto kMove = OpCapturePattern::Capture::kMove;
 static constexpr auto kClone = OpCapturePattern::Capture::kClone;
 
+template <typename T, typename... Ts>
+using MoveOp = OpCapture<kMove, T, Ts...>;
+template <typename T, typename... Ts>
+using CloneOp = OpCapture<kClone, T, Ts...>;
+
 // Capture gpu operations by moving them intp graph capture function.
-struct LaunchFuncOpCapture : public OpCapture<LaunchFuncOp, kMove> {};
+struct LaunchFuncOpCapture : public MoveOp<LaunchFuncOp> {};
+struct ConvOpCapture : public MoveOp<lmhlo_gpu::ConvForwardFusedOp> {};
 
 // Capture pure operations by cloning them into graph capture function.
-struct ConstantOpCapture : public OpCapture<arith::ConstantOp, kClone> {};
-struct ViewOpCapture : public OpCapture<memref::ViewOp, kClone> {};
+struct ConstantOpCapture : public CloneOp<arith::ConstantOp> {};
+struct ViewOpCapture : public CloneOp<memref::ViewOp> {};
 
 //===----------------------------------------------------------------------===//
 
@@ -320,6 +327,7 @@ void OutlineCudaGraphsPass::runOnOperation() {
 
   OpCapturePatternSet patterns;
   patterns.emplace_back(new LaunchFuncOpCapture());
+  patterns.emplace_back(new ConvOpCapture());
   patterns.emplace_back(new ConstantOpCapture());
   patterns.emplace_back(new ViewOpCapture());
 

xla/mlir/backends/gpu/transforms/passes.cc

@@ -32,18 +32,22 @@ void populateXlaGpuRuntimePasses(mlir::OpPassManager& pm,
   // Lower operations with registered IR emitters to Gpu launches.
   pm.addPass(createConvertLmhloToGpuLaunchPass(thunk_sequence));
 
+  // Clean up IR before converting it to the runtime operations.
+  pm.addPass(createCSEPass());
+  pm.addPass(createCanonicalizerPass());
+
   // Convert global memrefs corresponding to constant arguments.
   pm.addPass(createConvertMemrefGetGlobalToArgPass());
   pm.addPass(createSymbolDCEPass());  // Clean up unused global constants.
 
-  // Lower all Gpu operations to the XLA Gpu runtime custom calls.
-  pm.addPass(createConvertLmhloGpuToGpuRuntimePass());
-  pm.addPass(createConvertLmhloToGpuRuntimePass());
-
+  // Outline CUDA-Graph-compatible operations into graph capture functions.
   if (opts.enable_cuda_graphs) {
     pm.addPass(createOutlineCudaGraphsPass());
   }
 
+  // Lower all Gpu operations to the XLA Gpu runtime custom calls.
+  pm.addPass(createConvertLmhloGpuToGpuRuntimePass());
+  pm.addPass(createConvertLmhloToGpuRuntimePass());
   pm.addPass(createConvertGpuToGpuRuntimePass());
 
   // Add performance tracing annotations.

xla/service/gpu/runtime/BUILD

@@ -165,6 +165,7 @@ cc_library(
     srcs = ["graph_launch.cc"],
     hdrs = ["graph_launch.h"],
     deps = [
+        ":conv",
         ":kernel_launch",
         ":support",
         "//xla:types",

xla/service/gpu/runtime/graph_launch.cc

@@ -25,6 +25,7 @@ limitations under the License.
 
 #include "xla/runtime/custom_call.h"
 #include "xla/runtime/executable.h"
+#include "xla/service/gpu/runtime/conv.h"
 #include "xla/service/gpu/runtime/kernel_launch.h"
 #include "xla/service/gpu/runtime/support.h"
 #include "xla/service/service_executable_run_options.h"
@@ -204,9 +205,11 @@ static absl::StatusOr<OwnedGraph> CaptureGraph(
 //===----------------------------------------------------------------------===//
 
 static absl::Status LaunchGraph(
-    const ServiceExecutableRunOptions* run_options, const std::string* ptx,
+    const ServiceExecutableRunOptions* run_options,
+    const DebugOptions* debug_options, const std::string* ptx,
     const std::vector<uint8_t>* cubin, se::DeviceMemoryBase* temp_buffer,
     StreamExecutorKernels::Snapshot* kernels,
+    StreamExecutorConvRunners::Snapshot* convs,
     GraphInstances::Snapshot* instances, runtime::Executable* executable,
     CustomCall::RemainingArgs fwd_args, CustomCall::FunctionOrdinal capture) {
 #if GOOGLE_CUDA
@@ -220,8 +223,8 @@ static absl::Status LaunchGraph(
 
   // Forwards user data required for launching kernels.
   auto user_data = [&] {
-    return CustomCall::UserData(run_options, ptx, cubin, temp_buffer, kernels,
-                                executable);
+    return CustomCall::UserData(run_options, debug_options, ptx, cubin,
+                                temp_buffer, kernels, convs, executable);
   };
 
   absl::StatusOr<GraphInstance*> instance = instances->GetOrCreate(
@@ -294,10 +297,12 @@ XLA_RUNTIME_DEFINE_CUSTOM_CALL(
     Launch, FunctionWrapper<LaunchGraph>(), checks,
     CustomCall::Bind("xla.gpu.cuda.graph.launch")
         .UserData<const ServiceExecutableRunOptions*>()
+        .UserData<const DebugOptions*>()
         .UserData<const std::string*>()
         .UserData<const std::vector<uint8_t>*>()
         .UserData<se::DeviceMemoryBase*>()
         .UserData<StreamExecutorKernels::Snapshot*>()
+        .UserData<StreamExecutorConvRunners::Snapshot*>()
         .UserData<GraphInstances::Snapshot*>()
         .UserData<Executable*>()
         .RemainingArgs()