huajsj commented on a change in pull request #8655:
URL: https://github.com/apache/tvm/pull/8655#discussion_r684579481
##########
File path: tests/python/unittest/test_target_codegen_arm.py
##########
@@ -21,6 +21,93 @@
import ctypes
+def test_llvm_flip_pipeline_sve():
+ target = "llvm -device=arm_cpu -mtriple=aarch64-gnu-linux
-mattr=v8.2a,+sve"
+
+ def check_llvm(nn, base):
+ n = tvm.runtime.convert(nn)
+ A = te.placeholder((n + base), name="A")
+ C = te.compute((n,), lambda i: A(nn + base - i - 1), name="C")
+ s = te.create_schedule(C.op)
+ xo, xi = s[C].split(C.op.axis[0], factor=4)
+ s[C].parallel(xo)
+ s[C].vectorize_scalable(xi)
+
+ # build and invoke the kernel.
+ f = tvm.build(s, [A, C], target)
+
+ # ctx = remote.context(target)
+ # # launch the kernel.
+ # n = nn
+ # a = tvm.nd.array(np.random.uniform(size=(n +
base)).astype(A.dtype), ctx)
+ # c = tvm.nd.array(np.zeros(n, dtype=C.dtype), ctx)
+ # f(a, c)
+ # tvm.testing.assert_allclose(c.asnumpy(), a.asnumpy()[::-1][:n])
Review comment:
remove
##########
File path: include/tvm/runtime/data_type.h
##########
@@ -107,6 +119,13 @@ class DataType {
bool is_vector_bool() const { return is_vector() && bits() == 1; }
/*! \return whether type is a Void type. */
bool is_void() const { return code() == DataType::kHandle && bits() == 0 &&
lanes() == 0; }
+ bool is_scalable() const { return is_scalable_; }
+
+ DataType with_scalable_lanes() const {
+ int min_num_lanes = 128 / bits();
Review comment:
macro? and does here only support vector size 128 or should be
configurable between 128 - 2048?
##########
File path: include/tvm/runtime/data_type.h
##########
@@ -107,6 +119,13 @@ class DataType {
bool is_vector_bool() const { return is_vector() && bits() == 1; }
/*! \return whether type is a Void type. */
bool is_void() const { return code() == DataType::kHandle && bits() == 0 &&
lanes() == 0; }
+ bool is_scalable() const { return is_scalable_; }
+
+ DataType with_scalable_lanes() const {
Review comment:
how this function cooperate with is_scalable_? what happen if
is_scalable_ is false but this function get called?
##########
File path: include/tvm/runtime/data_type.h
##########
@@ -131,7 +150,7 @@ class DataType {
*/
bool operator==(const DataType& other) const {
return data_.code == other.data_.code && data_.bits == other.data_.bits &&
- data_.lanes == other.data_.lanes;
+ data_.lanes == other.data_.lanes; // && is_scalable_ ==
other.is_scalable_;
Review comment:
seem like no diff with original code, is the comment for "is_scalable_
== other.is_scalable" intend logic?
##########
File path: include/tvm/runtime/data_type.h
##########
@@ -336,6 +366,17 @@ inline std::string DLDataType2String(DLDataType t) {
return os.str();
}
+inline std::string VLADataType2String(DataType t) {
+ if (t.bits() == 0) return "";
+ std::ostringstream os;
+ os << t.operator DLDataType();
+ // auto const str_to_parse = os.str();
+ // auto pos = str_to_parse.find("x");
+ // auto stem= str_to_parse.substr(0, pos);
Review comment:
remove
##########
File path: include/tvm/runtime/data_type.h
##########
@@ -69,14 +69,26 @@ class DataType {
* \param bits The number of bits in the type.
* \param lanes The number of lanes.
*/
- DataType(int code, int bits, int lanes) {
+ DataType(int code, int bits, int lanes, bool is_scalable = false) {
data_.code = static_cast<uint8_t>(code);
data_.bits = static_cast<uint8_t>(bits);
data_.lanes = static_cast<uint16_t>(lanes);
+ is_scalable_ = is_scalable;
if (code == kBFloat) {
ICHECK_EQ(bits, 16);
}
}
+ // DataType(int code, int bits) {
+ // data_.code = static_cast<uint8_t>(code);
+ // data_.bits = static_cast<uint8_t>(bits);
+ // is_scalable_ = true;
+ // std::cout<<bits<<std::endl;
+ // data_.lanes = uint16_t(128) / static_cast<uint16_t>(8); // minimal
lanes
+ //
+ //// if (code == kBFloat) {
+ //// ICHECK_EQ(bits, 16);
+ //// }
+ // }
Review comment:
remove
##########
File path: tests/python/unittest/test_target_codegen_arm.py
##########
@@ -21,6 +21,93 @@
import ctypes
+def test_llvm_flip_pipeline_sve():
+ target = "llvm -device=arm_cpu -mtriple=aarch64-gnu-linux
-mattr=v8.2a,+sve"
+
+ def check_llvm(nn, base):
+ n = tvm.runtime.convert(nn)
+ A = te.placeholder((n + base), name="A")
+ C = te.compute((n,), lambda i: A(nn + base - i - 1), name="C")
+ s = te.create_schedule(C.op)
+ xo, xi = s[C].split(C.op.axis[0], factor=4)
+ s[C].parallel(xo)
+ s[C].vectorize_scalable(xi)
+
+ # build and invoke the kernel.
+ f = tvm.build(s, [A, C], target)
+
+ # ctx = remote.context(target)
+ # # launch the kernel.
+ # n = nn
+ # a = tvm.nd.array(np.random.uniform(size=(n +
base)).astype(A.dtype), ctx)
+ # c = tvm.nd.array(np.zeros(n, dtype=C.dtype), ctx)
+ # f(a, c)
+ # tvm.testing.assert_allclose(c.asnumpy(), a.asnumpy()[::-1][:n])
+
+ check_llvm(4, 0)
+ check_llvm(128, 8)
+ check_llvm(3, 0)
+ check_llvm(128, 1)
+
+
+def test_llvm_vadd_pipeline_sve():
+ target = "llvm -device=arm_cpu -mtriple=aarch64-gnu-linux
-mattr=v8.2a,+sve"
+
+ def check_llvm(n, lanes):
+ A = te.placeholder((n,), name="A", dtype="float32x%d" % lanes)
+ B = te.compute((n,), lambda i: A[i], name="B")
+ C = te.compute((n,), lambda i: B[i] + tvm.tir.const(1, A.dtype),
name="C")
+ s = te.create_schedule(C.op)
+ xo, xi = s[C].split(C.op.axis[0], nparts=2)
+ _, xi = s[C].split(xi, factor=2)
+ s[C].parallel(xo)
+ s[C].vectorize_scalable(xi)
+ s[B].compute_at(s[C], xo)
+ xo, xi = s[B].split(B.op.axis[0], factor=2)
+ s[B].vectorize_scalable(xi)
+ # build and invoke the kernel.
+ f = tvm.build(s, [A, C], target)
+
+ # ctx = remote.context(target)
+ # # launch the kernel.
+ # a = tvm.nd.empty((n,), A.dtype,
ctx).copyfrom(np.random.uniform(size=(n, lanes)))
+ # c = tvm.nd.empty((n,), C.dtype, ctx)
+ # f(a, c)
+ # tvm.testing.assert_allclose(c.asnumpy(), a.asnumpy() + 1)
Review comment:
remove
##########
File path: include/tvm/runtime/data_type.h
##########
@@ -107,6 +119,13 @@ class DataType {
bool is_vector_bool() const { return is_vector() && bits() == 1; }
/*! \return whether type is a Void type. */
bool is_void() const { return code() == DataType::kHandle && bits() == 0 &&
lanes() == 0; }
+ bool is_scalable() const { return is_scalable_; }
Review comment:
doxygen comments.
##########
File path: tests/python/unittest/test_target_codegen_arm.py
##########
@@ -21,6 +21,93 @@
import ctypes
+def test_llvm_flip_pipeline_sve():
+ target = "llvm -device=arm_cpu -mtriple=aarch64-gnu-linux
-mattr=v8.2a,+sve"
+
+ def check_llvm(nn, base):
+ n = tvm.runtime.convert(nn)
+ A = te.placeholder((n + base), name="A")
+ C = te.compute((n,), lambda i: A(nn + base - i - 1), name="C")
+ s = te.create_schedule(C.op)
+ xo, xi = s[C].split(C.op.axis[0], factor=4)
+ s[C].parallel(xo)
+ s[C].vectorize_scalable(xi)
+
+ # build and invoke the kernel.
+ f = tvm.build(s, [A, C], target)
+
+ # ctx = remote.context(target)
+ # # launch the kernel.
+ # n = nn
+ # a = tvm.nd.array(np.random.uniform(size=(n +
base)).astype(A.dtype), ctx)
+ # c = tvm.nd.array(np.zeros(n, dtype=C.dtype), ctx)
+ # f(a, c)
+ # tvm.testing.assert_allclose(c.asnumpy(), a.asnumpy()[::-1][:n])
+
+ check_llvm(4, 0)
+ check_llvm(128, 8)
+ check_llvm(3, 0)
+ check_llvm(128, 1)
+
+
+def test_llvm_vadd_pipeline_sve():
+ target = "llvm -device=arm_cpu -mtriple=aarch64-gnu-linux
-mattr=v8.2a,+sve"
+
+ def check_llvm(n, lanes):
+ A = te.placeholder((n,), name="A", dtype="float32x%d" % lanes)
+ B = te.compute((n,), lambda i: A[i], name="B")
+ C = te.compute((n,), lambda i: B[i] + tvm.tir.const(1, A.dtype),
name="C")
+ s = te.create_schedule(C.op)
+ xo, xi = s[C].split(C.op.axis[0], nparts=2)
+ _, xi = s[C].split(xi, factor=2)
+ s[C].parallel(xo)
+ s[C].vectorize_scalable(xi)
+ s[B].compute_at(s[C], xo)
+ xo, xi = s[B].split(B.op.axis[0], factor=2)
+ s[B].vectorize_scalable(xi)
+ # build and invoke the kernel.
+ f = tvm.build(s, [A, C], target)
+
+ # ctx = remote.context(target)
+ # # launch the kernel.
+ # a = tvm.nd.empty((n,), A.dtype,
ctx).copyfrom(np.random.uniform(size=(n, lanes)))
+ # c = tvm.nd.empty((n,), C.dtype, ctx)
+ # f(a, c)
+ # tvm.testing.assert_allclose(c.asnumpy(), a.asnumpy() + 1)
+
+ check_llvm(64, 2)
+ check_llvm(512, 2)
+
+
+def test_llvm_madd_pipeline_sve():
+ target = "llvm -device=arm_cpu -mtriple=aarch64-gnu-linux
-mattr=v8.2a,+sve"
+
+ def check_llvm(nn, base, stride):
+ n = tvm.runtime.convert(nn)
+ A = te.placeholder((n + base, stride), name="A")
+ C = te.compute((n, stride), lambda i, j: A(base + i, j) + 1, name="C")
+ s = te.create_schedule(C.op)
+ xo, xi = s[C].split(C.op.axis[0], factor=4)
+ s[C].parallel(xo)
+ s[C].vectorize_scalable(xi)
+ # build and invoke the kernel.
+ f = tvm.build(s, [A, C], target)
+
+ # ctx = remote.context(target)
+ # # launch the kernel.
+ # n = nn
+ # a = tvm.nd.array(np.random.uniform(size=(n + base,
stride)).astype(A.dtype), ctx)
+ # c = tvm.nd.array(np.zeros((n, stride), dtype=C.dtype), ctx)
+ # f(a, c)
+ # tvm.testing.assert_allclose(c.asnumpy(), a.asnumpy()[base:] + 1)
Review comment:
remove
##########
File path: src/target/llvm/codegen_aarch64.cc
##########
@@ -0,0 +1,298 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+/*!
+ * \file codegen_arm.cc
+ * \brief ARM specific code generator
+ */
+#ifdef TVM_LLVM_VERSION
+
+#include <tvm/runtime/registry.h>
+
+#include "codegen_cpu.h"
+
+namespace tvm {
+namespace codegen {
+
+// -------------------------
+// Utility functions to remove
+void print_LLVM_type(llvm::Type* type) {
+ std::string type_str;
+ llvm::raw_string_ostream rso(type_str);
+ type->print(rso);
+ std::cout << rso.str() << std::endl;
+ ;
+}
+
+void print_LLVM_val(llvm::Value* val) {
+ std::string type_str;
+ llvm::raw_string_ostream rso(type_str);
+ val->print(rso);
+ std::cout << rso.str() << std::endl;
+ ;
+}
+// -------------------------
+
+// AArch64 code generation
+class CodeGenAArch64 final : public CodeGenCPU {
+ public:
+ void InitTarget(llvm::TargetMachine* tm) final {
+ // set native vector bits.
+ native_vector_bits_ = 16 * 8;
+ CodeGenCPU::InitTarget(tm);
+ }
+ llvm::Value* VisitExpr_(const LoadNode* op);
+ void VisitStmt_(const ForNode* op);
+ void VisitStmt_(const StoreNode* op);
+
+ private:
+ // SVE LLVM intrinsics
+ llvm::Value* sve_stride(int min_lanes);
+ llvm::Value* sve_whilelt(llvm::Value* a, llvm::Value* b, int min_lanes);
+ llvm::Value* sve_store(llvm::Value* ptr, llvm::Value* val, DataType t);
+ llvm::Value* sve_load(llvm::Value* ptr, DataType t);
+ void CreateSVEFor(llvm::Value* begin, llvm::Value* end, llvm::Value* stride,
const Var& loop_var,
+ const Stmt& body, int min_lanes);
+
+ // Predicate
+ llvm::Value* mask_;
+};
+
+llvm::Value* CodeGenAArch64::sve_stride(int min_lanes) {
+ llvm::Intrinsic::ID cnt_id;
+
+ switch (min_lanes) {
+ case 16:
+ cnt_id = llvm::Function::lookupIntrinsicID("llvm.aarch64.sve.cntb");
+ break;
+ case 8: // half
+ cnt_id = llvm::Function::lookupIntrinsicID("llvm.aarch64.sve.cnth");
+ break;
+ case 4: // float
+ cnt_id = llvm::Function::lookupIntrinsicID("llvm.aarch64.sve.cntw");
+ break;
+ default: // double
+ cnt_id = llvm::Function::lookupIntrinsicID("llvm.aarch64.sve.cntd");
+ }
+
+ // All pattern
+ int all_pattern = 31;
Review comment:
use macro?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
