def visit_FieldDecl(self, node: cuir.FieldDecl, **kwargs): if "external_arg" in kwargs: if kwargs["external_arg"]: return "py::buffer {name}, std::array<gt::uint_t,{ndim}> {name}_origin".format( name=node.name, ndim=node.dimensions.count(True), ) else: num_dims = node.dimensions.count(True) sid_def = """gt::as_cuda_sid<{dtype}, {num_dims}, std::integral_constant<int, {unique_index}>>({name})""".format( name=node.name, dtype=self.visit(node.dtype), unique_index=self.unique_index(), num_dims=num_dims, ) if num_dims != 3: gt_dims = [ f"gt::stencil::dim::{dim}" for dim in gtc_utils.dimension_flags_to_names(node.dimensions) ] sid_def = "gt::sid::rename_numbered_dimensions<{gt_dims}>({sid_def})".format( gt_dims=", ".join(gt_dims), sid_def=sid_def ) return "gt::sid::shift_sid_origin({sid_def}, {name}_origin)".format( sid_def=sid_def, name=node.name, )
def make_args_data_from_gtir(pipeline: GtirPipeline) -> ModuleData: data = ModuleData() node = pipeline.full() field_extents = compute_legacy_extents(node) write_fields = ( node.iter_tree() .if_isinstance(gtir.ParAssignStmt) .getattr("left") .if_isinstance(gtir.FieldAccess) .getattr("name") .to_set() ) read_fields: Set[str] = set() for expr in node.iter_tree().if_isinstance(gtir.ParAssignStmt).getattr("right"): read_fields |= expr.iter_tree().if_isinstance(gtir.FieldAccess).getattr("name").to_set() referenced_field_params = [ param.name for param in node.params if isinstance(param, gtir.FieldDecl) ] for name in sorted(referenced_field_params): access = AccessKind.NONE if name in read_fields: access |= AccessKind.READ if name in write_fields: access |= AccessKind.WRITE data.field_info[name] = FieldInfo( access=access, boundary=field_extents[name].to_boundary(), axes=tuple(dimension_flags_to_names(node.symtable_[name].dimensions).upper()), data_dims=tuple(node.symtable_[name].data_dims), dtype=numpy.dtype(node.symtable_[name].dtype.name.lower()), ) referenced_scalar_params = [ param.name for param in node.params if param.name not in referenced_field_params ] for name in sorted(referenced_scalar_params): data.parameter_info[name] = ParameterInfo( dtype=numpy.dtype(node.symtable_[name].dtype.name.lower()) ) unref_params = get_unused_params_from_gtir(pipeline) for param in sorted(unref_params, key=lambda decl: decl.name): if isinstance(param, gtir.FieldDecl): data.field_info[param.name] = None elif isinstance(param, gtir.ScalarDecl): data.parameter_info[param.name] = None data.unreferenced = [*sorted(param.name for param in unref_params)] return data
def make_args_data_from_gtir(pipeline: GtirPipeline) -> ModuleData: """ Compute module data containing information about stencil arguments from gtir. This is no longer compatible with the legacy backends. """ data = ModuleData() # NOTE: pipeline.gtir has not had prune_unused_parameters applied. all_params = pipeline.gtir.params node = pipeline.full() oir = gtir_to_oir.GTIRToOIR().visit(node) field_extents = compute_fields_extents(oir) accesses = compute_access_kinds(oir) for decl in (param for param in all_params if isinstance(param, gtir.FieldDecl)): access = accesses[decl.name] dtype = numpy.dtype(decl.dtype.name.lower()) if access != AccessKind.NONE: k_boundary = compute_k_boundary(node)[decl.name] boundary = Boundary(*field_extents[decl.name].to_boundary()[0:2], k_boundary) else: boundary = Boundary.zeros(ndims=3) data.field_info[decl.name] = FieldInfo( access=access, boundary=boundary, axes=tuple(dimension_flags_to_names(decl.dimensions).upper()), data_dims=tuple(decl.data_dims), dtype=dtype, ) for decl in (param for param in all_params if isinstance(param, gtir.ScalarDecl)): access = accesses[decl.name] dtype = numpy.dtype(decl.dtype.name.lower()) data.parameter_info[decl.name] = ParameterInfo(access=access, dtype=dtype) data.unreferenced = [ *sorted(name for name in accesses if accesses[name] == AccessKind.NONE) ] return data
def visit_FieldDecl(self, node: gtcpp.FieldDecl, **kwargs): backend = kwargs["backend"] if "external_arg" in kwargs: domain_ndim = node.dimensions.count(True) data_ndim = len(node.data_dims) sid_ndim = domain_ndim + data_ndim if kwargs["external_arg"]: return "py::buffer {name}, std::array<gt::uint_t,{sid_ndim}> {name}_origin".format( name=node.name, sid_ndim=sid_ndim, ) else: make_layout_map = backend.storage_info["layout_map"] layout_map = [ x for x in make_layout_map(node.dimensions + (True, ) * data_ndim) if x is not None ] sid_def = """gt::as_{sid_type}<{dtype}, {sid_ndim}, gt::integral_constant<int, {unique_index}>, {unit_stride_dim}>({name})""".format( sid_type="cuda_sid" if backend.GT_BACKEND_T == "gpu" else "sid", name=node.name, dtype=self.visit(node.dtype), unique_index=self.unique_index(), sid_ndim=sid_ndim, unit_stride_dim=layout_map.index(max(layout_map)), ) sid_def = "gt::sid::shift_sid_origin({sid_def}, {name}_origin)".format( sid_def=sid_def, name=node.name, ) if domain_ndim != 3: gt_dims = [ f"gt::stencil::dim::{dim}" for dim in gtc_utils.dimension_flags_to_names(node.dimensions) ] if data_ndim: gt_dims += [ f"gt::integral_constant<int, {3 + dim}>" for dim in range(data_ndim) ] sid_def = "gt::sid::rename_numbered_dimensions<{gt_dims}>({sid_def})".format( gt_dims=", ".join(gt_dims), sid_def=sid_def) return sid_def
def test_make_args_data_from_gtir(backend_name, mode): backend_cls = backend_registry[backend_name] builder = StencilBuilder(stencil_def, backend=backend_cls).with_externals( {"MODE": mode}) args_data = make_args_data_from_gtir(builder.gtir_pipeline) assert set(args_data.unreferenced) == set(unreferenced_val[mode]) field_info_from_gtir = {( p.name, np.dtype(p.dtype.name.lower()), utils.dimension_flags_to_names(p.dimensions).upper(), p.data_dims, ) for p in builder.gtir.params if isinstance(p, gtir.FieldDecl)} field_info_from_args_data = {(name, d.dtype, "".join(d.axes), d.data_dims) for name, d in args_data.field_info.items() if name not in args_data.unreferenced} assert field_info_from_gtir == field_info_from_args_data param_info_from_gtir = {(p.name, np.dtype(p.dtype.name.lower())) for p in builder.gtir.params if isinstance(p, gtir.ScalarDecl)} param_info_from_args_data = { (name, d.dtype) for name, d in args_data.parameter_info.items() if name not in args_data.unreferenced } assert param_info_from_gtir == param_info_from_args_data for name, field_info in args_data.field_info.items(): if name == "out": access = AccessKind.WRITE elif name in field_info_val[mode]: access = AccessKind.READ else: access = AccessKind.NONE assert field_info.access == access for name, param_info in args_data.parameter_info.items(): if name in parameter_info_val[mode]: access = AccessKind.READ else: access = AccessKind.NONE assert param_info.access == access
def pybuffer_to_sid( *, name: str, ctype: str, domain_dim_flags: Tuple[bool, bool, bool], data_ndim: int, stride_kind_index: int, backend: Backend, check_layout: bool = True, ): domain_ndim = domain_dim_flags.count(True) sid_ndim = domain_ndim + data_ndim as_sid = "as_cuda_sid" if backend.storage_info[ "device"] == "gpu" else "as_sid" sid_def = """gt::{as_sid}<{ctype}, {sid_ndim}, gt::integral_constant<int, {unique_index}>{unit_stride_dim}>({name})""".format( name=name, ctype=ctype, unique_index=stride_kind_index, sid_ndim=sid_ndim, as_sid=as_sid, unit_stride_dim= f", {_get_unit_stride_dim(backend, domain_dim_flags, data_ndim)}" if check_layout else "", ) sid_def = "gt::sid::shift_sid_origin({sid_def}, {name}_origin)".format( sid_def=sid_def, name=name, ) if domain_ndim != 3: gt_dims = [ f"gt::stencil::dim::{dim}" for dim in gtc_utils.dimension_flags_to_names(domain_dim_flags) ] if data_ndim: gt_dims += [ f"gt::integral_constant<int, {3 + dim}>" for dim in range(data_ndim) ] sid_def = "gt::sid::rename_numbered_dimensions<{gt_dims}>({sid_def})".format( gt_dims=", ".join(gt_dims), sid_def=sid_def) return sid_def
def visit_AssignStmt( self, node: AssignStmt, *, loop_order: LoopOrder, symtable: Dict[str, Any], **kwargs: Any, ) -> None: if not isinstance(symtable[node.left.name], self.decl_type): return None allowed_flags = self._allowed_flags(loop_order) flags = symtable[node.left.name].dimensions if flags not in allowed_flags: dims = dimension_flags_to_names(flags) raise ValueError( f"Not allowed to assign to {dims}-field `{node.left.name}` in {loop_order.name}." ) return None
def make_args_data_from_gtir(pipeline: GtirPipeline) -> ModuleData: data = ModuleData() node = pipeline.full() field_extents = compute_legacy_extents(node) write_fields = (node.iter_tree().if_isinstance( gtir.ParAssignStmt).getattr("left").if_isinstance( gtir.FieldAccess).getattr("name").to_list()) referenced_field_params = { param.name for param in node.params if isinstance(param, gtir.FieldDecl) } for name in referenced_field_params: data.field_info[name] = FieldInfo( access=AccessKind.READ_WRITE if name in write_fields else AccessKind.READ_ONLY, boundary=field_extents[name].to_boundary(), axes=list( dimension_flags_to_names( node.symtable_[name].dimensions).upper()), dtype=numpy.dtype(node.symtable_[name].dtype.name.lower()), ) referenced_scalar_params = set( node.param_names).difference(referenced_field_params) for name in referenced_scalar_params: data.parameter_info[name] = ParameterInfo( dtype=numpy.dtype(node.symtable_[name].dtype.name.lower())) unref_params = get_unused_params_from_gtir(pipeline) for param in unref_params: if isinstance(param, gtir.FieldDecl): data.field_info[param.name] = None elif isinstance(param, gtir.ScalarDecl): data.parameter_info[param.name] = None data.unreferenced = {param.name for param in unref_params} return data
def visit_AssignStmt( self, node: AssignStmt, *, loop_order: LoopOrder, symtable: Dict[str, Any], **kwargs: Any, ) -> None: decl = symtable.get(node.left.name, None) if decl is None: raise ValueError("Symbol {} not found.".format(node.left.name)) if not isinstance(decl, self.decl_type): return None allowed_flags = self._allowed_flags(loop_order) flags = decl.dimensions if flags not in allowed_flags: dims = dimension_flags_to_names(flags) raise ValueError( f"Not allowed to assign to {dims}-field `{node.left.name}` in {loop_order.name}." ) return None
def visit_FieldDecl(self, node: gtcpp.FieldDecl, **kwargs): assert "gt_backend_t" in kwargs if "external_arg" in kwargs: domain_ndim = node.dimensions.count(True) data_ndim = len(node.data_dims) sid_ndim = domain_ndim + data_ndim if kwargs["external_arg"]: return "py::buffer {name}, std::array<gt::uint_t,{sid_ndim}> {name}_origin".format( name=node.name, sid_ndim=sid_ndim, ) else: sid_def = """gt::as_{sid_type}<{dtype}, {sid_ndim}, gt::integral_constant<int, {unique_index}>>({name})""".format( sid_type="cuda_sid" if kwargs["gt_backend_t"] == "gpu" else "sid", name=node.name, dtype=self.visit(node.dtype), unique_index=self.unique_index(), sid_ndim=sid_ndim, ) if domain_ndim != 3: gt_dims = [ f"gt::stencil::dim::{dim}" for dim in gtc_utils.dimension_flags_to_names(node.dimensions) ] if data_ndim: gt_dims += [ f"gt::integral_constant<int, {3 + dim}>" for dim in range(data_ndim) ] sid_def = "gt::sid::rename_numbered_dimensions<{gt_dims}>({sid_def})".format( gt_dims=", ".join(gt_dims), sid_def=sid_def) return "gt::sid::shift_sid_origin({sid_def}, {name}_origin)".format( sid_def=sid_def, name=node.name, )
def make_args_data_from_gtir(pipeline: GtirPipeline, legacy=False) -> ModuleData: """ Compute module data containing information about stencil arguments from gtir. Use `legacy` parameter to ensure equality with values from :func:`make_args_data_from_iir`. """ data = ModuleData() node = pipeline.full() write_fields = (node.iter_tree().if_isinstance( gtir.ParAssignStmt).getattr("left").if_isinstance( gtir.FieldAccess).getattr("name").to_set()) read_fields: Set[str] = set() for expr in node.iter_tree().if_isinstance( gtir.ParAssignStmt).getattr("right"): read_fields |= expr.iter_tree().if_isinstance( gtir.FieldAccess).getattr("name").to_set() referenced_field_params = [ param.name for param in node.params if isinstance(param, gtir.FieldDecl) ] field_extents = compute_legacy_extents(node, mask_inwards=legacy) k_boundary = (compute_k_boundary(node) if not legacy else {v: (0, 0) for v in referenced_field_params}) for name in sorted(referenced_field_params): access = AccessKind.NONE if name in read_fields: access |= AccessKind.READ if name in write_fields: access |= AccessKind.WRITE boundary = Boundary(*field_extents[name].to_boundary()[0:2], k_boundary[name]) data.field_info[name] = FieldInfo( access=access, boundary=boundary, axes=tuple( dimension_flags_to_names( node.symtable_[name].dimensions).upper()), data_dims=tuple(node.symtable_[name].data_dims), dtype=numpy.dtype(node.symtable_[name].dtype.name.lower()), ) referenced_scalar_params = [ param.name for param in node.params if param.name not in referenced_field_params ] for name in sorted(referenced_scalar_params): data.parameter_info[name] = ParameterInfo( dtype=numpy.dtype(node.symtable_[name].dtype.name.lower())) unref_params = get_unused_params_from_gtir(pipeline) for param in sorted(unref_params, key=lambda decl: decl.name): if isinstance(param, gtir.FieldDecl): data.field_info[param.name] = None elif isinstance(param, gtir.ScalarDecl): data.parameter_info[param.name] = None data.unreferenced = [*sorted(param.name for param in unref_params)] return data