[mlir][sparse] add casts to operations to lattice and exp builders

[aartbik]

Summary

Further enhance the set of operations that can be handled by the sparse compiler

[bixia]

mlir/include/mlir/Dialect/SparseTensor/Utils/Merger.h

86–88

Would it be more straightforward to say "or the destination type of a cast operation"?

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

669

Per offline conversation, it is cleaner to get vl from the type of v0.

[aartbik]

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

669

Yes, it makes this a bit more "black box" with less input from the outside world. Thanks!

[aartbik]

infer vl from local information only

mlir/include/mlir/Dialect/SparseTensor/Utils/Merger.h

enum Kind {
   kTensor = 0,
   kInvariant,
   // Unary operations.
   kAbsF,
   kCeilF,
   kFloorF,
   kNegF,
   kNegI,
+  kTruncF,
+  kExtF,
+  kCastFS, // signed
+  kCastFU, // unsigned
+  kCastSF, // signed
+  kCastUF, // unsigned
+  kCastS,  // signed
+  kCastU,  // unsigned
+  kTruncI,
+  kBitCast,
   // Binary operations.
   kMulF,
   kMulI,
   kDivF,
   kDivS, // signed
   kDivU, // unsigned
   kAddF,
   kAddI,
struct TensorExp {
   union {
     /// Expressions representing tensors simply have a tensor number.
     unsigned tensor;
 
     /// Tensor operations hold the indices of their children.
     Children children;
   };
 
-  /// Direct link to IR for an invariant. During code generation,
-  /// field is used to cache "hoisted" loop invariant tensor loads.
+  /// Direct link to IR for an invariant or the destination value (to
+  /// infer destination type) of a cast operation During code generation,
+  /// this field may be used to cache "hoisted" loop invariant tensor loads.
   Value val;
 };
 
 /// Lattice point. Each lattice point consists of a conjunction of tensor
 /// loop indices (encoded in a bitvector) and the index of the corresponding
 /// tensor expression.
 struct LatPoint {
   LatPoint(unsigned n, unsigned e, unsigned b);
public:
   /// additional synthetic tensor at the end of this set to represent all
   /// invariant expressions in the kernel.
   Merger(unsigned t, unsigned l)
       : outTensor(t - 1), syntheticTensor(t), numTensors(t + 1), numLoops(l),
         dims(t + 1, std::vector<Dim>(l, Dim::kUndef)) {}
 
   /// Adds a tensor expression. Returns its index.
   unsigned addExp(Kind k, unsigned e0, unsigned e1 = -1u, Value v = Value());
+  unsigned addExp(Kind k, unsigned e, Value v) { return addExp(k, e, -1u, v); }
   unsigned addExp(Kind k, Value v) { return addExp(k, -1u, -1u, v); }
 
   /// Adds an iteration lattice point. Returns its index.
   unsigned addLat(unsigned t, unsigned i, unsigned e);
 
   /// Adds a new, initially empty, set. Returns its index.
   unsigned addSet();
 
public:
 
   /// Disjunctive merge of two lattice sets L0 and L1 is (L0 /\_op L1, L0, L1).
   /// Returns the index of the new set.
   unsigned takeDisj(Kind kind, unsigned s0, unsigned s1);
 
   /// Maps the unary operator over the lattice set of the operand, i.e. each
   /// lattice point on an expression E is simply copied over, but with OP E
   /// as new expression. Returns the index of the new set.
-  unsigned mapSet(Kind kind, unsigned s0);
+  unsigned mapSet(Kind kind, unsigned s0, Value v = Value());
 
   /// Optimizes the iteration lattice points in the given set. This
   /// method should be called right before code generation to avoid
   /// generating redundant loops and conditions.
   unsigned optimizeSet(unsigned s0);
 
   /// Simplifies the conditions in a conjunction of a given lattice point
   /// within the given set using just two basic rules:
#endif
 
   /// Rebuilds SSA format from a tensor expression.
   Value buildExp(PatternRewriter &rewriter, Location loc, unsigned e, Value v0,
                  Value v1);
 
 private:
   bool maybeZero(unsigned e) const;
   bool isInvariant(unsigned e) const;
+  Type inferType(unsigned e, Value src);
 
   /// Traverses the SSA tree (possibly a DAG) to build a tensor expression.
   Optional<unsigned> buildTensorExp(linalg::GenericOp op, Value v);
 
   const unsigned outTensor;
   const unsigned syntheticTensor;
   const unsigned numTensors;
   const unsigned numLoops;

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

 //===- Merger.cpp - Implementation of iteration lattices ------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Dialect/SparseTensor/Utils/Merger.h"
 
 #include "mlir/IR/Operation.h"
 #include "llvm/Support/Debug.h"
 
 namespace mlir {
 namespace sparse_tensor {
 
-//
+//===----------------------------------------------------------------------===//
 // Constructors.
-//
+//===----------------------------------------------------------------------===//
 
 TensorExp::TensorExp(Kind k, unsigned x, unsigned y, Value v)
     : kind(k), val(v) {
   switch (kind) {
   case kTensor:
     assert(x != -1u && y == -1u && !v);
     tensor = x;
     break;
   case kInvariant:
     assert(x == -1u && y == -1u && v);
     break;
   case kAbsF:
   case kCeilF:
   case kFloorF:
   case kNegF:
   case kNegI:
     assert(x != -1u && y == -1u && !v);
     children.e0 = x;
     children.e1 = y;
     break;
+  case kTruncF:
+  case kExtF:
+  case kCastFS:
+  case kCastFU:
+  case kCastSF:
+  case kCastUF:
+  case kCastS:
+  case kCastU:
+  case kTruncI:
+  case kBitCast:
+    assert(x != -1u && y == -1u && v);
+    children.e0 = x;
+    children.e1 = y;
+    break;
   default:
     assert(x != -1u && y != -1u && !v);
     children.e0 = x;
     children.e1 = y;
     break;
   }
 }
 
 LatPoint::LatPoint(unsigned n, unsigned e, unsigned b)
     : bits(n, false), simple(), exp(e) {
   bits.set(b);
 }
 
 LatPoint::LatPoint(const llvm::BitVector &b, unsigned e)
     : bits(b), simple(), exp(e) {}
 
-//
+//===----------------------------------------------------------------------===//
 // Lattice methods.
-//
+//===----------------------------------------------------------------------===//
 
 unsigned Merger::addExp(Kind k, unsigned e0, unsigned e1, Value v) {
   unsigned e = tensorExps.size();
   tensorExps.push_back(TensorExp(k, e0, e1, v));
   return e;
 }
 
 unsigned Merger::addLat(unsigned t, unsigned i, unsigned e) {
unsigned Merger::takeDisj(Kind kind, unsigned s0, unsigned s1) {
   else if (kind == kSubI)
     s1 = mapSet(kNegI, s1);
   // Followed by all in s1.
   for (unsigned p : latSets[s1])
     latSets[s].push_back(p);
   return s;
 }
 
-unsigned Merger::mapSet(Kind kind, unsigned s0) {
-  assert(kAbsF <= kind && kind <= kNegI);
+unsigned Merger::mapSet(Kind kind, unsigned s0, Value v) {
+  assert(kAbsF <= kind && kind <= kBitCast);
   unsigned s = addSet();
   for (unsigned p : latSets[s0]) {
-    unsigned e = addExp(kind, latPoints[p].exp);
+    unsigned e = addExp(kind, latPoints[p].exp, v);
     latPoints.push_back(LatPoint(latPoints[p].bits, e));
     latSets[s].push_back(latPoints.size() - 1);
   }
   return s;
 }
 
 unsigned Merger::optimizeSet(unsigned s0) {
   unsigned s = addSet();
bool Merger::isConjunction(unsigned t, unsigned e) const {
   switch (tensorExps[e].kind) {
   case kTensor:
     return tensorExps[e].tensor == t;
   case kAbsF:
   case kCeilF:
   case kFloorF:
   case kNegF:
   case kNegI:
+  case kTruncF:
+  case kExtF:
+  case kCastFS:
+  case kCastFU:
+  case kCastSF:
+  case kCastUF:
+  case kCastS:
+  case kCastU:
+  case kTruncI:
+  case kBitCast:
     return isConjunction(t, tensorExps[e].children.e0);
   case kDivF: // note: x / c only
   case kDivS:
   case kDivU:
     assert(!maybeZero(tensorExps[e].children.e1));
     return isConjunction(t, tensorExps[e].children.e0);
   case kShrS: // note: x >> inv only
   case kShrU:
   case kShlI:
     assert(isInvariant(tensorExps[e].children.e1));
     return isConjunction(t, tensorExps[e].children.e0);
   case kMulF:
   case kMulI:
   case kAndI:
     return isConjunction(t, tensorExps[e].children.e0) ||
            isConjunction(t, tensorExps[e].children.e1);
   default:
     return false;
   }
 }
 
 #ifndef NDEBUG
 
-//
+//===----------------------------------------------------------------------===//
 // Print methods (for debugging).
-//
+//===----------------------------------------------------------------------===//
 
 static const char *kindToOpSymbol(Kind kind) {
   switch (kind) {
   case kTensor:
     return "tensor";
   case kInvariant:
     return "invariant";
   case kAbsF:
     return "abs";
   case kCeilF:
     return "ceil";
   case kFloorF:
     return "floor";
   case kNegF:
     return "-";
   case kNegI:
     return "-";
+  case kTruncF:
+  case kExtF:
+  case kCastFS:
+  case kCastFU:
+  case kCastSF:
+  case kCastUF:
+  case kCastS:
+  case kCastU:
+  case kTruncI:
+  case kBitCast:
+    return "cast";
   case kMulF:
     return "*";
   case kMulI:
     return "*";
   case kDivF:
     return "/";
   case kDivS:
     return "/";
void Merger::dumpExp(unsigned e) const {
   case kInvariant:
     llvm::dbgs() << "invariant";
     break;
   case kAbsF:
   case kCeilF:
   case kFloorF:
   case kNegF:
   case kNegI:
+  case kTruncF:
+  case kExtF:
+  case kCastFS:
+  case kCastFU:
+  case kCastSF:
+  case kCastUF:
+  case kCastS:
+  case kCastU:
+  case kTruncI:
+  case kBitCast:
     llvm::dbgs() << kindToOpSymbol(tensorExps[e].kind) << " ";
     dumpExp(tensorExps[e].children.e0);
     break;
   default:
     llvm::dbgs() << "(";
     dumpExp(tensorExps[e].children.e0);
     llvm::dbgs() << " " << kindToOpSymbol(tensorExps[e].kind) << " ";
     dumpExp(tensorExps[e].children.e1);
if (bits[b]) {
         break;
       }
     }
   }
 }
 
 #endif // NDEBUG
 
-//
+//===----------------------------------------------------------------------===//
 // Builder methods.
-//
+//===----------------------------------------------------------------------===//
 
 unsigned Merger::buildLattices(unsigned e, unsigned i) {
   Kind kind = tensorExps[e].kind;
   switch (kind) {
   case kTensor:
   case kInvariant: {
     // Either the index is really used in the tensor expression, or it is
     // set to the undefined index in that dimension. An invariant expression
     // is set to a synthetic tensor with undefined indices only.
     unsigned s = addSet();
     unsigned t = kind == kTensor ? tensorExps[e].tensor : syntheticTensor;
     latSets[s].push_back(addLat(t, i, e));
     return s;
   }
   case kAbsF:
   case kCeilF:
   case kFloorF:
   case kNegF:
   case kNegI:
+  case kTruncF:
+  case kExtF:
+  case kCastFS:
+  case kCastFU:
+  case kCastSF:
+  case kCastUF:
+  case kCastS:
+  case kCastU:
+  case kTruncI:
+  case kBitCast:
     // A zero preserving operation (viz. f(0) = 0, [Bik96,Ch5]) maps the
     // lattice set of the operand through the operator into a new set.
     //
     //  -y|!y | y |
     //  --+---+---+
     //    | 0 |-y |
-    return mapSet(kind, buildLattices(tensorExps[e].children.e0, i));
+    return mapSet(kind, buildLattices(tensorExps[e].children.e0, i),
+                  tensorExps[e].val);
   case kMulF:
   case kMulI:
   case kAndI:
     // A multiplicative operation only needs to be performed
     // for the conjunction of sparse iteration spaces.
     //
     //  x*y|!y | y |
     //  ---+---+---+
bool Merger::maybeZero(unsigned e) const {
   }
   return true;
 }
 
 bool Merger::isInvariant(unsigned e) const {
   return tensorExps[e].kind == kInvariant;
 }
 
+Type Merger::inferType(unsigned e, Value src) {
+  // Obtain the destination type from the cast node.
+  Type dtp = tensorExps[e].val.getType();
+  // Inspect source type. For vector types, apply the same
+  // vectorization to the destination type.
+  if (auto vtp = src.getType().dyn_cast<VectorType>())
+    return VectorType::get(vtp.getNumElements(), dtp);
+  return dtp;
+}
+
 Optional<unsigned> Merger::buildTensorExp(linalg::GenericOp op, Value v) {
   if (auto arg = v.dyn_cast<BlockArgument>()) {
     unsigned argN = arg.getArgNumber();
     // Any argument of the generic op that is not marked as a scalar
     // argument is considered a tensor, indexed by the implicit loop
     // bounds. This includes rank-0 tensor arguments.
     if (arg.getOwner()->getParentOp() == op) {
       OpOperand *t = op.getInputAndOutputOperands()[argN];
if (x.hasValue()) {
       unsigned e = x.getValue();
       if (isa<AbsFOp>(def))
         return addExp(kAbsF, e);
       if (isa<CeilFOp>(def))
         return addExp(kCeilF, e);
       if (isa<FloorFOp>(def))
         return addExp(kFloorF, e);
       if (isa<NegFOp>(def))
-        return addExp(kNegF, e);
-      // TODO: no negi in std?
+        return addExp(kNegF, e); // TODO: no negi in std?
+      if (isa<FPTruncOp>(def))
+        return addExp(kTruncF, e, v);
+      if (isa<FPExtOp>(def))
+        return addExp(kExtF, e, v);
+      if (isa<FPToSIOp>(def))
+        return addExp(kCastFS, e, v);
+      if (isa<FPToUIOp>(def))
+        return addExp(kCastFU, e, v);
+      if (isa<SIToFPOp>(def))
+        return addExp(kCastSF, e, v);
+      if (isa<UIToFPOp>(def))
+        return addExp(kCastUF, e, v);
+      if (isa<SignExtendIOp>(def))
+        return addExp(kCastS, e, v);
+      if (isa<ZeroExtendIOp>(def))
+        return addExp(kCastU, e, v);
+      if (isa<TruncateIOp>(def))
+        return addExp(kTruncI, e, v);
+      if (isa<BitcastOp>(def))
+        return addExp(kBitCast, e, v);
     }
   }
   // Construct binary operations if subexpressions can be built.
-  // TODO: see buildLattices() for an explanation of rejecting certain divisions
+  // TODO: see buildLattices() for an explanation of rejecting
+  //       certain division and shift operations
   if (def->getNumOperands() == 2) {
     auto x = buildTensorExp(op, def->getOperand(0));
     auto y = buildTensorExp(op, def->getOperand(1));
     if (x.hasValue() && y.hasValue()) {
       unsigned e0 = x.getValue();
       unsigned e1 = y.getValue();
       if (isa<MulFOp>(def))
         return addExp(kMulF, e0, e1);

 }
 
 Value Merger::buildExp(PatternRewriter &rewriter, Location loc, unsigned e,
                        Value v0, Value v1) {
   switch (tensorExps[e].kind) {
   case kTensor:
   case kInvariant:
     llvm_unreachable("unexpected non-op");
+  // Unary ops.
   case kAbsF:
     return rewriter.create<AbsFOp>(loc, v0);
   case kCeilF:
     return rewriter.create<CeilFOp>(loc, v0);
   case kFloorF:
     return rewriter.create<FloorFOp>(loc, v0);
   case kNegF:
     return rewriter.create<NegFOp>(loc, v0);
   case kNegI:
     assert(v1); // no negi in std
     return rewriter.create<SubIOp>(loc, v0, v1);
+  case kTruncF:
+    return rewriter.create<FPTruncOp>(loc, v0, inferType(e, v0));
+  case kExtF:
+    return rewriter.create<FPExtOp>(loc, v0, inferType(e, v0));
+  case kCastFS:
+    return rewriter.create<FPToSIOp>(loc, v0, inferType(e, v0));
+  case kCastFU:
+    return rewriter.create<FPToUIOp>(loc, v0, inferType(e, v0));
+  case kCastSF:
+    return rewriter.create<SIToFPOp>(loc, v0, inferType(e, v0));
+  case kCastUF:
+    return rewriter.create<UIToFPOp>(loc, v0, inferType(e, v0));
+  case kCastS:
+    return rewriter.create<SignExtendIOp>(loc, v0, inferType(e, v0));
+  case kCastU:
+    return rewriter.create<ZeroExtendIOp>(loc, v0, inferType(e, v0));
+  case kTruncI:
+    return rewriter.create<TruncateIOp>(loc, v0, inferType(e, v0));
+  case kBitCast:
+    return rewriter.create<BitcastOp>(loc, v0, inferType(e, v0));
+  // Binary ops.
   case kMulF:
     return rewriter.create<MulFOp>(loc, v0, v1);
   case kMulI:
     return rewriter.create<MulIOp>(loc, v0, v1);
   case kDivF:
     return rewriter.create<DivFOp>(loc, v0, v1);
   case kDivS:
     return rewriter.create<SignedDivIOp>(loc, v0, v1);

mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_cast.mlir

+// RUN: mlir-opt %s \
+// RUN:   --sparsification --sparse-tensor-conversion \
+// RUN:   --convert-vector-to-scf --convert-scf-to-std \
+// RUN:   --func-bufferize --tensor-constant-bufferize --tensor-bufferize \
+// RUN:   --std-bufferize --finalizing-bufferize --lower-affine \
+// RUN:   --convert-vector-to-llvm --convert-memref-to-llvm --convert-std-to-llvm | \
+// RUN: mlir-cpu-runner \
+// RUN:  -e entry -entry-point-result=void  \
+// RUN:  -shared-libs=%mlir_integration_test_dir/libmlir_c_runner_utils%shlibext | \
+// RUN: FileCheck %s
+//
+// Do the same run, but now with SIMDization as well. This should not change the outcome.
+//
+// RUN: mlir-opt %s \
+// RUN:   --sparsification="vectorization-strategy=2 vl=2 enable-simd-index32" --sparse-tensor-conversion \
+// RUN:   --convert-vector-to-scf --convert-scf-to-std \
+// RUN:   --func-bufferize --tensor-constant-bufferize --tensor-bufferize \
+// RUN:   --std-bufferize --finalizing-bufferize --lower-affine \
+// RUN:   --convert-vector-to-llvm --convert-memref-to-llvm --convert-std-to-llvm | \
+// RUN: TENSOR0="%mlir_integration_test_dir/data/test.mtx" \
+// RUN: mlir-cpu-runner \
+// RUN:  -e entry -entry-point-result=void  \
+// RUN:  -shared-libs=%mlir_integration_test_dir/libmlir_c_runner_utils%shlibext | \
+// RUN: FileCheck %s
+//
+
+#SV = #sparse_tensor.encoding<{ dimLevelType = [ "compressed" ] }>
+
+#trait_cast = {
+  indexing_maps = [
+    affine_map<(i) -> (i)>,  // A (in)
+    affine_map<(i) -> (i)>   // X (out)
+  ],
+  iterator_types = ["parallel"],
+  doc = "X(i) = cast A(i)"
+}
+
+//
+// Integration test that lowers a kernel annotated as sparse to actual sparse
+// code, initializes a matching sparse storage scheme from a dense vector,
+// and runs the resulting code with the JIT compiler.
+//
+module {
+  //
+  // Various kernels that cast a sparse vector from one type to another.
+  // Standard supports the following casts.
+  //   sitofp
+  //   uitofp
+  //   fptosi
+  //   fptoui
+  //   fpext
+  //   fptrunc
+  //   sexti
+  //   zexti
+  //   trunci
+  //   bitcast
+  // Since all casts are "zero preserving" unary operations, lattice computation
+  // and conversion to sparse code is straightforward.
+  //
+  func @sparse_cast_s32_to_f32(%arga: tensor<10xi32, #SV>) -> tensor<10xf32> {
+    %argx = constant dense<0.0> : tensor<10xf32>
+    %0 = linalg.generic #trait_cast
+      ins(%arga: tensor<10xi32, #SV>)
+      outs(%argx: tensor<10xf32>) {
+        ^bb(%a: i32, %x : f32):
+          %cst = sitofp %a : i32 to f32
+          linalg.yield %cst : f32
+    } -> tensor<10xf32>
+    return %0 : tensor<10xf32>
+  }
+  func @sparse_cast_u32_to_f32(%arga: tensor<10xi32, #SV>) -> tensor<10xf32> {
+    %argx = constant dense<0.0> : tensor<10xf32>
+    %0 = linalg.generic #trait_cast
+      ins(%arga: tensor<10xi32, #SV>)
+      outs(%argx: tensor<10xf32>) {
+        ^bb(%a: i32, %x : f32):
+          %cst = uitofp %a : i32 to f32
+          linalg.yield %cst : f32
+    } -> tensor<10xf32>
+    return %0 : tensor<10xf32>
+  }
+  func @sparse_cast_f32_to_s32(%arga: tensor<10xf32, #SV>) -> tensor<10xi32> {
+    %argx = constant dense<0> : tensor<10xi32>
+    %0 = linalg.generic #trait_cast
+      ins(%arga: tensor<10xf32, #SV>)
+      outs(%argx: tensor<10xi32>) {
+        ^bb(%a: f32, %x : i32):
+          %cst = fptosi %a : f32 to i32
+          linalg.yield %cst : i32
+    } -> tensor<10xi32>
+    return %0 : tensor<10xi32>
+  }
+  func @sparse_cast_f64_to_u32(%arga: tensor<10xf64, #SV>) -> tensor<10xi32> {
+    %argx = constant dense<0> : tensor<10xi32>
+    %0 = linalg.generic #trait_cast
+      ins(%arga: tensor<10xf64, #SV>)
+      outs(%argx: tensor<10xi32>) {
+        ^bb(%a: f64, %x : i32):
+          %cst = fptoui %a : f64 to i32
+          linalg.yield %cst : i32
+    } -> tensor<10xi32>
+    return %0 : tensor<10xi32>
+  }
+  func @sparse_cast_f32_to_f64(%arga: tensor<10xf32, #SV>) -> tensor<10xf64> {
+    %argx = constant dense<0.0> : tensor<10xf64>
+    %0 = linalg.generic #trait_cast
+      ins(%arga: tensor<10xf32, #SV>)
+      outs(%argx: tensor<10xf64>) {
+        ^bb(%a: f32, %x : f64):
+          %cst = fpext %a : f32 to f64
+          linalg.yield %cst : f64
+    } -> tensor<10xf64>
+    return %0 : tensor<10xf64>
+  }
+  func @sparse_cast_f64_to_f32(%arga: tensor<10xf64, #SV>) -> tensor<10xf32> {
+    %argx = constant dense<0.0> : tensor<10xf32>
+    %0 = linalg.generic #trait_cast
+      ins(%arga: tensor<10xf64, #SV>)
+      outs(%argx: tensor<10xf32>) {
+        ^bb(%a: f64, %x : f32):
+          %cst = fptrunc %a : f64 to f32
+          linalg.yield %cst : f32
+    } -> tensor<10xf32>
+    return %0 : tensor<10xf32>
+  }
+  func @sparse_cast_s32_to_u64(%arga: tensor<10xi32, #SV>) -> tensor<10xi64> {
+    %argx = constant dense<0> : tensor<10xi64>
+    %0 = linalg.generic #trait_cast
+      ins(%arga: tensor<10xi32, #SV>)
+      outs(%argx: tensor<10xi64>) {
+        ^bb(%a: i32, %x : i64):
+          %cst = sexti %a : i32 to i64
+          linalg.yield %cst : i64
+    } -> tensor<10xi64>
+    return %0 : tensor<10xi64>
+  }
+  func @sparse_cast_u32_to_s64(%arga: tensor<10xi32, #SV>) -> tensor<10xi64> {
+    %argx = constant dense<0> : tensor<10xi64>
+    %0 = linalg.generic #trait_cast
+      ins(%arga: tensor<10xi32, #SV>)
+      outs(%argx: tensor<10xi64>) {
+        ^bb(%a: i32, %x : i64):
+          %cst = zexti %a : i32 to i64
+          linalg.yield %cst : i64
+    } -> tensor<10xi64>
+    return %0 : tensor<10xi64>
+  }
+  func @sparse_cast_i32_to_i8(%arga: tensor<10xi32, #SV>) -> tensor<10xi8> {
+    %argx = constant dense<0> : tensor<10xi8>
+    %0 = linalg.generic #trait_cast
+      ins(%arga: tensor<10xi32, #SV>)
+      outs(%argx: tensor<10xi8>) {
+        ^bb(%a: i32, %x : i8):
+          %cst = trunci %a : i32 to i8
+          linalg.yield %cst : i8
+    } -> tensor<10xi8>
+    return %0 : tensor<10xi8>
+  }
+  func @sparse_cast_f32_as_s32(%arga: tensor<10xf32, #SV>) -> tensor<10xi32> {
+    %argx = constant dense<0> : tensor<10xi32>
+    %0 = linalg.generic #trait_cast
+      ins(%arga: tensor<10xf32, #SV>)
+      outs(%argx: tensor<10xi32>) {
+        ^bb(%a: f32, %x : i32):
+          %cst = bitcast %a : f32 to i32
+          linalg.yield %cst : i32
+    } -> tensor<10xi32>
+    return %0 : tensor<10xi32>
+  }
+
+  //
+  // Main driver that converts a dense tensor into a sparse tensor
+  // and then calls the sparse casting kernel.
+  //
+  func @entry() {
+    %z = constant 0 : index
+    %b = constant 0 : i8
+    %i = constant 0 : i32
+    %l = constant 0 : i64
+    %f = constant 0.0 : f32
+    %d = constant 0.0 : f64
+
+    // Initialize dense tensors, convert to a sparse vectors.
+    %0 = constant dense<[ -4, -3, -2, -1, 0, 1, 2, 3, 4, 305 ]> : tensor<10xi32>
+    %1 = sparse_tensor.convert %0 : tensor<10xi32> to tensor<10xi32, #SV>
+    %2 = constant dense<[ -4.4, -3.3, -2.2, -1.1, 0.0, 1.1, 2.2, 3.3, 4.4, 305.5 ]> : tensor<10xf32>
+    %3 = sparse_tensor.convert %2 : tensor<10xf32> to tensor<10xf32, #SV>
+    %4 = constant dense<[ -4.4, -3.3, -2.2, -1.1, 0.0, 1.1, 2.2, 3.3, 4.4, 305.5 ]> : tensor<10xf64>
+    %5 = sparse_tensor.convert %4 : tensor<10xf64> to tensor<10xf64, #SV>
+    %6 = constant dense<[ 4294967295.0, 4294967294.0, 4294967293.0, 4294967292.0,
+                          0.0, 1.1, 2.2, 3.3, 4.4, 305.5 ]> : tensor<10xf64>
+    %7 = sparse_tensor.convert %6 : tensor<10xf64> to tensor<10xf64, #SV>
+
+    //
+    // CHECK: ( -4, -3, -2, -1, 0, 1, 2, 3, 4, 305 )
+    //
+    %c0 = call @sparse_cast_s32_to_f32(%1) : (tensor<10xi32, #SV>) -> tensor<10xf32>
+    %m0 = memref.buffer_cast %c0 : memref<10xf32>
+    %v0 = vector.transfer_read %m0[%z], %f: memref<10xf32>, vector<10xf32>
+    vector.print %v0 : vector<10xf32>
+
+    //
+    // CHECK: ( 4.29497e+09, 4.29497e+09, 4.29497e+09, 4.29497e+09, 0, 1, 2, 3, 4, 305 )
+    //
+    %c1 = call @sparse_cast_u32_to_f32(%1) : (tensor<10xi32, #SV>) -> tensor<10xf32>
+    %m1 = memref.buffer_cast %c1 : memref<10xf32>
+    %v1 = vector.transfer_read %m1[%z], %f: memref<10xf32>, vector<10xf32>
+    vector.print %v1 : vector<10xf32>
+
+    //
+    // CHECK: ( -4, -3, -2, -1, 0, 1, 2, 3, 4, 305 )
+    //
+    %c2 = call @sparse_cast_f32_to_s32(%3) : (tensor<10xf32, #SV>) -> tensor<10xi32>
+    %m2 = memref.buffer_cast %c2 : memref<10xi32>
+    %v2 = vector.transfer_read %m2[%z], %i: memref<10xi32>, vector<10xi32>
+    vector.print %v2 : vector<10xi32>
+
+    //
+    // CHECK: ( 4294967295, 4294967294, 4294967293, 4294967292, 0, 1, 2, 3, 4, 305 )
+    //
+    %c3 = call @sparse_cast_f64_to_u32(%7) : (tensor<10xf64, #SV>) -> tensor<10xi32>
+    %m3 = memref.buffer_cast %c3 : memref<10xi32>
+    %v3 = vector.transfer_read %m3[%z], %i: memref<10xi32>, vector<10xi32>
+    %vu = vector.bitcast %v3 : vector<10xi32> to vector<10xui32>
+    vector.print %vu : vector<10xui32>
+
+    //
+    // CHECK: ( -4.4, -3.3, -2.2, -1.1, 0, 1.1, 2.2, 3.3, 4.4, 305.5 )
+    //
+    %c4 = call @sparse_cast_f32_to_f64(%3) : (tensor<10xf32, #SV>) -> tensor<10xf64>
+    %m4 = memref.buffer_cast %c4 : memref<10xf64>
+    %v4 = vector.transfer_read %m4[%z], %d: memref<10xf64>, vector<10xf64>
+    vector.print %v4 : vector<10xf64>
+
+    //
+    // CHECK: ( -4.4, -3.3, -2.2, -1.1, 0, 1.1, 2.2, 3.3, 4.4, 305.5 )
+    //
+    %c5 = call @sparse_cast_f64_to_f32(%5) : (tensor<10xf64, #SV>) -> tensor<10xf32>
+    %m5 = memref.buffer_cast %c5 : memref<10xf32>
+    %v5 = vector.transfer_read %m5[%z], %f: memref<10xf32>, vector<10xf32>
+    vector.print %v5 : vector<10xf32>
+
+    //
+    // CHECK: ( -4, -3, -2, -1, 0, 1, 2, 3, 4, 305 )
+    //
+    %c6 = call @sparse_cast_s32_to_u64(%1) : (tensor<10xi32, #SV>) -> tensor<10xi64>
+    %m6 = memref.buffer_cast %c6 : memref<10xi64>
+    %v6 = vector.transfer_read %m6[%z], %l: memref<10xi64>, vector<10xi64>
+    vector.print %v6 : vector<10xi64>
+
+    //
+    // CHECK: ( 4294967292, 4294967293, 4294967294, 4294967295, 0, 1, 2, 3, 4, 305 )
+    //
+    %c7 = call @sparse_cast_u32_to_s64(%1) : (tensor<10xi32, #SV>) -> tensor<10xi64>
+    %m7 = memref.buffer_cast %c7 : memref<10xi64>
+    %v7 = vector.transfer_read %m7[%z], %l: memref<10xi64>, vector<10xi64>
+    vector.print %v7 : vector<10xi64>
+
+    //
+    // CHECK: ( -4, -3, -2, -1, 0, 1, 2, 3, 4, 49 )
+    //
+    %c8 = call @sparse_cast_i32_to_i8(%1) : (tensor<10xi32, #SV>) -> tensor<10xi8>
+    %m8 = memref.buffer_cast %c8 : memref<10xi8>
+    %v8 = vector.transfer_read %m8[%z], %b: memref<10xi8>, vector<10xi8>
+    vector.print %v8 : vector<10xi8>
+
+    //
+    // CHECK: ( -1064514355, -1068289229, -1072902963, -1081291571, 0, 1066192077, 1074580685, 1079194419, 1082969293, 1134084096 )
+    //
+    %c9 = call @sparse_cast_f32_as_s32(%3) : (tensor<10xf32, #SV>) -> tensor<10xi32>
+    %m9 = memref.buffer_cast %c9 : memref<10xi32>
+    %v9 = vector.transfer_read %m9[%z], %i: memref<10xi32>, vector<10xi32>
+    vector.print %v9 : vector<10xi32>
+
+    return
+  }
+}