[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 ↗

(On Diff #540635)

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 ↗

(On Diff #546166)

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

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

47 ↗

(On Diff #546166)

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 ↗

(On Diff #546166)

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 ↗

(On Diff #546166)

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 ↗

(On Diff #546166)

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

castR
castC
castV
castX
castY

559 ↗

(On Diff #546166)

keep block together (unnecessary change)

571 ↗

(On Diff #546166)

no whiteline

621 ↗

(On Diff #546166)

same, just castR, castC etc.

733 ↗

(On Diff #546166)

no whiteline

734 ↗

(On Diff #546166)

castA, castB etc.

874 ↗

(On Diff #546166)

castB, castA

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

83 ↗

(On Diff #546166)

add tracker #

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

14 ↗

(On Diff #546166)

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 ↗

(On Diff #546203)

GPUDataTransferStrategy
(not necessary Sparse)

58 ↗

(On Diff #546203)

let's name this gpuDataTransferStrategy

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

13 ↗

(On Diff #546228)

the other RUN and NOTRUN?

[K-Wu]

remove unnecessary change

[K-Wu]

addressing comments

[K-Wu]

addressing comments

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:  %{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}
+// RUNNOT:  %{compile} enable-runtime-library=true data-transfer-strategy=zero-copy" | %{run}
 //
 // without RT lib (AoS COO): note, may fall back to CPU
 //
 // 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}
+// RUNNOT: %{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: %{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}
+// RUNNOT: %{compile} enable-runtime-library=true data-transfer-strategy=zero-copy"  | %{run}
 //
 // without RT lib (AoS COO): note, may fall back to CPU
 //
 // 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}
+// RUNNOT: %{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" ],

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: %{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}
+// RUNNOT: %{compile} enable-runtime-library=true data-transfer-strategy=zero-copy" | %{run}
 //
 // without RT lib:
 //
 // 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}
+// RUNNOT:  %{compile} enable-runtime-library=false data-transfer-strategy=zero-copy" | %{run}
 //
 
 !Filename = !llvm.ptr<i8>
 
 #CSR = #sparse_tensor.encoding<{
   lvlTypes = ["dense", "compressed"]
 }>