diff options
| -rw-r--r-- | mlir/include/mlir/Dialect/Affine/IR/AffineOps.td | 14 | ||||
| -rw-r--r-- | mlir/lib/Dialect/Affine/IR/AffineOps.cpp | 11 | ||||
| -rw-r--r-- | mlir/lib/Dialect/Affine/IR/CMakeLists.txt | 1 | ||||
| -rw-r--r-- | mlir/test/python/dialects/affine.py | 11 |
4 files changed, 25 insertions, 12 deletions
diff --git a/mlir/include/mlir/Dialect/Affine/IR/AffineOps.td b/mlir/include/mlir/Dialect/Affine/IR/AffineOps.td index f9578cf3..c638646 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 a7fc7dd..7f2f3c3 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.assign(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 9e3c116..7f7a01b 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 9e20279..07b1722 100644 --- a/mlir/test/python/dialects/affine.py +++ b/mlir/test/python/dialects/affine.py @@ -44,6 +44,17 @@ def testAffineStoreOp(): return mem +# 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: %{{.*}}:2 = affine.delinearize_index %[[C1:.*]] into (%[[C1:.*]], %[[C0:.*]]) : index, index + two_indices = affine.AffineDelinearizeIndexOp(c1, [c1, c0]) + + # CHECK-LABEL: TEST: testAffineLoadOp @constructAndPrintInModule def testAffineLoadOp(): |