def const_array_like(ndarray): fname = "ptx.cmem.arylike" from .descriptor import CUDATargetDesc aryty = CUDATargetDesc.typingctx.resolve_argument_type(ndarray) sig = typing.signature(aryty, aryty) return ir.Intrinsic(fname, sig, args=[ndarray]) raise NotImplementedError [aryarg] = args ary = aryarg.value count = reduce(operator.mul, ary.shape) dtype = types.from_dtype(numpy.dtype(ary.dtype)) def impl(context, args, argtys, retty): builder = context.builder lmod = builder.basic_block.function.module addrspace = nvvm.ADDRSPACE_CONSTANT data_t = dtype.llvm_as_value() flat = ary.flatten(order='A') # preserve order constvals = [dtype.llvm_const(flat[i]) for i in range(flat.size)] constary = lc.Constant.array(data_t, constvals) gv = lmod.add_global_variable(constary.type, "cmem", addrspace) gv.linkage = lc.LINKAGE_INTERNAL gv.global_constant = True gv.initializer = constary byte = lc.Type.int(8) byte_ptr_as = lc.Type.pointer(byte, addrspace) to_generic = nvvmutils.insert_addrspace_conv(lmod, byte, addrspace) rawdata = builder.call(to_generic, [builder.bitcast(gv, byte_ptr_as)]) data = builder.bitcast(rawdata, lc.Type.pointer(data_t)) llintp = types.intp.llvm_as_value() cshape = lc.Constant.array(llintp, map(types.const_intp, ary.shape)) cstrides = lc.Constant.array(llintp, map(types.const_intp, ary.strides)) res = lc.Constant.struct([lc.Constant.null(data.type), cshape, cstrides]) res = builder.insert_value(res, data, 0) return res if ary.flags['C_CONTIGUOUS']: contig = 'C' elif ary.flags['F_CONTIGUOUS']: contig = 'F' else: raise TypeError("array must be either C/F contiguous to be used as a " "constant") impl.codegen = True impl.return_type = types.arraytype(dtype, ary.ndim, 'A') return impl