[mlir][sparse] Better error handling when bufferizing sparse_tensor ops

[springerm]

Summary

sparse_tensor ops cannot be bufferized with One-Shot Bufferize. (They can only be analyzed.) The sparse compiler does the actual lowering to memref. Produce a proper error message instead of crashing.

This fixes #61311.

[aartbik]

Oh very neat. Thanks for this!

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

 
 using namespace mlir::bufferization;
 using namespace mlir::sparse_tensor;
 
 namespace mlir {
 namespace sparse_tensor {
 namespace {
 
+template <typename ConcreteModel, typename ConcreteOp>
+struct SparseBufferizableOpInterfaceExternalModel
+    : public BufferizableOpInterface::ExternalModel<ConcreteModel, ConcreteOp> {
+  LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
+                          const BufferizationOptions &options) const {
+    return op->emitError(
+        "sparse_tensor ops must be bufferized with the sparse compiler");
+  }
+};
+
 struct ConcatenateOpInterface
-    : public BufferizableOpInterface::ExternalModel<
-          ConcatenateOpInterface, sparse_tensor::ConcatenateOp> {
+    : SparseBufferizableOpInterfaceExternalModel<ConcatenateOpInterface,
+                                                 sparse_tensor::ConcatenateOp> {
   bool bufferizesToAllocation(Operation *op, Value value) const { return true; }
 
   bool bufferizesToMemoryRead(Operation *op, OpOperand &opOperand,
                               const AnalysisState &state) const {
     return true;
   }
 
   bool bufferizesToMemoryWrite(Operation *op, OpOperand &opOperand,
                                const AnalysisState &state) const {
     return false;
   }
 
   AliasingValueList getAliasingValues(Operation *op, OpOperand &opOperand,
                                       const AnalysisState &state) const {
     return {};
   }
 
   bool isWritable(Operation *op, Value value,
                   const AnalysisState &state) const {
     return true;
   }
 };
 
-struct ConvertOpInterface
-    : public BufferizableOpInterface::ExternalModel<ConvertOpInterface,
-                                                    sparse_tensor::ConvertOp> {
+struct ConvertOpInterface : public SparseBufferizableOpInterfaceExternalModel<
+                                ConvertOpInterface, sparse_tensor::ConvertOp> {
   bool bufferizesToAllocation(Operation *op, Value value) const {
     // ConvertOps may allocate. (Unless they convert between two identical
     // types, then they fold away.)
     return true;
   }
 
   bool bufferizesToMemoryRead(Operation *op, OpOperand &opOperand,
                               const AnalysisState &state) const {
struct ConvertOpInterface : public SparseBufferizableOpInterfaceExternalModel<
 
   bool isWritable(Operation *op, Value value,
                   const AnalysisState &state) const {
     return true;
   }
 };
 
 struct LoadOpInterface
-    : public BufferizableOpInterface::ExternalModel<LoadOpInterface,
+    : public SparseBufferizableOpInterfaceExternalModel<LoadOpInterface,
                                                         sparse_tensor::LoadOp> {
   bool bufferizesToMemoryRead(Operation *op, OpOperand &opOperand,
                               const AnalysisState &state) const {
     return false;
   }
 
   bool bufferizesToMemoryWrite(Operation *op, OpOperand &opOperand,
                                const AnalysisState &state) const {
     return false;
   }
 
   AliasingValueList getAliasingValues(Operation *op, OpOperand &opOperand,
                                       const AnalysisState &state) const {
     return {{op->getOpResult(0), BufferRelation::Equivalent}};
   }
 };
 
 struct NewOpInterface
-    : public BufferizableOpInterface::ExternalModel<NewOpInterface,
+    : public SparseBufferizableOpInterfaceExternalModel<NewOpInterface,
                                                         sparse_tensor::NewOp> {
   bool resultBufferizesToMemoryWrite(Operation *op, OpResult opResult,
                                      const AnalysisState &state) const {
     // NewOps allocate but do not write.
     return false;
   }
 
   bool bufferizesToAllocation(Operation *op, Value value) const { return true; }
 };
 
 struct PackOpInterface
-    : public BufferizableOpInterface::ExternalModel<PackOpInterface,
+    : public SparseBufferizableOpInterfaceExternalModel<PackOpInterface,
                                                         sparse_tensor::PackOp> {
   bool bufferizesToAllocation(Operation *op, Value value) const {
     // PackOp reuses all the buffers instead of allocating new ones
     return false;
   }
 
   bool bufferizesToMemoryRead(Operation *op, OpOperand &opOperand,
                               const AnalysisState &state) const {
     return true;
struct PackOpInterface
   }
 
   BufferRelation bufferRelation(Operation *oo, OpResult opResult,
                                 const AnalysisState &state) const {
     return BufferRelation::Unknown;
   }
 };
 
-struct UnpackOpInterface
-    : public BufferizableOpInterface::ExternalModel<UnpackOpInterface,
-                                                    sparse_tensor::UnpackOp> {
+struct UnpackOpInterface : public SparseBufferizableOpInterfaceExternalModel<
+                               UnpackOpInterface, sparse_tensor::UnpackOp> {
   bool bufferizesToAllocation(Operation *op, Value value) const {
     // The output buffer is pre-allocated by the user.
     return false;
   }
 
   bool bufferizesToMemoryRead(Operation *op, OpOperand &opOperand,
                               const AnalysisState &state) const {
     // The first operand is the sparse tensor that we are unpacking.
AliasingValueList getAliasingValues(Operation *op, OpOperand &opOperand,
     if (opOperand.getOperandNumber() == 0)
       return {};
     // We write directly into the output tensors and returns them.
     return {{op->getResult(opOperand.getOperandNumber() - 1),
              BufferRelation::Equivalent}};
   }
 };
 
-struct InsertOpInterface
-    : public BufferizableOpInterface::ExternalModel<InsertOpInterface,
-                                                    sparse_tensor::InsertOp> {
+struct InsertOpInterface : public SparseBufferizableOpInterfaceExternalModel<
+                               InsertOpInterface, sparse_tensor::InsertOp> {
   bool bufferizesToMemoryRead(Operation *op, OpOperand &opOperand,
                               const AnalysisState &state) const {
     return true;
   }
 
   bool bufferizesToMemoryWrite(Operation *op, OpOperand &opOperand,
                                const AnalysisState &state) const {
     // InsertOp writes to memory.
     return true;
   }
 
   AliasingValueList getAliasingValues(Operation *op, OpOperand &opOperand,
                                       const AnalysisState &state) const {
     // InsertOp returns an alias of its operand.
     assert(op->getNumResults() == 1);
     return {{op->getOpResult(0), BufferRelation::Equivalent}};
   }
 };
 
 struct NumberOfEntriesOpInterface
-    : public BufferizableOpInterface::ExternalModel<
+    : public SparseBufferizableOpInterfaceExternalModel<
           NumberOfEntriesOpInterface, sparse_tensor::NumberOfEntriesOp> {
   bool bufferizesToMemoryRead(Operation *op, OpOperand &opOperand,
                               const AnalysisState &state) const {
     return true;
   }
 
   bool bufferizesToMemoryWrite(Operation *op, OpOperand &opOperand,
                                const AnalysisState &state) const {
     return false;
   }
 
   AliasingValueList getAliasingValues(Operation *op, OpOperand &opOperand,
                                       const AnalysisState &state) const {
     return {};
   }
 };
 
 struct ToCoordinatesBufferOpInterface
-    : public BufferizableOpInterface::ExternalModel<
+    : public SparseBufferizableOpInterfaceExternalModel<
           ToCoordinatesBufferOpInterface,
           sparse_tensor::ToCoordinatesBufferOp> {
   bool bufferizesToMemoryRead(Operation *op, OpOperand &opOperand,
                               const AnalysisState &state) const {
     return true;
   }
 
   bool bufferizesToMemoryWrite(Operation *op, OpOperand &opOperand,
                                const AnalysisState &state) const {
     // Potential writes into memory through the result of
     // `sparse_tensor.coordinates` are not considered.
     return false;
   }
 
   AliasingValueList getAliasingValues(Operation *op, OpOperand &opOperand,
                                       const AnalysisState &state) const {
     return {};
   }
 };
 
 struct ToCoordinatesOpInterface
-    : public BufferizableOpInterface::ExternalModel<
+    : public SparseBufferizableOpInterfaceExternalModel<
           ToCoordinatesOpInterface, sparse_tensor::ToCoordinatesOp> {
   bool bufferizesToMemoryRead(Operation *op, OpOperand &opOperand,
                               const AnalysisState &state) const {
     return true;
   }
 
   bool bufferizesToMemoryWrite(Operation *op, OpOperand &opOperand,
                                const AnalysisState &state) const {
     // Potential writes into memory through the result of
     // `sparse_tensor.coordinates` are not considered.
     return false;
   }
 
   AliasingValueList getAliasingValues(Operation *op, OpOperand &opOperand,
                                       const AnalysisState &state) const {
     return {};
   }
 };
 
 struct ToPositionsOpInterface
-    : public BufferizableOpInterface::ExternalModel<
+    : public SparseBufferizableOpInterfaceExternalModel<
           ToPositionsOpInterface, sparse_tensor::ToPositionsOp> {
   bool bufferizesToMemoryRead(Operation *op, OpOperand &opOperand,
                               const AnalysisState &state) const {
     return true;
   }
 
   bool bufferizesToMemoryWrite(Operation *op, OpOperand &opOperand,
                                const AnalysisState &state) const {
     // Potential writes into memory through the result of
     // `sparse_tensor.positions` are not considered.
     return false;
   }
 
   AliasingValueList getAliasingValues(Operation *op, OpOperand &opOperand,
                                       const AnalysisState &state) const {
     return {};
   }
 };
 
 struct ToValuesOpInterface
-    : public BufferizableOpInterface::ExternalModel<ToValuesOpInterface,
-                                                    sparse_tensor::ToValuesOp> {
+    : public SparseBufferizableOpInterfaceExternalModel<
+          ToValuesOpInterface, sparse_tensor::ToValuesOp> {
   bool bufferizesToMemoryRead(Operation *op, OpOperand &opOperand,
                               const AnalysisState &state) const {
     return true;
   }
 
   bool bufferizesToMemoryWrite(Operation *op, OpOperand &opOperand,
                                const AnalysisState &state) const {
     // Potential writes into memory through the result of sparse_tensor.values

mlir/test/Dialect/SparseTensor/one_shot_bufferize_invalid.mlir

+// RUN: mlir-opt %s -one-shot-bufferize -verify-diagnostics
+
+#SparseVector = #sparse_tensor.encoding<{
+  lvlTypes = ["compressed"]
+}>
+
+func.func @sparse_tensor_op(%arg0: tensor<64xf32, #SparseVector>) -> tensor<64xf32, #SparseVector> {
+  // expected-error @below{{sparse_tensor ops must be bufferized with the sparse compiler}}
+  // expected-error @below{{failed to bufferize op}}
+  %0 = sparse_tensor.convert %arg0 : tensor<64xf32, #SparseVector> to tensor<64xf32, #SparseVector>
+  return %0 : tensor<64xf32, #SparseVector>
+}