Add AffineMaxOp

[merge_guards_bot]

Unit tests: pass. 62382 tests passed, 0 failed and 839 were skipped.

clang-tidy: pass.

clang-format: fail. Please format your changes with clang-format by running git-clang-format HEAD^ or applying this patch.

Build artifacts: diff.json, clang-tidy.txt, clang-format.patch, CMakeCache.txt, console-log.txt, test-results.xml

Pre-merge checks is in beta. Report issue. Please join beta or enable it for your project.

[nicolasvasilache]

please split out (or rebase, not sure) anything related to ReduceOp

[nicolasvasilache]

mlir/include/mlir/Dialect/AffineOps/AffineOps.td

237–250

this should be in a separate revision.

mlir/lib/Dialect/AffineOps/AffineOps.cpp

1941–1964

this should be in a separate revision.

2028

this should be in a separate revision and have a test in canonicalize.mlir.

[OuHangKresnik]

format

[OuHangKresnik]

Reply

mlir/include/mlir/Dialect/AffineOps/AffineOps.td

237–250

Since there are many common code between AffineMax and AffineMIn.
The introduction of reduceOp (class not def) here is to try to avoid repeated code.

I think the same code could be applied to Ops like avg / sum ops

Any better way ? If there is no such ReduceOp ?

mlir/lib/Dialect/AffineOps/AffineOps.cpp

2028

I wrote two tests in canonicalize.mlir which follows the tests of AffineMin, do you mean more tests needed there ?

[merge_guards_bot]

Unit tests: pass. 62382 tests passed, 0 failed and 839 were skipped.

clang-tidy: pass.

clang-format: pass.

Build artifacts: diff.json, clang-tidy.txt, clang-format.patch, CMakeCache.txt, console-log.txt, test-results.xml

Pre-merge checks is in beta. Report issue. Please join beta or enable it for your project.

[ftynse]

Thanks! I have a couple of comments inline. Major comments:

• the lowering looks incorrect;
• please call SomethingOp only the thing that is an actual Op in the IR, use OpBase or Base for common implementation details.

.gitignore

48

This should be a separate patch if you intend all of LLVM to follow this. Check out https://stackoverflow.com/questions/1753070/how-do-i-configure-git-to-ignore-some-files-locally if you want to make git ignore files without pushing your personal config upstream.

mlir/include/mlir/Dialect/AffineOps/AffineOps.td

237

This should not be called "Op", it sounds like it is a separate operation but it is not. I would suggest AffineMinMaxOpBase, a bit ugly but crystal clear.

237–250

Like I mentioned above, calling it an Op is misleading. But factoring out common code is the right approach.

I don't think we need affine.add, that can be perfectly expressed with affine.apply instead, unlike min and max.

238

I think we should mark both min and max as not having side effects (in a follow-up commit since min was not marked as such and the change may break something)

mlir/lib/Conversion/AffineToStandard/AffineToStandard.cpp

314

This looks suspicious: why are you calling "lowerAffineMap*Min*` in the pattern for max ?

mlir/lib/Dialect/AffineOps/AffineOps.cpp

1938

There's no such thing as AffineReduceOp.

1988

You can avoid this function by using something like return parseAffineReduceOp<$cppClass>(parser, result); in the ODS file.

2028

I'm fine with this change being in the same revision. It's small enough to review.

mlir/test/AffineOps/ops.mlir

84

Please don't use SSA names in CHECK lines, even if the file around does. See https://mlir.llvm.org/getting_started/TestingGuide/ for more details.

mlir/test/Transforms/lower-affine.mlir

618–619

This lowering looks like it computes min, not max.

[OuHangKresnik]

• Merge branch 'master' of https://github.com/llvm/llvm-project

[OuHangKresnik]

Thank you for the review! Replied with some questions.

mlir/include/mlir/Dialect/AffineOps/AffineOps.td

238

Ok, I will do this change in next commit?

mlir/test/AffineOps/ops.mlir

84

Here follows the affine_min test.
Change them both ? or I just change affine_max

[merge_guards_bot]

Unit tests: pass. 62414 tests passed, 0 failed and 839 were skipped.

clang-tidy: pass.

clang-format: pass.

Build artifacts: diff.json, clang-tidy.txt, clang-format.patch, CMakeCache.txt, console-log.txt, test-results.xml

Pre-merge checks is in beta. Report issue. Please join beta or enable it for your project.

[ftynse]

mlir/include/mlir/Dialect/AffineOps/AffineOps.td

238

Yes please.

mlir/test/AffineOps/ops.mlir

84

I'm not sure I understand the question. The source IR here is a "max" operation. The code you generate computes the minimum instead. It's incorrect and needs to be fixed. Existing tests are correct.

[ftynse]

Almost there, a couple more comments.

.gitignore

48

Spurious change, please remove this from the patch.

mlir/lib/Conversion/AffineToStandard/AffineToStandard.cpp

265–270

Please call lowerAffineMapMax here, the same way it is done for lowerAffineUpperBound and min below.

mlir/lib/Dialect/AffineOps/AffineOps.cpp

1952

I think it should be possible to do op.map() isntead of op.getAttr(T::getMapAttrName()).

1959

Use T::getMapAttrName() instead of hardcoded "map".

2031

Nit: update the comment to reflect what the code below does.

[OuHangKresnik]

• Merge branch 'master' of https://github.com/llvm/llvm-project

[OuHangKresnik]

Reply with fixes and several questions

mlir/lib/Dialect/AffineOps/AffineOps.cpp

1952

It will cause many test breaks such as :

/disk2/ouhang.oh/llvm-project/mlir/test/Dialect/Linalg/tile.mlir:257:12: error: TILE-2: expected string not found in input
// TILE-2: %[[szM:.*]] = affine.min #[[bound_map]](%[[C2]], %[[localM]], %[[I]])

and

<stdin>:37:18: note: with "MAP0" equal to "map0"

2031

Thank you, I will be more careful next time

mlir/test/AffineOps/ops.mlir

84

Here follows the affine_min test.
Change them both ? or I just change affine_max

I mean should I change affine_max here only or both affine_max and affine_min ?
Or keep them the same here and change both of them in next commit ?

The source IR here is a "max" operation. The code you generate computes the minimum instead

In my current understanding, the test here is correct ?

And I just fixed the lower-affine.mlir, due to misuse of the CmpIPredicate

[merge_guards_bot]

Unit tests: pass. 62414 tests passed, 0 failed and 839 were skipped.

clang-tidy: pass.

clang-format: pass.

Build artifacts: diff.json, clang-tidy.txt, clang-format.patch, CMakeCache.txt, console-log.txt, test-results.xml

Pre-merge checks is in beta. Report issue. Please join beta or enable it for your project.

[ftynse]

mlir/lib/Dialect/AffineOps/AffineOps.cpp

1952

It would be nice to understand _why_, but I'm fine with the current approach.

2031

It's generally better not to copy-paste ;)

mlir/test/AffineOps/ops.mlir

84

I mean should I change affine_max here only or both affine_max and affine_min ?
Or keep them the same here and change both of them in next commit ?

I don't see why do you need to change affine_min in any way here. It is generally bad practice to change the tests that are not related to the code changes you are making.

In my current understanding, the test here is correct ?

Looks okay to me.

.gitignore

 cscope.out
 autoconf/aclocal.m4
 autoconf/autom4te.cache
 /compile_commands.json
 # Visual Studio built-in CMake configuration
 /CMakeSettings.json
 # CLion project configuration
 /.idea
-
+/third_party
+/release
 #==============================================================================#
 # Directories to ignore (do not add trailing '/'s, they skip symlinks).
 #==============================================================================#
 # VS2017 and VSCode config files.
 .vscode
 .vs
 # clangd index
 .clangd

mlir/include/mlir/Dialect/AffineOps/AffineOps.td

OpBuilder getElseBodyBuilder() {
       Block &body = elseRegion().front();
       return OpBuilder(&body, std::prev(body.end()));
     }
   }];
 
   let hasFolder = 1;
 }
 
-def AffineMinOp : Affine_Op<"min"> {
+class Affine_ReduceOp <string mnemonic, list<OpTrait> traits = []> :
+    Op<Affine_Dialect, mnemonic, traits> {
+  let arguments = (ins AffineMapAttr:$map, Variadic<Index>:$operands);
+  let results = (outs Index);
+  let extraClassDeclaration = [{
+    static StringRef getMapAttrName() { return "map"; }
+  }];
+  let verifier = [{ return ::verifyReduceOp(*this); }];
+  let printer = [{ return ::printReduceOp(p, *this); }];
+  let parser = [{ return ::parse$cppClass(parser, result); }];
+  let hasFolder = 1;
+}
+
+def AffineMinOp : Affine_ReduceOp<"min"> {
   let summary = "min operation";
   let description = [{
     The "min" operation computes the minimum value result from a multi-result
     affine map.
 
     Example:
 
        %0 = affine.min (d0) -> (1000, d0 + 512) (%i0) : index
   }];
-  let arguments = (ins AffineMapAttr:$map, Variadic<Index>:$operands);
-  let results = (outs Index);
-  let extraClassDeclaration = [{
-    static StringRef getMapAttrName() { return "map"; }
+}
+
+def AffineMaxOp : Affine_ReduceOp<"max"> {
+  let summary = "max operation";
+  let description = [{
+    The "max" operation computes the maximum value result from a multi-result
+    affine map.
+
+    Example:
+
+       %0 = affine.max (d0) -> (1000, d0 + 512) (%i0) : index
   }];
-  let hasFolder = 1;
 }
 
 def AffinePrefetchOp : Affine_Op<"prefetch"> {
   let summary = "affine prefetch operation";
   let description = [{
     The "affine.prefetch" op prefetches data from a memref location described
     with an affine subscript similar to affine.load, and has three attributes:
     a read/write specifier, a locality hint, and a cache type specifier as shown

mlir/include/mlir/EDSC/Intrinsics.h

template <typename Op> struct OperationBuilder : public OperationHandle {
   OperationBuilder() : OperationHandle(OperationHandle::create<Op>()) {}
 };
 
 using addf = ValueBuilder<AddFOp>;
 using affine_apply = ValueBuilder<AffineApplyOp>;
 using affine_if = OperationBuilder<AffineIfOp>;
 using affine_load = ValueBuilder<AffineLoadOp>;
 using affine_min = ValueBuilder<AffineMinOp>;
+using affine_max = ValueBuilder<AffineMaxOp>;
 using affine_store = OperationBuilder<AffineStoreOp>;
 using alloc = ValueBuilder<AllocOp>;
 using call = OperationBuilder<mlir::CallOp>;
 using constant_float = ValueBuilder<ConstantFloatOp>;
 using constant_index = ValueBuilder<ConstantIndexOp>;
 using constant_int = ValueBuilder<ConstantIntOp>;
 using dealloc = OperationBuilder<DeallocOp>;
 using dim = ValueBuilder<DimOp>;

mlir/lib/Conversion/AffineToStandard/AffineToStandard.cpp

static Value buildMinMaxReductionSeq(Location loc, CmpIPredicate predicate,
 
   return value;
 }
 
 /// Emit instructions that correspond to the affine map in the lower bound
 /// applied to the respective operands, and compute the maximum value across
 /// the results.
 Value mlir::lowerAffineLowerBound(AffineForOp op, OpBuilder &builder) {
   auto lbValues = expandAffineMap(builder, op.getLoc(), op.getLowerBoundMap(),
                                   op.getLowerBoundOperands());
   if (!lbValues)
     return nullptr;
   return buildMinMaxReductionSeq(op.getLoc(), CmpIPredicate::sgt, *lbValues,
                                  builder);
 }
 
 /// Emit instructions that correspond to computing the minimum value amoung the
 /// values of a (potentially) multi-output affine map applied to `operands`.
 static Value lowerAffineMapMin(OpBuilder &builder, Location loc, AffineMap map,
                                ValueRange operands) {
   if (auto values = expandAffineMap(builder, loc, map, operands))
     return buildMinMaxReductionSeq(loc, CmpIPredicate::slt, *values, builder);
PatternMatchResult matchAndRewrite(AffineMinOp op,
     if (!reduced)
       return matchFailure();
 
     rewriter.replaceOp(op, reduced);
     return matchSuccess();
   }
 };
 
+class AffineMaxLowering : public OpRewritePattern<AffineMaxOp> {
+public:
+  using OpRewritePattern<AffineMaxOp>::OpRewritePattern;
+
+  PatternMatchResult matchAndRewrite(AffineMaxOp op,
+                                     PatternRewriter &rewriter) const override {
+    Value reduced =
+        lowerAffineMapMin(rewriter, op.getLoc(), op.map(), op.operands());
+    if (!reduced)
+      return matchFailure();
+
+    rewriter.replaceOp(op, reduced);
+    return matchSuccess();
+  }
+};
+
 /// Affine terminators are removed.
 class AffineTerminatorLowering : public OpRewritePattern<AffineTerminatorOp> {
 public:
   using OpRewritePattern<AffineTerminatorOp>::OpRewritePattern;
 
   PatternMatchResult matchAndRewrite(AffineTerminatorOp op,
                                      PatternRewriter &rewriter) const override {
     rewriter.replaceOpWithNewOp<loop::TerminatorOp>(op);
void mlir::populateAffineToStdConversionPatterns(
     OwningRewritePatternList &patterns, MLIRContext *ctx) {
   // clang-format off
   patterns.insert<
       AffineApplyLowering,
       AffineDmaStartLowering,
       AffineDmaWaitLowering,
       AffineLoadLowering,
       AffineMinLowering,
+      AffineMaxLowering,
       AffinePrefetchLowering,
       AffineStoreLowering,
       AffineForLowering,
       AffineIfLowering,
       AffineTerminatorLowering>(ctx);
   // clang-format on
 }
 

mlir/lib/Dialect/AffineOps/AffineOps.cpp

 
 LogicalResult AffineStoreOp::fold(ArrayRef<Attribute> cstOperands,
                                   SmallVectorImpl<OpFoldResult> &results) {
   /// store(memrefcast) -> store
   return foldMemRefCast(*this);
 }
 
 //===----------------------------------------------------------------------===//
-// AffineMinOp
+// AffineReduceOp
 //===----------------------------------------------------------------------===//
-//
-//   %0 = affine.min (d0) -> (1000, d0 + 512) (%i0)
-//
 
-static ParseResult parseAffineMinOp(OpAsmParser &parser,
+template <typename T>
+static LogicalResult verifyReduceOp(T op) {
+  // Verify that operand count matches affine map dimension and symbol count.
+  if (op.getNumOperands() != op.map().getNumDims() + op.map().getNumSymbols())
+    return op.emitOpError(
+        "operand count and affine map dimension and symbol count must match");
+  return success();
+}
+
+template <typename T>
+static void printReduceOp(OpAsmPrinter &p, T op) {
+  p << op.getOperationName() << ' ' << op.getAttr(T::getMapAttrName());
+  auto operands = op.getOperands();
+  unsigned numDims = op.map().getNumDims();
+  p << '(' << operands.take_front(numDims) << ')';
+
+  if (operands.size() != numDims)
+    p << '[' << operands.drop_front(numDims) << ']';
+  p.printOptionalAttrDict(op.getAttrs(), /*elidedAttrs=*/{"map"});
+}
+
+template <typename T>
+static ParseResult parseAffineReduceOp(OpAsmParser &parser,
                                        OperationState &result) {
   auto &builder = parser.getBuilder();
   auto indexType = builder.getIndexType();
   SmallVector<OpAsmParser::OperandType, 8> dim_infos;
   SmallVector<OpAsmParser::OperandType, 8> sym_infos;
   AffineMapAttr mapAttr;
   return failure(
-      parser.parseAttribute(mapAttr, AffineMinOp::getMapAttrName(),
-                            result.attributes) ||
+      parser.parseAttribute(mapAttr, T::getMapAttrName(), result.attributes) ||
       parser.parseOperandList(dim_infos, OpAsmParser::Delimiter::Paren) ||
       parser.parseOperandList(sym_infos,
                               OpAsmParser::Delimiter::OptionalSquare) ||
       parser.parseOptionalAttrDict(result.attributes) ||
       parser.resolveOperands(dim_infos, indexType, result.operands) ||
       parser.resolveOperands(sym_infos, indexType, result.operands) ||
       parser.addTypeToList(indexType, result.types));
 }
 
-static void print(OpAsmPrinter &p, AffineMinOp op) {
-  p << op.getOperationName() << ' '
-    << op.getAttr(AffineMinOp::getMapAttrName());
-  auto operands = op.getOperands();
-  unsigned numDims = op.map().getNumDims();
-  p << '(' << operands.take_front(numDims) << ')';
-
-  if (operands.size() != numDims)
-    p << '[' << operands.drop_front(numDims) << ']';
-  p.printOptionalAttrDict(op.getAttrs(), /*elidedAttrs=*/{"map"});
-}
+//===----------------------------------------------------------------------===//
+// AffineMinOp
+//===----------------------------------------------------------------------===//
+//
+//   %0 = affine.min (d0) -> (1000, d0 + 512) (%i0)
+//
 
-static LogicalResult verify(AffineMinOp op) {
-  // Verify that operand count matches affine map dimension and symbol count.
-  if (op.getNumOperands() != op.map().getNumDims() + op.map().getNumSymbols())
-    return op.emitOpError(
-        "operand count and affine map dimension and symbol count must match");
-  return success();
+static ParseResult parseAffineMinOp(OpAsmParser &parser,
+                                    OperationState &result) {
+  return parseAffineReduceOp<AffineMinOp>(parser, result);
 }
 
 OpFoldResult AffineMinOp::fold(ArrayRef<Attribute> operands) {
   // Fold the affine map.
   // TODO(andydavis, ntv) Fold more cases: partial static information,
   // min(some_affine, some_affine + constant, ...).
   SmallVector<Attribute, 2> results;
   if (failed(map().constantFold(operands, results)))
for (unsigned i = 0, e = results.size(); i < e; ++i) {
     }
   }
   if (minIndex < 0)
     return {};
   return results[minIndex];
 }
 
 //===----------------------------------------------------------------------===//
+// AffineMaxOp
+//===----------------------------------------------------------------------===//
+//
+//   %0 = affine.max (d0) -> (1000, d0 + 512) (%i0)
+//
+
+static ParseResult parseAffineMaxOp(OpAsmParser &parser,
+                                    OperationState &result) {
+  return parseAffineReduceOp<AffineMaxOp>(parser, result);
+}
+
+OpFoldResult AffineMaxOp::fold(ArrayRef<Attribute> operands) {
+  // Fold the affine map.
+  // TODO(andydavis, ntv, ouhang) Fold more cases: partial static information,
+  // max(some_affine, some_affine + constant, ...).
+  SmallVector<Attribute, 2> results;
+  if (failed(map().constantFold(operands, results)))
+    return {};
+
+  // Compute and return min of folded map results.
+  int64_t max = std::numeric_limits<int64_t>::min();
+  int maxIndex = -1;
+  for (unsigned i = 0, e = results.size(); i < e; ++i) {
+    auto intAttr = results[i].cast<IntegerAttr>();
+    if (intAttr.getInt() > max) {
+      max = intAttr.getInt();
+      maxIndex = i;
+    }
+  }
+  if (maxIndex < 0)
+    return {};
+  return results[maxIndex];
+}
+
+//===----------------------------------------------------------------------===//
 // AffinePrefetchOp
 //===----------------------------------------------------------------------===//
 
 //
 // affine.prefetch %0[%i, %j + 5], read, locality<3>, data : memref<400x400xi32>
 //
 static ParseResult parseAffinePrefetchOp(OpAsmParser &parser,
                                          OperationState &result) {

mlir/test/AffineOps/canonicalize.mlir

func @affine_min(%arg0 : index, %arg1 : index, %arg2 : index) {
   %c20 = constant 20 : index
   %0 = affine.min affine_map<(d0)[s0] -> (1000, d0 floordiv 4, (s0 mod 5) + 1)> (%c20)[%c3]
   "op0"(%0) : (index) -> ()
   // CHECK:       %[[CST:.*]] = constant 4 : index
   // CHECK-NEXT:  "op0"(%[[CST]]) : (index) -> ()
   // CHECK-NEXT:  return
   return
 }
+
+// -----
+
+func @affine_max(%arg0 : index, %arg1 : index, %arg2 : index) {
+  %c511 = constant 511 : index
+  %c1 = constant 0 : index
+  %0 = affine.max affine_map<(d0)[s0] -> (1000, d0 + 512, s0 + 1)> (%c1)[%c511]
+  "op0"(%0) : (index) -> ()
+  // CHECK:       %[[CST:.*]] = constant 1000 : index
+  // CHECK-NEXT:  "op0"(%[[CST]]) : (index) -> ()
+  // CHECK-NEXT:  return
+  return
+}
+
+// -----
+
+func @affine_max(%arg0 : index, %arg1 : index, %arg2 : index) {
+  %c3 = constant 3 : index
+  %c20 = constant 20 : index
+  %0 = affine.max affine_map<(d0)[s0] -> (1000, d0 floordiv 4, (s0 mod 5) + 1)> (%c20)[%c3]
+  "op0"(%0) : (index) -> ()
+  // CHECK:       %[[CST:.*]] = constant 1000 : index
+  // CHECK-NEXT:  "op0"(%[[CST]]) : (index) -> ()
+  // CHECK-NEXT:  return
+  return
+}
 No newline at end of file

mlir/test/AffineOps/invalid.mlir

 
 // CHECK-LABEL: @affine_min
 func @affine_min(%arg0 : index, %arg1 : index, %arg2 : index) {
   // expected-error@+1 {{operand count and affine map dimension and symbol count must match}}
   %0 = affine.min affine_map<(d0) -> (d0)> ()
 
   return
 }
+
+// -----
+
+// CHECK-LABEL: @affine_max
+func @affine_max(%arg0 : index, %arg1 : index, %arg2 : index) {
+  // expected-error@+1 {{operand count and affine map dimension and symbol count must match}}
+  %0 = affine.max affine_map<(d0) -> (d0)> (%arg0, %arg1)
+
+  return
+}
+
+// -----
+
+// CHECK-LABEL: @affine_max
+func @affine_max(%arg0 : index, %arg1 : index, %arg2 : index) {
+  // expected-error@+1 {{operand count and affine map dimension and symbol count must match}}
+  %0 = affine.max affine_map<()[s0] -> (s0)> (%arg0, %arg1)
+
+  return
+}
+
+// -----
+
+// CHECK-LABEL: @affine_max
+func @affine_max(%arg0 : index, %arg1 : index, %arg2 : index) {
+  // expected-error@+1 {{operand count and affine map dimension and symbol count must match}}
+  %0 = affine.max affine_map<(d0) -> (d0)> ()
+
+  return
+}
 No newline at end of file

mlir/test/AffineOps/ops.mlir

func @affine_min(%arg0 : index, %arg1 : index, %arg2 : index) {
   %1 = affine.min affine_map<(d0, d1)[s0] -> (d0 - d1, s0 + 512)> (%arg0, %arg1)[%arg2]
   // CHECK: affine.min #[[MAP2]]()[%arg1, %arg2]
   %2 = affine.min affine_map<()[s0, s1] -> (s0 - s1, 11)> ()[%arg1, %arg2]
   // CHECK: affine.min #[[MAP3]]()
   %3 = affine.min affine_map<()[] -> (77, 78, 79)> ()[]
   return
 }
 
+// CHECK-LABEL: @affine_max
+func @affine_max(%arg0 : index, %arg1 : index, %arg2 : index) {
+  // CHECK: affine.max #[[MAP0]](%arg0)[%arg1]
+  %0 = affine.max affine_map<(d0)[s0] -> (1000, d0 + 512, s0)> (%arg0)[%arg1]
+  // CHECK: affine.max #[[MAP1]](%arg0, %arg1)[%arg2]
+  %1 = affine.max affine_map<(d0, d1)[s0] -> (d0 - d1, s0 + 512)> (%arg0, %arg1)[%arg2]
+  // CHECK: affine.max #[[MAP2]]()[%arg1, %arg2]
+  %2 = affine.max affine_map<()[s0, s1] -> (s0 - s1, 11)> ()[%arg1, %arg2]
+  // CHECK: affine.max #[[MAP3]]()
+  %3 = affine.max affine_map<()[] -> (77, 78, 79)> ()[]
+  return
+}
+
 // -----
 
 func @valid_symbols(%arg0: index, %arg1: index, %arg2: index) {
   %c0 = constant 1 : index
   %c1 = constant 0 : index
   %0 = alloc(%arg0, %arg1) : memref<?x?xf32>
   affine.for %arg3 = 0 to %arg2 step 768 {
     %13 = dim %0, 1 : memref<?x?xf32>

mlir/test/Transforms/lower-affine.mlir

func @affine_min(%arg0: index, %arg1: index) -> index{
   // CHECK: %[[Cm2:.*]] = constant -1
   // CHECK: %[[neg2:.*]] = muli %[[ARG0]], %[[Cm2:.*]]
   // CHECK: %[[second:.*]] = addi %[[ARG1]], %[[neg2]]
   // CHECK: %[[cmp:.*]] = cmpi "slt", %[[first]], %[[second]]
   // CHECK: select %[[cmp]], %[[first]], %[[second]]
   %0 = affine.min affine_map<(d0,d1) -> (d0 - d1, d1 - d0)>(%arg0, %arg1)
   return %0 : index
 }
+
+// CHECK-LABEL: func @affine_max
+// CHECK-SAME: %[[ARG0:.*]]: index, %[[ARG1:.*]]: index
+func @affine_max(%arg0: index, %arg1: index) -> index{
+  // CHECK: %[[Cm1:.*]] = constant -1
+  // CHECK: %[[neg1:.*]] = muli %[[ARG1]], %[[Cm1:.*]]
+  // CHECK: %[[first:.*]] = addi %[[ARG0]], %[[neg1]]
+  // CHECK: %[[Cm2:.*]] = constant -1
+  // CHECK: %[[neg2:.*]] = muli %[[ARG0]], %[[Cm2:.*]]
+  // CHECK: %[[second:.*]] = addi %[[ARG1]], %[[neg2]]
+  // CHECK: %[[cmp:.*]] = cmpi "slt", %[[first]], %[[second]]
+  // CHECK: select %[[cmp]], %[[first]], %[[second]]
+  %0 = affine.max affine_map<(d0,d1) -> (d0 - d1, d1 - d0)>(%arg0, %arg1)
+  return %0 : index
+}
 No newline at end of file