From 71cb94268200003ecafad76788212df8fc61c824 Mon Sep 17 00:00:00 2001
From: Aart Bik <ajcbik@google.com>
Date: Tue, 17 Dec 2024 08:03:58 -0800
Subject: [PATCH] [torch-mlir][sparse] register sparse tensor dialect for all
 rewriting (#3918)

We incorrectly relied on the fact that StableHLO registers the sparse
tensor dialect, but when building for e.g. just LinAlg, the dependency
was missing. This fixes this shortcoming.

FIXES: https://github.com/llvm/torch-mlir/issues/3816
---
 lib/InitAll.cpp | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/lib/InitAll.cpp b/lib/InitAll.cpp
index c9638c8353b1..d9d7ef1a0cd4 100644
--- a/lib/InitAll.cpp
+++ b/lib/InitAll.cpp
@@ -16,6 +16,7 @@
 #include "mlir/Dialect/MLProgram/IR/MLProgram.h"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
 #include "mlir/Dialect/SCF/IR/SCF.h"
+#include "mlir/Dialect/SparseTensor/IR/SparseTensor.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
 #include "mlir/Dialect/Tensor/IR/TensorInferTypeOpInterfaceImpl.h"
 #include "mlir/Dialect/Tosa/IR/TosaOps.h"
@@ -52,7 +53,8 @@ void mlir::torch::registerOptionalInputDialects(
     mlir::DialectRegistry &registry) {
   registry.insert<complex::ComplexDialect, linalg::LinalgDialect,
                   memref::MemRefDialect, ml_program::MLProgramDialect,
-                  scf::SCFDialect, tensor::TensorDialect, tosa::TosaDialect>();
+                  scf::SCFDialect, sparse_tensor::SparseTensorDialect,
+                  tensor::TensorDialect, tosa::TosaDialect>();
 }
 
 void mlir::torch::registerAllPasses() {