[mlir][transform][sparse] introduce sparse tensor transform extension.

[Peiming]

blaze change WIP.

[nicolasvasilache]

mlir/include/mlir/Dialect/SparseTensor/TransformOps/SparseTensorTransformOps.td

28

I was in the process of making some ill-advised suggestions on how the transform should be more composable but then I realized that this is a actual match op and not a transform.

So I will defer to @ftynse as I have not yet had time to dig into that part of the system.
Offhand, it seems you want to test this as part of a transform.match.structured region rather than as a transform (see e.g. mlir/test/Dialect/Linalg/match-ops-interpreter.mlir)

[Peiming]

address comments.

[aartbik]

typo in title: instroduce

mlir/include/mlir/Dialect/SparseTensor/TransformOps/SparseTensorTransformOps.h

35

looks like this is a copy and paste error from the original file that inspired you ;-)

[wrengr]

mlir/include/mlir/Dialect/SparseTensor/TransformOps/SparseTensorTransformOps.td

2

should adjust this line to fit the 80-char limit

[Peiming]

address comments.

[ftynse]

mlir/include/mlir/Dialect/SparseTensor/TransformOps/SparseTensorTransformOps.td

35

Nit: this can use custom<SemiFunctionType>(type($target), type($result)) if desired so it is future-proof with potential syntax changes in other match ops.

mlir/lib/Dialect/SparseTensor/TransformOps/SparseTensorTransformOps.cpp

25

Super-nit: start error messages with a non-capital letter, these are usually appended to something.

41

It doesn't (yet?) generate anything, so there's no need to declare this. Okay to keep if there is an intention to add more ops quickly.

[Peiming]

mlir/lib/Dialect/SparseTensor/TransformOps/SparseTensorTransformOps.cpp

41

Yeah, that is what we expect, I will keep it.

[Peiming]

address comments.

[Peiming]

add bazel config.

[aartbik]

mlir/test/Integration/Dialect/SparseTensor/TransformOps/match_sparse_ops.mlir

2

Perhaps others feel strongly about this, but can we simply put this in test/Dialect/SparseTensor/TransformsOps (i.e. not as an integration test). It feels different from the end-to-end run that we define there

[Peiming]

move to non-integration test.

[Peiming]

mlir/test/Integration/Dialect/SparseTensor/TransformOps/match_sparse_ops.mlir

2

I think it make sense to move it out integration test.

[Peiming]

rebase.

[aartbik]

mlir/test/Integration/Dialect/SparseTensor/TransformOps/match_sparse_ops.mlir

2

Thanks, feels better to me too.

[Peiming]

Seems cause a buildbot error, fixing it...

[Peiming]

In D157025#4561691, @Peiming wrote:

Seems cause a buildbot error, fixing it...

https://reviews.llvm.org/D157134

mlir/include/mlir/Dialect/SparseTensor/CMakeLists.txt

 add_subdirectory(IR)
 add_subdirectory(Transforms)
+add_subdirectory(TransformOps)

mlir/include/mlir/Dialect/SparseTensor/TransformOps/CMakeLists.txt

+set(LLVM_TARGET_DEFINITIONS SparseTensorTransformOps.td)
+mlir_tablegen(SparseTensorTransformOps.h.inc -gen-op-decls)
+mlir_tablegen(SparseTensorTransformOps.cpp.inc -gen-op-defs)
+add_public_tablegen_target(MLIRSparseTensorTransformOpsIncGen)

mlir/include/mlir/Dialect/SparseTensor/TransformOps/SparseTensorTransformOps.h

+//===- SparseTensorTransformOps.h - sparse tensor transform ops -*- C++ -*-===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_DIALECT_SPARSETENSOR_TRANSFORMOPS_SPARSETENSORTRANSFORMOPS_H
+#define MLIR_DIALECT_SPARSETENSOR_TRANSFORMOPS_SPARSETENSORTRANSFORMOPS_H
+
+#include "mlir/Dialect/Transform/IR/MatchInterfaces.h"
+#include "mlir/Dialect/Transform/IR/TransformAttrs.h"
+#include "mlir/Dialect/Transform/IR/TransformDialect.h"
+#include "mlir/Dialect/Transform/IR/TransformInterfaces.h"
+#include "mlir/IR/OpImplementation.h"
+#include "mlir/IR/RegionKindInterface.h"
+
+namespace mlir {
+namespace transform {
+class TransformHandleTypeInterface;
+} // namespace transform
+} // namespace mlir
+
+namespace mlir {
+class DialectRegistry;
+
+namespace sparse_tensor {
+void registerTransformDialectExtension(DialectRegistry &registry);
+} // namespace sparse_tensor
+} // namespace mlir
+
+//===----------------------------------------------------------------------===//
+// SparseTensor Transform Operations
+//===----------------------------------------------------------------------===//
+
+#define GET_OP_CLASSES
+#include "mlir/Dialect/SparseTensor/TransformOps/SparseTensorTransformOps.h.inc"
+
+#endif // MLIR_DIALECT_SPARSETENSOR_TRANSFORMOPS_SPARSETENSORTRANSFORMOPS_H

mlir/include/mlir/Dialect/SparseTensor/TransformOps/SparseTensorTransformOps.td

+//=== SparseTensorTransformOps.td ----------------------------*-tablegen *-===//
+//
+// Tablegen rules for Sparse Tensor transformation operations.
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef SPARSETENSOR_TRANSFORM_OPS
+#define SPARSETENSOR_TRANSFORM_OPS
+
+include "mlir/Dialect/Transform/IR/MatchInterfaces.td"
+include "mlir/Dialect/Transform/IR/TransformAttrs.td"
+include "mlir/Dialect/Transform/IR/TransformDialect.td"
+include "mlir/Dialect/Transform/IR/TransformTypes.td"
+include "mlir/Interfaces/SideEffectInterfaces.td"
+
+//===----------------------------------------------------------------------===//
+// SparseTensorMatchOps
+//===----------------------------------------------------------------------===//
+
+def MatchSparseInOut : Op<Transform_Dialect, "sparse_tensor.match.sparse_inout", [
+     MatchOpInterface,
+     SingleOpMatcher,
+     MemoryEffectsOpInterface]> {
+  let description = [{
+     Checks if the payload op has any sparse inputs and/or outputs.
+  }];
+
+  let arguments = (ins TransformHandleTypeInterface:$target);
+  let results = (outs TransformHandleTypeInterface:$result);
+
+  let assemblyFormat = "$target attr-dict `:` type($target) `->` type($result)";
+  let extraClassDeclaration = SingleOpMatcher.extraDeclaration # [{
+    ::mlir::Value getOperandHandle() { return getTarget(); }
+  }];
+}
+
+#endif // SPARSETENSOR_TRANSFORM_OPS

mlir/include/mlir/InitAllDialects.h

 #include "mlir/Dialect/SCF/IR/SCF.h"
 #include "mlir/Dialect/SCF/IR/ValueBoundsOpInterfaceImpl.h"
 #include "mlir/Dialect/SCF/TransformOps/SCFTransformOps.h"
 #include "mlir/Dialect/SCF/Transforms/BufferizableOpInterfaceImpl.h"
 #include "mlir/Dialect/SPIRV/IR/SPIRVDialect.h"
 #include "mlir/Dialect/Shape/IR/Shape.h"
 #include "mlir/Dialect/Shape/Transforms/BufferizableOpInterfaceImpl.h"
 #include "mlir/Dialect/SparseTensor/IR/SparseTensor.h"
+#include "mlir/Dialect/SparseTensor/TransformOps/SparseTensorTransformOps.h"
 #include "mlir/Dialect/SparseTensor/Transforms/BufferizableOpInterfaceImpl.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
 #include "mlir/Dialect/Tensor/IR/TensorInferTypeOpInterfaceImpl.h"
 #include "mlir/Dialect/Tensor/IR/TensorTilingInterfaceImpl.h"
 #include "mlir/Dialect/Tensor/IR/ValueBoundsOpInterfaceImpl.h"
 #include "mlir/Dialect/Tensor/TransformOps/TensorTransformOps.h"
 #include "mlir/Dialect/Tensor/Transforms/BufferizableOpInterfaceImpl.h"
 #include "mlir/Dialect/Tosa/IR/TosaOps.h"
inline void registerAllDialects(DialectRegistry &registry) {
   // Register all dialect extensions.
   affine::registerTransformDialectExtension(registry);
   bufferization::registerTransformDialectExtension(registry);
   gpu::registerTransformDialectExtension(registry);
   linalg::registerTransformDialectExtension(registry);
   memref::registerTransformDialectExtension(registry);
   nvgpu::registerTransformDialectExtension(registry);
   scf::registerTransformDialectExtension(registry);
+  sparse_tensor::registerTransformDialectExtension(registry);
   tensor::registerTransformDialectExtension(registry);
   transform::registerPDLExtension(registry);
   vector::registerTransformDialectExtension(registry);
 
   // Register all external models.
   affine::registerValueBoundsOpInterfaceExternalModels(registry);
   arith::registerBufferizableOpInterfaceExternalModels(registry);
   arith::registerValueBoundsOpInterfaceExternalModels(registry);

mlir/lib/Dialect/SparseTensor/CMakeLists.txt

 add_subdirectory(IR)
 add_subdirectory(Transforms)
+add_subdirectory(TransformOps)
 add_subdirectory(Pipelines)
 add_subdirectory(Utils)

mlir/lib/Dialect/SparseTensor/TransformOps/CMakeLists.txt

+add_mlir_dialect_library(MLIRSparseTensorTransformOps
+  SparseTensorTransformOps.cpp
+
+  ADDITIONAL_HEADER_DIRS
+  ${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/SparseTensor/TransformOps
+
+  DEPENDS
+  MLIRSparseTensorTransformOpsIncGen
+
+  LINK_LIBS PUBLIC
+  MLIRIR
+  MLIRSparseTensorDialect
+  MLIRParser
+  MLIRSideEffectInterfaces
+  MLIRTransformDialect
+  MLIRTransformDialectUtils
+  )

mlir/lib/Dialect/SparseTensor/TransformOps/SparseTensorTransformOps.cpp

+//===- SparseTensorTransformOps.cpp - sparse tensor transform ops impl ----===//
+//
+// 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/TransformOps/SparseTensorTransformOps.h"
+#include "mlir/Dialect/SparseTensor/IR/SparseTensor.h"
+
+using namespace mlir;
+using namespace mlir::sparse_tensor;
+
+//===----------------------------------------------------------------------===//
+// Transform op implementation
+//===----------------------------------------------------------------------===//
+
+DiagnosedSilenceableFailure transform::MatchSparseInOut::matchOperation(
+    mlir::Operation *current, mlir::transform::TransformResults &results,
+    mlir::transform::TransformState &state) {
+  bool hasSparseInOut = hasAnySparseOperandOrResult(current);
+  if (!hasSparseInOut) {
+    return emitSilenceableFailure(current->getLoc(),
+                                  "Operation has no sparse input or output");
+  }
+  results.set(getResult().cast<OpResult>(), state.getPayloadOps(getTarget()));
+  return DiagnosedSilenceableFailure::success();
+}
+
+//===----------------------------------------------------------------------===//
+// Transform op registration
+//===----------------------------------------------------------------------===//
+
+namespace {
+class SparseTensorTransformDialectExtension
+    : public transform::TransformDialectExtension<
+          SparseTensorTransformDialectExtension> {
+public:
+  SparseTensorTransformDialectExtension() {
+    declareGeneratedDialect<sparse_tensor::SparseTensorDialect>();
+    registerTransformOps<
+#define GET_OP_LIST
+#include "mlir/Dialect/SparseTensor/TransformOps/SparseTensorTransformOps.cpp.inc"
+        >();
+  }
+};
+} // namespace
+
+#define GET_OP_CLASSES
+#include "mlir/Dialect/SparseTensor/TransformOps/SparseTensorTransformOps.cpp.inc"
+
+void mlir::sparse_tensor::registerTransformDialectExtension(
+    DialectRegistry &registry) {
+  registry.addExtensions<SparseTensorTransformDialectExtension>();
+}

mlir/test/Integration/Dialect/SparseTensor/TransformOps/match_sparse_ops.mlir

+// RUN: mlir-opt %s --test-transform-dialect-interpreter --verify-diagnostics --split-input-file
+
+module attributes { transform.with_named_sequence } {
+  transform.named_sequence @match_sparse_structured(%arg0: !transform.any_op {transform.readonly}) -> !transform.any_op {
+    %0 = transform.match.structured %arg0 : (!transform.any_op) -> !transform.any_op {
+    ^bb0(%struct: !transform.any_op):
+      %sp_kernel = transform.sparse_tensor.match.sparse_inout %struct
+          : !transform.any_op -> !transform.any_op
+      transform.match.structured.yield %sp_kernel : !transform.any_op
+    }
+    transform.yield %0 : !transform.any_op
+  }
+
+  transform.named_sequence @print_sparse_structured(%arg0: !transform.any_op {transform.readonly}) {
+    transform.test_print_remark_at_operand %arg0, "sparse_kernel" : !transform.any_op
+    transform.yield
+  }
+
+  // Entry point. Match any structured sparse operation and emit at remark.
+  transform.sequence failures(propagate) attributes { transform.target_tag = "transform" } {
+  ^bb0(%arg0: !transform.any_op):
+    transform.foreach_match in %arg0
+        @match_sparse_structured -> @print_sparse_structured
+        : (!transform.any_op) -> !transform.any_op
+  }
+}
+
+#CSR = #sparse_tensor.encoding<{lvlTypes = ["dense", "compressed"]}>
+
+func.func @payload(%lhs: tensor<10x20xf16>,
+                   %sp_lhs: tensor<10x20xf16, #CSR>,
+                   %rhs: tensor<20x15xf32>) -> tensor<10x15xf64>{
+  %cst = arith.constant 0.0 : f64
+  %empty = tensor.empty() : tensor<10x15xf64>
+  %fill = linalg.fill ins(%cst : f64) outs(%empty : tensor<10x15xf64>) -> tensor<10x15xf64>
+
+  %result = linalg.matmul ins(%lhs, %rhs: tensor<10x20xf16>, tensor<20x15xf32>)
+                         outs(%fill: tensor<10x15xf64>) -> tensor<10x15xf64>
+  // expected-remark @below {{sparse_kernel}}
+  %sp_in = linalg.matmul ins(%sp_lhs, %rhs: tensor<10x20xf16, #CSR>, tensor<20x15xf32>)
+                        outs(%fill: tensor<10x15xf64>) -> tensor<10x15xf64>
+
+  %sp_empty = tensor.empty() : tensor<10x15xf64, #CSR>
+  // expected-remark @below {{sparse_kernel}}
+  %sp_out = linalg.matmul ins(%lhs, %rhs: tensor<10x20xf16>, tensor<20x15xf32>)
+                         outs(%sp_empty: tensor<10x15xf64, #CSR>) -> tensor<10x15xf64, #CSR>
+  return %result : tensor<10x15xf64>
+}