diff options
author | Aart Bik <39774503+aartbik@users.noreply.github.com> | 2023-11-16 17:17:41 -0800 |
---|---|---|
committer | GitHub <noreply@github.com> | 2023-11-16 17:17:41 -0800 |
commit | 83cf0dc98234bbd8cb0d0959baa570477a8daf92 (patch) | |
tree | 2654c69f37e36c94b34e336dc50967700e710ba9 | |
parent | d97981c98a70ebeaa8901f34921b0a69a068ff5d (diff) | |
download | llvm-83cf0dc98234bbd8cb0d0959baa570477a8daf92.zip llvm-83cf0dc98234bbd8cb0d0959baa570477a8daf92.tar.gz llvm-83cf0dc98234bbd8cb0d0959baa570477a8daf92.tar.bz2 |
[mlir][sparse] implement direct IR alloc/empty/new for non-permutations (#72585)
This change implements the correct *level* sizes set up for the direct
IR codegen fields in the sparse storage scheme. This brings libgen and
codegen together again.
This is step 3 out of 3 to make sparse_tensor.new work for BSR
6 files changed, 128 insertions, 121 deletions
diff --git a/mlir/include/mlir/ExecutionEngine/SparseTensor/File.h b/mlir/include/mlir/ExecutionEngine/SparseTensor/File.h index 85bbfe0..6b4a174 100644 --- a/mlir/include/mlir/ExecutionEngine/SparseTensor/File.h +++ b/mlir/include/mlir/ExecutionEngine/SparseTensor/File.h @@ -1,4 +1,4 @@ -//===- File.h - Reading sparse tensors from files --------------*- C++ -*-===// +//===- File.h - Reading sparse tensors from files ---------------*- C++ -*-===// // // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. // See https://llvm.org/LICENSE.txt for license information. diff --git a/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorCodegen.cpp b/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorCodegen.cpp index cfc8eb1..9f41db7 100644 --- a/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorCodegen.cpp +++ b/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorCodegen.cpp @@ -159,41 +159,46 @@ static Value createAllocation(OpBuilder &builder, Location loc, return buffer; } +/// Creates the dim sizes array, filling in from dynamic sizes. +static void createDimSizes(OpBuilder &builder, Location loc, + SparseTensorType stt, ValueRange dynSizes, + /*out*/ SmallVectorImpl<Value> &dimSizesValues) { + const Dimension dimRank = stt.getDimRank(); + dimSizesValues.clear(); + dimSizesValues.reserve(dimRank); + unsigned i = 0; + for (const Size sz : stt.getDimShape()) + dimSizesValues.push_back(ShapedType::isDynamic(sz) + ? dynSizes[i++] + : constantIndex(builder, loc, sz)); +} + /// Creates allocation for each field in sparse tensor type. Note that /// for all dynamic memrefs in the sparse tensor stroage layout, the /// memory size is really the capacity of the "vector", while the actual /// size resides in the sizes array. static void createAllocFields(OpBuilder &builder, Location loc, - SparseTensorType stt, ValueRange dynSizes, - bool enableInit, SmallVectorImpl<Value> &fields, - Value sizeHint) { - // Build original sizes. - assert((dynSizes.size() == static_cast<size_t>(stt.getNumDynamicDims())) && - "Got wrong number of dynamic sizes"); - const Dimension dimRank = stt.getDimRank(); - SmallVector<Value> dimSizes; - dimSizes.reserve(dimRank); - unsigned i = 0; // cumulative index into `dynSizes`. - for (const Size sh : stt.getDimShape()) - dimSizes.push_back(ShapedType::isDynamic(sh) - ? dynSizes[i++] - : constantIndex(builder, loc, sh)); - + SparseTensorType stt, bool enableInit, + Value sizeHint, + SmallVectorImpl<Value> &lvlSizesValues, + /*out*/ SmallVectorImpl<Value> &fields) { + Level lvlRank = stt.getLvlRank(); // Set up some heuristic sizes. We try to set the initial // size based on available information. Otherwise we just // initialize a few elements to start the reallocation chain. // TODO: refine this Value posHeuristic, crdHeuristic, valHeuristic; if (stt.isAllDense()) { - valHeuristic = dimSizes[0]; - for (const Value sz : ArrayRef<Value>{dimSizes}.drop_front()) - valHeuristic = builder.create<arith::MulIOp>(loc, valHeuristic, sz); + valHeuristic = lvlSizesValues[0]; + for (Level lvl = 1; lvl < lvlRank; lvl++) + valHeuristic = + builder.create<arith::MulIOp>(loc, valHeuristic, lvlSizesValues[lvl]); } else if (sizeHint) { if (getCOOStart(stt.getEncoding()) == 0) { posHeuristic = constantIndex(builder, loc, 2); crdHeuristic = builder.create<arith::MulIOp>( - loc, constantIndex(builder, loc, dimRank), sizeHint); // AOS - } else if (dimRank == 2 && stt.isDenseLvl(0) && stt.isCompressedLvl(1)) { + loc, constantIndex(builder, loc, lvlRank), sizeHint); // AOS + } else if (lvlRank == 2 && stt.isDenseLvl(0) && stt.isCompressedLvl(1)) { posHeuristic = builder.create<arith::AddIOp>( loc, sizeHint, constantIndex(builder, loc, 1)); crdHeuristic = sizeHint; @@ -205,7 +210,6 @@ static void createAllocFields(OpBuilder &builder, Location loc, posHeuristic = crdHeuristic = valHeuristic = constantIndex(builder, loc, 16); } - // Initializes all fields. An initial storage specifier and allocated // positions/coordinates/values memrefs (with heuristic capacity). foreachFieldAndTypeInSparseTensor( @@ -237,16 +241,13 @@ static void createAllocFields(OpBuilder &builder, Location loc, // Returns true to continue the iteration. return true; }); - // Initialize the storage scheme to an empty tensor. Sets the lvlSizes // and gives all position fields an initial zero entry, so that it is // easier to maintain the "linear + 1" length property. MutSparseTensorDescriptor desc(stt, fields); Value posZero = constantZero(builder, loc, stt.getPosType()); for (Level lvl = 0, lvlRank = stt.getLvlRank(); lvl < lvlRank; lvl++) { - // FIXME: `toOrigDim` is deprecated. - desc.setLvlSize(builder, loc, lvl, - dimSizes[toOrigDim(stt.getEncoding(), lvl)]); + desc.setLvlSize(builder, loc, lvl, lvlSizesValues[lvl]); const auto dlt = stt.getLvlType(lvl); if (isCompressedDLT(dlt) || isLooseCompressedDLT(dlt)) createPushback(builder, loc, desc, SparseTensorFieldKind::PosMemRef, lvl, @@ -371,8 +372,8 @@ static void genEndInsert(OpBuilder &builder, Location loc, SparseTensorDescriptor desc) { const SparseTensorType stt(desc.getRankedTensorType()); const Level lvlRank = stt.getLvlRank(); - for (Level l = 0; l < lvlRank; l++) { - const auto dlt = stt.getLvlType(l); + for (Level lvl = 0; lvl < lvlRank; lvl++) { + const auto dlt = stt.getLvlType(lvl); if (isCompressedDLT(dlt)) { // Compressed dimensions need a position cleanup for all entries // that were not visited during the insertion pass. @@ -380,10 +381,10 @@ static void genEndInsert(OpBuilder &builder, Location loc, // TODO: avoid cleanup and keep compressed scheme consistent at all // times? // - if (l > 0) { + if (lvl > 0) { Type posType = stt.getPosType(); - Value posMemRef = desc.getPosMemRef(l); - Value hi = desc.getPosMemSize(builder, loc, l); + Value posMemRef = desc.getPosMemRef(lvl); + Value hi = desc.getPosMemSize(builder, loc, lvl); Value zero = constantIndex(builder, loc, 0); Value one = constantIndex(builder, loc, 1); // Vector of only one, but needed by createFor's prototype. @@ -729,7 +730,6 @@ public: }; /// Sparse codegen rule for the alloc operator. -/// TODO(springerm): remove when bufferization.alloc_tensor is gone class SparseTensorAllocConverter : public OpConversionPattern<bufferization::AllocTensorOp> { public: @@ -746,6 +746,8 @@ public: if (!resType.hasEncoding()) return failure(); Location loc = op.getLoc(); + + // Deal with copy. if (op.getCopy()) { auto desc = getDescriptorFromTensorTuple(adaptor.getCopy()); SmallVector<Value> fields; @@ -766,17 +768,22 @@ public: return success(); } + // Construct the dim/lvl sizes and the (unused) dim2lvl/lvl2dim buffers. + SmallVector<Value> dimSizesValues; + SmallVector<Value> lvlSizesValues; + Value dimSizesBuffer; + Value dim2lvlBuffer; + Value lvl2dimBuffer; + createDimSizes(rewriter, loc, resType, adaptor.getDynamicSizes(), + dimSizesValues); + genMapBuffers(rewriter, loc, resType, dimSizesValues, dimSizesBuffer, + lvlSizesValues, dim2lvlBuffer, lvl2dimBuffer); + // Construct allocation for each field. Value sizeHint = op.getSizeHint(); - ValueRange dynSizes = adaptor.getDynamicSizes(); - const size_t found = dynSizes.size(); - const int64_t expected = resType.getNumDynamicDims(); - if (found != static_cast<size_t>(expected)) - return rewriter.notifyMatchFailure(op, - "Got wrong number of dynamic sizes"); SmallVector<Value> fields; - createAllocFields(rewriter, loc, resType, dynSizes, - enableBufferInitialization, fields, sizeHint); + createAllocFields(rewriter, loc, resType, enableBufferInitialization, + sizeHint, lvlSizesValues, fields); // Replace operation with resulting memrefs. rewriter.replaceOp(op, genTuple(rewriter, loc, resType, fields)); @@ -788,7 +795,6 @@ private: }; /// Sparse codegen rule for the empty tensor operator. -/// TODO(springerm): remove when bufferization.alloc_tensor is gone class SparseTensorEmptyConverter : public OpConversionPattern<tensor::EmptyOp> { public: using OpConversionPattern::OpConversionPattern; @@ -803,19 +809,24 @@ public: const auto resType = getSparseTensorType(op); if (!resType.hasEncoding()) return failure(); + Location loc = op.getLoc(); + + // Construct the dim/lvl sizes and the (unused) dim2lvl/lvl2dim buffers. + SmallVector<Value> dimSizesValues; + SmallVector<Value> lvlSizesValues; + Value dimSizesBuffer; + Value dim2lvlBuffer; + Value lvl2dimBuffer; + createDimSizes(rewriter, loc, resType, adaptor.getDynamicSizes(), + dimSizesValues); + genMapBuffers(rewriter, loc, resType, dimSizesValues, dimSizesBuffer, + lvlSizesValues, dim2lvlBuffer, lvl2dimBuffer); // Construct allocation for each field. - Location loc = op.getLoc(); Value sizeHint; // none - const ValueRange dynSizes = adaptor.getDynamicSizes(); - const size_t found = dynSizes.size(); - const int64_t expected = resType.getNumDynamicDims(); - if (found != static_cast<size_t>(expected)) - return rewriter.notifyMatchFailure(op, - "Got wrong number of dynamic sizes"); SmallVector<Value> fields; - createAllocFields(rewriter, loc, resType, dynSizes, - enableBufferInitialization, fields, sizeHint); + createAllocFields(rewriter, loc, resType, enableBufferInitialization, + sizeHint, lvlSizesValues, fields); // Replace operation with resulting memrefs. rewriter.replaceOp(op, genTuple(rewriter, loc, resType, fields)); @@ -1461,10 +1472,10 @@ struct SparseNewConverter : public OpConversionPattern<NewOp> { // if (! %isSorted) sparse_tensor.sort_coo(%nse, %coordinates, %values) // update storage specifier // @delSparseTensorReader(%reader) - SmallVector<Value> dimShapesValues; + SmallVector<Value> dimSizesValues; Value dimSizesBuffer; Value reader = genReader(rewriter, loc, dstTp, adaptor.getOperands()[0], - dimShapesValues, dimSizesBuffer); + dimSizesValues, dimSizesBuffer); // Get the number of stored entries. const Type indexTp = rewriter.getIndexType(); @@ -1472,25 +1483,19 @@ struct SparseNewConverter : public OpConversionPattern<NewOp> { {indexTp}, {reader}, EmitCInterface::Off) .getResult(0); - // Construct allocation for each field. - SmallVector<Value> dynSizes; - if (dstTp.hasDynamicDimShape()) { - for (const auto &d : llvm::enumerate(dstTp.getDimShape())) - if (ShapedType::isDynamic(d.value())) - dynSizes.push_back(rewriter.create<memref::LoadOp>( - loc, dimSizesBuffer, constantIndex(rewriter, loc, d.index()))); - } - SmallVector<Value> fields; - createAllocFields(rewriter, loc, dstTp, dynSizes, /*enableInit=*/false, - fields, nse); - - // Now construct the lvl sizes and the dim2lvl/lvl2dim buffers. + // Construct the lvl sizes and the dim2lvl/lvl2dim buffers. SmallVector<Value> lvlSizesValues; Value dim2lvlBuffer; Value lvl2dimBuffer; - genMapBuffers(rewriter, loc, dstTp, dimShapesValues, dimSizesBuffer, + genMapBuffers(rewriter, loc, dstTp, dimSizesValues, dimSizesBuffer, lvlSizesValues, dim2lvlBuffer, lvl2dimBuffer); + // Construct allocation for each field. + Value sizeHint = nse; + SmallVector<Value> fields; + createAllocFields(rewriter, loc, dstTp, /*enableInit=*/false, sizeHint, + lvlSizesValues, fields); + // Read the COO tensor data. MutSparseTensorDescriptor desc(dstTp, fields); Value xs = desc.getAOSMemRef(); diff --git a/mlir/test/Dialect/SparseTensor/codegen.mlir b/mlir/test/Dialect/SparseTensor/codegen.mlir index 12d77ec..e63595b 100644 --- a/mlir/test/Dialect/SparseTensor/codegen.mlir +++ b/mlir/test/Dialect/SparseTensor/codegen.mlir @@ -252,7 +252,7 @@ func.func @sparse_values_coo(%arg0: tensor<?x?x?xf64, #ccoo>) -> memref<?xf64> { } -// CHECK-LABEL: func.func @sparse_indices_coo( +// CHECK-LABEL: func.func @sparse_indices_coo( // CHECK-SAME: %[[A0:.*0]]: memref<?xindex>, // CHECK-SAME: %[[A1:.*1]]: memref<?xindex>, // CHECK-SAME: %[[A2:.*2]]: memref<?xindex>, @@ -270,7 +270,7 @@ func.func @sparse_indices_coo(%arg0: tensor<?x?x?xf64, #ccoo>) -> memref<?xindex return %0 : memref<?xindex, strided<[?], offset: ?>> } -// CHECK-LABEL: func.func @sparse_indices_buffer_coo( +// CHECK-LABEL: func.func @sparse_indices_buffer_coo( // CHECK-SAME: %[[A0:.*0]]: memref<?xindex>, // CHECK-SAME: %[[A1:.*1]]: memref<?xindex>, // CHECK-SAME: %[[A2:.*2]]: memref<?xindex>, @@ -676,26 +676,26 @@ func.func @sparse_convert_element_type(%arg0: tensor<32xf32, #SparseVector>) -> // CHECK: memref.store %[[VAL_4]], %[[VAL_6]]{{\[}}%[[VAL_3]]] : memref<2xindex> // CHECK: %[[VAL_8:.*]] = call @createCheckedSparseTensorReader(%[[A0]], %[[VAL_7]], %[[VAL_2]]) : (!llvm.ptr, memref<?xindex>, i32) -> !llvm.ptr // CHECK: %[[VAL_9:.*]] = call @getSparseTensorReaderDimSizes(%[[VAL_8]]) : (!llvm.ptr) -> memref<?xindex> -// CHECK: %[[VAL_10:.*]] = call @getSparseTensorReaderNSE(%[[VAL_8]]) : (!llvm.ptr) -> index -// CHECK: %[[VAL_11:.*]] = memref.load %[[VAL_9]]{{\[}}%[[VAL_4]]] : memref<?xindex> -// CHECK: %[[VAL_12:.*]] = memref.load %[[VAL_9]]{{\[}}%[[VAL_3]]] : memref<?xindex> -// CHECK: %[[VAL_13:.*]] = arith.muli %[[VAL_10]], %[[VAL_5]] : index -// CHECK: %[[VAL_14:.*]] = memref.alloc() : memref<2xindex> -// CHECK: %[[VAL_15:.*]] = memref.cast %[[VAL_14]] : memref<2xindex> to memref<?xindex> -// CHECK: %[[VAL_16:.*]] = memref.alloc(%[[VAL_13]]) : memref<?xindex> -// CHECK: %[[VAL_17:.*]] = memref.alloc(%[[VAL_10]]) : memref<?xf32> -// CHECK: %[[VAL_18:.*]] = sparse_tensor.storage_specifier.init -// CHECK: %[[VAL_19:.*]] = sparse_tensor.storage_specifier.set %[[VAL_18]] lvl_sz at 0 with %[[VAL_11]] -// CHECK: %[[VAL_20:.*]] = sparse_tensor.storage_specifier.get %[[VAL_19]] pos_mem_sz at 0 -// CHECK: %[[VAL_21:.*]], %[[VAL_22:.*]] = sparse_tensor.push_back %[[VAL_20]], %[[VAL_15]], %[[VAL_4]] -// CHECK: %[[VAL_23:.*]] = sparse_tensor.storage_specifier.set %[[VAL_19]] pos_mem_sz at 0 with %[[VAL_22]] -// CHECK: %[[VAL_24:.*]] = sparse_tensor.storage_specifier.set %[[VAL_23]] lvl_sz at 1 with %[[VAL_12]] -// CHECK: %[[VAL_25:.*]], %[[VAL_26:.*]] = sparse_tensor.push_back %[[VAL_22]], %[[VAL_21]], %[[VAL_4]], %[[VAL_3]] -// CHECK: %[[VAL_27:.*]] = sparse_tensor.storage_specifier.set %[[VAL_24]] pos_mem_sz at 0 with %[[VAL_26]] -// CHECK: %[[VAL_28:.*]] = memref.alloca() : memref<2xindex> -// CHECK: %[[VAL_29:.*]] = memref.cast %[[VAL_28]] : memref<2xindex> to memref<?xindex> -// CHECK: memref.store %[[VAL_4]], %[[VAL_28]]{{\[}}%[[VAL_4]]] : memref<2xindex> -// CHECK: memref.store %[[VAL_3]], %[[VAL_28]]{{\[}}%[[VAL_3]]] : memref<2xindex> +// CHECK-DAG: %[[VAL_10:.*]] = call @getSparseTensorReaderNSE(%[[VAL_8]]) : (!llvm.ptr) -> index +// CHECK-DAG: %[[VAL_11:.*]] = memref.load %[[VAL_9]]{{\[}}%[[VAL_4]]] : memref<?xindex> +// CHECK-DAG: %[[VAL_12:.*]] = memref.load %[[VAL_9]]{{\[}}%[[VAL_3]]] : memref<?xindex> +// CHECK-DAG: %[[VAL_13:.*]] = arith.muli %[[VAL_10]], %[[VAL_5]] : index +// CHECK-DAG: %[[VAL_14:.*]] = memref.alloc() : memref<2xindex> +// CHECK-DAG: %[[VAL_15:.*]] = memref.cast %[[VAL_14]] : memref<2xindex> to memref<?xindex> +// CHECK-DAG: %[[VAL_16:.*]] = memref.alloc(%[[VAL_13]]) : memref<?xindex> +// CHECK-DAG: %[[VAL_17:.*]] = memref.alloc(%[[VAL_10]]) : memref<?xf32> +// CHECK-DAG: %[[VAL_18:.*]] = sparse_tensor.storage_specifier.init +// CHECK-DAG: %[[VAL_19:.*]] = sparse_tensor.storage_specifier.set %[[VAL_18]] lvl_sz at 0 with %[[VAL_11]] +// CHECK-DAG: %[[VAL_20:.*]] = sparse_tensor.storage_specifier.get %[[VAL_19]] pos_mem_sz at 0 +// CHECK-DAG: %[[VAL_21:.*]], %[[VAL_22:.*]] = sparse_tensor.push_back %[[VAL_20]], %[[VAL_15]], %[[VAL_4]] +// CHECK-DAG: %[[VAL_23:.*]] = sparse_tensor.storage_specifier.set %[[VAL_19]] pos_mem_sz at 0 with %[[VAL_22]] +// CHECK-DAG: %[[VAL_24:.*]] = sparse_tensor.storage_specifier.set %[[VAL_23]] lvl_sz at 1 with %[[VAL_12]] +// CHECK-DAG: %[[VAL_25:.*]], %[[VAL_26:.*]] = sparse_tensor.push_back %[[VAL_22]], %[[VAL_21]], %[[VAL_4]], %[[VAL_3]] +// CHECK-DAG: %[[VAL_27:.*]] = sparse_tensor.storage_specifier.set %[[VAL_24]] pos_mem_sz at 0 with %[[VAL_26]] +// CHECK-DAG: %[[VAL_28:.*]] = memref.alloca() : memref<2xindex> +// CHECK-DAG: %[[VAL_29:.*]] = memref.cast %[[VAL_28]] : memref<2xindex> to memref<?xindex> +// CHECK-DAG: memref.store %[[VAL_4]], %[[VAL_28]]{{\[}}%[[VAL_4]]] : memref<2xindex> +// CHECK-DAG: memref.store %[[VAL_3]], %[[VAL_28]]{{\[}}%[[VAL_3]]] : memref<2xindex> // CHECK: %[[VAL_30:.*]] = call @getSparseTensorReaderReadToBuffers0F32(%[[VAL_8]], %[[VAL_29]], %[[VAL_29]], %[[VAL_16]], %[[VAL_17]]) : (!llvm.ptr, memref<?xindex>, memref<?xindex>, memref<?xindex>, memref<?xf32>) -> i1 // CHECK: %[[VAL_31:.*]] = arith.cmpi eq, %[[VAL_30]], %[[VAL_1]] : i1 // CHECK: scf.if %[[VAL_31]] { @@ -722,31 +722,31 @@ func.func @sparse_new_coo(%arg0: !llvm.ptr) -> tensor<?x?xf32, #Coo> { // CHECK: memref.store %[[VAL_3]], %[[VAL_5]]{{\[}}%[[VAL_3]]] : memref<2xindex> // CHECK: memref.store %[[VAL_3]], %[[VAL_5]]{{\[}}%[[VAL_2]]] : memref<2xindex> // CHECK: %[[VAL_7:.*]] = call @createCheckedSparseTensorReader(%[[A0]], %[[VAL_6]], %[[VAL_1]]) : (!llvm.ptr, memref<?xindex>, i32) -> !llvm.ptr -// CHECK: %[[VAL_8:.*]] = call @getSparseTensorReaderDimSizes(%[[VAL_7]]) : (!llvm.ptr) -> memref<?xindex> -// CHECK: %[[VAL_9:.*]] = call @getSparseTensorReaderNSE(%[[VAL_7]]) : (!llvm.ptr) -> index -// CHECK: %[[VAL_10:.*]] = memref.load %[[VAL_8]]{{\[}}%[[VAL_3]]] : memref<?xindex> -// CHECK: %[[VAL_11:.*]] = memref.load %[[VAL_8]]{{\[}}%[[VAL_2]]] : memref<?xindex> -// CHECK: %[[VAL_12:.*]] = arith.muli %[[VAL_9]], %[[VAL_4]] : index -// CHECK: %[[VAL_13:.*]] = memref.alloc() : memref<2xindex> -// CHECK: %[[VAL_14:.*]] = memref.cast %[[VAL_13]] : memref<2xindex> to memref<?xindex> -// CHECK: %[[VAL_15:.*]] = memref.alloc(%[[VAL_12]]) : memref<?xindex> -// CHECK: %[[VAL_16:.*]] = memref.alloc(%[[VAL_9]]) : memref<?xf32> -// CHECK: %[[VAL_17:.*]] = sparse_tensor.storage_specifier.init -// CHECK: %[[VAL_18:.*]] = sparse_tensor.storage_specifier.set %[[VAL_17]] lvl_sz at 0 with %[[VAL_11]] -// CHECK: %[[VAL_19:.*]] = sparse_tensor.storage_specifier.get %[[VAL_18]] pos_mem_sz at 0 -// CHECK: %[[VAL_20:.*]], %[[VAL_21:.*]] = sparse_tensor.push_back %[[VAL_19]], %[[VAL_14]], %[[VAL_3]] -// CHECK: %[[VAL_22:.*]] = sparse_tensor.storage_specifier.set %[[VAL_18]] pos_mem_sz at 0 with %[[VAL_21]] -// CHECK: %[[VAL_23:.*]] = sparse_tensor.storage_specifier.set %[[VAL_22]] lvl_sz at 1 with %[[VAL_10]] -// CHECK: %[[VAL_24:.*]], %[[VAL_25:.*]] = sparse_tensor.push_back %[[VAL_21]], %[[VAL_20]], %[[VAL_3]], %[[VAL_2]] -// CHECK: %[[VAL_26:.*]] = sparse_tensor.storage_specifier.set %[[VAL_23]] pos_mem_sz at 0 with %[[VAL_25]] -// CHECK: %[[VAL_27:.*]] = memref.alloca() : memref<2xindex> -// CHECK: %[[VAL_28:.*]] = memref.cast %[[VAL_27]] : memref<2xindex> to memref<?xindex> -// CHECK: memref.store %[[VAL_2]], %[[VAL_27]]{{\[}}%[[VAL_3]]] : memref<2xindex> -// CHECK: memref.store %[[VAL_3]], %[[VAL_27]]{{\[}}%[[VAL_2]]] : memref<2xindex> -// CHECK: %[[VAL_29:.*]] = memref.alloca() : memref<2xindex> -// CHECK: %[[VAL_30:.*]] = memref.cast %[[VAL_29]] : memref<2xindex> to memref<?xindex> -// CHECK: memref.store %[[VAL_2]], %[[VAL_29]]{{\[}}%[[VAL_3]]] : memref<2xindex> -// CHECK: memref.store %[[VAL_3]], %[[VAL_29]]{{\[}}%[[VAL_2]]] : memref<2xindex> +// CHECK-DAG: %[[VAL_8:.*]] = call @getSparseTensorReaderDimSizes(%[[VAL_7]]) : (!llvm.ptr) -> memref<?xindex> +// CHECK-DAG: %[[VAL_9:.*]] = call @getSparseTensorReaderNSE(%[[VAL_7]]) : (!llvm.ptr) -> index +// CHECK-DAG: %[[VAL_10:.*]] = memref.load %[[VAL_8]]{{\[}}%[[VAL_3]]] : memref<?xindex> +// CHECK-DAG: %[[VAL_11:.*]] = memref.load %[[VAL_8]]{{\[}}%[[VAL_2]]] : memref<?xindex> +// CHECK-DAG: %[[VAL_12:.*]] = arith.muli %[[VAL_9]], %[[VAL_4]] : index +// CHECK-DAG: %[[VAL_13:.*]] = memref.alloc() : memref<2xindex> +// CHECK-DAG: %[[VAL_14:.*]] = memref.cast %[[VAL_13]] : memref<2xindex> to memref<?xindex> +// CHECK-DAG: %[[VAL_15:.*]] = memref.alloc(%[[VAL_12]]) : memref<?xindex> +// CHECK-DAG: %[[VAL_16:.*]] = memref.alloc(%[[VAL_9]]) : memref<?xf32> +// CHECK-DAG: %[[VAL_17:.*]] = sparse_tensor.storage_specifier.init +// CHECK-DAG: %[[VAL_18:.*]] = sparse_tensor.storage_specifier.set %[[VAL_17]] lvl_sz at 0 with %[[VAL_11]] +// CHECK-DAG: %[[VAL_19:.*]] = sparse_tensor.storage_specifier.get %[[VAL_18]] pos_mem_sz at 0 +// CHECK-DAG: %[[VAL_20:.*]], %[[VAL_21:.*]] = sparse_tensor.push_back %[[VAL_19]], %[[VAL_14]], %[[VAL_3]] +// CHECK-DAG: %[[VAL_22:.*]] = sparse_tensor.storage_specifier.set %[[VAL_18]] pos_mem_sz at 0 with %[[VAL_21]] +// CHECK-DAG: %[[VAL_23:.*]] = sparse_tensor.storage_specifier.set %[[VAL_22]] lvl_sz at 1 with %[[VAL_10]] +// CHECK-DAG: %[[VAL_24:.*]], %[[VAL_25:.*]] = sparse_tensor.push_back %[[VAL_21]], %[[VAL_20]], %[[VAL_3]], %[[VAL_2]] +// CHECK-DAG: %[[VAL_26:.*]] = sparse_tensor.storage_specifier.set %[[VAL_23]] pos_mem_sz at 0 with %[[VAL_25]] +// CHECK-DAG: %[[VAL_27:.*]] = memref.alloca() : memref<2xindex> +// CHECK-DAG: %[[VAL_28:.*]] = memref.cast %[[VAL_27]] : memref<2xindex> to memref<?xindex> +// CHECK-DAG: memref.store %[[VAL_2]], %[[VAL_27]]{{\[}}%[[VAL_3]]] : memref<2xindex> +// CHECK-DAG: memref.store %[[VAL_3]], %[[VAL_27]]{{\[}}%[[VAL_2]]] : memref<2xindex> +// CHECK-DAG: %[[VAL_29:.*]] = memref.alloca() : memref<2xindex> +// CHECK-DAG: %[[VAL_30:.*]] = memref.cast %[[VAL_29]] : memref<2xindex> to memref<?xindex> +// CHECK-DAG: memref.store %[[VAL_2]], %[[VAL_29]]{{\[}}%[[VAL_3]]] : memref<2xindex> +// CHECK-DAG: memref.store %[[VAL_3]], %[[VAL_29]]{{\[}}%[[VAL_2]]] : memref<2xindex> // CHECK: %[[VAL_31:.*]] = call @getSparseTensorReaderReadToBuffers0F32(%[[VAL_7]], %[[VAL_28]], %[[VAL_30]], %[[VAL_15]], %[[VAL_16]]) : (!llvm.ptr, memref<?xindex>, memref<?xindex>, memref<?xindex>, memref<?xf32>) -> i1 // CHECK: memref.store %[[VAL_9]], %[[VAL_24]]{{\[}}%[[VAL_2]]] : memref<?xindex> // CHECK: %[[VAL_32:.*]] = sparse_tensor.storage_specifier.set %[[VAL_26]] crd_mem_sz at 0 with %[[VAL_12]] diff --git a/mlir/test/Dialect/SparseTensor/codegen_buffer_initialization.mlir b/mlir/test/Dialect/SparseTensor/codegen_buffer_initialization.mlir index 14a8f0e..e6d9700 100644 --- a/mlir/test/Dialect/SparseTensor/codegen_buffer_initialization.mlir +++ b/mlir/test/Dialect/SparseTensor/codegen_buffer_initialization.mlir @@ -4,9 +4,9 @@ // CHECK-LABEL: func.func @empty_sparse_vector( // CHECK-SAME: %[[VAL_0:.*]]: index) -> (memref<?xindex>, memref<?xindex>, memref<?xf64>, !sparse_tensor.storage_specifier -// CHECK: %[[VAL_1:.*]] = arith.constant 1 : index -// CHECK: %[[VAL_2:.*]] = arith.constant 0.000000e+00 : f64 -// CHECK: %[[VAL_3:.*]] = arith.constant 0 : index +// CHECK-DAG: %[[VAL_1:.*]] = arith.constant 1 : index +// CHECK-DAG: %[[VAL_2:.*]] = arith.constant 0.000000e+00 : f64 +// CHECK-DAG: %[[VAL_3:.*]] = arith.constant 0 : index // CHECK: %[[VAL_4:.*]] = memref.alloc() : memref<16xindex> // CHECK: %[[VAL_5:.*]] = memref.cast %[[VAL_4]] : memref<16xindex> to memref<?xindex> // CHECK: linalg.fill ins(%[[VAL_3]] : index) outs(%[[VAL_4]] : memref<16xindex>) diff --git a/mlir/test/Integration/Dialect/SparseTensor/CPU/block.mlir b/mlir/test/Integration/Dialect/SparseTensor/CPU/block.mlir index d92165e..6468c4b 100644 --- a/mlir/test/Integration/Dialect/SparseTensor/CPU/block.mlir +++ b/mlir/test/Integration/Dialect/SparseTensor/CPU/block.mlir @@ -3,7 +3,7 @@ // // 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. +// 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} @@ -20,10 +20,13 @@ // REDEFINE: %{env} = TENSOR0="%mlir_src_dir/test/Integration/data/block.mtx" // RUN: %{compile} | env %{env} %{run} | FileCheck %s // -// TODO: enable! // Do the same run, but now with direct IR generation. // REDEFINE: %{sparsifier_opts} = enable-runtime-library=false -// R_UN: %{compile} | env %{env} %{run} | FileCheck %s +// RUN: %{compile} | env %{env} %{run} | FileCheck %s +// +// Do the same run, but now with direct IR generation and vectorization. +// REDEFINE: %{sparsifier_opts} = enable-runtime-library=false enable-buffer-initialization=true vl=2 reassociate-fp-reductions=true enable-index-optimizations=true +// RUN: %{compile} | env %{env} %{run} | FileCheck %s !Filename = !llvm.ptr diff --git a/mlir/test/Integration/Dialect/SparseTensor/GPU/CUDA/sparse-sddmm-lib.mlir b/mlir/test/Integration/Dialect/SparseTensor/GPU/CUDA/sparse-sddmm-lib.mlir index fb0da0c..db5c154 100644 --- a/mlir/test/Integration/Dialect/SparseTensor/GPU/CUDA/sparse-sddmm-lib.mlir +++ b/mlir/test/Integration/Dialect/SparseTensor/GPU/CUDA/sparse-sddmm-lib.mlir @@ -16,8 +16,7 @@ // // without RT lib: // -// TODO: make this work -// R_U_N: %{compile} enable-runtime-library=false" | %{run} +// RUN: %{compile} enable-runtime-library=false" | %{run} !Filename = !llvm.ptr |