Commit 3d323e71 authored by Nguyen, Thien Minh's avatar Nguyen, Thien Minh
Browse files

Handle array slicing by range 



Signed-off-by: default avatarThien Nguyen <nguyentm@ornl.gov>
parent 7b976dbb

mlir/parsers/qasm3/visitor_handlers/alias_handler.cpp

+2 −2
Original line number Diff line number Diff line
@@ -91,14 +91,14 @@ antlrcpp::Any qasm3_visitor::visitAliasStatement( 
      printErrorMessage("Invalid range: step size must be non-zero.");

    }



    std::cout << "Range: Start = " << range_start << "; Step = " << range_step << "; Stop = " << range_stop << "\n";

    // std::cout << "Range: Start = " << range_start << "; Step = " <<

    // range_step << "; Stop = " << range_stop << "\n";

    auto range_start_mlir_val = get_or_create_constant_integer_value(

        range_start, location, builder.getI64Type(), symbol_table, builder);

    auto range_step_mlir_val = get_or_create_constant_integer_value(

        range_step, location, builder.getI64Type(), symbol_table, builder);

    auto range_stop_mlir_val = get_or_create_constant_integer_value(

        range_stop, location, builder.getI64Type(), symbol_table, builder);

    // I think we can handle RANGE with a memref<3xi64>...

    mlir::Value array_slice = builder.create<mlir::quantum::ArraySliceOp>(

        location, array_type, allocated_symbol,

        llvm::makeArrayRef(std::vector<mlir::Value>{

mlir/qir_qrt/qir-qrt.hpp

+12 −1
Original line number Diff line number Diff line
@@ -72,7 +72,18 @@ Array *__quantum__rt__array_concatenate(Array *head, Array *tail); 
// Creates and returns an array that is a slice of an existing array. 

// The int dim which dimension the slice is on (0 for 1d arrays). 

// The Range range specifies the slice.

Array *__quantum__rt__array_slice(Array *array, int32_t dim, Range range);

// Note: QIR defines a Range as type { i64, i64, i64 }

// i.e. a struct of 3 int64_t

// and define an API at the *LLVM IR* level of passing this by value

// i.e. the signature is %Range, not "%struct.Range* byval(%struct.Range)" 

// Hence, it is actually equivalent to an expanded list of struct member.

// https://lists.llvm.org/pipermail/llvm-dev/2018-April/122714.html

// Until the spec. is updated (see https://github.com/microsoft/qsharp-language/issues/31)

// this is actually the C-ABI that will match the QIR IR.

Array *__quantum__rt__array_slice(Array *array, int32_t dim, int64_t range_start, int64_t range_step, int64_t range_end);

// Note: Overloading is not possible in C, so just keep the implementation in this local func.

Array *quantum__rt__array_slice(Array *array, int32_t dim, Range range);



// Ref. counting

void __quantum__rt__array_update_alias_count(Array *array, int64_t increment);

void __quantum__rt__array_update_reference_count(Array *aux, int64_t count);

mlir/qir_qrt/qir-range-util.cpp

0 → 100644
+52 −0
Original line number Diff line number Diff line 
#include "qir-qrt.hpp"



namespace qcor {

std::vector<int64_t> getRangeValues(Array *in_array, const Range &in_range) {

  const bool is_fwd_range = in_range.step > 0;



  const auto convertIndex = [&](int64_t in_rawIdx) -> int64_t {

    if (in_rawIdx >= 0) {

      return in_rawIdx;

    }

    // Negative-based index:

    // in_rawIdx = -1 => size - 1 (last element)

    int64_t result = in_array->size() + in_rawIdx;

    if (result < 0) {

      throw std::invalid_argument("range");

    }

    return result;

  };



  // Convert to absolute index.

  const auto start_idx = convertIndex(in_range.start);

  const auto end_idx = convertIndex(in_range.end);

  // start == end

  if (start_idx == end_idx) {

    return {end_idx};

  }



  if (is_fwd_range) {

    if (start_idx > end_idx) {

      return {};

    } else {

      assert(in_range.step > 0);

      std::vector<int64_t> result;

      for (int64_t i = start_idx; i <= end_idx; i += in_range.step) {

        result.emplace_back(i);

      }

      return result;

    }

  } else {

    if (start_idx < end_idx) {

      return {};

    } else {

      std::vector<int64_t> result;

      assert(in_range.step < 0);

      for (int64_t i = start_idx; i >= end_idx; i += in_range.step) {

        result.emplace_back(i);

      }

      return result;

    }

  }

}

} // namespace qcor
 
 No newline at end of file

mlir/qir_qrt/qir-types.hpp

+8 −2
Original line number Diff line number Diff line
@@ -59,8 +59,9 @@ struct Array { 
                     other.m_storage.end());

  }



  int64_t size() { return m_storage.size() / m_itemSizeInBytes; }

  int64_t size() const { return m_storage.size() / m_itemSizeInBytes; }

  void clear() { m_storage.clear(); }

  int64_t element_size() const { return m_itemSizeInBytes; }



private:

  // Must be const, i.e. changing the element size is NOT allowed.
@@ -85,3 +86,8 @@ struct Range { 
  int64_t step;

  int64_t end;

};



namespace qcor {

// Helper func.

std::vector<int64_t> getRangeValues(Array *in_array, const Range &in_range);

} // namespace qcor

mlir/qir_qrt/quantum_rt_array.cpp

+31 −6
Original line number Diff line number Diff line 
#include "qir-qrt.hpp"

#include <iostream>



#include <cstring>

extern "C" {

Array *__quantum__rt__array_create_1d(int32_t itemSizeInBytes,

                                      int64_t count_items) {
@@ -15,13 +15,38 @@ int64_t __quantum__rt__array_get_size_1d(Array *state1) { 
  return state1->size();

}



Array *__quantum__rt__array_slice(Array *array, int32_t dim, Range range) {

Array *__quantum__rt__array_slice(Array *array, int32_t dim,

                                  int64_t range_start, int64_t range_step,

                                  int64_t range_end) {

  return quantum__rt__array_slice(array, dim,

                                  {range_start, range_step, range_end});

}



Array *quantum__rt__array_slice(Array *array, int32_t dim, Range range) {

  if (verbose)

    std::cout << "[qir-qrt] Extract array slice for the range [" << range.start

              << ":" << range.step << ":" << range.end << "].\n";

    std::cout << "[qir-qrt] Extract array slice (dim = " << dim

              << ") for the range [" << range.start << ":" << range.step << ":"

              << range.end << "].\n";



  // TODO:

  return array;

  const std::vector<int64_t> range_idxs = qcor::getRangeValues(array, range);

  if (verbose) {

    std::cout << "[qir-qrt] Resolve range indices:";

    for (const auto &idx : range_idxs) {

      std::cout << idx << " ";

    }

    std::cout << "\n";

  }



  auto resultArray = new Array(range_idxs.size(), array->element_size());

  int64_t counter = 0;

  for (const auto &idx : range_idxs) {

    int8_t *src_ptr = array->getItemPointer(idx);

    int8_t *dest_ptr = resultArray->getItemPointer(counter);

    memcpy(dest_ptr, src_ptr, array->element_size());

    counter++;

  }



  return resultArray;

}



void __quantum__rt__array_update_alias_count(Array *array, int64_t increment) {