4afd0cf553
After --sparse-tensor-codegen, sparse tensor arguments are replaced by memrefs and \!sparse_tensor.storage_specifier types. The subsequent --sparse-assembler pass calls getSparseTensorEncoding() to identify sparse arguments to wrap/unwrap. However, getSparseTensorEncoding() returns non-null for StorageSpecifierType as well as for sparse RankedTensorType. Since StorageSpecifierType is not a RankedTensorType, the subsequent cast<RankedTensorType> in convTypes() and convVals() would crash with an assertion failure. Fix by also checking isa<RankedTensorType>(type) in the passthrough condition in both convTypes() and convVals(), so that StorageSpecifierType arguments pass through unchanged. Fixes #183776
31 lines
1.5 KiB
MLIR
31 lines
1.5 KiB
MLIR
// RUN: mlir-opt %s --sparse-tensor-codegen --sparse-assembler | FileCheck %s
|
|
|
|
// Regression test for https://github.com/llvm/llvm-project/issues/183776:
|
|
// Running --sparse-assembler after --sparse-tensor-codegen must not crash.
|
|
// After codegen, sparse tensor arguments are replaced by memrefs and
|
|
// \!sparse_tensor.storage_specifier types. getSparseTensorEncoding() returns
|
|
// non-null for StorageSpecifierType, but convTypes()/convVals() must not
|
|
// attempt cast<RankedTensorType> on it. Instead, non-RankedTensorType types
|
|
// with a sparse encoding should pass through unchanged.
|
|
|
|
#CSR = #sparse_tensor.encoding<{ map = (d0, d1) -> (d0: dense, d1: compressed) }>
|
|
|
|
// Storage_specifier types from codegen must pass through sparse-assembler
|
|
// unchanged (not be treated as sparse tensor arguments to wrap).
|
|
// CHECK-LABEL: func.func @storage_specifier_passthrough(
|
|
// CHECK-SAME: storage_specifier
|
|
// CHECK-SAME: storage_specifier
|
|
// CHECK: return %{{.*}} : tensor<32x32xf32>
|
|
func.func @storage_specifier_passthrough(%arg0: tensor<32x32xf32, #CSR>,
|
|
%arg1: tensor<32x32xf32, #CSR>)
|
|
-> tensor<32x32xf32> {
|
|
%cst = arith.constant 0.0 : f32
|
|
%init = tensor.empty() : tensor<32x32xf32>
|
|
%out = linalg.fill ins(%cst : f32) outs(%init : tensor<32x32xf32>)
|
|
-> tensor<32x32xf32>
|
|
%3 = linalg.add
|
|
ins(%arg0, %arg1 : tensor<32x32xf32, #CSR>, tensor<32x32xf32, #CSR>)
|
|
outs(%out : tensor<32x32xf32>) -> tensor<32x32xf32>
|
|
return %3 : tensor<32x32xf32>
|
|
}
|