[mlir][sparse] Factoring out type-based function-name suffixes

[wrengr]

Summary

Depends On D115010

This changes a couple of places that used to return failure(); to now use llvm_unreachable() instead. However, Transforms/Sparsification.cpp should be doing the necessary type checks to ensure that those cases are in fact unreachable.

[aartbik]

It is okay to make all cases unreachable instead of failure provided that we check up front for the right types in the sparsification pass, which I was planning to do anyway to limit the possibilities of primary storage type up front
(note that later we may want to add e.g. complex and some of the lower fp types)

[wrengr]

In D115012#3170030, @aartbik wrote:

It is okay to make all cases unreachable instead of failure provided that we check up front for the right types in the sparsification pass, which I was planning to do anyway to limit the possibilities of primary storage type up front
(note that later we may want to add e.g. complex and some of the lower fp types)

Cool cool. I just wanted to be sure, since this differential changes some prior behavior (not that any of our invalidity tests seem to care); and cuz I want to maximize the helpfulness of error messages for users.

[wrengr]

Resolved all my fixme notes. Also removed a spurious code cleanup that's already handled by D115004.

[aartbik]

Good to go when parent of this chain goes in.

[wrengr]

rebasing to updated D115008

[wrengr]

Rebasing to updated D115008

mlir/lib/Dialect/SparseTensor/Transforms/CodegenUtils.h

 
 //===----------------------------------------------------------------------===//
 // ExecutionEngine/SparseTensorUtils helper functions.
 //===----------------------------------------------------------------------===//
 
 /// Converts an overhead storage bitwidth to its internal type-encoding.
 OverheadType overheadTypeEncoding(unsigned width);
 
+/// Converts an overhead storage type to its internal type-encoding.
+OverheadType overheadTypeEncoding(Type tp);
+
 /// Converts the internal type-encoding for overhead storage to an mlir::Type.
 Type getOverheadType(Builder &builder, OverheadType ot);
 
 /// Returns the mlir::Type for pointer overhead storage.
 Type getPointerOverheadType(Builder &builder,
                             const SparseTensorEncodingAttr &enc);
 
 /// Returns the mlir::Type for index overhead storage.
 Type getIndexOverheadType(Builder &builder,
                           const SparseTensorEncodingAttr &enc);
 
+/// Convert OverheadType to its function-name suffix.
+StringRef overheadTypeFunctionSuffix(OverheadType ot);
+
+/// Converts an overhead storage type to its function-name suffix.
+StringRef overheadTypeFunctionSuffix(Type overheadTp);
+
 /// Converts a primary storage type to its internal type-encoding.
 PrimaryType primaryTypeEncoding(Type elemTp);
 
+/// Convert PrimaryType to its function-name suffix.
+StringRef primaryTypeFunctionSuffix(PrimaryType pt);
+
+/// Converts a primary storage type to its function-name suffix.
+StringRef primaryTypeFunctionSuffix(Type elemTp);
+
 /// Converts the IR's dimension level type to its internal type-encoding.
 DimLevelType dimLevelTypeEncoding(SparseTensorEncodingAttr::DimLevelType dlt);
 
 //===----------------------------------------------------------------------===//
 // Misc code generators.
 //
 // TODO: both of these should move upstream to their respective classes.
 // Once RFCs have been created for those changes, list them here.

mlir/lib/Dialect/SparseTensor/Transforms/CodegenUtils.cpp

OverheadType mlir::sparse_tensor::overheadTypeEncoding(unsigned width) {
   case 8:
     return OverheadType::kU8;
   case 0:
     return OverheadType::kIndex;
   }
   llvm_unreachable("Unsupported overhead bitwidth");
 }
 
+OverheadType mlir::sparse_tensor::overheadTypeEncoding(Type tp) {
+  if (tp.isIndex())
+    return OverheadType::kIndex;
+  if (auto intTp = tp.dyn_cast<IntegerType>())
+    return overheadTypeEncoding(intTp.getWidth());
+  llvm_unreachable("Unknown overhead type");
+}
+
 Type mlir::sparse_tensor::getOverheadType(Builder &builder, OverheadType ot) {
   switch (ot) {
   case OverheadType::kIndex:
     return builder.getIndexType();
   case OverheadType::kU64:
     return builder.getIntegerType(64);
   case OverheadType::kU32:
     return builder.getIntegerType(32);
return getOverheadType(builder,
                          overheadTypeEncoding(enc.getPointerBitWidth()));
 }
 
 Type mlir::sparse_tensor::getIndexOverheadType(
     Builder &builder, const SparseTensorEncodingAttr &enc) {
   return getOverheadType(builder, overheadTypeEncoding(enc.getIndexBitWidth()));
 }
 
+StringRef mlir::sparse_tensor::overheadTypeFunctionSuffix(OverheadType ot) {
+  switch (ot) {
+  case OverheadType::kIndex:
+    return "";
+  case OverheadType::kU64:
+    return "64";
+  case OverheadType::kU32:
+    return "32";
+  case OverheadType::kU16:
+    return "16";
+  case OverheadType::kU8:
+    return "8";
+  }
+  llvm_unreachable("Unknown OverheadType");
+}
+
+StringRef mlir::sparse_tensor::overheadTypeFunctionSuffix(Type tp) {
+  return overheadTypeFunctionSuffix(overheadTypeEncoding(tp));
+}
+
 PrimaryType mlir::sparse_tensor::primaryTypeEncoding(Type elemTp) {
   if (elemTp.isF64())
     return PrimaryType::kF64;
   if (elemTp.isF32())
     return PrimaryType::kF32;
   if (elemTp.isInteger(64))
     return PrimaryType::kI64;
   if (elemTp.isInteger(32))
     return PrimaryType::kI32;
   if (elemTp.isInteger(16))
     return PrimaryType::kI16;
   if (elemTp.isInteger(8))
     return PrimaryType::kI8;
   llvm_unreachable("Unknown primary type");
 }
 
+StringRef mlir::sparse_tensor::primaryTypeFunctionSuffix(PrimaryType pt) {
+  switch (pt) {
+  case PrimaryType::kF64:
+    return "F64";
+  case PrimaryType::kF32:
+    return "F32";
+  case PrimaryType::kI64:
+    return "I64";
+  case PrimaryType::kI32:
+    return "I32";
+  case PrimaryType::kI16:
+    return "I16";
+  case PrimaryType::kI8:
+    return "I8";
+  }
+  llvm_unreachable("Unknown PrimaryType");
+}
+
+StringRef mlir::sparse_tensor::primaryTypeFunctionSuffix(Type elemTp) {
+  return primaryTypeFunctionSuffix(primaryTypeEncoding(elemTp));
+}
+
 DimLevelType mlir::sparse_tensor::dimLevelTypeEncoding(
     SparseTensorEncodingAttr::DimLevelType dlt) {
   switch (dlt) {
   case SparseTensorEncodingAttr::DimLevelType::Dense:
     return DimLevelType::kDense;
   case SparseTensorEncodingAttr::DimLevelType::Compressed:
     return DimLevelType::kCompressed;
   case SparseTensorEncodingAttr::DimLevelType::Singleton:

mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorConversion.cpp

 /// Generates a call that adds one element to a coordinate scheme.
 /// In particular, this generates code like the following:
 ///   val = a[i1,..,ik];
 ///   if val != 0
 ///     t->add(val, [i1,..,ik], [p1,..,pk]);
 static void genAddEltCall(ConversionPatternRewriter &rewriter, Operation *op,
                           Type eltType, Value ptr, Value val, Value ind,
                           Value perm) {
-  StringRef name;
-  if (eltType.isF64())
-    name = "addEltF64";
-  else if (eltType.isF32())
-    name = "addEltF32";
-  else if (eltType.isInteger(64))
-    name = "addEltI64";
-  else if (eltType.isInteger(32))
-    name = "addEltI32";
-  else if (eltType.isInteger(16))
-    name = "addEltI16";
-  else if (eltType.isInteger(8))
-    name = "addEltI8";
-  else
-    llvm_unreachable("Unknown element type");
+  SmallString<9> name{"addElt", primaryTypeFunctionSuffix(eltType)};
   SmallVector<Value, 4> params{ptr, val, ind, perm};
   Type pTp = getOpaquePointerType(rewriter);
   createFuncCall(rewriter, op, name, pTp, params, EmitCInterface::On);
 }
 
 /// Generates a call to `iter->getNext()`.  If there is a next element,
 /// then it is copied into the out-parameters `ind` and `elemPtr`,
 /// and the return value is true.  If there isn't a next element, then
 /// the memory for `iter` is freed and the return value is false.
 static Value genGetNextCall(ConversionPatternRewriter &rewriter, Operation *op,
                             Value iter, Value ind, Value elemPtr) {
   Type elemTp = elemPtr.getType().cast<ShapedType>().getElementType();
-  StringRef name;
-  if (elemTp.isF64())
-    name = "getNextF64";
-  else if (elemTp.isF32())
-    name = "getNextF32";
-  else if (elemTp.isInteger(64))
-    name = "getNextI64";
-  else if (elemTp.isInteger(32))
-    name = "getNextI32";
-  else if (elemTp.isInteger(16))
-    name = "getNextI16";
-  else if (elemTp.isInteger(8))
-    name = "getNextI8";
-  else
-    llvm_unreachable("Unknown element type");
+  SmallString<10> name{"getNext", primaryTypeFunctionSuffix(elemTp)};
   SmallVector<Value, 3> params{iter, ind, elemPtr};
   Type i1 = rewriter.getI1Type();
   return createFuncCall(rewriter, op, name, i1, params, EmitCInterface::On)
       .getResult(0);
 }
 
 /// If the tensor is a sparse constant, generates and returns the pair of
 /// the constants for the indices and the values.

 class SparseTensorToPointersConverter
     : public OpConversionPattern<ToPointersOp> {
 public:
   using OpConversionPattern::OpConversionPattern;
   LogicalResult
   matchAndRewrite(ToPointersOp op, OpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override {
     Type resType = op.getType();
-    Type eltType = resType.cast<ShapedType>().getElementType();
-    StringRef name;
-    if (eltType.isIndex())
-      name = "sparsePointers";
-    else if (eltType.isInteger(64))
-      name = "sparsePointers64";
-    else if (eltType.isInteger(32))
-      name = "sparsePointers32";
-    else if (eltType.isInteger(16))
-      name = "sparsePointers16";
-    else if (eltType.isInteger(8))
-      name = "sparsePointers8";
-    else
-      return failure();
+    Type ptrType = resType.cast<ShapedType>().getElementType();
+    SmallString<16> name{"sparsePointers", overheadTypeFunctionSuffix(ptrType)};
     replaceOpWithFuncCall(rewriter, op, name, resType, adaptor.getOperands(),
                           EmitCInterface::On);
     return success();
   }
 };
 
 /// Sparse conversion rule for index accesses.
 class SparseTensorToIndicesConverter : public OpConversionPattern<ToIndicesOp> {
 public:
   using OpConversionPattern::OpConversionPattern;
   LogicalResult
   matchAndRewrite(ToIndicesOp op, OpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override {
     Type resType = op.getType();
-    Type eltType = resType.cast<ShapedType>().getElementType();
-    StringRef name;
-    if (eltType.isIndex())
-      name = "sparseIndices";
-    else if (eltType.isInteger(64))
-      name = "sparseIndices64";
-    else if (eltType.isInteger(32))
-      name = "sparseIndices32";
-    else if (eltType.isInteger(16))
-      name = "sparseIndices16";
-    else if (eltType.isInteger(8))
-      name = "sparseIndices8";
-    else
-      return failure();
+    Type indType = resType.cast<ShapedType>().getElementType();
+    SmallString<15> name{"sparseIndices", overheadTypeFunctionSuffix(indType)};
     replaceOpWithFuncCall(rewriter, op, name, resType, adaptor.getOperands(),
                           EmitCInterface::On);
     return success();
   }
 };
 
 /// Sparse conversion rule for value accesses.
 class SparseTensorToValuesConverter : public OpConversionPattern<ToValuesOp> {
 public:
   using OpConversionPattern::OpConversionPattern;
   LogicalResult
   matchAndRewrite(ToValuesOp op, OpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override {
     Type resType = op.getType();
     Type eltType = resType.cast<ShapedType>().getElementType();
-    StringRef name;
-    if (eltType.isF64())
-      name = "sparseValuesF64";
-    else if (eltType.isF32())
-      name = "sparseValuesF32";
-    else if (eltType.isInteger(64))
-      name = "sparseValuesI64";
-    else if (eltType.isInteger(32))
-      name = "sparseValuesI32";
-    else if (eltType.isInteger(16))
-      name = "sparseValuesI16";
-    else if (eltType.isInteger(8))
-      name = "sparseValuesI8";
-    else
-      return failure();
+    SmallString<15> name{"sparseValues", primaryTypeFunctionSuffix(eltType)};
     replaceOpWithFuncCall(rewriter, op, name, resType, adaptor.getOperands(),
                           EmitCInterface::On);
     return success();
   }
 };
 
 /// Sparse conversion rule for tensor rematerialization.
 class SparseTensorLoadConverter : public OpConversionPattern<LoadOp> {

 
 /// Sparse conversion rule for inserting in lexicographic index order.
 class SparseTensorLexInsertConverter : public OpConversionPattern<LexInsertOp> {
 public:
   using OpConversionPattern::OpConversionPattern;
   LogicalResult
   matchAndRewrite(LexInsertOp op, OpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override {
-    Type srcType = op.tensor().getType();
-    Type eltType = srcType.cast<ShapedType>().getElementType();
-    StringRef name;
-    if (eltType.isF64())
-      name = "lexInsertF64";
-    else if (eltType.isF32())
-      name = "lexInsertF32";
-    else if (eltType.isInteger(64))
-      name = "lexInsertI64";
-    else if (eltType.isInteger(32))
-      name = "lexInsertI32";
-    else if (eltType.isInteger(16))
-      name = "lexInsertI16";
-    else if (eltType.isInteger(8))
-      name = "lexInsertI8";
-    else
-      llvm_unreachable("Unknown element type");
+    Type elemTp = op.tensor().getType().cast<ShapedType>().getElementType();
+    SmallString<12> name{"lexInsert", primaryTypeFunctionSuffix(elemTp)};
     TypeRange noTp;
     replaceOpWithFuncCall(rewriter, op, name, noTp, adaptor.getOperands(),
                           EmitCInterface::On);
     return success();
   }
 };
 
 class SparseTensorExpandConverter : public OpConversionPattern<ExpandOp> {
public:
   using OpConversionPattern::OpConversionPattern;
   LogicalResult
   matchAndRewrite(CompressOp op, OpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override {
     // Note that this method call resets the values/filled-switch back to
     // all-zero/false by only iterating over the set elements, so the
     // complexity remains proportional to the sparsity of the expanded
     // access pattern.
-    Type srcType = op.tensor().getType();
-    Type eltType = srcType.cast<ShapedType>().getElementType();
-    StringRef name;
-    if (eltType.isF64())
-      name = "expInsertF64";
-    else if (eltType.isF32())
-      name = "expInsertF32";
-    else if (eltType.isInteger(64))
-      name = "expInsertI64";
-    else if (eltType.isInteger(32))
-      name = "expInsertI32";
-    else if (eltType.isInteger(16))
-      name = "expInsertI16";
-    else if (eltType.isInteger(8))
-      name = "expInsertI8";
-    else
-      return failure();
+    Type elemTp = op.tensor().getType().cast<ShapedType>().getElementType();
+    SmallString<12> name{"expInsert", primaryTypeFunctionSuffix(elemTp)};
     TypeRange noTp;
     replaceOpWithFuncCall(rewriter, op, name, noTp, adaptor.getOperands(),
                           EmitCInterface::On);
     return success();
   }
 };
 
 } // namespace