[Mlir-commits] [mlir] [mlir][sparse] external entry method wrapper for sparse tensors (PR #80326)

Aart Bik llvmlistbot at llvm.org
Thu Feb 1 11:37:00 PST 2024


================
@@ -0,0 +1,236 @@
+//===- SparseAssembler.cpp - adds wrapper method around sparse types ------===//
+//
+// 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 "Utils/CodegenUtils.h"
+
+#include "mlir/Dialect/SparseTensor/IR/SparseTensor.h"
+#include "mlir/Dialect/SparseTensor/IR/SparseTensorStorageLayout.h"
+#include "mlir/Dialect/SparseTensor/IR/SparseTensorType.h"
+#include "mlir/Dialect/SparseTensor/Transforms/Passes.h"
+#include "mlir/Dialect/Tensor/IR/Tensor.h"
+#include "llvm/Support/FormatVariadic.h"
+
+using namespace mlir;
+using namespace sparse_tensor;
+
+//===----------------------------------------------------------------------===//
+// Helper methods.
+//===----------------------------------------------------------------------===//
+
+// Convert type range to new types range, with sparse tensors externalized.
+void convTypes(TypeRange types, SmallVectorImpl<Type> &convTypes,
+               SmallVectorImpl<Type> *extraTypes = nullptr) {
+  for (auto type : types) {
+    if (auto rtp = dyn_cast<RankedTensorType>(type)) {
+      const SparseTensorType stt(rtp);
+      if (stt.hasEncoding()) {
+        auto shape = {ShapedType::kDynamic};
+        // Convert the external representation of the values array.
+        auto vtp = RankedTensorType::get(shape, stt.getElementType());
----------------
aartbik wrote:

The idea is, of course, similar, but the details are very different (we could perhaps use one of the forall  loops  eventually). However, it is possible that external formats and internal memref formats will differ for different frameworks. In that case, this can be controlled by "framework options" eventually, doing slightly more subtle things.

https://github.com/llvm/llvm-project/pull/80326


More information about the Mlir-commits mailing list