[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/include/mlir/Dialect/SparseTensor/Utils/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.
 //===----------------------------------------------------------------------===//
 
 /// Generates a 1-valued attribute of the given type.  This supports

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) {
-  Location loc = op->getLoc();
-  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");
+  // FIXME(wrengr): this one used to throw llvm_unreachable (and still does)
+  SmallString<9> name{"addElt", primaryTypeFunctionSuffix(eltType)};
   SmallVector<Value, 8> params;
   params.push_back(ptr);
   params.push_back(val);
   params.push_back(ind);
   params.push_back(perm);
   Type pTp = LLVM::LLVMPointerType::get(rewriter.getI8Type());
   auto fn = getFunc(op, name, pTp, params, /*emitCInterface=*/true);
-  rewriter.create<CallOp>(loc, pTp, fn, params);
+  rewriter.create<CallOp>(op->getLoc(), pTp, fn, params);
 }
 
 /// 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) {
   Location loc = op->getLoc();
   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");
+  // FIXME(wrengr): this one used to throw llvm_unreachable (and still does)
+  SmallString<10> name{"getNext", primaryTypeFunctionSuffix(elemTp)};
   SmallVector<Value, 3> params;
   params.push_back(iter);
   params.push_back(ind);
   params.push_back(elemPtr);
   Type i1 = rewriter.getI1Type();
   auto fn = getFunc(op, name, i1, params, /*emitCInterface=*/true);
   auto call = rewriter.create<CallOp>(loc, i1, fn, params);
   return call.getResult(0);

 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();
+    // FIXME(wrengr): this one used to return failure(); but now throws
+    // llvm_unreachable
+    SmallString<16> name{"sparsePointers", overheadTypeFunctionSuffix(ptrType)};
     auto fn = getFunc(op, name, resType, adaptor.getOperands(),
                       /*emitCInterface=*/true);
     rewriter.replaceOpWithNewOp<CallOp>(op, resType, fn, adaptor.getOperands());
     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();
+    // FIXME(wrengr): this one used to return failure(); but now throws
+    // llvm_unreachable
+    SmallString<15> name{"sparseIndices", overheadTypeFunctionSuffix(indType)};
     auto fn = getFunc(op, name, resType, adaptor.getOperands(),
                       /*emitCInterface=*/true);
     rewriter.replaceOpWithNewOp<CallOp>(op, resType, fn, adaptor.getOperands());
     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();
+    // FIXME(wrengr): this one used to return failure(); but now throws
+    // llvm_unreachable
+    SmallString<15> name{"sparseValues", primaryTypeFunctionSuffix(eltType)};
     auto fn = getFunc(op, name, resType, adaptor.getOperands(),
                       /*emitCInterface=*/true);
     rewriter.replaceOpWithNewOp<CallOp>(op, resType, fn, adaptor.getOperands());
     return success();
   }
 };
 
 /// Sparse conversion rule for tensor rematerialization.

 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");
+    // FIXME(wrengr): this one used to throw llvm_unreachable (and still does)
+    SmallString<12> name{"lexInsert", primaryTypeFunctionSuffix(eltType)};
     TypeRange noTp;
     auto fn =
         getFunc(op, name, noTp, adaptor.getOperands(), /*emitCInterface=*/true);
     rewriter.replaceOpWithNewOp<CallOp>(op, noTp, fn, adaptor.getOperands());
     return success();
   }
 };
 

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

case 16:
     return OverheadType::kU16;
   case 8:
     return OverheadType::kU8;
   case 0:
     return OverheadType::kIndex;
   }
 }
 
+OverheadType 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 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 getIndexOverheadType(Builder &builder,
                           const SparseTensorEncodingAttr &enc) {
   return getOverheadType(builder, overheadTypeEncoding(enc.getIndexBitWidth()));
 }
 
+StringRef 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 overheadTypeFunctionSuffix(Type tp) {
+  return overheadTypeFunctionSuffix(overheadTypeEncoding(tp));
+}
+
 PrimaryType 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 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 primaryTypeFunctionSuffix(Type elemTp) {
+  return primaryTypeFunctionSuffix(primaryTypeEncoding(elemTp));
+}
+
 DimLevelType dimLevelTypeEncoding(SparseTensorEncodingAttr::DimLevelType dlt) {
   switch (dlt) {
   case SparseTensorEncodingAttr::DimLevelType::Dense:
     return DimLevelType::kDense;
   case SparseTensorEncodingAttr::DimLevelType::Compressed:
     return DimLevelType::kCompressed;
   case SparseTensorEncodingAttr::DimLevelType::Singleton:
     return DimLevelType::kSingleton;