[Mlir-commits] [mlir] [mlir][SparseTensor] Use explicit attribute APIs (PR #218904)
Mehdi Amini
llvmlistbot at llvm.org
Wed Aug 26 04:47:07 PDT 2026
https://github.com/joker-eph created https://github.com/llvm/llvm-project/pull/218904
Migrate SparseTensor IR, lowering, code generation, and loop emission to explicit discardable or operation-specific attribute access.
Assisted-by: Codex
>From b9320232fdebd69dd5198cea24549d286aedefad Mon Sep 17 00:00:00 2001
From: Mehdi Amini <joker.eph at gmail.com>
Date: Thu, 20 Aug 2026 06:32:53 -0700
Subject: [PATCH] [mlir][SparseTensor] Use explicit attribute APIs
Migrate SparseTensor IR, lowering, code generation, and loop emission to
explicit discardable or operation-specific attribute access.
Assisted-by: Codex
---
mlir/lib/Dialect/SparseTensor/IR/SparseTensorDialect.cpp | 2 +-
.../Dialect/SparseTensor/Transforms/SparseAssembler.cpp | 9 +++++----
.../Dialect/SparseTensor/Transforms/SparseGPUCodegen.cpp | 7 ++++---
.../SparseTensor/Transforms/SparseReinterpretMap.cpp | 4 ++--
.../SparseTensor/Transforms/SparseVectorization.cpp | 9 +++++----
.../Dialect/SparseTensor/Transforms/Sparsification.cpp | 4 ++--
.../SparseTensor/Transforms/Utils/CodegenUtils.cpp | 4 ++--
.../Dialect/SparseTensor/Transforms/Utils/LoopEmitter.h | 3 ++-
8 files changed, 23 insertions(+), 19 deletions(-)
diff --git a/mlir/lib/Dialect/SparseTensor/IR/SparseTensorDialect.cpp b/mlir/lib/Dialect/SparseTensor/IR/SparseTensorDialect.cpp
index 6c26a31bf88a0..781189210ecd3 100644
--- a/mlir/lib/Dialect/SparseTensor/IR/SparseTensorDialect.cpp
+++ b/mlir/lib/Dialect/SparseTensor/IR/SparseTensorDialect.cpp
@@ -1799,7 +1799,7 @@ static LogicalResult verifyNumBlockArgs(T *op, Region ®ion,
}
LogicalResult BinaryOp::verify() {
- NamedAttrList attrs = (*this)->getAttrs();
+ NamedAttrList attrs = (*this)->getDiscardableAttrDictionary().getValue();
Type leftType = getX().getType();
Type rightType = getY().getType();
Type outputType = getOutput().getType();
diff --git a/mlir/lib/Dialect/SparseTensor/Transforms/SparseAssembler.cpp b/mlir/lib/Dialect/SparseTensor/Transforms/SparseAssembler.cpp
index 4428339b38c5d..a5c74f6941554 100644
--- a/mlir/lib/Dialect/SparseTensor/Transforms/SparseAssembler.cpp
+++ b/mlir/lib/Dialect/SparseTensor/Transforms/SparseAssembler.cpp
@@ -233,12 +233,13 @@ struct SparseFuncAssembler : public OpRewritePattern<func::FuncOp> {
func::ReturnOp::create(rewriter, loc, outputs);
// Finally, migrate a potential c-interface property.
- if (funcOp->getAttrOfType<UnitAttr>(
+ if (funcOp->getDiscardableAttrOfType<UnitAttr>(
LLVM::LLVMDialect::getEmitCWrapperAttrName())) {
- func->setAttr(LLVM::LLVMDialect::getEmitCWrapperAttrName(),
- UnitAttr::get(context));
+ func->setDiscardableAttr(LLVM::LLVMDialect::getEmitCWrapperAttrName(),
+ UnitAttr::get(context));
rewriter.modifyOpInPlace(funcOp, [&]() {
- funcOp->removeAttr(LLVM::LLVMDialect::getEmitCWrapperAttrName());
+ funcOp->removeDiscardableAttr(
+ LLVM::LLVMDialect::getEmitCWrapperAttrName());
});
}
return success();
diff --git a/mlir/lib/Dialect/SparseTensor/Transforms/SparseGPUCodegen.cpp b/mlir/lib/Dialect/SparseTensor/Transforms/SparseGPUCodegen.cpp
index 14418b09f8354..8243bf9317c22 100644
--- a/mlir/lib/Dialect/SparseTensor/Transforms/SparseGPUCodegen.cpp
+++ b/mlir/lib/Dialect/SparseTensor/Transforms/SparseGPUCodegen.cpp
@@ -50,8 +50,8 @@ enum class CuSparseFormat {
/// Marks the given top module as a GPU container module.
static void markAsGPUContainer(ModuleOp topModule) {
- topModule->setAttr(gpu::GPUDialect::getContainerModuleAttrName(),
- UnitAttr::get(topModule->getContext()));
+ topModule->setDiscardableAttr(gpu::GPUDialect::getContainerModuleAttrName(),
+ UnitAttr::get(topModule->getContext()));
}
/// Constructs a new GPU module (for GPU kernels) inside the given top module,
@@ -1189,7 +1189,8 @@ struct ForallRewriter : public OpRewritePattern<scf::ParallelOp> {
// of the form
// forall (i = 0; i < N; i++)
// so that cyclic scheduling over the threads is easy.
- if (!forallOp->hasAttr(LoopEmitter::getLoopEmitterLoopAttrName()) ||
+ if (!forallOp->hasDiscardableAttr(
+ LoopEmitter::getLoopEmitterLoopAttrName()) ||
forallOp.getNumReductions() != 0 || forallOp.getNumLoops() != 1 ||
!matchPattern(forallOp.getLowerBound()[0], m_Zero()) ||
!matchPattern(forallOp.getStep()[0], m_One()))
diff --git a/mlir/lib/Dialect/SparseTensor/Transforms/SparseReinterpretMap.cpp b/mlir/lib/Dialect/SparseTensor/Transforms/SparseReinterpretMap.cpp
index f280377f2cfc2..87c715c425866 100644
--- a/mlir/lib/Dialect/SparseTensor/Transforms/SparseReinterpretMap.cpp
+++ b/mlir/lib/Dialect/SparseTensor/Transforms/SparseReinterpretMap.cpp
@@ -426,7 +426,7 @@ struct GenericOpScheduler : public OpRewritePattern<linalg::GenericOp> {
}
const StringRef sorted = "sorted";
- if (linalgOp->hasAttr(sorted))
+ if (linalgOp->hasDiscardableAttr(sorted))
return failure();
// Pass strategy to IterationGraphSorter.
@@ -468,7 +468,7 @@ struct GenericOpScheduler : public OpRewritePattern<linalg::GenericOp> {
// Marks the GenericOp to avoid recursive matching.
rewriter.modifyOpInPlace(linalgOp, [&]() {
- linalgOp->setAttr(sorted, rewriter.getBoolAttr(true));
+ linalgOp->setDiscardableAttr(sorted, rewriter.getBoolAttr(true));
});
// Already sorted.
diff --git a/mlir/lib/Dialect/SparseTensor/Transforms/SparseVectorization.cpp b/mlir/lib/Dialect/SparseTensor/Transforms/SparseVectorization.cpp
index c60ca523d81f3..548eef578c4ff 100644
--- a/mlir/lib/Dialect/SparseTensor/Transforms/SparseVectorization.cpp
+++ b/mlir/lib/Dialect/SparseTensor/Transforms/SparseVectorization.cpp
@@ -585,9 +585,9 @@ static bool vectorizeStmt(PatternRewriter &rewriter, scf::ForOp forOp, VL vl,
scf::ForOp::create(rewriter, loc, forOp.getLowerBound(),
forOp.getUpperBound(), step, vinit,
/*bodyBuilder=*/nullptr, forOp.getUnsignedCmp());
- forOpNew->setAttr(
+ forOpNew->setDiscardableAttr(
LoopEmitter::getLoopEmitterLoopAttrName(),
- forOp->getAttr(LoopEmitter::getLoopEmitterLoopAttrName()));
+ forOp->getDiscardableAttr(LoopEmitter::getLoopEmitterLoopAttrName()));
rewriter.setInsertionPointToStart(forOpNew.getBody());
} else {
rewriter.modifyOpInPlace(forOp, [&]() { forOp.setStep(step); });
@@ -666,7 +666,7 @@ struct ForOpRewriter : public OpRewritePattern<scf::ForOp> {
// sparsifier, which means no data dependence analysis is required,
// and its loop-body is very restricted in form.
if (!op.getRegion().hasOneBlock() || !isOneInteger(op.getStep()) ||
- !op->hasAttr(LoopEmitter::getLoopEmitterLoopAttrName()))
+ !op->hasDiscardableAttr(LoopEmitter::getLoopEmitterLoopAttrName()))
return failure();
// Analyze (!codegen) and rewrite (codegen) loop-body.
if (vectorizeStmt(rewriter, op, vl, /*codegen=*/false) &&
@@ -683,7 +683,8 @@ static LogicalResult cleanReducChain(PatternRewriter &rewriter, Operation *op,
Value inp) {
if (auto redOp = inp.getDefiningOp<vector::ReductionOp>()) {
if (auto forOp = redOp.getVector().getDefiningOp<scf::ForOp>()) {
- if (forOp->hasAttr(LoopEmitter::getLoopEmitterLoopAttrName())) {
+ if (forOp->hasDiscardableAttr(
+ LoopEmitter::getLoopEmitterLoopAttrName())) {
rewriter.replaceOp(op, redOp.getVector());
return success();
}
diff --git a/mlir/lib/Dialect/SparseTensor/Transforms/Sparsification.cpp b/mlir/lib/Dialect/SparseTensor/Transforms/Sparsification.cpp
index 6004ab26f4663..860ea68029c3a 100644
--- a/mlir/lib/Dialect/SparseTensor/Transforms/Sparsification.cpp
+++ b/mlir/lib/Dialect/SparseTensor/Transforms/Sparsification.cpp
@@ -876,7 +876,7 @@ static void finalizeWhileOp(CodegenEnv &env, OpBuilder &builder,
while (auto ifOp = dyn_cast_or_null<scf::IfOp>(
builder.getInsertionBlock()->getParentOp())) {
// Break on IfOp for slicing filtering.
- if (ifOp->getAttr(LoopEmitter::getLoopEmitterLoopAttrName()) ==
+ if (ifOp->getDiscardableAttr(LoopEmitter::getLoopEmitterLoopAttrName()) ==
StringAttr::get(ifOp->getContext(), "slice"))
break;
@@ -1411,7 +1411,7 @@ struct GenericOpSparsifier : public OpRewritePattern<linalg::GenericOp> {
return failure();
// Only accept scheduled loops.
- if (!op->hasAttr("sorted")) {
+ if (!op->hasDiscardableAttr("sorted")) {
return rewriter.notifyMatchFailure(
op, "Loops not yet scheduled, try run --sparse-reinterpret-map "
"before sparsification.");
diff --git a/mlir/lib/Dialect/SparseTensor/Transforms/Utils/CodegenUtils.cpp b/mlir/lib/Dialect/SparseTensor/Transforms/Utils/CodegenUtils.cpp
index 3c6a905ebb696..1d9bf51f8d25a 100644
--- a/mlir/lib/Dialect/SparseTensor/Transforms/Utils/CodegenUtils.cpp
+++ b/mlir/lib/Dialect/SparseTensor/Transforms/Utils/CodegenUtils.cpp
@@ -332,8 +332,8 @@ FlatSymbolRefAttr mlir::sparse_tensor::getFunc(ModuleOp module, StringRef name,
FunctionType::get(context, operands.getTypes(), resultType));
func.setPrivate();
if (static_cast<bool>(emitCInterface))
- func->setAttr(LLVM::LLVMDialect::getEmitCWrapperAttrName(),
- UnitAttr::get(context));
+ func->setDiscardableAttr(LLVM::LLVMDialect::getEmitCWrapperAttrName(),
+ UnitAttr::get(context));
}
return result;
}
diff --git a/mlir/lib/Dialect/SparseTensor/Transforms/Utils/LoopEmitter.h b/mlir/lib/Dialect/SparseTensor/Transforms/Utils/LoopEmitter.h
index 355804bcc33fb..a47062987b4b6 100644
--- a/mlir/lib/Dialect/SparseTensor/Transforms/Utils/LoopEmitter.h
+++ b/mlir/lib/Dialect/SparseTensor/Transforms/Utils/LoopEmitter.h
@@ -257,7 +257,8 @@ class LoopEmitter {
: tidLvls(tidLvls), loop(loop), userCodeBlock(userBlock), iv(iv) {
// Attached a special tag to loop emitter generated loop.
if (loopTag)
- loop->setAttr(LoopEmitter::getLoopEmitterLoopAttrName(), loopTag);
+ loop->setDiscardableAttr(LoopEmitter::getLoopEmitterLoopAttrName(),
+ loopTag);
}
// The set of <tensor, lvl>, with *only* trivial index expressions, that are
// used as the condition for the generated loop. Extra information is
More information about the Mlir-commits
mailing list