1 files changed, 16 insertions, 18 deletions
diff --git a/mlir/lib/Transforms/DmaGeneration.cpp b/mlir/lib/Transforms/DmaGeneration.cpp
index 29cc435a8a9..06aa3758712 100644
--- a/mlir/lib/Transforms/DmaGeneration.cpp
+++ b/mlir/lib/Transforms/DmaGeneration.cpp
@@ -42,14 +42,19 @@ using llvm::SmallMapVector;
 
 static llvm::cl::OptionCategory clOptionsCategory(DEBUG_TYPE " options");
 
-static llvm::cl::opt<unsigned> clFastMemorySpace(
-    "dma-fast-mem-space", llvm::cl::Hidden,
-    llvm::cl::desc("Set fast memory space id for DMA generation"),
+static llvm::cl::opt<unsigned> clFastMemoryCapacity(
+    "dma-fast-mem-capacity",
+    llvm::cl::init(std::numeric_limits<uint64_t>::max() / 1024),
+    llvm::cl::desc(
+        "Set fast memory space capacity in KiB (default: unlimited)"),
     llvm::cl::cat(clOptionsCategory));
 
-static llvm::cl::opt<unsigned> clFastMemoryCapacity(
-    "dma-fast-mem-capacity", llvm::cl::Hidden,
-    llvm::cl::desc("Set fast memory space capacity in KiB"),
+static const unsigned kDefaultFastMemorySpace = 1;
+
+static llvm::cl::opt<unsigned> clFastMemorySpace(
+    "dma-fast-mem-space", llvm::cl::init(kDefaultFastMemorySpace),
+    llvm::cl::desc(
+        "Fast memory space identifier for DMA generation (default: 1)"),
     llvm::cl::cat(clOptionsCategory));
 
 namespace {
@@ -60,10 +65,11 @@ namespace {
 // TODO(bondhugula): We currently can't generate DMAs correctly when stores are
 // strided. Check for strided stores.
 struct DmaGeneration : public FunctionPass {
-  explicit DmaGeneration(
-      unsigned slowMemorySpace = 0, unsigned fastMemorySpace = 1,
-      int minDmaTransferSize = 1024,
-      uint64_t fastMemCapacityBytes = std::numeric_limits<uint64_t>::max())
+  explicit DmaGeneration(unsigned slowMemorySpace = 0,
+                         unsigned fastMemorySpace = clFastMemorySpace,
+                         int minDmaTransferSize = 1024,
+                         uint64_t fastMemCapacityBytes = clFastMemoryCapacity *
+                                                         1024)
       : FunctionPass(&DmaGeneration::passID), slowMemorySpace(slowMemorySpace),
         fastMemorySpace(fastMemorySpace),
         minDmaTransferSize(minDmaTransferSize),
@@ -684,14 +690,6 @@ PassResult DmaGeneration::runOnFunction(Function *f) {
   FuncBuilder topBuilder(f);
   zeroIndex = topBuilder.create<ConstantIndexOp>(f->getLoc(), 0);
 
-  if (clFastMemorySpace.getNumOccurrences() > 0) {
-    fastMemorySpace = clFastMemorySpace;
-  }
-
-  if (clFastMemoryCapacity.getNumOccurrences() > 0) {
-    fastMemCapacityBytes = clFastMemoryCapacity * 1024;
-  }
-
   for (auto &block : *f) {
     runOnBlock(&block, /*consumedCapacityBytes=*/0);
   }