[mlir][Arithmetic] Add `arith.delinearize_index` operation

[christopherbate]

Summary

This change adds a new DelinearizeIndexOp to the arith 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 arith.divui and arith.remui operations is also given.

[nicolasvasilache]

Generally looks good and is an important step forward.

I would either:

• move this op to arith and keep the operands as they are, or
• keep it in tensor but then take an actual tensor as the basis and abstract away the fact that you get the Dims as a lowering detail.

Unless you have a stronger argument for keeping the semantics as is that I don't see yet?

Rest of the impl LGTM so I'll approve and let you decide which way to go (or to justify the current form with an extra argument).

mlir/test/Dialect/Tensor/ops.mlir

270

nit: nl here and below

[nicolasvasilache]

mlir/test/Dialect/Tensor/ops.mlir

270

and above especially :p

[christopherbate]

In D129697#3662671, @nicolasvasilache wrote:

Generally looks good and is an important step forward.

I would either:

• move this op to arith and keep the operands as they are, or
• keep it in tensor but then take an actual tensor as the basis and abstract away the fact that you get the Dims as a lowering detail.

Unless you have a stronger argument for keeping the semantics as is that I don't see yet?

Rest of the impl LGTM so I'll approve and let you decide which way to go (or to justify the current form with an extra argument).

Ok, I think moving to Arithmetic is probably the best way forward right now.

[christopherbate]

Move to Arithmetic dialect.

Diff is simplified because the lowering transform is moved to the existing -arith-expand-ops pass.

[nicolasvasilache]

mlir/test/Dialect/Arithmetic/expand-ops.mlir

254 ↗

(On Diff #446256)

nit: nl here and above.

[jpienaar]

mlir/lib/Dialect/Arithmetic/Utils/Utils.cpp

131 ↗

(On Diff #446256)

When does this happen and should the op be considered valid/pass verification where this is true?

[christopherbate]

mlir/lib/Dialect/Arithmetic/Utils/Utils.cpp

131 ↗

(On Diff #446256)

Is it possible to add a "non empty" constraint for the Variadic<Index> parameter via the Tablegen spec? If not I'll add a verifier

[christopherbate]

Add missing end-of-file newlines.
Add a verifier for the op + invalid tests.

[christopherbate]

Fix ArithmeticTransforms missing link dependency.

[chongxing]

mlir/include/mlir/Dialect/Arithmetic/IR/ArithmeticOps.td

1242 ↗

(On Diff #446895)

is this small typo? -> %v1 = arith.muli %2, %3 : index

[rriddle]

This revision is really concerning to me because this operation is not close to any of the others in this dialect in terms of abstraction level.

mlir/include/mlir/Dialect/Arithmetic/IR/ArithmeticOps.td

1226 ↗

(On Diff #446895)

This should be using ArithOp, also keep the operations in this file sorted as much as possible.

1230–1231 ↗

(On Diff #446895)

Adding this operation to arith dialect is honestly significantky concerning to me, given that representationally this feels like a completely different level of abstraction compared to the rest of the dialect.

[rriddle]

Can you revert this and raise an RFC that adds additional context and justification? The direction this takes the arith dialect isn't clear to me.

[christopherbate]

In D129697#3674751, @rriddle wrote:

Can you revert this and raise an RFC that adds additional context and justification? The direction this takes the arith dialect isn't clear to me.

I'm about to land a diff that depends on this: D129699 That diff is the main context.

Happy to move to another dialect that makes sense. Originally I created this diff to add the op to tensor, but was asked by @nicolasvasilache to move it to arith or change the spec:

In D129697#3662671, @nicolasvasilache wrote:

Generally looks good and is an important step forward.

I would either:

• move this op to arith and keep the operands as they are, or
• keep it in tensor but then take an actual tensor as the basis and abstract away the fact that you get the Dims as a lowering detail.

Unless you have a stronger argument for keeping the semantics as is that I don't see yet?

Rest of the impl LGTM so I'll approve and let you decide which way to go (or to justify the current form with an extra argument).

Is your request to raise an RFC just about this op's positioning in the arith dialect? If so, I can revert and do option number 2 above to put the op back in Tensor, then no RFC needed?

[rriddle]

In D129697#3674796, @christopherbate wrote:

In D129697#3674751, @rriddle wrote:

Can you revert this and raise an RFC that adds additional context and justification? The direction this takes the arith dialect isn't clear to me.

I'm about to land a diff that depends on this: D129699 That diff is the main context.

Happy to move to another dialect that makes sense. Originally I created this diff to add the op to tensor, but was asked by @nicolasvasilache to move it to arith or change the spec:

In D129697#3662671, @nicolasvasilache wrote:

Generally looks good and is an important step forward.

I would either:

• move this op to arith and keep the operands as they are, or
• keep it in tensor but then take an actual tensor as the basis and abstract away the fact that you get the Dims as a lowering detail.

Unless you have a stronger argument for keeping the semantics as is that I don't see yet?

Rest of the impl LGTM so I'll approve and let you decide which way to go (or to justify the current form with an extra argument).

Is your request to raise an RFC just about this op's positioning in the arith dialect? If so, I can revert and do option number 2 above to put the op back in Tensor, then no RFC needed?

The op from an abstraction point-of-view would make more sense in tensor, given that (at least at-present) the operation fundamentally breaks away from the purpose of the arith dialect (as defined e.g. by it's docs).
Going beyond that though, there is still a huge amount of context missing here and the documentation for the operation is quite lacking. Neither the operation documentation or the commit description really dig into the "why" behind this commit.

[christopherbate]

In D129697#3674826, @rriddle wrote:

In D129697#3674796, @christopherbate wrote:

In D129697#3674751, @rriddle wrote:

Can you revert this and raise an RFC that adds additional context and justification? The direction this takes the arith dialect isn't clear to me.

I'm about to land a diff that depends on this: D129699 That diff is the main context.

Happy to move to another dialect that makes sense. Originally I created this diff to add the op to tensor, but was asked by @nicolasvasilache to move it to arith or change the spec:

In D129697#3662671, @nicolasvasilache wrote:

Generally looks good and is an important step forward.

I would either:

• move this op to arith and keep the operands as they are, or
• keep it in tensor but then take an actual tensor as the basis and abstract away the fact that you get the Dims as a lowering detail.

Unless you have a stronger argument for keeping the semantics as is that I don't see yet?

Rest of the impl LGTM so I'll approve and let you decide which way to go (or to justify the current form with an extra argument).

Is your request to raise an RFC just about this op's positioning in the arith dialect? If so, I can revert and do option number 2 above to put the op back in Tensor, then no RFC needed?

The op from an abstraction point-of-view would make more sense in tensor, given that (at least at-present) the operation fundamentally breaks away from the purpose of the arith dialect (as defined e.g. by it's docs).
Going beyond that though, there is still a huge amount of context missing here and the documentation for the operation is quite lacking. Neither the operation documentation or the commit description really dig into the "why" behind this commit.

Very well, I will revert the commit and post an RFC as requested.

mlir/include/mlir/Dialect/Tensor/IR/TensorOps.td

def Tensor_YieldOp : Tensor_Op<"yield",
   let arguments = (ins AnyType:$value);
   let assemblyFormat = "$value attr-dict `:` type($value)";
 
   // Dummy builder to appease code in templated ensureTerminator that
   // GenerateOp's auto-generated parser calls.
   let builders = [OpBuilder<(ins), [{ /* nothing to do */ }]>];
 }
 
+//===----------------------------------------------------------------------===//
+// DelinearizeIndexOp
+//===----------------------------------------------------------------------===//
+
+def DelinearizeIndexOp : Tensor_Op<"delinearize_index", [NoSideEffect,
+  RangedTypesMatchWith<"result type matches operand", "basis", "multi_index",
+                         "llvm::make_range($_self.begin(), $_self.end())">
+]> {
+  let summary = "delinearize an index";
+  let description = [{
+    The `tensor.delinearize_index` operation takes a single index value and 
+    calculates the multi-index according to the given basis.
+
+    Example:
+    ```mlir
+    %indices:3 = tensor.delinearize_index %linear_index (3, 4, 5)
+    ```
+
+    In the above example, `%indices:3` conceptually holds the following:
+    ```
+    %indices#0 = floorDiv(%linear_index , 20)
+    %indices#1 = floorDiv(remander(%linear_index , 20), 5)
+    %indices#2 = remainder(remainder(%linear_idnex, 20), 5)
+    ```
+  }];
+
+  let arguments = (ins Index:$linear_index, Variadic<Index>:$basis);
+  let results = (outs Variadic<Index>:$multi_index);
+
+  let assemblyFormat = [{
+    $linear_index `(` $basis `:` type($basis)`)` attr-dict 
+  }];
+
+  let builders = [
+    OpBuilder<(ins "Value":$linear_index, "ArrayRef<OpFoldResult>":$basis)>
+  ];
+}
+
 #endif // TENSOR_OPS

mlir/include/mlir/Dialect/Tensor/Transforms/Passes.h

 #define MLIR_DIALECT_TENSOR_TRANSFORMS_PASSES_H_
 
 #include "mlir/Pass/Pass.h"
 
 namespace mlir {
 /// Creates an instance of `tensor` dialect bufferization pass.
 std::unique_ptr<Pass> createTensorBufferizePass();
 
+/// Create a TensorLowerDelinearizeIndexPass.
+std::unique_ptr<Pass> createTensorLowerDelinearizeIndexPass();
+
 //===----------------------------------------------------------------------===//
 // Registration
 //===----------------------------------------------------------------------===//
 
 namespace tensor {
 /// Generate the code for registering passes.
 #define GEN_PASS_REGISTRATION
 #include "mlir/Dialect/Tensor/Transforms/Passes.h.inc"
 } // namespace tensor
 
 } // namespace mlir
 
 #endif // MLIR_DIALECT_TENSOR_TRANSFORMS_PASSES_H_

mlir/include/mlir/Dialect/Tensor/Transforms/Passes.td

 
 include "mlir/Pass/PassBase.td"
 
 def TensorBufferize : Pass<"tensor-bufferize", "func::FuncOp"> {
   let summary = "Bufferize the `tensor` dialect";
   let constructor = "mlir::createTensorBufferizePass()";
 }
 
+def TensorLowerDelinearizeIndexPass
+    : Pass<"tensor-lower-delinearize-index", ""> {
+  let summary = "Lower delinearize indices.";
+  let constructor = "mlir::createTensorLowerDelinearizeIndexPass()";
+  let dependentDialects = ["::mlir::arith::ArithmeticDialect"];
+  let options = [];
+}
+
 #endif // MLIR_DIALECT_TENSOR_TRANSFORMS_PASSES

mlir/include/mlir/Dialect/Tensor/Transforms/Transforms.h

 /// Pattern to swap an `tensor.extract_slice` with its producer when the
 /// producer implements the `TilingInterface`. The pattern itself does not
 /// provide a mechanism to control where the application happens. With use of
 /// transform dialect that control is done within the transform dialect. Other
 /// use cases can inherit from this pattern and add necessary controls.
 FailureOr<Value> replaceExtractSliceWithTiledProducer(
     OpBuilder &builder, tensor::ExtractSliceOp sliceOp, OpResult producerOp);
 
+/// Populates rewrite patterns that lower `tensor.delinearize_index`.
+void populateLowerDelinearizeIndexPatterns(RewritePatternSet &patterns);
+
 } // namespace tensor
 } // namespace mlir
 
 #endif // MLIR_DIALECT_TENSOR_TRANSFORMS_TRANSFORMS_H

mlir/include/mlir/Dialect/Tensor/Utils/Utils.h

PadOp createPadScalarOp(Type type, Value source, Value pad,
                         ArrayRef<OpFoldResult> low, ArrayRef<OpFoldResult> high,
                         bool nofold, Location loc, OpBuilder &builder);
 
 // Creates dim ops for each dynamic dimension of the raked tensor argument and
 // returns these as values.
 SmallVector<Value> createDynamicDimValues(OpBuilder &b, Location loc,
                                           Value rankedTensor);
 
+/// Holds the result of (div a, b)  and (mod a, b)
+struct DivModValue {
+  Value quotient;
+  Value remainder;
+};
+
+/// Create IR to calculate (div a, b)  and (mod a, b)
+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> dimSizes);
+
 } // namespace tensor
 } // namespace mlir
 
 #endif // MLIR_DIALECT_TENSOR_UTILS_UTILS_H_

mlir/lib/Dialect/Linalg/Transforms/Transforms.cpp

 //===- Transforms.cpp - Linalg transformations as patterns ----------------===//
 //
 // 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 logic and helpers to expose Linalg transforms as rewrite
 // patterns.
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Dialect/Linalg/Transforms/Transforms.h"
+#include "PassDetail.h"
 #include "mlir/Dialect/Affine/IR/AffineOps.h"
 #include "mlir/Dialect/Arithmetic/IR/Arithmetic.h"
 #include "mlir/Dialect/Func/IR/FuncOps.h"
 #include "mlir/Dialect/Linalg/Analysis/DependenceAnalysis.h"
 #include "mlir/Dialect/Linalg/IR/Linalg.h"
+#include "mlir/Dialect/Linalg/Passes.h"
 #include "mlir/Dialect/Linalg/Transforms/HoistPadding.h"
 #include "mlir/Dialect/Linalg/Utils/Utils.h"
 #include "mlir/Dialect/SCF/Transforms/Transforms.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
 #include "mlir/Dialect/Tensor/IR/TensorTilingInterfaceImpl.h"
 #include "mlir/Dialect/Utils/StaticValueUtils.h"
 #include "mlir/Dialect/Utils/StructuredOpsUtils.h"
 #include "mlir/Dialect/Vector/IR/VectorOps.h"
 #include "mlir/IR/AffineExpr.h"
 #include "mlir/IR/Matchers.h"
+#include "mlir/IR/PatternMatch.h"
 #include "mlir/Pass/Pass.h"
 #include "mlir/Support/LLVM.h"
 #include "mlir/Transforms/GreedyPatternRewriteDriver.h"
 #include "llvm/ADT/ScopeExit.h"
 #include "llvm/ADT/TypeSwitch.h"
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/raw_ostream.h"
+#include <memory>
 #include <type_traits>
 #include <utility>
 
 #define DEBUG_TYPE "linalg-transforms"
 
 using namespace mlir;
 using namespace mlir::linalg;
 

mlir/lib/Dialect/Tensor/IR/TensorOps.cpp

OpFoldResult SplatOp::fold(ArrayRef<Attribute> operands) {
   if (!constOperand.isa_and_nonnull<IntegerAttr, FloatAttr>())
     return {};
 
   // SplatElementsAttr::get treats single value for second arg as being a splat.
   return SplatElementsAttr::get(getType(), {constOperand});
 }
 
 //===----------------------------------------------------------------------===//
+// DelinarizeIndexOp
+//===----------------------------------------------------------------------===//
+void DelinearizeIndexOp::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.hasValue())
+          return builder.create<arith::ConstantIndexOp>(result.location,
+                                                        *staticDim);
+        return ofr.dyn_cast<Value>();
+      }));
+  result.addOperands(basisValues);
+}
+
+//===----------------------------------------------------------------------===//
 // TableGen'd op method definitions
 //===----------------------------------------------------------------------===//
 
 #define GET_OP_CLASSES
 #include "mlir/Dialect/Tensor/IR/TensorOps.cpp.inc"

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

 add_mlir_dialect_library(MLIRTensorTransforms
   BufferizableOpInterfaceImpl.cpp
   Bufferize.cpp
+  DelinearizeIndex.cpp
   SplitPadding.cpp
   SwapExtractSliceWithProducer.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/Tensor/Transforms
 
   DEPENDS
   MLIRTensorTransformsIncGen

mlir/lib/Dialect/Tensor/Transforms/DelinearizeIndex.cpp

+//===- DelinearizeIndex.cpp - Implementation of delinearize_index lowering-===//
+//
+// 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 transformations and lowerings for the
+// `tensor.delinearize_index` operation.
+//
+//===----------------------------------------------------------------------===//
+#include "PassDetail.h"
+#include "mlir/Dialect/Arithmetic/IR/Arithmetic.h"
+#include "mlir/Dialect/Tensor/IR/Tensor.h"
+#include "mlir/Dialect/Tensor/Transforms/Passes.h"
+#include "mlir/Dialect/Tensor/Transforms/Transforms.h"
+#include "mlir/Dialect/Tensor/Utils/Utils.h"
+#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
+
+using namespace mlir;
+using namespace mlir::tensor;
+
+namespace {
+struct LowerDelinearizeIndexOps : public OpRewritePattern<DelinearizeIndexOp> {
+  using OpRewritePattern<DelinearizeIndexOp>::OpRewritePattern;
+  LogicalResult matchAndRewrite(DelinearizeIndexOp op,
+                                PatternRewriter &rewriter) const override {
+    FailureOr<SmallVector<Value>> multiIndex = delinearizeIndex(
+        rewriter, op->getLoc(), op.linear_index(), llvm::to_vector(op.basis()));
+    if (failed(multiIndex))
+      return failure();
+    rewriter.replaceOp(op, *multiIndex);
+    return success();
+  }
+};
+
+class LowerDelinearizeIndexPass
+    : public TensorLowerDelinearizeIndexPassBase<LowerDelinearizeIndexPass> {
+public:
+  LowerDelinearizeIndexPass() = default;
+  void runOnOperation() override {
+    MLIRContext *ctx = &getContext();
+    RewritePatternSet patterns(ctx);
+    patterns.add<LowerDelinearizeIndexOps>(ctx);
+    if (failed(
+            applyPatternsAndFoldGreedily(getOperation(), std::move(patterns))))
+      return signalPassFailure();
+  }
+};
+} // namespace
+
+void tensor::populateLowerDelinearizeIndexPatterns(
+    RewritePatternSet &patterns) {
+  patterns.add<LowerDelinearizeIndexOps>(patterns.getContext());
+}
+
+std::unique_ptr<Pass> mlir::createTensorLowerDelinearizeIndexPass() {
+  return std::make_unique<LowerDelinearizeIndexPass>();
+}
 No newline at end of file

mlir/lib/Dialect/Tensor/Transforms/PassDetail.h

 namespace memref {
 class MemRefDialect;
 } // namespace memref
 
 namespace scf {
 class SCFDialect;
 } // namespace scf
 
+namespace arith {
+class ArithmeticDialect;
+}
+
 #define GEN_PASS_CLASSES
 #include "mlir/Dialect/Tensor/Transforms/Passes.h.inc"
 
 } // namespace mlir
 
 #endif // DIALECT_TENSOR_TRANSFORMS_PASSDETAIL_H_

mlir/lib/Dialect/Tensor/Utils/Utils.cpp

SmallVector<Value> mlir::tensor::createDynamicDimValues(OpBuilder &b,
   SmallVector<Value> dynamicDims;
   for (const auto &en : llvm::enumerate(tensorTy.getShape())) {
     if (en.value() == ShapedType::kDynamicSize)
       dynamicDims.push_back(
           b.create<tensor::DimOp>(loc, rankedTensor, en.index()));
   }
   return dynamicDims;
 }
+
+DivModValue mlir::tensor::getDivMod(OpBuilder &b, Location loc, Value lhs,
+                                    Value rhs) {
+  DivModValue result;
+  result.quotient = b.create<arith::DivUIOp>(loc, lhs, rhs);
+  result.remainder = b.create<arith::RemUIOp>(loc, lhs, rhs);
+  return result;
+}
+
+/// Create IR that computes the product of all elements in the set.
+static FailureOr<Value> getIndexProduct(OpBuilder &b, Location loc,
+                                        ArrayRef<Value> set) {
+  if (set.empty())
+    return failure();
+  Value result = set[0];
+  for (unsigned i = 1; i < set.size(); i++)
+    result = b.create<arith::MulIOp>(loc, result, set[i]);
+  return result;
+}
+
+FailureOr<SmallVector<Value>>
+mlir::tensor::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.begin() + i, dimSizes.end());
+    FailureOr<Value> prod = getIndexProduct(b, loc, slice);
+    if (failed(prod))
+      return failure();
+    divisors.push_back(*prod);
+  }
+
+  SmallVector<Value> results;
+  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;
+}
 No newline at end of file

mlir/test/Dialect/Tensor/delinearize-index.mlir

+// RUN: mlir-opt -split-input-file -tensor-lower-delinearize-index %s | FileCheck %s
+
+// CHECK: @static_basis(%[[IDX:.+]]: index)
+// CHECK-DAG: %[[c50176:.+]] = arith.constant 50176 : index
+// CHECK-DAG: %[[c224:.+]] = arith.constant 224 : index
+// CHECK: %[[N:.+]] = arith.divui %[[IDX]], %[[c50176]] : index
+// CHECK: %[[RES:.+]] = arith.remui %[[IDX]], %[[c50176]] : index
+// CHECK: %[[P:.+]] = arith.divui %[[RES]], %[[c224]] : index
+// CHECK: %[[Q:.+]] = arith.remui %[[RES]], %[[c224]] : index
+// 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 = tensor.delinearize_index %linear_index (%b0, %b1, %b2 : index, index, index)
+  return %1#0, %1#1, %1#2 : index, index, index
+}
 No newline at end of file

mlir/test/Dialect/Tensor/ops.mlir

 func.func @test_splat_op(%s : f32) {
   // CHECK: tensor.splat [[S]] : tensor<8xf32>
   %v = tensor.splat %s : tensor<8xf32>
   
   // CHECK: tensor.splat [[S]] : tensor<4xf32>
   %u = "tensor.splat"(%s) : (f32) -> tensor<4xf32>
   return
 }
+
+// -----
+
+func.func @delinearize(%idx: index, %basis0: index, %basis1 :index) -> (index, index) {
+  %1:2 = tensor.delinearize_index %idx (%basis0, %basis1 : index, index)
+  return %1#0, %1#1 : index, index
+}
 No newline at end of file