diff --git a/include/cudaq/Optimizer/Dialect/CC/CCOps.td b/include/cudaq/Optimizer/Dialect/CC/CCOps.td index 10accb100c..e5a9d3a340 100644 --- a/include/cudaq/Optimizer/Dialect/CC/CCOps.td +++ b/include/cudaq/Optimizer/Dialect/CC/CCOps.td @@ -1066,10 +1066,10 @@ def cc_ComputePtrOp : CCOp<"compute_ptr", [Pure]> { def cc_SizeOfOp : CCOp<"sizeof", [Pure]> { let summary = "Return the size, in bytes, of a type."; let description = [{ - This operation is the size, in bytes, of a type that can be reified. It can - be used to determine how many bytes to allocate or move when a value is in - memory, for example. If the type cannot be reified, it returns a size of - poison. + This operation returns the size, in bytes, of a type that can be reified. + It can be used to determine how many bytes to allocate or move when a value + is in memory, for example. If the type cannot be reified, it returns a size + of poison. Example: @@ -1089,6 +1089,38 @@ def cc_SizeOfOp : CCOp<"sizeof", [Pure]> { }]; } +def cc_OffsetOfOp : CCOp<"offsetof", [Pure]> { + let summary = "Return the interior offset, in bytes, of a subelement."; + let description = [{ + This operation returns the offset, in bytes, of a subelement of a type that + can be reified. It can be used to determine the byte offset of a particular + field in a struct type, for example. If the type cannot be reified, this op + returns a value of poison. + + Example: + ```mlir + // the offset of the i64 in the struct + %0 = cc.offsetof !cc.struct [2] : i64 + // the offset of of the 3rd i32 element in the array field + %1 = cc.offsetof !cc.struct x 4>, f32> + [0, 2, 0] : i64 + ``` + }]; + + let arguments = (ins + TypeAttr:$inputType, + DenseI32ArrayAttr:$constantIndices + ); + let results = (outs AnySignlessInteger); + + let hasVerifier = 1; + let hasCanonicalizer = 1; + + let assemblyFormat = [{ + $inputType $constantIndices `:` type(results) attr-dict + }]; +} + def cc_ConstantArrayOp : CCOp<"const_array", [Pure]> { let summary = "A one dimensional array of constant values."; let description = [{ diff --git a/lib/Optimizer/Dialect/CC/CCOps.cpp b/lib/Optimizer/Dialect/CC/CCOps.cpp index 49a686dfe7..9035113581 100644 --- a/lib/Optimizer/Dialect/CC/CCOps.cpp +++ b/lib/Optimizer/Dialect/CC/CCOps.cpp @@ -345,7 +345,8 @@ LogicalResult cudaq::cc::ComputePtrOp::verify() { auto resultTy = cast(getResult().getType()); for (std::int32_t i : getRawConstantIndices()) { if (auto arrTy = dyn_cast(eleTy)) { - if (i != kDynamicIndex && !arrTy.isUnknownSize() && i > arrTy.getSize()) { + if (i != kDynamicIndex && !arrTy.isUnknownSize() && + (i < 0 || i > arrTy.getSize())) { // Note: allow indexing of last element + 1 so we can compute a // pointer to `end()` of a stdvec buffer. Consider adding a flag // to cc.compute_ptr explicitly for this or using casts. @@ -355,9 +356,13 @@ LogicalResult cudaq::cc::ComputePtrOp::verify() { } else if (auto strTy = dyn_cast(eleTy)) { if (i == kDynamicIndex) return emitOpError("struct cannot have non-constant index"); - if (static_cast(i) >= strTy.getMembers().size()) + if (i < 0 || static_cast(i) >= strTy.getMembers().size()) return emitOpError("struct cannot index out of bounds members"); eleTy = strTy.getMember(i); + } else if (auto complexTy = dyn_cast(eleTy)) { + if (!(i == 0 || i == 1 || i == kDynamicIndex)) + return emitOpError("complex index is out of bounds"); + eleTy = complexTy.getElementType(); } else { return emitOpError("too many indices (" + std::to_string(getRawConstantIndices().size()) + @@ -658,15 +663,19 @@ LogicalResult cudaq::cc::ExtractValueOp::verify() { if (auto arrTy = dyn_cast(eleTy)) { if (arrTy.isUnknownSize()) return emitOpError("array must have constant size"); - if (i != kDynamicIndex && i >= arrTy.getSize()) + if (i != kDynamicIndex && (i < 0 || i >= arrTy.getSize())) return emitOpError("array cannot index out of bounds elements"); eleTy = arrTy.getElementType(); } else if (auto strTy = dyn_cast(eleTy)) { if (i == kDynamicIndex) return emitOpError("struct cannot have non-constant index"); - if (static_cast(i) >= strTy.getMembers().size()) + if (i < 0 || static_cast(i) >= strTy.getMembers().size()) return emitOpError("struct cannot index out of bounds members"); eleTy = strTy.getMember(i); + } else if (auto complexTy = dyn_cast(eleTy)) { + if (!(i == 0 || i == 1)) + return emitOpError("complex index is out of bounds"); + eleTy = complexTy.getElementType(); } else { return emitOpError("too many indices (" + std::to_string(getRawConstantIndices().size()) + @@ -1788,6 +1797,72 @@ MutableOperandRange cudaq::cc::ConditionOp::getMutableSuccessorOperands( return getResultsMutable(); } +//===----------------------------------------------------------------------===// +// OffsetOfOp +//===----------------------------------------------------------------------===// + +LogicalResult cudaq::cc::OffsetOfOp::verify() { + Type ty = getInputType(); + for (std::int32_t i : getConstantIndices()) { + if (auto strTy = dyn_cast(ty)) { + if (i < 0 || static_cast(i) >= strTy.getMembers().size()) + return emitOpError("out of bounds for struct"); + ty = strTy.getMembers()[i]; + continue; + } + if (auto arrTy = dyn_cast(ty)) { + if (arrTy.isUnknownSize()) + return emitOpError("array must have bounds"); + if (i < 0 || i >= arrTy.getSize()) + return emitOpError("out of bounds for array"); + ty = arrTy.getElementType(); + continue; + } + if (auto complexTy = dyn_cast(ty)) { + if (i < 0 || i > 1) + return emitOpError("out of bounds for complex"); + ty = complexTy.getElementType(); + continue; + } + // Walking off the end of the type. + return failure(); + } + return success(); +} + +namespace { +struct FoldTrivialOffsetOf : public OpRewritePattern { + using Base = OpRewritePattern; + using Base::Base; + + LogicalResult matchAndRewrite(cudaq::cc::OffsetOfOp offOp, + PatternRewriter &rewriter) const override { + // If there are no offsets, the offset is 0. + if (offOp.getConstantIndices().empty()) { + rewriter.replaceOpWithNewOp(offOp, 0, + offOp.getType()); + return success(); + } + + // If all indices are 0, then the offset is necessarily 0. + if (std::all_of(offOp.getConstantIndices().begin(), + offOp.getConstantIndices().end(), + [](std::int32_t i) { return i == 0; })) { + rewriter.replaceOpWithNewOp(offOp, 0, + offOp.getType()); + return success(); + } + + return failure(); + } +}; +} // namespace + +void cudaq::cc::OffsetOfOp::getCanonicalizationPatterns( + RewritePatternSet &patterns, MLIRContext *context) { + patterns.add(context); +} + //===----------------------------------------------------------------------===// // ReturnOp //===----------------------------------------------------------------------===// diff --git a/test/Quake/roundtrip-ops.qke b/test/Quake/roundtrip-ops.qke index 742923856b..fe643c5ab5 100644 --- a/test/Quake/roundtrip-ops.qke +++ b/test/Quake/roundtrip-ops.qke @@ -760,3 +760,20 @@ func.func @phys_hoops() { // CHECK: quake.return_wire %[[VAL_2]]#1 : !quake.wire // CHECK: return // CHECK: } + +func.func @offsets() { + %0 = cc.offsetof !cc.struct<{i32, f64}> [1] : i64 + %1 = cc.offsetof !cc.struct<{i32, f64}> [1] : i64 + %2 = cc.offsetof !cc.struct<{i32, f64, complex}> [2,1] : i64 + %3 = cc.offsetof !cc.array x 100> [86, 1] : i64 + %4 = cc.offsetof !cc.array x 100> [0, 0] : i64 + return +} + +// CHECK-LABEL: func.func @offsets() { +// CHECK: %[[VAL_0:.*]] = cc.offsetof !cc.struct<{i32, f64}> [1] : i64 +// CHECK: %[[VAL_1:.*]] = cc.offsetof !cc.struct<{i32, f64}> [1] : i64 +// CHECK: %[[VAL_2:.*]] = cc.offsetof !cc.struct<{i32, f64, complex}> [2, 1] : i64 +// CHECK: %[[VAL_3:.*]] = cc.offsetof !cc.array x 100> [86, 1] : i64 +// CHECK: %[[VAL_4:.*]] = cc.offsetof !cc.array x 100> [0, 0] : i64 +// CHECK: return