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
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
#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"
}
//===----------------------------------------------------------------------===//
-// 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.has_value())
- 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
//===----------------------------------------------------------------------===//
LINK_LIBS PUBLIC
MLIRDialect
- MLIRDialectUtils
MLIRInferIntRangeInterface
MLIRInferTypeOpInterface
MLIRIR
LINK_LIBS PUBLIC
MLIRAnalysis
MLIRArithmeticDialect
- MLIRArithmeticUtils
MLIRBufferizationDialect
MLIRBufferizationTransforms
MLIRInferIntRangeInterface
#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"
}
};
-/// 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 {
arith::MaxUIOp,
arith::MinFOp,
arith::MinSIOp,
- arith::MinUIOp,
- arith::DelinearizeIndexOp
+ arith::MinUIOp
>();
// clang-format on
if (failed(applyPartialConversion(getOperation(), target,
MaxMinIOpConverter<MaxSIOp, arith::CmpIPredicate::sgt>,
MaxMinIOpConverter<MaxUIOp, arith::CmpIPredicate::ugt>,
MaxMinIOpConverter<MinSIOp, arith::CmpIPredicate::slt>,
- MaxMinIOpConverter<MinUIOp, arith::CmpIPredicate::ult>,
- LowerDelinearizeIndexOps
+ MaxMinIOpConverter<MinUIOp, arith::CmpIPredicate::ult>
>(patterns.getContext());
// clang-format on
}
#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;
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;
-}
}
// 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
-}
%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
-}
%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
-}