diff --git a/mlir/include/mlir/EDSC/Helpers.h b/mlir/include/mlir/EDSC/Helpers.h index 5974439a2660..d653c7f4d359 100644 --- a/mlir/include/mlir/EDSC/Helpers.h +++ b/mlir/include/mlir/EDSC/Helpers.h @@ -62,6 +62,9 @@ public: unsigned rank() const { return lbs.size(); } unsigned fastestVarying() const { return rank() - 1; } + IndexHandle lb(unsigned idx) { return lbs[idx]; } + IndexHandle ub(unsigned idx) { return ubs[idx]; } + int64_t step(unsigned idx) { return steps[idx]; } std::tuple range(unsigned idx) { return std::make_tuple(lbs[idx], ubs[idx], steps[idx]); } @@ -93,15 +96,18 @@ ValueHandle operator/(ValueHandle v, IndexedValue i); /// Assigning to an IndexedValue emits an actual store operation, while using /// converting an IndexedValue to a ValueHandle emits an actual load operation. struct IndexedValue { - explicit IndexedValue(MemRefView &v, llvm::ArrayRef indices = {}) - : view(v), indices(indices.begin(), indices.end()) {} + explicit IndexedValue(Type t) : base(t) {} + explicit IndexedValue(Value *v, llvm::ArrayRef indices = {}) + : IndexedValue(ValueHandle(v), indices) {} + explicit IndexedValue(ValueHandle v, llvm::ArrayRef indices = {}) + : base(v), indices(indices.begin(), indices.end()) {} IndexedValue(const IndexedValue &rhs) = default; IndexedValue &operator=(const IndexedValue &rhs) = default; /// Returns a new `IndexedValue`. IndexedValue operator()(llvm::ArrayRef indices = {}) { - return IndexedValue(view, indices); + return IndexedValue(base, indices); } /// Emits a `store`. @@ -110,7 +116,7 @@ struct IndexedValue { return intrinsics::STORE(rhs, getBase(), indices); } - ValueHandle getBase() const { return view.base; } + ValueHandle getBase() const { return base; } /// Emits a `load` when converting to a ValueHandle. explicit operator ValueHandle() { @@ -152,7 +158,7 @@ struct IndexedValue { } private: - MemRefView &view; + ValueHandle base; llvm::SmallVector indices; }; diff --git a/mlir/test/EDSC/builder-api-test.cpp b/mlir/test/EDSC/builder-api-test.cpp index f2a60273785c..3bb22f0d8aa6 100644 --- a/mlir/test/EDSC/builder-api-test.cpp +++ b/mlir/test/EDSC/builder-api-test.cpp @@ -326,13 +326,16 @@ TEST_FUNC(builder_helpers) { // clang-format off ValueHandle f7( ValueHandle::create(llvm::APFloat(7.0f), f32Type)); - MemRefView vA(f->getArgument(0)), vB(f->getArgument(1)), vC(f->getArgument(2)); - IndexedValue A(vA), B(vB), C(vC); + MemRefView vA(f->getArgument(0)), vB(f->getArgument(1)), + vC(f->getArgument(2)); + IndexedValue A(f->getArgument(0)), B(f->getArgument(1)), C(f->getArgument(2)); IndexHandle i, j, k1, k2, lb0, lb1, lb2, ub0, ub1, ub2; int64_t step0, step1, step2; std::tie(lb0, ub0, step0) = vA.range(0); std::tie(lb1, ub1, step1) = vA.range(1); - std::tie(lb2, ub2, step2) = vA.range(2); + lb2 = vA.lb(2); + ub2 = vA.ub(2); + step2 = vA.step(2); LoopNestBuilder({&i, &j}, {lb0, lb1}, {ub0, ub1}, {step0, step1})({ LoopBuilder(&k1, lb2, ub2, step2)({ C({i, j, k1}) = f7 + A({i, j, k1}) + B({i, j, k1}),