[mlir][sparse][gpu] introduce flag that controls host to device copy strategies (regular dma default)

[K-Wu]

Summary

The zero-copy is disabled in this revision with a tracker https://github.com/llvm/llvm-project/issues/64316

[K-Wu]

disabling zero copy test for the time being

[K-Wu]

recover

[K-Wu]

abc

[K-Wu]

mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.td

126

TODO: unnecessary change:

[K-Wu]

updating

[K-Wu]

upd tests

[K-Wu]

add assertion to disable zero-copy at this moment

[aartbik]

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

55

can we make this a bit more precise, e.g. transfer -> gpuDataTransfer

mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.h

47

add a TODO that ZeroCopy is not supported yet, together with a tracker # that describes the problem so it can be fixed in the future (and then enabled here)

mlir/lib/Dialect/SparseTensor/Transforms/SparseGPUCodegen.cpp

494

this change changes the order in all our tests; can't we keep the original order, and to the !regularDMA part later?
(bit hard to judge, so this is a soft request, okay to keep if needed)

496

keep block together (unnecessary change)

or, start new block with comment if you want to have whitespace
(but given content above, that makes less sense)

497

to avoid the underscore and keep naming consistent, can't we just use

castR
castC
castV
castX
castY

559

keep block together (unnecessary change)

571

no whiteline

621

same, just castR, castC etc.

733

no whiteline

734

castA, castB etc.

874

castB, castA

mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorPasses.cpp

83

add tracker #

mlir/test/Integration/Dialect/SparseTensor/GPU/CUDA/sm80-lt/sparse-matmul-2-4-lib-from-linalg.mlir

14

add tracker #

mlir/test/Integration/Dialect/SparseTensor/GPU/CUDA/sparse-matmul-lib.mlir

17

tracker

mlir/test/Integration/Dialect/SparseTensor/GPU/CUDA/sparse-matvec-lib.mlir

16

tracker

mlir/test/Integration/Dialect/SparseTensor/GPU/CUDA/sparse-sampled-matmul-lib.mlir

17

tracker

[K-Wu]

addressing comments

[K-Wu]

addressing some comments

[K-Wu]

addressing comments

[K-Wu]

reverting order changes and addressing comemnts

[K-Wu]

fix tests

[aartbik]

mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.h

48

GPUDataTransferStrategy
(not necessary Sparse)

57–59

let's name this gpuDataTransferStrategy

mlir/test/Integration/Dialect/SparseTensor/GPU/CUDA/sm80-lt/sparse-matmul-2-4-lib.mlir

13

the other RUN and NOTRUN?

[K-Wu]

remove unnecessary change

[K-Wu]

addressing comments

[K-Wu]

addressing comments

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

PassOptions::Option<mlir::SparseParallelizationStrategy> parallelization{
                      "the outer loop."),
           clEnumValN(mlir::SparseParallelizationStrategy::kDenseAnyLoop,
                      "dense-any-loop",
                      "Enable dense parallelization for any loop."),
           clEnumValN(
               mlir::SparseParallelizationStrategy::kAnyStorageAnyLoop,
               "any-storage-any-loop",
               "Enable sparse parallelization for any storage and loop."))};
+  PassOptions::Option<mlir::SparseDataTransferStrategy> transfer{
+      *this, "data-transfer-strategy",
+      ::llvm::cl::desc(
+          "Set the data transfer strategy between the host and the GPUs"),
+      ::llvm::cl::init(mlir::SparseDataTransferStrategy::kRegularDMA),
+      llvm::cl::values(
+          clEnumValN(mlir::SparseDataTransferStrategy::kRegularDMA,
+                     "regular-dma",
+                     "Default option: malloc on host without additional "
+                     "options or care and then use DMA to copy the data"),
+          clEnumValN(mlir::SparseDataTransferStrategy::kPinnedDMA, "pinned-dma",
+                     "Based on the default option, pin the host memory to "
+                     "accelerate the data transfer"),
+          clEnumValN(mlir::SparseDataTransferStrategy::kZeroCopy, "zero-copy",
+                     "Use zero-copy to perform the data transfer from the host "
+                     "to the GPU"))};
 
   PassOptions::Option<bool> enableIndexReduction{
       *this, "enable-index-reduction",
       desc("Enable dependent index reduction based algorithm to handle "
            "non-trivial index expressions on sparse inputs (experimental "
            "features)"),
       init(false)};
 
struct SparseCompilerOptions
   /// This option is used to enable GPU library generation.
   PassOptions::Option<bool> enableGPULibgen{
       *this, "enable-gpu-libgen",
       desc("Enables GPU acceleration by means of direct library calls (like "
            "cuSPARSE)")};
 
   /// Projects out the options for `createSparsificationPass`.
   SparsificationOptions sparsificationOptions() const {
-    return SparsificationOptions(parallelization, enableIndexReduction,
-                                 enableGPULibgen, enableRuntimeLibrary);
+    return SparsificationOptions(parallelization, transfer,
+                                 enableIndexReduction, enableGPULibgen,
+                                 enableRuntimeLibrary);
   }
 
   /// Projects out the options for `createSparseTensorConversionPass`.
   SparseTensorConversionOptions sparseTensorConversionOptions() const {
     return SparseTensorConversionOptions(
         sparseToSparseConversionStrategy(sparseToSparse));
   }
 

mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.h

enum class SparseParallelizationStrategy {
   kNone,
   kDenseOuterLoop,
   kAnyStorageOuterLoop,
   kDenseAnyLoop,
   kAnyStorageAnyLoop
   // TODO: support reduction parallelization too?
 };
 
+enum class SparseDataTransferStrategy { kRegularDMA, kZeroCopy, kPinnedDMA };
+
 #define GEN_PASS_DECL
 #include "mlir/Dialect/SparseTensor/Transforms/Passes.h.inc"
 
 /// Options for the Sparsification pass.
 struct SparsificationOptions {
-  SparsificationOptions(SparseParallelizationStrategy p, bool idxReduc,
+  SparsificationOptions(SparseParallelizationStrategy p,
+                        SparseDataTransferStrategy t, bool idxReduc,
                         bool gpuLibgen, bool enableRT)
-      : parallelizationStrategy(p), enableIndexReduction(idxReduc),
-        enableGPULibgen(gpuLibgen), enableRuntimeLibrary(enableRT) {}
+      : parallelizationStrategy(p), dataTransferStrategy(t),
+        enableIndexReduction(idxReduc), enableGPULibgen(gpuLibgen),
+        enableRuntimeLibrary(enableRT) {}
   SparsificationOptions()
-      : SparsificationOptions(SparseParallelizationStrategy::kNone, false,
+      : SparsificationOptions(SparseParallelizationStrategy::kNone,
+                              SparseDataTransferStrategy::kRegularDMA, false,
                               false, true) {}
   SparseParallelizationStrategy parallelizationStrategy;
+  SparseDataTransferStrategy dataTransferStrategy;
   bool enableIndexReduction;
   bool enableGPULibgen;
   bool enableRuntimeLibrary;
 };
 
 /// Sets up sparsification rewriting rules with the given options.
 void populateSparsificationPatterns(
     RewritePatternSet &patterns,

 std::unique_ptr<Pass> createSparseVectorizationPass();
 std::unique_ptr<Pass> createSparseVectorizationPass(unsigned vectorLength,
                                                     bool enableVLAVectorization,
                                                     bool enableSIMDIndex32);
 
 void populateSparseGPUCodegenPatterns(RewritePatternSet &patterns,
                                       unsigned numThreads);
 
-void populateSparseGPULibgenPatterns(RewritePatternSet &patterns,
-                                     bool enableRT);
+void populateSparseGPULibgenPatterns(RewritePatternSet &patterns, bool enableRT,
+                                     SparseDataTransferStrategy transfer);
 
 std::unique_ptr<Pass> createSparseGPUCodegenPass();
 std::unique_ptr<Pass> createSparseGPUCodegenPass(unsigned numThreads);
 
 //===----------------------------------------------------------------------===//
 // Registration.
 //===----------------------------------------------------------------------===//
 
 /// Generate the code for registering passes.
 #define GEN_PASS_REGISTRATION
 #include "mlir/Dialect/SparseTensor/Transforms/Passes.h.inc"
 
 } // namespace mlir
 
 #endif // MLIR_DIALECT_SPARSETENSOR_TRANSFORMS_PASSES_H_

mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.td

Option<"parallelization", "parallelization-strategy", "mlir::SparseParallelizationStrategy",
                         "any-storage-outer-loop",
                         "Enable sparse parallelization regardless of storage for the outer loop."),
              clEnumValN(mlir::SparseParallelizationStrategy::kDenseAnyLoop,
                         "dense-any-loop",
                         "Enable dense parallelization for any loop."),
              clEnumValN(mlir::SparseParallelizationStrategy::kAnyStorageAnyLoop,
                         "any-storage-any-loop",
                         "Enable sparse parallelization for any storage and loop."))}]>,
+    Option<"transfer", "data-transfer-strategy", "mlir::SparseDataTransferStrategy",
+            "mlir::SparseDataTransferStrategy::kRegularDMA",
+            "Set the data transfer strategy", [{llvm::cl::values(
+               clEnumValN(mlir::SparseDataTransferStrategy::kRegularDMA,
+                     "regular-dma",
+                     "Default option: malloc on host without additional "
+                     "options or care and then use DMA to copy the data"),
+          clEnumValN(mlir::SparseDataTransferStrategy::kPinnedDMA, "pinned-dma",
+                     "Based on the default option, pin the host memory to "
+                     "accelerate the data transfer"),
+          clEnumValN(mlir::SparseDataTransferStrategy::kZeroCopy, "zero-copy",
+                     "Use zero-copy to perform the data transfer from the host "
+                     "to the GPU"))}]>,
     Option<"enableGPULibgen", "enable-gpu-libgen", "bool",
            "false",
            "Enable GPU acceleration by means of direct library calls (like cuSPARSE)">,
     Option<"enableRuntimeLibrary", "enable-runtime-library", "bool",
            "true", "Enable runtime library for manipulating sparse tensors">,
   ];
 }
 
+
 def PostSparsificationRewrite : Pass<"post-sparsification-rewrite", "ModuleOp"> {
   let summary = "Applies sparse tensor rewriting rules after sparsification";
   let description = [{
     A pass that applies rewriting rules to sparse tensor operations after
     running the actual sparsification pass.
   }];
   let constructor = "mlir::createPostSparsificationRewritePass()";
   let dependentDialects = [

mlir/lib/Dialect/SparseTensor/Transforms/SparseGPUCodegen.cpp

 #endif
   }
   assert(colA);
   return builder.create<gpu::CreateCsrOp>(loc, handleTp, tokenTp, token, sz1,
                                           sz2, nseA, rowA, colA, valA);
 }
 
 /// Match and rewrite SpMV kernel.
-static LogicalResult rewriteSpMV(PatternRewriter &rewriter,
-                                 linalg::GenericOp op, bool enableRT) {
+static LogicalResult
+rewriteSpMV(PatternRewriter &rewriter, linalg::GenericOp op, bool enableRT,
+            SparseDataTransferStrategy dataTransferStrategy) {
   Location loc = op.getLoc();
   Value a = op.getOperand(0);
   Value x = op.getOperand(1);
   Value y = op.getOperand(2); // we have y = Ax
   SmallVector<Value> tokens;
 
+  bool isZeroCopy =
+      dataTransferStrategy == SparseDataTransferStrategy::kZeroCopy;
+
   // Only admissible sparse matrix format and dense vectors.
   bool isCOO = false;
   SparseTensorType aTp = getSparseTensorType(a);
   SparseTensorType xTp = getSparseTensorType(x);
   SparseTensorType yTp = getSparseTensorType(y);
   if (!areAdmissibleTypes(aTp, xTp, yTp, enableRT, /*isMatVec=*/true, isCOO))
     return failure();
 
   // Start sparse kernel and copy data from host to device.
   //   a : memR/memC/memV -> rowA,colA,valA
   //   x : memX           -> vecX
   //   y : memY           -> vecY
   Value nseA = rewriter.create<NumberOfEntriesOp>(loc, a);
   Value szY = linalg::createOrFoldDimOp(rewriter, loc, a, 0);
   Value szX = linalg::createOrFoldDimOp(rewriter, loc, a, 1);
   Value memR = genFirstPosOrCrds(rewriter, loc, a, isCOO, enableRT);
   Value memC = genSecondCrds(rewriter, loc, a, isCOO, enableRT);
   Value memV = genToValues(rewriter, loc, a);
+  Value memX = genTensorToMemref(rewriter, loc, x);
+  Value memY = genTensorToMemref(rewriter, loc, y);
+
+  Value memR_cast, memC_cast, memV_cast, memX_cast, memY_cast;
+  if (dataTransferStrategy != SparseDataTransferStrategy::kRegularDMA) {
+    memR_cast = genHostRegisterMemref(rewriter, loc, memR);
+    if (memC)
+      memC_cast = genHostRegisterMemref(rewriter, loc, memC);
+    memV_cast = genHostRegisterMemref(rewriter, loc, memV);
+    memX_cast = genHostRegisterMemref(rewriter, loc, memX);
+    memY_cast = genHostRegisterMemref(rewriter, loc, memY);
+  }
+
   Value rowA = genAllocCopy(rewriter, loc, memR, tokens);
   Value colA = memC ? genAllocCopy(rewriter, loc, memC, tokens) : Value();
   Value valA = genAllocCopy(rewriter, loc, memV, tokens);
-  Value memX = genTensorToMemref(rewriter, loc, x);
-  Value vecX = genAllocCopy(rewriter, loc, memX, tokens);
-  Value memY = genTensorToMemref(rewriter, loc, y);
+  Value vecX = isZeroCopy ? memX : genAllocCopy(rewriter, loc, memX, tokens);
   Value vecY = genAllocCopy(rewriter, loc, memY, tokens);
   genBlockingWait(rewriter, loc, tokens);
   tokens.clear();
 
   // Create sparse environment and sparse matrix/dense vector handles.
   Type indexTp = rewriter.getIndexType();
   Type dnTensorHandleTp = rewriter.getType<gpu::SparseDnTensorHandleType>();
   Type spmatHandleTp = rewriter.getType<gpu::SparseSpMatHandleType>();
rewriteSpMV(PatternRewriter &rewriter, linalg::GenericOp op, bool enableRT,
 
   // Copy data back to host and free all the resoures.
   token = rewriter.create<gpu::DestroySpMatOp>(loc, tokenTp, token, spMatA)
               .getAsyncToken();
   token = rewriter.create<gpu::DestroyDnTensorOp>(loc, tokenTp, token, dnX)
               .getAsyncToken();
   token = rewriter.create<gpu::DestroyDnTensorOp>(loc, tokenTp, token, dnY)
               .getAsyncToken();
+
   token = genDeallocMemRef(rewriter, loc, rowA, token);
   if (colA)
     token = genDeallocMemRef(rewriter, loc, colA, token);
   token = genDeallocMemRef(rewriter, loc, valA, token);
   token = genDeallocMemRef(rewriter, loc, buffer, token);
+  if (!isZeroCopy)
     token = genDeallocMemRef(rewriter, loc, vecX, token);
   token = genCopyMemRef(rewriter, loc, memY, vecY, token);
   token = genDeallocMemRef(rewriter, loc, vecY, token);
   tokens.push_back(token);
   genBlockingWait(rewriter, loc, tokens);
+
+  if (dataTransferStrategy != SparseDataTransferStrategy::kRegularDMA) {
+    genHostUnregisterMemref(rewriter, loc, memR_cast);
+    if (memC)
+      genHostUnregisterMemref(rewriter, loc, memC_cast);
+    genHostUnregisterMemref(rewriter, loc, memV_cast);
+    genHostUnregisterMemref(rewriter, loc, memX_cast);
+    genHostUnregisterMemref(rewriter, loc, memY_cast);
+  }
   tokens.clear();
 
   // Done.
   rewriter.replaceOpWithNewOp<bufferization::ToTensorOp>(op, memY);
   return success();
 }
 
 /// Match and rewrite SpMM kernel.
-static LogicalResult rewriteSpMM(PatternRewriter &rewriter,
-                                 linalg::GenericOp op, bool enableRT) {
+static LogicalResult
+rewriteSpMM(PatternRewriter &rewriter, linalg::GenericOp op, bool enableRT,
+            SparseDataTransferStrategy dataTransferStrategy) {
   Location loc = op.getLoc();
   Value a = op.getOperand(0);
   Value b = op.getOperand(1);
   Value c = op.getOperand(2); // we have C = AB
   SmallVector<Value> tokens;
 
+  bool isZeroCopy =
+      dataTransferStrategy == SparseDataTransferStrategy::kZeroCopy;
+
   // Only admissible sparse matrix format and dense matrices.
   bool isCOO = false;
   SparseTensorType aTp = getSparseTensorType(a);
   SparseTensorType bTp = getSparseTensorType(b);
   SparseTensorType cTp = getSparseTensorType(c);
   if (!areAdmissibleTypes(aTp, bTp, cTp, enableRT, /*isMatVec=*/false, isCOO))
     return failure();
 
   // Start sparse kernel and copy data from host to device.
   //   a : memR/memC/memV -> rowA,colA,valA
   //   b : bufB           -> matA
   //   c : bufC           -> matC
   Value nseA = rewriter.create<NumberOfEntriesOp>(loc, a);
   Value szm = linalg::createOrFoldDimOp(rewriter, loc, a, 0);
   Value szk = linalg::createOrFoldDimOp(rewriter, loc, a, 1);
   Value szn = linalg::createOrFoldDimOp(rewriter, loc, b, 1);
   Value memR = genFirstPosOrCrds(rewriter, loc, a, isCOO, enableRT);
   Value memC = genSecondCrds(rewriter, loc, a, isCOO, enableRT);
   Value memV = genToValues(rewriter, loc, a);
+  Value bufB = genTensorToMemref(rewriter, loc, b);
+  Value bufC = genTensorToMemref(rewriter, loc, c);
+  Value memR_cast, memC_cast, memV_cast, bufB_cast, bufC_cast;
+  if (dataTransferStrategy != SparseDataTransferStrategy::kRegularDMA) {
+    memR_cast = genHostRegisterMemref(rewriter, loc, memR);
+    if (memC)
+      memC_cast = genHostRegisterMemref(rewriter, loc, memC);
+    memV_cast = genHostRegisterMemref(rewriter, loc, memV);
+    bufB_cast = genHostRegisterMemref(rewriter, loc, bufB);
+    bufC_cast = genHostRegisterMemref(rewriter, loc, bufC);
+  }
+
   Value rowA = genAllocCopy(rewriter, loc, memR, tokens);
   Value colA = memC ? genAllocCopy(rewriter, loc, memC, tokens) : Value();
   Value valA = genAllocCopy(rewriter, loc, memV, tokens);
-  Value bufB = genTensorToMemref(rewriter, loc, b);
-  Value matB = genAllocCopy(rewriter, loc, bufB, tokens);
-  Value bufC = genTensorToMemref(rewriter, loc, c);
+  Value matB = isZeroCopy ? bufB : genAllocCopy(rewriter, loc, bufB, tokens);
   Value matC = genAllocCopy(rewriter, loc, bufC, tokens);
   genBlockingWait(rewriter, loc, tokens);
   tokens.clear();
 
   // Create sparse environment and sparse matrix/dense matrix handles.
   Type indexTp = rewriter.getIndexType();
   Type dnTensorHandleTp = rewriter.getType<gpu::SparseDnTensorHandleType>();
   Type spMatHandleTp = rewriter.getType<gpu::SparseSpMatHandleType>();
token = rewriter.create<gpu::DestroyDnTensorOp>(loc, tokenTp, token, dnB)
               .getAsyncToken();
   token = rewriter.create<gpu::DestroyDnTensorOp>(loc, tokenTp, token, dnC)
               .getAsyncToken();
   token = genDeallocMemRef(rewriter, loc, rowA, token);
   if (colA)
     token = genDeallocMemRef(rewriter, loc, colA, token);
   token = genDeallocMemRef(rewriter, loc, valA, token);
   token = genDeallocMemRef(rewriter, loc, buffer, token);
+  if (!isZeroCopy)
     token = genDeallocMemRef(rewriter, loc, matB, token);
   token = genCopyMemRef(rewriter, loc, bufC, matC, token);
   token = genDeallocMemRef(rewriter, loc, matC, token);
   tokens.push_back(token);
   genBlockingWait(rewriter, loc, tokens);
+  if (dataTransferStrategy != SparseDataTransferStrategy::kRegularDMA) {
+    genHostUnregisterMemref(rewriter, loc, memR_cast);
+    if (memC)
+      genHostUnregisterMemref(rewriter, loc, memC_cast);
+    genHostUnregisterMemref(rewriter, loc, memV_cast);
+    genHostUnregisterMemref(rewriter, loc, bufB_cast);
+    genHostUnregisterMemref(rewriter, loc, bufC_cast);
+  }
   tokens.clear();
 
   // Done.
   rewriter.replaceOpWithNewOp<bufferization::ToTensorOp>(op, bufC);
   return success();
 }
 
 // Match and rewrite 2:4 SpMM kernels.
-static LogicalResult rewrite2To4SpMM(PatternRewriter &rewriter,
-                                     linalg::GenericOp op) {
+static LogicalResult
+rewrite2To4SpMM(PatternRewriter &rewriter, linalg::GenericOp op,
+                SparseDataTransferStrategy dataTransferStrategy) {
   Location loc = op.getLoc();
   Value A = op.getOperand(0);
   Value B = op.getOperand(1);
   Value C = op.getOperand(2); // we have C = AB
   SmallVector<Value> tokens;
 
+  bool isZeroCopy =
+      dataTransferStrategy == SparseDataTransferStrategy::kZeroCopy;
+
   // All input should be dense tensors.
   if (!isDenseTensor(A) || !isDenseTensor(B) || !isDenseTensor(C))
     return failure();
 
   Value bufA = genTensorToMemref(rewriter, loc, A);
-  Value matA = genAllocCopy(rewriter, loc, bufA, tokens);
   Value bufB = genTensorToMemref(rewriter, loc, B);
-  Value matB = genAllocCopy(rewriter, loc, bufB, tokens);
   Value bufC = genTensorToMemref(rewriter, loc, C);
+
+  Value bufA_cast, bufB_cast, bufC_cast;
+  if (dataTransferStrategy != SparseDataTransferStrategy::kRegularDMA) {
+    bufA_cast = genHostRegisterMemref(rewriter, loc, bufA);
+    bufB_cast = genHostRegisterMemref(rewriter, loc, bufB);
+    bufC_cast = genHostRegisterMemref(rewriter, loc, bufC);
+  }
+
+  Value matA = isZeroCopy ? bufA : genAllocCopy(rewriter, loc, bufA, tokens);
+  Value matB = isZeroCopy ? bufB : genAllocCopy(rewriter, loc, bufB, tokens);
   Value matC = genAllocCopy(rewriter, loc, bufC, tokens);
   genBlockingWait(rewriter, loc, tokens);
   tokens.clear();
   Value szm = linalg::createOrFoldDimOp(rewriter, loc, matA, 0);
   Value szk = linalg::createOrFoldDimOp(rewriter, loc, matB, 0);
   Value szn = linalg::createOrFoldDimOp(rewriter, loc, matC, 1);
 
   Type indexTp = rewriter.getIndexType();
token = rewriter.create<gpu::DestroyDnTensorOp>(loc, tokenTp, token, dnB)
               .getAsyncToken();
   token = rewriter.create<gpu::DestroyDnTensorOp>(loc, tokenTp, token, dnC)
               .getAsyncToken();
   SmallVector<Value> newDynamicSizes;
 
   token = genDeallocMemRef(rewriter, loc, buffer, token);
   token = genDeallocMemRef(rewriter, loc, buffer2, token);
   token = genDeallocMemRef(rewriter, loc, buffer3, token);
+
+  if (!isZeroCopy)
     token = genDeallocMemRef(rewriter, loc, matA, token);
+  if (!isZeroCopy)
     token = genDeallocMemRef(rewriter, loc, matB, token);
   token = genCopyMemRef(rewriter, loc, bufC, matC, token);
   token = genDeallocMemRef(rewriter, loc, matC, token);
   tokens.push_back(token);
   genBlockingWait(rewriter, loc, tokens);
+  if (dataTransferStrategy != SparseDataTransferStrategy::kRegularDMA) {
+    genHostUnregisterMemref(rewriter, loc, bufA_cast);
+    genHostUnregisterMemref(rewriter, loc, bufB_cast);
+    genHostUnregisterMemref(rewriter, loc, bufC_cast);
+  }
   tokens.clear();
   rewriter.replaceOpWithNewOp<bufferization::ToTensorOp>(op, bufC);
   return success();
 }
 
 /// Match and rewrite SDDMM kernel.
-static LogicalResult rewriteSDDMM(PatternRewriter &rewriter,
-                                  linalg::GenericOp op, bool enableRT) {
+static LogicalResult
+rewriteSDDMM(PatternRewriter &rewriter, linalg::GenericOp op, bool enableRT,
+             SparseDataTransferStrategy dataTransferStrategy) {
   Location loc = op.getLoc();
   Value a = op.getOperand(0);
   Value b = op.getOperand(1);
   Value c = op.getOperand(2);
   SmallVector<Value> tokens;
 
+  bool isZeroCopy =
+      dataTransferStrategy == SparseDataTransferStrategy::kZeroCopy;
+
   // Only admissible sparse matrix format and dense matrices, no COO.
   bool isCOO = false;
   SparseTensorType aTp = getSparseTensorType(a);
   SparseTensorType bTp = getSparseTensorType(b);
   SparseTensorType cTp = getSparseTensorType(c);
   if (!areAdmissibleTypes(cTp, bTp, aTp, enableRT, false, isCOO))
     return failure();
   if (isCOO)
     return failure();
 
   // The SDDMM does the in-place operation.
   // Start sparse kernel and copy data from host to device.
   //   a : bufA           -> matA
   //   b : bufB           -> matA
   //   c : memR/memC/memV -> rowC,colC,valC
   Value nseC = rewriter.create<NumberOfEntriesOp>(loc, c);
   Value szm = linalg::createOrFoldDimOp(rewriter, loc, a, 0);
   Value szk = linalg::createOrFoldDimOp(rewriter, loc, a, 1);
   Value szn = linalg::createOrFoldDimOp(rewriter, loc, b, 1);
   Value bufA = genTensorToMemref(rewriter, loc, a);
-  Value matA = genAllocCopy(rewriter, loc, bufA, tokens);
   Value bufB = genTensorToMemref(rewriter, loc, b);
-  Value matB = genAllocCopy(rewriter, loc, bufB, tokens);
   Value memR = genFirstPosOrCrds(rewriter, loc, c, isCOO, enableRT);
   Value memC = genSecondCrds(rewriter, loc, c, isCOO, enableRT);
   Value memV = genToValues(rewriter, loc, c);
+
+  Value bufB_cast, bufA_cast, memR_cast, memC_cast, memV_cast;
+  if (dataTransferStrategy != SparseDataTransferStrategy::kRegularDMA) {
+    bufB_cast = genHostRegisterMemref(rewriter, loc, bufB);
+    bufA_cast = genHostRegisterMemref(rewriter, loc, bufA);
+    memR_cast = genHostRegisterMemref(rewriter, loc, memR);
+    if (memC)
+      memC_cast = genHostRegisterMemref(rewriter, loc, memC);
+    memV_cast = genHostRegisterMemref(rewriter, loc, memV);
+  }
+
+  Value matA = isZeroCopy ? bufA : genAllocCopy(rewriter, loc, bufA, tokens);
+  Value matB = isZeroCopy ? bufB : genAllocCopy(rewriter, loc, bufB, tokens);
   Value rowC = genAllocCopy(rewriter, loc, memR, tokens);
   Value colC = memC ? genAllocCopy(rewriter, loc, memC, tokens) : Value();
   Value valC = genAllocCopy(rewriter, loc, memV, tokens);
   genBlockingWait(rewriter, loc, tokens);
   tokens.clear();
 
   // Create sparse environment and sparse matrix/dense matrix handles.
   Type indexTp = rewriter.getIndexType();
rewriteSDDMM(PatternRewriter &rewriter, linalg::GenericOp op, bool enableRT,
   // Copy data back to host and free all the resoures.
   token = rewriter.create<gpu::DestroyDnTensorOp>(loc, tokenTp, token, dnA)
               .getAsyncToken();
   token = rewriter.create<gpu::DestroyDnTensorOp>(loc, tokenTp, token, dnB)
               .getAsyncToken();
   token = rewriter.create<gpu::DestroySpMatOp>(loc, tokenTp, token, spMatC)
               .getAsyncToken();
   token = genDeallocMemRef(rewriter, loc, buffer, token);
+  if (!isZeroCopy) {
     token = genDeallocMemRef(rewriter, loc, matA, token);
     token = genDeallocMemRef(rewriter, loc, matB, token);
+  }
   token = genDeallocMemRef(rewriter, loc, rowC, token);
   if (colC)
     token = genDeallocMemRef(rewriter, loc, colC, token);
   token = genCopyMemRef(rewriter, loc, memV, valC, token);
   token = genDeallocMemRef(rewriter, loc, valC, token);
   tokens.push_back(token);
   genBlockingWait(rewriter, loc, tokens);
+  if (dataTransferStrategy != SparseDataTransferStrategy::kRegularDMA) {
+    genHostUnregisterMemref(rewriter, loc, bufB_cast);
+    genHostUnregisterMemref(rewriter, loc, bufA_cast);
+    genHostUnregisterMemref(rewriter, loc, memR_cast);
+    if (memC)
+      genHostUnregisterMemref(rewriter, loc, memC_cast);
+    genHostUnregisterMemref(rewriter, loc, memV_cast);
+  }
   tokens.clear();
 
   // Done.
   rewriter.replaceOpWithNewOp<sparse_tensor::LoadOp>(op, c);
   return success();
 }
 
 //===----------------------------------------------------------------------===//

 // Rewriting rules for library recognition and code generation.
 //===----------------------------------------------------------------------===//
 
 /// Proof-of-concept rewriter. This rule recognizes certain math kernels
 /// and replaces these with corresponding calls into a sparse library.
 struct LinalgOpRewriter : public OpRewritePattern<linalg::GenericOp> {
   using OpRewritePattern<linalg::GenericOp>::OpRewritePattern;
 
-  LinalgOpRewriter(MLIRContext *context, bool rt)
-      : OpRewritePattern(context), enableRT(rt) {}
+  LinalgOpRewriter(MLIRContext *context, bool rt, SparseDataTransferStrategy t)
+      : OpRewritePattern(context), enableRT(rt), dataTransferStrategy(t) {}
 
   LogicalResult matchAndRewrite(linalg::GenericOp op,
                                 PatternRewriter &rewriter) const override {
     if (op.getNumDpsInits() != 1)
       return failure(); // reject multi-output
 
     const unsigned numLoops = op.getNumLoops();
     const unsigned numTensors = op->getNumOperands();
LogicalResult matchAndRewrite(linalg::GenericOp op,
     // TODO: identify alpha and beta and pass them to the CUDA calls
 
     // Recognize a SpMV kernel.
     if (numLoops == 2 && numTensors == 3 &&
         linalg::isParallelIterator(iteratorTypes[0]) &&
         linalg::isReductionIterator(iteratorTypes[1]) &&
         // TODO: add transposed {i, j}
         maps == infer({{i, j}, {j}, {i}}) && matchSumOfMultOfArgs(op)) {
-      return rewriteSpMV(rewriter, op, enableRT);
+      return rewriteSpMV(rewriter, op, enableRT, dataTransferStrategy);
     }
 
     // Recognize a SpMM kernel.
     if (numLoops == 3 && numTensors == 3 &&
         linalg::isParallelIterator(iteratorTypes[0]) &&
         linalg::isParallelIterator(iteratorTypes[1]) &&
         linalg::isReductionIterator(iteratorTypes[2]) &&
         // TODO: add transposed {i, k}, {k, j}
         // TODO: maybe add transposed {i, j} in future
         maps == infer({{i, k}, {k, j}, {i, j}}) && matchSumOfMultOfArgs(op)) {
       if (op->getAttr("DENSE24"))
-        return rewrite2To4SpMM(rewriter, op);
+        return rewrite2To4SpMM(rewriter, op, dataTransferStrategy);
 
-      return rewriteSpMM(rewriter, op, enableRT);
+      return rewriteSpMM(rewriter, op, enableRT, dataTransferStrategy);
     }
 
     // Recognize a SDDMM kernel.
     if (numLoops == 3 && numTensors == 3 &&
         linalg::isParallelIterator(iteratorTypes[0]) &&
         linalg::isParallelIterator(iteratorTypes[1]) &&
         linalg::isReductionIterator(iteratorTypes[2]) &&
         // TODO: add transposed {i, k}, {k, j}
         // TODO: maybe add transposed {i, j} in future
         maps == infer({{i, k}, {k, j}, {i, j}}) &&
         matchSumReductionOfMulUnary(op)) {
-      return rewriteSDDMM(rewriter, op, enableRT);
+      return rewriteSDDMM(rewriter, op, enableRT, dataTransferStrategy);
     }
 
     return failure();
   }
 
 private:
   bool enableRT;
+  SparseDataTransferStrategy dataTransferStrategy;
 };
 
 } // namespace
 
 //===----------------------------------------------------------------------===//
 // Public method for populating GPU rewriting rules.
 //
 // Currently two set of rewriting rules are made available. The first set
 // implements direct code generation, currently by means of convering the
 // outermost paralell loop into GPU threads. The second set implements
 // libary recognition of a set of sparse operations. Eventually, the right
 // combination of these two approaches has to be found.
 //===----------------------------------------------------------------------===//
 
 void mlir::populateSparseGPUCodegenPatterns(RewritePatternSet &patterns,
                                             unsigned numThreads) {
   patterns.add<ForallRewriter>(patterns.getContext(), numThreads);
 }
 
-void mlir::populateSparseGPULibgenPatterns(RewritePatternSet &patterns,
-                                           bool enableRT) {
-  patterns.add<LinalgOpRewriter>(patterns.getContext(), enableRT);
+void mlir::populateSparseGPULibgenPatterns(
+    RewritePatternSet &patterns, bool enableRT,
+    SparseDataTransferStrategy transfer) {
+  patterns.add<LinalgOpRewriter>(patterns.getContext(), enableRT, transfer);
 }

mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorPasses.cpp

 
 struct SparsificationPass
     : public impl::SparsificationPassBase<SparsificationPass> {
 
   SparsificationPass() = default;
   SparsificationPass(const SparsificationPass &pass) = default;
   SparsificationPass(const SparsificationOptions &options) {
     parallelization = options.parallelizationStrategy;
+    transfer = options.dataTransferStrategy;
     enableIndexReduction = options.enableIndexReduction;
     enableGPULibgen = options.enableGPULibgen;
     enableRuntimeLibrary = options.enableRuntimeLibrary;
   }
 
   void runOnOperation() override {
     auto *ctx = &getContext();
     // Translate strategy flags to strategy options.
-    SparsificationOptions options(parallelization, enableIndexReduction,
-                                  enableGPULibgen, enableRuntimeLibrary);
+    SparsificationOptions options(parallelization, transfer,
+                                  enableIndexReduction, enableGPULibgen,
+                                  enableRuntimeLibrary);
     // Apply GPU libgen (if requested), sparsification, and cleanup rewriting.
     RewritePatternSet patterns(ctx);
     if (enableGPULibgen) {
-      populateSparseGPULibgenPatterns(patterns, enableRuntimeLibrary);
+      populateSparseGPULibgenPatterns(patterns, enableRuntimeLibrary, transfer);
     }
     populateSparsificationPatterns(patterns, options);
     scf::ForOp::getCanonicalizationPatterns(patterns, ctx);
     (void)applyPatternsAndFoldGreedily(getOperation(), std::move(patterns));
   }
 };
 
 struct PostSparsificationRewritePass

mlir/test/Integration/Dialect/SparseTensor/GPU/CUDA/sm80-lt/sparse-matmul-2-4-lib-from-linalg.mlir

 //
 // NOTE: this test requires gpu-sm80 and cusparselt
 //
-// RUN: mlir-opt %s \
-// RUN: --sparse-compiler="enable-runtime-library=true enable-gpu-libgen gpu-triple=nvptx64-nvidia-cuda gpu-chip=sm_80 gpu-features=+ptx71" \
-// RUN: | mlir-cpu-runner \
-// RUN:   --shared-libs=%mlir_cuda_runtime \
-// RUN:   --shared-libs=%mlir_c_runner_utils \
-// RUN:   --e main --entry-point-result=void \
-// RUN: | FileCheck %s
+// DEFINE: %{compile} = mlir-opt %s \
+// DEFINE: --sparse-compiler="enable-runtime-library=true enable-gpu-libgen gpu-triple=nvptx64-nvidia-cuda gpu-chip=sm_80 gpu-features=+ptx71
+// DEFINE: %{run} = mlir-cpu-runner \
+// DEFINE: --shared-libs=%mlir_cuda_runtime \
+// DEFINE: --shared-libs=%mlir_c_runner_utils \
+// DEFINE: --e main --entry-point-result=void \
+// DEFINE: | FileCheck %s
+
+//  RUN:  %{compile}" | %{run}
+//  RUN:  %{compile} data-transfer-strategy=pinned-dma" | %{run}
+//   %{compile} data-transfer-strategy=zero-copy" | %{run}
 
 #map = affine_map<(d0, d1, d2) -> (d0, d2)>
 #map1 = affine_map<(d0, d1, d2) -> (d2, d1)>
 #map2 = affine_map<(d0, d1, d2) -> (d0, d1)>
 
 module {
   llvm.func @mgpuCreateSparseLtEnv()
   llvm.func @mgpuDestroySparseLtEnv()

mlir/test/Integration/Dialect/SparseTensor/GPU/CUDA/sm80-lt/sparse-matmul-2-4-lib.mlir

 //
 // NOTE: this test requires gpu-sm80 and cusparselt
 //
-// RUN: mlir-opt --convert-scf-to-cf -convert-cf-to-llvm --convert-vector-to-llvm \
-// RUN:          --convert-arith-to-llvm --gpu-to-llvm --reconcile-unrealized-casts \
-// RUN:          %s \
-// RUN: | mlir-cpu-runner \
-// RUN:   --shared-libs=%mlir_cuda_runtime \
-// RUN:   --shared-libs=%mlir_c_runner_utils \
-// RUN:   --e main --entry-point-result=void \
-// RUN: | FileCheck %s
+// DEFINE: %{compile} = mlir-opt --convert-scf-to-cf -convert-cf-to-llvm --convert-vector-to-llvm \
+// DEFINE: --convert-arith-to-llvm --gpu-to-llvm --reconcile-unrealized-casts \
+// DEFINE: %s 
+// DEFINE: %{run} = mlir-cpu-runner \
+// DEFINE:   --shared-libs=%mlir_cuda_runtime \
+// DEFINE:   --shared-libs=%mlir_c_runner_utils \
+// DEFINE:   --e main --entry-point-result=void \
+// DEFINE: | FileCheck %s
+
+// RUN: %{compile} | %{run}
 
 module {
   llvm.func @mgpuCreateSparseLtEnv()
   llvm.func @mgpuDestroySparseLtEnv()
 
   func.func @sampled_matmul(%a : memref<16x32xf16>,
                             %b : memref<32x16xf16>,
                             %c : memref<16x16xf16>) {

mlir/test/Integration/Dialect/SparseTensor/GPU/CUDA/sparse-matmul-lib.mlir

 //
 // NOTE: this test requires gpu-sm80
 //
+// DEFINE: %{compile} = mlir-opt %s \
+// DEFINE:    --sparse-compiler="enable-gpu-libgen gpu-triple=nvptx64-nvidia-cuda gpu-chip=sm_80 gpu-features=+ptx71
+// DEFINE: %{run} = mlir-cpu-runner \
+// DEFINE:   --shared-libs=%mlir_cuda_runtime \
+// DEFINE:   --shared-libs=%mlir_c_runner_utils \
+// DEFINE:   --e main --entry-point-result=void \
+// DEFINE: | FileCheck %s
+//
+//
 // with RT lib (SoA COO):
 //
-// RUN: mlir-opt %s \
-// RUN:   --sparse-compiler="enable-runtime-library=true enable-gpu-libgen gpu-triple=nvptx64-nvidia-cuda gpu-chip=sm_80 gpu-features=+ptx71"  \
-// RUN: | mlir-cpu-runner \
-// RUN:   --shared-libs=%mlir_cuda_runtime \
-// RUN:   --shared-libs=%mlir_c_runner_utils \
-// RUN:   --e main --entry-point-result=void \
-// RUN: | FileCheck %s
+// RUN:  %{compile} enable-runtime-library=true" | %{run}
+// RUN:  %{compile} enable-runtime-library=true data-transfer-strategy=pinned-dma" | %{run}
+//  %{compile} enable-runtime-library=true data-transfer-strategy=zero-copy" | %{run}
 //
 // without RT lib (AoS COO): note, may fall back to CPU
 //
-// RUN: mlir-opt %s \
-// RUN:   --sparse-compiler="enable-runtime-library=false enable-gpu-libgen gpu-triple=nvptx64-nvidia-cuda gpu-chip=sm_80 gpu-features=+ptx71"  \
-// RUN: | mlir-cpu-runner \
-// RUN:   --shared-libs=%mlir_cuda_runtime \
-// RUN:   --shared-libs=%mlir_c_runner_utils \
-// RUN:   --e main --entry-point-result=void \
-// RUN: | FileCheck %s
+// RUN: %{compile} enable-runtime-library=false" | %{run}
+// RUN: %{compile} enable-runtime-library=false data-transfer-strategy=pinned-dma" | %{run}
+//  %{compile} enable-runtime-library=false data-transfer-strategy=zero-copy" | %{run}
 
 #SortedCOO = #sparse_tensor.encoding<{
   lvlTypes = [ "compressed-nu", "singleton" ]
 }>
 
 #CSR = #sparse_tensor.encoding<{
   lvlTypes = [ "dense", "compressed" ],
   posWidth = 32,

mlir/test/Integration/Dialect/SparseTensor/GPU/CUDA/sparse-matvec-lib.mlir

 //
 // NOTE: this test requires gpu-sm80
 //
+// DEFINE: %{compile} = mlir-opt %s \
+// DEFINE:   --sparse-compiler="enable-gpu-libgen gpu-triple=nvptx64-nvidia-cuda gpu-chip=sm_80 gpu-features=+ptx71
+// DEFINE: %{run} = mlir-cpu-runner \
+// DEFINE:   --shared-libs=%mlir_cuda_runtime \
+// DEFINE:   --shared-libs=%mlir_c_runner_utils \
+// DEFINE:   --e main --entry-point-result=void \
+// DEFINE: | FileCheck %s
+//
 // with RT lib (SoA COO):
 //
-// RUN: mlir-opt %s \
-// RUN:   --sparse-compiler="enable-runtime-library=true enable-gpu-libgen gpu-triple=nvptx64-nvidia-cuda gpu-chip=sm_80 gpu-features=+ptx71"  \
-// RUN: | mlir-cpu-runner \
-// RUN:   --shared-libs=%mlir_cuda_runtime \
-// RUN:   --shared-libs=%mlir_c_runner_utils \
-// RUN:   --e main --entry-point-result=void \
-// RUN: | FileCheck %s
+// RUN: %{compile} enable-runtime-library=true"  | %{run}
+// RUN: %{compile} enable-runtime-library=true data-transfer-strategy=pinned-dma" | %{run}
+// %{compile} enable-runtime-library=true data-transfer-strategy=zero-copy"  | %{run}
 //
 // without RT lib (AoS COO): note, may fall back to CPU
 //
-// RUN: mlir-opt %s \
-// RUN:   --sparse-compiler="enable-runtime-library=false enable-gpu-libgen gpu-triple=nvptx64-nvidia-cuda gpu-chip=sm_80 gpu-features=+ptx71"  \
-// RUN: | mlir-cpu-runner \
-// RUN:   --shared-libs=%mlir_cuda_runtime \
-// RUN:   --shared-libs=%mlir_c_runner_utils \
-// RUN:   --e main --entry-point-result=void \
-// RUN: | FileCheck %s
+// RUN: %{compile} enable-runtime-library=false"  | %{run}
+// RUN: %{compile} enable-runtime-library=false data-transfer-strategy=pinned-dma" | %{run}
+// %{compile} enable-runtime-library=false data-transfer-strategy=zero-copy"  | %{run}
+//
 
 #SortedCOO = #sparse_tensor.encoding<{
   lvlTypes = [ "compressed-nu", "singleton" ]
 }>
 
 #CSR = #sparse_tensor.encoding<{
   lvlTypes = [ "dense", "compressed" ],
   posWidth = 32,

mlir/test/Integration/Dialect/SparseTensor/GPU/CUDA/sparse-sampled-matmul-lib.mlir

 //
 // NOTE: this test requires gpu-sm80
 //
+// DEFINE: %{compile} = mlir-opt %s \
+// DEFINE:   --sparse-compiler="enable-gpu-libgen gpu-triple=nvptx64-nvidia-cuda gpu-chip=sm_80 gpu-features=+ptx71
+// DEFINE: %{run} = TENSOR0="%mlir_src_dir/test/Integration/data/test.mtx" \
+// DEFINE:   mlir-cpu-runner \
+// DEFINE:   --shared-libs=%mlir_cuda_runtime \
+// DEFINE:   --shared-libs=%mlir_c_runner_utils \
+// DEFINE:   --e entry --entry-point-result=void \
+// DEFINE: | FileCheck %s
+//
 // with RT lib:
 //
-// RUN: mlir-opt %s \
-// RUN:   --sparse-compiler="enable-runtime-library=true enable-gpu-libgen gpu-triple=nvptx64-nvidia-cuda gpu-chip=sm_80 gpu-features=+ptx71" \
-// RUN: | TENSOR0="%mlir_src_dir/test/Integration/data/test.mtx" \
-// RUN:   mlir-cpu-runner \
-// RUN:   --shared-libs=%mlir_cuda_runtime \
-// RUN:   --shared-libs=%mlir_c_runner_utils \
-// RUN:   --e entry --entry-point-result=void \
-// RUN: | FileCheck %s
+//  RUN:  %{compile} enable-runtime-library=true" | %{run}
+//  RUN:  %{compile} enable-runtime-library=true data-transfer-strategy=pinned-dma" | %{run}
+//   %{compile} enable-runtime-library=true data-transfer-strategy=zero-copy" | %{run}
 //
 // without RT lib:
 //
-// RUN: mlir-opt %s \
-// RUN:   --sparse-compiler="enable-runtime-library=false enable-gpu-libgen gpu-triple=nvptx64-nvidia-cuda gpu-chip=sm_80 gpu-features=+ptx71" \
-// RUN: | TENSOR0="%mlir_src_dir/test/Integration/data/test.mtx" \
-// RUN:   mlir-cpu-runner \
-// RUN:   --shared-libs=%mlir_cuda_runtime \
-// RUN:   --shared-libs=%mlir_c_runner_utils \
-// RUN:   --e entry --entry-point-result=void \
-// RUN: | FileCheck %s
-//
+// RUN:  %{compile} enable-runtime-library=false" | %{run}
+// RUN:  %{compile} enable-runtime-library=false data-transfer-strategy=pinned-dma" | %{run}
+// %{compile} enable-runtime-library=false data-transfer-strategy=zero-copy" | %{run}
+// 
 
 !Filename = !llvm.ptr<i8>
 
 #CSR = #sparse_tensor.encoding<{
   lvlTypes = ["dense", "compressed"]
 }>
 
 #trait_sampled_dense_dense = {