[mlir][Affine] Add affine.delinearize_index operation

[christopherbate]

Summary

This change adds a new AffineDelinearizeIndexOp to the affine dialect.
The operation accepts an index type as well as a basis (array of index
values) representing how the index should be decomposed into a
multi-index. The decomposition obeys a canonical semantic that treats
the final basis element as "fastest varying" and the first basis element
as "slowest varying". A naive lowering of the operation using a sequence
of AffineApplyOps is given.

RFC was discussed on discourse here: https://discourse.llvm.org/t/rfc-tensor-extracting-slices-from-tensor-collapse-shape/64034

[bondhugula]

mlir/include/mlir/Dialect/Affine/IR/AffineOps.td

1078

It's perhaps more readable if we added a keyword like "into" or "to" between %linear_index and the parenthetical list.

mlir/test/Dialect/Affine/affine-expand-index-ops.mlir

6–12

We should just be using affine.apply ops here, which are a higher and a better level abstraction than arith ops; the former will enable better optimization: for eg., the delinearize_index results will typically flow into affine load/store ops and affine.apply ops will compose into them. And the affine.apply ops will anyway be converted to the arith dialect ones with lower-affine.

arith.divui -> floordiv
arith.remui -> mod

[christopherbate]

Address comments

[nicolasvasilache]

Thanks @christopherbate, I consider the comments unambiguously addressed.
Please update the commit message that still mentions arith ops to use affine_apply instead.

LGTM predicated on the commit message update!

mlir/test/Dialect/Affine/affine-expand-index-ops.mlir

6–12

+1 to avoid jumping the abstraction gap and using affine.apply.

mlir/test/Dialect/Affine/invalid.mlir

504

newline

[nicolasvasilache]

mlir/include/mlir/Dialect/Affine/IR/AffineOps.td

1078

the asm format was updated as requested but not the doc, please udpate it too before landing

[bondhugula]

LGTM - but several minor comments.

mlir/include/mlir/Dialect/Affine/IR/AffineOps.td

1104

Nit: drop extra blank line.

mlir/include/mlir/Dialect/Affine/Utils.h

307

Comments to be terminated with a period - here and below.

313

Nit: Extra space before and.

313

a -> lhs, b -> rhs?

316

basis is absent.

mlir/lib/Dialect/Affine/Transforms/AffineExpandIndexOps.cpp

14–20

Can combine these together.

56

Empty private section; please drop.

mlir/lib/Dialect/Affine/Utils/Utils.cpp

1844

Use ... e = ...; i < e form.

1856

You can simply use dimSizes.drop_front(i).

1863

results.reserve(divisors.size());

[christopherbate]

Address comments

mlir/include/mlir/Dialect/Affine/IR/AffineOps.td

let extraClassDeclaration = extraClassDeclarationBase # [{
     }
   }];
 
   let hasCanonicalizer = 1;
   let hasCustomAssemblyFormat = 1;
   let hasVerifier = 1;
 }
 
+//===----------------------------------------------------------------------===//
+// AffineDelinearizeIndexOp
+//===----------------------------------------------------------------------===//
+
+def AffineDelinearizeIndexOp : Affine_Op<"delinearize_index",
+    [NoSideEffect]> {
+  let summary = "delinearize an index";
+  let description = [{
+    The `affine.delinearize_index` operation takes a single index value and
+    calculates the multi-index according to the given basis.
+
+    Example:
+
+    ```
+    %indices:3 = affine.delinearize_index %linear_index into (%c16, %c224, %c224) : index, index, index
+    ```
+
+    In the above example, `%indices:3` conceptually holds the following:
+
+    ```
+    #map0 = affine_map<()[s0] -> (s0 floordiv 50176)>
+    #map1 = affine_map<()[s0] -> ((s0 mod 50176) floordiv 224)>
+    #map2 = affine_map<()[s0] -> (s0 mod 224)>
+    %indices_0 = affine.apply #map0()[%linear_index]
+    %indices_1 = affine.apply #map1()[%linear_index]
+    %indices_2 = affine.apply #map2()[%linear_index]
+    ```
+  }];
+
+  let arguments = (ins Index:$linear_index, Variadic<Index>:$basis);
+  let results = (outs Variadic<Index>:$multi_index);
+
+  let assemblyFormat = [{
+    $linear_index `into` ` ` `(` $basis `)` attr-dict `:` type($multi_index)
+  }];
+
+  let builders = [
+    OpBuilder<(ins "Value":$linear_index, "ArrayRef<OpFoldResult>":$basis)>
+  ];
+
+  let hasVerifier = 1;
+}
+
 #endif // AFFINE_OPS

mlir/include/mlir/Dialect/Affine/Passes.h

 
 /// Creates a pass to vectorize loops, operations and data types using a
 /// target-independent, n-D super-vector abstraction.
 std::unique_ptr<OperationPass<func::FuncOp>>
 createSuperVectorizePass(ArrayRef<int64_t> virtualVectorSize);
 /// Overload relying on pass options for initialization.
 std::unique_ptr<OperationPass<func::FuncOp>> createSuperVectorizePass();
 
+/// Populate patterns that expand affine index operations into more fundamental
+/// operations (not necessarily restricted to Affine dialect).
+void populateAffineExpandIndexOpsPatterns(RewritePatternSet &patterns);
+
+/// Creates a pass to expand affine index operations into more fundamental
+/// operations (not necessarily restricted to Affine dialect).
+std::unique_ptr<Pass> createAffineExpandIndexOpsPass();
+
 //===----------------------------------------------------------------------===//
 // Registration
 //===----------------------------------------------------------------------===//
 
 /// Generate the code for registering passes.
 #define GEN_PASS_REGISTRATION
 #include "mlir/Dialect/Affine/Passes.h.inc"
 
 } // namespace mlir
 
 #endif // MLIR_DIALECT_AFFINE_PASSES_H

mlir/include/mlir/Dialect/Affine/Passes.td

 }
 
 def SimplifyAffineStructures : Pass<"affine-simplify-structures", "func::FuncOp"> {
   let summary = "Simplify affine expressions in maps/sets and normalize "
                 "memrefs";
   let constructor = "mlir::createSimplifyAffineStructuresPass()";
 }
 
+def AffineExpandIndexOps : Pass<"affine-expand-index-ops"> {
+  let summary = "Lower affine operations operating on indices into more fundamental operations";
+  let constructor = "mlir::createAffineExpandIndexOpsPass()";
+}
+
 #endif // MLIR_DIALECT_AFFINE_PASSES

mlir/include/mlir/Dialect/Affine/Utils.h

 
 /// Create a sequence of operations that implement the `affineMap` applied to
 /// the given `operands` (as it it were an AffineApplyOp).
 Optional<SmallVector<Value, 8>> expandAffineMap(OpBuilder &builder,
                                                 Location loc,
                                                 AffineMap affineMap,
                                                 ValueRange operands);
 
+/// Holds the result of (div a, b)  and (mod a, b).
+struct DivModValue {
+  Value quotient;
+  Value remainder;
+};
+
+/// Create IR to calculate (div lhs, rhs) and (mod lhs, rhs).
+DivModValue getDivMod(OpBuilder &b, Location loc, Value lhs, Value rhs);
+
+/// Generate the IR to delinearize `linearIndex` given the `basis` and return
+/// the multi-index.
+FailureOr<SmallVector<Value>> delinearizeIndex(OpBuilder &b, Location loc,
+                                               Value linearIndex,
+                                               ArrayRef<Value> basis);
+
 } // namespace mlir
 
 #endif // MLIR_DIALECT_AFFINE_UTILS_H

mlir/lib/Dialect/Affine/IR/AffineOps.cpp

LogicalResult AffineVectorStoreOp::verify() {
 
   if (failed(verifyVectorMemoryOp(*this, memrefType, getVectorType())))
     return failure();
 
   return success();
 }
 
 //===----------------------------------------------------------------------===//
+// DelinearizeIndexOp
+//===----------------------------------------------------------------------===//
+
+void AffineDelinearizeIndexOp::build(OpBuilder &builder, OperationState &result,
+                                     Value linear_index,
+                                     ArrayRef<OpFoldResult> basis) {
+  result.addTypes(SmallVector<Type>(basis.size(), builder.getIndexType()));
+  result.addOperands(linear_index);
+  SmallVector<Value> basisValues =
+      llvm::to_vector(llvm::map_range(basis, [&](OpFoldResult ofr) -> Value {
+        Optional<int64_t> staticDim = getConstantIntValue(ofr);
+        if (staticDim.has_value())
+          return builder.create<arith::ConstantIndexOp>(result.location,
+                                                        *staticDim);
+        return ofr.dyn_cast<Value>();
+      }));
+  result.addOperands(basisValues);
+}
+
+LogicalResult AffineDelinearizeIndexOp::verify() {
+  if (getBasis().empty())
+    return emitOpError("basis should not be empty");
+  if (getNumResults() != getBasis().size())
+    return emitOpError("should return an index for each basis element");
+  return success();
+}
+
+//===----------------------------------------------------------------------===//
 // TableGen'd op method definitions
 //===----------------------------------------------------------------------===//
 
 #define GET_OP_CLASSES
 #include "mlir/Dialect/Affine/IR/AffineOps.cpp.inc"

mlir/lib/Dialect/Affine/IR/CMakeLists.txt

 add_mlir_dialect_library(MLIRAffineDialect
   AffineMemoryOpInterfaces.cpp
   AffineOps.cpp
   AffineValueMap.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/Affine
 
   DEPENDS
   MLIRAffineMemoryOpInterfacesIncGen
   MLIRAffineOpsIncGen
 
   LINK_LIBS PUBLIC
   MLIRArithmeticDialect
+  MLIRDialectUtils
   MLIRIR
   MLIRLoopLikeInterface
   MLIRMemRefDialect
   MLIRSideEffectInterfaces
   MLIRTensorDialect
   )

mlir/lib/Dialect/Affine/Transforms/AffineExpandIndexOps.cpp

+//===- AffineExpandIndexOps.cpp - Affine expand index ops pass ------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// This file implements a pass to expand affine index ops into one or more more
+// fundamental operations.
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Dialect/Affine/IR/AffineOps.h"
+#include "mlir/Dialect/Affine/Passes.h"
+
+#include "PassDetail.h"
+#include "mlir/Dialect/Affine/Utils.h"
+#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
+
+using namespace mlir;
+
+namespace {
+/// Lowers `affine.delinearize_index` into a sequence of division and remainder
+/// operations.
+struct LowerDelinearizeIndexOps
+    : public OpRewritePattern<AffineDelinearizeIndexOp> {
+  using OpRewritePattern<AffineDelinearizeIndexOp>::OpRewritePattern;
+  LogicalResult matchAndRewrite(AffineDelinearizeIndexOp op,
+                                PatternRewriter &rewriter) const override {
+    FailureOr<SmallVector<Value>> multiIndex =
+        delinearizeIndex(rewriter, op->getLoc(), op.getLinearIndex(),
+                         llvm::to_vector(op.getBasis()));
+    if (failed(multiIndex))
+      return failure();
+    rewriter.replaceOp(op, *multiIndex);
+    return success();
+  }
+};
+
+class ExpandAffineIndexOpsPass
+    : public AffineExpandIndexOpsBase<ExpandAffineIndexOpsPass> {
+public:
+  ExpandAffineIndexOpsPass() = default;
+
+  void runOnOperation() override {
+    MLIRContext *context = &getContext();
+    RewritePatternSet patterns(context);
+    populateAffineExpandIndexOpsPatterns(patterns);
+    if (failed(
+            applyPatternsAndFoldGreedily(getOperation(), std::move(patterns))))
+      return signalPassFailure();
+  }
+};
+
+} // namespace
+
+void mlir::populateAffineExpandIndexOpsPatterns(RewritePatternSet &patterns) {
+  patterns.insert<LowerDelinearizeIndexOps>(patterns.getContext());
+}
+
+std::unique_ptr<Pass> mlir::createAffineExpandIndexOpsPass() {
+  return std::make_unique<ExpandAffineIndexOpsPass>();
+}

mlir/lib/Dialect/Affine/Transforms/CMakeLists.txt

 add_mlir_dialect_library(MLIRAffineTransforms
   AffineDataCopyGeneration.cpp
+  AffineExpandIndexOps.cpp
   AffineLoopInvariantCodeMotion.cpp
   AffineLoopNormalize.cpp
   AffineParallelize.cpp
   AffineScalarReplacement.cpp
   LoopCoalescing.cpp
   LoopFusion.cpp
   LoopTiling.cpp
   LoopUnroll.cpp

mlir/lib/Dialect/Affine/Utils/CMakeLists.txt

 add_mlir_dialect_library(MLIRAffineUtils
   LoopFusionUtils.cpp
   LoopUtils.cpp
   Utils.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/Affine
 
   LINK_LIBS PUBLIC
   MLIRAffineDialect
   MLIRAffineAnalysis
   MLIRAnalysis
+  MLIRArithmeticUtils
   MLIRMemRefDialect
   MLIRTransformUtils
   )

mlir/lib/Dialect/Affine/Utils/Utils.cpp

 //===----------------------------------------------------------------------===//
 
 #include "mlir/Dialect/Affine/Utils.h"
 
 #include "mlir/Dialect/Affine/Analysis/Utils.h"
 #include "mlir/Dialect/Affine/IR/AffineOps.h"
 #include "mlir/Dialect/Affine/IR/AffineValueMap.h"
 #include "mlir/Dialect/Affine/LoopUtils.h"
+#include "mlir/Dialect/Arithmetic/Utils/Utils.h"
 #include "mlir/Dialect/Func/IR/FuncOps.h"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
 #include "mlir/IR/AffineExprVisitor.h"
 #include "mlir/IR/BlockAndValueMapping.h"
 #include "mlir/IR/Dominance.h"
 #include "mlir/IR/IntegerSet.h"
 #include "mlir/Transforms/GreedyPatternRewriteDriver.h"
 
MemRefType mlir::normalizeMemRefType(MemRefType memrefType, OpBuilder b,
   // Create the new memref type after trivializing the old layout map.
   MemRefType newMemRefType =
       MemRefType::Builder(memrefType)
           .setShape(newShape)
           .setLayout(AffineMapAttr::get(b.getMultiDimIdentityMap(newRank)));
 
   return newMemRefType;
 }
+
+DivModValue mlir::getDivMod(OpBuilder &b, Location loc, Value lhs, Value rhs) {
+  DivModValue result;
+  AffineExpr d0, d1;
+  bindDims(b.getContext(), d0, d1);
+  result.quotient =
+      makeComposedAffineApply(b, loc, d0.floorDiv(d1), {lhs, rhs});
+  result.remainder = makeComposedAffineApply(b, loc, d0 % d1, {lhs, rhs});
+  return result;
+}
+
+/// Create IR that computes the product of all elements in the set.
+static FailureOr<OpFoldResult> getIndexProduct(OpBuilder &b, Location loc,
+                                               ArrayRef<Value> set) {
+  if (set.empty())
+    return failure();
+  OpFoldResult result = set[0];
+  AffineExpr s0, s1;
+  bindSymbols(b.getContext(), s0, s1);
+  for (unsigned i = 1, e = set.size(); i < e; i++)
+    result = makeComposedFoldedAffineApply(b, loc, s0 * s1, {result, set[i]});
+  return result;
+}
+
+FailureOr<SmallVector<Value>> mlir::delinearizeIndex(OpBuilder &b, Location loc,
+                                                     Value linearIndex,
+                                                     ArrayRef<Value> dimSizes) {
+  unsigned numDims = dimSizes.size();
+
+  SmallVector<Value> divisors;
+  for (unsigned i = 1; i < numDims; i++) {
+    ArrayRef<Value> slice = dimSizes.drop_front(i);
+    FailureOr<OpFoldResult> prod = getIndexProduct(b, loc, slice);
+    if (failed(prod))
+      return failure();
+    divisors.push_back(getValueOrCreateConstantIndexOp(b, loc, *prod));
+  }
+
+  SmallVector<Value> results;
+  results.reserve(divisors.size() + 1);
+  Value residual = linearIndex;
+  for (Value divisor : divisors) {
+    DivModValue divMod = getDivMod(b, loc, residual, divisor);
+    results.push_back(divMod.quotient);
+    residual = divMod.remainder;
+  }
+  results.push_back(residual);
+  return results;
+}

mlir/test/Dialect/Affine/affine-expand-index-ops.mlir

+// RUN: mlir-opt %s -affine-expand-index-ops -split-input-file | FileCheck %s
+
+//   CHECK-DAG:   #[[$map0:.+]] = affine_map<()[s0] -> (s0 floordiv 50176)>
+//   CHECK-DAG:   #[[$map1:.+]] = affine_map<()[s0] -> ((s0 mod 50176) floordiv 224)>
+//   CHECK-DAG:   #[[$map2:.+]] = affine_map<()[s0] -> (s0 mod 224)>
+
+// CHECK-LABEL: @static_basis
+//  CHECK-SAME:    (%[[IDX:.+]]: index)
+//       CHECK:   %[[N:.+]] = affine.apply #[[$map0]]()[%[[IDX]]]
+//       CHECK:   %[[P:.+]] = affine.apply #[[$map1]]()[%[[IDX]]]
+//       CHECK:   %[[Q:.+]] = affine.apply #[[$map2]]()[%[[IDX]]]
+//       CHECK:   return %[[N]], %[[P]], %[[Q]]
+func.func @static_basis(%linear_index: index) -> (index, index, index) {
+  %b0 = arith.constant 16 : index
+  %b1 = arith.constant 224 : index
+  %b2 = arith.constant 224 : index
+  %1:3 = affine.delinearize_index %linear_index into (%b0, %b1, %b2) : index, index, index
+  return %1#0, %1#1, %1#2 : index, index, index
+}
+
+// -----
+
+//   CHECK-DAG:   #[[$map0:.+]] = affine_map<()[s0, s1, s2] -> (s2 floordiv (s0 * s1))>
+//   CHECK-DAG:   #[[$map1:.+]] = affine_map<()[s0, s1, s2] -> ((s2 mod (s0 * s1)) floordiv s1)>
+//   CHECK-DAG:   #[[$map2:.+]] = affine_map<()[s0, s1, s2] -> ((s2 mod (s0 * s1)) mod s1)>
+
+// CHECK-LABEL: @dynamic_basis
+//  CHECK-SAME:    (%[[IDX:.+]]: index, %[[MEMREF:.+]]: memref
+//   CHECK-DAG:   %[[C1:.+]] = arith.constant 1 : index
+//   CHECK-DAG:   %[[C2:.+]] = arith.constant 2 : index
+//        CHECK:  %[[DIM1:.+]] = memref.dim %[[MEMREF]], %[[C1]] :
+//        CHECK:  %[[DIM2:.+]] = memref.dim %[[MEMREF]], %[[C2]] :
+//       CHECK:   %[[N:.+]] = affine.apply #[[$map0]]()[%[[DIM1]], %[[DIM2]], %[[IDX]]]
+//       CHECK:   %[[P:.+]] = affine.apply #[[$map1]]()[%[[DIM1]], %[[DIM2]], %[[IDX]]]
+//       CHECK:   %[[Q:.+]] = affine.apply #[[$map2]]()[%[[DIM1]], %[[DIM2]], %[[IDX]]]
+//       CHECK:   return %[[N]], %[[P]], %[[Q]]
+func.func @dynamic_basis(%linear_index: index, %src: memref<?x?x?xf32>) -> (index, index, index) {
+  %c0 = arith.constant 0 : index
+  %c1 = arith.constant 1 : index
+  %c2 = arith.constant 2 : index
+  %b0 = memref.dim %src, %c0 : memref<?x?x?xf32>
+  %b1 = memref.dim %src, %c1 : memref<?x?x?xf32>
+  %b2 = memref.dim %src, %c2 : memref<?x?x?xf32>
+  %1:3 = affine.delinearize_index %linear_index into (%b0, %b1, %b2) : index, index, index
+  return %1#0, %1#1, %1#2 : index, index, index
+}

mlir/test/Dialect/Affine/invalid.mlir

 // -----
 
 func.func @missing_for_min(%arg0: index, %arg1: index, %arg2: memref<100xf32>) {
   // expected-error @+1 {{upper loop bound affine map with multiple results requires 'min' prefix}}
   affine.for %i0 = %arg0 to affine_map<()[s]->(100,s+1)>()[%arg1] {
   }
   return
 }
+
+// -----
+
+func.func @delinearize(%idx: index, %basis0: index, %basis1 :index) {
+  // expected-error@+1 {{'affine.delinearize_index' op should return an index for each basis element}}
+  %1 = affine.delinearize_index %idx into (%basis0, %basis1) : index
+  return
+}
+
+// -----
+
+func.func @delinearize(%idx: index, %basis0: index, %basis1 :index) {
+  // expected-error@+1 {{'affine.delinearize_index' op basis should not be empty}}
+  affine.delinearize_index %idx into () : index
+  return
+}

mlir/test/Dialect/Affine/ops.mlir

func.func @affine_for_multiple_yield(%buffer: memref<1024xf32>) -> (f32, f32) {
   return %res1, %res2 : f32, f32
 }
 // CHECK:      %[[const_0:.*]] = arith.constant 0.000000e+00 : f32
 // CHECK-NEXT: %[[output:[0-9]+]]:2 = affine.for %{{.*}} = 0 to 10 step 2 iter_args(%[[iter_arg1:.*]] = %[[const_0]], %[[iter_arg2:.*]] = %[[const_0]]) -> (f32, f32) {
 // CHECK:        %[[res1:.*]] = arith.addf %{{.*}}, %[[iter_arg1]] : f32
 // CHECK-NEXT:   %[[res2:.*]] = arith.addf %{{.*}}, %[[iter_arg2]] : f32
 // CHECK-NEXT:   affine.yield %[[res1]], %[[res2]] : f32, f32
 // CHECK-NEXT: }
+
+// -----
+
+// CHECK-LABEL: func @delinearize
+func.func @delinearize(%linear_idx: index, %basis0: index, %basis1 :index) -> (index, index) {
+  // CHECK: affine.delinearize_index %{{.+}} into (%{{.+}}, %{{.+}}) : index, index
+  %1:2 = affine.delinearize_index %linear_idx into (%basis0, %basis1) : index, index
+  return %1#0, %1#1 : index, index
+}