17 #include "llvm/Support/FormatVariadic.h"
20 using namespace sparse_tensor;
30 for (
auto type : types) {
47 auto rtp = cast<ShapedType>(t);
51 extraTypes->push_back(rtp);
66 for (
auto type : types) {
69 toVals.push_back(fromVals[idx++]);
73 auto rtp = cast<RankedTensorType>(type);
79 inputs.push_back(fromVals[idx++]);
89 inputs.push_back(fromVals[idx++]);
90 }
else if (directOut) {
93 mem = builder.
create<sparse_tensor::ToPositionsOp>(loc, inputs[0],
96 mem = builder.
create<sparse_tensor::ToCoordinatesOp>(loc, inputs[0],
99 mem = builder.
create<sparse_tensor::ToValuesOp>(loc, inputs[0]);
100 toVals.push_back(mem);
102 ShapedType rtp = cast<ShapedType>(t);
104 inputs.push_back(extraVals[extra++]);
105 retTypes.push_back(rtp);
114 auto a = builder.
create<sparse_tensor::AssembleOp>(loc, rtp, inputs);
115 toVals.push_back(a.getResult());
116 }
else if (!directOut) {
119 unsigned len = retTypes.size();
120 retTypes.append(cntTypes);
122 builder.
create<sparse_tensor::DisassembleOp>(loc, retTypes, inputs);
123 for (
unsigned i = 0; i < len; i++)
124 toVals.push_back(d.getResult(i));
171 LogicalResult matchAndRewrite(func::FuncOp funcOp,
174 if (funcOp.isPrivate())
181 bool hasAnnotation =
false;
182 convTypes(hasAnnotation, funcOp.getArgumentTypes(), inputTypes,
nullptr,
184 convTypes(hasAnnotation, funcOp.getResultTypes(), outputTypes, &extraTypes,
192 auto orgName = funcOp.getName();
193 std::string wrapper = llvm::formatv(
"_internal_{0}", orgName).str();
194 funcOp.setName(wrapper);
199 ModuleOp modOp = funcOp->getParentOfType<ModuleOp>();
201 OpBuilder moduleBuilder(modOp.getBodyRegion());
202 unsigned extra = inputTypes.size();
203 inputTypes.append(extraTypes);
204 auto func = moduleBuilder.create<func::FuncOp>(
210 Block *body = func.addEntryBlock();
221 auto call = rewriter.
create<func::CallOp>(loc, funcOp.getResultTypes(), org,
226 convVals(rewriter, loc, funcOp.getResultTypes(), call.getResults(),
227 body->
getArguments(), outputs, extra,
false, directOut);
228 rewriter.
create<func::ReturnOp>(loc, outputs);
231 if (funcOp->getAttrOfType<UnitAttr>(
232 LLVM::LLVMDialect::getEmitCWrapperAttrName())) {
233 func->
setAttr(LLVM::LLVMDialect::getEmitCWrapperAttrName(),
235 funcOp->removeAttr(LLVM::LLVMDialect::getEmitCWrapperAttrName());
241 const bool directOut;
static void convTypes(bool &hasAnnotation, TypeRange types, SmallVectorImpl< Type > &convTypes, SmallVectorImpl< Type > *extraTypes, bool directOut)
static void convVals(OpBuilder &builder, Location loc, TypeRange types, ValueRange fromVals, ValueRange extraVals, SmallVectorImpl< Value > &toVals, unsigned extra, bool isIn, bool directOut)
Block represents an ordered list of Operations.
BlockArgListType getArguments()
This class defines the main interface for locations in MLIR and acts as a non-nullable wrapper around...
MLIRContext is the top-level object for a collection of MLIR operations.
RAII guard to reset the insertion point of the builder when destroyed.
This class helps build Operations.
void setInsertionPointToStart(Block *block)
Sets the insertion point to the start of the specified block.
Operation * create(const OperationState &state)
Creates an operation given the fields represented as an OperationState.
void setAttr(StringAttr name, Attribute value)
If the an attribute exists with the specified name, change it to the new value.
A special type of RewriterBase that coordinates the application of a rewrite pattern on the current I...
This class provides an abstraction over the various different ranges of value types.
Instances of the Type class are uniqued, have an immutable identifier and an optional mutable compone...
This class provides an abstraction over the different types of ranges over Values.
This class represents an instance of an SSA value in the MLIR system, representing a computable value...
A wrapper around RankedTensorType, which has three goals:
void foreachFieldAndTypeInSparseTensor(SparseTensorType, llvm::function_ref< bool(Type, FieldIndex, SparseTensorFieldKind, Level, LevelType)>)
unsigned FieldIndex
The type of field indices.
uint64_t Level
The type of level identifiers and level-ranks.
SparseTensorEncodingAttr getSparseTensorEncoding(Type type)
Convenience method to get a sparse encoding attribute from a type.
SparseTensorFieldKind
===-------------------------------------------------------------------—===// The sparse tensor storag...
Include the generated interface declarations.
const FrozenRewritePatternSet & patterns
void populateSparseAssembler(RewritePatternSet &patterns, bool directOut)
auto get(MLIRContext *context, Ts &&...params)
Helper method that injects context only if needed, this helps unify some of the attribute constructio...
OpRewritePattern is a wrapper around RewritePattern that allows for matching and rewriting against an...
OpRewritePattern(MLIRContext *context, PatternBenefit benefit=1, ArrayRef< StringRef > generatedNames={})
Patterns must specify the root operation name they match against, and can also specify the benefit of...
This enum defines all the sparse representations supportable by the SparseTensor dialect.