From 96cab659a19e9565d34647da756461060050f518 Mon Sep 17 00:00:00 2001 From: Aart Bik Date: Wed, 19 Oct 2022 10:37:25 -0700 Subject: [PATCH] [mlir][sparse] end-to-end sparse vector insertion codegen Reviewed By: Peiming Differential Revision: https://reviews.llvm.org/D136275 --- .../Transforms/SparseTensorCodegen.cpp | 4 +- mlir/test/Dialect/SparseTensor/codegen.mlir | 8 +-- .../Dialect/SparseTensor/CPU/sparse_insert_1d.mlir | 65 ++++++++++++++++++++++ 3 files changed, 71 insertions(+), 6 deletions(-) create mode 100644 mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_insert_1d.mlir diff --git a/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorCodegen.cpp b/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorCodegen.cpp index 707c6c9..a090e60 100644 --- a/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorCodegen.cpp +++ b/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorCodegen.cpp @@ -282,13 +282,13 @@ static scf::ForOp createFor(OpBuilder &builder, Location loc, Value count, static void createPushback(OpBuilder &builder, Location loc, SmallVectorImpl &fields, unsigned field, Value value) { - assert(field < fields.size()); + assert(2 <= field && field < fields.size()); Type etp = fields[field].getType().cast().getElementType(); if (value.getType() != etp) value = builder.create(loc, etp, value); fields[field] = builder.create(loc, fields[field].getType(), fields[1], - fields[field], value, APInt(64, field)); + fields[field], value, APInt(64, field - 2)); } /// Generates insertion code. diff --git a/mlir/test/Dialect/SparseTensor/codegen.mlir b/mlir/test/Dialect/SparseTensor/codegen.mlir index 6a32c72..bb2eb88 100644 --- a/mlir/test/Dialect/SparseTensor/codegen.mlir +++ b/mlir/test/Dialect/SparseTensor/codegen.mlir @@ -373,8 +373,8 @@ func.func @sparse_expansion3(%arg0: index, %arg1: index) -> memref { // CHECK: %[[R:.*]]:2 = scf.for %[[I:.*]] = %[[C0]] to %[[A8]] step %[[C1]] iter_args(%[[P0:.*]] = %[[A3]], %[[P1:.*]] = %[[A4]]) -> (memref, memref) { // CHECK: %[[T1:.*]] = memref.load %[[A7]][%[[I]]] : memref // CHECK: %[[T2:.*]] = memref.load %[[A5]][%[[T1]]] : memref -// CHECK: %[[T3:.*]] = sparse_tensor.push_back %[[A1]], %[[P0]], %[[T1]] {idx = 3 : index} : memref<3xindex>, memref, index -// CHECK: %[[T4:.*]] = sparse_tensor.push_back %[[A1]], %[[P1]], %[[T2]] {idx = 4 : index} : memref<3xindex>, memref, f64 +// CHECK: %[[T3:.*]] = sparse_tensor.push_back %[[A1]], %[[P0]], %[[T1]] {idx = 1 : index} : memref<3xindex>, memref, index +// CHECK: %[[T4:.*]] = sparse_tensor.push_back %[[A1]], %[[P1]], %[[T2]] {idx = 2 : index} : memref<3xindex>, memref, f64 // CHECK: memref.store %[[F0]], %arg5[%[[T1]]] : memref // CHECK: memref.store %[[B0]], %arg6[%[[T1]]] : memref // CHECK: scf.yield %[[T3]], %[[T4]] : memref, memref @@ -383,8 +383,8 @@ func.func @sparse_expansion3(%arg0: index, %arg1: index) -> memref { // CHECK: memref.dealloc %[[A6]] : memref // CHECK: memref.dealloc %[[A7]] : memref // CHECK: %[[LL:.*]] = memref.load %[[A1]][%[[C2]]] : memref<3xindex> -// CHECK: %[[P1:.*]] = sparse_tensor.push_back %[[A1]], %[[A2]], %[[C0]] {idx = 2 : index} : memref<3xindex>, memref, index -// CHECK: %[[P2:.*]] = sparse_tensor.push_back %[[A1]], %[[P1]], %[[LL]] {idx = 2 : index} : memref<3xindex>, memref, index +// CHECK: %[[P1:.*]] = sparse_tensor.push_back %[[A1]], %[[A2]], %[[C0]] {idx = 0 : index} : memref<3xindex>, memref, index +// CHECK: %[[P2:.*]] = sparse_tensor.push_back %[[A1]], %[[P1]], %[[LL]] {idx = 0 : index} : memref<3xindex>, memref, index // CHECK: return %[[A0]], %[[A1]], %[[P2]], %[[R]]#0, %[[R]]#1 : memref<1xindex>, memref<3xindex>, memref, memref, memref func.func @sparse_compression_1d(%tensor: tensor<100xf64, #SV>, %values: memref, diff --git a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_insert_1d.mlir b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_insert_1d.mlir new file mode 100644 index 0000000..6b8576e --- /dev/null +++ b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_insert_1d.mlir @@ -0,0 +1,65 @@ +// RUN: mlir-opt %s --sparse-compiler=enable-runtime-library=false | \ +// RUN: mlir-cpu-runner \ +// RUN: -e entry -entry-point-result=void \ +// RUN: -shared-libs=%mlir_lib_dir/libmlir_c_runner_utils%shlibext | \ +// RUN: FileCheck %s + +// Insertion example using pure codegen (no sparse runtime support lib). + +#SparseVector = #sparse_tensor.encoding<{ dimLevelType = [ "compressed" ] }> + +#trait_mul_s = { + indexing_maps = [ + affine_map<(i) -> (i)> // x (out) + ], + iterator_types = ["parallel"], + doc = "x(i) = x(i) * 2.0" +} + +module { + + // Dumps pointers, indices, values for verification. + func.func @dump(%argx: tensor<1024xf32, #SparseVector>) { + %c0 = arith.constant 0 : index + %cu = arith.constant 99 : index + %fu = arith.constant 99.0 : f32 + %p = sparse_tensor.pointers %argx { dimension = 0 : index } + : tensor<1024xf32, #SparseVector> to memref + %i = sparse_tensor.indices %argx { dimension = 0 : index } + : tensor<1024xf32, #SparseVector> to memref + %v = sparse_tensor.values %argx + : tensor<1024xf32, #SparseVector> to memref + %vp = vector.transfer_read %p[%c0], %cu: memref, vector<8xindex> + %vi = vector.transfer_read %i[%c0], %cu: memref, vector<8xindex> + %vv = vector.transfer_read %v[%c0], %fu: memref, vector<8xf32> + vector.print %vp : vector<8xindex> + vector.print %vi : vector<8xindex> + vector.print %vv : vector<8xf32> + return + } + + func.func @entry() { + %f1 = arith.constant 1.0 : f32 + %f2 = arith.constant 2.0 : f32 + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index + %c3 = arith.constant 3 : index + %c1023 = arith.constant 1023 : index + + // Build the sparse vector from code. + %0 = bufferization.alloc_tensor() : tensor<1024xf32, #SparseVector> + %1 = sparse_tensor.insert %f1 into %0[%c0] : tensor<1024xf32, #SparseVector> + %2 = sparse_tensor.insert %f2 into %1[%c1] : tensor<1024xf32, #SparseVector> + %3 = sparse_tensor.insert %f1 into %2[%c3] : tensor<1024xf32, #SparseVector> + %4 = sparse_tensor.insert %f2 into %3[%c1023] : tensor<1024xf32, #SparseVector> + %5 = sparse_tensor.load %4 hasInserts : tensor<1024xf32, #SparseVector> + + // CHECK: ( 0, 4, 99, 99, 99, 99, 99, 99 ) + // CHECK-NEXT: ( 0, 1, 3, 1023, 99, 99, 99, 99 ) + // CHECK-NEXT: ( 1, 2, 1, 2, 99, 99, 99, 99 ) + call @dump(%5) : (tensor<1024xf32, #SparseVector>) -> () + + bufferization.dealloc_tensor %5 : tensor<1024xf32, #SparseVector> + return + } +} -- 2.7.4