Add transform::Simplify

Performs basic peephole optimizations on the AST.

Use in transform::Hlsl.

Required to have the DecomposeStorageAccess transform operate correctly with the output of InlinePointerLets transform, specifically when declaring `let` pointer expressions to storage buffers.

Fixed: tint:221
Fixed: tint:492
Fixed: tint:829
Change-Id: I536390921a6492378104e9c3c100d9e761294a27
Reviewed-on: https://dawn-review.googlesource.com/c/tint/+/51921
Commit-Queue: Ben Clayton <bclayton@chromium.org>
Reviewed-by: David Neto <dneto@google.com>

src/BUILD.gn

@@ -519,6 +519,8 @@ libtint_source_set("libtint_core_all_src") {
     "transform/manager.h",
     "transform/renamer.cc",
     "transform/renamer.h",
+    "transform/simplify.cc",
+    "transform/simplify.h",
     "transform/single_entry_point.cc",
     "transform/single_entry_point.h",
     "transform/transform.cc",

src/CMakeLists.txt

@@ -282,6 +282,8 @@ set(TINT_LIB_SRCS
   transform/manager.h
   transform/renamer.cc
   transform/renamer.h
+  transform/simplify.cc
+  transform/simplify.h
   transform/single_entry_point.cc
   transform/single_entry_point.h
   transform/transform.cc
@@ -810,7 +812,8 @@ if(${TINT_BUILD_TESTS})
       transform/first_index_offset_test.cc
       transform/inline_pointer_lets_test.cc
       transform/renamer_test.cc
-      transform/single_entry_point.cc
+      transform/simplify_test.cc
+      transform/single_entry_point_test.cc
       transform/test_helper.h
       transform/var_for_dynamic_index_test.cc
       transform/vertex_pulling_test.cc

src/transform/hlsl.cc

@@ -29,6 +29,7 @@
 #include "src/transform/external_texture_transform.h"
 #include "src/transform/inline_pointer_lets.h"
 #include "src/transform/manager.h"
+#include "src/transform/simplify.h"
 
 namespace tint {
 namespace transform {
@@ -39,10 +40,16 @@ Hlsl::~Hlsl() = default;
 Output Hlsl::Run(const Program* in, const DataMap& data) {
   Manager manager;
   manager.Add<CanonicalizeEntryPointIO>();
+  manager.Add<InlinePointerLets>();
+  // Simplify cleans up messy `*(&(expr))` expressions from InlinePointerLets.
+  manager.Add<Simplify>();
+  // DecomposeStorageAccess must come after InlinePointerLets as we cannot take
+  // the address of calls to DecomposeStorageAccess::Intrinsic. Must also come
+  // after Simplify, as we need to fold away the address-of and defers of
+  // `*(&(intrinsic_load()))` expressions.
   manager.Add<DecomposeStorageAccess>();
   manager.Add<CalculateArrayLength>();
   manager.Add<ExternalTextureTransform>();
-  manager.Add<InlinePointerLets>();
   auto out = manager.Run(in, data);
   if (!out.program.IsValid()) {
     return out;

src/transform/simplify.cc

@@ -0,0 +1,63 @@
+// Copyright 2021 The Tint Authors.
+//
+// Licensed 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.
+
+#include "src/transform/simplify.h"
+
+#include <memory>
+#include <unordered_map>
+#include <utility>
+
+#include "src/program_builder.h"
+#include "src/sem/block_statement.h"
+#include "src/sem/function.h"
+#include "src/sem/statement.h"
+#include "src/sem/variable.h"
+#include "src/utils/scoped_assignment.h"
+
+namespace tint {
+namespace transform {
+
+Simplify::Simplify() = default;
+
+Simplify::~Simplify() = default;
+
+Output Simplify::Run(const Program* in, const DataMap&) {
+  ProgramBuilder out;
+  CloneContext ctx(&out, in);
+
+  ctx.ReplaceAll([&](ast::Expression* expr) -> ast::Expression* {
+    if (auto* outer = expr->As<ast::UnaryOpExpression>()) {
+      if (auto* inner = outer->expr()->As<ast::UnaryOpExpression>()) {
+        if (outer->op() == ast::UnaryOp::kAddressOf &&
+            inner->op() == ast::UnaryOp::kIndirection) {
+          // &(*(expr)) => expr
+          return ctx.Clone(inner->expr());
+        }
+        if (outer->op() == ast::UnaryOp::kIndirection &&
+            inner->op() == ast::UnaryOp::kAddressOf) {
+          // *(&(expr)) => expr
+          return ctx.Clone(inner->expr());
+        }
+      }
+    }
+    return nullptr;
+  });
+
+  ctx.Clone();
+
+  return Output(Program(std::move(out)));
+}
+
+}  // namespace transform
+}  // namespace tint

src/transform/simplify.h

@@ -0,0 +1,49 @@
+// Copyright 2021 The Tint Authors.
+//
+// Licensed 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.
+
+#ifndef SRC_TRANSFORM_SIMPLIFY_H_
+#define SRC_TRANSFORM_SIMPLIFY_H_
+
+#include <string>
+#include <unordered_map>
+
+#include "src/transform/transform.h"
+
+namespace tint {
+namespace transform {
+
+/// Simplify is a peephole optimizer Transform that simplifies ASTs by removing
+/// unnecessary operations.
+/// Simplify currently optimizes the following:
+/// `&(*(expr))` => `expr`
+/// `*(&(expr))` => `expr`
+class Simplify : public Transform {
+ public:
+  /// Constructor
+  Simplify();
+
+  /// Destructor
+  ~Simplify() override;
+
+  /// Runs the transform on `program`, returning the transformation result.
+  /// @param program the source program to transform
+  /// @param data optional extra transform-specific input data
+  /// @returns the transformation result
+  Output Run(const Program* program, const DataMap& data = {}) override;
+};
+
+}  // namespace transform
+}  // namespace tint
+
+#endif  // SRC_TRANSFORM_SIMPLIFY_H_

src/transform/simplify_test.cc

@@ -0,0 +1,112 @@
+// Copyright 2021 The Tint Authors.
+//
+// Licensed 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.
+
+#include "src/transform/simplify.h"
+
+#include <memory>
+#include <utility>
+#include <vector>
+
+#include "src/transform/test_helper.h"
+
+namespace tint {
+namespace transform {
+namespace {
+
+using SimplifyTest = TransformTest;
+
+TEST_F(SimplifyTest, EmptyModule) {
+  auto* src = "";
+  auto* expect = "";
+
+  auto got = Run<Simplify>(src);
+
+  EXPECT_EQ(expect, str(got));
+}
+
+TEST_F(SimplifyTest, AddressOfDeref) {
+  auto* src = R"(
+fn f() {
+  var v : i32;
+  let p : ptr<function, i32> = &(v);
+  let x : ptr<function, i32> = &(*(p));
+  let y : ptr<function, i32> = &(*(&(*(p))));
+  let z : ptr<function, i32> = &(*(&(*(&(*(&(*(p))))))));
+}
+)";
+
+  auto* expect = R"(
+fn f() {
+  var v : i32;
+  let p : ptr<function, i32> = &(v);
+  let x : ptr<function, i32> = p;
+  let y : ptr<function, i32> = p;
+  let z : ptr<function, i32> = p;
+}
+)";
+
+  auto got = Run<Simplify>(src);
+
+  EXPECT_EQ(expect, str(got));
+}
+
+TEST_F(SimplifyTest, DerefAddressOf) {
+  auto* src = R"(
+fn f() {
+  var v : i32;
+  let x : i32 = *(&(v));
+  let y : i32 = *(&(*(&(v))));
+  let z : i32 = *(&(*(&(*(&(*(&(v))))))));
+}
+)";
+
+  auto* expect = R"(
+fn f() {
+  var v : i32;
+  let x : i32 = v;
+  let y : i32 = v;
+  let z : i32 = v;
+}
+)";
+
+  auto got = Run<Simplify>(src);
+
+  EXPECT_EQ(expect, str(got));
+}
+
+TEST_F(SimplifyTest, NoChange) {
+  auto* src = R"(
+fn f() {
+  var v : i32;
+  let p : ptr<function, i32> = &(v);
+  let x : i32 = *(p);
+}
+)";
+
+  auto* expect = R"(
+fn f() {
+  var v : i32;
+  let p : ptr<function, i32> = &(v);
+  let x : i32 = *(p);
+}
+)";
+
+  auto got = Run<Simplify>(src);
+
+  EXPECT_EQ(expect, str(got));
+}
+
+}  // namespace
+}  // namespace transform
+}  // namespace tint

test/BUILD.gn

@@ -304,6 +304,7 @@ tint_unittests_source_set("tint_unittests_core_src") {
     "../src/transform/first_index_offset_test.cc",
     "../src/transform/inline_pointer_lets_test.cc",
     "../src/transform/renamer_test.cc",
+    "../src/transform/simplify_test.cc",
     "../src/transform/single_entry_point_test.cc",
     "../src/transform/transform_test.cc",
     "../src/transform/var_for_dynamic_index_test.cc",

test/bug/tint/221.wgsl

@@ -0,0 +1,25 @@
+type Arr = array<u32, 50>;
+
+[[block]]
+struct Buf{
+  count : u32;
+  data : Arr;
+};
+
+[[group(0), binding (0)]] var<storage>  b : [[access(read)]] Buf;
+
+[[stage(compute)]]
+fn main() {
+  var i : u32 = 0u;
+  loop {
+    if (i >= b.count) { break; }
+    let p : ptr<storage, u32> = &b.data[i];
+    if ((i % 2u) == 0u) { continue; }
+    *p = 0u;              // Set odd elements of the array to 0
+    continuing {
+      *p = *p * 2u;        // Double the value in the array entry
+      i = i + 1u;
+    }
+  }
+  return;
+}

test/bug/tint/221.wgsl.expected.hlsl

@@ -0,0 +1,27 @@
+
+ByteAddressBuffer b : register(t0, space0);
+
+[numthreads(1, 1, 1)]
+void main() {
+  uint i = 0u;
+  while (true) {
+    if ((i >= b.Load(0u))) {
+      break;
+    }
+    const uint p_save = i;
+    if (((i % 2u) == 0u)) {
+      {
+        b.Store((4u + (4u * p_save)), asuint((b.Load((4u + (4u * p_save))) * 2u)));
+        i = (i + 1u);
+      }
+      continue;
+    }
+    b.Store((4u + (4u * p_save)), asuint(0u));
+    {
+      b.Store((4u + (4u * p_save)), asuint((b.Load((4u + (4u * p_save))) * 2u)));
+      i = (i + 1u);
+    }
+  }
+  return;
+}
+

test/bug/tint/221.wgsl.expected.msl

@@ -0,0 +1,33 @@
+#include <metal_stdlib>
+
+using namespace metal;
+struct Buf {
+  /* 0x0000 */ uint count;
+  /* 0x0004 */ uint data[50];
+};
+
+kernel void tint_symbol(const device Buf& b [[buffer(0)]]) {
+  uint i = 0u;
+  {
+    bool tint_msl_is_first_1 = true;
+    device uint* const p;
+    for(;;) {
+      if (!tint_msl_is_first_1) {
+        *(p) = (*(p) * 2u);
+        i = (i + 1u);
+      }
+      tint_msl_is_first_1 = false;
+
+      if ((i >= b.count)) {
+        break;
+      }
+      p = &(b.data[i]);
+      if (((i % 2u) == 0u)) {
+        continue;
+      }
+      *(p) = 0u;
+    }
+  }
+  return;
+}
+

test/bug/tint/221.wgsl.expected.spvasm

@@ -0,0 +1,78 @@
+; SPIR-V
+; Version: 1.3
+; Generator: Google Tint Compiler; 0
+; Bound: 44
+; Schema: 0
+               OpCapability Shader
+               OpMemoryModel Logical GLSL450
+               OpEntryPoint GLCompute %main "main"
+               OpExecutionMode %main LocalSize 1 1 1
+               OpName %Buf "Buf"
+               OpMemberName %Buf 0 "count"
+               OpMemberName %Buf 1 "data"
+               OpName %b "b"
+               OpName %main "main"
+               OpName %i "i"
+               OpDecorate %Buf Block
+               OpMemberDecorate %Buf 0 Offset 0
+               OpMemberDecorate %Buf 1 Offset 4
+               OpDecorate %_arr_uint_uint_50 ArrayStride 4
+               OpDecorate %b NonWritable
+               OpDecorate %b DescriptorSet 0
+               OpDecorate %b Binding 0
+       %uint = OpTypeInt 32 0
+    %uint_50 = OpConstant %uint 50
+%_arr_uint_uint_50 = OpTypeArray %uint %uint_50
+        %Buf = OpTypeStruct %uint %_arr_uint_uint_50
+%_ptr_StorageBuffer_Buf = OpTypePointer StorageBuffer %Buf
+          %b = OpVariable %_ptr_StorageBuffer_Buf StorageBuffer
+       %void = OpTypeVoid
+          %7 = OpTypeFunction %void
+     %uint_0 = OpConstant %uint 0
+%_ptr_Function_uint = OpTypePointer Function %uint
+         %14 = OpConstantNull %uint
+%_ptr_StorageBuffer_uint = OpTypePointer StorageBuffer %uint
+       %bool = OpTypeBool
+     %uint_1 = OpConstant %uint 1
+     %uint_2 = OpConstant %uint 2
+       %main = OpFunction %void None %7
+         %10 = OpLabel
+          %i = OpVariable %_ptr_Function_uint Function %14
+               OpStore %i %uint_0
+               OpBranch %15
+         %15 = OpLabel
+               OpLoopMerge %16 %17 None
+               OpBranch %18
+         %18 = OpLabel
+         %19 = OpLoad %uint %i
+         %21 = OpAccessChain %_ptr_StorageBuffer_uint %b %uint_0
+         %22 = OpLoad %uint %21
+         %23 = OpUGreaterThanEqual %bool %19 %22
+               OpSelectionMerge %25 None
+               OpBranchConditional %23 %26 %25
+         %26 = OpLabel
+               OpBranch %16
+         %25 = OpLabel
+         %29 = OpLoad %uint %i
+         %30 = OpAccessChain %_ptr_StorageBuffer_uint %b %uint_1 %29
+         %31 = OpLoad %uint %i
+         %33 = OpUMod %uint %31 %uint_2
+         %34 = OpIEqual %bool %33 %uint_0
+               OpSelectionMerge %35 None
+               OpBranchConditional %34 %36 %35
+         %36 = OpLabel
+               OpBranch %17
+         %35 = OpLabel
+               OpStore %30 %uint_0
+               OpBranch %17
+         %17 = OpLabel
+         %40 = OpLoad %uint %30
+         %41 = OpIMul %uint %40 %uint_2
+               OpStore %30 %41
+         %42 = OpLoad %uint %i
+         %43 = OpIAdd %uint %42 %uint_1
+               OpStore %i %43
+               OpBranch %15
+         %16 = OpLabel
+               OpReturn
+               OpFunctionEnd

test/bug/tint/221.wgsl.expected.wgsl

@@ -0,0 +1,30 @@
+type Arr = array<u32, 50>;
+
+[[block]]
+struct Buf {
+  count : u32;
+  data : Arr;
+};
+
+[[group(0), binding(0)]] var<storage> b : [[access(read)]] Buf;
+
+[[stage(compute)]]
+fn main() {
+  var i : u32 = 0u;
+  loop {
+    if ((i >= b.count)) {
+      break;
+    }
+    let p : ptr<storage, u32> = &(b.data[i]);
+    if (((i % 2u) == 0u)) {
+      continue;
+    }
+    *(p) = 0u;
+
+    continuing {
+      *(p) = (*(p) * 2u);
+      i = (i + 1u);
+    }
+  }
+  return;
+}

test/bug/tint/492.wgsl

@@ -0,0 +1,7 @@
+[[block]] struct S { a : i32; };
+[[group(0), binding(0)]] var<storage> buf : [[access(read_write)]] S;
+
+[[stage(compute)]] fn main() {
+  let p : ptr<storage, i32> = &buf.a;
+  *p = 12;
+}

test/bug/tint/492.wgsl.expected.hlsl

@@ -0,0 +1,9 @@
+
+RWByteAddressBuffer buf : register(u0, space0);
+
+[numthreads(1, 1, 1)]
+void main() {
+  buf.Store(0u, asuint(12));
+  return;
+}
+

test/bug/tint/492.wgsl.expected.msl

@@ -0,0 +1,13 @@
+#include <metal_stdlib>
+
+using namespace metal;
+struct S {
+  /* 0x0000 */ int a;
+};
+
+kernel void tint_symbol(device S& buf [[buffer(0)]]) {
+  device int* const p = &(buf.a);
+  *(p) = 12;
+  return;
+}
+

test/bug/tint/492.wgsl.expected.spvasm

@@ -0,0 +1,33 @@
+; SPIR-V
+; Version: 1.3
+; Generator: Google Tint Compiler; 0
+; Bound: 16
+; Schema: 0
+               OpCapability Shader
+               OpMemoryModel Logical GLSL450
+               OpEntryPoint GLCompute %main "main"
+               OpExecutionMode %main LocalSize 1 1 1
+               OpName %S "S"
+               OpMemberName %S 0 "a"
+               OpName %buf "buf"
+               OpName %main "main"
+               OpDecorate %S Block
+               OpMemberDecorate %S 0 Offset 0
+               OpDecorate %buf DescriptorSet 0
+               OpDecorate %buf Binding 0
+        %int = OpTypeInt 32 1
+          %S = OpTypeStruct %int
+%_ptr_StorageBuffer_S = OpTypePointer StorageBuffer %S
+        %buf = OpVariable %_ptr_StorageBuffer_S StorageBuffer
+       %void = OpTypeVoid
+          %5 = OpTypeFunction %void
+       %uint = OpTypeInt 32 0
+     %uint_0 = OpConstant %uint 0
+%_ptr_StorageBuffer_int = OpTypePointer StorageBuffer %int
+     %int_12 = OpConstant %int 12
+       %main = OpFunction %void None %5
+          %8 = OpLabel
+         %13 = OpAccessChain %_ptr_StorageBuffer_int %buf %uint_0
+               OpStore %13 %int_12
+               OpReturn
+               OpFunctionEnd

test/bug/tint/492.wgsl.expected.wgsl

@@ -0,0 +1,12 @@
+[[block]]
+struct S {
+  a : i32;
+};
+
+[[group(0), binding(0)]] var<storage> buf : [[access(read_write)]] S;
+
+[[stage(compute)]]
+fn main() {
+  let p : ptr<storage, i32> = &(buf.a);
+  *(p) = 12;
+}

test/ptr_ref/load/local/ptr_storage.wgsl.expected.hlsl

@@ -1 +1,9 @@
-SKIP: error: cannot take the address of expression
+
+RWByteAddressBuffer v : register(u0, space0);
+
+[numthreads(1, 1, 1)]
+void main() {
+  const int use = (asint(v.Load(0u)) + 1);
+  return;
+}
+