[mlir][sparse] Fix a problem in handling data type conversion.

[bixia]

Summary

Previously, the genCast function generates arith.trunci for converting f32 to
i32. Fix the function to use mlir::convertScalarToDtype to correctly handle
conversion cases beyond index casting.

Add a test case for codegen the sparse_tensor.convert op.

[Peiming]

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

213

Thanks!

I was not aware the existence of this!

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

Value sparse_tensor::genCast(OpBuilder &builder, Location loc, Value value,
   const Type srcTp = value.getType();
   if (srcTp == dstTp)
     return value;
 
   // int <=> index
   if (srcTp.isa<IndexType>() || dstTp.isa<IndexType>())
     return builder.create<arith::IndexCastOp>(loc, dstTp, value);
 
-  const bool ext =
-      srcTp.getIntOrFloatBitWidth() < dstTp.getIntOrFloatBitWidth();
-
-  // float => float.
-  if (srcTp.isa<FloatType>() && dstTp.isa<FloatType>()) {
-    if (ext)
-      return builder.create<arith::ExtFOp>(loc, dstTp, value);
-    return builder.create<arith::TruncFOp>(loc, dstTp, value);
-  }
-
-  // int => int
-  const auto srcIntTp = srcTp.dyn_cast<IntegerType>();
-  if (srcIntTp && dstTp.isa<IntegerType>()) {
-    if (!ext)
-      return builder.create<arith::TruncIOp>(loc, dstTp, value);
-    if (srcIntTp.isUnsigned())
-      return builder.create<arith::ExtUIOp>(loc, dstTp, value);
-    if (srcIntTp.isSigned())
-      return builder.create<arith::ExtSIOp>(loc, dstTp, value);
-  }
-
-  llvm_unreachable("unhandled type casting");
+  const auto srcIntTp = srcTp.dyn_cast_or_null<IntegerType>();
+  const bool isUnsignedCast = srcIntTp ? srcIntTp.isUnsigned() : false;
+  return mlir::convertScalarToDtype(builder, loc, value, dstTp, isUnsignedCast);
 }
 
 mlir::Attribute mlir::sparse_tensor::getOneAttr(Builder &builder, Type tp) {
   if (tp.isa<FloatType>())
     return builder.getFloatAttr(tp, 1.0);
   if (tp.isa<IndexType>())
     return builder.getIndexAttr(1);
   if (auto intTp = tp.dyn_cast<IntegerType>())

mlir/test/Dialect/SparseTensor/codegen.mlir

 //  CHECK-SAME: %[[A4:.*]]: !sparse_tensor.storage_specifier
 //       CHECK: return  %[[A1]], %[[A2]], %[[A3]], %[[A4]] :
 //  CHECK-SAME:   memref<?xi32>, memref<?xi64>, memref<?xf32>, !sparse_tensor.storage_specifier
 func.func @sparse_nop_convert(%arg0: tensor<32xf32, #SparseVector>) -> tensor<?xf32, #SparseVector> {
   %0 = sparse_tensor.convert %arg0 : tensor<32xf32, #SparseVector> to tensor<?xf32, #SparseVector>
   return %0 : tensor<?xf32, #SparseVector>
 }
 
+// CHECK-LABEL: func.func @sparse_convert_element_type(
+//  CHECK-SAME: %[[A1:.*]]: memref<?xi32>,
+//  CHECK-SAME: %[[A2:.*]]: memref<?xi64>,
+//  CHECK-SAME: %[[A3:.*]]: memref<?xf32>,
+//  CHECK-SAME: %[[A4:.*]]: !sparse_tensor.storage_specifier
+//       CHECK: scf.for
+//       CHECK:   %[[FValue:.*]] = memref.load
+//       CHECK:   %[[IValue:.*]] = arith.fptosi %[[FValue]]
+//       CHECK:   memref.store %[[IValue]]
+//       CHECK: return  %{{.*}}, %{{.*}}, %{{.*}}, %[[A4]] :
+//  CHECK-SAME:   memref<?xi32>, memref<?xi64>, memref<?xi32>, !sparse_tensor.storage_specifier
+func.func @sparse_convert_element_type(%arg0: tensor<32xf32, #SparseVector>) -> tensor<?xi32, #SparseVector> {
+  %0 = sparse_tensor.convert %arg0 : tensor<32xf32, #SparseVector> to tensor<?xi32, #SparseVector>
+  return %0 : tensor<?xi32, #SparseVector>
+}
+
 // CHECK-LABEL: func.func @sparse_new_coo(
 // CHECK-SAME:  %[[A0:.*]]: !llvm.ptr<i8>) -> (memref<?xindex>, memref<?xindex>, memref<?xf32>, !sparse_tensor.storage_specifier<#sparse_tensor.encoding<{ dimLevelType = [ "compressed", "singleton" ] }>>) {
 //   CHECK-DAG: %[[A1:.*]] = arith.constant false
 //   CHECK-DAG: %[[A2:.*]] = arith.constant 1 : index
 //   CHECK-DAG: %[[A3:.*]] = arith.constant 0 : index
 //   CHECK-DAG: %[[A4:.*]] = arith.constant 2 : index
 //       CHECK: %[[A5:.*]] = call @createSparseTensorReader(%[[A0]])
 //       CHECK: %[[A6:.*]] = memref.alloca() : memref<2xindex>