EDSC: emit composed affine maps again

The recent rework of MLIREmitter switched to using the generic call to
`builder.createOperation` from OperationState instead of individual customized
calls to `builder.create<>`.  As a result, regular non-composed affine apply
operations where emitted.  Introduce a special case in Expr::build to always
create composed affine maps instead, as it used to be the case before the
rework.

Such special-casing goes against the idea of EDSC generality and extensibility.
Instead, we should consider declaring the composed form canonical for
affine.apply operations and using the builder support for creating operations
and canonicalizing them immediately (ongoing effort).

PiperOrigin-RevId: 234790129

diff --git a/mlir/lib/EDSC/Types.cpp b/mlir/lib/EDSC/Types.cpp
index a07d7649af0af850f7a0a787564622b4ed407332..102ec923748a842a6ba7c0f71e19bd145dacb5cc 100644 (file)
--- a/mlir/lib/EDSC/Types.cpp
+++ b/mlir/lib/EDSC/Types.cpp
@@ -165,7 +165,6 @@ Expr::build(FuncBuilder &b,
     return {it->second};
 
   auto *impl = static_cast<ImplType *>(storage);
-  auto state = OperationState(b.getContext(), b.getUnknownLoc(), impl->opName);
   SmallVector<Value *, 4> operandValues;
   operandValues.reserve(impl->operands.size());
   for (auto child : impl->operands) {
@@ -174,6 +173,19 @@ Expr::build(FuncBuilder &b,
            "expected single-result expression as operand");
     operandValues.push_back(subResults.front());
   }
+
+  // Special case for emitting composed affine.applies.
+  // FIXME: this should not be a special case, instead, define composed form as
+  // canonical for the affine.apply operator and expose a generic createAndFold
+  // operation on builder that canonicalizes all operations that we emit here.
+  if (is_op<AffineApplyOp>()) {
+    auto affInstr = makeComposedAffineApply(
+        &b, b.getUnknownLoc(),
+        getAttribute("map").cast<AffineMapAttr>().getValue(), operandValues);
+    return {affInstr->getResult()};
+  }
+
+  auto state = OperationState(b.getContext(), b.getUnknownLoc(), impl->opName);
   state.addOperands(operandValues);
   state.addTypes(impl->resultTypes);
   for (const auto &attr : impl->attributes)

diff --git a/mlir/test/EDSC/for-loops.mlir b/mlir/test/EDSC/for-loops.mlir
index 9444a56cdd77a4cc7aecac0db9441effc3fe3f0a..3356ed390bc68bbaa6c894a9218f451c0c4c4ce2 100644 (file)
--- a/mlir/test/EDSC/for-loops.mlir
+++ b/mlir/test/EDSC/for-loops.mlir
@@ -4,8 +4,9 @@
 // CHECK-DAG: #[[idmap:.*]] = (d0) -> (d0)
 // CHECK-DAG: #[[diffmap:.*]] = (d0, d1) -> (d0 - d1)
 // CHECK-DAG: #[[addmap:.*]] = (d0, d1) -> (d0 + d1)
-// CHECK-DAG: #[[prodconstmap:.*]] = (d0, d1) -> (d0 * 3)
-// CHECK-DAG: #[[addconstmap:.*]] = (d0, d1) -> (d1 + 3)
+// CHECK-DAG: #[[prodconstmap:.*]] = (d0) -> (d0 * 3)
+// CHECK-DAG: #[[addconstmap:.*]] = (d0) -> (d0 + 3)
+// CHECK-DAG: #[[composedmap:.*]] = (d0, d1) -> (d0 * 3 + d1)
 // CHECK-DAG: #[[id2dmap:.*]] = (d0, d1) -> (d0, d1)
 
 // This function will be detected by the test pass that will insert
@@ -27,9 +28,9 @@ func @blocks() {
 // before the `return` instruction.
 // CHECK-LABEL: func @dynamic_for_func_args(%arg0: index, %arg1: index) {
 // CHECK:  for %i0 = #[[idmap]](%arg0) to #[[idmap]](%arg1) step 3 {
-// CHECK:  %[[res:.*]] = affine.apply #[[prodconstmap]](%arg0, %c3)
-// CHECK:  {{.*}} = affine.apply #[[addmap]](%[[res]], %arg1)
-// CHECK:  {{.*}} = affine.apply #[[addconstmap]](%c3, %arg0)
+// CHECK:  {{.*}} = affine.apply #[[prodconstmap]](%arg0)
+// CHECK:  {{.*}} = affine.apply #[[composedmap]](%arg0, %arg1)
+// CHECK:  {{.*}} = affine.apply #[[addconstmap]](%arg0)
 func @dynamic_for_func_args(%arg0 : index, %arg1 : index) {
   return
 }