From ec499985341f5e803378586d9f8c034caae845e2 Mon Sep 17 00:00:00 2001
From: Aart Bik <ajcbik@google.com>
Date: Wed, 24 Apr 2024 18:14:16 -0700
Subject: [PATCH 1/2] [torch-mlir][sparse] sparse tensor dialect is a legal
 dialect

---
 .../Transforms/VerifyLinalgOnTensorsBackendContract.cpp      | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/lib/Dialect/TorchConversion/Transforms/VerifyLinalgOnTensorsBackendContract.cpp b/lib/Dialect/TorchConversion/Transforms/VerifyLinalgOnTensorsBackendContract.cpp
index 5ad3fa1c9f4f..700c7a85a41a 100644
--- a/lib/Dialect/TorchConversion/Transforms/VerifyLinalgOnTensorsBackendContract.cpp
+++ b/lib/Dialect/TorchConversion/Transforms/VerifyLinalgOnTensorsBackendContract.cpp
@@ -18,6 +18,7 @@
 #include "mlir/Dialect/MLProgram/IR/MLProgram.h"
 #include "mlir/Dialect/Math/IR/Math.h"
 #include "mlir/Dialect/SCF/IR/SCF.h"
+#include "mlir/Dialect/SparseTensor/IR/SparseTensor.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
 #include "mlir/IR/OpDefinition.h"
 #include "mlir/Transforms/DialectConversion.h"
@@ -76,8 +77,10 @@ class VerifyLinalgOnTensorsBackendContractPass
     target.addDynamicallyLegalDialect<arith::ArithDialect>(isLegalScalarOp);
     target.addDynamicallyLegalDialect<complex::ComplexDialect>(isLegalScalarOp);
 
-    // Tensor operations should go through linalg and the tensor dialect.
+    // Tensor operations should go through linalg and the (sparse) tensor dialect.
     target.addDynamicallyLegalDialect<linalg::LinalgDialect>(opHasLegalTypes);
+    target.addDynamicallyLegalDialect<sparse_tensor::SparseTensorDialect>(
+        opHasLegalTypes);
     target.addDynamicallyLegalDialect<tensor::TensorDialect>(opHasLegalTypes);
     target.addDynamicallyLegalDialect<affine::AffineDialect>(opHasLegalTypes);
     target.addDynamicallyLegalDialect<cf::ControlFlowDialect>(opHasLegalTypes);

From a3bdae578c0a120b49a0be8c76a0bc5165328760 Mon Sep 17 00:00:00 2001
From: Aart Bik <ajcbik@google.com>
Date: Wed, 24 Apr 2024 18:29:47 -0700
Subject: [PATCH 2/2] lint

---
 .../Transforms/VerifyLinalgOnTensorsBackendContract.cpp         | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/lib/Dialect/TorchConversion/Transforms/VerifyLinalgOnTensorsBackendContract.cpp b/lib/Dialect/TorchConversion/Transforms/VerifyLinalgOnTensorsBackendContract.cpp
index 700c7a85a41a..e8789a05a3be 100644
--- a/lib/Dialect/TorchConversion/Transforms/VerifyLinalgOnTensorsBackendContract.cpp
+++ b/lib/Dialect/TorchConversion/Transforms/VerifyLinalgOnTensorsBackendContract.cpp
@@ -77,7 +77,7 @@ class VerifyLinalgOnTensorsBackendContractPass
     target.addDynamicallyLegalDialect<arith::ArithDialect>(isLegalScalarOp);
     target.addDynamicallyLegalDialect<complex::ComplexDialect>(isLegalScalarOp);
 
-    // Tensor operations should go through linalg and the (sparse) tensor dialect.
+    // Tensor operations should go through linalg and the tensor dialect.
     target.addDynamicallyLegalDialect<linalg::LinalgDialect>(opHasLegalTypes);
     target.addDynamicallyLegalDialect<sparse_tensor::SparseTensorDialect>(
         opHasLegalTypes);