Enable end to end DPS testing

mlir-tensorrt/compiler/include/mlir-tensorrt/Dialect/Plan/Transforms/Passes.td

@@ -248,8 +248,8 @@ def StablehloClusteringPass : Pass<"stablehlo-clustering", "::mlir::ModuleOp"> {
     Option<"entrypoint", "entrypoint", "std::string", "\"\"",
       "the name of the entrypoint function; if empty then the clustering runs"
       " on all functions">,
-    Option<"enableNonDPSReturns",
-      "enable-non-dps-returns", "bool", "false",
+    Option<"forceEntrypointsReturnAllocs",
+      "force-entrypoints-return-allocs", "bool", "false",
       "allow backend clusters to directly allocate outputs">,
     Option<"disableCreateShapeFuncPass", "disable-create-shape-func-pass", "bool", "false",
       "don't apply create shape to func pass in TensorRT clusters">
@@ -331,7 +331,7 @@ def CreateClosedRegionsPass : Pass<"plan-create-closed-regions", "::mlir::Module
       "(used only in testing) specifies to outline regions by walking in "
       " pre-order; used for verifying results are not sensitive "
       "to traversal order">,
-    Option<"enableNonDPSReturns", "enable-non-dps-returns", "bool",
+    Option<"forceEntrypointsReturnAllocs", "force-entrypoints-return-allocs", "bool",
            /*default=*/"false",
            "Allow backend clusters to directly allocate outputs">
   ];

mlir-tensorrt/compiler/lib/Compiler/StablehloToExecutable/StablehloToExecutable.cpp

@@ -130,7 +130,10 @@ void StablehloToExecutableTask::buildPostClusteringPipeline(
 
   // Perform bufferization.
   pm.addPass(createMemRefCastEliminationPass());
-  pm.addPass(plan::createPlanAllocTensorsPass());
+  plan::PlanAllocTensorsPassOptions allocTensorOpts{};
+  allocTensorOpts.forceEntrypointsReturnAllocs =
+      opts.forceEntrypointsReturnAllocs;
+  pm.addPass(plan::createPlanAllocTensorsPass(allocTensorOpts));
   pm.addPass(plan::createPlanBufferizePass());
   pm.addPass(createMemRefCastEliminationPass());
   pm.addPass(createCanonicalizerPass());

mlir-tensorrt/compiler/lib/Dialect/Plan/Transforms/CMakeLists.txt

@@ -37,6 +37,7 @@ add_mlir_tensorrt_library(MLIRTensorRTPlanTransforms
   MLIRTensorRTStablehloScalarToArith
   MLIRTensorRTStablehloToTensorRT
   MLIRTensorRTTensorRTRuntimeDialect
+  MLIRBufferizationToMemRef
   MLIRTransforms
   StablehloOps
 )

mlir-tensorrt/compiler/lib/Dialect/Plan/Transforms/CreateClosedRegions.cpp

@@ -561,12 +561,12 @@ createInlineClosedAllocGroupOp(RewriterBase &rewriter, plan::InlineGroupOp op,
 static LogicalResult createClosedGroupOp(RewriterBase &rewriter,
                                          plan::InlineGroupOp op,
                                          DataFlowSolver &solver,
-                                         bool enableNonDPSReturns) {
+                                         bool forceEntrypointsReturnAllocs) {
   OpBuilder::InsertionGuard g(rewriter);
 
   // Materialize destination operands if not using non-DPS call convention.
   SmallVector<DestinationOperandMaterializationResult> destinationOperands;
-  if (!enableNonDPSReturns)
+  if (!forceEntrypointsReturnAllocs)
     if (failed(materializeDestinationOperands(rewriter, op, solver,
                                               destinationOperands)))
       return failure();
@@ -581,7 +581,7 @@ static LogicalResult createClosedGroupOp(RewriterBase &rewriter,
 
   // Create and populate the appropriate closed group op based on call
   // convention.
-  if (!enableNonDPSReturns)
+  if (!forceEntrypointsReturnAllocs)
     return createInlineClosedGroupOp(rewriter, op, solver, inputs,
                                      destinationOperands);
   return createInlineClosedAllocGroupOp(rewriter, op, solver, inputs);
@@ -629,7 +629,7 @@ class CreateClosedRegionsPass
     IRRewriter rewriter(ctx);
     for (InlineGroupOp groupOp : llvm::make_early_inc_range(groupOps)) {
       if (failed(createClosedGroupOp(rewriter, groupOp, solver,
-                                     enableNonDPSReturns)))
+                                     forceEntrypointsReturnAllocs)))
         return signalPassFailure();
     }
   }

mlir-tensorrt/compiler/lib/Dialect/Plan/Transforms/Passes.cpp

@@ -24,6 +24,7 @@
 //===----------------------------------------------------------------------===//
 #include "mlir-tensorrt/Dialect/Plan/Transforms/Passes.h"
 #include "mlir-tensorrt/Transforms/Passes.h"
+#include "mlir/Conversion/BufferizationToMemRef/BufferizationToMemRef.h"
 #include "mlir/Dialect/Bufferization/IR/BufferDeallocationOpInterface.h"
 #include "mlir/Dialect/Bufferization/Pipelines/Passes.h"
 #include "mlir/Dialect/Bufferization/Transforms/Passes.h"
@@ -48,7 +49,8 @@ void plan::buildPlanSegmentationPipeline(
       plan::createPlanPopulateFunctionBoundsAttributesPass());
   pm.addPass(plan::createStablehloClusteringPass(opts));
   plan::CreateClosedRegionsPassOptions closedRegionOptions{};
-  closedRegionOptions.enableNonDPSReturns = opts.enableNonDPSReturns;
+  closedRegionOptions.forceEntrypointsReturnAllocs =
+      opts.forceEntrypointsReturnAllocs;
   pm.addPass(plan::createCreateClosedRegionsPass(closedRegionOptions));
   pm.addPass(plan::createOutlineClustersPass());
   pm.addPass(mlir::createFuncExtDuplicateFunctionEliminationPass());
@@ -80,6 +82,7 @@ void plan::buildPlanBufferDeallocationPipeline(
   pm.addPass(createCanonicalizerPass());
   pm.addPass(bufferization::createBufferDeallocationSimplificationPass());
   pm.addPass(bufferization::createLowerDeallocationsPass());
+  pm.addPass(mlir::createBufferizationToMemRefPass());
   pm.addPass(createCSEPass());
   pm.addPass(createCanonicalizerPass());
 }

mlir-tensorrt/compiler/lib/Dialect/Plan/Transforms/StablehloClustering.cpp

@@ -277,11 +277,11 @@ class StablehloClusteringPass
                });
 
     for (func::FuncOp func : funcs) {
-      if (failed(
-              applyClusteringToFunc(rewriter, func, solver, schedule,
-                                    StablehloClusteringPassOptions{
-                                        entrypoint, enableNonDPSReturns,
-                                        /*disableCreateShapeFuncPass=*/false})))
+      if (failed(applyClusteringToFunc(
+              rewriter, func, solver, schedule,
+              StablehloClusteringPassOptions{
+                  entrypoint, forceEntrypointsReturnAllocs,
+                  /*disableCreateShapeFuncPass=*/false})))
         return signalPassFailure();
     }
 

mlir-tensorrt/compiler/test/Dialect/Plan/create-closed-regions.mlir

@@ -1,6 +1,6 @@
 // RUN: mlir-tensorrt-opt %s -plan-create-closed-regions -split-input-file | FileCheck %s
 // RUN: mlir-tensorrt-opt %s -plan-create-closed-regions=test-pre-walk-order=true -split-input-file | FileCheck %s
-// RUN: mlir-tensorrt-opt %s -plan-create-closed-regions=enable-non-dps-returns=true -split-input-file | FileCheck %s --check-prefix=CHECK-ALLOC
+// RUN: mlir-tensorrt-opt %s -plan-create-closed-regions=force-entrypoints-return-allocs=true -split-input-file | FileCheck %s --check-prefix=CHECK-ALLOC
 
 func.func @test_simple_static(%arg0: tensor<10xf32>, %arg1: tensor<10xf32>) -> tensor<10xf32> {
   %0 = plan.inline_group target(#plan.tensorrt_cluster<disallow_shape_tensor_calculations = false, benefit = 1>) -> tensor<10xf32> {