[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 ↗

(On Diff #444426)

nit: nl here and below

[nicolasvasilache]

mlir/test/Dialect/Tensor/ops.mlir

270 ↗

(On Diff #444426)

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

nit: nl here and above.

[jpienaar]

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

131

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

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

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

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

1230–1231

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/Arithmetic/IR/ArithmeticOps.td

def SelectOp : Arith_Op<"select", [
   let hasCanonicalizer = 1;
   let hasFolder = 1;
   let hasVerifier = 1;
 
   // FIXME: Switch this to use the declarative assembly format.
   let hasCustomAssemblyFormat = 1;
 }
 
+//===----------------------------------------------------------------------===//
+// DelinearizeIndexOp
+//===----------------------------------------------------------------------===//
+
+def DelinearizeIndexOp : Op<Arithmetic_Dialect, "delinearize_index",
+    [NoSideEffect]> {
+  let summary = "delinearize an index";
+  let description = [{
+    The `arith.delinearize_index` operation takes a single index value and
+    calculates the multi-index according to the given basis.
+
+    Example:
+
+    ```
+    %indices:3 = arith.delinearize_index %linear_index (%1, %2, %3) : index, index, index
+    ```
+
+    In the above example, `%indices:3` conceptually holds the following:
+
+    ```
+    %v1 = arith.muli %1, %2 : index
+    %indices#0 = floorDiv(%linear_index , %v1)
+    %indices#1 = floorDiv(remander(%linear_index , %v1), %3)
+    %indices#2 = remainder(remainder(%linear_idnex, %v1), %3)
+    ```
+  }];
+
+  let arguments = (ins Index:$linear_index, Variadic<Index>:$basis);
+  let results = (outs Variadic<Index>:$multi_index);
+
+  let assemblyFormat = [{
+    $linear_index `(` $basis `)` attr-dict `:` type($multi_index)
+  }];
+
+  let builders = [
+    OpBuilder<(ins "Value":$linear_index, "ArrayRef<OpFoldResult>":$basis)>
+  ];
+
+  let hasVerifier = 1;
+}
+
+
+
 #endif // ARITHMETIC_OPS

mlir/include/mlir/Dialect/Arithmetic/Utils/Utils.h

struct ArithBuilder {
   Value select(Value cmp, Value lhs, Value rhs);
   Value sgt(Value lhs, Value rhs);
   Value slt(Value lhs, Value rhs);
 
 private:
   OpBuilder &b;
   Location loc;
 };
+
+/// 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 mlir
 
 #endif // MLIR_DIALECT_ARITHMETIC_UTILS_UTILS_H

mlir/lib/Dialect/Arithmetic/IR/ArithmeticOps.cpp

 //===- ArithmeticOps.cpp - MLIR Arithmetic dialect ops implementation -----===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 
 #include <utility>
 
 #include "mlir/Dialect/Arithmetic/IR/Arithmetic.h"
 #include "mlir/Dialect/CommonFolders.h"
+#include "mlir/Dialect/Utils/StaticValueUtils.h"
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/Matchers.h"
 #include "mlir/IR/OpImplementation.h"
 #include "mlir/IR/PatternMatch.h"
 #include "mlir/IR/TypeUtilities.h"
 #include "llvm/ADT/SmallString.h"
 
 #include "llvm/ADT/APSInt.h"
Value mlir::arith::getReductionOp(AtomicRMWKind op, OpBuilder &builder,
   default:
     (void)emitOptionalError(loc, "Reduction operation type not supported");
     break;
   }
   return nullptr;
 }
 
 //===----------------------------------------------------------------------===//
+// DelinearizeIndexOp
+//===----------------------------------------------------------------------===//
+
+void arith::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);
+}
+
+LogicalResult arith::DelinearizeIndexOp::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/Arithmetic/IR/ArithmeticOps.cpp.inc"
 
 //===----------------------------------------------------------------------===//
 // TableGen'd enum attribute definitions
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Dialect/Arithmetic/IR/ArithmeticOpsEnums.cpp.inc"

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

add_mlir_dialect_library(MLIRArithmeticDialect
   ADDITIONAL_HEADER_DIRS
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/Arithmetic
 
   DEPENDS
   MLIRArithmeticOpsIncGen
 
   LINK_LIBS PUBLIC
   MLIRDialect
+  MLIRDialectUtils
   MLIRInferIntRangeInterface
   MLIRInferTypeOpInterface
   MLIRIR
   )

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

 add_mlir_dialect_library(MLIRArithmeticTransforms
   BufferizableOpInterfaceImpl.cpp
   Bufferize.cpp
   ExpandOps.cpp
   UnsignedWhenEquivalent.cpp
 
   ADDITIONAL_HEADER_DIRS
   {$MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/Arithmetic/Transforms
 
   DEPENDS
   MLIRArithmeticTransformsIncGen
 
   LINK_LIBS PUBLIC
   MLIRAnalysis
   MLIRArithmeticDialect
+  MLIRArithmeticUtils
   MLIRBufferizationDialect
   MLIRBufferizationTransforms
   MLIRInferIntRangeInterface
   MLIRIR
   MLIRMemRefDialect
   MLIRPass
   MLIRTransforms
   )

mlir/lib/Dialect/Arithmetic/Transforms/ExpandOps.cpp

 //===- ExpandOps.cpp - Pass to legalize Arithmetic ops for LLVM 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
 //
 //===----------------------------------------------------------------------===//
 
 #include "PassDetail.h"
 #include "mlir/Dialect/Arithmetic/IR/Arithmetic.h"
 #include "mlir/Dialect/Arithmetic/Transforms/Passes.h"
+#include "mlir/Dialect/Arithmetic/Utils/Utils.h"
 #include "mlir/IR/TypeUtilities.h"
 #include "mlir/Transforms/DialectConversion.h"
 
 using namespace mlir;
 
 /// Create an integer or index constant.
 static Value createConst(Location loc, Type type, int value,
                          PatternRewriter &rewriter) {
LogicalResult matchAndRewrite(OpTy op,
 
     Location loc = op.getLoc();
     Value cmp = rewriter.create<arith::CmpIOp>(loc, pred, lhs, rhs);
     rewriter.replaceOpWithNewOp<arith::SelectOp>(op, cmp, lhs, rhs);
     return success();
   }
 };
 
+/// Lowers `arith.delinearize_index` into a sequence of division and remainder
+/// operations.
+struct LowerDelinearizeIndexOps
+    : public OpRewritePattern<arith::DelinearizeIndexOp> {
+  using OpRewritePattern<arith::DelinearizeIndexOp>::OpRewritePattern;
+  LogicalResult matchAndRewrite(arith::DelinearizeIndexOp 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();
+  }
+};
+
 struct ArithmeticExpandOpsPass
     : public ArithmeticExpandOpsBase<ArithmeticExpandOpsPass> {
   void runOnOperation() override {
     RewritePatternSet patterns(&getContext());
     ConversionTarget target(getContext());
 
     arith::populateArithmeticExpandOpsPatterns(patterns);
 
     target.addLegalDialect<arith::ArithmeticDialect>();
     // clang-format off
     target.addIllegalOp<
       arith::CeilDivSIOp,
       arith::CeilDivUIOp,
       arith::FloorDivSIOp,
       arith::MaxFOp,
       arith::MaxSIOp,
       arith::MaxUIOp,
       arith::MinFOp,
       arith::MinSIOp,
-      arith::MinUIOp
+      arith::MinUIOp,
+      arith::DelinearizeIndexOp
     >();
     // clang-format on
     if (failed(applyPartialConversion(getOperation(), target,
                                       std::move(patterns))))
       signalPassFailure();
   }
 };
 
 } // namespace
 
 void mlir::arith::populateArithmeticExpandOpsPatterns(
     RewritePatternSet &patterns) {
   // clang-format off
   patterns.add<
     CeilDivSIOpConverter,
     CeilDivUIOpConverter,
     FloorDivSIOpConverter,
     MaxMinFOpConverter<MaxFOp, arith::CmpFPredicate::UGT>,
     MaxMinFOpConverter<MinFOp, arith::CmpFPredicate::ULT>,
     MaxMinIOpConverter<MaxSIOp, arith::CmpIPredicate::sgt>,
     MaxMinIOpConverter<MaxUIOp, arith::CmpIPredicate::ugt>,
     MaxMinIOpConverter<MinSIOp, arith::CmpIPredicate::slt>,
-    MaxMinIOpConverter<MinUIOp, arith::CmpIPredicate::ult>
+    MaxMinIOpConverter<MinUIOp, arith::CmpIPredicate::ult>,
+    LowerDelinearizeIndexOps
    >(patterns.getContext());
   // clang-format on
 }
 
 std::unique_ptr<Pass> mlir::arith::createArithmeticExpandOpsPass() {
   return std::make_unique<ArithmeticExpandOpsPass>();
 }

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

 //===- Utils.cpp - Utilities to support the Linalg dialect ----------------===//
 //
 // 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 utilities for the Linalg dialect.
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Dialect/Arithmetic/Utils/Utils.h"
 #include "mlir/Dialect/Arithmetic/IR/Arithmetic.h"
+#include "mlir/IR/OpDefinition.h"
 #include "llvm/ADT/SmallBitVector.h"
 
 using namespace mlir;
 
 /// Matches a ConstantIndexOp.
 /// TODO: This should probably just be a general matcher that uses matchConstant
 /// and checks the operation for an index type.
 detail::op_matcher<arith::ConstantIndexOp> mlir::matchConstantIndex() {

 Value ArithBuilder::slt(Value lhs, Value rhs) {
   if (lhs.getType().isa<IndexType, IntegerType>())
     return b.create<arith::CmpIOp>(loc, arith::CmpIPredicate::slt, lhs, rhs);
   return b.create<arith::CmpFOp>(loc, arith::CmpFPredicate::OLT, lhs, rhs);
 }
 Value ArithBuilder::select(Value cmp, Value lhs, Value rhs) {
   return b.create<arith::SelectOp>(loc, cmp, lhs, rhs);
 }
+
+DivModValue mlir::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<OpFoldResult> getIndexProduct(OpBuilder &b, Location loc,
+                                               ArrayRef<Value> set) {
+  if (set.empty())
+    return failure();
+  OpFoldResult result = set[0];
+  for (unsigned i = 1; i < set.size(); i++)
+    result = b.createOrFold<arith::MulIOp>(
+        loc, getValueOrCreateConstantIndexOp(b, loc, 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.begin() + i, dimSizes.end());
+    FailureOr<OpFoldResult> prod = getIndexProduct(b, loc, slice);
+    if (failed(prod))
+      return failure();
+    divisors.push_back(getValueOrCreateConstantIndexOp(b, loc, *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;
+}

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

 
 // CHECK-LABEL: func @minui
 func.func @minui(%a: i32, %b: i32) -> i32 {
   %result = arith.minui %a, %b : i32
   return %result : i32
 }
 // CHECK-SAME: %[[LHS:.*]]: i32, %[[RHS:.*]]: i32)
 // CHECK-NEXT: %[[CMP:.*]] = arith.cmpi ult, %[[LHS]], %[[RHS]] : i32
+
+// -----
+
+// CHECK-LABEL: @static_basis
+//  CHECK-SAME:    (%[[IDX:.+]]: index)
+//       CHECK:   arith.constant
+//       CHECK:   arith.constant
+//   CHECK-DAG:   %[[c224:.+]] = arith.constant 224 : index
+//   CHECK-DAG:   %[[c50176:.+]] = arith.constant 50176 : 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 = arith.delinearize_index %linear_index (%b0, %b1, %b2) : index, index, index
+  return %1#0, %1#1, %1#2 : index, index, index
+}

mlir/test/Dialect/Arithmetic/invalid.mlir

 
 // -----
 
 func.func @func() {
   %c0 = arith.constant  // expected-error {{expected attribute value}}
 
   %x = arith.constant 1 : i32
 }
+
+// -----
+
+func.func @delinearize(%idx: index, %basis0: index, %basis1 :index) {
+  // expected-error@+1 {{'arith.delinearize_index' op should return an index for each basis element}}
+  %1 = arith.delinearize_index %idx (%basis0, %basis1) : index
+  return
+}
+
+// -----
+
+func.func @delinearize(%idx: index, %basis0: index, %basis1 :index) {
+  // expected-error@+1 {{'arith.delinearize_index' op basis should not be empty}}
+  arith.delinearize_index %idx () : index
+  return
+}

mlir/test/Dialect/Arithmetic/ops.mlir

func.func @minimum(%v1: vector<4xf32>, %v2: vector<4xf32>,
                %i1: i32, %i2: i32) {
   %min_vector = arith.minf %v1, %v2 : vector<4xf32>
   %min_scalable_vector = arith.minf %sv1, %sv2 : vector<[4]xf32>
   %min_float = arith.minf %f1, %f2 : f32
   %min_signed = arith.minsi %i1, %i2 : i32
   %min_unsigned = arith.minui %i1, %i2 : i32
   return
 }
+
+// CHECK-LABEL: func @delinearize
+func.func @delinearize(%idx: index, %basis0: index, %basis1 :index) -> (index, index) {
+  %1:2 = arith.delinearize_index %idx (%basis0, %basis1) : index, index
+  return %1#0, %1#1 : index, index
+}