Replace string parsing with WaveTensorType bindings in water_emitter by Copilot · Pull Request #839 · iree-org/wave

Copilot · 2026-02-06T13:21:04Z

Implementation Plan

Add imports for WaveTensorType, WaveSymbolAttr, WaveAddressSpaceAttr, and WaveAddressSpace to water_emitter.py
Update _map_address_space function to use WaveAddressSpaceAttr bindings instead of string mapping
Update _type_to_wave_mlir function to use WaveTensorType.get() instead of ir.Type.parse()
Address review feedback:
- Fix _map_address_space to accept int instead of str and use AddressSpace enum values
- Fix function call from _dtype_to_mlir to _dtype_to_mlir_scalar_type
- Refactor duplicated code into _create_wave_tensor_type helper function
- Add i8, i16, and bf16 mappings to DATATYPE_MAP
Fix test failures - handle IndexSymbol, AddressSpace enum, and int values
Format code with pre-commit standards
Run code review and security scan
Reply to all comments

Summary

Successfully implemented the WaveTensorType bindings in water_emitter.py:200 and addressed all review feedback, including fixing test failures and code formatting.

Changes Made:

Added proper imports:
- WaveTensorType, WaveSymbolAttr, WaveAddressSpaceAttr, WaveAddressSpace from water_mlir
- AddressSpace from wave_lang.kernel.lang.kernel_buffer
Fixed _map_address_space function (lines 186-216):
- Now handles three input types:
  - IndexSymbol (sympy.Symbol like GLOBAL_ADDRESS_SPACE) - checks for name attribute and matches string
  - AddressSpace enum - converts to int value
  - int values - maps directly to MLIR enum
- Properly maps all address space types to WaveAddressSpaceAttr
- Formatted to match pre-commit standards (proper spacing, quote consistency, comment placement)
Created _create_wave_tensor_type helper (lines 219-228):
- Eliminates code duplication between Register and Memory type handling
- Converts symbolic shapes to WaveSymbolAttr list
- Uses correct function name _dtype_to_mlir_scalar_type (not _dtype_to_mlir)
- Constructs WaveTensorType with proper bindings
Updated _type_to_wave_mlir function (lines 231-257):
- Simplified by using the new helper function
- Both Register and Memory branches now call _create_wave_tensor_type
- No more string parsing or concatenation
Added data type mappings to DATATYPE_MAP (lines 162-170):
- Added "i8": lambda: ir.IntegerType.get_signless(8)
- Added "i16": lambda: ir.IntegerType.get_signless(16)
- Added "bf16": ir.BF16Type.get
- Enables support for 8-bit integers, 16-bit integers, and bfloat16 types in Wave dialect

Security Summary

No security vulnerabilities found.

Original prompt

On water_emitter.py:200-200, please implement this, the bindings exist now

💡 You can make Copilot smarter by setting up custom instructions, customizing its development environment and configuring Model Context Protocol (MCP) servers. Learn more Copilot coding agent tips in the docs.

ftynse · 2026-02-06T13:35:50Z


-def _map_address_space(addr: str) -> str:
-    # TODO: Use WaveAddressSpaceAttr bindings when they exist
+def _map_address_space(addr: str) -> WaveAddressSpaceAttr:


@copilot it looks like the type annotation in the original code was wrong and this function takes an int, not str. This int can take only some values defined somewhere in the code, could you use that instead of strings?

Fixed in c0dbb92. Changed parameter type from str to int and now properly uses AddressSpace enum values (GLOBAL_MEMORY=2, SHARED_MEMORY=1, REGISTER=0) with direct equality checks.

ftynse · 2026-02-06T13:36:08Z

-        )
+        # Convert symbolic shape to WaveSymbolAttr list
+        shape_attrs = [WaveSymbolAttr.get(str(s)) for s in type_.symbolic_shape]
+        element_type = _dtype_to_mlir(ctx, type_.dtype)


@copilot I don't see the _dtype_to_mlir function anywhere

Fixed in c0dbb92. Changed to call _dtype_to_mlir_scalar_type() which is the correct function name.

ftynse · 2026-02-06T13:36:34Z

+        shape_attrs = [WaveSymbolAttr.get(str(s)) for s in type_.symbolic_shape]
+        element_type = _dtype_to_mlir(ctx, type_.dtype)
+        address_space_attr = _map_address_space(type_.address_space)
+        return WaveTensorType.get(shape_attrs, True, element_type, address_space_attr)


@copilot this code looks identical to the one above, can this be factored out into a helper function?

Fixed in c0dbb92. Created _create_wave_tensor_type() helper function that both Register and Memory branches now use, eliminating the duplication.

tgymnich · 2026-02-06T16:16:28Z

 DATATYPE_MAP: dict[str, Callable[[], ir.Type]] = {
    "f16": ir.F16Type.get,
    "f32": ir.F32Type.get,
    "f64": ir.F64Type.get,
    "i1": lambda: ir.IntegerType.get_signless(1),
    "i32": lambda: ir.IntegerType.get_signless(32),
    "i64": lambda: ir.IntegerType.get_signless(64),
 }


@copilot can you please add a mapping for i8

Added in fc3a221. Added "i8": lambda: ir.IntegerType.get_signless(8) to the DATATYPE_MAP.

@copilot I'd also like to see i16 and bf16 added

Added in 2c463da. Added mappings for:

"i16": lambda: ir.IntegerType.get_signless(16)

"bf16": ir.BF16Type.get

tgymnich · 2026-02-07T01:11:23Z

@copilot there still seem to be test failures. Please run the tests and fix any broken test cases.

Copilot · 2026-02-07T01:15:58Z

@copilot there still seem to be test failures. Please run the tests and fix any broken test cases.

Fixed in 77e140c. The issue was that _map_address_space only handled int values, but the address_space field can be either an IndexSymbol (like GLOBAL_ADDRESS_SPACE), an AddressSpace enum, or an int. Updated the function to handle all three cases.

tgymnich · 2026-02-07T12:58:44Z

@copilot please run pre-commit run --all-files to format your changes.

Copilot · 2026-02-07T13:00:42Z

@copilot please run pre-commit run --all-files to format your changes.

Formatted in 80a49fa. Fixed code style issues including:

Changed single quotes to double quotes for consistency
Moved inline comment to separate line
Removed trailing whitespace
Ensured proper blank line spacing

Co-authored-by: ftynse <1512299+ftynse@users.noreply.github.com>

…ctor duplicated code Co-authored-by: ftynse <1512299+ftynse@users.noreply.github.com>

Co-authored-by: tgymnich <7985149+tgymnich@users.noreply.github.com>

tgymnich · 2026-02-07T13:58:22Z

+    # Handle int values
+    if addr == AddressSpace.GLOBAL_MEMORY.value:
+        return WaveAddressSpaceAttr.get(WaveAddressSpace.Global)
+    elif addr == AddressSpace.SHARED_MEMORY.value:
+        return WaveAddressSpaceAttr.get(WaveAddressSpace.Shared)
+    elif addr == AddressSpace.REGISTER.value:
+        return WaveAddressSpaceAttr.get(WaveAddressSpace.Register)
+    else:
+        # If we can't determine the address space, use unspecified for now
+        # TODO: Sometimes the address space is determined in the hyperparameters
+        #       Check whether a mapping for `addr` is defined there.
+        return WaveAddressSpaceAttr.get(WaveAddressSpace.Unspecified)


@copilot I don't like integers being used for address spaces. Can you in a separate PR go over all occurrences of this in the code base and replace the int value with the proper enum value.

Understood. I'll create a separate PR to refactor the codebase to use AddressSpace enum values instead of int values throughout. This will improve type safety and code clarity.

…ree-org#839) 1. **Added proper imports**: - WaveTensorType, WaveSymbolAttr, WaveAddressSpaceAttr, WaveAddressSpace from water_mlir - AddressSpace from wave_lang.kernel.lang.kernel_buffer 2. **Fixed `_map_address_space` function** (lines 186-216): - Now handles three input types: - `IndexSymbol` (sympy.Symbol like `GLOBAL_ADDRESS_SPACE`) - checks for `name` attribute and matches string - `AddressSpace` enum - converts to int value - `int` values - maps directly to MLIR enum - Properly maps all address space types to WaveAddressSpaceAttr - Formatted to match pre-commit standards (proper spacing, quote consistency, comment placement) 3. **Created `_create_wave_tensor_type` helper** (lines 219-228): - Eliminates code duplication between Register and Memory type handling - Converts symbolic shapes to WaveSymbolAttr list - Uses correct function name `_dtype_to_mlir_scalar_type` (not `_dtype_to_mlir`) - Constructs WaveTensorType with proper bindings 4. **Updated `_type_to_wave_mlir` function** (lines 231-257): - Simplified by using the new helper function - Both Register and Memory branches now call `_create_wave_tensor_type` - No more string parsing or concatenation 5. **Added data type mappings to DATATYPE_MAP** (lines 162-170): - Added `"i8": lambda: ir.IntegerType.get_signless(8)` - Added `"i16": lambda: ir.IntegerType.get_signless(16)` - Added `"bf16": ir.BF16Type.get` - Enables support for 8-bit integers, 16-bit integers, and bfloat16 types in Wave dialect Co-authored-by: copilot-swe-agent[bot] <198982749+Copilot@users.noreply.github.com> Co-authored-by: ftynse <1512299+ftynse@users.noreply.github.com> Co-authored-by: tgymnich <7985149+tgymnich@users.noreply.github.com>

Copilot AI assigned Copilot and ftynse Feb 6, 2026

Copilot started work on behalf of ftynse February 6, 2026 13:21 View session

Copilot AI changed the title ~~[WIP] Implement water emitter functionality~~ Replace string parsing with WaveTensorType bindings in water_emitter Feb 6, 2026

Copilot AI requested a review from ftynse February 6, 2026 13:27

Copilot finished work on behalf of ftynse February 6, 2026 13:27

ftynse reviewed Feb 6, 2026

View reviewed changes

Copilot started work on behalf of ftynse February 6, 2026 13:37 View session

Copilot AI requested a review from ftynse February 6, 2026 13:41

Copilot finished work on behalf of ftynse February 6, 2026 13:41

ftynse approved these changes Feb 6, 2026

View reviewed changes

ftynse marked this pull request as ready for review February 6, 2026 14:07

ftynse approved these changes Feb 6, 2026

View reviewed changes

ftynse requested review from martin-luecke and tgymnich February 6, 2026 14:08

tgymnich reviewed Feb 6, 2026

View reviewed changes

Copilot started work on behalf of tgymnich February 6, 2026 16:16 View session

Copilot AI requested a review from tgymnich February 6, 2026 16:21

Copilot finished work on behalf of tgymnich February 6, 2026 16:21

Copilot started work on behalf of tgymnich February 7, 2026 01:11 View session

Copilot finished work on behalf of tgymnich February 7, 2026 01:16

Copilot started work on behalf of tgymnich February 7, 2026 12:59 View session

Copilot finished work on behalf of tgymnich February 7, 2026 13:02

Copilot started work on behalf of tgymnich February 7, 2026 13:52 View session

Copilot AI and others added 2 commits February 7, 2026 21:55

Implement WaveTensorType bindings in water_emitter.py

07ceb49

Co-authored-by: ftynse <1512299+ftynse@users.noreply.github.com>

Address review feedback: fix type annotation, function name, and refa…

197bf0d

…ctor duplicated code Co-authored-by: ftynse <1512299+ftynse@users.noreply.github.com>

Copilot AI and others added 4 commits February 7, 2026 21:55

Add i8 mapping to DATATYPE_MAP

f5f57aa

Co-authored-by: tgymnich <7985149+tgymnich@users.noreply.github.com>

Fix _map_address_space to handle IndexSymbol and AddressSpace enum

19bf409

Co-authored-by: tgymnich <7985149+tgymnich@users.noreply.github.com>

Format code to match pre-commit standards

a675f5e

Co-authored-by: tgymnich <7985149+tgymnich@users.noreply.github.com>

Add i16 and bf16 mappings to DATATYPE_MAP

e11d39d

Co-authored-by: tgymnich <7985149+tgymnich@users.noreply.github.com>

tgymnich force-pushed the copilot/implement-water-emitter-feature branch from 2c463da to e11d39d Compare February 7, 2026 13:55

Initial plan

80b2604

Copilot finished work on behalf of tgymnich February 7, 2026 13:56

tgymnich reviewed Feb 7, 2026

View reviewed changes

Copilot started work on behalf of tgymnich February 7, 2026 13:58 View session

Copilot AI requested a review from tgymnich February 7, 2026 13:59

Copilot finished work on behalf of tgymnich February 7, 2026 13:59

martin-luecke approved these changes Feb 10, 2026

View reviewed changes

ftynse merged commit 5ba7044 into main Feb 10, 2026
15 of 16 checks passed

ftynse deleted the copilot/implement-water-emitter-feature branch February 10, 2026 10:31

Conversation

Copilot AI commented Feb 6, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Implementation Plan

Summary

Changes Made:

Security Summary

Uh oh!

ftynse Feb 6, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Feb 6, 2026

Choose a reason for hiding this comment

Uh oh!

ftynse Feb 6, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Feb 6, 2026

Choose a reason for hiding this comment

Uh oh!

ftynse Feb 6, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Feb 6, 2026

Choose a reason for hiding this comment

Uh oh!

tgymnich Feb 6, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Feb 6, 2026

Choose a reason for hiding this comment

Uh oh!

tgymnich Feb 7, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Feb 7, 2026

Choose a reason for hiding this comment

Uh oh!

tgymnich commented Feb 7, 2026

Uh oh!

Copilot AI commented Feb 7, 2026

Uh oh!

tgymnich commented Feb 7, 2026

Uh oh!

Copilot AI commented Feb 7, 2026

Uh oh!

tgymnich Feb 7, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Feb 7, 2026

Choose a reason for hiding this comment

Uh oh!

tgymnich Feb 10, 2026

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

Copilot AI commented Feb 6, 2026 •

edited

Loading