disable zero-copy for the time being as it causes wrong result

[K-Wu]

abc

[aartbik]

The title of this revision is a bit misleading. The revision introduces three alternatives for copying data from host, where the dma method currently used is the default
It introduces two other ways, however. As written, the current title implies we disable something that is already there.

Also, needs tags.

I propose

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

[aartbik]

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

126

unnecessary change

[K-Wu]

WIP in D155352. This diff is submitted in error. I will incorporate your advice .SOrry for the confusion.

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);
+  if (dataTransferStrategy != SparseDataTransferStrategy::kRegularDMA)
+    genHostRegisterMemref(rewriter, loc, memR);
   Value memC = genSecondCrds(rewriter, loc, a, isCOO, enableRT);
+  if (memC && dataTransferStrategy != SparseDataTransferStrategy::kRegularDMA)
+    genHostRegisterMemref(rewriter, loc, memC);
   Value memV = genToValues(rewriter, loc, a);
-  Value rowA = genAllocCopy(rewriter, loc, memR, tokens);
-  Value colA = memC ? genAllocCopy(rewriter, loc, memC, tokens) : Value();
-  Value valA = genAllocCopy(rewriter, loc, memV, tokens);
+  if (dataTransferStrategy != SparseDataTransferStrategy::kRegularDMA)
+    genHostRegisterMemref(rewriter, loc, memV);
+  Value rowA = isZeroCopy ? memR : genAllocCopy(rewriter, loc, memR, tokens);
+  Value colA =
+      memC ? (isZeroCopy ? memC : genAllocCopy(rewriter, loc, memC, tokens))
+           : Value();
+  Value valA = isZeroCopy ? memV : genAllocCopy(rewriter, loc, memV, tokens);
   Value memX = genTensorToMemref(rewriter, loc, x);
-  Value vecX = genAllocCopy(rewriter, loc, memX, tokens);
+  if (dataTransferStrategy != SparseDataTransferStrategy::kRegularDMA)
+    genHostRegisterMemref(rewriter, loc, memX);
+  Value vecX = isZeroCopy ? memX : genAllocCopy(rewriter, loc, memX, tokens);
   Value memY = genTensorToMemref(rewriter, loc, y);
-  Value vecY = genAllocCopy(rewriter, loc, memY, tokens);
+  if (dataTransferStrategy != SparseDataTransferStrategy::kRegularDMA)
+    genHostRegisterMemref(rewriter, loc, memY);
+  Value vecY = isZeroCopy ? memY : 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>();
   Type tokenTp = rewriter.getType<gpu::AsyncTokenType>();
rewriteSpMV(PatternRewriter &rewriter, linalg::GenericOp op, bool enableRT,
   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);
+  if (dataTransferStrategy != SparseDataTransferStrategy::kRegularDMA)
+    genHostRegisterMemref(rewriter, loc, memR);
   Value memC = genSecondCrds(rewriter, loc, a, isCOO, enableRT);
+  if (memC && dataTransferStrategy != SparseDataTransferStrategy::kRegularDMA)
+    genHostRegisterMemref(rewriter, loc, memC);
   Value memV = genToValues(rewriter, loc, a);
-  Value rowA = genAllocCopy(rewriter, loc, memR, tokens);
-  Value colA = memC ? genAllocCopy(rewriter, loc, memC, tokens) : Value();
-  Value valA = genAllocCopy(rewriter, loc, memV, tokens);
+  if (dataTransferStrategy != SparseDataTransferStrategy::kRegularDMA)
+    genHostRegisterMemref(rewriter, loc, memV);
+  Value rowA = isZeroCopy ? memR : genAllocCopy(rewriter, loc, memR, tokens);
+  Value colA =
+      memC ? (isZeroCopy ? memC : genAllocCopy(rewriter, loc, memC, tokens))
+           : Value();
+  Value valA = isZeroCopy ? memV : genAllocCopy(rewriter, loc, memV, tokens);
   Value bufB = genTensorToMemref(rewriter, loc, b);
-  Value matB = genAllocCopy(rewriter, loc, bufB, tokens);
+  if (dataTransferStrategy != SparseDataTransferStrategy::kRegularDMA)
+    genHostRegisterMemref(rewriter, loc, bufB);
+  Value matB = isZeroCopy ? bufB : genAllocCopy(rewriter, loc, bufB, tokens);
   Value bufC = genTensorToMemref(rewriter, loc, c);
-  Value matC = genAllocCopy(rewriter, loc, bufC, tokens);
+  if (dataTransferStrategy != SparseDataTransferStrategy::kRegularDMA)
+    genHostRegisterMemref(rewriter, loc, bufC);
+  Value matC = isZeroCopy ? bufC : 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>();
   Type tokenTp = rewriter.getType<gpu::AsyncTokenType>();
rewriteSpMM(PatternRewriter &rewriter, linalg::GenericOp op, bool enableRT,
   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);
+  if (dataTransferStrategy != SparseDataTransferStrategy::kRegularDMA)
+    genHostRegisterMemref(rewriter, loc, bufA);
+  Value matA = isZeroCopy ? bufA : genAllocCopy(rewriter, loc, bufA, tokens);
   Value bufB = genTensorToMemref(rewriter, loc, B);
-  Value matB = genAllocCopy(rewriter, loc, bufB, tokens);
+  if (dataTransferStrategy != SparseDataTransferStrategy::kRegularDMA)
+    genHostRegisterMemref(rewriter, loc, bufB);
+  Value matB = isZeroCopy ? bufB : genAllocCopy(rewriter, loc, bufB, tokens);
   Value bufC = genTensorToMemref(rewriter, loc, C);
-  Value matC = genAllocCopy(rewriter, loc, bufC, tokens);
+  if (dataTransferStrategy != SparseDataTransferStrategy::kRegularDMA)
+    genHostRegisterMemref(rewriter, loc, bufC);
+  Value matC = isZeroCopy ? bufC : 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();
   Type dnTensorHandleTp = rewriter.getType<gpu::SparseDnTensorHandleType>();
rewrite2To4SpMM(PatternRewriter &rewriter, linalg::GenericOp op,
   tokens.push_back(token);
   genBlockingWait(rewriter, loc, tokens);
   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);
+  if (dataTransferStrategy != SparseDataTransferStrategy::kRegularDMA)
+    genHostRegisterMemref(rewriter, loc, bufA);
+  Value matA = isZeroCopy ? bufA : genAllocCopy(rewriter, loc, bufA, tokens);
   Value bufB = genTensorToMemref(rewriter, loc, b);
-  Value matB = genAllocCopy(rewriter, loc, bufB, tokens);
+  if (dataTransferStrategy != SparseDataTransferStrategy::kRegularDMA)
+    genHostRegisterMemref(rewriter, loc, bufB);
+  Value matB = isZeroCopy ? bufB : genAllocCopy(rewriter, loc, bufB, tokens);
   Value memR = genFirstPosOrCrds(rewriter, loc, c, isCOO, enableRT);
+  if (dataTransferStrategy != SparseDataTransferStrategy::kRegularDMA)
+    genHostRegisterMemref(rewriter, loc, memR);
   Value memC = genSecondCrds(rewriter, loc, c, isCOO, enableRT);
+  if (memC && dataTransferStrategy != SparseDataTransferStrategy::kRegularDMA)
+    genHostRegisterMemref(rewriter, loc, memC);
   Value memV = genToValues(rewriter, loc, c);
-  Value rowC = genAllocCopy(rewriter, loc, memR, tokens);
-  Value colC = memC ? genAllocCopy(rewriter, loc, memC, tokens) : Value();
-  Value valC = genAllocCopy(rewriter, loc, memV, tokens);
+  if (dataTransferStrategy != SparseDataTransferStrategy::kRegularDMA)
+    genHostRegisterMemref(rewriter, loc, memV);
+  Value rowC = isZeroCopy ? memR : genAllocCopy(rewriter, loc, memR, tokens);
+  Value colC =
+      memC ? (isZeroCopy ? memC : genAllocCopy(rewriter, loc, memC, tokens))
+           : Value();
+  Value valC = isZeroCopy ? memV : genAllocCopy(rewriter, loc, memV, tokens);
   genBlockingWait(rewriter, loc, tokens);
   tokens.clear();
 
   // Create sparse environment and sparse matrix/dense matrix handles.
   Type indexTp = rewriter.getIndexType();
   Type dnMatHandleTp = rewriter.getType<gpu::SparseDnTensorHandleType>();
   Type spMatHandleTp = rewriter.getType<gpu::SparseSpMatHandleType>();
   Type tokenTp = rewriter.getType<gpu::AsyncTokenType>();

 // 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.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}
+// 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}
+// 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-memory"  | %{run}
+// 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-memory"  | %{run}
+// 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-memory" | %{run}
+// 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-memory" | %{run}
+// 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 = {