[Mlir-commits] [mlir] 98d8dce - [mlir][affine] implement inferType for delinearize (#74644)

[llvmlistbot at llvm.org]

Author: Maksim Levental
Date: 2023-12-07T15:59:52-06:00
New Revision: 98d8dce6e9e21a995f6a06fa4485fa529931be37

URL: https://github.com/llvm/llvm-project/commit/98d8dce6e9e21a995f6a06fa4485fa529931be37
DIFF: https://github.com/llvm/llvm-project/commit/98d8dce6e9e21a995f6a06fa4485fa529931be37.diff

LOG: [mlir][affine] implement inferType for delinearize (#74644)

Added: 
    

Modified: 
    mlir/include/mlir/Dialect/Affine/IR/AffineOps.td
    mlir/lib/Dialect/Affine/IR/AffineOps.cpp
    mlir/lib/Dialect/Affine/IR/CMakeLists.txt
    mlir/test/python/dialects/affine.py

Removed: 
    


################################################################################
diff  --git a/mlir/include/mlir/Dialect/Affine/IR/AffineOps.td b/mlir/include/mlir/Dialect/Affine/IR/AffineOps.td
index f9578cf37d5d76..c638646b9c3277 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 a7fc7ddec26e61..7f2f3c3410c33b 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 9e3c1161fd92a0..7f7a01be891e05 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 9e2027937a09fe..07b1722f124024 100644
--- a/mlir/test/python/dialects/affine.py
+++ b/mlir/test/python/dialects/affine.py
@@ -44,6 +44,17 @@ def affine_store_test(arg0):
         return mem
 
 
+# CHECK-LABEL: TEST: testAffineDelinearizeInfer
+ at 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():