diff --git a/mlir/lib/Dialect/SparseTensor/Transforms/CodegenUtils.h b/mlir/lib/Dialect/SparseTensor/Transforms/CodegenUtils.h --- a/mlir/lib/Dialect/SparseTensor/Transforms/CodegenUtils.h +++ b/mlir/lib/Dialect/SparseTensor/Transforms/CodegenUtils.h @@ -259,6 +259,45 @@ /// of the given type, and returns the `memref<$tp>`. Value genAllocaScalar(OpBuilder &builder, Location loc, Type tp); +/// Generates code to allocate a buffer of the given type, and zero +/// initialize it. If the buffer type has any dynamic sizes, then the +/// `sizes` parameter should be as filled by sizesFromPtr(); that way +/// we can reuse the genDimSizeCall() results generated by sizesFromPtr(). +Value allocDenseTensor(OpBuilder &builder, Location loc, + RankedTensorType tensorTp, ValueRange sizes); + +/// Generates the code to read the value from tensor[ivs].The generated code +/// looks like the following and the insertion point after this routine is +/// inside the if-then branch behind the assignment to ind. +/// if (tensor[ivs] != 0) +/// insert_point +Value genValueForDense(OpBuilder &builder, Location loc, Value tensor, + ValueRange ivs); + +/// Generates the loop structure to iterate over a dense tensor or a sparse +/// tensor constant to support the lowering of dense-to-sparse convert operator. +// +// The loop to iterate a dense tensor: +// for i1 in dim1 +// .. +// for ik in dimk +// val = a[i1,..,ik] +// if val != 0 +// loop-body +// +// The loop to iterate a sparse tensor constant: +// for i in range(NNZ) +// val = values[i] +// [i1,..,ik] = indices[i] +// loop-body +void genDenseTensorOrSparseConstantIterLoop( + OpBuilder &builder, Location loc, Value src, unsigned rank, + function_ref bodyBuilder); + +/// Populates given sizes array from dense tensor or sparse tensor constant. +void sizesFromSrc(OpBuilder &builder, SmallVector &sizes, + Location loc, Value src); + //===----------------------------------------------------------------------===// // Inlined constant generators. // diff --git a/mlir/lib/Dialect/SparseTensor/Transforms/CodegenUtils.cpp b/mlir/lib/Dialect/SparseTensor/Transforms/CodegenUtils.cpp --- a/mlir/lib/Dialect/SparseTensor/Transforms/CodegenUtils.cpp +++ b/mlir/lib/Dialect/SparseTensor/Transforms/CodegenUtils.cpp @@ -13,6 +13,7 @@ #include "mlir/Dialect/Linalg/IR/Linalg.h" #include "mlir/Dialect/Linalg/Utils/Utils.h" #include "mlir/Dialect/MemRef/IR/MemRef.h" +#include "mlir/Dialect/Tensor/IR/Tensor.h" #include "mlir/Dialect/Vector/IR/VectorOps.h" #include "mlir/IR/Matchers.h" #include "mlir/IR/Types.h" @@ -40,6 +41,56 @@ return load; } +/// If the tensor is a sparse constant, generates and returns the pair of +/// the constants for the indices and the values. +static Optional> +genSplitSparseConstant(OpBuilder &builder, Location loc, Value tensor) { + if (auto constOp = tensor.getDefiningOp()) { + if (auto attr = constOp.getValue().dyn_cast()) { + DenseElementsAttr indicesAttr = attr.getIndices(); + Value indices = builder.create(loc, indicesAttr); + DenseElementsAttr valuesAttr = attr.getValues(); + Value values = builder.create(loc, valuesAttr); + return std::make_pair(indices, values); + } + } + return {}; +} + +/// Generates the code to copy the index at indices[ivs] to ind, and return +/// the value at value[ivs]. +static Value genIndexAndValueForSparse(OpBuilder &builder, Location loc, + Value indices, Value values, + SmallVectorImpl &indicesArray, + ValueRange ivs, unsigned rank) { + for (unsigned i = 0; i < rank; i++) { + Value idx = constantIndex(builder, loc, i); + Value val = builder.create(loc, indices, + ValueRange{ivs[0], idx}); + val = builder.create(loc, builder.getIndexType(), val); + // builder.create(loc, val, ind, idx); + indicesArray.push_back(val); + } + return builder.create(loc, values, ivs[0]); +} + +/// Generates the code to read the value from tensor[ivs], and conditionally +/// stores the indices ivs to the memory in ind. The generated code looks like +/// the following and the insertion point after this routine is inside the +/// if-then branch behind the assignment to ind. This is to ensure that the +/// code that uses the ind, such as a addEltX call generated after, is inside +/// the if-then branch. +/// if (tensor[ivs] != 0) +/// ind = ivs +static Value genIndexAndValueForDense(OpBuilder &builder, Location loc, + Value tensor, + SmallVectorImpl &indicesArray, + ValueRange ivs) { + Value val = genValueForDense(builder, loc, tensor, ivs); + indicesArray.append(ivs.begin(), ivs.end()); + return val; +} + //===----------------------------------------------------------------------===// // Sparse tensor loop emitter class implementations //===----------------------------------------------------------------------===// @@ -574,3 +625,79 @@ Type tp) { return builder.create(loc, MemRefType::get({}, tp)); } + +Value mlir::sparse_tensor::allocDenseTensor(OpBuilder &builder, Location loc, + RankedTensorType tensorTp, + ValueRange sizes) { + Type elemTp = tensorTp.getElementType(); + auto shape = tensorTp.getShape(); + auto memTp = MemRefType::get(shape, elemTp); + SmallVector dynamicSizes; + for (unsigned i = 0, rank = tensorTp.getRank(); i < rank; i++) { + if (shape[i] == ShapedType::kDynamicSize) + dynamicSizes.push_back(sizes[i]); + } + Value mem = builder.create(loc, memTp, dynamicSizes); + Value zero = constantZero(builder, loc, elemTp); + builder.create(loc, ValueRange{zero}, ValueRange{mem}); + return mem; +} + +Value mlir::sparse_tensor::genValueForDense(OpBuilder &builder, Location loc, + Value tensor, ValueRange ivs) { + Value val = builder.create(loc, tensor, ivs); + Value cond = genIsNonzero(builder, loc, val); + scf::IfOp ifOp = builder.create(loc, cond, /*else*/ false); + builder.setInsertionPointToStart(&ifOp.getThenRegion().front()); + return val; +} + +void mlir::sparse_tensor::genDenseTensorOrSparseConstantIterLoop( + OpBuilder &builder, Location loc, Value src, unsigned rank, + function_ref bodyBuilder) { + SmallVector indicesArray; + SmallVector lo; + SmallVector hi; + SmallVector st; + Value zero = constantIndex(builder, loc, 0); + Value one = constantIndex(builder, loc, 1); + auto indicesValues = genSplitSparseConstant(builder, loc, src); + bool isCOOConstant = indicesValues.has_value(); + Value indices; + Value values; + if (isCOOConstant) { + indices = indicesValues->first; + values = indicesValues->second; + lo.push_back(zero); + hi.push_back(linalg::createOrFoldDimOp(builder, loc, values, 0)); + st.push_back(one); + } else { + for (unsigned i = 0; i < rank; i++) { + lo.push_back(zero); + hi.push_back(linalg::createOrFoldDimOp(builder, loc, src, i)); + st.push_back(one); + } + } + + scf::buildLoopNest( + builder, loc, lo, hi, st, {}, + [&](OpBuilder &builder, Location loc, ValueRange ivs, + ValueRange args) -> scf::ValueVector { + Value val; + if (isCOOConstant) + val = genIndexAndValueForSparse(builder, loc, indices, values, + indicesArray, ivs, rank); + else + val = genIndexAndValueForDense(builder, loc, src, indicesArray, ivs); + bodyBuilder(builder, loc, val, indicesArray); + return {}; + }); +} + +void mlir::sparse_tensor::sizesFromSrc(OpBuilder &builder, + SmallVector &sizes, + Location loc, Value src) { + unsigned rank = src.getType().cast().getRank(); + for (unsigned i = 0; i < rank; i++) + sizes.push_back(linalg::createOrFoldDimOp(builder, loc, src, i)); +} diff --git a/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorConversion.cpp b/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorConversion.cpp --- a/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorConversion.cpp +++ b/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorConversion.cpp @@ -110,14 +110,6 @@ } } -/// Populates given sizes array from source. -static void sizesFromSrc(OpBuilder &builder, SmallVector &sizes, - Location loc, Value src) { - unsigned rank = src.getType().cast().getRank(); - for (unsigned i = 0; i < rank; i++) - sizes.push_back(linalg::createOrFoldDimOp(builder, loc, src, i)); -} - /// Populates the given sizes array for concatenation from type (for static /// sizes) and from an already-converted opaque pointer source (for dynamic /// sizes). @@ -213,38 +205,6 @@ params.push_back(ptr); } -/// Generates the code to read the value from tensor[ivs].The generated code -/// looks like the following and the insertion point after this routine is -/// inside the if-then branch behind the assignment to ind. -/// if (tensor[ivs] != 0) -/// insert_point -static Value genValueForDense(OpBuilder &builder, Location loc, Value tensor, - ValueRange ivs) { - Value val = builder.create(loc, tensor, ivs); - Value cond = genIsNonzero(builder, loc, val); - scf::IfOp ifOp = builder.create(loc, cond, /*else*/ false); - builder.setInsertionPointToStart(&ifOp.getThenRegion().front()); - return val; -} - -/// Generates the code to read the value from tensor[ivs], and conditionally -/// stores the indices ivs to the memory in ind. The generated code looks like -/// the following and the insertion point after this routine is inside the -/// if-then branch behind the assignment to ind. This is to ensure that the -/// addEltX call generated after is inside the if-then branch. -/// if (tensor[ivs] != 0) -/// ind = ivs -static Value genIndexAndValueForDense(OpBuilder &builder, Location loc, - Value tensor, Value ind, ValueRange ivs) { - Value val = genValueForDense(builder, loc, tensor, ivs); - unsigned i = 0; - for (auto iv : ivs) { - Value idx = constantIndex(builder, loc, i++); - builder.create(loc, iv, ind, idx); - } - return val; -} - /// Generates a call to release/delete a `SparseTensorCOO`. static void genDelCOOCall(OpBuilder &builder, Location loc, Type elemTp, Value coo) { @@ -287,57 +247,6 @@ .getResult(0); } -/// If the tensor is a sparse constant, generates and returns the pair of -/// the constants for the indices and the values. -static Optional> -genSplitSparseConstant(OpBuilder &builder, Location loc, Value tensor) { - if (auto constOp = tensor.getDefiningOp()) { - if (auto attr = constOp.getValue().dyn_cast()) { - DenseElementsAttr indicesAttr = attr.getIndices(); - Value indices = builder.create(loc, indicesAttr); - DenseElementsAttr valuesAttr = attr.getValues(); - Value values = builder.create(loc, valuesAttr); - return std::make_pair(indices, values); - } - } - return {}; -} - -/// Generates the code to copy the index at indices[ivs] to ind, and return -/// the value at value[ivs]. -static Value genIndexAndValueForSparse(OpBuilder &builder, Location loc, - Value indices, Value values, Value ind, - ValueRange ivs, unsigned rank) { - for (unsigned i = 0; i < rank; i++) { - Value idx = constantIndex(builder, loc, i); - Value val = builder.create(loc, indices, - ValueRange{ivs[0], idx}); - val = builder.create(loc, builder.getIndexType(), val); - builder.create(loc, val, ind, idx); - } - return builder.create(loc, values, ivs[0]); -} - -/// Generates code to allocate a buffer of the given type, and zero -/// initialize it. If the buffer type has any dynamic sizes, then the -/// `sizes` parameter should be as filled by sizesFromPtr(); that way -/// we can reuse the genDimSizeCall() results generated by sizesFromPtr(). -static Value allocDenseTensor(OpBuilder &builder, Location loc, - RankedTensorType tensorTp, ValueRange sizes) { - Type elemTp = tensorTp.getElementType(); - auto shape = tensorTp.getShape(); - auto memTp = MemRefType::get(shape, elemTp); - SmallVector dynamicSizes; - for (unsigned i = 0, rank = tensorTp.getRank(); i < rank; i++) { - if (shape[i] == ShapedType::kDynamicSize) - dynamicSizes.push_back(sizes[i]); - } - Value mem = builder.create(loc, memTp, dynamicSizes); - Value zero = constantZero(builder, loc, elemTp); - builder.create(loc, ValueRange{zero}, ValueRange{mem}); - return mem; -} - /// Generates code to deallocate a dense buffer. static void deallocDenseTensor(OpBuilder &builder, Location loc, Value buffer) { builder.create(loc, buffer); @@ -905,43 +814,17 @@ Value coo = genNewCall(rewriter, loc, params); Value ind = genAlloca(rewriter, loc, rank, rewriter.getIndexType()); Value perm = params[2]; - SmallVector lo; - SmallVector hi; - SmallVector st; - Value zero = constantIndex(rewriter, loc, 0); - Value one = constantIndex(rewriter, loc, 1); - auto indicesValues = genSplitSparseConstant(rewriter, loc, src); - bool isCOOConstant = indicesValues.has_value(); - Value indices; - Value values; - if (isCOOConstant) { - indices = indicesValues->first; - values = indicesValues->second; - lo.push_back(zero); - hi.push_back(linalg::createOrFoldDimOp(rewriter, loc, values, 0)); - st.push_back(one); - } else { - for (unsigned i = 0; i < rank; i++) { - lo.push_back(zero); - hi.push_back(linalg::createOrFoldDimOp(rewriter, loc, src, i)); - st.push_back(one); - } - } Type eltType = stp.getElementType(); Value elemPtr = genAllocaScalar(rewriter, loc, eltType); - scf::buildLoopNest( - rewriter, op.getLoc(), lo, hi, st, {}, - [&](OpBuilder &builder, Location loc, ValueRange ivs, - ValueRange args) -> scf::ValueVector { - Value val; - if (isCOOConstant) - val = genIndexAndValueForSparse(rewriter, loc, indices, values, ind, - ivs, rank); - else - val = genIndexAndValueForDense(rewriter, loc, src, ind, ivs); + genDenseTensorOrSparseConstantIterLoop( + rewriter, loc, src, rank, + [&](OpBuilder &builder, Location loc, Value val, ValueRange indices) { + for (unsigned i = 0; i < rank; i++) { + Value idx = constantIndex(builder, loc, i); + builder.create(loc, indices[i], ind, idx); + } builder.create(loc, val, elemPtr); - genAddEltCall(rewriter, loc, eltType, coo, elemPtr, ind, perm); - return {}; + genAddEltCall(builder, loc, eltType, coo, elemPtr, ind, perm); }); // Final call to construct sparse tensor storage. params[6] = constantAction(rewriter, loc, Action::kFromCOO); diff --git a/mlir/test/Dialect/SparseTensor/convert_dense2sparse.mlir b/mlir/test/Dialect/SparseTensor/convert_dense2sparse.mlir --- a/mlir/test/Dialect/SparseTensor/convert_dense2sparse.mlir +++ b/mlir/test/Dialect/SparseTensor/convert_dense2sparse.mlir @@ -123,9 +123,9 @@ // CHECK: %[[N:.*]] = memref.cast %[[M]] : memref<2xindex> to memref // CHECK: %[[BUF:.*]] = memref.alloca() : memref // CHECK: scf.for %[[I:.*]] = %[[C0]] to %[[C2]] step %[[C1]] { -// CHECK: memref.store %{{.*}}, %[[M]][%[[C0]]] : memref<2xindex> -// CHECK: memref.store %{{.*}}, %[[M]][%[[C1]]] : memref<2xindex> -// CHECK: %[[V:.*]] = tensor.extract %{{.*}}[%[[I]]] : tensor<2xf32> +// CHECK-DAG: memref.store %{{.*}}, %[[M]][%[[C0]]] : memref<2xindex> +// CHECK-DAG: memref.store %{{.*}}, %[[M]][%[[C1]]] : memref<2xindex> +// CHECK-DAG: %[[V:.*]] = tensor.extract %{{.*}}[%[[I]]] : tensor<2xf32> // CHECK: memref.store %[[V]], %[[BUF]][] : memref // CHECK: call @addEltF32(%{{.*}}, %[[BUF]], %[[N]], %{{.*}}) // CHECK: }