[Mlir-commits] [mlir] [mlir][affine] implement inferType for delinearize (PR #74644)
Maksim Levental
llvmlistbot at llvm.org
Thu Dec 7 11:26:30 PST 2023
https://github.com/makslevental updated https://github.com/llvm/llvm-project/pull/74644
>From 8effd1a5ebf7e8225b82aca3b1dffeee7a7bc14c Mon Sep 17 00:00:00 2001
From: max <maksim.levental at gmail.com>
Date: Wed, 6 Dec 2023 13:01:46 -0600
Subject: [PATCH 1/2] [mlir][affine] implement inferType for delinearize
---
mlir/include/mlir/Dialect/Affine/IR/AffineOps.td | 14 ++------------
mlir/lib/Dialect/Affine/IR/AffineOps.cpp | 11 +++++++++++
mlir/lib/Dialect/Affine/IR/CMakeLists.txt | 1 +
mlir/test/python/dialects/affine.py | 10 ++++++++++
4 files changed, 24 insertions(+), 12 deletions(-)
diff --git a/mlir/include/mlir/Dialect/Affine/IR/AffineOps.td b/mlir/include/mlir/Dialect/Affine/IR/AffineOps.td
index f9578cf37d5d7..c638646b9c327 100644
--- a/mlir/include/mlir/Dialect/Affine/IR/AffineOps.td
+++ b/mlir/include/mlir/Dialect/Affine/IR/AffineOps.td
@@ -16,6 +16,7 @@
include "mlir/Dialect/Arith/IR/ArithBase.td"
include "mlir/Dialect/Affine/IR/AffineMemoryOpInterfaces.td"
include "mlir/Interfaces/ControlFlowInterfaces.td"
+include "mlir/Interfaces/InferTypeOpInterface.td"
include "mlir/Interfaces/LoopLikeInterface.td"
include "mlir/Interfaces/SideEffectInterfaces.td"
@@ -63,10 +64,6 @@ def AffineApplyOp : Affine_Op<"apply", [Pure]> {
// has a constant builder. That way we wouldn't need to explicitly specify the
// result types here.
let builders = [
- OpBuilder<(ins "AffineMap":$map, "ValueRange":$mapOperands),
- [{
- build($_builder, $_state, $_builder.getIndexType(), map, mapOperands);
- }]>,
OpBuilder<(ins "ArrayRef<AffineExpr> ":$exprList,"ValueRange":$mapOperands),
[{
build($_builder, $_state, $_builder.getIndexType(),
@@ -541,13 +538,6 @@ class AffineMinMaxOpBase<string mnemonic, list<Trait> traits = []> :
let arguments = (ins AffineMapAttr:$map, Variadic<Index>:$operands);
let results = (outs Index);
- let builders = [
- OpBuilder<(ins "AffineMap":$affineMap, "ValueRange":$mapOperands),
- [{
- build($_builder, $_state, $_builder.getIndexType(), affineMap, mapOperands);
- }]>
- ];
-
let extraClassDeclaration = [{
static StringRef getMapAttrStrName() { return "map"; }
AffineMap getAffineMap() { return getMap(); }
@@ -1068,7 +1058,7 @@ def AffineVectorStoreOp : AffineStoreOpBase<"vector_store"> {
//===----------------------------------------------------------------------===//
def AffineDelinearizeIndexOp : Affine_Op<"delinearize_index",
- [Pure]> {
+ [Pure, DeclareOpInterfaceMethods<InferTypeOpInterface>]> {
let summary = "delinearize an index";
let description = [{
The `affine.delinearize_index` operation takes a single index value and
diff --git a/mlir/lib/Dialect/Affine/IR/AffineOps.cpp b/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
index a7fc7ddec26e6..b9a8176f32a42 100644
--- a/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
+++ b/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
@@ -4474,6 +4474,17 @@ LogicalResult AffineVectorStoreOp::verify() {
// DelinearizeIndexOp
//===----------------------------------------------------------------------===//
+LogicalResult AffineDelinearizeIndexOp::inferReturnTypes(
+ MLIRContext *context, std::optional<::mlir::Location> location,
+ ValueRange operands, DictionaryAttr attributes, OpaqueProperties properties,
+ RegionRange regions, SmallVectorImpl<Type> &inferredReturnTypes) {
+ AffineDelinearizeIndexOpAdaptor adaptor(operands, attributes, properties,
+ regions);
+ inferredReturnTypes =
+ SmallVector<Type>(adaptor.getBasis().size(), IndexType::get(context));
+ return success();
+}
+
void AffineDelinearizeIndexOp::build(OpBuilder &builder, OperationState &result,
Value linearIndex,
ArrayRef<OpFoldResult> basis) {
diff --git a/mlir/lib/Dialect/Affine/IR/CMakeLists.txt b/mlir/lib/Dialect/Affine/IR/CMakeLists.txt
index 9e3c1161fd92a..7f7a01be891e0 100644
--- a/mlir/lib/Dialect/Affine/IR/CMakeLists.txt
+++ b/mlir/lib/Dialect/Affine/IR/CMakeLists.txt
@@ -15,6 +15,7 @@ add_mlir_dialect_library(MLIRAffineDialect
MLIRArithDialect
MLIRDialectUtils
MLIRIR
+ MLIRInferTypeOpInterface
MLIRLoopLikeInterface
MLIRMemRefDialect
MLIRShapedOpInterfaces
diff --git a/mlir/test/python/dialects/affine.py b/mlir/test/python/dialects/affine.py
index df42f8fcf1a57..cb3457934dbb5 100644
--- a/mlir/test/python/dialects/affine.py
+++ b/mlir/test/python/dialects/affine.py
@@ -5,6 +5,7 @@
from mlir.dialects import arith
from mlir.dialects import memref
from mlir.dialects import affine
+import mlir.extras.types as T
def constructAndPrintInModule(f):
@@ -43,6 +44,15 @@ def affine_store_test(arg0):
return mem
+# CHECK-LABEL: TEST: testAffineDelinearizeInfer
+ at constructAndPrintInModule
+def testAffineDelinearizeInfer():
+ c0 = arith.ConstantOp(T.index(), 0)
+ c1 = arith.ConstantOp(T.index(), 1)
+ # CHECK: %0:2 = affine.delinearize_index %c1 into (%c1, %c0) : index, index
+ two_indices = affine.AffineDelinearizeIndexOp(c1, [c1, c0])
+
+
# CHECK-LABEL: TEST: testAffineLoadOp
@constructAndPrintInModule
def testAffineLoadOp():
>From 7ace5a593e1e4f961b29e5e46406c659cee607c8 Mon Sep 17 00:00:00 2001
From: max <maksim.levental at gmail.com>
Date: Thu, 7 Dec 2023 13:26:20 -0600
Subject: [PATCH 2/2] incorporate comments
---
mlir/lib/Dialect/Affine/IR/AffineOps.cpp | 4 ++--
mlir/test/python/dialects/affine.py | 4 +++-
2 files changed, 5 insertions(+), 3 deletions(-)
diff --git a/mlir/lib/Dialect/Affine/IR/AffineOps.cpp b/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
index b9a8176f32a42..7f2f3c3410c33 100644
--- a/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
+++ b/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
@@ -4480,8 +4480,8 @@ LogicalResult AffineDelinearizeIndexOp::inferReturnTypes(
RegionRange regions, SmallVectorImpl<Type> &inferredReturnTypes) {
AffineDelinearizeIndexOpAdaptor adaptor(operands, attributes, properties,
regions);
- inferredReturnTypes =
- SmallVector<Type>(adaptor.getBasis().size(), IndexType::get(context));
+ inferredReturnTypes.assign(adaptor.getBasis().size(),
+ IndexType::get(context));
return success();
}
diff --git a/mlir/test/python/dialects/affine.py b/mlir/test/python/dialects/affine.py
index cb3457934dbb5..bfbf7f80cd4a6 100644
--- a/mlir/test/python/dialects/affine.py
+++ b/mlir/test/python/dialects/affine.py
@@ -47,9 +47,11 @@ def affine_store_test(arg0):
# CHECK-LABEL: TEST: testAffineDelinearizeInfer
@constructAndPrintInModule
def testAffineDelinearizeInfer():
+ # CHECK: %[[C0:.*]] = arith.constant 0 : index
c0 = arith.ConstantOp(T.index(), 0)
+ # CHECK: %[[C1:.*]] = arith.constant 1 : index
c1 = arith.ConstantOp(T.index(), 1)
- # CHECK: %0:2 = affine.delinearize_index %c1 into (%c1, %c0) : index, index
+ # CHECK: %{{.*}}:2 = affine.delinearize_index %[[C1:.*]] into (%[[C1:.*]], %[[C0:.*]]) : index, index
two_indices = affine.AffineDelinearizeIndexOp(c1, [c1, c0])
More information about the Mlir-commits
mailing list