[mlir][sparse] fix asan issue
authorAart Bik <ajcbik@google.com>
Thu, 16 Jun 2022 20:25:23 +0000 (13:25 -0700)
committerAart Bik <ajcbik@google.com>
Thu, 16 Jun 2022 21:49:02 +0000 (14:49 -0700)
The LinalgElementwiseOpFusion pass has become smarter, and converts
the simple conversion linalg operation into a sparse dialect convert
operation. However, since our current bufferization does not take the
new semantics into consideration, we leak memory of the allocation.
For now, this has been fixed by making the operation less trivial.

Reviewed By: bixia

Differential Revision: https://reviews.llvm.org/D128002

mlir/test/Integration/Dialect/SparseTensor/CPU/dense_output.mlir

index 0018c07..ca02b40 100644 (file)
@@ -23,7 +23,7 @@
     affine_map<(i,j) -> (i,j)>  // X (out)
   ],
   iterator_types = ["parallel", "parallel"],
-  doc = "X(i,j) = A(i,j)"
+  doc = "X(i,j) = A(i,j) * 2"
 }
 
 //
 // library.
 module {
   //
-  // A kernel that assigns elements from A to X.
+  // A kernel that assigns multiplied elements from A to X.
   //
   func.func @dense_output(%arga: tensor<?x?xf64, #SparseMatrix>) -> tensor<?x?xf64, #DenseMatrix> {
     %c0 = arith.constant 0 : index
     %c1 = arith.constant 1 : index
+    %c2 = arith.constant 2.0 : f64
     %d0 = tensor.dim %arga, %c0 : tensor<?x?xf64, #SparseMatrix>
     %d1 = tensor.dim %arga, %c1 : tensor<?x?xf64, #SparseMatrix>
     %init = bufferization.alloc_tensor(%d0, %d1) : tensor<?x?xf64, #DenseMatrix>
@@ -51,7 +52,8 @@ module {
        ins(%arga: tensor<?x?xf64, #SparseMatrix>)
       outs(%init: tensor<?x?xf64, #DenseMatrix>) {
       ^bb(%a: f64, %x: f64):
-        linalg.yield %a : f64
+        %0 = arith.mulf %a, %c2 : f64
+        linalg.yield %0 : f64
     } -> tensor<?x?xf64, #DenseMatrix>
     return %0 : tensor<?x?xf64, #DenseMatrix>
   }
@@ -78,7 +80,7 @@ module {
     //
     // Print the linearized 5x5 result for verification.
     //
-    // CHECK: ( 1, 0, 0, 1.4, 0, 0, 2, 0, 0, 2.5, 0, 0, 3, 0, 0, 4.1, 0, 0, 4, 0, 0, 5.2, 0, 0, 5 )
+    // CHECK: ( 2, 0, 0, 2.8, 0, 0, 4, 0, 0, 5, 0, 0, 6, 0, 0, 8.2, 0, 0, 8, 0, 0, 10.4, 0, 0, 10 )
     //
     %m = sparse_tensor.values %0
       : tensor<?x?xf64, #DenseMatrix> to memref<?xf64>