[Mlir-commits] [mlir] [mlir][sparse] add assemble test for Batched-CSR and CSR-Dense (PR #81660)

[Aart Bik]

https://github.com/aartbik created https://github.com/llvm/llvm-project/pull/81660

These are formats supported by PyTorch sparse, so good to make sure that our assemble instructions work on these.

>From 487f28ad4504a42fd1a3ee2d22f1554d496a1853 Mon Sep 17 00:00:00 2001
From: Aart Bik <ajcbik at google.com>
Date: Tue, 13 Feb 2024 11:55:29 -0800
Subject: [PATCH] [mlir][sparse] add assemble test for Batched-CSR and
 CSR-Dense

These are formats supported by PyTorch sparse, so good to make
sure that our assemble instructions work on these.
---
 .../SparseTensor/CPU/sparse_pack_d.mlir       | 95 +++++++++++++++++++
 1 file changed, 95 insertions(+)
 create mode 100755 mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_pack_d.mlir

diff --git a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_pack_d.mlir b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_pack_d.mlir
new file mode 100755
index 00000000000000..55585a7c997430
--- /dev/null
+++ b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_pack_d.mlir
@@ -0,0 +1,95 @@
+//--------------------------------------------------------------------------------------------------
+// WHEN CREATING A NEW TEST, PLEASE JUST COPY & PASTE WITHOUT EDITS.
+//
+// Set-up that's shared across all tests in this directory. In principle, this
+// config could be moved to lit.local.cfg. However, there are downstream users that
+// do not use these LIT config files. Hence why this is kept inline.
+//
+// DEFINE: %{sparsifier_opts} = enable-runtime-library=true
+// DEFINE: %{sparsifier_opts_sve} = enable-arm-sve=true %{sparsifier_opts}
+// DEFINE: %{compile} = mlir-opt %s --sparsifier="%{sparsifier_opts}"
+// DEFINE: %{compile_sve} = mlir-opt %s --sparsifier="%{sparsifier_opts_sve}"
+// DEFINE: %{run_libs} = -shared-libs=%mlir_c_runner_utils,%mlir_runner_utils
+// DEFINE: %{run_opts} = -e entry -entry-point-result=void
+// DEFINE: %{run} = mlir-cpu-runner %{run_opts} %{run_libs}
+// DEFINE: %{run_sve} = %mcr_aarch64_cmd --march=aarch64 --mattr="+sve" %{run_opts} %{run_libs}
+//
+// DEFINE: %{env} =
+//--------------------------------------------------------------------------------------------------
+
+// RUN: %{compile} | %{run} | FileCheck %s
+//
+// Do the same run, but now with direct IR generation.
+// REDEFINE: %{sparsifier_opts} = enable-runtime-library=false
+// RUN: %{compile} | %{run} | FileCheck %s
+
+#BatchedCSR = #sparse_tensor.encoding<{
+  map = (d0, d1, d2) -> (d0 : dense, d1 : dense, d2 : compressed),
+  posWidth = 64,
+  crdWidth = 32
+}>
+
+#CSRDense = #sparse_tensor.encoding<{
+  map = (d0, d1, d2) -> (d0 : dense, d1 : compressed, d2 : dense),
+  posWidth = 64,
+  crdWidth = 32
+}>
+
+// Test with batched-CSR and CSR-dense.
+module {
+  //
+  // Main driver.
+  //
+  func.func @entry() {
+    %c0 = arith.constant 0 : index
+    %f0 = arith.constant 0.0 : f32
+
+    //
+    // Setup BatchedCSR.
+    //
+
+    %data1 = arith.constant dense<
+       [ 1.0,  2.0,  3.0,  4.0,  5.0,  6.0,  7.0,  8.0,
+         9.0, 10.0, 11.0, 12.0, 13.0, 14.0, 15.0, 16.0 ]> : tensor<16xf32>
+    %pos1 = arith.constant dense<
+       [ 0, 2, 3, 4, 6, 6, 7, 9, 11, 13, 14, 15, 16 ]> : tensor<13xi64>
+    %crd1 = arith.constant dense<
+       [ 0, 1, 1, 0, 0, 1, 1, 0, 1, 0, 1, 0, 1, 0, 1, 1]> : tensor<16xi32>
+
+    %s1 = sparse_tensor.assemble %data1, %pos1, %crd1 : tensor<16xf32>, tensor<13xi64>, tensor<16xi32> to tensor<4x3x2xf32, #BatchedCSR>
+
+    //
+    // Setup CSRDense.
+    //
+
+    %data2 = arith.constant dense<
+      [ 1.0,  2.0,  0.0,  3.0,  4.0,  0.0, 5.0, 6.0,  0.0, 7.0,  8.0,
+        9.0, 10.0, 11.0, 12.0, 13.0, 14.0, 0.0, 0.0, 15.0, 0.0, 16.0 ]> : tensor<22xf32>
+    %pos2 = arith.constant dense<
+      [ 0, 3, 5, 8, 11 ]> : tensor<5xi64>
+    %crd2 = arith.constant dense<
+      [ 0, 1, 2, 0, 2, 0, 1, 2, 0, 1, 2 ]> : tensor<11xi32>
+
+    %s2 = sparse_tensor.assemble %data2, %pos2, %crd2 : tensor<22xf32>, tensor<5xi64>, tensor<11xi32> to tensor<4x3x2xf32, #CSRDense>
+
+    //
+    // Verify.
+    //
+    // CHECK: ( ( ( 1, 2 ), ( 0, 3 ), ( 4, 0 ) ), ( ( 5, 6 ), ( 0, 0 ), ( 0, 7 ) ), ( ( 8, 9 ), ( 10, 11 ), ( 12, 13 ) ), ( ( 14, 0 ), ( 0, 15 ), ( 0, 16 ) ) )
+    // CHECK: ( ( ( 1, 2 ), ( 0, 3 ), ( 4, 0 ) ), ( ( 5, 6 ), ( 0, 0 ), ( 0, 7 ) ), ( ( 8, 9 ), ( 10, 11 ), ( 12, 13 ) ), ( ( 14, 0 ), ( 0, 15 ), ( 0, 16 ) ) )
+    //
+
+    %d1 = sparse_tensor.convert %s1 : tensor<4x3x2xf32, #BatchedCSR> to tensor<4x3x2xf32>
+    %v1 = vector.transfer_read %d1[%c0, %c0, %c0], %f0 : tensor<4x3x2xf32>, vector<4x3x2xf32>
+    vector.print %v1 : vector<4x3x2xf32>
+
+    %d2 = sparse_tensor.convert %s2 : tensor<4x3x2xf32, #CSRDense> to tensor<4x3x2xf32>
+    %v2 = vector.transfer_read %d1[%c0, %c0, %c0], %f0 : tensor<4x3x2xf32>, vector<4x3x2xf32>
+    vector.print %v2 : vector<4x3x2xf32>
+
+    // FIXME: doing this explicitly crashes runtime
+    // bufferization.dealloc_tensor %s1 : tensor<4x3x2xf32, #BatchedCSR>
+    // bufferization.dealloc_tensor %s2 : tensor<4x3x2xf32, #CSRDense>
+    return
+  }
+}