halide · steven-johnson · Sep 24, 2020 · Sep 24, 2020 · Sep 24, 2020 · Sep 24, 2020
diff --git a/python_bindings/src/PyEnums.cpp b/python_bindings/src/PyEnums.cpp
@@ -36,7 +36,10 @@ void define_enums(py::module &m) {
         .value("Heap", MemoryType::Heap)
         .value("Stack", MemoryType::Stack)
         .value("Register", MemoryType::Register)
-        .value("GPUShared", MemoryType::GPUShared);
+        .value("GPUShared", MemoryType::GPUShared)
+        .value("GPUTexture", MemoryType::GPUTexture)
+        .value("LockedCache", MemoryType::LockedCache)
+        .value("VTCM", MemoryType::VTCM);
 
     py::enum_<NameMangling>(m, "NameMangling")
         .value("Default", NameMangling::Default)

diff --git a/python_bindings/src/PyImageParam.cpp b/python_bindings/src/PyImageParam.cpp
@@ -31,6 +31,7 @@ void define_image_param(py::module &m) {
             .def("host_alignment", &OutputImageParam::host_alignment)
             .def("set_estimates", &OutputImageParam::set_estimates, py::arg("estimates"))
             .def("set_host_alignment", &OutputImageParam::set_host_alignment)
+            .def("store_in", &OutputImageParam::store_in, py::arg("memory_type"))
             .def("dimensions", &OutputImageParam::dimensions)
             .def("left", &OutputImageParam::left)
             .def("right", &OutputImageParam::right)

diff --git a/src/Closure.h b/src/Closure.h
@@ -55,11 +55,14 @@ class Closure : public IRVisitor {
         /** The buffer is written to. */
         bool write;
 
+        /** The buffer is a texture */
+        bool texture;
+
         /** The size of the buffer if known, otherwise zero. */
         size_t size;
 
         Buffer()
-            : dimensions(0), read(false), write(false), size(0) {
+            : dimensions(0), read(false), write(false), texture(false), size(0) {
         }
     };
 

diff --git a/src/CodeGen_OpenCL_Dev.cpp b/src/CodeGen_OpenCL_Dev.cpp
@@ -1,4 +1,5 @@
 #include <algorithm>
+#include <array>
 #include <sstream>
 #include <utility>
 
@@ -174,6 +175,21 @@ string CodeGen_OpenCL_Dev::CodeGen_OpenCL_C::get_memory_space(const string &buf)
     }
 }
 
+namespace {
+std::string image_type_suffix(const Type &type) {
+    if (type.is_int()) {
+        return "i";
+    } else if (type.is_uint()) {
+        return "ui";
+    } else if (type.is_float()) {
+        return "f";
+    } else {
+        internal_error << "Invalid type for image: " << type << "\n";
+    }
+    return "";
+}
+}  // namespace
+
 void CodeGen_OpenCL_Dev::CodeGen_OpenCL_C::visit(const Call *op) {
     if (op->is_intrinsic(Call::bool_to_mask)) {
         if (op->args[0].type().is_vector()) {
@@ -242,6 +258,107 @@ void CodeGen_OpenCL_Dev::CodeGen_OpenCL_C::visit(const Call *op) {
         } else {
             CodeGen_C::visit(op);
         }
+    } else if (op->is_intrinsic(Call::image_load_texture)) {
+        // image_load(<image name>, <buffer>, <x>, <x-extent>, <y>,
+        // <y-extent>, <z>, <z-extent>)
+        int dims = (op->args.size() - 2) / 2;
+        internal_assert(dims >= 1 && dims <= 3);
+        const StringImm *string_imm = op->args[0].as<StringImm>();
+        if (!string_imm) {
+            internal_assert(op->args[0].as<Broadcast>());
+            string_imm = op->args[0].as<Broadcast>()->value.as<StringImm>();
+        }
+        internal_assert(string_imm);
+        Type arg_type = op->args[2].type();
+        internal_assert(arg_type.lanes() <= 16);
+        internal_assert(arg_type.lanes() == op->type.lanes());
+
+        std::array<string, 3> coord;
+        for (int i = 0; i < dims; i++) {
+            coord[i] = print_expr(op->args[i * 2 + 2]);
+        }
+        vector<string> results(arg_type.lanes());
+        // For vectorized reads, codegen as a sequence of read_image calls
+        for (int i = 0; i < arg_type.lanes(); i++) {
+            ostringstream rhs;
+            rhs << "read_image" << image_type_suffix(op->type) << "(" << print_name(string_imm->value) << ", ";
+            string idx = arg_type.is_vector() ? string(".s") + vector_elements[i] : "";
+            switch (dims) {
+            case 1:
+                rhs << coord[0] << idx << ").s0";
+                break;
+            case 2:
+                rhs << "(int2)(" << coord[0] << idx << ", " << coord[1] << idx << ")).s0";
+                break;
+            case 3:
+                rhs << "(int4)(" << coord[0] << idx << ", " << coord[1] << idx
+                    << ", " << coord[2] << idx << ", 0)).s0";
+                break;
+            }
+            print_assignment(op->type.with_bits(32).with_lanes(1), rhs.str());
+            results[i] = id;
+        }
+
+        if (op->type.is_vector()) {
+            // Combine all results into a single vector
+            ostringstream rhs;
+            rhs << "(" << print_type(op->type) << ")(";
+            for (int i = 0; i < op->type.lanes(); i++) {
+                rhs << results[i];
+                if (i < op->type.lanes() - 1) {
+                    rhs << ", ";
+                }
+            }
+            rhs << ")";
+            print_assignment(op->type, rhs.str());
+        }
+        if (op->type.bits() != 32) {
+            // Widen to the correct type
+            print_assignment(op->type, "convert_" + print_type(op->type) + "(" + id + ")");
+        }
+    } else if (op->is_intrinsic(Call::image_store_texture)) {
+        // image_store(<image name>, <buffer>, <x>, <y>, <z>, <value>)
+        const StringImm *string_imm = op->args[0].as<StringImm>();
+        if (!string_imm) {
+            internal_assert(op->args[0].as<Broadcast>());
+            string_imm = op->args[0].as<Broadcast>()->value.as<StringImm>();
+        }
+        internal_assert(string_imm);
+        int dims = op->args.size() - 3;
+        internal_assert(dims >= 1 && dims <= 3);
+        Type arg_type = op->args[2].type();
+        internal_assert(arg_type.lanes() <= 16);
+        Type value_type = op->args.back().type();
+        internal_assert(arg_type.lanes() == value_type.lanes());
+
+        std::array<string, 3> coord;
+        for (int i = 0; i < dims; i++) {
+            coord[i] = print_expr(op->args[i + 2]);
+        }
+        string value = print_expr(op->args.back());
+        // For vectorized writes, codegen as a sequence of write_image calls
+        for (int i = 0; i < arg_type.lanes(); i++) {
+            ostringstream write_image;
+            write_image << "write_image" << image_type_suffix(op->type)
+                        << "(" << print_name(string_imm->value) << ", ";
+            string idx = arg_type.is_vector() ? string(".s") + vector_elements[i] : "";
+            switch (dims) {
+            case 1:
+                write_image << coord[0] << idx;
+                break;
+            case 2:
+                write_image << "(int2)(" << coord[0] << idx << ", " << coord[1] << idx << ")";
+                break;
+            case 3:
+                write_image << "(int4)(" << coord[0] << idx << ", " << coord[1] << idx
+                            << ", " << coord[2] << idx << ", 0)";
+                break;
+            }
+            write_image << ", (" << print_type(value_type.with_bits(32).with_lanes(4))
+                        << ")(" << value << idx << ", 0, 0, 0));\n";
+            //  do_indent();
+            stream << write_image.str();
+        }
     } else {
         CodeGen_C::visit(op);
     }
@@ -765,11 +882,27 @@ void CodeGen_OpenCL_Dev::CodeGen_OpenCL_C::add_kernel(Stmt s,
     stream << "__kernel void " << name << "(\n";
     for (size_t i = 0; i < args.size(); i++) {
         if (args[i].is_buffer) {
-            stream << " " << get_memory_space(args[i].name) << " ";
-            if (!args[i].write) stream << "const ";
-            stream << print_type(args[i].type) << " *"
-                   << "restrict "
-                   << print_name(args[i].name);
+            if (args[i].is_texture) {
+                int dims = args[i].dimensions;
+                internal_assert(dims >= 1 && dims <= 3) << "dims = " << dims << "\n";
+                if (args[i].read && args[i].write) {
+                    stream << " __read_write ";
+                } else if (args[i].read) {
+                    stream << " __read_only ";
+                } else if (args[i].write) {
+                    stream << " __write_only ";
+                } else {
+                    internal_error << "CL Image argument " << args[i].name
+                                   << " is neither read nor write";
+                }
+                stream << "image" << dims << "d_t ";
+            } else {
+                stream << " " << get_memory_space(args[i].name) << " ";
+                if (!args[i].write) stream << "const ";
+                stream << print_type(args[i].type) << " *"
+                       << "restrict ";
+            }
+            stream << print_name(args[i].name);
             Allocation alloc;
             alloc.type = args[i].type;
             allocations.push(args[i].name, alloc);
@@ -797,7 +930,7 @@ void CodeGen_OpenCL_Dev::CodeGen_OpenCL_C::add_kernel(Stmt s,
         void visit(const Allocate *op) override {
             if (op->memory_type == MemoryType::GPUShared) {
                 internal_assert(alloc == nullptr)
-                    << "Found multiple shared allocations in metal kernel\n";
+                    << "Found multiple shared allocations in opencl kernel\n";
                 alloc = op;
             }
         }

diff --git a/src/DeviceArgument.cpp b/src/DeviceArgument.cpp
@@ -15,15 +15,17 @@ std::vector<DeviceArgument> HostClosure::arguments() {
     std::vector<DeviceArgument> res;
     for (const auto &v : vars) {
         debug(2) << "var: " << v.first << "\n";
-        res.emplace_back(v.first, false, v.second, 0);
+        res.emplace_back(v.first, false, false, v.second, 0);
     }
     for (const auto &b : buffers) {
         debug(2) << "buffer: " << b.first << " " << b.second.size;
         if (b.second.read) debug(2) << " (read)";
         if (b.second.write) debug(2) << " (write)";
+        if (b.second.texture) debug(2) << " <texture>";
+        debug(2) << " dims=" << (int)b.second.dimensions;
         debug(2) << "\n";
 
-        DeviceArgument arg(b.first, true, b.second.type, b.second.dimensions, b.second.size);
+        DeviceArgument arg(b.first, true, b.second.texture, b.second.type, b.second.dimensions, b.second.size);
         arg.read = b.second.read;
         arg.write = b.second.write;
         res.push_back(arg);
@@ -34,8 +36,10 @@ std::vector<DeviceArgument> HostClosure::arguments() {
 void HostClosure::visit(const Call *op) {
     if (op->is_intrinsic(Call::glsl_texture_load) ||
         op->is_intrinsic(Call::image_load) ||
+        op->is_intrinsic(Call::image_load_texture) ||
         op->is_intrinsic(Call::glsl_texture_store) ||
-        op->is_intrinsic(Call::image_store)) {
+        op->is_intrinsic(Call::image_store) ||
+        op->is_intrinsic(Call::image_store_texture)) {
 
         // The argument to the call is either a StringImm or a broadcasted
         // StringImm if this is part of a vectorized expression
@@ -51,14 +55,19 @@ void HostClosure::visit(const Call *op) {
         std::string bufname = string_imm->value;
         Buffer &ref = buffers[bufname];
         ref.type = op->type;
-        // TODO: do we need to set ref.dimensions?
+        ref.texture = op->is_intrinsic(Call::image_load_texture) ||
+                      op->is_intrinsic(Call::image_store_texture);
 
         if (op->is_intrinsic(Call::glsl_texture_load) ||
-            op->is_intrinsic(Call::image_load)) {
+            op->is_intrinsic(Call::image_load) ||
+            op->is_intrinsic(Call::image_load_texture)) {
             ref.read = true;
+            ref.dimensions = (op->args.size() - 2) / 2;
         } else if (op->is_intrinsic(Call::glsl_texture_store) ||
-                   op->is_intrinsic(Call::image_store)) {
+                   op->is_intrinsic(Call::image_store) ||
+                   op->is_intrinsic(Call::image_store_texture)) {
             ref.write = true;
+            ref.dimensions = op->args.size() - 3;
         }
 
         // The Func's name and the associated .buffer are mentioned in the

diff --git a/src/DeviceArgument.h b/src/DeviceArgument.h
@@ -36,6 +36,12 @@ struct DeviceArgument {
      */
     bool is_buffer;
 
+    /** If is_buffer == true and is_texture == true, this argument should be
+     * passed and accessed through texture sampler operations instead of
+     * directly as a memory array
+     */
+    bool is_texture;
+
     /** If is_buffer is true, this is the dimensionality of the buffer.
      * If is_buffer is false, this value is ignored (and should always be set to zero) */
     uint8_t dimensions;
@@ -66,6 +72,7 @@ struct DeviceArgument {
 
     DeviceArgument()
         : is_buffer(false),
+          is_texture(false),
           dimensions(0),
           size(0),
           packed_index(0),
@@ -75,11 +82,13 @@ struct DeviceArgument {
 
     DeviceArgument(const std::string &_name,
                    bool _is_buffer,
+                   bool _is_texture,
                    Type _type,
                    uint8_t _dimensions,
                    size_t _size = 0)
         : name(_name),
           is_buffer(_is_buffer),
+          is_texture(_is_texture),
           dimensions(_dimensions),
           type(_type),
           size(_size),

diff --git a/src/DeviceInterface.cpp b/src/DeviceInterface.cpp
@@ -162,7 +162,7 @@ DeviceAPI get_default_device_api_for_target(const Target &target) {
 }
 
 namespace Internal {
-Expr make_device_interface_call(DeviceAPI device_api) {
+Expr make_device_interface_call(DeviceAPI device_api, bool texture) {
     if (device_api == DeviceAPI::Host) {
         return make_zero(type_of<const halide_device_interface_t *>());
     }
@@ -173,7 +173,11 @@ Expr make_device_interface_call(DeviceAPI device_api) {
         interface_name = "halide_cuda_device_interface";
         break;
     case DeviceAPI::OpenCL:
-        interface_name = "halide_opencl_device_interface";
+        if (texture) {
+            interface_name = "halide_opencl_image_device_interface";
+        } else {
+            interface_name = "halide_opencl_device_interface";
+        }
         break;
     case DeviceAPI::Metal:
         interface_name = "halide_metal_device_interface";

diff --git a/src/DeviceInterface.h b/src/DeviceInterface.h
@@ -37,7 +37,7 @@ bool host_supports_target_device(const Target &t);
 
 namespace Internal {
 /** Get an Expr which evaluates to the device interface for the given device api at runtime. */
-Expr make_device_interface_call(DeviceAPI device_api);
+Expr make_device_interface_call(DeviceAPI device_api, bool texture = false);
 }  // namespace Internal
 
 }  // namespace Halide

diff --git a/src/Expr.h b/src/Expr.h
@@ -368,6 +368,10 @@ enum class MemoryType {
      * across GPU threads within the same block. */
     GPUShared,
 
+    /** Allocation is stored in GPU texture memory and accessed through
+     * hardware sampler */
+    GPUTexture,
+
     /** Allocate Locked Cache Memory to act as local memory */
     LockedCache,
     /** Vector Tightly Coupled Memory. HVX (Hexagon) local memory available on

diff --git a/src/FuseGPUThreadLoops.cpp b/src/FuseGPUThreadLoops.cpp
@@ -458,7 +458,8 @@ class ExtractSharedAndHeapAllocations : public IRMutator {
 
         if ((fixed_size_thread_allocation &&
              op->memory_type != MemoryType::Heap &&
-             op->memory_type != MemoryType::GPUShared) ||
+             op->memory_type != MemoryType::GPUShared &&
+             op->memory_type != MemoryType::GPUTexture) ||
             op->memory_type == MemoryType::Register ||
             op->memory_type == MemoryType::Stack) {
             // These allocations go in register or local memory
@@ -467,6 +468,7 @@ class ExtractSharedAndHeapAllocations : public IRMutator {
 
         user_assert(op->memory_type == MemoryType::Auto ||
                     op->memory_type == MemoryType::GPUShared ||
+                    op->memory_type == MemoryType::GPUTexture ||
                     op->memory_type == MemoryType::Heap)
             << "Allocation " << op->name << " must live in shared or heap memory, "
             << "but is scheduled to live in " << op->memory_type << " memory.\n";
@@ -1263,6 +1265,7 @@ class InjectThreadBarriers : public IRMutator {
             break;
         case MemoryType::Auto:
         case MemoryType::Heap:
+        case MemoryType::GPUTexture:
             debug(4) << "   memory type is heap or auto\n";
             device_stores.insert(op->name);
             break;
@@ -1286,6 +1289,7 @@ class InjectThreadBarriers : public IRMutator {
             break;
         case MemoryType::Auto:
         case MemoryType::Heap:
+        case MemoryType::GPUTexture:
             debug(4) << "   memory type is heap or auto\n";
             device_loads.insert(op->name);
             break;

diff --git a/src/Generator.cpp b/src/Generator.cpp
@@ -747,7 +747,7 @@ std::string halide_type_to_c_type(const Type &t) {
 
 int generate_filter_main_inner(int argc, char **argv, std::ostream &cerr) {
     const char kUsage[] =
-        "gengen \n"
+        "gengen  \n"
         "  [-g GENERATOR_NAME] [-f FUNCTION_NAME] [-o OUTPUT_DIR] [-r RUNTIME_NAME] [-d 1|0]\n"
         "  [-e EMIT_OPTIONS] [-n FILE_BASE_NAME] [-p PLUGIN_NAME] [-s AUTOSCHEDULER_NAME]\n"
         "       target=target-string[,target-string...] [generator_arg=value [...]]\n"